PaperlayerView exact revision on arXiv

Source-linked AI summary

Making Expert Reasoning Learnable with Self-Distillation

Ethan Mendes, Jungsoo Park, Alan Ritter

arXiv:2602.02405v2cs.LGcs.AI

TL;DR

DAIL addresses the difficulty of learning from problems that current models cannot solve and the distribution mismatch that makes direct expert-solution imitation harmful. It transforms expert solutions into in-distribution reasoning traces and applies contrastive learning against shortcut-prone references. The method improves challenging mathematics reasoning, efficiency, and out-of-domain generalization using a small expert dataset.

  • Problem

    RLVR cannot provide training signals when models fail to sample correct solutions, while direct imitation of didactic expert traces can collapse reasoning performance because they omit model-critical intermediate steps.

  • Method

    DAIL uses self-distillation to transform expert solutions into detailed student-distribution traces, then applies contrastive learning against references conditioned on partial solutions.

  • Results

    DAIL improves pass@k across challenging mathematics benchmarks, generalizes to GPQA-diamond, and matches or exceeds untrained models with 2× fewer tokens.

  • Takeaways & Limitations

    DAIL enables reasoning improvement from fewer than 1000 high-quality expert solutions, including for difficult non-verifiable problems where standard RLVR fails.

  • Takeaways & Limitations

    For reflective LRMs such as Qwen3, direct sampling can produce solution references and fewer self-correction steps, motivating mixed-policy rollouts.

Abstract

from arXiv · show

Improving the reasoning capabilities of large language models (LLMs) typically relies either on the model's ability to sample a correct solution to be reinforced or the existence of a stronger model able to solve the problem. However, many difficult problems remain intractable for even current frontier models, preventing the extraction of valid training signals. A promising alternative is to leverage high-quality expert human solutions, yet naive imitation of this data fails because it is fundamentally out-of-distribution: expert solutions are typically didactic, containing implicit reasoning gaps intended for human readers rather than computational models. Furthermore, high-quality expert solutions are expensive, necessitating generalizable, sample-efficient training methods. We propose Distribution Aligned Imitation Learning (DAIL), a two-step self-distillation method that bridges the distributional gap by first transforming expert solutions into detailed, in-distribution reasoning traces and then applying a contrastive objective to focus learning on expert insights and methodologies. We find that DAIL can leverage fewer than 1000 high-quality expert solutions to achieve up to 31% pass@128 gains on Qwen2.5-Instruct and Qwen3, double reasoning efficiency, and enable out-of-domain generalization.

1. Introduction

DAIL addresses the failure of standard reinforcement learning and direct imitation on difficult problems by transforming expert solutions into in-distribution traces and learning selectively from them. With fewer than 1000 expert solutions, it improves challenging mathematics reasoning, efficiency, and out-of-domain generalization.

  • Motivation: RLVR provides no learning signal when models cannot sample a correct rollout, creating an exploration bottleneck on the hardest problems.Human expert solutions offer an alternative signal, but direct fine-tuning can substantially degrade reasoning performance.
  • Motivation: Direct imitation fails because didactic expert solutions omit granular steps and search dynamics that models may need for successful reasoning.Behavioral cloning can force the model to shortcut its post-training reasoning process.
  • DAIL: DAIL transforms expert solutions into detailed, in-distribution reasoning traces through mixed-policy generation with a teacher conditioned on the ground-truth solution.The method fills skipped steps and implied calculations while targeting the student’s own reasoning distribution.
  • DAIL: DAIL adds a contrastive objective to suppress rationalization shortcuts that may appear in generated traces.These shortcuts can arise when a model forces derivations toward known expert results rather than establishing each step.
  • Evidence: Fewer than 1000 expert solutions yield improvements across non-reasoning and long-chain-of-thought model settings.The experiments use difficult competition mathematics problems, including cases where the model had pass@k = 0 under repeated sampling.
  • Evidence: DAIL improves pass@k on challenging mathematics benchmarks, generalizes to GPQA-diamond, and matches or exceeds untrained models with 2× fewer tokens.The reported improvements scale with model size and test-time compute for long-chain-of-thought reasoning models.

2. Distribution Aligned Imitation Learning

DAIL uses a two-stage offline pipeline: it generates student-distribution reasoning traces grounded in expert solutions, then contrasts full traces against shortcut-prone references. The approach is designed to preserve authentic reasoning while reducing imitation of deficient logical bridges.

  • Trace generation: DAIL first synthesizes expanded reasoning traces that are in-distribution for the student while remaining grounded in expert solutions.This addresses didactic shortcuts such as omitted granular steps and implied calculations.
  • Trace generation: The teacher generates a synthetic dataset Dsyn of reasoning traces r for expert solution pairs (x, s), using frozen reference weights conditioned on the problem and solution.Direct sampling is effective for Qwen2.5-Instruct when prompted to fill reasoning gaps and solve without guidance.
  • Trace generation: Direct sampling is challenging for Qwen3 because expert conditioning can induce references to the solution and reduce authentic self-correction.Mixed-policy rollouts address these effects by letting the student generate tokens and deferring to the teacher only when necessary.
  • Contrastive learning: NLL indiscriminately imitates every token in expanded traces, including rationalization shortcuts that jump between known intermediate results.These shortcuts arise because the generator has access to the expert solution while producing the trace.
  • Contrastive learning: DAIL constructs a negative reference from partial expert solutions containing coarse-grained waypoints, which favors bypassing step-by-step logical progression.The contrastive objective then reduces the likelihood of tokens favored by this shortcut-laden reference relative to the full-information teacher.
  • Efficiency: The offline generation process is decoupled from optimization, enabling distributed dataset construction before training.Shared frozen base parameters and LoRA adapters allow teacher and student computations with one stored model-weight copy.

3. Results

DAIL improves mathematical reasoning and generalization by training on curated expert-grounded datasets, with evaluations spanning harder mathematical benchmarks, out-of-domain science reasoning, efficiency, and ablations of its learning objective.

  • Experimental setup: DAIL is evaluated using e1-verifiable, e1-proof, three mathematical benchmarks, and GPQA-Diamond for out-of-domain science reasoning.The training datasets contain 417 verifiable AIME problems and 669 non-verifiable Olympiad proof problems.
  • Mathematical reasoning: DAIL consistently outperforms the base instruction model across benchmarks, with larger gains at higher k and improvements on problems harder than e1-verifiable.RLVR baselines instead degrade on difficult settings, including marginal gains from GRPO trained on more than 40K samples.
  • Baseline comparisons: Direct SFT and STaR rationalization degrade performance, indicating that difficult tasks require in-distribution reasoning traces grounded in external expert solutions.Self-generated rationalization is insufficient when problem difficulty exceeds the model’s baseline capabilities.
  • Test-time efficiency: DAIL reaches roughly the same performance as the untrained model with 2× fewer tokens across benchmarks, especially under lower budgets on harder problems.The efficiency evaluation uses coverage (pass@128) across varying reasoning token limits.
  • Out-of-domain generalization: DAIL generally preserves or improves performance on GPQA-Diamond, suggesting mathematics-focused training does not cause catastrophic forgetting or domain overfitting.The out-of-domain comparison covers most models and inference settings.
  • Ablations and analysis: The contrastive objective consistently outperforms NLL, with gains in pass@1 and pass@128 across direct sampling and mixed policy rollouts.The pass@128 difference between generation methods is roughly 1%, while direct sampling yields larger pass@1 gains over NLL.

4. Related Works

Prior reasoning-distillation methods either require a stronger teacher or rely on synthetic and self-generated traces, while RLVR struggles when models cannot produce correct rollouts. DAIL instead expands compressed expert solutions into constructive traces for difficult problems where these approaches are insufficient.

  • Reasoning distillation transfers a stronger teacher’s outputs to a smaller student, but this dependence limits use at the frontier when no stronger teacher exists.
  • Synthetic reasoning methods bootstrap model rationales, synthesize new problems, or rewrite solutions into richer traces, often using strong closed-source teachers.
  • Self-distillation rewrites solutions with the same model, but prior work primarily targets simpler datasets, while concurrent methods target harder coding and mathematics tasks.
  • RLVR improves trajectories achieving correct outcomes, but difficult problems can yield zero advantage when on-policy models fail to generate useful rollouts.
  • DAIL differs by expanding compressed expert solutions rather than repairing incorrect student rollouts, providing constructive traces when models cannot generate partially correct reasoning paths.

5. Conclusion

DAIL converts didactic expert solutions into in-distribution constructive reasoning traces and uses contrastive learning to suppress shortcuts. It improves reasoning on challenging non-verifiable mathematics, increases efficiency, and supports out-of-domain generalization with relatively few expert solutions.

  • DAIL converts didactic expert solutions into in-distribution constructive traces, enabling models to learn from high-value data where standard RLVR fails.
  • DAIL’s contrastive objective suppresses learning from reasoning shortcuts while training on a realistically small set of high-quality expert solutions.
  • DAIL improves pass@k on challenging, non-verifiable mathematics problems and produces models that reason more efficiently and generalize to out-of-domain tasks.

Impact Statement

The paper identifies broader applications for DAIL beyond mathematics and science reasoning, including supervised fine-tuning of reasoning models and potential safety-related reasoning tasks. These applications are presented as future possibilities rather than demonstrated results.

  • DAIL is studied on difficult mathematics and science reasoning problems, while its broader applicability remains a proposed direction.
  • The authors suggest applying DAIL to expand supervised input-output datasets into in-distribution samples before fine-tuning reasoning models such as Qwen3.
  • Potential safety applications include training models to reason about refusals and privacy policies in a human-like manner.

A. Dataset Curation

The dataset curation process filters difficult AIME problems, collects Olympiad proof problems with expert solutions, and addresses potential source overlap through post-hoc deduplication analysis.

  • The curation section expands the procedures used to construct the e1-verifiable and e1-proof datasets.
  • The authors filter 903 AIME problems from 1985–2023 to retain those Qwen2.5-7B-Instruct could not solve in 32 attempts, yielding 417 problems.
  • For e1-verifiable, multiple AoPS solutions are concatenated into s and supplied to the teacher model, including figures drawn in Asymptote.
  • The e1-proof dataset contains 683 Olympiad problems and expert solutions, with figures included when present.
  • Post-hoc analysis found that 15 IMO-AnswerBench problems appeared to share underlying sources with e1-verifiable instances.

B. Results on Non-Mathematics Domains

DAIL is evaluated beyond mathematics on maze-solving, grid-diffusion, healthcare, and tool-use reasoning tasks. The supplied passages specify task setups and identify the corresponding evaluation tables.

  • Non-mathematics benchmarks: Maze evaluates the minimum steps needed to reach a goal in a 2×2 grid, using 500 training examples and a 50-sample test set.Privileged information provides the specific steps required for each maze.
  • Non-mathematics benchmarks: Rotten Oranges evaluates the minimum minutes required for rot to spread through an orange grid under the same training and testing configuration as Maze.Privileged information consists of the correct grid at each time step.
  • Non-mathematics benchmarks: HealthBench assesses responses to healthcare questions using 1,529 consensus-subset training examples and evaluation on the hard subset.The evaluation uses gpt-4o-mini instead of gpt-4o to reduce cost.

C.1. A Note on an Updated Evaluation Configuration.

The test-time scaling evaluation was revised because answer parsing from the allotted answer tokens produced too many non-responses. A final-answer prompt was added, reducing non-response rates to negligible levels.

  • Evaluation update: Parsing answers from 2048 answer tokens initially produced non-response rates above 10%, especially for base reasoning models.This issue affected the test-time scaling evaluation shown in Figure 4.
  • Evaluation update: The Rotten Oranges and HealthBench results are presented in their respective evaluation tables.Table 3 reports Rotten Oranges pass@k results, while Table 4 reports HealthBench Hard performance.
  • Evaluation update: The revised evaluation adds the prompt “The final answer (in boxed)” to force a parseable answer.The experiments were rerun with this final-answer instruction.
  • Evaluation update: Non-response rates fell to ≤0.2% after adding the final-answer prompt.The authors describe the resulting rates as negligible.

C.2. Data Generation

DAIL generates student-distribution reasoning traces from expert solutions, constructs contrastive negative references, and applies calibrated training procedures. The generated traces are substantially longer than the raw expert solutions, while offline training is markedly cheaper than RL training.

  • Generated data: The data-generation pipeline expands expert solutions into traces that are 4× as long on average.Figure 12 compares raw expert-solution lengths with the expanded traces used for Qwen2.5-7B-Instruct on e1-verifiable.
  • Negative references: The contrastive objective constructs negative references from unstructured intermediate results, encouraging the model to hallucinate logical connectives between them.Conditioning the negative reference model on answer-leak context increases rationalization shortcuts as a negative signal.
  • Waypoint extraction: Final answers are extracted from the last boxed expression, while intermediate waypoints are collected by regular-expression patterns and the final answer is removed from the waypoint set.For numerical e1-verifiable problems, only numbers are included in the waypoint set.
  • Generation settings: Mixed policy rollout temperature is calibrated across candidate τ values, with τ = 0.8 selected across experiments.The calibration uses training-set performance because joint validation tuning is computationally infeasible.
  • Training efficiency: DAIL costs 0.1 H100 hours per epoch after synthetic-data generation, compared with 277.6 and 8.7 H100 hours per epoch for GRPO on e1-verifiable and DeepScaleR.The comparison reflects DAIL’s offline formulation and uses a conservative 3× A40-to-H100 conversion factor.

C.7. Exploratory Experiments and Connections to On-Policy Methods.

Exploratory on-policy DAIL training converged very slowly on the intentionally small dataset, whereas the authors associate off-policy training with more efficient updates and parallel generation. They leave further investigation to future work.

  • Exploratory experiments: The on-policy DAIL variant converged extremely slowly and required many epochs on the intentionally small dataset.The authors speculate that the difficulty of the questions contributed to this discrepancy.
  • Exploratory experiments: Off-policy training enables more efficient updates and parallel dataset generation but likely sacrifices some performance gains.This is presented as a trade-off relative to on-policy methods.
  • Exploratory experiments: Further investigation of the on-policy variant is left for future work.The exploratory comparison does not constitute a full study of on-policy DAIL.

C.8. Other Evaluation Details and Ablations

The evaluations test DAIL’s contrastive objective, comparisons with standard distillation, backbone transfer, and performance boundaries on weaker models. Results support the contrastive design and broader applicability, while showing that model competence remains important.

  • Effect of contrastive objective: The contrastive objective outperforms optimizing only DKL between the student and teacher on BeyondAIME with mixed policy rollouts.The ablation directly supports the intended role of the contrastive objective.
  • Effect of contrastive objective: Minimizing divergence between the student and negative reference degrades performance compared with the untrained model at k = 128.This result matches the motivation for the objective and its negative reference.
  • Comparing to standard distillation: DAIL and standard supervised distillation from Qwen3-235-A22B enable similar levels of pass@128 coverage, although DAIL does not require a larger distillation model.Both methods use the same number of tokens per example in this comparison.
  • Competent models are necessary for consistent improvement with DAIL: Gains are modest on smaller reasoning models, and DAIL shows no improvement on Llama-3.1-8B-Instruct with e1-verifiable.The paper identifies a baseline level of reasoning competence as necessary for maximum improvement.
  • Performance on other backbones: DAIL consistently improves performance on a distilled Llama DeepSeek-R1 variant, indicating transfer beyond the Qwen model family.The reported results are from Table 12 on the Llama-based model.
Loading 2602.02405v2…