You Can’t Reweight a Dead End: TRD and the Prefix Failure Problem

TL;DR for operators

The paper’s main message is simple: if a reasoning model has already walked into a dead end, per-token distillation often keeps supervising it from inside the dead end. A clever loss cap is not a map. A top-k filter is not a tow truck.

Trajectory-Refined Distillation, or TRD, repairs the student’s own rollout before using it for distillation. The pipeline is: sample the student’s attempt, ask a teacher or privileged self-teacher to rewrite the trajectory into a better one, then train on the refined trajectory rather than on the original failed rollout. The technical contribution is not “better prompting”, although prompts are used. It is the shift from token-level correction to trajectory-level correction.

The strongest evidence is on Qwen3 math reasoning. In on-policy self-distillation, TRD is reported as best across the evaluated math benchmarks at two model scales, with particularly large Pass@16 gains on AMOBench: Qwen3-8B rises from 41.0 to 61.5, a +20.5 point absolute gain. In on-policy distillation with a separate Qwen3-8B teacher, TRD also improves most math results and preserves stronger base capabilities better than raw token-level KL baselines. Code is less tidy: HumanEval+ and MBPP+ improve or match strong baselines, but LiveCodeBench does not beat the base model.

For business use, the inference is practical but bounded. If an organisation is post-training reasoning, code, planning, or agentic models, the operational question should not be only “which KL direction?” or “what clipping threshold?” It should be “are we training on broken trajectories?” TRD suggests a post-training factory where failed attempts are repaired upstream before dense supervision. The boundary: this depends on verifier quality, refinement quality, teacher competence, and the cost of extra rollout generation. Tedious constraints, yes. Reality does enjoy showing up uninvited.

The comforting story is wrong: the teacher is not always the bottleneck

The obvious reader misconception is that a stronger teacher should fix a weaker student. The student produces a flawed reasoning trace; the teacher evaluates each prefix; the loss pulls the student toward the teacher. More intelligence goes in, better behaviour comes out. Very civilised.

The paper argues that this story breaks at the level where reasoning actually fails. In on-policy distillation, the teacher is not supervising an abstract solution. It is supervising the prefixes that the student already generated. If the student’s partial solution has entered a contradictory path, the teacher is being asked to provide token-level guidance from a context that may no longer have a clean continuation to the correct answer. That is prefix failure: a wrong prefix that cannot reach the right solution without backtracking, contradiction, or reflection.

This matters because dense per-token supervision is usually treated as the main advantage of on-policy distillation. Reinforcement learning with verifiable rewards gives sparse sequence-level feedback. Supervised fine-tuning gives off-policy expert traces. OPD sits in the middle: it keeps the student on-policy while giving dense teacher signals along the student’s own rollout. That sounds efficient, until the student’s own rollout becomes the problem.

Jiang, Xu, Ding, and Zhang’s Trajectory-Refined Distillation paper formalises this failure and proposes TRD as the repair mechanism.¹ The important move is conceptual before it is algorithmic: the paper changes the unit of intervention. Existing fixes adjust per-token losses on the original rollout. TRD rewrites the rollout first.

Prefix failure turns dense supervision into a bad GPS

A wrong reasoning prefix creates two competing pressures for the teacher.

One pressure is local consistency. Given the student’s existing prefix, the teacher may continue in a way that is coherent with what has already been written, even if the path is globally doomed. The other pressure is correction. The teacher may try to pivot back toward the right solution, but that pivot may require a correction-onset token that is low-probability under the student and awkward under the existing context.

The paper describes this as a bimodal teacher distribution. Under prefix failure, the teacher distribution can split between a wrong-continuation mode and a correction mode. This creates different pathologies depending on the KL direction.

Forward KL is mode-covering. It can be dominated by the correction-onset region, which may force the student toward an out-of-distribution pivot. Reverse KL is mode-seeking. It can be dominated by the wrong-continuation region, because that is where the student already places probability mass. One version risks unstable correction pressure; the other quietly keeps polishing the failure. A pleasing menu of bad options.

The deeper issue is gradient fragmentation. The paper contrasts the ideal correction path with what dense KL actually supervises. Ideally, once the teacher suggests a corrective token, later supervision should unfold along the corrected path. If the correction starts at token $\bar{y}^{\ast}_t$, the next teacher signal should be conditioned on a context that includes $\bar{y}^{\ast}_t$.

But standard dense KL does not do that. It evaluates the teacher along the frozen student rollout $y_o$. At the next position, the context still contains the student’s original wrong token, not the teacher’s corrective token. The teacher may keep recommending the same correction onset again and again, each time from a deeper wrong-continuation context. The correction does not unfold. It stutters.

In business language, this is the difference between advising a failed process and redesigning it. If a procurement workflow took the wrong approval route three steps ago, adding more comments to the current form may not fix it. You need to reroute the workflow. Otherwise, everyone is very well informed while still being wrong.

TRD repairs the route before the lesson begins

TRD is a simple pipeline with a non-trivial implication.

First, the student samples a raw on-policy rollout $y_o$ for a problem $x$. Second, a refinement step produces a revised trajectory $y_r$. Third, the student is trained using dense KL on $y_r$, not on $y_o$.

The implementation differs by setting.

In OPD, the teacher is a separate Qwen3-8B model. The refinement prompt asks the teacher to rewrite the initial solution, preserve useful structure, and fix errors. The appendix prompt does not show the reference solution to the OPD teacher. In OPSD, the teacher and student share the same backbone, but the teacher branch receives privileged information: the reference solution $y^{\ast}$. The OPSD refinement prompt explicitly asks the model to rewrite the solution using that reference as guidance.

That distinction is important. TRD is not merely “ask a bigger model for the answer”. Its design objective is to remain close to the student’s support while improving the trajectory. The raw rollout anchors the refinement to reasoning patterns the student has already demonstrated. The teacher or privileged context then corrects the parts that break the path.

The paper frames this as a trajectory-construction problem upstream of standard OPD. The objective resembles RLVR in that it wants trajectories with higher verifier pass rates, but TRD does not directly optimise that objective with reinforcement learning. It constructs better trajectories and then applies standard distribution matching. That is a practical compromise: keep dense learning, but stop feeding it the wrong contexts.

The experiments mostly test whether repair beats reweighting

The experimental design is more useful when read by purpose, not by table order.

The models are from the Qwen3 family. In OPD, the teacher is Qwen3-8B and the students are Qwen3-1.7B and Qwen3-4B-Instruct-2507. In OPSD, Qwen3-4B-Instruct-2507 and Qwen3-8B are used as shared teacher-student backbones under privileged conditioning. Math training uses roughly 40,000 DeepScaleR problems. Code training uses roughly 25,000 TACO problems. Evaluation draws 16 completions per problem and reports Avg@16 and Pass@16.

Avg@16 measures average sample quality. Pass@16 measures whether at least one of the 16 sampled completions solves the problem. For operators, this distinction is not cosmetic. Avg@16 says how good the model’s average attempt is. Pass@16 says whether a sampling budget can find a solution. The first is exploitation; the second is coverage. TRD’s most interesting gains are often in coverage.

The math results say “consistent”, not “magical”

In OPD, TRD improves Avg@16 over the base at both student scales and is best or tied-best on seven of eight benchmarks in each block. For Qwen3-1.7B, TRD’s Avg@16 reaches 49.4 on AIME24, a +4.6 point gain over base, and improves AIME25 by +1.7 points. For Qwen3-4B-Instruct-2507, TRD reaches 65.4 on AIME24, 47.9 on AIME25, and 33.2 on HMMT25, again improving over base.

The more revealing OPD case is the stronger 4B student. Several raw-rollout KL baselines regress below base on multiple math benchmarks. TRD generally avoids that damage. This fits the mechanism: token-level pressure on the original rollout can disrupt an already competent student, while trajectory-level refinement gives it a cleaner target.

Pass@16 tells the harder story. In OPD, TRD improves AMOBench Pass@16 by +5.1 points for Qwen3-1.7B and +12.8 points for Qwen3-4B-Instruct-2507. AIME24 and AIME25 are comparatively saturated, so the gains are less visible there. This is exactly where Pass@16 is useful: it exposes whether the method broadens the set of reachable solutions rather than only raising the average probability of already reachable ones.

In OPSD, the math evidence is stronger. TRD is reported as best across the evaluated Avg@16 benchmarks at both scales and never below base. On Qwen3-8B, TRD reaches 69.2 on AIME25, 44.5 on HMMT25, 42.8 on BeyondAIME, and 17.3 on AMOBench. The Pass@16 result on AMOBench is the headline: Qwen3-8B rises from 41.0 to 61.5, an absolute gain of +20.5 points and roughly 50% relative improvement. HMMT25 also rises from 66.7 to 76.3.

That does not make TRD a universal post-training wand. It makes it a strong candidate mechanism for reasoning tasks where failed prefixes are common, verifiers are available, and reference-guided or teacher-guided repair can produce better trajectories. That is a narrower claim. It is also a more useful one.

Code is the warning label, not a footnote

The code results are mixed, and they should be treated as part of the lesson rather than as an inconvenience to be shoved into the methodological attic.

In OPD code evaluation, TRD improves HumanEval+ and MBPP+ in several rows and is best on MBPP+. For Qwen3-1.7B, TRD raises MBPP+ Avg@16 from 50.1 to 51.2 and Pass@16 from 60.8 to 62.7. For Qwen3-4B-Instruct-2507, TRD improves MBPP+ Avg@16 from 64.6 to 65.2.

LiveCodeBench is different. TRD does not beat the base model there. In the Qwen3-4B-Instruct-2507 OPD Pass@16 table, the base is 55.2 while TRD is 54.1. For Qwen3-1.7B, the base is 48.5 while TRD is 46.8. The paper’s own interpretation is restrained: the current teacher may not provide effective refinements on these harder code tasks.

That boundary matters commercially. Code repair is not the same as math solution repair. A mathematical derivation can often be rewritten toward a reference structure with a rule-based final-answer verifier. Code has hidden edge cases, interface constraints, and execution semantics. A refinement that looks cleaner may still fail tests. Anyone selling “trajectory repair” for code without stronger verification has discovered marketing, not engineering.

The appendix shows where the gain comes from

The appendix is not decorative. It tells us why the benchmark gains are plausible.

First, the training-trajectory analysis shows that refinement changes the data distribution. On Qwen3-8B under OPSD, the verifier pass rate rises from 65.8% on raw rollouts to 81.4% on refined trajectories. Median length compresses from 7.7K tokens to 0.88K, moving closer to the reference median of about 0.49K. The same compression appears even on already-correct raw rollouts, which means TRD is not only fixing failures. It is also exposing the student to shorter alternative derivations.

Second, the AMOBench bucket analysis localises the coverage gain. The paper partitions AMOBench questions by how often the base model solves them in 16 samples. On the hardest bucket, B0, the base model solves zero out of 16 attempts by construction. There are 23 such questions. TRD achieves Pass@16 of 0.39 on this bucket, compared with 0.22 for the strongest forward-KL comparison. In ordinary terms: TRD solves 9 of 23 base-unreachable AMOBench questions under that sampling budget.

Third, the subset-filter ablation prevents a lazy conclusion. One might assume TRD should train only on failed rollouts, or only on cases where refinement turns failure into success. The paper tests this. Filtering hurts TRD, especially Pass@16. On Qwen3-8B, AMOBench Pass@16 drops by about 10 to 15 points under the tested subset filters. The full refined corpus works better because failures and successes carry complementary signals: hard cases teach reachability, while already-solved cases teach alternative and shorter paths.

This is the operationally interesting part. TRD is not just a failure-correction filter. It is a trajectory-distribution shaping method.

The business interpretation is trajectory repair, not better loss garnish

What the paper directly shows is bounded: TRD improves post-training outcomes on the evaluated Qwen3 math setups, compares favourably with common dense-KL baselines, and has mixed results on code. It also provides mechanism diagnostics and ablations that support the prefix-failure explanation.

What Cognaptus infers for business use is broader but still disciplined. In model post-training pipelines, the unit of quality control should increasingly move from tokens to trajectories. This is especially true for reasoning-heavy systems: quantitative agents, research assistants, coding copilots, compliance analysis tools, planning agents, and operations copilots. In these systems, the business risk is rarely one bad token. It is a plausible-looking path that becomes unrecoverable five steps later.

A TRD-like production pattern would look like this:

The most direct business value is not “higher benchmark score”. It is cheaper recovery from model-generated failure patterns during post-training. If a company is already paying for rollouts, verifiers, and teacher models, TRD suggests that the expensive mistake is training on raw failures and hoping the loss function will develop common sense. It may not. It is a loss function, not a therapist.

Boundaries that affect deployment

The first boundary is teacher quality. TRD relies on the teacher or privileged self-teacher producing refinements that are both more correct and still close enough to the student’s support. A weak teacher gives weak repair. A teacher that rewrites too aggressively can destroy the on-policy character that makes OPD attractive.

The second boundary is verifier quality. The paper’s strongest results are in competition math, where final-answer verification is relatively clean. The code results show what happens when the repair problem becomes more semantically demanding. For enterprise workflows, this means TRD-like systems need domain-specific checks: tests, constraints, audits, simulations, or human review. “Looks better” is not a verifier. It is how demos happen.

The third boundary is compute. TRD adds an extra sampling pass to construct refined trajectories. The paper reports that shorter refined trajectories can offset some of this cost. On Qwen3-8B OPSD, total training-pipeline wall-clock is nearly matched: 9:20 for TRD versus 9:40 for vanilla OPSD on a single 8×H100 80GB node. But smaller OPD setups show more visible overhead, such as 8:10 for TRD versus 4:40 for vanilla OPD on Qwen3-1.7B. The efficiency story is therefore scale- and setup-dependent.

The fourth boundary is domain transfer. The mechanism is general, but the evidence is not equally strong everywhere. The paper gives its cleanest support for math reasoning under OPD and OPSD, with Qwen3-family models and specific datasets. It does not prove that the same recipe automatically improves legal reasoning, financial modelling, medical triage, multi-agent workflows, or production code repair. Those domains need their own verifier and refinement studies. The universe remains inconveniently specialised.

The operational lesson: repair before you distil

TRD’s contribution is not that it invents teacher guidance. It is that it places guidance at the right level of abstraction.

Dense token supervision is useful when the context is useful. When the context is a failed prefix, the teacher’s signal may become fragmented, bimodal, or locally aligned with a doomed continuation. Token-level clipping and reweighting can moderate the damage, but they still operate inside the wrong trajectory. TRD asks a more practical question: why not repair the trajectory first?

For AI operators, that is the durable lesson. The post-training stack should not treat generated reasoning traces as passive data. It should inspect them, repair them, and only then use them as training material. In other words: stop teaching the model from the wreckage and start teaching it from the reroute.

Cognaptus: Automate the Present, Incubate the Future.