Failure Profiling and Reachable Trajectory Selection for Reasoning Distillation

AI summary

Overview

Research area: Natural Language Processing — specifically reasoning distillation, the practice of fine-tuning small language models on reasoning traces produced by much larger teacher models. It also touches on data selection, error-aware data construction, and out-of-distribution generalization.

Technical level: Intermediate. The paper assumes familiarity with supervised fine-tuning (SFT), token-level log-probabilities, rejection sampling, and the teacher–student distillation framing, but its core idea is explained without requiring deep mathematical background.

Scope: The paper proposes a three-stage data-construction pipeline (behavioral failure profiling, mode-specific correction generation, and localized "branching point reachability" trajectory selection) and evaluates it on two student backbones across five math benchmarks.

What This Paper Is About

When a small model is trained on reasoning solutions written by a much stronger teacher, it learns to imitate the teacher's style rather than its own — so once the student makes a mistake at inference time, it has no learned way to recover. The paper argues that the fix is to build training data from both what the student actually gets wrong and what it can realistically reproduce. It does this by repeatedly sampling the student to diagnose how it fails, prompting the teacher for corrections tailored to that specific failure mode, and then picking the correction whose first few corrective steps the student itself finds most probable.

Key Contributions

1. Behavioral failure profiling. The paper introduces a two-statistic diagnosis of student errors — error rate α (fraction of K samples with a wrong final answer) and error consistency κ (share of incorrect samples accounted for by the single most frequent wrong answer) — which classifies each problem into one of three failure modes: stubborn (α ≥ 0.6, κ ≥ 0.8), inconsistent (α ≥ 0.6, κ < 0.8), or sporadic (0 < α < 0.6). This differs from prior error taxonomies that classify by semantic error type rather than by the student's behavioral pattern.

2. Mode-specific adaptive correction generation. Each diagnosed mode triggers a distinct teacher prompt: a multi-angle attack on the locked-in reasoning path for stubborn errors, a complete structured solution framework for inconsistent errors, and a minimal surgical fix preserving the longest correct prefix for sporadic errors. Each candidate is decomposed into a reflection r, a preserved prefix p, and a corrective continuation c.

3. Branching point reachability (BPR) scoring. Instead of reranking a whole trajectory by its likelihood under the student, BPR averages the student's log-probabilities only over the first L sentences of the corrective continuation, immediately after the point of divergence. This avoids dilution from later tokens that already sit in a corrected context, and the averaging (rather than summing) removes a length bias toward shorter candidates.

4. A full generate-then-select pipeline with empirical validation. The pipeline is validated on Llama-3.1-8B-Instruct and DeepSeek-Math-7B-Instruct against SFT, RFT, ISC, S³C-Math, GPTAug, RefAug, and LEMMA, with the code and trained weights released.

Main Findings

• Best average accuracy on both backbones at 88.90K examples. FPRT reaches 79.7% ± 0.4 average on Llama-3.1-8B (GSM8K 86.2, MATH 46.2, SVAMP 86.8, ASDiv 86.1, MAWPS 93.3) and 80.6% ± 0.3 on DeepSeek-Math-7B (84.0, 48.9, 85.4, 89.8, 94.7), exceeding the strongest baseline, LEMMA, by 2.4% and 1.8% respectively.

• The lead holds at smaller corpus sizes. At the first-pass sizes of 62.74K (Llama) and 65.76K (DeepSeek), the averages are 79.0% and 80.0%, still above every comparison method with a reported five-benchmark average — surpassing LEMMA by 1.7% and 1.2% at roughly 70% of LEMMA's corpus size.

• Gains concentrate under distribution shift. Out-of-distribution averages of 88.7% (Llama) and 90.0% (DeepSeek) exceed LEMMA by 3.0% and 2.6%. MAWPS alone shows leads of 4.4% and 3.6%.

• MATH is the exception. FPRT trails LEMMA by 2.6% on Llama-3.1-8B and 1.8% on DeepSeek-Math-7B on MATH, ranking second on both. The authors suggest LEMMA's error-type-grounded augmentation may suit MATH better and that correctness filtering may reduce candidate availability on hard problems; they state the current results do not isolate construction from selection as the cause.

• Profiling drives most of the gain; BPR refines it. On the Llama first-pass configuration, replacing profiling with a single uniform minimal-correction prompt drops the average from 79.0% to 72.9% (GSM8K falls 8.0 points, from 85.5 to 77.5). Replacing BPR with global log-probability scoring drops it to 77.3%. Removing both gives 72.4%, only 0.5% below the BPR-only variant.

• BPR beats both controls. Against direct teacher-trace SFT (75.9% Llama, 78.1% DeepSeek) and unscored random selection from the same correct candidate pool (76.6%, 79.0%), BPR reaches 79.0% and 80.0% — gains of 3.1% and 1.9% over the direct control.

• Window size has a moderate optimum. Sweeping L over {2, 4, 8, 16, 32, full} sentences, accuracy peaks at 79.0% at L = 8. Below that it is 77.1% and 76.7%; above it decays to 78.6% at L = 16, 78.2% at L = 32, and 77.3% at the global setting — matching the 1.7% BPR-versus-global gap.

• Boundary localization is reliable. An audit of 50 traces from the 65.76K DeepSeek configuration finds 88% accurate first-error boundaries; perturbing boundaries on 12% of questions with preserved prefixes changes average accuracy from 80.0% to 79.7% in a single run.

• Data scaling shapes differ by backbone. From 16.5K to 88.90K examples, Llama-3.1-8B rises from 71.8% to 79.7% (a 7.9% gain, with final increments of 0.3% and 0.4%), while DeepSeek-Math-7B rises from 77.9% to 80.6% (a 2.7% gain, with final increments of 0.5% and 0.1%). Both curves flatten at the top.

• Teacher swap is a minor effect. Substituting DeepSeek-V4-Flash for DeepSeek-V3.2 in a matched single-run comparison yields 81.0% versus 80.6% on DeepSeek-Math-7B.

Methodology in Plain English

The pipeline runs in three stages, then a standard fine-tuning step.

Stage 1 — Diagnose the student. For each of the 394.99K MetaMathQA problems, the student is sampled K = 5 times. Two numbers summarize the outcome: how often it is wrong (α), and how concentrated those wrong answers are on one particular wrong answer (κ). High α with high κ means the model is confidently stuck on a specific misconception ("stubborn"). High α with low κ means it is flailing in different directions each time ("inconsistent"). Low-but-nonzero α means it occasionally slips ("sporadic"). Problems the student always solves (α = 0) are dropped. One representative wrong solution is kept per problem as the concrete failure to fix.

Stage 2 — Ask the teacher for tailored fixes. The teacher (DeepSeek-V3.2) receives the problem, the representative error, and a prompt chosen by the diagnosed mode, and returns N = 5 candidates in one structured call. Stubborn errors get prompts that attack the locked path from different angles (direct fix, alternative method, conceptual reframe) plus a first-person reflection naming the trap. Inconsistent errors get prompts that supply full problem-solving frameworks. Sporadic errors get prompts demanding minimal edits that keep the student's correct prefix intact. Each candidate splits into a reflection, the verbatim correct prefix of the student's own attempt, and a new correct continuation.

Stage 3 — Pick what the student can actually do. Candidates with the wrong final answer are discarded. Among the survivors, the paper computes the student's average log-probability over only the first L = 8 sentences of the new continuation — the stretch right after the student's first wrong step. The highest-scoring candidate is kept. The intuition: later steps are easy precisely because the context has already been repaired, so scoring them adds noise; the first corrective moves are where learnability is actually decided. This is a single forward pass with no gradients. The paper is explicit that BPR is a deliberately biased proxy for the ideal objective of maximizing the student's full-trajectory probability, justified by a dilution argument and by the empirical window sweep.

Training. The selected trajectories form an SFT corpus. 90% of examples take the problem plus the student's erroneous attempt as input and the reflection plus corrected reasoning as target; 10% take the problem alone as input and the corrected solution as target, without the reflection or error trace. All inputs are loss-masked. Two construction passes produce the 88.90K corpora.

Why This Matters

Impact on research. The paper reframes distillation data selection as a two-part question — what does the student fail on, and what can it learn — rather than a single correctness filter. Its finding that localized scoring beats whole-trajectory reranking (79.0% versus 77.3%) is a concrete, testable claim against the common practice of scoring full responses, and connects to prior observations (cited via SCoRe) that SFT on offline correction traces suffers from distribution mismatch. The result that profiling contributes more than selection (a 6.1-point drop versus 1.7-point) is a useful signal about where effort should go in data-construction pipelines.

Real-world applications:

• Cheap model deployment for math tutoring. Tutoring systems that walk a student through a solution step by step benefit from small models that can recover from their own missteps rather than derailing.
• On-device and cost-constrained reasoning assistants. Building competitive 7–8B models reduces the need to serve frontier-scale models for structured quantitative tasks.
• Synthetic data curation pipelines. The profiling-plus-localized-scoring recipe generalizes to any setting where a strong generator supplies data for a weaker consumer with a known capability gap.
• Robustness-focused evaluation. The out-of-distribution emphasis (SVAMP, ASDiv, MAWPS) reflects deployment conditions where inputs rarely match training templates.

Industry relevance. The pipeline requires no reinforcement learning, no process reward models, and no preference optimization — only repeated sampling, teacher prompting, one forward pass per candidate for scoring, and standard SFT. It also demonstrably works at roughly 62.74K–65.76K examples, about a sixth of a 394.99K-example SFT run, which is directly relevant to teams watching data-generation budgets.

Future Directions

• Closing the MATH gap. FPRT trails LEMMA on MATH by 2.6% (Llama) and 1.8% (DeepSeek). The authors flag two untested hypotheses — error-type coverage and candidate retention after correctness filtering — and explicitly note that current results do not separate construction from selection as the source of the deficit.

• Scaling the profiling budget K. With K = 5, α can only take values in {0, 0.2, 0.4, 0.6, 0.8, 1.0} and κ is similarly coarse. The authors state that larger K would give finer profiling at proportionally higher sampling cost but that they did not explore this trade-off systematically.

• Reducing dependence on a strong teacher. Since the teacher both generates corrections and localizes the divergent step, problems the teacher itself cannot solve are dropped entirely. Better handling of teacher failures — or weaker teachers combined with self-localization — is a natural extension.

• Generalizing beyond mathematical reasoning. The paper evaluates only on math (GSM8K, MATH, SVAMP, ASDiv, MAWPS) with two student backbones. Whether the stubborn/inconsistent/sporadic taxonomy and the L = 8-sentence window transfer to code generation, logical reasoning, or other multi-step domains is untested, and the window optimum is likely domain-dependent given its asymmetric curve.

Target Audience

Most useful to researchers and engineers working on knowledge distillation, synthetic data construction, and small-model reasoning — particularly those building SFT corpora from teacher models and looking for selection criteria better than "is the final answer correct." It will also interest practitioners running inference-time self-correction or process-supervision work, since the paper positions itself against RL- and DPO-based correction training. Readers wanting a first introduction to distillation may find the setup accessible, but the value is greatest for someone already choosing among rejection sampling, DART-Math-style rebalancing, and likelihood-based reranking methods.

Authors’ abstract

Training smaller language models on reasoning solutions generated by a larger teacher model is the prevailing approach for transferring mathematical reasoning capabilities. However, the student tends to overfit the output distribution of the teacher and generalizes poorly when its own behavior at inference time deviates from the training data. Recent methods attempt to bridge this gap by incorporating error data, yet the corrective trajectories still reflect the reasoning…