Distillation: Licensing, Speculative & Reasoning

The goal is not to build a smaller model, but to build a smaller model that remembers what the bigger one learned.

Distill, Memory-Retentive AI Agent

This continuation of Section 17.5 picks up the practical distillation pipeline and extends it to three production-relevant topics: the licensing constraints that decide whether you are even allowed to distill from a given teacher, speculative distillation that uses a small student to accelerate inference of a large model, and reasoning-trace distillation that transfers chain-of-thought capability into a model that otherwise could not afford it.

17.6.1 Distillation Licensing and Usage Policies

Before building a distillation pipeline, you must understand the legal constraints imposed by model providers. Most major LLM providers restrict or prohibit using their model outputs to train competing models. Violating these terms can result in service termination, legal liability, or both. The landscape varies significantly across providers.

17.6.1.1 Provider Policies (as of 2026)

The practical implication is clear: if you plan to distill knowledge from a proprietary API into your own model, verify that the provider's terms permit it. For many production use cases, starting with an open-weight teacher (Llama, Mistral, DeepSeek) avoids licensing uncertainty entirely. When using a proprietary API, the safest approach is to use the provider's own fine-tuning service rather than extracting outputs for external training.

17.6.2 Speculative Distillation

Speculative decoding uses a small "draft" model to propose multiple tokens at once, which a larger model then verifies in a single forward pass. Speculative distillation trains the draft model specifically to mimic the larger model's token distribution, improving the acceptance rate (how often the large model agrees with the draft) and thus the overall throughput. This technique turns distillation into a serving-time optimization rather than just a training-time technique. Figure 17.6.3a shows this propose-and-verify workflow.

Figure 17.6.3b: A distilled draft model proposes tokens that the target model verifies in parallel, multiplying throughput.

17.6.3 Distillation for Reasoning: Chain-of-Thought Preservation

Standard distillation transfers a teacher's input-output mapping to a smaller student, but this approach discards the teacher's intermediate reasoning. When a 70B model solves a math problem, it internally constructs a chain of logical steps before arriving at the answer. A student trained only on final answers learns what to output but not how to reason, limiting its ability to generalize to novel problems. Chain-of-Thought (CoT) distillation addresses this by training the student to reproduce the teacher's explicit reasoning traces alongside the final answers.

Formally, given a question $q$, the teacher produces a reasoning trace $r = (r_1, r_2, \ldots, r_k)$ followed by an answer $a$. Plain answer-only distillation trains the student to maximize $\log p_S(a \mid q)$, marginalizing the reasoning away. CoT distillation instead trains on the joint sequence $(r, a)$:

where the trajectory $y_{1:k+|a|}$ concatenates reasoning steps and the answer. The extra $\log p_S(r \mid q)$ term is what carries the procedural knowledge from teacher to student.

The procedure is straightforward. First, the teacher model generates solutions to a large set of training problems using chain-of-thought prompting, producing step-by-step reasoning followed by a final answer. Second, these complete reasoning traces become the training targets for the student. The student learns to generate the full chain of thought, not just the conclusion. This approach was central to the Distilled DeepSeek-R1 models, where 800K reasoning traces from the R1 teacher were used to fine-tune Qwen and Llama base models. The resulting students (1.5B to 70B parameters) achieved reasoning performance competitive with models several times their size.

CoT distillation introduces a design choice: should the student be trained on the teacher's complete reasoning trace, or should the traces be filtered and curated? In practice, filtering is essential. Teacher models sometimes produce redundant steps, circular reasoning, or incorrect intermediate conclusions that happen to reach the right final answer. A robust pipeline filters traces by (1) verifying the final answer against a ground truth, (2) checking that intermediate steps are logically consistent, and (3) removing excessively long or repetitive reasoning chains. The Orca series demonstrated that curating high-quality explanations (with explicit reasoning structure: "First, ..., Therefore, ..., Finally, ...") produces better students than training on raw, unfiltered teacher outputs.

A key finding from reasoning distillation research is that the student's reasoning format does not need to match the teacher's. A teacher that uses verbose, natural-language reasoning can be distilled into a student that produces compact, structured reasoning (numbered steps or symbolic notation) as long as the training data is reformatted accordingly. This flexibility allows practitioners to optimize the student's reasoning format for their deployment constraints: verbose reasoning for user-facing explanations, compact reasoning for latency-sensitive applications where the chain of thought is used internally but not shown to the user.

# Load teacher (gpt2-medium, frozen) and student (gpt2, trainable).
# The student will learn to mimic the teacher's output distribution.
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
device = "cuda" if torch.cuda.is_available() else "cpu"
teacher = AutoModelForCausalLM.from_pretrained("gpt2-medium").to(device)
teacher.eval()
for p in teacher.parameters():
    p.requires_grad = False
    student = AutoModelForCausalLM.from_pretrained("distilgpt2").to(device)
    tokenizer = AutoTokenizer.from_pretrained("gpt2")
    tokenizer.pad_token = tokenizer.eos_token
    t_params = sum(p.numel() for p in teacher.parameters())
    s_params = sum(p.numel() for p in student.parameters())
    print(f"Teacher: {t_params:,} params")
    print(f"Student: {s_params:,} params")
    print(f"Compression: {t_params / s_params:.1f}x")

import torch
# Distillation loss: KL divergence on temperature-scaled soft targets
# blended with cross-entropy on hard (ground-truth) labels.
import torch.nn.functional as F
def distillation_loss(student_logits, teacher_logits, labels,
temperature=3.0, alpha=0.5):
 """Combined distillation loss with soft and hard targets."""
# TODO: Compute soft targets from teacher (temperature-scaled softmax)
# TODO: Compute soft predictions from student (temperature-scaled log-softmax)
# TODO: KL divergence between student and teacher distributions
# TODO: Standard cross-entropy on hard labels
# TODO: Combine: alpha * T^2 * KL + (1-alpha) * CE
pass
# Quick test
b, s, v = 2, 10, 50257
loss = distillation_loss(torch.randn(b,s,v), torch.randn(b,s,v),
torch.randint(0, v, (b,s)))
print(f"Test loss: {loss.item():.4f}")

Test loss: 11.2438

import torch
# Training loop: forward both teacher and student on each batch,
# compute the combined distillation loss, and update student only.
from datasets import load_dataset
from torch.utils.data import DataLoader
from tqdm import tqdm
dataset = load_dataset("wikitext", "wikitext-2-raw-v1", split="train")
dataset = dataset.filter(lambda x: len(x['text'].strip()) > 50)
def tokenize(examples):
    return tokenizer(examples['text'], truncation=True,
        max_length=128, padding='max_length')
    tokenized = dataset.select(range(2000)).map(
        tokenize, batched=True, remove_columns=['text'])
    tokenized.set_format('torch')
    loader = DataLoader(tokenized, batch_size=8, shuffle=True)
    optimizer = torch.optim.AdamW(student.parameters(), lr=5e-5)
    student.train()
    for epoch in range(2):
        total_loss = 0
        for batch in tqdm(loader, desc=f"Epoch {epoch+1}"):
            input_ids = batch['input_ids'].to(device)
            attn = batch['attention_mask'].to(device)
            student_out = student(input_ids, attention_mask=attn)
            with torch.no_grad():
                teacher_out = teacher(input_ids, attention_mask=attn)
                loss = distillation_loss(
                    student_out.logits, teacher_out.logits, input_ids)
                loss.backward()
                optimizer.step()
                optimizer.zero_grad()
                total_loss += loss.item()
                print(f"Epoch {epoch+1} avg loss: {total_loss / len(loader):.4f}")

Epoch 1: 100%|##########| 250/250 [01:42<00:00, 2.44it/s]
Epoch 1 avg loss: 4.7261
Epoch 2: 100%|##########| 250/250 [01:38<00:00, 2.53it/s]
Epoch 2 avg loss: 3.8934

from transformers import AutoModelForCausalLM
from torch.utils.data import DataLoader
from datasets import load_dataset
import torch
# Evaluate perplexity: compare original student, distilled student,
# and teacher to quantify how much knowledge transferred.
import math
def compute_perplexity(model, eval_loader):
    model.eval()
    total_loss, total_tokens = 0, 0
    with torch.no_grad():
        for batch in eval_loader:
            ids = batch['input_ids'].to(device)
            attn = batch['attention_mask'].to(device)
            out = model(ids, attention_mask=attn, labels=ids)
            total_loss += out.loss.item() * attn.sum().item()
            total_tokens += attn.sum().item()
            return math.exp(total_loss / total_tokens)
        eval_data = load_dataset("wikitext", "wikitext-2-raw-v1", split="validation")
        eval_data = eval_data.filter(lambda x: len(x['text'].strip()) > 50)
        eval_tok = eval_data.select(range(500)).map(
            tokenize, batched=True, remove_columns=['text'])
        eval_tok.set_format('torch')
        eval_loader = DataLoader(eval_tok, batch_size=8)
        original = AutoModelForCausalLM.from_pretrained("distilgpt2").to(device)
        print(f"Teacher PPL: {compute_perplexity(teacher, eval_loader):.2f}")
        print(f"Original student PPL: {compute_perplexity(original, eval_loader):.2f}")
        print(f"Distilled student PPL: {compute_perplexity(student, eval_loader):.2f}")

Teacher PPL: 24.87
Original student PPL: 46.13
Distilled student PPL: 36.41

# Complete distillation lab: load teacher/student, train with KL
# divergence loss, evaluate perplexity before and after distillation.
from transformers import AutoModelForCausalLM, AutoTokenizer
from datasets import load_dataset
from torch.utils.data import DataLoader
import torch, torch.nn.functional as F, math
from tqdm import tqdm
device = "cuda" if torch.cuda.is_available() else "cpu"
teacher = AutoModelForCausalLM.from_pretrained("gpt2-medium").to(device)
teacher.eval()
for p in teacher.parameters(): p.requires_grad = False
student = AutoModelForCausalLM.from_pretrained("distilgpt2").to(device)
tokenizer = AutoTokenizer.from_pretrained("gpt2")
tokenizer.pad_token = tokenizer.eos_token
def distillation_loss(s_logits, t_logits, labels, T=3.0, alpha=0.5):
    soft_t = F.softmax(t_logits / T, dim=-1)
    soft_s = F.log_softmax(s_logits / T, dim=-1)
    kl = F.kl_div(soft_s, soft_t, reduction='batchmean')
    ce = F.cross_entropy(s_logits.view(-1, s_logits.size(-1)), labels.view(-1))
    return alpha * (T**2) * kl + (1 - alpha) * ce
ds = load_dataset("wikitext", "wikitext-2-raw-v1", split="train")
ds = ds.filter(lambda x: len(x['text'].strip()) > 50)
def tok(ex): return tokenizer(ex['text'], truncation=True, max_length=128, padding='max_length')
train_tok = ds.select(range(2000)).map(tok, batched=True, remove_columns=['text'])
train_tok.set_format('torch')
loader = DataLoader(train_tok, batch_size=8, shuffle=True)
opt = torch.optim.AdamW(student.parameters(), lr=5e-5)
student.train()
for epoch in range(2):
    total = 0
    for batch in tqdm(loader, desc=f"Epoch {epoch+1}"):
        ids, attn = batch['input_ids'].to(device), batch['attention_mask'].to(device)
        s_out = student(ids, attention_mask=attn)
        with torch.no_grad(): t_out = teacher(ids, attention_mask=attn)
        loss = distillation_loss(s_out.logits, t_out.logits, ids)
        loss.backward(); opt.step(); opt.zero_grad(); total += loss.item()
        print(f"Epoch {epoch+1}: {total/len(loader):.4f}")
def ppl(model, dl):
    model.eval(); tl, tt = 0, 0
    with torch.no_grad():
        for b in dl:
            i, a = b['input_ids'].to(device), b['attention_mask'].to(device)
            o = model(i, attention_mask=a, labels=i)
            tl += o.loss.item()*a.sum().item(); tt += a.sum().item()
            return math.exp(tl/tt)
        ev = load_dataset("wikitext","wikitext-2-raw-v1",split="validation")
        ev = ev.filter(lambda x: len(x['text'].strip())>50)
        ev_t = ev.select(range(500)).map(tok, batched=True, remove_columns=['text'])
        ev_t.set_format('torch')
        el = DataLoader(ev_t, batch_size=8)
        orig = AutoModelForCausalLM.from_pretrained("distilgpt2").to(device)
        print(f"Teacher: {ppl(teacher,el):.2f}")
        print(f"Original: {ppl(orig,el):.2f}")
        print(f"Distilled: {ppl(student,el):.2f}")

Epoch 1: 100%|##########| 250/250 [01:42<00:00, 2.44it/s]
Epoch 1: 4.7261
Epoch 2: 100%|##########| 250/250 [01:38<00:00, 2.53it/s]
Epoch 2: 3.8934
Teacher: 24.87
Original: 46.13
Distilled: 36.41