Where Did This Sentence Come From? Tracing Provenance in LLM Reasoning Distillation

1. Problem Statement

Reasoning distillation involves using a large "teacher" model to generate reasoning paths (Chain-of-Thought) to fine-tune a smaller "student" model. While this approach has improved student performance, a critical gap exists in understanding provenance:

• The Core Question: When a distilled model generates a response in a new (test-time) context, does it genuinely follow the teacher's reasoning logic, or does it revert to its original, pre-distillation output patterns?
• The Limitation: Previous work focused on performance metrics and heuristic data selection (e.g., selecting data that matches the student's current distribution) but lacked a fine-grained analysis of where specific reasoning steps originate. It remains unclear if generalization gains are due to transferred reasoning capabilities or merely the reinforcement of pre-existing student patterns.

2. Methodology: Reasoning Distillation Provenance Tracing

The authors propose a framework to trace the origin of every action (sentence) produced by a distilled model by comparing predictive probabilities across three models under the same context:

1. Teacher Model ($M_T$): The original large model.
2. Original Student Model ($M_S$): The student before distillation.
3. Distilled Model ($M_D$): The student after fine-tuning.

The Process:

1. Sampling: The distilled model ($M_D$) generates a trajectory on a test set.
2. Re-evaluation: Each sentence (action) in this trajectory is fed back into $M_T$, $M_S$, and $M_D$ to obtain the probability of generating that specific sentence given the preceding context.
3. Classification: Based on the probability differences ($\Delta$), each sentence is categorized into four types:
   • Teacher Sentence: The teacher assigns a significantly higher probability than the original student ($\Delta_{TS} > \beta$). This indicates the action originates from the teacher's knowledge.
   • Student Sentence: The original student assigns a significantly higher probability than the teacher ($\Delta_{TS} < -\beta$). This indicates the action reverts to the student's original patterns.
   • Shared Sentence: Probabilities are similar across all models ($|\Delta| \le \alpha$). These are pre-existing patterns in both models not enhanced by distillation.
   • Boosted Sentence: Probabilities are similar between teacher and student, but the distilled model assigns a significantly higher probability. This indicates the student's internal patterns were amplified by distillation.

Thresholds:

• $\alpha$: Filters out negligible probability differences.
• $\beta$: Determines the threshold for distinguishing between Teacher/Student origins.

3. Key Contributions

A. The Provenance Tracing Framework

The paper introduces a systematic method to disentangle the origins of reasoning steps. This moves beyond "black box" performance evaluation to a granular analysis of how the model reasons. It provides a principled way to determine if distillation transfers new capabilities or merely reinforces old ones.

B. Empirical Insights on Generalization

Through analysis on benchmarks like AIME24 and GPQA-D using models like DeepSeek-Distill-Qwen-7B, DeepSeek-R1-0528-Qwen3-8B, and LIMO-v2, the authors found:

• Teacher Sentences Correlate with Correctness: In correct responses, "Teacher Sentences" appear with higher probability, especially in the early stages of reasoning (planning and analysis).
• Activation of Latent Patterns: Distillation often activates "Boosted Sentences" (student-internal patterns). However, these are not uniformly beneficial; in smaller models, an over-reliance on Boosted Sentences in later stages correlated with incorrect answers.
• Early vs. Late Reasoning: High-probability Teacher Sentences in the early inference steps are a strong indicator of successful reasoning transfer.

C. Teacher-Guided Data Selection Strategy

Based on the insight that "Teacher Sentences" drive performance, the authors propose a new data selection method for training:

• Hypothesis: Training data rich in Teacher Sentences leads to better generalization.
• Method: Instead of selecting data based on how well it fits the student's current distribution (as in GRAPE), the method selects training samples where the Teacher and Student models diverge the most (i.e., samples containing the highest count of Teacher Sentences).
• Cost: This requires only a single forward pass to extract logits, making it computationally efficient compared to token-by-token generation.

4. Experimental Results

The authors validated their approach across multiple teacher-student pairs (e.g., DeepSeek-R1, QwQ, GPT-OSS-120B as teachers; Qwen2.5/3 variants as students) and datasets (AceReason, OpenThoughts).

• Performance Gains: The proposed "Teacher-Guided" selection strategy outperformed both vanilla distillation and the state-of-the-art GRAPE method.
  • Average Improvement: 1.7% to 2.5% across benchmarks (AIME24, AIME25, MATH500, OlympiadBench).
  • Specific Example: On DeepSeek-R1 + Qwen3-4B-Base, the method improved AIME24 accuracy from 44.4% (Vanilla) to 49.3%.
• Ablation Studies:
  • Maximizing the absolute count of Teacher Sentences in training data yielded better results than maximizing relative proportion or length.
  • The method showed robustness across different $\beta$ values and domains (Science and Math), suggesting that reasoning patterns (Teacher Sentences) are transferable across domains.
• Provenance Shift: Post-training analysis confirmed that models trained with the new strategy generated a higher proportion of Teacher Sentences and fewer Student Sentences compared to the baseline.

5. Significance and Impact

• Explainability: The paper shifts the focus from whether distillation works to why it works, providing a quantitative explanation for generalization gains via the "Teacher Sentence" metric.
• Data Curation: It challenges the prevailing heuristic of "selecting data the student likes" (GRAPE) and proposes a "selecting data the teacher knows best" approach, which is more effective for transferring complex reasoning.
• Model Auditing: The framework offers a tool to audit distilled models, potentially identifying which teacher model contributed to specific capabilities or detecting if a model has reverted to its original patterns.
• Future Directions: The authors suggest this framework can be extended to ensemble teacher settings and larger model scales, offering a new paradigm for reliable reasoning distillation protocols.

In summary, this work provides a rigorous, probability-based lens to analyze reasoning distillation, proving that explicitly prioritizing teacher-originated reasoning steps during data selection significantly enhances the generalization capabilities of small language models.