Why we built AgentGrad.

For the past year, our agent-optimization engagements have run on an internal library we had not talked about publicly. Client work kept following the same arc: an agent stuck below the accuracy bar, a prompt full of hand-patched rules, and a team with no way to tell which change helped. Each time, the fix was the same loop: define the metric, trace failures backward through the runtime, update the responsible components, validate on held-out examples. After enough engagements ended with that loop producing the result, we gave it a name. This post is about why AgentGrad exists and what it does differently.

The patch log problem

Wiring an agent together is a solved problem. What is not solved is what happens after the demo: accuracy plateaus below the bar the product needs, and the only tool most teams reach for is editing prompts by hand.

Hand-tuning fails in a specific, predictable way. When a multi-step agent gets something wrong, someone has to guess which component caused it: the planner prompt, a tool description, a bad few-shot example, the temperature. The guess usually lands on the last prompt someone touched. Then the fix is appended as one more rule, and the prompt grows into a history of incidents: every line a patch for one failure, no record of which patches still earn their place, and no way to know whether today's fix quietly broke last week's behavior.

Prompt engineering, done this way, is gradient descent performed by a human: one example at a time, no batching, no validation, no convergence guarantee. Our position is simple: if that is the loop, automate it properly.

The ideas we built on

We did not start from zero, and this is not a takedown of the work we built on. DSPy established that a language-model program can be compiled against a metric, with signatures, modules, and optimizers as first-class concepts.¹ TextGrad established automatic differentiation through text: an LLM can generate natural-language feedback and propagate it backward through a computation graph to improve its variables. It also adopted a PyTorch-like interface.² GEPA showed how reflection over complete trajectories, combined with evolutionary search and a Pareto frontier of candidates, can turn a small number of rollouts into strong prompt updates.³ These are good ideas. AgentGrad exists because we wanted to operate them as one training system.

The distinction matters. AgentGrad is not novel merely because it uses language as a gradient, borrows names from PyTorch, or optimizes prompts against a metric. The bet is that the complete training model matters: arbitrary runtime graphs, multiple parameter types, batch-level updates, contract-preserving steps, inspectable state, held-out retention, and tool-heavy pipelines as a first-class case.

Teams that take textual optimization into production report why those controls matter. Snowflake's engineering team, optimizing LLM judges, found that out-of-the-box DSPy few-shot optimization did not generalize well to new inputs, and that TextGrad's gradients across their training data were not producing consistent improvement; even choosing the loss was a challenge.⁴ Those findings are not reasons to abandon textual gradients. They are reasons to design the surrounding training system to detect and reject regressions.

Train the runtime like a model

AgentGrad's design bet is that the right abstractions for this problem already exist, and every ML engineer already knows them. A pipeline node is a Module. Anything worth learning is a Parameter: a prompt, a tool definition, a set of few-shot examples, a sampling parameter. Running the pipeline is the forward pass, and it builds a computation graph implicitly, exactly the way PyTorch does. Your metric is the loss: an LLM judge, exact match, schema validity, tool-call correctness. None of it needs to be differentiable, because the gradient is language.

The mapping is literal

The PyTorch parallel is structural rather than decorative. The same concepts have the same jobs, so the framework feels familiar before anyone learns an AgentGrad-specific abstraction.

A complete training loop

This is a condensed but representative loop: two learnable stages, a batched forward pass, one aggregate loss, one backward pass, and one optimizer step. The call chain builds the graph implicitly.

1class RewriteQA(Module):2    # two learnable system prompts, one per stage3    def __init__(self, engine):4        super().__init__()5        self.rewriter = BlackboxLLM("rewrite the query to be searchable.", engine)6        self.answerer = BlackboxLLM("answer concisely.", engine)78    def forward(self, question):9        # the graph builds itself from the call chain10        return self.answerer(self.rewriter(question))1112engine    = get_llm_provider("gpt-4.1")13model     = RewriteQA(engine)14criterion = LLMJudgeLoss(backward_engine=engine,15                         criteria="correct language and style")16optimizer = TextualGradientDescent(model.parameters(), engine)1718for epoch in range(num_epochs):19    optimizer.zero_grad()20    preds = model.forward_batch({"question": questions}, max_workers=8)21    loss  = AggregateLoss(criterion.forward_batch(questions, labels, preds))22    if loss.num_failures == 0:23        break                            # the whole batch passed24    loss.backward(backward_engine=engine, max_workers=8)25    optimizer.step()                     # fit the pattern, not the noise

What backward actually does

When a loss records a failure, backward() walks the graph in reverse and generates targeted feedback for every learnable parameter, with an LLM deciding, per input, whether feedback is relevant, how it should transform as it passes through a node, and where it should stop. Failure attribution, in other words, is part of the framework rather than a guess made by whoever is on call. And because the parallel to PyTorch is structural, not cosmetic, any forward computation you can express (branching pipelines, multi-step reasoning, tool loops) is optimizable without a DSL and without re-architecting the agent to fit the optimizer.

The feedback is not a numerical derivative and does not pretend to be one. It is a structured diagnosis: what failed, which upstream decision contributed, what should change, and what must remain invariant. As it crosses a node, the backward engine decides whether that diagnosis is relevant to the node's inputs, transforms it into parameter-specific feedback, or stops it. The graph provides the path; the language model performs the attribution.

The controls that make it production-grade

Gradients accumulate before anything changes. The optimizer collects feedback across a batch of failures, extracts the common failure mode, and discards one-off noise, the textual analogue of large-batch SGD. This is the control that separates "learns the pattern" from "appends another brittle exception to the prompt."

Steps preserve your contracts. Constrained optimization guarantees that a learned prompt keeps what production depends on: required placeholders, input/output schemas, mandated structure. The optimizer improves a template without ever breaking the code that fills it.

Prompts project onto the target model. Different model families want differently shaped prompts. AgentGrad implements the equivalent of projected gradient descent: after each step, learned prompts are projected onto model-family guidelines or your own constraints. Switching providers becomes re-projection and re-optimization, not a rewrite.

One prompt is not forced to cover every mode. When an input does not fit any learned prompt, the optimizer can detect the new mode and learn a specialized prompt for it: a mixture of prompts. We think of this as skill discovery: the pipeline finds the distinct competencies a workload demands and writes one prompt per competency, instead of straining a single template across all of them.

State stays inspectable. Parameters, gradients, optimizer state, and checkpoints can be inspected and serialized. A team can see why a prompt changed, compare it with the prior version, reproduce the run that produced it, and deploy it through the same review path as a model change.

Underneath all five sits the retention rule our tool-call optimization work runs on: a candidate update is kept only if it improves held-out and regression examples. The newest prompt is never assumed to be the best one.

Where AgentGrad fits

DSPy is a broad framework for programming and compiling language-model systems. GEPA is a strong reflective optimizer that evolves candidate texts from complete trajectories. TextGrad is the clearest foundation for automatic differentiation through text and graph-shaped systems. AgentGrad borrows from all three, but makes a different unit of abstraction primary: the training loop around an existing agent runtime.

That means prompts are not the only learnable object, one optimizer is not the only permitted search strategy, and a pre-defined module vocabulary does not constrain the forward pass. The same loop can optimize a system prompt, the argument contract for one tool, a set of few-shot examples, or a mixture of specialized skills, while retaining only changes that improve the end-to-end metric. For a single prompt, that machinery may be more than you need. For a branching production agent with several interacting failure surfaces, it is the point.

What it has done so far

AgentGrad is not a research prototype; it has been the engine in engagements and benchmark work for the past year.

• 62% → 91% tool-call argument accuracy. A legal-AI client engagement used this loop end to end: per-tool argument builders as modules, field-level failures traced back to the responsible prompt sections, updates retained only when held-out accuracy improved.
• Competition math. On AIME 2025, mixture-of-prompts skill discovery raised GPT-4.1 mini from 48% to 64% accuracy, with GPT-5 as the optimization teacher. The agent learned which solving skills it was missing and wrote them into its own prompts.
• Closing an architecture gap. Prompt optimization alone brought DiffusionGemma 26B to parity with its autoregressive counterpart on AIME 2026: 70.8% → 75.2% avg@4 against the AR twin's 75.0%, from 20 training samples and zero weight updates, while keeping the diffusion model's 3× speed advantage.
• Head-to-head against GEPA. On the same AIME 2025 task, AgentGrad reached 64% against GEPA's 57%. A full comparison writeup is coming.
• Distillation without fine-tuning. A frontier model diagnoses what a smaller model's reasoning is missing; the smaller model updates its own prompts and skills. Frontier-level behavior at small-model cost, with no training infrastructure and no weight access.

What the abstraction does not solve

The loss is still the ceiling. An exact-match metric can reward the wrong behavior; an LLM judge can be noisy; synthetic traffic can miss the edge case that matters in production. AgentGrad makes the optimization loop systematic, but it cannot rescue a metric or dataset that does not represent the product.

Textual gradients are diagnoses and proposals, not mathematical guarantees. They can attribute a failure to the wrong component or suggest a plausible change that does not generalize. This is why updates accumulate across a batch, why contracts constrain the step, and why held-out and regression examples decide what survives. The evaluator, not the optimizer's confidence, has the final word.

The loop also costs model calls. Simulation, backward attribution, candidate generation, and validation are intended as offline optimization work, not latency added to every production request. In practice the trade is more calls during training for fewer failures, cheaper models, and less manual patching after deployment.

What's next

Today AgentGrad optimizes every parameter of a static pipeline: values can change, but the graph connecting them is defined by the engineer. The next step is making the graph itself learnable. Instead of only improving the prompts and tools inside a runtime, the optimizer could propose a new branch, split one overloaded step into specialists, remove a redundant call, or choose which subgraph should handle a newly discovered mode. It is the pipeline analogue of neural architecture search, with the same held-out metric deciding which structure survives.

AgentGrad powers our enterprise engagements today. If you have an agent metric that is stuck, the fastest way to see the loop on your own pipeline is the AgentGrad product page. Bring one failing metric, and we will scope the optimization in a 30-minute call. Early access to the library is open there as well.