🤖 AI

Making Expert Reasoning Learnable with Self-Distillation

The paper proposes Distribution Aligned Imitation Learning (DAIL), a self-distillation method that transforms didactic expert solutions into in-distribution reasoning traces to efficiently enhance LLM reasoning capabilities and generalization with minimal data.

Original authors: Ethan Mendes, Jungsoo Park, Alan Ritter

Published 2026-06-04

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Ethan Mendes, Jungsoo Park, Alan Ritter

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Core Problem: The "Textbook" vs. The "Thinker"

Imagine you are trying to teach a student how to solve incredibly difficult math problems—like those found in the International Math Olympiad.

You have two main ways to teach them:

Trial and Error (Reinforcement Learning): You give the student a problem. If they get it right, you give them a cookie. If they get it wrong, they get nothing.
- The Problem: For super-hard problems, the student almost never gets it right on their own. So, they never get a cookie, and they never learn. It’s like trying to teach someone to play the piano by only rewarding them if they play a perfect concerto on the first try. They’ll just quit.
Copying the Master (Imitation Learning): You show the student a solution written by a world-class expert.
- The Problem: Expert solutions are written for humans, not computers. They are "didactic," meaning they skip the boring, obvious steps to save space. An expert might write, "Using the quadratic formula, we find x = 5." They don’t show the messy algebra in between.
- When the AI tries to copy this, it learns to "skip steps" too. It learns to guess the answer or take shortcuts rather than actually understanding the logic. It’s like a student memorizing the answers in the back of the book without learning how to do the math.

The Solution: DAIL (Distribution Aligned Imitation Learning)

The authors propose a new method called DAIL. Think of it as a two-step translation process that turns "Expert Human Notes" into "AI-Friendly Lessons."

Step 1: The "Mixed Policy" Expansion (Filling in the Blanks)

The AI needs to see the full thought process, not just the summary. But if you just ask the AI to rewrite the expert’s solution, it might cheat or reference the answer too much.

So, the authors use a technique called Mixed Policy Rollouts. Imagine two people working together to write a story:

The Student (The AI): Tries to write the next sentence based on what it knows.
The Teacher (A frozen version of the same AI, but looking at the expert’s answer): Checks if the Student’s sentence is on the right track.

If the Student’s sentence is good, they keep it. If it’s off-track, the Teacher steps in and writes the correct next step. This creates a new, detailed version of the solution that looks like the AI’s natural thinking style but is guided by the expert’s correct path. It fills in all the skipped steps the human expert left out.

Step 2: The "Contrastive" Filter (Stopping the Cheating)

Even with the expanded steps, the AI might still learn bad habits. For example, it might see that the expert got to the number "175" and just force its own math to equal "175" without proper logic. This is called a "rationalization shortcut." It’s like a student seeing the answer is 10, so they just write "2 + 8 = 10" even if the problem was about apples and oranges.

To stop this, DAIL uses a Contrastive Objective. Think of it as a "Good Cop / Bad Cop" training method:

The Good Cop (Teacher): Shows the AI the full, correct, detailed reasoning. The AI is rewarded for matching this.
The Bad Cop (Negative Reference): Shows the AI a "cheated" version of the solution. This version only has the final answer and a few random intermediate numbers, with no logical connection between them.

The AI is trained to avoid sounding like the "Bad Cop." It learns to distinguish between real logical steps and fake shortcuts that just happen to lead to the right number.

The Results: Small Data, Big Gains

The authors tested this on very difficult math problems that the AI couldn't solve on its own.

Efficiency: They used fewer than 1,000 expert solutions. That’s a tiny amount of data for AI training.
Performance: The AI’s ability to solve these hard problems improved significantly (up to 31% better in some tests).
Speed: The AI learned to reason more efficiently. It could solve problems using half the "thinking time" (fewer tokens) compared to before, because it learned the direct, expert paths rather than wandering around in circles.
Generalization: The AI didn’t just get better at the specific math problems it was trained on; it also got better at unrelated science questions (like biology and physics), showing it learned how to think, not just what to answer.

Summary Analogy

Imagine you are a novice chef trying to learn a complex recipe from a master chef.

Old Method 1 (Trial and Error): You try to cook it from scratch. You burn it. You try again. You burn it. You never get it right, so you never learn.
Old Method 2 (Copying): You look at the master’s recipe card. It says, "Sauté until golden, then add spices." You don’t know what "golden" looks like or which spices. You guess, and your dish tastes weird.
DAIL Method:
1. Expansion: A teaching assistant stands next to you. When you hesitate, they gently guide your hand to chop the onions correctly, explaining why you’re doing it. You write down this detailed, step-by-step process.
2. Contrastive Filter: The assistant also shows you a "fake" recipe where someone just dumped random ingredients into a pot and got lucky. They tell you, "Do not cook like this person. Notice how they skipped the heating step? That’s bad. Cook like the master, who explains every step."

By the end, you have a detailed, honest recipe that you can follow, and you’ve learned to avoid lazy shortcuts. You can now cook the dish efficiently and even apply those skills to other recipes.

Technical Summary: Making Expert Reasoning Learnable with Self-Distillation

Problem Statement
The paper addresses the limitation of current methods for improving the reasoning capabilities of Large Language Models (LLMs), particularly Large Reasoning Models (LRMs). Standard Reinforcement Learning with Verifiable Rewards (RLVR) relies on the model sampling a correct solution to generate a training signal. However, for extremely difficult problems where even frontier models fail to find a correct answer, the reward signal is zero, preventing learning. While high-quality human expert solutions exist, directly fine-tuning models on this data via Supervised Fine-Tuning (SFT) or Behavioral Cloning (BC) leads to severe performance degradation. This failure arises because expert solutions are "didactic" and out-of-distribution (OOD) relative to the model’s internal reasoning process; they omit granular intermediate steps and lack explicit search dynamics (e.g., backtracking) that are essential for computational reasoning. Furthermore, naive imitation forces the model to adopt "rationalization shortcuts"—superficial logical leaps that fit the expert’s intermediate results but do not generalize.

Methodology: Distribution Aligned Imitation Learning (DAIL)
The authors propose Distribution Aligned Imitation Learning (DAIL), a two-step self-distillation framework designed to bridge the distributional gap between expert solutions and the student model’s reasoning process.

Generating In-Distribution Learnable Reasoning:
- Teacher Definition: A "teacher" model is defined as the frozen initial weights of the student model, conditioned on the problem and the expert solution.
- Mixed Policy Rollouts: To generate training data that is both grounded in the expert solution and in-distribution for the student, the authors employ a mixed policy rollout process. The student model generates tokens, which are then verified against the teacher’s probability distribution. If the student’s token probability is below a threshold $\tau$ , the teacher’s token is used instead. This approach preserves the student’s natural reasoning flow and artifacts (such as self-correction) while anchoring the generation to the expert’s logic. For non-reasoning instruction models, direct sampling from the teacher with specific prompting is also effective.
- Expansion: This process transforms compressed, didactic expert solutions into detailed, in-distribution reasoning traces that fill in omitted steps.
Contrastive Objective for Learning:
- Problem with NLL: Standard Negative Log-Likelihood (NLL) loss indiscriminately imitates all tokens in the generated traces, including "rationalization shortcuts" where the model forces a derivation to match an intermediate expert result without valid logical progression.
- Negative Reference: To mitigate this, DAIL introduces a contrastive objective. A "negative reference" model is constructed by conditioning the frozen student model on a partial solution $\tilde{s}$ , which consists only of coarse-grained intermediate waypoints (e.g., key numerical or symbolic results) extracted from the expert solution. This conditioning encourages the model to jump between results, characterizing the shortcut behavior.
- Loss Function: The training objective minimizes the KL divergence between the student and the teacher (positive signal) while maximizing the divergence between the student and the negative reference (negative signal). This penalizes the student from learning superficial shortcuts and encourages robust, step-by-step reasoning.

Key Contributions

DAIL Framework: A novel offline learning method that enables models to learn from high-quality expert solutions by first transforming them into in-distribution traces and then applying a contrastive objective to prevent shortcut learning.
Efficiency and Scalability: DAIL is sample-efficient, requiring fewer than 1,000 expert solutions. It is computationally efficient, allowing for asynchronous data generation and requiring only a single copy of model weights in memory during training via LoRA adapters.
Generalization to Non-Verifiable Problems: Unlike RLVR, which requires verifiable rewards, DAIL can train on non-verifiable problems, such as Olympiad proofs, by leveraging expert human solutions.
Public Dataset: The authors release e1-proof, a dataset of 669 non-verifiable Olympiad proof problems with expert solutions.

Results

Performance Gains: On Qwen2.5-7B-Instruct, DAIL achieved up to 31% pass@128 gains on benchmarks like AIME 2024/2025, BeyondAIME, and IMO-AnswerBench compared to the base model. It significantly outperformed RLVR baselines (GRPO, NuRL), which often degraded performance on these difficult tasks due to overfitting on stochastic successes or reward sparsity.
Reasoning Efficiency: On Qwen3 models, DAIL-enabled models achieved comparable or better performance than untrained models with 2× fewer reasoning tokens, indicating more efficient reasoning paths.
Out-of-Domain Generalization: DAIL maintained or improved performance on the GPQA-Diamond benchmark (graduate-level science questions), demonstrating that improvements in mathematical reasoning did not come at the cost of catastrophic forgetting in other domains.
Ablations: The contrastive objective consistently outperformed standard NLL loss. Mixed policy rollouts were crucial for reasoning models (like Qwen3) to avoid artifacts from direct teacher sampling, while direct sampling was sufficient for non-reasoning instruction models.

Significance
The paper claims that DAIL unlocks the value of expert human solutions for problems where standard RLVR fails due to zero reward signals. By converting didactic expert solutions into constructive, in-distribution reasoning traces and suppressing shortcut learning, DAIL enables frontier models to improve their reasoning capabilities using small, high-quality datasets. This approach extends the capabilities of LRMs to non-verifiable domains and promotes more efficient, generalizable reasoning.