Technical Summary: Distribution Corrected Offline Data Distillation for Large Language Models

Problem Statement

The paper addresses the challenge of transferring complex reasoning capabilities from large, strong "teacher" Large Language Models (LLMs) to smaller, resource-constrained "student" models. While offline distillation (training a student on fixed reasoning traces generated by a teacher) is sample-efficient and provides high-quality supervision, it suffers from a fundamental distributional drift.

During training, the student model is conditioned on teacher-generated prefixes. However, during inference, the student must autoregressively generate its own prefixes. As the student's generated tokens deviate from the teacher's trajectory, it enters regions of the token space underrepresented in the training data. This leads to compounding errors and unreliable behavior over long reasoning trajectories, a limitation that existing refined distillation objectives (e.g., bidirectional KL divergence minimization) do not fully resolve because they still optimize against a teacher-provided distribution.

Conversely, on-policy methods (e.g., GRPO, RLHF) mitigate this drift by training on self-generated data but incur high computational costs for online sampling and often struggle with low-quality traces in early training stages.

Methodology: DISCORD

The authors propose DISCORD (Distribution-Corrected Distillation), a principled offline framework that retains the efficiency of teacher-generated data while correcting the distribution mismatch. The core mechanism involves adaptively reweighting teacher-token supervision based on the student's current support (its likelihood of generating that token).

Key Technical Components:

1. Student-Support-Aware Reweighting:
   The method introduces a density ratio to reweight training trajectories $\{x, y \sim \pi_T(x; c)\}$:
   $$\frac{\bar{\pi}_\theta(y|x)}{\bar{\pi}_T(y|x;c)}$$
   where $\pi_\theta$ is the student, $\pi_T$ is the teacher, and the bar denotes a stop-gradient operation. This ratio prioritizes teacher tokens that the student is likely to encounter during inference.

2. Stabilized Transformation:
   Direct application of the density ratio causes training instability due to unbounded values and high variance over long sequences. The authors apply two stabilizations:

   • Sigmoid Transformation: The log-density ratio is passed through a sigmoid function $\sigma(\tau \cdot \log(\dots))$ to map the correction into a bounded range $(0, 1)$, preserving the ordering of the ratio while preventing explosion.
   • Token-Level Approximation: To mitigate length bias (where longer sequences accumulate larger aggregate weights), the correction is applied at the token level. For each token $y_t$ in a trajectory, the weight is:
     $$r_t = \sigma\left(\tau \log \frac{\bar{\pi}_\theta(y_t | x; y_{<t})}{\bar{\pi}_T(y_t | x; c; y_{<t})}\right)$$
     This weight is applied as a gradient multiplier during the training of the base objective (e.g., SFT or KL divergence).

3. Interpretation:
   The correction weight $r_t$ (with $\tau=1$) is interpreted as the posterior probability that a token was drawn from the student distribution rather than the teacher distribution, assuming a uniform prior. This connects the method to contrastive decoding principles but applies them during training to prioritize learning in regions the student can reliably inhabit.

Key Contributions

• A Principled Offline Framework: DISCORD is the first method to effectively correct teacher-student distribution drift using only offline, teacher-provided data, eliminating the need for costly online student rollouts.
• Flexible Integration: The token-level reweighting mechanism is compatible with various offline distillation backbones, including standard Supervised Fine-Tuning (SFT) and KL-divergence-based objectives.
• Stabilization Techniques: The introduction of sigmoid transformation and token-level approximation addresses the variance and instability issues inherent in density-ratio-based corrections for long sequences.

Experimental Results

The method was evaluated on mathematical reasoning benchmarks (GSM8K, MATH, MATH500) and harder held-out competition-style tasks (AMC, AIME, OlympiadBench, Omni-Math) using Qwen2.5 and Llama-3.2 models.

• Performance Gains: DISCORD consistently outperforms strong offline baselines (SFT, ToDi, Symmetric KL) on in-domain and out-of-domain (OOD) tasks. For instance, in the Qwen2.5-7B $\to$ 3B distillation, DISCORD improved MATH500 accuracy by 4.60 points over standard SFT and 2.60 points over Symmetric KL.
• Robustness to Prefix Drift: Evaluations using ExAccErr (distributional divergence) show that while standard SFT performance degrades significantly as the prefix length increases (student-generated prefixes diverge from teacher), DISCORD maintains stable performance, indicating reduced compounding errors.
• Efficiency vs. Online Methods: DISCORD achieves performance comparable to online student-rollout methods (like GKD) but with significantly lower computational cost. In the 7B $\to$ 3B setting, DISCORD required only 3.12 GPU-hours compared to 49.56 GPU-hours for GKD, while using no online student tokens.
• Reasoning Quality: Qualitative analysis reveals that DISCORD produces cleaner, more concise, and logically coherent traces. It reduces common failure modes such as repetition, post-answer continuation, and multiple distinct final answers. Crucially, in "failure-recovery" cases, DISCORD preserves the underlying mathematical structure of the teacher's solution even when it does not mimic the teacher's surface-level formatting.

Significance and Claims

The paper claims that lightweight, distribution-correction-aware training can substantially strengthen offline reasoning distillation without requiring online interaction.

• Bridging the Gap: The work demonstrates that it is possible to narrow the performance gap between offline and online distillation methods by intelligently reweighting existing offline data, rather than by generating new data.
• Generalizability: While the experiments focus on mathematical reasoning, the authors posit that the method is applicable to other long-horizon generation settings, including code generation and scientific reasoning.
• Practical Deployment: By avoiding the high compute costs of online sampling, DISCORD offers a viable path for deploying capable reasoning models in resource-constrained environments.

The authors remain modest, noting that while their method improves upon prior offline baselines, online methods like GKD still hold an edge in some specific in-domain benchmarks due to their direct observation of the student's rollout distribution. However, DISCORD offers a superior trade-off between performance and computational efficiency.