Technical Summary: Structural Rationale Distillation via Reasoning Space Compression (D-RPC)

Problem Statement

Knowledge distillation aims to train Small Language Models (SLMs) to reason by imitating Large Language Model (LLM) teachers. However, a critical bottleneck exists: teachers often generate rationales for similar problems with wildly varying structures and strategies. This inconsistency acts as noisy supervision, making it difficult for students to internalize reusable reasoning patterns. Furthermore, covering all valid solution paths requires excessive teacher generations, inflating costs, and highly variable supervision can amplify hallucinations during fine-tuning. Existing attempts to mitigate this via structured reasoning (e.g., hierarchical templates) often impose heavy context loads that SLMs struggle to utilize due to limited capacity and context windows. The core challenge is compressing the teacher's reasoning space into a compact, consistency-preserving structure before distillation without relying on verbose templates.

Methodology: Distillation through Reasoning Path Compression (D-RPC)

The authors propose D-RPC, a framework that constrains the teacher to follow a compact, dynamically maintained bank of reusable high-level reasoning paths. The process operates in three stages:

1. Reasoning Bank Initialization:

   • A seed set of questions (5% of training data) is solved by the teacher.
   • The teacher assigns category and intent labels to each question.
   • High-level reasoning paths (abstract sequences of steps) are extracted from the seed set.
   • Intents are embedded and clustered (using DBSCAN) to form "canonical intents."
   • A reasoning path bank is constructed, mapping each (category, canonical intent) pair to a set of observed reasoning paths.

2. Bank-Guided Supervision:

   • For each new training question, the teacher predicts the category and intent.
   • The system retrieves the top-$K$ candidate paths from the bank that best match the question's intent (using question embeddings rather than surface intent strings to handle noise).
   • The teacher is conditioned on these retrieved paths to generate a detailed rationale. If a novel, correct path is discovered, it is buffered and periodically merged back into the bank via re-clustering, allowing the bank to grow dynamically.
   • The resulting supervision tuples (Question, Path, Rationale, Answer) are collected.

3. Student Fine-Tuning:

   • The student SLM is fine-tuned using LoRA (Low-Rank Adaptation) on the collected supervision tuples.
   • Because related questions share reasoning paths, the supervision signal exhibits lower variance, facilitating the internalization of reusable strategies.

Theoretical Analysis

The paper provides a PAC-Bayes generalization bound to formalize the trade-off between bank size ($K_{bank}$) and coverage.

• The Trade-off: A smaller bank reduces supervision entropy (improving consistency) but risks coverage gaps (increasing approximation bias). A larger bank improves coverage but increases supervision entropy and reduces parameter sharing.
• The Bound: The generalization gap scales with $\sqrt{M \cdot N}$, where $M$ contains a term $\log K_{bank}$. The analysis predicts an optimal intermediate bank size that minimizes the deviation term.
• LoRA Specifics: The analysis leverages the LoRA parameterization to show that banking can reduce the frozen-model loss ($\hat{b}_{LS}(0)$), thereby tightening the complexity term of the bound.

Key Contributions

1. D-RPC Framework: A novel method that compresses the teacher's reasoning space into a dynamic bank of canonical paths, routing similar questions to shared strategies to reduce rationale divergence.
2. Theoretical Formalization: A PAC-Bayes bound explicitly linking bank size to generalization performance, identifying an optimal intermediate bank size that balances reuse and coverage.
3. Empirical Superiority: Demonstrated consistent improvements over Chain-of-Thought (CoT), freeform generation, direct distillation, and structured-supervision baselines across five benchmarks (GSM8K, AQUA, StrategyQA, AI2ARC, MATH) and two student models (Llama 3.1 8B, Qwen 3 1.7B).
4. Efficiency: Achieves higher accuracy while using fewer tokens than template-heavy alternatives like SuperCorrect.

Experimental Results

• Performance: D-RPC achieved the highest accuracy on all five benchmarks for both student models. Notably, on the challenging MATH benchmark, D-RPC outperformed the next-best method by +3.53 points (Llama) and +10.51 points (Qwen) over standard CoT.
• Ablation Studies:
  • Bank Size: Performance peaked at an intermediate bank size (75 paths for GSM8K), confirming the theoretical trade-off. Too few paths caused coverage gaps; too many diluted reuse.
  • Initialization: Heterogeneous initialization (diverse seed set) provided benefits at low data volumes, but the gap vanished with sufficient training data.
  • Mechanism: Disabling the bank (reverting to freeform generation) consistently degraded performance, confirming that path-guided supervision provides a lower-variance signal.
• Theoretical Validation: Empirical measurements of the bound's components confirmed that banking reduces the frozen-model NLL and generalization gap, aligning with the theoretical predictions.

Significance and Claims

The paper claims that D-RPC addresses the upstream variance in teacher rationales, providing a cleaner, more consistent signal for student training. Unlike prior work that restructures rationales on the student side or uses verbose templates that burden SLMs, D-RPC compresses the reasoning space before generation. The authors assert that effective distillation depends on compressing and reusing the reasoning space rather than generating longer or more verbose rationales. The results suggest that structuring the supervision distribution prior to training is a lightweight and effective strategy for improving SLM reasoning capabilities, particularly in domains with high solution heterogeneity. The work remains modest regarding its scope, noting that experiments were limited to specific teacher/student pairs and reasoning domains (math and commonsense), and the theoretical bounds are currently specific to LoRA-based fine-tuning.