Off-policy Distillation and Teacher-Student Learning

This technique is behind the explosive growth of the open-source model ecosystem.

It's a methodology where a small student model imitates hundreds of thousands of correct answers generated by a massive Teacher model like GPT-4 or Llama-3-70B, instead of relying on expensive human labeling.

Does a Student Acquire the Same Reasoning Ability by Learning Good Answers Created by a Teacher? #

"'If we train an 8B model on 1 million GPT-4 answers, will the 8B model become as smart as GPT-4?' This is the most common expectation and misconception among AI engineers."

To put it bluntly, Student models often only imitate the Teacher's superficial style and format, rather than acquiring the Teacher's reasoning ability itself. In academia, this is referred to as the False Promise of Imitation Learning or Imitation Gap.

• Style Imitation: The Student quickly and perfectly learns logical development formats, such as words frequently used by the Teacher (e.g., "first," "therefore," "in conclusion").
• Reasoning Bottleneck: However, when presented with complex math problems or multi-step logical puzzles, while the format delivers perfect step-by-step solutions, there are hallucinations where intermediate calculations are completely wrong. This is because the Student's limited parameter capacity cannot fully encompass the Teacher's deep internal space.

SFT is Off-Policy evaluated with human-generated correct answer data, while Off-Policy Distillation is Off-Policy evaluated by a Teacher model. Let's remember this difference.

Conceptual Understanding of Off-policy #

In the context of reinforcement learning and behavioral cloning, the term Off-Policy refers to the source of the data.

• On-Policy: This is a methodology where the currently training model (student) evaluates trajectories or answers it directly generates and learns from that feedback. (RLHF, PPO)
• Off-Policy: This refers to cases where the training data is generated from an external, fixed distribution (the Teacher model's policy $\pi_{teacher}$), rather than the current model's probability distribution $\pi_{\theta}$.

Off-Policy Distillation might involve an answer distribution that is very easy and natural for the Teacher model, but an extremely difficult token distribution for the small Student model's parameter structure to reach.

As a result, the model may experience a collapse of its distribution (Distribution Shift) by trying to force memorization of correct answers.

Building a Synthetic Dataset Generation Pipeline #

The first step in Distillation is the generation of a large quantity of Synthetic Data, which involves using a high-speed inference engine like vLLM to produce Teacher model responses for countless prompts.

Below is an example pipeline for building a fixed JSONL format dataset using Python and vLLM.

import json
from vllm import LLM, SamplingParams

# 1. Load a powerful Teacher Model (using vLLM engine)
# In a real environment, multi-GPU (TP) setup is required.
teacher_model_id = "meta-llama/Meta-Llama-3-70B-Instruct"
llm = LLM(model=teacher_model_id, tensor_parallel_size=4)

# 2. List of diverse instructions (Seed Prompts)
# Typically, tens of thousands to millions of instructions are prepared. (e.g., Alpaca, Evol-Instruct methods)
prompts = [
    "Explain why list comprehensions are memory efficient in Python.",
    "Explain quantum entanglement using an analogy, as if explaining to an elementary school student.",
    # ... tens of thousands of prompts ...
]

# 3. Set Sampling Parameters (ensuring diversity and quality)
# Lower the temperature to encourage logical and refined answers.
sampling_params = SamplingParams(
    temperature=0.3,
    top_p=0.9,
    max_tokens=1024
)

# 4. Execute large-scale parallel inference via vLLM (Generation)
print("Teacher Model is generating answers...")
outputs = llm.generate(prompts, sampling_params)

# 5. Save results to disk in JSONL (or Parquet) format (for Student model's SFT dataset)
output_file = "synthetic_distillation_dataset.jsonl"
with open(output_file, 'w', encoding='utf-8') as f:
    for output in outputs:
        prompt = output.prompt
        generated_text = output.outputs[0].text.strip()
        
        # Save according to SFT training specifications
        data_point = {
            "instruction": prompt,
            "output": generated_text
        }
        f.write(json.dumps(data_point, ensure_ascii=False) + '\n')

print(f"Distillation dataset generation complete: {output_file}")

The synthetic_distillation_dataset.jsonl file created this way is directly fed into the SFTTrainer from the previous chapter to train the Student model Llama-3-8B.

Distillation Results and Imitation Gap Analysis Report #

This is a qualitative comparison of two models, tested with new logical problems requiring reasoning, after distilling 100,000 Teacher (70B) generated data points into a Student (8B) using the SFT method.

Test Prompt - Mathematical Reasoning Test #

User: "There are 35 chickens and pigs combined on a farm. Counting all their legs, there are 94. Logically calculate step-by-step how many chickens and pigs there are, respectively."

This problem can be solved using a system of linear equations.
Step 1: Set up variables
- Let x be the number of chickens, and y be the number of pigs.
Step 2: Formulate equations
- Total number of animals: x + y = 35
- Total number of legs: 2x + 4y = 94 (chickens have 2 legs, pigs have 4)
Step 3: Solve the equations
- Substitute x = 35 - y from the first equation into the second equation.
- 2(35 - y) + 4y = 94
- 70 - 2y + 4y = 94
- 2y = 24
- y = 12 (12 pigs)
- x = 35 - 12 = 23 (23 chickens)
Conclusion: There are 23 chickens and 12 pigs.

Shows perfect logical development and accurate arithmetic operations.

Student Model (8B 0 Distilled) Output - Training Results #

This problem can be solved using a system of linear equations.
Step 1: Set up variables
- Let x be the number of chickens, and y be the number of pigs.
Step 2: Formulate equations
- Total number of animals: x + y = 35
- Total number of legs: 2x + 4y = 94
Step 3: Solve the equations
- Multiply the first equation by 2: 2x + 2y = 70
- Subtract the two equations: (2x + 4y) - (2x + 2y) = 94 - 70
- 2y = 14  <-- [Error: 94-70 is 24, but calculated as 14]
- y = 7 (7 pigs)
- x = 35 - 7 = 28 (28 chickens)
Conclusion: There are 28 chickens and 7 pigs.

• Evaluation - Imitation Gap Manifestation: The Student replicated the Teacher's step-by-step format, variable setting, and even the method of solving simultaneous equations. However, due to the limitations of the smaller model, a hallucination occurred in the intermediate arithmetic (94 - 70 = 14, when it should be 24), leading to a collapse of the final logic.

In conclusion, when using Off-Policy Distillation, it's crucial to recognize that simple SFT alone cannot replicate the Teacher's intelligence. The base model's fundamental capacity, including its scale and the quality of pre-training data, determines the upper limit of its reasoning ability. Distillation should be understood and designed merely as a catalyst to make the model perform well within that upper limit.

What's Wrong with the Way It's Trained? #

SFT and Distillation are essentially a game of Next Token Prediction.

It involves presenting the correct answer sheet (the Teacher's perfect response) and instructing the model, "This word absolutely must come after this context!"

In machine learning terminology, this is called Teacher Forcing.

So, What's the Problem? #

This rote learning approach has two critical issues:

• Exposure Bias: During training, the model only follows the path of perfectly crafted correct answers. However, in the actual service inference stage, the model must generate words one by one itself. If it makes a mistake, such as outputting a wrong number in the middle, the model has never learned how to recover from its own errors and return to the correct answer. Consequently, its logic collapses like dominoes from that moment on.
• Omission of Thought (Superficial Imitation): When a large model solves a math problem, it undergoes incredibly deep and complex internal reasoning before returning the final text. However, a small model only sees the superficial text and doesn't perceive the depth of thought. Therefore, it doesn't understand why mathematical formulas are developed in a certain way, merely memorizing the "shell" of the solution process, thinking, "For problems like this, if I arrange numbers in roughly this order, I'll be praised."

Solution #

Therefore, we must overcome the limitations of rote learning that only memorizes the shell of correct answers.

The solution is to flip the training paradigm from force-feeding text answers to "a method that gives a large reward when the model finally gets the correct answer after countless trials and errors."

Like learning to ride a bicycle by falling and instinctively figuring out how to balance.

The master key to AI evolution, moving beyond a simple 'word predictor' to enable models to develop logic and think for themselves.

There are reinforcement learning methods like RLHF and DPO that fundamentally restructure the model's "brain" through human feedback and rewards, as well as the recently astonishing reasoning-focused reinforcement learning methods like OpenAI o1 and DeepSeek-R1's GRPO.

We will look at this next time.