16 Trajectory Fine-Tuning
Purpose
Agent stack focus for this chapter.
Why can training on a useful recovery record teach the model to repeat the mistake instead of the correction?
While prompt engineering and runtime guardrails guide an off-the-shelf model, fine-tuning modifies the neural network weights directly, embedding specialized decision policies and tool calling conventions into model parameters. However, training on multi-turn trajectories requires strict separation between agent proposals and external environment observations, so that imitation loss applies only to selected policy responses while environment feedback remains conditioning context. Naive fine-tuning can teach a policy to fabricate its own tool results or overfit to static workspace layouts, degrading general reasoning and error recovery. Systems engineering governs how datasets are partitioned, how loss masks are applied, and how updated checkpoints are validated against regression suites before deployment. An update is worth releasing only when closed-loop task acceptance improves under the same runtime that still bounds every proposal the model makes, however far its training loss has fallen.
Learning Objectives
- Assign agent behavior to weights, context, or runtime, and decide when a failure warrants a weight change
- Construct model inputs that preserve intended conditioning and price per-invocation serialization in training tokens
- Map admitted response spans to causal token targets, observation masks, and termination policy
- Compare loss weighting and construct batches without unintended cross-example dependencies
- Calculate imitation regret bounds on reference and learner states to decide when on-policy labeling pays
- Calculate preference margins and gradient weights, and distinguish what a comparison measures from task acceptance
- Select behavior sources and parameter artifacts using explicit resource and version accounting
- Evaluate a versioned policy for release on closed-loop acceptance, cost per accepted task, and schema transfer, while preserving runtime enforcement
16.1 Recovery Records as Training Data
Fine-tuning reshapes each model call but leaves every exposure to the runtime.
Consider a recovery demonstration admitted by the curation pipeline for the repair agent (Four lighthouse workloads) into repair_demonstrations_v1. During the recorded session, the agent generated a malformed edit proposal that the runtime validator refused prior to tool dispatch. In response to the failure notification, the model inspected the tool schema, generated a corrected edit, and produced candidate patch ws_r20. From this multi-turn trajectory, the target-selection policy recovery_targets_v1 marks three assistant turns as supervised targets: the schema inspection request (e3), the corrected edit (e4), and the scoped final summary (e6). The malformed proposal (e2) and the runtime’s refusal remain in the input history as conditioning context. The member that entered the corpus is the redacted view trace_r19_01_redacted_v1, and the corpus manifest binds it to those role assignments, its lineage, and its split membership.
The raw event log of an admitted record is not yet in a form a training loop can safely ingest, because it interleaves model generations, runtime observations, and tool responses in the order they occurred. If a data loader treats every assistant message as a supervised target, it computes loss on the malformed edit proposal alongside the correction and reinforces the initial mistake. If the training pipeline computes loss across all sequence tokens indiscriminately, the model spends parameter capacity predicting deterministic environment replies, such as compiler diagnostics or directory listings, that it will never emit during inference. Finally, if the batching pipeline packs multiple trajectories into a shared context window without attention masking, self-attention heads attend across sample boundaries and leak private workspace contents into unrelated training examples. Standard cross-entropy loss fails silently under each of these mistakes: the training loss falls steadily as the optimizer fits the data, hiding the corrupted behavior from offline metrics. While curation contracts decide which actions show desirable behavior and evaluation contracts measure the deployed agent, neither specifies how to turn a multi-turn execution trace into isolated, causal gradient updates.
The path from an admitted record to an evaluated policy update begins with a question curation leaves open, namely which parts of the repair behavior belong in weights at all (section 16.2). The answer fixes what the data pipeline must protect. Serialization reconstructs each model input exactly as the agent received it (section 16.3), loss masking places the gradient on the selected responses and nowhere else (section 16.4), and batching shares buffers without letting examples read one another or skew one another’s weight (section 16.5). Masking and batching guard directly against the three silent failures above, and serialization guards the premise both rely on, that each prefix is the one the model actually saw.
A faithful loader still trains on recorded prefixes, and the deployed policy acts on prefixes its own choices produce. Exposure bias prices that gap over the horizon and decides when labeling the learner’s own states pays (section 16.6), and schema robustness asks the same question of the tool interfaces the policy is shown (section 16.7). The remaining sections vary the update itself along three independent choices: the objective, where preference optimization contrasts valid alternatives that imitation cannot rank (section 16.8); the behavior source, where distillation moves a costlier configuration’s behavior into a cheaper one (section 16.9); and the parameter artifact, where adapters change what a run must store (section 16.10). Closed-loop evaluation decides whether any of them earned a release (section 16.11), and a training manifest for the repair policy assembles the whole chain (section 16.12).
16.2 Scope of Supervised Fine-Tuning
Fine-tuning a model to work around a missing prompt instruction or a broken tool schema bakes a temporary workaround into billions of weights, only for the behavior to break again the next time the interface changes. Modifying model weights is an expensive, high-blast-radius change reserved for behavioral failures that survive cheaper fixes. The intervention ladder orders changes by cost and reversibility, so engineers test prompt, schema, and runtime fixes before spending compute on parameter updates, while diagnostic evidence can justify a direct weight update when the policy itself is at fault. Verified trajectory post-training (principle 13) requires curation to tie recorded evidence to a declared learning use and separate admissible targets from contextual or rejected actions. While this ordering is the default because of cost and blast radius, section 16.9 examines the complementary case where fine-tuning proves cheaper than paying the recurring token latency and cost of carrying lengthy procedural guidance in every prompt.
↩︎ The intervention ladder, which orders context, schema, runtime, and parameter interventions by cost and reversibility, is established in The intervention ladder.
Supervised fine-tuning (SFT) adjusts trainable model parameters to maximize the conditional likelihood of selected target tokens given their causal prefix context. When processing an admitted trajectory record such as trace_r19_01_redacted_v1, the training loop must follow the recovery sequence defined in recovery_targets_v1 rather than computing loss over every recorded token. The model must see the historical context and tool outputs that justify a chosen action, yet the data pipeline must exclude later steps that would reveal future observations or solutions ahead of time. Trajectory fine-tuning therefore inherits two data contracts: a target-selection contract that designates which tokens contribute to the loss gradient, and a conditioning contract that governs what causal history the model sees when generating those tokens.
The two contracts separate learning improved behavior from overfitting an idiosyncratic run. In an error-recovery trajectory, the behavioral goal is teaching the model to inspect an updated schema after a refusal, issue a properly formed corrective call, and interpret what the runtime verification check establishes. Memorizing the file paths of a specific test fixture or predicting hard-coded test outputs lowers training loss without teaching generalizable recovery skills. Whether the updated policy transfers these habits to unfamiliar environments and new tasks is what section 16.11 measures.
16.2.1 Procedure, contracts, and enforcement
Every response the repair agent generates comes from two inputs: the model’s static weights and the dynamic context window assembled for that call. Because fine-tuning updates only the weights, a tool-interface change—such as adding a required parameter or revoking a permission—immediately breaks a model that memorized the old signature in its parameters instead of reading the prompt. Engineers therefore divide agent behavior across three architectural homes, summarized in table 1:
- Procedure is behavior that stays stable across tasks and changes slowly, such as inspecting diagnostic output before modifying source code or checking unresolved issues before reporting completion. Procedure is what model weights can reliably absorb.
- Contracts are the dynamic operational facts that vary across deployments, user sessions, and individual turns, such as API tool definitions, schema requirements, environment parameters, and granted permissions. Contracts belong in the prompt context; a model that memorized an earlier schema will keep emitting obsolete arguments when an interface changes.
- Enforcement covers the invariants that must hold even when the model proposes an invalid or adversarial action, such as schema validation, authorization checks, sandbox containment, and execution timeouts. Enforcement stays in the deterministic runtime harness under invariant closure (principle 1).
| Home | Intended responsibility | Consequence for fine-tuning |
|---|---|---|
| Weights | Learned procedures and response tendencies, including tool selection and recovery behavior | Update the proposal distribution, then measure whether the changed behavior helps under the declared task conditions |
| Context | Current task facts, tool schemas, observations, and represented scope | Supply the information the policy needs now; do not rely on memorized versions as the authority for current facts or permissions |
| Runtime | Validation, authorization, containment, resource controls, and evidence collection | Preserve enforcement independently of the learned policy; a model’s predicted result is not an observed external effect |
Relying on parameters to store interface contracts makes a deployment brittle, because a revoked permission stays revoked regardless of how often the action appeared in past training demonstrations. Fine-tuning cannot replace dynamic contracts; what it can teach is the habit of reading and following whichever contract the context window supplies.
This boundary disproves the view that fine-tuning simply moves the prompt into the weights. Only the procedural parts of a prompt (structured reasoning guidelines, diagnostic checklists, or few-shot error-recovery examples) belong in the parameters, where they save token overhead on every turn (section 16.9). The contractual parts must stay in the prompt context, because omitting them forces the model to guess the current execution environment. The deterministic checks that enforce authorization and boundary invariants cannot live in either prompts or weights, because a probability distribution over tokens cannot guarantee an invariant.
Fine-tuning changes the policy’s proposal distribution: which tools it requests, how thoroughly it inspects diagnostic evidence, and which execution paths it tries. Those shifts change task success rates, end-to-end latency, and token use, yet every proposed action still runs inside the horizons, state transitions, and authority boundaries enforced by the runtime harness. This separation is the proposal boundary of The Model Call Boundary. Training reshapes what the policy proposes, while the runtime decides which actions may execute, and that boundary governs every adaptation choice that follows.
16.2.2 Objective, behavior source, and artifact
Once a team decides that a procedural habit belongs in the weights, “fine-tuning” still spans three independent choices that a single label can obscure:
- Training Objective: The loss function that defines the learning signal. Supervised imitation maximizes the likelihood of curated target demonstrations, whereas preference methods such as direct preference optimization (section 16.8) adjust weights by contrasting preferred trajectories against dispreferred alternatives under an explicit ranking rule.
- Behavior Source: Where the demonstrations and conditioning inputs come from. Traces may be collected from instrumented production runs and filtered by offline verification, generated by a teacher configuration with privileged context or a larger search budget (section 16.9), labeled on the states the learner itself reaches, or constructed by injecting faults the recordings never contained (section 16.6.2).
- Parameter Artifact: Which weights the run updates and stores. Teams can update all model parameters directly or train compact adapter matrices, such as low-rank adaptation (LoRA) modules, that attach to a frozen base checkpoint (section 16.10).
Because these three choices are independent, they combine freely; a deployment can, for example, train a low-rank adapter with a preference objective on synthetic teacher demonstrations. Whichever combination a team chooses, the run still trains on trajectory records converted into model inputs, so the first pipeline step is reconstructing the input prefix and conditioning window for each selected response.
16.3 Serializing Trajectories
Consider the repair demonstration admitted during curation: it contains a malformed edit proposal, a runtime refusal, a schema inspection, and a corrected edit. A training loader that strips out the refusal changes the problem the subsequent correction answered, and a loader that places a later passing test check before the edit leaks future evidence that the model did not have when proposing the patch. While curation selects which model completions to reinforce, serialization must reconstruct the exact prompt and observation history that prompted them, including any runtime redactions—such as synthetic placeholders substituted for private file paths—so the training example matches the context the model actually saw.
16.3.1 Reconstructing each model input
Runtime event logs record state transitions across the harness, but an event log alone does not show what was inside the model’s context window at a given call. System prompts change, tool schemas update between versions, and earlier messages are often summarized or truncated to fit token budgets. To reconstruct each call’s input faithfully, the training pipeline reads the delivered-input manifests, which record the exact messages, tool definitions, summaries, and bounded observations sent to the model along with their schema versions.
↩︎ The delivered-input manifest that records the exact payload sent to the model, rather than mutable paths or references, is established in The log event schema.
When an example requires historical fidelity, the loader serializes the selected invocation directly from its delivered-input manifest, preserving partial buffers, truncated outputs, and execution errors rather than substituting a cleaner result obtained later in the run. When a pipeline instead transforms an example deliberately—for instance, shortening prompts for context distillation (section 16.9)—it applies the reviewed transformation to both the prompt and the paired response, keeps the parent trace reference, and rechecks schema compatibility. Table 2 shows how the curated responses in the running repair demonstration map to their conditioning evidence; in the rendered training example, each row also carries the system instructions and active tool definitions delivered at that turn.
| Selected response | Relevant conditioning evidence | Response content selected by curation |
|---|---|---|
e3 |
Earlier delivered task context and the malformed e2 proposal with its pre-dispatch refusal |
Request to inspect the available schema |
e4 |
The corresponding returned schema and the retained context supporting the edit | Corrected edit with its version condition |
e6 |
The candidate-specific check summary delivered at e5 and retained evidence for the report’s scope |
Final report of the supported result and limitations |
The refusal at e2 shows that runtime validation caught the malformed proposal before dispatch, so no tool ran on the host. Because the check summary at e5 arrived only after the agent issued the corrected edit, it can condition the final report at e6 but cannot appear in the prefix for e4. Selecting the schema request at e3 likewise trains the model’s query tokens while leaving the returned schema as an external observation, so that loss masking in the next stage can separate model decisions from environment replies.
16.3.2 Rendering with versioned templates
Once the delivered messages for e3, e4, and e6 are reconstructed, the loader must turn those structured message objects—system instructions, user prompts, assistant reasoning, tool calls, and environment returns—into the exact token sequence the model expects at serving time. Chat templates differ across model families, and even checkpoints from the same base model often format tool schemas, parallel calls, or message boundaries differently; rendering a trajectory with delimiters that do not match the serving parser teaches the policy a syntax that production will reject. A serialized dataset artifact therefore records the tokenizer version, template configuration, tool-formatting schema, and response-boundary policy alongside the token arrays, because the message sequence and its rendering configuration together form the model input.
In schematic form, a multi-turn tool-calling example looks like the sequence below. The delimiter strings are illustrative placeholders, since each model family defines its own special tokens: the system segment gives the task contract and tool definitions, the assistant segment emits a reasoning trace followed by two parallel tool calls with unique identifiers, and each tool segment returns an observation bound to its call identifier.
<|system|> task contract; schemas for read_file(target_uri), run_tests(target) <|end|>
<|user|> Fix the connection-pool leak in src/pool.py <|end|>
<|assistant|> <|reasoning|> Reproduce the failure, then read the module. <|end_reasoning|>
<|call id=c1|> {"name": "read_file", "arguments": {"target_uri": "src/pool.py"}} <|end_call|>
<|call id=c2|> {"name": "run_tests", "arguments": {"target": "tests/test_pool.py"}} <|end_call|>
<|end_turn|>
<|tool id=c2|> FAILED test_release: pool exhausted <|end|>
<|tool id=c1|> ...file contents... <|end|>
<|assistant|> ...In this sequence, c2 finishes before c1, and because the runtime delivers observations in arrival order, the listing places the c2 reply first to match the prompt the model actually received on its next turn. The call identifiers keep each observation bound to its original call regardless of arrival order. Re-sorting the replies into call order (c1, then c2) for tidiness would invent a prompt the model never saw, and reordering calls whose arguments or observations depend on one another changes the causal problem into a synthetic transformation that needs its own validation. The <|end_turn|> token marks where the policy handed control back to the runtime loop, while all tokens between <|tool|> and <|end|> are external observations.
The training pipeline packages this rendered dataset as serialization_v1, whose manifest pairs the admitted execution records with the active rendering configuration and maps every selected response span to its token offsets. Because tokenization, template formatting, and string redaction can shift token boundaries—for example, when leading whitespace merges into an adjacent special token—the pipeline re-verifies all token offsets against the rendered token array and decodes completed examples to check task prompts, tool bindings, span boundaries, and end-of-sequence markers. If an example fails boundary alignment or drops content, the pipeline discards the record rather than patching it with synthetic filler.
Role delimiters define protocol syntax, but they carry no runtime authority: if a model emits text that mimics a <|tool|> delimiter, that text is still model output under the tokenizer’s grammar, and the serialization pipeline never promotes unverified generations into environment observations. Once the template fixes how each turn is rendered, the pipeline still has to decide whether a multi-turn trajectory can be stored as a single rendered sequence or must be split into separate per-invocation examples.
↩︎ The rule that model output is a proposal carrying no authority until the runtime validates it is established in The Model Call Boundary.
16.3.3 One example per invocation
Whether a trajectory such as the repair run can be stored as a single rendered transcript depends on how the serving runtime manages history between turns. Some runtimes append every turn to a growing transcript, while others prune earlier steps, compact intermediate context, or drop reasoning scratchpads between turns to reduce KV-cache memory. A single concatenated transcript is a faithful training example only when its causal prefixes match what the model actually received at each turn during serving. Because the repair agent’s runtime compacts history once the working set crosses its watermark, a turn delivered after compaction saw a digest rather than the raw earlier turns that a single concatenated transcript would place before it. Whenever serving prunes or compacts history between steps, the loader must serialize each invocation as its own training example containing that call’s exact delivered prefix.
↩︎ Watermark-triggered compaction, which replaces earlier turns with a digest between invocations, is established in Context Compaction.
Serializing each call separately carries the same triangular token cost that already priced an append-only trajectory at serving time (equation 2.9), now paid in stored training tokens. If a trajectory spans \(\TurnHorizon\) turns and each turn adds roughly \(\TurnAppendLength\) tokens, the example for turn \(\TurnIndex\) spans about \(\TurnIndex\,\TurnAppendLength\) tokens. Serializing all \(\TurnHorizon\) invocations separately stores \(\TurnAppendLength\,\TurnHorizon(\TurnHorizon+1)/2\) tokens, compared with \(\TurnHorizon\,\TurnAppendLength\) tokens for a single transcript—an expansion factor of \((\TurnHorizon+1)/2\), which over a stipulated 30-turn horizon multiplies the stored token volume by 15.5×. An agent’s history-retention policy therefore directly governs training storage and compute. The same fidelity rule applies to the model’s own reasoning text: if serving drops an earlier turn’s rationale before the next call, the reconstructed prefix for that call omits it too, whereas whether retained rationale is trained as a target or kept as context is a separate masking choice (section 16.4.4).
With serialization_v1 assembled, the pipeline has the tokenized sequences, their conditioning prefixes, and the response spans selected during curation. The next step is deciding which token positions contribute to the training loss, because computing cross-entropy over an entire multi-turn sequence would penalize the model for failing to predict tool outputs and user messages.
16.4 Observation Loss Masking
In serialization_v1, the token array for the repair trajectory still contains the prompt, the malformed proposal e2, the runtime refusal, the returned schema, and the check summary e5 alongside the selected responses e3, e4, and e6. A loader that selects every assistant token trains the model on the malformed proposal e2 alongside its correction e4, undoing curation, while a loader that computes loss across all tokens trains the policy to predict environment replies. Serialization turns the interaction into tokens; the loss mask decides which of those tokens drive gradient updates.
16.4.1 The masked objective
Let \(\SerializedTrainingExample=(\SequenceToken_1,\ldots,\SequenceToken_{\SequenceLength})\) be one serialized example, with an initial context or boundary token at \(\SequenceToken_1\), and let \(\EligibleTargetPositionSet\subseteq\{2,\ldots,\SequenceLength\}\) contain the policy-token positions selected for imitation. Which positions belong to \(\EligibleTargetPositionSet\) depends on the admitted response spans and the training objective rather than token authorship alone. With the binary mask \(\PositionLossMask_{t}=\IndicatorFunction\{\TokenIndex\in\EligibleTargetPositionSet\}\) and \(\PositionLossMask_1=0\), the mean masked cross-entropy for an example with at least one selected target is
\[\MaskedFinetuningLoss(\ModelParameters)=-\frac{1}{\ActiveTargetTokenCount}\sum_{t=2}^{\SequenceLength}\PositionLossMask_{t}\log\Policy(\SequenceToken_t\mid \TokenPrefix),\qquad \ActiveTargetTokenCount=\sum_{t=2}^{\SequenceLength}\PositionLossMask_{t}>0. \tag{1}\]
Each selected token is predicted from the recorded prefix \(\TokenPrefix\), which includes earlier actions, observations, and admitted mistakes, rather than from fresh rollouts sampled during training. This regime is teacher forcing1 (Williams and Zipser 1989), and it lets recovery prefixes that contain earlier failures serve as conditioning context whenever the trajectory satisfies the curation contract.
1 Teacher Forcing: Williams and Zipser (1989) named the technique for recurrent networks, feeding the recorded target back as the next input instead of the network’s own output. Because every prefix is recorded rather than generated, a causal model computes all of an example’s next-token predictions in one parallel forward pass, which is why imitating a trajectory costs one pass per example rather than one sampled rollout per decision.
Williams, Ronald J., and David Zipser. 1989. “A Learning Algorithm for Continually Running Fully Recurrent Neural Networks.” Neural Computation 1 (2): 270–80. https://doi.org/10.1162/neco.1989.1.2.270.
In the running repair trajectory, target_mask_v1 maps the reviewed response spans onto the token array of serialization_v1, separating target positions from conditioning context (table 3). Because changing a tokenizer vocabulary, chat template, or text transformation can shift token boundaries, the pipeline regenerates and checks masks against the final serialized tokens rather than copying offsets from an earlier stage.
| Token segment | Target policy for this imitation objective | What remains as conditioning input |
|---|---|---|
| Prompt and supplied tool contracts | \(\PositionLossMask_{t}=0\) | The actual task, current interfaces, and supplied constraints |
| Selected policy response | \(\PositionLossMask_{t}=1\) on its admitted target positions | Earlier tokens of the response and its available prefix |
| Unselected policy response, including the refused proposal | \(\PositionLossMask_{t}=0\) | History needed to interpret a selected correction |
| Tool or environment observation | \(\PositionLossMask_{t}=0\) | Delivered results, errors, and evidence available before the next response |
| Optional authorized rationale | Chosen by the declared target policy | Retained rationale remains visible whether or not it is selected |
| Termination boundary | Select an authentic policy-generated terminator when teaching the selected response to end | Do not substitute a tool-role boundary or invent a terminator for a cut-off response |
When an example has no selected targets (\(\ActiveTargetTokenCount=0\)), the loss in equation 1 is undefined, so the loader must drop the example or route it to a declared fallback rather than counting it as a completed demonstration. How examples of different lengths are weighted against one another inside a batch is a separate choice developed in section 16.5.
Definition 16.1: Observation loss masking
Observation loss masking excludes supplied observations from the direct token-prediction loss while retaining the permitted observation context needed to train selected policy responses.
- Significance: Focuses the optimization objective on the policy’s conditional distribution \(\Policy(\TrajectoryAction_{t}\mid \TrajectoryState_{t})\), so verbose environment output cannot dilute the gradient on policy actions regardless of how much context observations occupy (notebook 16.1 quantifies this dilution).
- Distinction: Masked tokens stay in the input context and participate in self-attention; only their next-token loss terms are zeroed. Policy authorship alone does not make a token a target, because refused attempts or unsafe actions are also masked even when emitted in the assistant role.
- Common pitfall: Selecting an entire multi-turn tool exchange simply because the tool call succeeded, which trains the model to predict external environment responses.
A deployed agent trajectory interleaves two different generators: the agent policy and the external environment. Training without observation masking forces one set of weights to fit both distributions at once:
- The agent policy \(\Policy(\TrajectoryAction_{t}\mid \TrajectoryState_{t})\), which maps prior observations and instructions to tool calls or answers.
- An environment dynamics model \(\EnvironmentTransitionDistribution(\TrajectoryState_{t+1}\mid \TrajectoryState_{t}, \TrajectoryAction_{t})\), which predicts the outputs returned by shells, databases, or remote APIs.
Because real tools run during deployment and return actual outputs, spending model capacity to predict environment returns is wasted work, and supervising the environment distribution causes two concrete failures:
- Loss dilution: Under token-averaged cross-entropy, unmasking expands the denominator from \(\ActiveTargetTokenCount\) to all \(\SequenceLength-1\) causal prediction terms (every stored position after the initial boundary). Each decision token’s gradient weight shrinks by the ratio of selected targets to predicted positions (computed in notebook 16.1), while the rest of the loss budget goes to file paths, stack traces, compiler dumps, and JSON payloads. Repetitive boilerplate pulls the average loss down, while novel stack traces or file contents dominate the gradient; in either case, the average stops measuring the policy.
- Role confusion: When trained to predict the tokens that follow a tool call, the policy learns that tool outputs continue the same autoregressive stream. At serving time, it can keep generating past its own tool call and write a fabricated tool result instead of emitting the stop token that hands control back to the runtime. Because the runtime rejects unauthenticated observations (section 16.3), those fabricated tokens are discarded, wasting output budget and stalling the turn.
The full repair record selects three responses (e3, e4, and e6). In the simplified view of figure 1, which omits the schema request e3 and its returned schema, only two of the seven segments shown—the corrected tool call and the final response—carry \(\PositionLossMask_{t}=1\) and receive direct imitation loss. The prompt, task statement, refused proposal, rejection message, and tool result carry \(\PositionLossMask_{t}=0\) and stay as conditioning context, where gradients from later selected targets still update the weights that encode them.
Observation masking prevents both loss dilution and role confusion at the objective level. Because an offline validation curve cannot show how often an unmasked model invents environment outputs during multi-step execution, that failure appears only in closed-loop trials.
16.4.2 Causal label alignment
In a causal autoregressive model,2 the hidden state at input position \(\TokenIndex-1\) produces the logits that predict token \(\SequenceToken_t\). The data loader must therefore pair the logits at \(\TokenIndex-1\) with target label \(\SequenceToken_t\) and target mask \(\PositionLossMask_{t}\), rather than with the mask of input token \(\SequenceToken_{t-1}\) (table 4). Because many training frameworks shift labels by one position inside their own loss function, shifting them once in the loader and a second time in the loss silently misaligns every target mask by one token across the corpus; checking the aligned labels and masks on a short synthetic example before launching a run catches this bug immediately.
2 Causal Self-Attention: Causal masking in Transformer decoders restricts each token from attending to subsequent positions (Vaswani et al. 2017). The same triangular structure enforces the label shift: the network predicts \(\SequenceToken_t\) using only the representation formed at \(\TokenIndex-1\).
Vaswani, Ashish, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. 2017. “Attention Is All You Need.” Advances in Neural Information Processing Systems (NeurIPS) 30: 5998–6008.
| Logits computed from a prefix ending at… | Next-token label | Target mask |
|---|---|---|
| The final supplied prompt token | First token of a selected policy response | \(1\) |
| The final content token of that selected response | Its authentic policy-generated end token | \(1\) |
| The end token before a delivered tool observation | First token of that observation | \(0\) |
Let \(\VocabularyTokenLogit_{t-1,\VocabularyToken}\) denote the logit for vocabulary token \(\VocabularyToken\) produced from \(\TokenPrefix\), using an ordinary softmax and no label smoothing. Differentiating equation 1 gives
\[\frac{\partial\MaskedFinetuningLoss}{\partial \VocabularyTokenLogit_{t-1,\VocabularyToken}}=\frac{\PositionLossMask_{t}}{\ActiveTargetTokenCount}\left(\Policy(\VocabularyToken\mid \TokenPrefix)-\IndicatorFunction\{\VocabularyToken=\SequenceToken_t\}\right).\]
When \(\PositionLossMask_{t}=0\), the prediction at position \(\TokenIndex-1\) contributes zero direct gradient, yet token \(\SequenceToken_t\) still shapes later representations through causal self-attention, so gradients from later active targets update the weights that read it. The model conditions on tool outputs without spending capacity trying to generate them, which carries verified trajectory post-training (principle 13) down to individual token positions.
16.4.3 Selected targets and compute
Observations fill 77.1 percent of the input; the loss sits on 4.9 percent.
Because every masked observation token still passes through self-attention, zeroing its loss term removes almost none of the forward-pass compute. The number of selected positions \(\ActiveTargetTokenCount\) sets the denominator of the loss and the gradient weight on each target token, whereas the full stored sequence length \(\SequenceLength\) sets the attention and projection FLOPs. Comparing \(\ActiveTargetTokenCount\) against the \(\SequenceLength-1\) causal prediction terms gives the dilution factor an unmasked loss would suffer, while \(\ActiveTargetTokenCount/\SequenceLength\) gives the fraction of stored positions that carry a target. Notebook 16.1 computes both quantities under two target policies on the same serialized trajectory.
Napkin Math 16.1: Stored tokens and selected targets
Problem: A hypothetical serialized example has the segment counts in table 5. Its curation record excludes some policy-response positions from imitation while retaining them as context. If the optional rationale is either selected or kept as context only, how many positions receive direct loss, what fraction of the stored input do they occupy, and by how much would an unmasked loss dilute each of them?
Table 5: Hypothetical Input Inventory: These disjoint counts include every stored position. They are not measurements of the shared repair run.
Table 6: Two Target Policies on One Input: Excluding a segment from the objective does not shorten the conditioning sequence.
| Input segment | Stored tokens |
|---|---|
| System prompt and schemas, including the initial boundary | 1,200 |
| Task statement | 184 |
| Optional authorized rationale | 1,400 |
| Other policy responses, including authentic end tokens | 960 |
| Delivered observations | 12,640 |
Variables: Exactly 160 of the other-policy-response positions belong to the refused proposal and stay context; every other response position is selected. The rationale is reviewed for either target policy, both policies see the same input, the initial position belongs to the prompt, and there is no padding.
Math: The disjoint segments sum to 16,384 stored tokens, of which observations occupy 77.1 percent. Removing the refused positions leaves 960 \(-\) 160 \(=\) 800 selected response targets, and selecting the rationale adds 1,400 more. The first stored boundary has no preceding-position prediction, so an unmasked causal average divides by 16,383 prediction terms instead of \(\ActiveTargetTokenCount\). Each selected target’s coefficient therefore shrinks by 16,383/800, a factor of 20.5× (7.4× when the rationale is selected). Under each target policy, table 6 sets \(\ActiveTargetTokenCount\) against the stored input it is drawn from.
| Target policy | Selected positions, \(\ActiveTargetTokenCount\) | Selected share of stored input |
|---|---|---|
| Selected responses plus reviewed rationale | 2,200 | 13.4% |
| Selected responses; rationale retained as context only | 800 | 4.9% |
Result: Both masks place the whole loss on a small minority of stored positions and leave the context identical; they differ only in whether the rationale is a target.
Systems insight: Because the dilution factor scales inversely with the selected fraction of predicted positions, it expands rapidly as environment interactions accumulate. Doubling the observation volume while holding response lengths constant escalates dilution from 20.5× to 36.3×. Consequently, observation masking becomes increasingly essential as agents tackle deep, tool-intensive reasoning traces. Masking, however, yields negligible reductions in FLOPs: every masked token must still pass through causal attention layers to form context for subsequent predictions. Only the final language model head projection and cross-entropy loss computation can be bypassed at masked positions.
Recorded repair trajectories show that the hypothetical inventory’s proportions are realistic. Figure 2 splits the transcripts of five published agent configurations on the same repository-repair tasks (Jimenez et al. 2024) by the role that produced each message.
Jimenez, Carlos E., John Yang, Alexander Wettig, Shunyu Yao, Kexin Pei, Ofir Press, and Karthik Narasimhan. 2024. “SWE-Bench: Can Language Models Resolve Real-World GitHub Issues?” International Conference on Learning Representations (ICLR).
Across the 5 configurations and 2,100 runs, observations make up 56 to 72 percent of the transcript, and model output only 12 to 31 percent. These shares count each message once. Serializing one example per invocation repeats the growing prefix in every example, so the observation share of the stored training input is larger still, approaching the inventory’s proportion. The mask is therefore not a refinement applied to a few verbose tool calls; in recorded repair work it decides the fate of most of the sequence.
16.4.4 Optional rationale as target or context
When a dataset includes intermediate reasoning traces or tool-use plans alongside the observations and tool calls, the loader must choose among three masking regimes (table 7) that match how the model will be served. Masking out a chain-of-thought scratchpad during training while expecting the deployed policy to act without any scratchpad at serving time trains the model on conditioning context it will not have in production. Supervising the rationale trains the model to generate its own reasoning before invoking tools; retaining the rationale under a loss mask treats the reasoning as supplied conditioning context, training the model to act on plans provided by an external orchestrator or user; and stripping the rationale entirely changes the conditional distribution so the policy acts from context alone.
Formally, given context prefix \(\TurnContext_{t}\), rationale tokens \(\RationaleTokens\), and action tokens \(\ResponseActionTokens=(\ResponseActionTokens_1,\ldots,\ResponseActionTokens_n)\), response-only supervision fits the conditional distribution \(\Policy(\ResponseActionTokens_j\mid \TurnContext_{t}, \RationaleTokens, \ActionTokenPrefix)\). If the production serving environment does not supply \(\RationaleTokens\), the deployed model instead evaluates \(\Policy(\ResponseActionTokens_j\mid \TurnContext_{t}, \ActionTokenPrefix)\). Because those are distinct conditional distributions, omitting rationale at serving time is a distillation step that must be evaluated directly against the unconditioned target.
| Treatment | Input and direct targets | Serving question to evaluate |
|---|---|---|
| Select reviewed rationale | Retain \([\TurnContext_{t},\RationaleTokens,\ResponseActionTokens]\); selected \(\RationaleTokens\) and \(\ResponseActionTokens\) positions receive loss | Does generated rationale improve the declared outcomes enough to justify its output and latency cost? |
| Retain rationale as context only | Retain \([\TurnContext_{t},\RationaleTokens,\ResponseActionTokens]\); only selected \(\ResponseActionTokens\) positions receive loss | How will appropriate rationale become available, and how do errors in it affect the response? |
| Remove rationale under a reviewed transformation | Use \([\TurnContext_{t},\ResponseActionTokens]\); selected \(\ResponseActionTokens\) positions receive loss | Does the policy perform adequately without the removed conditioning information? |
At serving time, either an external planner or the policy’s own autoregressive decoder supplies the reasoning text. When the model generates its own rationale, that text can drift from the training demonstrations, so closed-loop evaluation must measure whether the generated rationale helps enough to earn its output tokens and latency.
16.4.5 Termination and truncated responses
Whether a response carries rationale or only a tool call, its mask must also handle how the turn ends. To teach an agent when to yield control, the training loss must include the authentic policy-generated stop token (such as <|eom_id|> or <|im_end|>) at the end of each complete selected response, rather than confusing it with host delimiters that mark incoming tool payloads. Without direct supervision on the stop token, the policy can keep generating past the end of a valid tool call until it hits the output limit, returning TRUNCATED instead of COMPLETED so the runtime’s quarantining rule withholds the proposal.
↩︎ The normalized TRUNCATED status and the quarantining rule that withholds an incomplete candidate are established in The normalized status envelope.
Appending a synthetic end-of-turn token to a response cut short by a context window or timeout is the training-side twin of the hazard the quarantining rule blocks at serving time, where a truncated rm -rf argument still parses as a complete and broader command. The training-side version is worse, because a policy that has learned to emit the terminator after a partial command produces a COMPLETED status, which the quarantining rule admits. When a turn was truncated in the recorded trace, the loader has two sound options:
- Conservative exclusion: Mask the truncated turn entirely (\(\PositionLossMask_{t}=0\)) while keeping earlier complete turns in the trajectory.
- Prefix-only supervision: Supervise the valid initial tokens of the truncated turn up to the truncation point, and omit any terminating token.
Whenever trajectories undergo truncation, sliding-window packing, or template changes, the data pipeline recomputes target masks and checks that \(\ActiveTargetTokenCount > 0\). Automated assertions verify the first token of each agent response, the closing turn terminator, the transition into tool observations, refused recovery attempts, and truncated turns before the serialized dataset enters the training loop.
Checkpoint 16.1: Defend the target mask
With per-example loss masking in place, the loader separates conditioning context from imitation targets inside each sequence. Training at scale, however, combines trajectories of widely varying lengths into batches across device buffers, and the batching pipeline must preserve that per-example contract without leaking context across examples or silently skewing their weights.
16.5 Batching Variable-Length Trajectories
Consider a batch where a 30-turn debugging and code-repair trace shares a training buffer with a brief, one-turn file lookup. Under an ordinary lower-triangular causal mask, the lookup can attend directly to the repair trace’s intermediate file contents and tool outputs—context that the lookup policy will never have at serving time. Under token averaging, the multi-turn repair trace also outweighs the brief lookup by the ratio of their active target tokens.
Constructing a batch of variable-length trajectories therefore pairs two systems choices: how much gradient weight each trajectory receives in the batch loss, and how multiple sequences share physical buffer memory without leaking context across sequence boundaries. Sequence packing is only a memory layout; a faithful implementation keeps both the declared loss weighting and each example’s isolated predictive context.
16.5.1 Loss normalization across variable-length examples
Let a batch contain \(\BatchSize\) examples, where example \(\BatchExampleIndex\) has stored length \(\SequenceLength^{(\BatchExampleIndex)}\) (including context) and \({\ActiveTargetTokenCount}_{\BatchExampleIndex}>0\) selected target tokens under the mask of section 16.4, with mean masked negative log-likelihood \(\LossFunction_{\BatchExampleIndex}\). The two standard batch objectives weight the long repair trace and the one-turn file lookup differently:
\[\TokenAveragedBatchLoss= \sum_{\BatchExampleIndex=1}^{\BatchSize}\frac{{\ActiveTargetTokenCount}_{\BatchExampleIndex}}{\sum_{k=1}^{\BatchSize}{\ActiveTargetTokenCount}_k}\LossFunction_{\BatchExampleIndex}, \tag{2}\]
\[\ExampleAveragedBatchLoss=\frac{1}{\BatchSize}\sum_{\BatchExampleIndex=1}^{\BatchSize}\LossFunction_{\BatchExampleIndex}. \tag{3}\]
Token averaging in equation 2 gives equal weight to every selected token loss, whereas example averaging in equation 3 gives equal weight to every example mean. Both count only selected targets, never context or padding, and any record with no selected targets is excluded from both objectives. Because a serialized example may be a single invocation rather than a full task, example averaging equalizes serialized records rather than underlying tasks, while token averaging weights each record in proportion to its active target count.
The difference becomes critical when per-invocation serialization splits one multi-turn interaction into separate training records. If one of \(\BatchSize\) initial examples is split into \(\TurnHorizon\) per-invocation records while every other example stays unsplit, example averaging raises that trajectory’s share of the batch loss from \(1/\BatchSize\) to \(\TurnHorizon/(\BatchSize-1+\TurnHorizon)\), giving it \(\TurnHorizon\) times the weight of each unsplit example. Under token averaging, the trajectory’s aggregate share stays unchanged as long as its selected targets are partitioned disjointly across the \(\TurnHorizon\) records—each turn’s response selected once—because repeating prefix context does not repeat active targets. A corpus mixture specified in selected target tokens, the unit in which Records, stored tokens, and selected targets fixes class shares, therefore reaches the parameter update unchanged only under token averaging; example averaging, or a per-invocation split that selects the same target twice, changes the realized mixture (table 8).
| Property | Token averaging | Example averaging |
|---|---|---|
| Coefficient on example mean \(\LossFunction_{\BatchExampleIndex}\) | \({\ActiveTargetTokenCount}_{\BatchExampleIndex}/\sum_k {\ActiveTargetTokenCount}_k\) | \(1/\BatchSize\) |
| Coefficient on a selected token loss | \(1/\sum_k {\ActiveTargetTokenCount}_k\) | \(1/(\BatchSize {\ActiveTargetTokenCount}_{\BatchExampleIndex})\) |
| Quantity treated equally | Selected token losses | Example mean losses |
Preserving the coefficients in table 8 across a cluster requires matching distributed reduction. Because a large batch is split into microbatches (slices processed sequentially before a parameter step) and spread across data-parallel workers (each holding a model copy and processing its shard of the batch), computing a true global token average requires summing active token losses and active token counts across all microbatches and workers before dividing. Averaging local microbatch means directly matches the global token average only when every microbatch has the exact same active token count; when variable-length trajectories make those counts differ, workers must reduce the raw numerators and denominators before taking their ratio.
16.5.2 Example isolation in packed buffers
Packing the file lookup and the repair trace from the section’s opening into one buffer under an ordinary causal mask lets every query in the lookup read the repair trace’s file contents and tool observations. That leak across packed sequences is cross-contamination,3 and preventing it is the primary correctness obligation of sequence packing (Krell et al. 2021).
3 Cross-Contamination: Krell et al. (2021) used the term for attention and label leakage between sequences packed into one buffer, and showed that masks, position resets, and per-sequence loss handling make a packed model mathematically equivalent to the unpacked one. The equivalence is the property to test: a packed loader that trains faster but lets examples read one another has changed the learning problem, and no loss curve will show it.
Krell, Mario Michael, Matej Kosec, Sergio P. Perez, and Andrew Fitzgibbon. 2021. “Efficient Sequence Packing Without Cross-Contamination: Accelerating Large Language Models Without Impacting Performance.” arXiv Preprint arXiv:2107.02027, ahead of print. https://doi.org/10.48550/arXiv.2107.02027.
Definition 16.2: Sequence packing
Sequence packing places several variable-length training examples in shared buffers while preserving their intended information boundaries and loss targets.
- Significance: Uniform padding allocates the longest example’s length to every row, reducing batch occupancy to the ratio of mean to maximum sequence length. Sequence packing reclaims this wasted capacity (notebook 16.2).
- Distinction: Independent trajectories share physical memory buffers without sharing predictive context; multi-turn interactions within an individual trajectory retain access to their complete history.
- Common pitfall: Applying a standard causal mask across concatenated sequences, allowing later examples to attend illegitimately to earlier ones.
Let \(\QueryBufferPosition\) and \(\KeyBufferPosition\) denote buffer positions, \(\QueryVector_i\) and \(\KeyVector_j\) their query and key vectors, and \(\AttentionHeadDimension\) the attention-head dimension. The pre-softmax attention scores \(\text{score}_{i,j}\) enforce both causality and example membership:
\[\text{score}_{i,j}=\begin{cases} \QueryVector_i\KeyVector_j^T/\sqrt{\AttentionHeadDimension} & \text{if }\operatorname{example}(\QueryBufferPosition)=\operatorname{example}(\KeyBufferPosition)\text{ and }\KeyBufferPosition\le \QueryBufferPosition,\\ -\infty & \text{otherwise}. \end{cases} \tag{4}\]
Under an ordinary causal mask, every query in a later example can score keys in an earlier one; the block-diagonal mask sets each of those cross-example entries to \(-\infty\) before the softmax, and excludes trailing padding keys the same way (figure 3). Variable-length attention kernels (Dao 2023) enforce this without materializing a dense \(-\infty\) matrix. The loader hands the kernel the example boundaries as cumulative offsets, \(0\), \(\PackedExampleLength_1\), \(\PackedExampleLength_1+\PackedExampleLength_2\), and so on through the buffer’s total length (\(\BatchSize+1\) entries for \(\BatchSize\) examples), so that example \(\PackedExampleIndex\) occupies one contiguous interval of positions, and the kernel computes scores only for query and key pairs that fall inside the same interval. Whether a given kernel also skips the computation for cross-example blocks of the score matrix, and how much time that saves, is an implementation property to measure on the target accelerator. The mask is the correctness requirement; skipped work is a throughput bonus that must be shown, not assumed.
Dao, Tri. 2023. “FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning.” arXiv Preprint arXiv:2307.08691.
One long trajectory holds padded occupancy to 19.1 percent; packing reaches 76.6 percent.
Loss computation needs the same boundary isolation: shifting next-token labels across a concatenated buffer without resetting at example boundaries trains the final token of one example to predict the initial prompt token of the next. Each packed example’s initial prompt context must stay unselected (\(\PositionLossMask = 0\)), and every shifted prediction-target pair must stay inside its own example. The block-diagonal attention mask isolates visible context, while the per-example target mask isolates the loss gradient.
Napkin Math 16.2: Stored occupancy and target weighting
Problem: A batch holds one long trajectory among short ones. How much allocated sequence space does packing save over uniform padding, and how differently do the two loss normalizations weight a long and a short example?
Variables: Eight hypothetical examples have stored lengths \([32{,}768,\ 8{,}192,\ 4{,}096,\ 2{,}048,\ 1{,}024,\ 1{,}024,\ 512,\ 512]\). Each packed buffer holds \(32{,}768\) positions. Separately, two examples have \(256\) and \(8{,}192\) selected targets.
Math: Padding eight rows to the maximum allocates \(8\times32{,}768\) positions, or 262,144, for 50,176 stored tokens. Packing places the longest example alone in one buffer and the remaining 17,408 tokens in a second buffer. The two buffers allocate \(2\times32{,}768\) positions, or 65,536. Dividing stored tokens by allocated positions gives an occupancy of 19.1 percent with padding (the mean length of 6,272 divided by the maximum of 32,768) and 76.6 percent with packing. Packing allocates 4× fewer positions.
For the separate target-count illustration, token averaging gives coefficients \(256/(256+8{,}192)\) and \(8{,}192/(256+8{,}192)\), or 3.03 percent and 96.97 percent on the example means. Their ratio, \(8{,}192/256\), is 32×. Example averaging instead assigns each mean \(1/2\).
Systems insight: Padded occupancy is the ratio of mean to maximum length, so the single longest trajectory sets the waste: one 32,768-token example holds this batch to 19.1 percent, and packing raises it to 76.6 percent. Occupancy counts context and targets together, so the throughput it buys is measured on the actual attention kernel. Normalization is a separate knob with an equally large effect. Token averaging gives the long example 32× the short one’s weight, and example averaging gives them equal weight. Choose token averaging when every selected token should count equally, and example averaging when a few long trajectories should not dominate the update.
16.5.3 Position conventions and equivalence checks
Attention and label masks settle what each packed position reads and predicts, but each token also carries a position index. If the loader leaves buffer-wide position indices unreset, the second packed example starts at offset \(\PackedExampleLength_1\) instead of zero, which shifts its logits under any position scheme that is not shift-invariant. Resetting positions to zero at each example boundary avoids the shift directly: for zero-based buffer offsets \(\PackedExampleOffset_k\) and \(\PackedExampleOffset_{k+1}\) bounding example \(\PackedExampleIndex\),
\[\operatorname{pos}[\TokenIndex]=\TokenIndex-\PackedExampleOffset_k\quad\text{for }\TokenIndex\in[\PackedExampleOffset_k,\PackedExampleOffset_{k+1}).\]
Rotary position embeddings (RoPE), a widely used scheme that encodes position by rotating query and key vectors through position-dependent angles, are the case where the offset provably makes no difference, so for them resetting is an operational convention rather than a mathematical requirement. In the fixed-frequency formulation, position \(\QueryBufferPosition\) applies a rotation \(\RotaryRotationMatrix_i\) to a query or key vector. Shifting both positions by a common offset \(\CommonPositionShift\) leaves their relative rotation unchanged:
\[\RotaryRotationMatrix_{i+\CommonPositionShift}^T\RotaryRotationMatrix_{j+\CommonPositionShift}=\RotaryRotationMatrix_{j-i}=\RotaryRotationMatrix_i^T\RotaryRotationMatrix_j.\]
A common offset therefore leaves within-example attention scores unchanged, a direct consequence of the relative-position construction of RoPE (Su et al. 2024). Absolute-position features, position-dependent frequency scaling, and floating-point non-associativity can still compromise equivalence for a particular architecture, so no position scheme is exempt from checking.
Su, Jianlin, Murtadha Ahmed, Yu Lu, Shengfeng Pan, Wen Bo, and Yunfeng Liu. 2024. “RoFormer: Enhanced Transformer with Rotary Position Embedding.” Neurocomputing 568: 127063. https://doi.org/10.1016/j.neucom.2023.127063.
The general answer is therefore a test rather than a proof. The correctness test compares selected-token logits and the evaluated loss for the same examples processed separately and packed, with stochastic layers disabled and numerical tolerances explicitly stated. Perturbing one example must leave every other example’s logits unchanged. This test exercises attention boundaries, label shifting, position IDs, and the chosen denominator together. When it passes, the loader trains the demonstrated problem cleanly. The trained policy, however, still has to act in contexts its own choices produce.
16.6 Exposure Bias
Suppose a repair policy trained on this corpus is placed in a fresh bounded trial and meets a version conflict instead of the malformed-argument rejection recorded in trace_r19_01_redacted_v1. A plausible failure is that the policy repeats its previously accepted edit without first querying compatible state evidence. The runtime’s conditional-update guard then blocks the stale write, and the trial exhausts its execution allowance without completing the repair. The training corpus never included that state, so nothing in it taught a correction. While serialization and loss masking preserve the demonstrated decision problem, production deployment dictates which states the agent actually encounters. A policy trained only under teacher forcing never observes its own compounding errors during optimization. In standard text generation, an unexpected token merely drifts stylistic phrasing. In an autonomous agent, an errant action mutates the external environment and pushes the agent into states that the demonstration corpus never covered.
Definition 16.3: Exposure bias
Exposure bias is the mismatch between training a sequential policy on reference prefixes and running it on prefixes produced by its own earlier outputs (Bengio et al. 2015; Ranzato et al. 2016).
- Significance: It converts a per-step error rate into a horizon-dependent cost. Unchecked horizon compounding (Reliability and Cost over a Horizon) already makes the failure probability roughly \(\TurnHorizon\PerStepErrorRate\) for a small per-step rate \(\PerStepErrorRate\). When that rate is measured solely on reference states, the best available bound on excess task cost is \(\TurnHorizon^2\PerStepErrorRate\) (theorem 16.1, stated below), one further factor of \(\TurnHorizon\).
- Distinction: It reflects a structural property of the training distribution rather than an isolated mistake. A policy can err on states the corpus covers and recover on states it never showed.
- Common pitfall: Interpreting low offline validation loss on recorded prefixes as evidence that deployed execution will succeed. Offline loss samples only the reference trajectory, revealing nothing about the states the deployed policy visits once it deviates.
4 Exposure Bias: Ranzato et al. (2016) introduced the name for text generators trained on reference words and run on their own. Control researchers had met the same effect earlier as the compounding error of behavioral cloning, which Ross and Bagnell (2010) bounded. The two literatures describe one mechanism, and agents inherit the control version: a wrong action changes the environment, not just the phrasing.
Ranzato, Marc’Aurelio, Sumit Chopra, Michael Auli, and Wojciech Zaremba. 2016. “Sequence Level Training with Recurrent Neural Networks.” 4th International Conference on Learning Representations (ICLR).
Ross, Stéphane, and J. Andrew Bagnell. 2010. “Efficient Reductions for Imitation Learning.” Proceedings of the Thirteenth International Conference on Artificial Intelligence and Statistics (AISTATS) 9: 661–68.
Offline training corpora can include recovery trajectories, allowing an agent that departs from an ideal demonstration to re-enter a familiar state or discover an alternate valid solution. Exposure bias,4 however, concerns how thoroughly the corpus covers the state distribution generated by the deployed policy itself, including unexpected edge states like the version conflict.
16.6.1 State distributions and imitation regret
An idealized finite-horizon model makes the cost of missing states like the version conflict exact. Following the imitation-learning notation locally in this derivation, let \(\TurnIndex\) count discrete decision steps rather than token positions, let \(\DecisionProcessState\) denote the complete environment state of the decision process (including interaction history and system variables the agent cannot directly observe, whereas \(\TrajectoryState_{t}\) in definition 16.1 named only the policy’s visible token context), and view the trained policy \(\pi\) (corresponding to \(\Policy\)) and the deterministic expert reference \(\pi^*\) as maps from states to actions over a \(\TurnHorizon\)-step horizon under identical initial states and transition dynamics. If \(\StateDistribution^t\) is the distribution of states visited before step \(\TurnIndex\) under \(\pi\), the horizon-averaged state distribution is \(\StateDistribution = \TurnHorizon^{-1}\sum_{t=1}^\TurnHorizon \StateDistribution^t\).
Let \(\ImmediateTaskCostScore(\DecisionProcessState,\PolicyAction)\in[0,1]\) be a normalized immediate task cost and \(\ExpectedTaskCostSum(\pi)\) its expected sum over the horizon. Let \(\LossFunction(\DecisionProcessState,\pi)=\Pr_{\PolicyAction\sim\pi(\cdot\mid \DecisionProcessState)}[\PolicyAction\ne\pi^*(\DecisionProcessState)]\) measure the probability of choosing an action different from the reference at decision time—a conservative proxy for task error, since several distinct actions may be acceptable, and a decision-step quantity distinct from token-level cross-entropy. Finally, the cost-to-go \(\ReferenceCostToGo(\DecisionProcessState,\PolicyAction)\) is the expected cost from step \(\TurnIndex\) through \(\TurnHorizon\) after taking \(\PolicyAction\) in \(\DecisionProcessState\) and following the reference \(\pi^*\) thereafter.
Theorem 16.1: Imitation regret bounds
The horizon penalty of imitation depends on which states the error is measured on. For a policy \(\pi\) imitating a deterministic reference \(\pi^*\) over \(\TurnHorizon\) steps with per-step cost in \([0,1]\), Ross et al. (2011) bound the excess task cost in two ways.
Reference-state error (their Theorem 2.1, first proved by Ross and Bagnell (2010)). With disagreement measured on the reference policy’s states, \(\ReferenceStateDisagreementRate=\mathbb E_{\DecisionProcessState\sim \ReferenceStateDistribution}[\LossFunction(\DecisionProcessState,\pi)]\),
\[\ExpectedTaskCostSum(\pi)-\ExpectedTaskCostSum(\pi^*)\le \min\{\TurnHorizon,\;\TurnHorizon^2\ReferenceStateDisagreementRate\}. \tag{5}\]
Learner-state error (their Theorem 2.2). With disagreement measured on the learner’s states, \(\LearnerStateDisagreementRate=\mathbb E_{\DecisionProcessState\sim \StateDistribution}[\LossFunction(\DecisionProcessState,\pi)]\), and every deviation followed by reference recovery costing at most \(\DeviationDisadvantage\) extra, \(\ReferenceCostToGo(\DecisionProcessState,\PolicyAction)-\ReferenceCostToGo(\DecisionProcessState,\pi^*(\DecisionProcessState))\le \DeviationDisadvantage\) at every learner-reachable state,
\[\ExpectedTaskCostSum(\pi)-\ExpectedTaskCostSum(\pi^*)\le \min\{\TurnHorizon,\;\DeviationDisadvantage \TurnHorizon\LearnerStateDisagreementRate\}. \tag{6}\]
- Horizon (\(\TurnHorizon\)): decision steps in the task; because each step costs at most one, \(\TurnHorizon\) also caps both bounds.
- State distributions (\(\ReferenceStateDistribution\), \(\StateDistribution\)): the environment states the reference or the learner visits, averaged over the horizon.
- Disagreement rates (\(\ReferenceStateDisagreementRate\), \(\LearnerStateDisagreementRate\)): probability of an action other than the reference’s, averaged over reference-visited vs. learner-visited states.
- Deviation disadvantage (\(\DeviationDisadvantage\ge0\)): the largest extra expected cost of one deviation followed by reference recovery, an assumption about the task rather than a measured retry count.
The same error rate therefore carries different horizon penalties. At \(\TurnHorizon=\) 30 and a 2 percent disagreement rate, the reference-state bound is 18 cost units and the learner-state bound with \(\DeviationDisadvantage=\) 1 is 0.6, a gap of 30×, which is \(\TurnHorizon/\DeviationDisadvantage\). The linear bound holds only while \(\DeviationDisadvantage\) stays independent of \(\TurnHorizon\); one action that can spoil the rest of the task makes \(\DeviationDisadvantage\) grow with \(\TurnHorizon\) and restores the quadratic dependence.
The quadratic term has an intuitive physical reading. Couple the environment randomness of the two policies until they first choose different actions. The probability of any such divergence is at most the sum of disagreement probabilities along the reference process, \(\TurnHorizon\ReferenceStateDisagreementRate\), and a divergence can corrupt the entire remaining horizon. This coupling argument makes no assumptions regarding the learner’s error rate once it leaves the reference trajectory. Conditional error bounds in trajectory imitation derives the reference-state bound by coupling and the learner-state bound through telescoping policy comparisons.
Unchecked runtime step errors compound as \((1-\PerStepErrorRate)^\TurnHorizon\) (equation 1.1), and for small \(\PerStepErrorRate\) this is close to \(1-\TurnHorizon\PerStepErrorRate\), so the chance of some failure grows roughly as \(\TurnHorizon\PerStepErrorRate\). Theorem 16.1 adds the training-side consequence, since imitation whose error is measured on the reference distribution is bounded only by \(\TurnHorizon^2\PerStepErrorRate\), a second factor of \(\TurnHorizon\). The two rates are analogous rather than identical. The introduction’s \(\PerStepErrorRate\) is a probability of error, while the disagreement rate counts every departure from the reference, harmless ones included, and so upper-bounds that error probability.
Same 2 percent error: 18 cost units on reference states, 0.6 on learner states.
The bounds are stated in normalized cost units. If each step spent off the reference path costs one unit and every other step costs nothing, the reference-state bound of 18 units caps the expected number of wasted turns at 18 of 30, while the learner-state bound of 0.6 caps it below one. Wasted turns reach the familiar evaluation quantities directly, since they raise the expected attempt cost in \(\CostPerAcceptedTask\) and, once they exhaust the turn budget, lower the acceptance probability. The deviation disadvantage \(\DeviationDisadvantage\) is an assumption about the task and its reference, not a property of the trained learner; it asks how much extra cost one slip adds when the reference takes over afterward. A repair edit that the runtime refuses before dispatch, or that a version guard blocks, costs a small, fixed number of extra turns to correct (in the recorded trace, one schema inspection and one reissued call) however long the task is, which is why a small \(\DeviationDisadvantage\) that does not grow with \(\TurnHorizon\) is plausible for the repair agent. Retrospective recovery records help only to the extent that they populate the training corpus with learner-visited states and keep \(\DeviationDisadvantage\) small, but an offline dataset alone cannot guarantee that either assumption holds during deployment.
Napkin Math 16.3: The same error rate under different assumptions
Problem: An illustrative task has horizon 30 and immediate cost bounded between zero and one. Compare a 2 percent disagreement rate on reference-policy states with the same numerical rate on learner-policy states. What changes if the latter also satisfies a deviation-disadvantage bound of \(\DeviationDisadvantage=\) 1, and how accurate would offline imitation have to be to match it?
Math: The reference-state bound multiplies \(\TurnHorizon^2=\) 900 by \(\ReferenceStateDisagreementRate=\) 0.02, giving 18 normalized cost units, below the cap of 30. The learner-state bound multiplies \(\DeviationDisadvantage \TurnHorizon=\) 30 by the same rate, giving 0.6. If the disadvantage instead reaches \(\DeviationDisadvantage=\TurnHorizon\), the learner bound returns to 18. Matching the learner-state bound on reference states requires \(\ReferenceStateDisagreementRate\) equal to that bound divided by \(\TurnHorizon^2\), about 0.07 percent.
Result: At this horizon the two bounds differ by 30× for the same measured error rate.
Systems insight: The gap is \(\TurnHorizon/\DeviationDisadvantage\), so it grows linearly with the horizon. Offline imitation must be \(\TurnHorizon/\DeviationDisadvantage\) times more accurate on reference states to carry the guarantee that labeling learner-visited states provides. The reference-state bound also saturates at its cap once \(\TurnHorizon\ge 1/\ReferenceStateDisagreementRate\), here 50 steps, beyond which it constrains nothing. Long-horizon agents with cheap, recoverable deviations are therefore where on-policy labeling pays most. When one deviation can spoil the task (\(\DeviationDisadvantage\) near \(\TurnHorizon\)), the learner-state bound is no better than the quadratic one, and the runtime checks that invariant closure (principle 1) requires carry the reliability.
16.6.2 Dataset aggregation
Because a fixed offline corpus never shows the policy how to recover from states that only the learner’s own actions reach—such as the version conflict—Dataset Aggregation (DAgger) closes the coverage gap by rolling out trajectories under successive learned policies, querying an authorized reference labeler at the states they actually visit, adding those labeled transitions to the training corpus, and retraining the policy5 (Ross et al. 2011). The resulting dataset is on-policy data, drawn from the state distribution of the policy being trained. Executing a learner action in the environment and generating its reference label are separate operations, and early collection rounds mix them: in round \(\LabelingRoundIndex\), the collection worker executes the reference action with probability \(\ReferenceMixingProbability_k\) and the learner’s action otherwise, decaying \(\ReferenceMixingProbability_k\) toward zero across rounds. This schedule shifts the collected states toward the learner’s own distribution, where theorem 16.1 measures \(\LearnerStateDisagreementRate\). Because aggregation broadens coverage without making the training and learner distributions identical, and because the linear bound holds only where \(\DeviationDisadvantage\) stays small, each retrained round is checked on fresh rollouts rather than assumed to satisfy the bound.
5 DAgger’s Mixing Schedule: Ross et al. (2011) require only that the average reference-execution probability tend to zero over rounds, and report that the simplest schedule, executing the reference only in the first round, often performs best. The systems cost sits elsewhere: every visited state needs a reference label on demand, so the labeler is an online service in the collection loop rather than a one-time annotation pass.
Ross, Stéphane, Geoffrey Gordon, and J. Andrew Bagnell. 2011. “A Reduction of Imitation Learning and Structured Prediction to No-Regret Online Learning.” Proceedings of the Fourteenth International Conference on Artificial Intelligence and Statistics (AISTATS) 15: 627–35.
When the learner reaches an off-distribution state like the version conflict, the labeler—typically a larger foundation model or a search procedure with an expanded budget—proposes a corrective action, which can itself be wrong. To avoid training recovery in the wrong direction, the curation pipeline executes the proposed step inside an isolated sandbox and admits the label only if automated verification passes under the predicates of Admission predicates. A passing continuation confirms that the label works, not that it is the only valid action. Because steps that look superficially valid can still rely on ungrounded assumptions or violated preconditions, the pipeline records ground-truth observations, enforces tool contracts, and tracks lineage for disputed labels; after retraining shifts the policy’s visited states, the updated checkpoint is evaluated on fresh bounded rollouts and held-out suites rather than on the accumulated training buffer.
Fault injection supplies off-path states deliberately rather than waiting for the learner to reach them. The fault conditions a collection fixture must reproduce faithfully, and the causal distinction each one preserves, are established in Fault injection. The training-specific question is where the mask falls. In a fault-injected record, the synthesized error is presented as an observation with loss mask \(\PositionLossMask_{t} = 0\), and the diagnosis and corrective tool invocation that follow are targets with \(\PositionLossMask_{t} = 1\), so the gradient lands on decision logic that nominal demonstrations never exercise. The selected targets inherit the fixture’s causal distinctions, so a record that answers a lost acknowledgment with a blind retry teaches the duplicate effect that settlement before retry (principle 8) exists to prevent. Fault injection covers the failures the harness designers anticipated, whereas on-policy collection uncovers the ones the model’s own behavior produces (figure 4). Both are choices of behavior source in the sense of section 16.2.2, and both still need a step-level label for every selected decision. Reinforcement learning from verifiable rewards drops that requirement by scoring fresh rollouts against terminal outcomes instead.
↪ Optimizing a policy against verifiable outcomes on fresh rollouts, without step-level labels, is developed in Foundations of Verifiable Rewards.
Checkpoint 16.2: Coverage after faithful imitation
State coverage encompasses both environment variables and the tool contracts presented at each decision step. An agent may recover reliably from familiar runtime exceptions yet fail persistently by emitting deprecated argument names following an interface update. Schema robustness tests this narrower sensitivity to structured tool definitions.
16.7 Schema Robustness
When a training trajectory captures an agent inspecting a tool schema before repairing a failed invocation, fine-tuning teaches the policy to consult tool definitions before emitting an action. Consulting a schema and adhering to an altered interface contract, however, are two distinct capabilities: if a training corpus exposes the model to only a single static tool signature across thousands of examples, the policy can memorize specific parameter keys and argument types in its weights instead of reading the contract supplied in the prompt context. Testing whether a model genuinely conditions on the presented schema requires perturbing the interface definition while holding the underlying task fixed.
Consider a read-tool interface update. The active runtime interface requires a string argument named target_uri, forbids extraneous keys, and resolves paths relative to an authorized fixture directory. A policy that memorized an earlier interface version instead emits the legacy parameter name file_path:
{
"proposed_call": {
"name": "read_file",
"arguments": {"file_path": "src/pool.py"}
},
"current_argument_schema": {
"type": "object",
"properties": {"target_uri": {"type": "string"}},
"required": ["target_uri"],
"additionalProperties": false
}
}This proposal violates the presented schema because it omits the required target_uri property and supplies the disallowed file_path key. Because the agent runtime validates invocations before dispatch, the schema validator rejects the call immediately. The compatible payload is {"target_uri": "src/pool.py"}. Beyond syntactic validity, whether that path exists and whether the agent holds authorization to access it remain independent downstream runtime checks.
Emitting stale parameter keys is a clear failure mode of interface memorization, but attributing it to trajectory fine-tuning requires controlled comparison, because base foundation models also make schema errors when prompted zero-shot. Isolating the effect of fine-tuning demands evaluating both the base and fine-tuned policies under identical prompts, sampling parameters, and tool schemas; section 16.11 formalizes the metrics and comparison protocol that establish this attribution.
16.7.1 Schema augmentation
To keep the policy from memorizing file_path, schema augmentation constructs training trajectories with declared interface mutations so the model must read the argument names from the prompt context rather than from its weights. That transformation requires more than a string find-and-replace on tool-call arguments: renaming file_path to target_uri in the call while leaving file_path in the prompt schema or in an earlier error message trains the policy on a contradictory transcript. Every part of the decision context—schema definitions, selected calls, mock observations, and error messages—must change together, and the resulting derived trajectories are tracked under the lineage rules of Derived records. Conditioning policies directly on delivered interface contracts is well established: incorporating retrieved API documentation into training prompts lets fine-tuned models adapt when documentation changes after deployment (Patil et al. 2023). Table 9 pairs each schema transformation with the consistency check that keeps the resulting trajectory coherent.
Patil, Shishir G., Tianjun Zhang, Xin Wang, and Joseph E. Gonzalez. 2023. “Gorilla: Large Language Model Connected with Massive APIs.” arXiv Preprint arXiv:2305.15334, ahead of print. https://doi.org/10.48550/arXiv.2305.15334.
↩︎ The rule that a transformed trajectory becomes a new derived version bound to its parent trace is established in Derived records.
| Transformation | Intended coverage | Required consistency check |
|---|---|---|
| Reorder declarations or object fields | Dependence on one textual order where order has no semantic role | Preserve represented values and dependencies; do not shuffle dependent tool operations |
| Rename tool or argument identifiers | Binding the call to the presented interface | Update schema, selected calls, documentation, and relevant embedded references; provide compatible adapter behavior |
| Add unused tool declarations | Tool selection amid additional choices | Keep distractors coherent with available capabilities and authority; account for longer or ambiguous context |
| Remove an optional argument | Behavior when a previously available option is absent | Establish the replacement default or behavior, then recheck observations and target adequacy; deletion need not preserve the outcome |
In the read-tool example, a meaning-preserving rename needs an adapter that maps target_uri to the underlying file operation, and every prior tool output, system prompt, or turn that mentions file_path must be updated across the context so the derived record shows a coherent interaction without claiming those bytes appeared in the original historical run. Whenever a transformation changes return values, ordering, defaults, or side effects, the pipeline generates fresh execution evidence—either by replaying the trajectory in an instrumented sandbox or synthesizing verified responses—rather than splicing stale historical observations onto modified calls.
These transformations carry direct storage and compute costs: renaming identifiers shifts subword boundaries and token counts, adding distractor tools lengthens the prompt at every turn (raising training FLOPs and inference KV-cache footprints), and every synthetic variant expands the corpus that verification pipelines must check. Teams therefore version each transformation pipeline, track token expansion ratios, and test whether the schema transfer gains pay for the added training and review cost.
16.7.2 Transfer evaluation
Measuring whether schema augmentation actually prevented memorization requires separating three evaluation regimes: new tasks on familiar tool schemas, familiar tasks on renamed or restructured interfaces, and new tasks on unfamiliar tool families. Mixing these regimes in one score hides where a policy breaks. Evaluation splits must hold out task answers and their synthetic variants, whereas reusing familiar tool schemas across train and test splits is not contamination when the benchmark is testing task transfer under a known interface.
↩︎ The isolation of held-out task answers, and the rule that transformed records inherit their parent’s split group, are established in Held-out evidence protection.
Within each regime, the evaluation harness separates structural schema conformance, tool selection, and task completion, because a call can pass the JSON Schema validator while invoking the wrong tool, targeting an unauthorized path, or supplying arguments that do not fix the bug. Comparing augmented and unaugmented checkpoints under an equal token budget across both familiar and renamed schemas shows whether the model gained interface robustness without losing task performance, while the runtime continues to validate every call before execution. Supervised imitation and schema augmentation together teach an agent to emit valid calls that match demonstrated behavior; when several valid calls all pass the schema and solve the task at different costs or risks, the objective must move from imitation to preference comparison.
16.8 Trajectory Preference Optimization
Consider two independently collected continuations of the connection-pool repair that start from identical diagnostic context and both resolve the leak. One continuation runs the entire integration test suite after every minor edit, while the other runs targeted unit tests and executes the full suite only upon completion. Supervised imitation of both trajectories teaches the agent wasteful verification habits alongside disciplined ones; training only on the efficient path discards the signal that the slower trajectory was functionally correct yet needlessly expensive. A preference comparison captures this trade-off directly, steering the policy toward the efficient trajectory without treating the slower candidate as a failed repair. Not every pair of recorded interactions forms an admissible preference comparison: in the recovery demonstration, for instance, the malformed proposal e2 and the subsequent correction e4 cannot be paired as preferences, because the correction conditioned on environmental error messages that e2 never saw. Valid comparisons require alternative decisions evaluated under identical available information.
↩︎ The preference-pair admission predicates, which require competing trajectories to branch from identical conditioning contexts, are established in Admission predicates.
16.8.1 The trajectory preference loss
Let \(\ComparisonContext\) identify the shared comparison context and evaluation criterion, and let \(\PreferredTrajectory\) and \(\DispreferredTrajectory\) denote the preferred and dispreferred trajectory records. Unlike a supervised fine-tuning mask that zeros out dispreferred attempts, a preference comparison selects the divergent policy spans in both records and contrasts their log-probability ratios against a frozen starting checkpoint \(\pi_{\text{ref}}\) (typically the supervised fine-tuned checkpoint). These subscripts reflect the relative ranking assigned during curation rather than binary correctness, and the word “reference” now names a different role than it did for \(\pi^*\) in section 16.6: there, \(\pi^*\) was the expert whose actions defined correct behavior, whereas here \(\pi_{\text{ref}}\) is only the starting checkpoint the update is measured against. Both the active policy \(\Policy\) and \(\pi_{\text{ref}}\) must assign nonzero probability to all scored tokens under the specified input prefixes.
For each trajectory, let \(\InteractionHistory_t\) denote the rendered context prefix used to predict action token \(\SequenceToken_t\), and let \(\PreferenceScoredPositionSet(\Trajectory)\) denote the set of policy-generated token indices selected by the preference objective. Define the reference-relative score as
\[\ReferenceRelativeScore(\Trajectory;\ComparisonContext)=\sum_{t\in\PreferenceScoredPositionSet(\Trajectory)} \log\frac{\Policy(\SequenceToken_t\mid \InteractionHistory_t)}{\pi_{\text{ref}}(\SequenceToken_t\mid \InteractionHistory_t)}.\]
Environment observations interleaved within \(\InteractionHistory_t\) condition subsequent predictions but carry no policy likelihood terms. The token selection set \(\PreferenceScoredPositionSet(\Trajectory)\) is customized specifically for each comparison. The dispreferred trajectory is evaluated solely to construct a negative contrast; it never serves as an imitation target.
Earlier reinforcement learning from human feedback pipelines fit a separate reward model to pairwise comparisons and then optimized the policy against that learned surrogate using reinforcement learning (Christiano et al. 2017; Ouyang et al. 2022). Direct preference optimization (DPO) bypasses the reward model entirely by optimizing a logistic loss directly over the difference of policy-to-reference log-probability ratios (Rafailov et al. 2023). Using the standard sigmoid function \(\SigmoidFunction(\PreferenceMargin)=1/(1+e^{-\PreferenceMargin})\), a positive preference scale \(\PreferenceScale\) that sets how strongly the objective resists drift from \(\pi_{\text{ref}}\), and a dataset \(\PreferencePairDataset\) of curated comparison pairs, the resulting objective over selected tokens is
\[\DpoLoss=-\mathbb E_{(\ComparisonContext,\PreferredTrajectory,\DispreferredTrajectory)\sim\PreferencePairDataset} \log\SigmoidFunction\!\left(\PreferenceScale[\ReferenceRelativeScore(\PreferredTrajectory;\ComparisonContext)-\ReferenceRelativeScore(\DispreferredTrajectory;\ComparisonContext)]\right). \tag{7}\]
The classic DPO derivation applies to single-turn responses generated for an identical prompt under the Bradley–Terry preference model6 (Bradley and Terry 1952) and penalizes drift from \(\pi_{\text{ref}}\) with a Kullback–Leibler divergence, the expected log-ratio of two distributions’ probabilities, which is zero only when they match. The penalty is taken in the “reverse” direction, with the expectation over the trained policy’s own outputs rather than the reference’s. In outline, the best policy under that penalty is the reference reweighted by the exponentiated reward and divided by a per-context normalizing term (the partition function) that makes its probabilities sum to one. Solving that relation for the reward expresses it as \(\PreferenceScale\) times the policy-to-reference log ratio plus \(\PreferenceScale\) times the logarithm of the normalizer. Because both members of a pair share the same context, they share the same normalizer, and it drops out when one score is subtracted from the other, which is all the Bradley–Terry model ever uses. Consequently, the implicit reward is identified only up to an additive constant per prompt context. Applying this framework to multi-step agent trajectories by scoring selected token spans \(\PreferenceScoredPositionSet(\Trajectory)\) is an engineering surrogate built on that foundation, and closed-loop evaluation must validate it.
6 Bradley–Terry Model: The paired-comparison model of Bradley and Terry (1952) gives the probability that item \(\FirstComparedItem\) is preferred to item \(\SecondComparedItem\) as a logistic function of the difference of their latent scores. Only score differences are identified, which is why DPO’s implied reward carries an input-dependent constant that cancels within a pair, and why a DPO score has no meaning outside the comparison that produced it.
Bradley, Ralph Allan, and Milton E. Terry. 1952. “Rank Analysis of Incomplete Block Designs: I. The Method of Paired Comparisons.” Biometrika 39 (3/4): 324–45. https://doi.org/10.2307/2334029.
↩︎ The Kullback–Leibler divergence, used there to compare a baseline’s and a candidate’s tool-call distributions, is established in Offline comparison and shadow evidence.
When an objective scores every policy-generated token across a complete interaction, stationary environment dynamics cancel out in the policy-to-reference ratio for that specific recorded trajectory, because both models evaluate the actions against the identical sequence of historical observations. Between two alternative trajectories, however, the observed states diverge. Those divergent observations alter the conditioning context \(\InteractionHistory_t\) for all downstream actions, ensuring that environmental feedback continues to shape the cumulative score even though the policy assigns no explicit probability to environment outputs.
For a single comparison pair, define the preference margin \(\PreferenceMargin=\PreferenceScale(\PreferredReferenceRelativeScore-\DispreferredReferenceRelativeScore)\) (a new use of \(\PreferenceMargin\), unrelated to the rationale tokens of section 16.4.4), where \(\PreferredReferenceRelativeScore\) and \(\DispreferredReferenceRelativeScore\) represent the reference-relative log-likelihood sums for the preferred and dispreferred trajectories. Differentiating the loss in equation 7 with respect to the policy parameters \(\ModelParameters\) while holding the reference policy fixed yields
\[\nabla_\ModelParameters\DpoLoss=-\PreferenceScale\SigmoidFunction(-\PreferenceMargin) \left(\sum_{t\in\PreferenceScoredPositionSet(\PreferredTrajectory)}\nabla_\ModelParameters\log\Policy(\SequenceToken_t\mid \InteractionHistory_t) -\sum_{t\in\PreferenceScoredPositionSet(\DispreferredTrajectory)}\nabla_\ModelParameters\log\Policy(\SequenceToken_t\mid \InteractionHistory_t)\right). \tag{8}\]
The gradient has three systems consequences. First, identical selected tokens sharing an identical conditioning prefix contribute equal gradient vectors to both sums in equation 8 and cancel exactly. A shared prefix therefore suffers no negative gradient updates simply because it also appears at the start of the dispreferred record. Once the paths diverge, identical tokens no longer cancel because their conditioning histories differ. Second, because shared model parameters parameterize all tokens, updating weights on selected positions also shifts likelihoods on unselected positions. Third, the full scalar coefficient multiplying the difference of gradient sums is \(\PreferenceScale\SigmoidFunction(-\PreferenceMargin)\); \(\SigmoidFunction(-\PreferenceMargin)=1-\SigmoidFunction(\PreferenceMargin)\) is its saturation factor. At the reference policy, both scores are zero (\(\PreferenceMargin=0\)), so the factor is \(0.5\) and the coefficient is \(\PreferenceScale/2\). At fixed \(\PreferenceScale\), increasing a positive margin reduces both quantities, dampening the contribution of pairs the policy already distinguishes. The coefficient alone does not determine the full gradient norm or the optimizer step, which also depend on the gradient vectors and optimization state.
Napkin Math 16.4: How hard one comparison pulls
Problem: After some training, the policy has moved relative to the reference on one admitted comparison. What are the pair’s margin, loss, saturation factor, and full gradient coefficient, and how do they change if \(\PreferenceScale\) doubles?
Variables: The selected-token log-probability sums are -8 for the active policy and -10 for the reference on the preferred record; the dispreferred record’s sums are -11 and -10. Use \(\PreferenceScale=\) 0.2. These are hypothetical, dimensionless natural-log quantities.
Math: The reference-relative scores are 2 and -1, giving margin \(\PreferenceMargin=\) 0.6, logistic score \(\SigmoidFunction(\PreferenceMargin)=\) 0.646, and pair loss 0.437, down from 0.693 at \(\PreferenceMargin=0\). The saturation factor \(\SigmoidFunction(-\PreferenceMargin)\) falls from 0.5 to 0.354, while the full coefficient \(\PreferenceScale\SigmoidFunction(-\PreferenceMargin)\) falls from 0.1 to 0.0709. At \(\PreferenceScale=\) 0.4, the same log-probabilities give \(\PreferenceMargin=\) 1.2, a logistic score of 0.769, a saturation factor of 0.231, and a full coefficient of 0.0926.
Systems insight: Changing \(\PreferenceScale\) changes both the margin and the gradient prefactor. In this fixed-record comparison, doubling it reduces the saturation factor from 0.354 to 0.231 but increases the full coefficient from 0.0709 to 0.0926. This local coefficient comparison establishes neither the eventual policy’s divergence from the reference nor the size of its optimizer step. The logistic score \(\SigmoidFunction(\PreferenceMargin)\) rises with \(\PreferenceScale\) for the same log-probabilities, reflecting the model’s fit to this comparison rather than the operational success probability of the repair. Furthermore, the margin can increase either by boosting the likelihood of the preferred trajectory or by suppressing the dispreferred trajectory; only closed-loop execution can verify whether the resulting policy generalizes or merely avoids the dispreferred path.
Scoring an entire trajectory spreads credit and blame across every divergent token in the interaction, whereas narrowing \(\PreferenceScoredPositionSet(\Trajectory)\) to the branch point where the competing tool calls or edits first diverge focuses the gradient on that decision. Because log-probability ratios can be positive or negative, summing over longer sequences does not bias the margin in one direction, but it does increase score variance. Teams commonly test two modifications to control length sensitivity: dividing each score by its selected-token count \(|\PreferenceScoredPositionSet(\Trajectory)|\), and adding an auxiliary supervised imitation loss on the preferred record \(\PreferredTrajectory\) to keep the policy from drifting into low-likelihood outputs. Both choices are validated on closed-loop trials.
16.8.2 Reference log-probability caching
Evaluating \(\ReferenceRelativeScore(\Trajectory;\ComparisonContext)\) online requires holding a frozen copy of \(\pi_{\text{ref}}\) in accelerator memory alongside the active policy and running reference forward passes on every preferred and dispreferred record. Because \(\pi_{\text{ref}}\) and the training dataset are fixed across a DPO run, the pipeline can precompute and cache the reference policy’s log-probabilities for all selected action tokens ahead of time—and when token masks and aggregation rules are also fixed, storing only the scalar sum \(\ReferenceLogProbabilitySum(\Trajectory)\) avoids saving per-token arrays or materializing full-vocabulary logits. To keep cached sums from going silently stale across experiments, the cache key must bind every input that affects the number: the reference weights and any adapter configuration (section 16.10), the tokenizer and chat template, the exact token sequences and selected indices, the attention, position, and truncation settings, and the floating-point precision mode (such as FP16 vs. BF16). Changing data preprocessing or token selection invalidates cached sums even when the reference weights have not changed, so pipelines check cached values against live recomputations on sampled batches within a declared numerical tolerance (for example, \(10^{-5}\)) rather than assuming bitwise identity across accelerators or kernels.
Caching drops the reference weights and their forward passes from the training loop, trading preprocessing compute, disk space, and host-to-device bandwidth for accelerator memory. That removes the reference model’s weight and activation footprint, while the active policy’s weights, optimizer state (the running moments an adaptive optimizer keeps for each trained parameter), gradient buffers, and forward activations remain on the device; section 16.10 itemizes those remaining arrays. Preference optimization changes the loss function applied to a curated corpus, while on-policy labeling and fault injection change which states the corpus covers. When the bottleneck at serving time is instead the cost of carrying a long procedural prompt or running a large teacher model on every turn, policy distillation moves that behavior into a smaller student (section 16.9).
16.9 Policy Distillation
Consider an automated repair agent whose system prompt carries a detailed debugging checklist and worked examples showing how to inspect a stack trace before revising a tool call. In production, the agent resends this entire prompt prefix on every turn of every trajectory, consuming context-window capacity and raising time-to-first-token latency even though the procedural workflow belongs in the weights under table 1. Training the model to internalize the procedure lets the runtime drop the checklist from the prompt and reclaim those tokens on every call. Policy distillation transfers behavior from an expensive teacher configuration—whether a larger model family, a long procedural prompt, or a broader test-time search budget—into a cheaper student configuration, provided the teacher’s trajectories pass the same target admission filters as any other demonstration before entering training.
16.9.1 Context distillation
Context distillation7 moves stable procedural guidance from the prompt into the weights without removing the dynamic evidence the policy still needs at serving time. The pipeline first samples candidate demonstrations under the expanded prompt containing the full checklist, diagnostic routines, or few-shot examples; once admission filters verify the trajectories, the training loader strips the instructional guidance from the prompt and trains the student to emit the same selected actions from the shortened prefix. Because this transformation intentionally changes the conditioning input rather than replaying the teacher’s exact prompt, the lineage record preserves both prompt versions and the transformation rule (section 16.3).
7 Context Distillation: Askell et al. (2021) introduced the term for training a model without a long prompt to match the output distribution the same model produced with it, motivated by context-window limits and the compute required to process prompt prefixes repeatedly.
Askell, Amanda, Yuntao Bai, Anna Chen, et al. 2021. “A General Language Assistant as a Laboratory for Alignment.” arXiv Preprint arXiv:2112.00861, ahead of print. https://doi.org/10.48550/arXiv.2112.00861.
Stripping static procedural guidance is very different from stripping dynamic environment feedback such as an authorization refusal or a schema error that justified a corrective call: removing causal observations leaves the corrected action unmotivated by its prefix and trains the student to guess state changes without evidence. When a recorded trace such as trace_r19_01_redacted_v1 is shortened for context distillation, the pipeline records a new derived artifact with its own target review while keeping the historical runtime verdict attached to the original long-prompt trace.
The division of responsibilities in table 1 decides both what may be distilled and when to retrain. Stable procedures can move into weights, whereas current task facts, tool schemas, and authorization rules must stay in the prompt so an interface update does not break the model. By the same token, any rule that changes faster than the team can curate data, run fine-tuning, and clear the release gate—such as a business policy updated weekly—behaves as a contract in practice and stays in context.
The distilled student is evaluated with the abbreviated production prompt it will actually receive at serving time, checking that task acceptance and tool-calling accuracy hold while prompt tokens fall. Because end-to-end latency also depends on KV-cache prefix reuse, output length, tool execution time, and recovery turns, latency savings must be measured on closed-loop runs rather than inferred from prompt token counts alone: a distilled policy that reads fewer input tokens but takes an extra tool round-trip to recover from a bad argument can lose the prefill saving immediately.
16.9.2 Teacher distillation
When the teacher is a stronger model family or a multi-sample search harness rather than a longer prompt on the same model, teacher distillation transfers its verified trajectories into a smaller student. In classical knowledge distillation, a student trains on the teacher’s full probability distribution over classes8 (Buciluă et al. 2006; Hinton et al. 2015). Because trajectory collectors that call a frontier model or search harness over an API receive discrete sampled tool calls and messages rather than full-vocabulary logits, agent teacher distillation works by filtered imitation on verified trajectories rather than soft-target matching.
8 Soft Targets: Buciluă et al. (2006) compressed large ensembles by labeling data with the ensemble’s predictions and training a small network on those labels; Hinton et al. (2015) named the method distillation and trained on the teacher’s temperature-softened probabilities over every class. Soft targets require the teacher’s logits; a collector that sees only sampled messages has one token per position instead of a distribution.
Buciluă, Cristian, Rich Caruana, and Alexandru Niculescu-Mizil. 2006. “Model Compression.” Proceedings of the 12th ACM SIGKDD International Conference on Knowledge Discovery and Data Mining (KDD), 535–41. https://doi.org/10.1145/1150402.1150464.
Hinton, Geoffrey, Oriol Vinyals, and Jeff Dean. 2015. “Distilling the Knowledge in a Neural Network.” arXiv Preprint.
Zelikman, Eric, Yuhuai Wu, Jesse Mu, and Noah D. Goodman. 2022. “STaR: Bootstrapping Reasoning with Reasoning.” Advances in Neural Information Processing Systems (NeurIPS) 35: 15476–88.
Whether the teacher is a larger model, a test-time search loop (such as tree search or best-of-\(\CandidatePoolSize\) sampling), or the student’s own base model run with a long prompt, its candidate rollouts must pass the consumer’s admission checks before training; fluent reasoning or high self-reported confidence is never a substitute for verification. Training smaller models on verified trajectories from stronger teachers reliably improves tool-use accuracy and multi-step stability over zero-shot prompting (Chen et al. 2023; Zeng et al. 2023). The same filtered-imitation loop can even use the model’s own verified rollouts, as in STaR (Zelikman et al. 2022), subject to the self-selection narrowing risks of The Self-Improvement Loop. Because curation keeps only verified successes, the student trains on a cleaner distribution than the teacher’s raw rollouts and is not capped at the unguided teacher’s pass rate—though it will also absorb any inefficient tool habits or biases in the teacher that slip past the admission filter.
Distillation settles where the training trajectories come from, independently of the loss objective and the parameter artifact (section 16.2.2). Whether the run updates every weight or a compact adapter, the learned behavior is only as good as the admitted demonstrations; what the artifact choice changes is the memory and storage cost of holding the update, especially when a deployment maintains several specialized agent roles.
16.10 Parameter-Efficient Fine-Tuning
Full fine-tuning of an agent policy updates and stores a complete replica of every base parameter, requiring the training system to allocate memory for base weights, gradients, and optimizer state across the entire model. For an adaptive-moment optimizer under mixed precision, that state is a higher-precision master copy of each weight plus two running moments, so each trained parameter carries several times the bytes of its weight. Training a second specialized agent role, such as a diagnostic planner or an execution monitor, duplicates that resource footprint in full. Parameter-efficient fine-tuning avoids this multiplicative overhead by freezing the base model and training only a compact set of auxiliary parameters known as an adapter9 (Houlsby et al. 2019), which is distributed alongside the immutable base checkpoint. For systems operating multiple specialized agent roles, lightweight adapters streamline storage and version management. Adapters give modularity, not isolation. An adapter is only a set of weight matrices; like the rest of the weights, it shapes proposals and holds no permissions of its own, and roles served from one process share that process’s failure domain (section 16.10.2).
9 Adapter Modules: Houlsby et al. (2019) introduced adapters as small bottleneck layers inserted inside each Transformer block, updating only those bottleneck parameters per task. However, inserting sequential layers introduces serial computational dependencies during inference. Low-rank adaptation avoids this serving latency by parameterizing updates to match the exact dimensions of existing weight matrices, enabling operators to fold the adapter directly into the base weights when serving without adding serial layer overhead.
Houlsby, Neil, Andrei Giurgiu, Stanislaw Jastrzebski, Bruna Morrone, Quentin de Laroussilhe, Andrea Gesmundo, Mohammad Attariyan, and Sylvain Gelly. 2019. “Parameter-Efficient Transfer Learning for NLP.” Proceedings of the 36th International Conference on Machine Learning (ICML), Proceedings of machine learning research, vol. 97: 2790–99.
Low-rank adaptation (LoRA) expresses the weight update to a selected projection matrix through two low-rank factor matrices (Hu et al. 2021). Let \(\WeightMatrix_0\in\mathbb{R}^{\LayerOutputWidth\times \LayerInputWidth}\) denote a frozen base weight matrix acting on column-vector inputs. LoRA models the adapted matrix as
\[\WeightMatrix=\WeightMatrix_0+\frac{\AdapterScalingFactor}{\AdapterRank}\AdapterUpProjectionMatrix\AdapterDownProjectionMatrix,\qquad \AdapterUpProjectionMatrix\in\mathbb{R}^{\LayerOutputWidth\times \AdapterRank},\quad \AdapterDownProjectionMatrix\in\mathbb{R}^{\AdapterRank\times \LayerInputWidth}.\]
Here, the rank hyperparameter \(\AdapterRank \ll \min(\LayerInputWidth, \LayerOutputWidth)\) bounds the rank of the update, and the scaling factor \(\AdapterScalingFactor/\AdapterRank\) stabilizes optimization across different choices of \(\AdapterRank\). The two factor matrices contain \(\AdapterRank(\LayerInputWidth+\LayerOutputWidth)\) trainable parameters. When sizing a run, engineers must sum this parameter count over all targeted matrices; attention projections and multi-layer perceptron projections frequently differ in dimensionality and need not be square. Target matrix selection and rank operate as primary capacity knobs, tuned directly against empirical validation performance.10
10 Low Intrinsic Rank: Hu and colleagues motivated LoRA by hypothesizing that weight updates during downstream adaptation lie on a manifold with low intrinsic dimension, showing that small ranks (\(\AdapterRank \in \{4, 8\}\)) matched full fine-tuning on standard language benchmarks (Hu et al. 2021). That empirical finding reflects the characteristics of those benchmark tasks rather than a universal theoretical guarantee. Fine-tuning an agent policy on novel tool invocation protocols or structured error recovery often demands specialized capacity, making systematic rank sweeps essential under the evaluation contract of this chapter.
Hu, Edward J., Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, and Weizhu Chen. 2021. “LoRA: Low-Rank Adaptation of Large Language Models.” arXiv Preprint arXiv:2106.09685.
Initializing \(\AdapterDownProjectionMatrix\) from a Gaussian distribution and setting \(\AdapterUpProjectionMatrix=0\) ensures that \(\FiniteChange\WeightMatrix = \frac{\AdapterScalingFactor}{\AdapterRank}\AdapterUpProjectionMatrix\AdapterDownProjectionMatrix\) is initially zero, allowing training to begin from the exact forward mapping of the pretrained base model. On the initial backward step, the loss gradient backpropagates directly into \(\AdapterUpProjectionMatrix\), whereas the initial gradient with respect to \(\AdapterDownProjectionMatrix\) evaluates to zero because it is scaled by \(\AdapterUpProjectionMatrix=0\). Freezing \(\WeightMatrix_0\) eliminates the memory required for base gradients and optimizer state, yet the backward pass still evaluates matrix-vector products through \(\WeightMatrix_0^T\) to propagate activation gradients down to earlier transformer layers. The resulting gradient with respect to the layer input combines contributions from both the frozen base path and the adapter bypass, as given in LoRA matrix decomposition.
16.10.1 Parameter-state inventory
Training 0.10 percent of the parameters still keeps 12.6 percent of the state.
Because \(\WeightMatrix_0\) participates in every forward and backward pass, freezing the base model does not shrink GPU memory in proportion to the trainable parameter count: training 0.1 percent of the parameters still keeps 100 percent of the frozen base weights in device memory. Engineers size training hardware by counting the bytes of each resident array—frozen base weights, adapter weights, gradients, master weights, and optimizer moments—as notebook 16.5 does for a full fine-tuning run versus a LoRA run under mixed precision. Dynamic activations and temporary kernel workspaces are budgeted separately, because their size scales with sequence length, batch layout, and attention implementation rather than parameter count.
Napkin Math 16.5: What adapter training removes from the inventory
Problem: A hypothetical model contains 8,000,000,000 base parameters. The run attaches rank-16 adapters to 2 matrices in each of 32 layers. Each selected matrix maps an input width of 4,096 to an output width of 4,096. No other parameters are trained in the adapter case. What parameter-state memory does this save compared with training every base parameter?
Table 10: A declared parameter-state inventory: The count includes only the listed resident arrays, not the full peak memory of a training step.
For each trainable parameter, assume a weight using 2 bytes, a gradient using 2 bytes, a separate higher-precision master weight that the optimizer updates, using 4 bytes, and two optimizer moments (running averages of the gradient and its square) each using 4 bytes. Frozen parameters retain only their 2 bytes weight copy. This is the mixed-precision accounting of adaptive-moment optimization (Kingma and Ba 2015; Micikevicius et al. 2018), totaling 16 bytes per trainable parameter as modeled by Rajbhandari et al. (2020); it is one declared implementation convention, not a requirement of every optimizer or framework. Count one unsharded resident copy of each array; display capacities in decimal MB and GB.
Math: Each adapted matrix adds \(\AdapterRank(\LayerInputWidth+\LayerOutputWidth)\) parameters, so the adapters contain 32 \(\times\) 2 \(\times\) 16 \(\times\) (4,096 \(+\) 4,096) \(=\) 8,388,608 parameters, 0.10 percent of the base. The convention assigns 16 bytes of state to each trainable parameter and 2 bytes to each frozen one. Table 10 inventories the resulting arrays.
| Resident component | Adapter-training inventory |
|---|---|
| Frozen base weights | 16 GB |
| Trainable adapter weights | 16.78 MB |
| Adapter gradients | 16.78 MB |
| Adapter master weights | 33.55 MB |
| Two adapter optimizer moments | 67.11 MB |
Result: The adapter arrays total 134.22 MB. Including the frozen base gives approximately 16.13 GB, compared with 128 GB for full-parameter training under the same convention, a reduction of approximately 111.87 GB, or 7.9×. The saving comes from dropping gradients, master weights, and optimizer moments for the frozen base.
Systems insight: Even though adapter training updates only 0.10 percent of the parameters, the system must still allocate 12.6 percent of the full-parameter memory state. The frozen base weights remain resident throughout execution and constitute 99.2 percent of what remains. Consequently, total parameter-state savings are bounded by the ratio of trainable to frozen bytes per parameter (here 16 bytes against 2 bytes, so no choice of rank can push the reduction past that ratio). Rank adjustments exert virtually no influence on this footprint; doubling the rank increases memory by only 134.22 MB. The primary engineering lever for shrinking adapter memory state is therefore the numerical precision of the frozen base itself. Dynamic activations and execution workspaces, required under both training paradigms, must be provisioned on top of this subtotal.
Kingma, Diederik P., and Jimmy Ba. 2015. “Adam: A Method for Stochastic Optimization.” 3rd International Conference on Learning Representations (ICLR).
Beyond persistent parameter state, a complete memory plan must budget for forward activations (whether retained in high-bandwidth memory or selectively recomputed), temporary kernel scratchpads, collective communication buffers, allocator fragmentation, and operating headroom. Sharding and offloading distribute these resident arrays across cluster nodes or host memory tiers, trading network or PCIe interconnect bandwidth for device memory headroom. When on-device capacity is severely constrained, quantizing the frozen base weights directly targets the dominant state component, as demonstrated by QLoRA (quantized LoRA) (Dettmers et al. 2023), trading runtime dequantization arithmetic and metadata storage for lower memory occupancy. Because dynamic activation footprints and kernel workspaces scale with context length, batch structure, and attention kernels, engineers should profile these dynamic terms directly on the target hardware cluster before finalizing resource sizing.
16.10.2 Adapter release and compatibility
Because low-rank factors \(\AdapterUpProjectionMatrix\) and \(\AdapterDownProjectionMatrix\) express weight deltas relative to a specific base matrix \(\WeightMatrix_0\), mounting a trained adapter onto a silently updated base checkpoint corrupts model outputs without raising a tensor-shape error. Releasing an adapter for a specialized agent role therefore requires shipping a configuration manifest alongside the low-rank weights that pins the parent base model version, target projection names, rank \(\AdapterRank\), scaling factor \(\AdapterScalingFactor\), tokenizer revision, prompt templates, numerical precision, and links to the training corpus and evaluation audit. At serving time, the runtime must also enforce KV-cache prefix compatibility: cached prefix states computed under one adapter cannot be reused by another adapter that modifies the layers that produced them.
↩︎ The rule that cached key-value prefix state is reusable only under the same base weights and active adapters is established in Prefix compatibility.
Sheng, Ying, Shiyi Cao, Dacheng Li, Coleman Hooper, Shuo Chen, Daxin Jiang, Christopher Rodriguez, et al. 2024. “S-LoRA: Serving Thousands of Concurrent LoRA Adapters.” Proceedings of the 7th Conference on Machine Learning and Systems (MLSys).
Chen, Liyuan, Zihao Ye, Yonghao Wu, Danyang Zhuo, Luis Ceze, and Arvind Krishnamurthy. 2024. “Punica: Multi-Tenant LoRA Serving.” Proceedings of the 7th Conference on Machine Learning and Systems (MLSys).
Keeping adapters decoupled from the base weights lets a serving engine host several agent roles over one shared base model in accelerator memory, batching concurrent requests across distinct adapters while running the shared base projections in common matrix multiplications (Sheng et al. 2024; Chen et al. 2024). While this layout amortizes base-weight memory across roles, each role still pays for its own factor matrices and dynamic activations, and sharing one serving process across roles provides no fault or authority isolation between them: a bad adapter or anomalous input shares the process’s resource limits, and distinct roles still need distinct runtime capability grants. When a deployment serves only a single role, folding the adapter directly into the base weights (\(\WeightMatrix = \WeightMatrix_0 + \frac{\AdapterScalingFactor}{\AdapterRank}\AdapterUpProjectionMatrix\AdapterDownProjectionMatrix\)) removes per-token adapter kernel overhead and yields a single dense checkpoint.
If a newly deployed adapter regresses in production, rolling back the routing configuration sends new task admissions back to the prior verified release. Model rollback, however, cannot undo tool side effects already committed to external systems (database writes, network calls, or file edits) or settle in-flight trajectories on its own, so the release policy must declare whether active runs finish under their starting adapter or hand over to the fallback policy under the runtime’s recovery rules. Whether an update is packaged as an adapter or a dense checkpoint, deciding if it is safe to promote requires closed-loop evaluation.
16.11 Evaluating a Fine-Tuned Agent
A fine-tuned checkpoint can reach a low loss on a repair corpus yet fail to resolve actual connection-pool leaks in production. Training configurations and data manifests record how an artifact was produced, but only closed-loop trials reveal how the policy navigates multi-step execution when its own actions dictate future observations. The task acceptance contract defined in The Task Acceptance Contract governs these trials, addressing the specific failure modes and behavioral shifts that weight updates introduce.
16.11.1 Comparable loss and closed-loop performance
Held-out token loss measures prediction accuracy under teacher forcing, where recorded prefixes are supplied at every step, whereas closed-loop task completion runs the policy on the context history its own actions create under a fixed time and token budget. Even as an offline check for optimization bugs, token losses from two training runs are comparable only when computed under an identical measurement contract: if one run uses an all-token loss and another masks observations, their dashboard losses average over different token sets. In notebook 16.1, an unmasked causal loss averages across 16,383 prediction terms—of which read-only observations make up 77.1 percent of the stored input—whereas the masked objective averages over the 800 selected response tokens. Comparing raw training losses across those two runs rewards the unmasked checkpoint for predicting repetitive compiler boilerplate rather than choosing better actions. Valid offline comparison requires recomputing held-out loss on a common validation split under the exact same prompt contexts, tokenizer, serialization template, target mask, and normalization rule, logging both the active-target count and the split version; if two checkpoints use different tokenizers, their per-token losses do not even share a common unit.
Example 16.1: A lower loss is not yet a release decision
Scenario: In an evaluation bake-off, two candidate checkpoints are trained on the same curated trajectories: one with an all-token loss and the other with a selected-target response mask. The run trained on all tokens displays a substantially lower dashboard loss.
Diagnosis: The two training metrics calculate averages over incompatible target sets. The all-token run lowers its average by predicting the repetitive parts of environment observations (echoed schemas, boilerplate, repeated paths), whereas the masked run evaluates only policy tokens. Evaluating both checkpoints on an identical held-out validation set using matched response masks makes their loss values comparable. Even a checkpoint with superior held-out loss, however, is evaluated solely on teacher-forced prefixes; during autonomous rollouts, compounding errors can still cause it to invoke invalid tools, misinterpret unexpected exceptions, or blow past its step budget.
Systems lesson: Rank candidate checkpoints offline using an identical held-out evaluation loss, but base promotion and release decisions exclusively on closed-loop task acceptance. A checkpoint boasting lower loss remains unpromotable if it fails the recovery criteria specified in the release contract. When modifying training data and loss masks simultaneously, run controlled ablation trials to isolate the source of any observed performance shift.
16.11.2 Task and behavior measurements
Because neither held-out loss nor a single syntax check can show whether a weight update improved repair behavior without breaking other tasks, the release harness evaluates a candidate against the full policy stack it would replace across five complementary measurements. Budgeted task acceptance is the primary release outcome, measured under the declared criteria, runtime feedback rules, pinned fixtures, and resource envelopes of the task acceptance contract. Every initiated trial is counted, so runs that stall or time out at the observation cutoff are recorded as failures or unresolved verdicts rather than dropped. When the production deployment wraps the model in a multi-candidate selection or repair loop, the harness evaluates that composite loop alongside diagnostic views such as single-attempt pass rate (pass@1) and repeated-run consistency (pass\(^\IndependentAttemptCount\), Potential and repeated-run consistency).
Call validity checks whether generated tool calls satisfy the active interface contract by running them through the production parser and schema validator. Rather than collapsing every bad call into a single failure rate, the harness separates unparsable outputs, missing required fields, type mismatches, and truncated proposals from downstream tool execution errors (table 11).
| Observed category | Evidence to inspect | Competing explanations to test |
|---|---|---|
| Malformed serialization | Delivered template, raw output, parser contract, and selected training targets | Defective examples, template mismatch, generation behavior, or capture failure |
| Schema violation | Schema actually delivered for this invocation and the complete proposed arguments | Stale context, inadequate schema coverage, a serialization mismatch, or an incorrect validator |
| Undeclared tool | Tool inventory delivered to the model and any intervening configuration change | Learned tool-name preference, missing interface context, or inconsistent serving configuration |
| Truncated proposal | Output limit, stop status, transport evidence, and target termination policy | Resource exhaustion, interruption, decoding configuration, or unsuitable targets |
Schema robustness tests whether the fine-tuned policy generalizes to modified tool interfaces instead of falling back on memorized argument names, evaluating familiar schemas separately from renamed schemas and unseen tool families (Qin et al. 2024) with mock tools and call-site checks updated in lockstep (section 16.7). Over \(\ScoredToolCallCount\) proposed calls \(\ScoredToolCall_i\), each checked against its active schema set \(\ActiveSchemaSet_i\), schema-following accuracy is the fraction of calls that name a declared tool and satisfy its type schema:
Qin, Yujia, Shihao Liang, Yining Ye, Kunlun Zhu, Lan Yan, Yaxi Lu, Yankai Lin, et al. 2024. “ToolLLM: Facilitating Large Language Models to Master 16000+ Real-World APIs.” Proceedings of the 12th International Conference on Learning Representations (ICLR).
\[\text{SFA}=\frac{1}{\ScoredToolCallCount}\sum_{i=1}^{\ScoredToolCallCount}\EventIndicator\{\operatorname{tool}(\ScoredToolCall_i)\in\ActiveSchemaSet_i\ \wedge\ \ScoredToolCall_i\models\ActiveSchemaSet_i\}.\]
Parametric recall error is the fraction of calls that emit parameter names from the training key set \(\TrainingSchemaKeySet\) that do not appear in the delivered schema \(\ActiveSchemaSet_i\):
\[\text{PRE}=\frac{1}{\ScoredToolCallCount}\sum_{i=1}^{\ScoredToolCallCount}\EventIndicator\{\exists \ToolArgumentKey\in\operatorname{keys}(\ScoredToolCall_i):\ToolArgumentKey\in\TrainingSchemaKeySet\wedge \ToolArgumentKey\notin\operatorname{params}(\ActiveSchemaSet_i)\}.\]
When an adapted policy’s SFA rises on familiar schemas while its PRE climbs on renamed schemas, the weights are recalling training-set argument keys rather than reading the prompt. Before blaming the weights, however, the release audit checks the serving prompt to confirm that the old key was not accidentally left in the template.
Resource use prices what the policy spends to achieve that acceptance rate, tracking wall-clock latency, turn counts, generated tokens, and API charges through the effective cost per accepted task \(\CostPerAcceptedTask\) (equation 1.2), which divides the total spend across all initiated attempts (including rejected calls, retries, verifier runs, and sandbox cleanup) by the number of accepted solutions. Because \(\CostPerAcceptedTask\) divides by the acceptance rate, a fine-tuned candidate that takes an extra diagnostic turn per attempt can still lower \(\CostPerAcceptedTask\) if those turns prevent failed runs; if a candidate completes zero tasks in a trial set, the report lists total spend beside the zero-acceptance count rather than an undefined ratio.
↩︎ Comparing policies under one fixed resource envelope, with time and tokens reported separately from monetary charges, is established in Matched-budget comparison.
11 Catastrophic Forgetting: McCloskey and Cohen (1989) showed that a connectionist network trained sequentially on a new task can lose an earlier one almost entirely, because both share the same weights. A repair-trained policy shares its weights with every other behavior the agent still needs, so retention is measured on the release’s required task set rather than assumed from a narrow training corpus.
McCloskey, Michael, and Neal J. Cohen. 1989. “Catastrophic Interference in Connectionist Networks: The Sequential Learning Problem.” Psychology of Learning and Motivation 24: 109–65. https://doi.org/10.1016/S0079-7421(08)60536-8.
Capability retention checks that specializing the weights on repair traces has not degraded the other tasks the deployment serves, using a regression suite drawn from the agent’s full operational scope rather than generic language benchmarks alone. When performance drops on non-repair tasks, engineers check both the serving template and the training mixture before attributing the drop to catastrophic forgetting11 (McCloskey and Cohen 1989), and re-evaluate any mitigation—such as mixing in general trajectories, lowering the learning rate, or switching to a LoRA adapter—on the same five-part scorecard.
16.11.3 Controlled comparison and conditional release
To attribute changes on those five measurements to the weight update rather than to a coincidental prompt or schema edit, the release protocol varies one factor at a time under the fixture contract of Task fixtures. Each comparison pins task fixtures, acceptance criteria, resource budgets, and random seeds across repeated trials, binding every score to an artifact manifest that records the checkpoint or adapter hash, base model hash, tokenizer version, serialization template, decoding parameters, and harness commit.
The leakage controls of Evaluation leakage govern contamination here as well. Re-prompted attempts, redacted traces, and paraphrased variants inherit their source task’s split; checkpoints and tuning mixtures are selected on validation splits alone; and a test suite whose failures were fed back into training or prompt tuning has become development data and must be replaced by a fresh split. Benchmarking on familiar repositories or APIs reflects legitimate deployment requirements, but a claim that the fine-tuned policy generalizes demands held-out schema and environment families.
Controlled ablations compare the fine-tuned checkpoint against the incumbent policy, as well as against the prompt-engineering, schema-refactoring, and runtime-guardrail alternatives prioritized by the intervention ladder. Because the candidate and the incumbent can run on the same tasks, the comparison is a paired one, and the paired-outcome discipline (an interval on the difference, never two overlapping single-arm intervals) applies unchanged. Reading several ablations at once adds one correction beyond a single pair. If each comparison is tested at the same significance level, the chance that at least one of them shows a false win grows with their number, so the per-comparison level must be tightened; dividing it by the number of comparisons is the simplest such rule. A release claim resting on one seed or one favorable run is a sample of one.
↩︎ Paired comparison of two configurations on the same tasks, with intervals on the difference and the rule for splitting trials between new tasks and repeats, is established in Statistical Evaluation Rigor.
While a lower common validation loss confirms that training optimized the intended objective, only closed-loop task acceptance, recovery resilience, and load-tested serving configurations provide defensible release evidence. A fine-tuned candidate enters the same declared release gate as any other change (The declared decision rule), approved for an exact artifact hash and operating envelope under principle 12. Even after a fine-tuned model passes release qualification, external safety boundaries remain non-negotiable: runtime schema validators, permission sandboxes, token budgets, and supervisory recovery monitors must remain active in production.
16.12 Case Study: Training the Repair Policy
The recovery record that opened the chapter can now become a training artifact. The behavior to teach is narrow. When the runtime refuses an invalid tool proposal, the policy should inspect the schema, issue a corrected invocation that satisfies the interface, and end with a bounded report. Consider an illustrative supervised fine-tuning run that ingests repair_demonstrations_v1 to produce candidate policy trained_policy_v1. That artifact is a set of weights, distinct from ws_r20, the workspace the demonstration produced.
Each of the three adaptation choices is settled by evidence the chapter has already developed. The objective is imitation because the record supplies a demonstrated correction, not an admissible comparison; e2 and e4 were produced under different information, so they cannot form a preference pair (section 16.8). The behavior source is the admitted recovery corpus. It covers the malformed-argument refusal the recording happened to contain, not the version conflict that section 16.6 showed a policy trained on it would meet, so that uncovered state is the first candidate for a labeling round or a fault-injected record, and the first evaluation of trained_policy_v1 must include trials that reach it. The artifact is a full-parameter update, chosen for two reasons the chapter has developed. The repair policy is a single role, so none of the multi-role sharing that makes adapters attractive applies yet. The lesson is also a new recovery behavior, the kind for which a small rank is not known in advance to suffice, so the full update serves as the reference that any later adapter candidate must match. The choice has a price. On the hypothetical model of notebook 16.5, it keeps 128 GB of parameter state where an adapter run would keep 16.13 GB. A rank sweep whose adapter matches the full update’s paired acceptance, or a second role that must share the base, would flip the decision toward the adapter, which would then be a different candidate with its own evidence. Changing any one of the three choices produces a configuration that needs its own evidentiary trail.
Table 12 details the formal manifest bindings for this candidate policy. In production workflows, recording version identifiers alone is insufficient; a complete training record must capture the exact artifact contents, serialization parameters, and optimization settings required to reproduce and audit the resulting model.
| Manifest entry | Binding for this candidate |
|---|---|
| Corpus and member lineage | repair_demonstrations_v1, including trace_r19_01_redacted_v1 and its admission, transformation, task-group, and split records |
| Inputs and response selection | serialization_v1 implements recovery_targets_v1 with the named tokenizer, template, tool formatting, input transformations, and response boundaries |
| Prediction targets | target_mask_v1 aligns the selected policy-response tokens with their preceding-position predictions; e2, e5, and tool observations remain context rather than imitation targets |
| Update configuration | The identified base checkpoint, imitation objective, loss aggregation, sampling policy, attention and position configuration, optimizer settings, and training-state provenance |
| Candidate and comparison | trained_policy_v1, its serving configuration, the identified baseline, and the held-out evaluation protocol and findings required for promotion |
The mask and the serialization follow directly from the earlier sections. target_mask_v1 selects the policy-response spans of e3 (the schema query), e4 (the corrected tool call), and e6 (the final report), each with its authentic end-of-turn token. The malformed proposal and its pre-dispatch refusal at e2, the schema returned to e3, and the check summary at e5 carry \(\PositionLossMask_{t}=0\); they stay in the prefix as conditioning and receive no direct loss (table 2). Selecting e2 would teach the malformed call, and selecting the returned schema or e5 would teach the policy to write tool results instead of requesting them. Whether serialization_v1 may use one concatenated transcript or must serialize each selected invocation separately depends on whether the runtime compacted history between e2 and e6. If it did not, each selected response’s prefix is a prefix of one transcript and a single example is faithful. The manifest records which case held, from the delivered-input manifests, rather than assuming it, because a record that crossed the watermark would need one example per invocation at up to \((\TurnHorizon+1)/2\) the stored tokens (section 16.3.3). If the run packs these examples with others, block-diagonal attention and per-example label shifts keep each example’s context its own (section 16.5).
The release gate for trained_policy_v1 follows from the lesson it was meant to learn. Lower held-out loss on the repair corpus shows only that the objective was optimized. Promotion needs budgeted task acceptance on held-out repair tasks under the same runtime and fixtures as the incumbent, compared as paired trials, with \(\CostPerAcceptedTask\) reported for both so that a slower but more reliable candidate is credited correctly. Because the lesson is “consult the delivered schema,” the most direct test is schema-following accuracy and parametric recall error on a renamed-argument set: a candidate whose PRE rises there may be relying on the recorded interface rather than the delivered schema. The gate also needs trials that reach the version-conflict state the corpus does not cover, and a retention suite for the agent’s other required tasks. A candidate missing any of these measurements is held, however its loss compares (section 16.11).
Even when the evaluations justify a staged release, the proposal boundary stays where it was. Fine-tuning makes the policy’s proposals more reliable without making the model an execution authority. The runtime keeps enforcing schema validation, least-privilege tool permissions, process isolation, and hard timeout budgets, and the agent’s final report remains a claim the harness checks against workspace artifacts. The manifest is what turns one admitted recovery into a reproducible, auditable policy update; the runtime is what keeps that update’s mistakes bounded.
16.13 Fallacies and Pitfalls
A fine-tuning pipeline can optimize its nominal loss objective with complete numerical fidelity while departing from the intended operational behavior. In agent systems, where training corpora interweave multi-turn plans, tool invocations, and environment feedback, subtle flaws emerge at the interfaces between admitted evidence, context assembly, loss masking, and deployment evaluation. Each of the following mistakes produces a concrete failure mode, and each carries a quantifiable engineering penalty illustrated in this chapter’s worked examples.
Fallacy: Every model response in an accepted trajectory is a desirable training target.
When an agent successfully recovers from an error, the recorded trajectory necessarily contains the initial flawed or refused proposal alongside the eventual correction. That initial misstep is essential as conditioning context, because it explains why the correction was required, but it is the one action the policy must never learn to emit. In the inventory of notebook 16.1, the refused positions number 160 tokens, 20 percent as many as the 800 selected response targets. An unselective mask over all assistant turns therefore puts 160 of the 960 response positions to work reproducing the malformed call. That share is a minority of the response weight, yet all of it rewards the refused proposal directly, in every recovery record that contains one. Systems must instead implement target admission masks (section 16.4): retain rejected proposals and environment observations in the prefix context so the agent learns to condition on failure signals, but restrict direct imitation loss to admitted response targets.
Pitfall: Treating efficient batching as a change with no learning consequences.
Packing multiple variable-length trajectories into a single contiguous buffer raises occupancy, but it alters training dynamics unless the execution kernel accounts for sequence boundaries. In notebook 16.2, sequence packing raises token occupancy from 19.1 percent to 76.6 percent. However, applying a standard causal attention mask across that packed buffer allows subsequent examples to attend across episode boundaries to unrelated earlier sequences. Loss normalization alters optimization just as profoundly: standard token-level averaging gives the long example 32× the short one’s weight, whereas per-example averaging weights both examples equally. Always configure block-diagonal attention with per-example position coordinates and reset label shifts (section 16.5), explicitly declare loss reduction semantics, and verify real execution throughput on target accelerator kernels before claiming a speedup.
Fallacy: Recovery examples and schema perturbations make a fine-tuned agent robust to the states it will meet.
Synthesized recovery trajectories and perturbed schemas broaden training coverage, but they do not provide a blanket guarantee against distribution shift. Theoretical bounds on policy regret depend on where state error rates are evaluated. Over an execution horizon of 30 steps with a 2 percent per-step disagreement rate, theorem 16.1 bounds excess cost at 18 units on reference demonstration states, but at 0.6 units on states generated by the learner itself. Furthermore, that tighter bound holds only while the disadvantage of each individual deviation remains near one, a condition that fails as soon as one unexpected tool error can spoil the rest of the task. Robust agent systems require sampling and labeling states actually reached by the learner, maintaining end-to-end coherence between schema transformations, tool calls, and parser observations (section 16.7), and stress-testing the policy under held-out operational shifts.
Pitfall: Reading a preference pair’s logistic score as the probability that the preferred repair succeeds.
The implicit reward margin computed in direct preference optimization reflects pairwise discrimination on historical data, not absolute operational success. In notebook 16.4, identical trajectory log-probabilities yield a logistic score of 0.646 at \(\PreferenceScale=\) 0.2, yet that score rises to 0.769 when the preference scale doubles to \(\PreferenceScale=\) 0.4. A metric that shifts this much with an optimization hyperparameter measures relative fitting margin rather than task competence. Furthermore, that margin can widen because the objective suppresses the dispreferred trajectory rather than improving the execution likelihood of the preferred one. Reliable acceptance rates cannot be read off training curves; they must be measured directly through closed-loop environment trials (section 16.11).
Fallacy: Training a small fraction of the parameters shrinks training memory in proportion.
Parameter-efficient fine-tuning cuts the number of updated weights to a small fraction, but accelerator memory does not scale down linearly with the parameter fraction. In notebook 16.5, low-rank adapters account for 0.10 percent of total model parameters, yet the training job still requires 12.6 percent of the parameter-associated memory consumed by a full fine-tuning run. This 7.9× reduction is governed by bytes per parameter across base and adapter states rather than the parameter count alone. The frozen base weights must remain resident in memory throughout forward and backward passes, accounting for 99.2 percent of what remains, while dynamic activation buffers and workspace memory scale with sequence length for both approaches. Dimension adapter training hardware by accounting for frozen base representations and empirically measured activation peaks (section 16.10).
Pitfall: Promoting the artifact with the lowest recorded loss.
Selecting an agent checkpoint based purely on validation loss conflates prefix prediction with closed-loop execution competence. In notebook 16.1, an unmasked causal evaluation averages across 16,383 prediction terms per example. Tool outputs that the policy never emits occupy 77.1 percent of the stored input. In contrast, a selective objective evaluates only the 800 target positions where the model generates actions. A checkpoint with lower overall loss may simply fit the repetitive parts of that environment text more tightly without improving decision quality. Even when measured exclusively on action tokens, offline likelihood improvements do not guarantee that the policy will avoid deadlocks or handle schema variations during execution. Always qualify production candidates by measuring scoped task acceptance, effective cost \(\CostPerAcceptedTask\), schema compliance, and retained foundational capabilities (section 16.11), while maintaining strict runtime validation checks regardless of training performance.
16.14 Summary
Trajectory fine-tuning changes the policy itself, and every decision in this chapter serves one requirement, that the update teach the lesson curation admitted and nothing else. A trajectory mixes the policy’s responses with the environment’s replies, so a loss computed over the wrong positions, under the wrong prefixes, or with the wrong batch weighting can fall while the behavior it teaches drifts. Even a faithful update is fitted on recorded prefixes. Its value is settled only by closed-loop evidence on the states the policy reaches on its own, under the same runtime that bounds every proposal it makes.
The quantitative decisions of the chapter rest on the relations collected in table 13. Each row pairs a formal relation with the training, packaging, or release choice it governs and points back to the section where it is derived.
| Equation | What it bounds or decides | Where derived |
|---|---|---|
| \(\MaskedFinetuningLoss=-\frac{1}{\ActiveTargetTokenCount}\sum_{t}\PositionLossMask_{t}\log \Policy(\SequenceToken_t\mid \TokenPrefix)\) | Which positions receive direct loss and the coefficient \(1/\ActiveTargetTokenCount\) on each; an unmasked average dilutes it by \((\SequenceLength-1)/\ActiveTargetTokenCount\) | equation 1 |
| \(\TurnAppendLength\,\TurnHorizon(\TurnHorizon+1)/2\) stored tokens against \(\TurnHorizon\,\TurnAppendLength\) | Training-token price of one example per invocation, an expansion of \((\TurnHorizon+1)/2\) under approximately linearly growing delivered prefixes | section 16.3 |
| \(\TokenAveragedBatchLoss\) weights \(\LossFunction_{\BatchExampleIndex}\) by \({\ActiveTargetTokenCount}_{\BatchExampleIndex}/\sum_k {\ActiveTargetTokenCount}_k\); \(\ExampleAveragedBatchLoss\) by \(1/\BatchSize\) | How far a long trajectory outweighs a short one, and how segmentation shifts task weighting | equation 2, equation 3 |
| \(\text{score}_{i,j}=-\infty\) unless \(\operatorname{example}(\QueryBufferPosition)=\operatorname{example}(\KeyBufferPosition)\) and \(\KeyBufferPosition\le \QueryBufferPosition\) | Whether a packed buffer preserves each example’s visible context | equation 4 |
| \(\ExpectedTaskCostSum(\pi)-\ExpectedTaskCostSum(\pi^*)\le\min\{\TurnHorizon,\;\TurnHorizon^2\ReferenceStateDisagreementRate\}\) | Excess task cost of imitation whose error is measured on reference states | equation 5 |
| \(\ExpectedTaskCostSum(\pi)-\ExpectedTaskCostSum(\pi^*)\le\min\{\TurnHorizon,\;\DeviationDisadvantage \TurnHorizon\LearnerStateDisagreementRate\}\) | Excess cost when error is measured on learner states and each deviation costs at most \(\DeviationDisadvantage\); decides when on-policy labeling pays | equation 6 |
| \(\DpoLoss=-\mathbb E\log\SigmoidFunction\!\left(\PreferenceScale[\ReferenceRelativeScore(\PreferredTrajectory;\ComparisonContext)-\ReferenceRelativeScore(\DispreferredTrajectory;\ComparisonContext)]\right)\) | The preference objective over reference-relative scores of the selected spans | equation 7 |
| Pair gradient coefficient \(\PreferenceScale\SigmoidFunction(-\PreferenceMargin)\) with \(\PreferenceMargin=\PreferenceScale(\PreferredReferenceRelativeScore-\DispreferredReferenceRelativeScore)\) | The full scalar on the gradient difference; \(\SigmoidFunction(-\PreferenceMargin)\) separately measures saturation | equation 8 |
| \(\AdapterRank(\LayerInputWidth+\LayerOutputWidth)\) trainable parameters per adapted matrix | Adapter size; the run’s parameter state is then set by the bytes of the frozen base | section 16.10 |
| \(\text{SFA}\) and \(\text{PRE}\) over \(\ScoredToolCallCount\) proposed calls | Conformance to the delivered schema, and recall of training-schema keys the current schema lacks | section 16.11.2 |
Key Takeaways: Preserve the lesson from record to release
- Weights hold procedure, never contracts: Training changes what the policy proposes, not what the runtime permits. Stable procedure can move into the weights, while schemas, permissions, and any guidance that changes faster than a training and evaluation cycle stay in context, and enforcement stays in the runtime.
- A trajectory has two authors; train one: Observations stay in the input as context but receive no direct loss. An unmasked average dilutes each selected target by the inverse of its selected share of predicted positions, 20.5× in the worked inventory, and teaches the policy to invent tool results.
- A loader can change the lesson without changing the loss: Later evidence in a prefix, a shifted label, or a causal mask across packed examples each redefines the prediction problem. Normalization is a choice too, since token averaging gives a long trajectory 32× the weight of a short one in the worked batch.
- Where error is measured sets the horizon penalty: The same 2 percent disagreement over 30 steps bounds excess cost at 18 units on reference states and 0.6 on learner states when the deviation disadvantage is bounded by one, so long-horizon agents need labels on the states they actually visit.
- A preference score measures one comparison: A pair’s logistic score moves with \(\PreferenceScale\) and can rise by suppressing the dispreferred record, so it measures fit to that comparison, never the probability that a repair is accepted.
- Adapters cut state by bytes, not parameter fraction: Training 0.10 percent of the parameters still keeps 12.6 percent of the parameter state, because the frozen base stays resident. The lever is the frozen base’s representation, not the adapter rank.
- Release on closed-loop acceptance, not loss: A candidate is promoted on task acceptance, \(\CostPerAcceptedTask\), schema behavior, and retained capability under its full serving configuration. Rollback selects an earlier policy for future work and leaves effects already dispatched to settlement and compensation.
Verified trajectory post-training (principle 13) now has a concrete implementation obligation at the token level: preserve useful conditioning while aligning the objective with the selected predictions, rather than equating every recorded policy token with a target. At the record level, it connects the admitted lesson to a versioned update whose usefulness remains an evaluation question. Invariant closure (principle 1) stays outside that update, which is why the proposal boundary holds from the first training example to the release decision.
16.15 Historical Perspective and Further Reading
The core mechanism of this chapter is older than language models. Williams and Zipser named teacher forcing in 1989, training recurrent networks by feeding the recorded target back as the next input (Williams and Zipser 1989). At about the same time, ALVINN learned road following as a supervised mapping from camera images to a steering direction (Pomerleau 1988), the template for what control researchers came to call behavioral cloning. Both generations met the same obstacle. A policy trained on recorded prefixes is run on its own, and its first error moves it into inputs the recording never showed. Ross and Bagnell turned that obstacle into a quadratic bound in 2010 (Ross and Bagnell 2010), and DAgger answered it the following year by labeling the states the learner itself visits (Ross et al. 2011). Sequence-generation research reached the same diagnosis independently, with scheduled sampling (Bengio et al. 2015) and the name exposure bias (Ranzato et al. 2016).
Pomerleau, Dean A. 1988. “ALVINN: An Autonomous Land Vehicle in a Neural Network.” Advances in Neural Information Processing Systems (NIPS) 1: 305–13.
Bengio, Samy, Oriol Vinyals, Navdeep Jaitly, and Noam Shazeer. 2015. “Scheduled Sampling for Sequence Prediction with Recurrent Neural Networks.” Advances in Neural Information Processing Systems (NeurIPS 2015) 28.
Christiano, P. F., J. Leike, T. B. Brown, M. Martic, S. Legg, and D. Amodei. 2017. “Deep Reinforcement Learning from Human Preferences.” In Advances in Neural Information Processing Systems 30: Annual Conference on Neural Information Processing Systems 2017, December 4-9, 2017, Long Beach, CA, USA, edited by Isabelle Guyon, Ulrike von Luxburg, Samy Bengio, Hanna M. Wallach, Rob Fergus, S. V. N. Vishwanathan, and Roman Garnett. Curran Associates.
Preference learning has a statistical root. Bradley and Terry’s 1952 model of paired comparisons (Bradley and Terry 1952) underlies the reward models that Christiano and colleagues fit to human comparisons of agent behavior (Christiano et al. 2017), and InstructGPT combined supervised demonstrations with such a reward model to align a language model (Ouyang et al. 2022). DPO then showed that the separate reward model and reinforcement learning stage could be folded into a classification loss on policy/reference log ratios (Rafailov et al. 2023), the form this chapter extends from responses to selected trajectory spans.
The cost of adaptation has its own lineage. Model compression (Buciluă et al. 2006) and knowledge distillation (Hinton et al. 2015) made a large model’s behavior a training signal for a cheaper one, and context distillation applied the same idea to a prompt (Askell et al. 2021). Adapter modules (Houlsby et al. 2019), low-rank adaptation (Hu et al. 2021), and quantized frozen bases (Dettmers et al. 2023) shrank what each adaptation stores and trains, while the systems accounting of mixed-precision training (Micikevicius et al. 2018) and ZeRO (Rajbhandari et al. 2020) made the per-parameter memory of that training explicit. Agent fine-tuning joined these threads once models emitted tool calls inside interleaved reasoning and action traces (Yao et al. 2023; Schick et al. 2023), and trajectories generated by stronger models became training data for smaller ones (Chen et al. 2023; Zeng et al. 2023). What changes for agents is the environment: a trajectory has two authors, which is why masking, serialization, and on-policy coverage become engineering obligations rather than details.
Dettmers, Tim, Artidoro Pagnoni, Ari Holtzman, and Luke Zettlemoyer. 2023. “QLoRA: Efficient Finetuning of Quantized LLMs.” Advances in Neural Information Processing Systems (NeurIPS) 36: 10088–115.
Micikevicius, Paulius, Sharan Narang, Jonah Alben, Gregory Diamos, Erich Elsen, David Garcia, Boris Ginsburg, et al. 2018. “Mixed Precision Training.” International Conference on Learning Representations (ICLR).
Yao, Shunyu, Jeffrey Zhao, Dian Yu, Nan Du, Izhak Shafran, Karthik Narasimhan, and Yuan Cao. 2023. “ReAct: Synergizing Reasoning and Acting in Language Models.” International Conference on Learning Representations (ICLR).
Schick, Timo, Jane Dwivedi-Yu, Roberto Dessı̀, Roberta Raileanu, Maria Lomeli, Luke Zettlemoyer, Nicola Cancedda, and Thomas Scialom. 2023. “Toolformer: Language Models Can Teach Themselves to Use Tools.” Advances in Neural Information Processing Systems (NeurIPS) 36: 68539–51.
Chen, Baian, Chang Shu, Ehsan Shareghi, Nigel Collier, Karthik Narasimhan, and Shunyu Yao. 2023. “FireAct: Toward Language Agent Fine-Tuning.” arXiv Preprint arXiv:2310.05915, ahead of print. https://doi.org/10.48550/arXiv.2310.05915.
Zeng, Aohan, Mingdao Liu, Rui Lu, Bowen Wang, Xiao Liu, Yuxiao Dong, and Jie Tang. 2023. “AgentTuning: Enabling Generalized Agent Capabilities for LLMs.” arXiv Preprint arXiv:2310.12823, ahead of print. https://doi.org/10.48550/arXiv.2310.12823.
For further reading:
- Ross, Gordon, and Bagnell’s DAgger paper (Ross et al. 2011) states both regret bounds of theorem 16.1 and the reduction to online learning behind them; it is the shortest route to seeing why where error is measured matters more than how small it is.
- Rafailov and colleagues’ DPO paper (Rafailov et al. 2023) derives the preference loss from a reference-regularized objective, and reading the derivation shows exactly which assumptions a trajectory-span surrogate relaxes.
- Ouyang and colleagues’ InstructGPT paper (Ouyang et al. 2022) documents a complete demonstration-then-preference pipeline, including the human evaluations that compared the model after each stage.
- Hu and colleagues’ LoRA paper (Hu et al. 2021) motivates low-rank updates and reports experiments on which matrices to adapt and at what rank, the knobs this chapter says to set by evaluation.
- Rajbhandari and colleagues’ ZeRO paper (Rajbhandari et al. 2020) itemizes training memory by owner and precision, the same inventory discipline notebook 16.5 applies to an adapter run.
- Krell and colleagues’ packing paper (Krell et al. 2021) treats packing as an equivalence to prove rather than a throughput trick, the stance the chapter’s packed-versus-separate logit test takes.
Rafailov, Rafael, Archit Sharma, Eric Mitchell, Stefano Ermon, Christopher D. Manning, and Chelsea Finn. 2023. “Direct Preference Optimization: Your Language Model Is Secretly a Reward Model.” Advances in Neural Information Processing Systems (NeurIPS) 36: 53728–41. https://doi.org/10.52202/075280-2338.
Ouyang, L., J. Wu, X. Jiang, D. Almeida, C. L. Wainwright, P. Mishkin, C. Zhang, et al. 2022. “Training Language Models to Follow Instructions with Human Feedback.” Advances in Neural Information Processing Systems (NeurIPS) 35: 27730–44. https://doi.org/10.52202/068431-2011.
Rajbhandari, Samyam, Jeff Rasley, Olatunji Ruwase, and Yuxiong He. 2020. “ZeRO: Memory Optimizations Toward Training Trillion Parameter Models.” SC20: International Conference for High Performance Computing, Networking, Storage and Analysis, 1–16. https://doi.org/10.1109/sc41405.2020.00024.
16.16 Exercises
Every given below is a hypothetical problem input chosen for round arithmetic, not a measurement of any system. Each problem asks for a calculation and the training or release decision it supports.
Per-invocation serialization under a token budget (Basic). The serving policy of the repair agent omits earlier generated text before each later invocation, so every turn of a trajectory becomes its own training example (section 16.3). Assume that the delivered prefix at each turn still grows by approximately \(\TurnAppendLength\) tokens despite the omitted generated text. Trajectories run \(\TurnHorizon=40\) turns, each turn adds about \(\TurnAppendLength=500\) tokens, the corpus holds 2,000 admitted trajectories, and the run may store at most 200 million training tokens.
- Compute the stored tokens per trajectory under per-invocation serialization and under a single transcript, and the expansion factor between them.
- Compute how many complete trajectories fit the budget under each serialization, and what fraction of the corpus the per-invocation run can use.
- Name the serving-side change that would make a single transcript a faithful example, and state what the training pipeline must verify before relying on it.
Target count and dilution (Basic). A serialized repair example stores 2,000 prompt and schema tokens, 18,000 observation tokens, and 1,500 policy-response tokens. Of the response tokens, 300 belong to a refused proposal that curation retains as context only. The prompt and schema count includes the initial boundary; no preceding external seed is supplied. The example has no rationale and no padding.
- Compute \(\ActiveTargetTokenCount\) under equation 1, the selected share of the stored input, and the factor by which an unmasked token average would dilute each selected target’s coefficient.
- A loader instead selects every policy-response token. Compute the fraction of the imitation weight that lands on the refused proposal, and state what behavior that weight trains.
- A colleague expects masking to cut this example’s training compute by the dilution factor from part (a). Explain why the saving is far smaller, and identify the only work a loader can skip at masked positions.
Normalization and segmentation (Basic). A batch holds one long repair trajectory with 6,000 selected targets and 15 short file-lookup examples with 200 selected targets each.
- Compute the long trajectory’s share of the batch objective under token averaging (equation 2) and under example averaging (equation 3).
- Serialization instead splits the long trajectory into 20 per-invocation examples of 300 selected targets each, every example selecting only its own turn’s response. Recompute both shares.
- The short lookups exercise a capability the release must retain. Choose a normalization and a segmentation, and state what the choice costs the long task’s influence on the update.
Packing and the cross-example work a kernel can skip (Intermediate). Ten independent examples have stored lengths of 16,384, 8,192, 6,144, 4,096, 4,096, 2,048, 2,048, 1,024, 1,024, and 512 tokens. Each buffer holds 16,384 positions.
- Compute the occupancy of a padded batch that allocates the maximum length to every row.
- Pack the examples by first-fit decreasing, placing each example, longest first, in the first buffer with room. Report the buffers, the packed occupancy, and the factor by which packing reduces allocated positions.
- For the buffer holding the 8,192-, 6,144-, and 2,048-token examples, compute the fraction of causal score entries (the lower triangle, diagonal included) that block-diagonal attention (equation 4) sets to \(-\infty\). Explain what that fraction means for a variable-length attention kernel, and why the mask, not the skipped work, is the correctness requirement.
When on-policy labeling pays (Intermediate). Offline imitation of the repair agent reaches a disagreement rate of \(\ReferenceStateDisagreementRate=0.01\) on reference states over a horizon of \(\TurnHorizon=50\) decision steps. Three DAgger rounds would each collect 400 trajectories that run the full horizon and label all 50 visited states of each, and each label costs $0.05 to obtain and $0.25 to verify by a sandbox continuation. After aggregation, the retrained policy’s disagreement on its own states is \(\LearnerStateDisagreementRate=0.02\), and one deviation followed by reference recovery costs at most \(\DeviationDisadvantage=2\).
- Compute both bounds of theorem 16.1, their ratio, and the reference-state error offline imitation would need to match the learner-state bound.
- Compute the labeling cost of the three rounds, and the horizon at which the reference-state bound saturates at its cap.
- A remediation step for the remediation agent changes live configuration with no adequate reversal, so one deviation there can spoil the rest of the task. Recompute the learner-state bound with \(\DeviationDisadvantage=\TurnHorizon\) at the same \(\LearnerStateDisagreementRate\), and state whether further labeling or a runtime mechanism should carry that step’s reliability.
Reading preference margins (Intermediate). Two checkpoints trained with equation 7 at \(\PreferenceScale=0.1\) are compared on one admitted repair pair. On the selected spans, the reference log-probability sums are \(-22\) for the preferred record and \(-24\) for the dispreferred one. Checkpoint A assigns \(-20\) and \(-25\); checkpoint B assigns \(-23\) and \(-30\). The sums exclude a shared 1,000-token prefix with identical conditioning, whose inclusion would add \(+0.5\) to each record’s score under either checkpoint.
- Compute each checkpoint’s scores \(\PreferredReferenceRelativeScore\) and \(\DispreferredReferenceRelativeScore\), margin \(\PreferenceMargin\), pair loss, saturation factor \(\SigmoidFunction(-\PreferenceMargin)\), and full gradient coefficient \(\PreferenceScale\SigmoidFunction(-\PreferenceMargin)\).
- Show that including the shared prefix leaves every quantity in part (a) unchanged.
- A team proposes promoting checkpoint B because its pair loss is lower. Critique the proposal from the scores, and name the measurement that decides between the checkpoints.
Sizing an adapter run on a fixed memory budget (Intermediate). A hypothetical base model has 20 billion parameters. The run attaches rank-32 adapters to four matrices in each of 48 layers, and each adapted matrix has 6,144 input and 6,144 output coordinates. Use the chapter’s convention of 16 bytes of state per trainable parameter and 2 bytes per frozen parameter, with decimal GB. The accelerator offers 48 GB, of which 16 GB is reserved for activations and workspace.
- Compute the adapter parameter count, its fraction of the base, and the parameter state of full fine-tuning and of the adapter run. Decide whether either fits the remaining budget.
- Store the frozen base at 4 bits per parameter plus 0.5 bits per parameter of quantization metadata. Recompute the adapter run’s parameter state, decide whether it fits, and compute what doubling the rank would add.
- Five agent roles each need their own adaptation. With every weight stored at 2 bytes per parameter, compare five full-weight copies with one shared base plus five adapters, and state what the shared base does not provide between roles.
Context distillation for the support agent (Advanced). The support agent runs short sessions under concurrency, so its binding constraint is per-turn input and response latency rather than the long observations that dominate the repair corpus. Each turn’s input carries a 3,000-token procedural checklist, a 1,500-token refund policy that changes weekly, and about 1,500 tokens of conversation and customer record. Sessions average four turns, and the deployment serves 2 million sessions per month. Input tokens cost $2 per million, and a cached prefix is billed at one tenth of that rate. Each redistillation, a training run plus the release evaluation, costs $15,000, and the checklist is revised once a month.
- Using table 1, decide which of the three input segments may be distilled into the weights and which must stay in context.
- Compute the monthly input tokens and spend that distilling the checklist removes, first with no prefix caching and then with the stable prefix always billed at the cached rate. Compute the net monthly saving after redistillation in each case, and the number of checklist revisions per month at which the uncached saving reaches zero.
- Decide whether to distill under each serving configuration. State what must be measured, rather than inferred from the token reduction, before claiming a latency gain, and contrast the constraint that decides this case with the one that governs masking for the repair agent.
Critique a release proposal (Advanced). A team proposes promoting a fine-tuned repair policy on the following evidence. Its all-token training loss fell. On 500 closed-loop trials each, the baseline was accepted in 60 percent of trials at a mean attempt cost of $1.20, and the candidate in 66 percent at $1.50. On familiar interfaces, schema-following accuracy rose from 0.92 to 0.95. On a renamed-argument test set, the baseline scored an SFA of 0.90 with a PRE of 0.02, and the candidate an SFA of 0.81 with a PRE of 0.12.
- Compute \(\CostPerAcceptedTask\) for each policy in the per-attempt form of equation 1.2, and the largest mean attempt cost at which the candidate would match the baseline.
- Using a normal approximation for independent trials, compute the standard error of the difference in acceptance rates and the difference measured in standard errors.
- Interpret the pattern in SFA and PRE, and redesign the release gate. Name what the decision should be now, which measurements must be added or repeated, and which training change to test first.
- Suppose the two arms in fact ran on the same 500 tasks. State what record the paired comparison of section 16.11 needs beyond the two acceptance rates, and explain how the count of discordant tasks could make the difference in part (b) more or less convincing than the independent-trials calculation suggests.
What’s Next: From admitted responses to reward-scored interaction
When a task offers an outcome check but no demonstrated response for each decision, the adaptation framework must evolve beyond supervised imitation and static comparisons. Every update in this chapter needed a demonstrated or compared response, so the version conflict that
trained_policy_v1 never saw could be taught only by labeling that state or injecting it. Reinforcement Learning from Verifiable Rewards starts from that checkpoint and lets the repair agent’s acceptance check score fresh attempts on the states the policy reaches by itself, turning each verdict into a weight update. The check then stops being only a release gate and becomes the objective. That shift raises questions this chapter’s objectives never posed: whether the reward resists the shortcuts optimization will find, how one terminal verdict assigns credit across many turns, how rollout records stay faithful to the policy that sampled them, and what the rollout and verifier fleet costs. Containment and closed-loop release evidence apply as before.