Knowledge Distillation: Training Small Models to Match Large Ones

A 70B parameter language model achieves remarkable quality. A 7B model is 10x cheaper to serve. Knowledge distillation is the technique that bridges this gap: train the small model to mimic the large one, recovering 85-95% of the teacher’s quality at a fraction of the inference cost. This post covers the mathematics, implementation, and empirical results of distillation applied to large language models.

The core idea is simple. Training on hard labels (one-hot targets from the dataset) discards most of the information in the teacher’s output distribution. When a teacher model assigns probability 0.7 to “cat”, 0.15 to “kitten”, 0.05 to “feline”, and 0.001 to “table”, those relative probabilities encode knowledge about semantic similarity that hard labels throw away. Distillation trains the student on these soft probability distributions, transferring the teacher’s learned structure.

1. The Teacher-Student Framework

1.1 Setup

Given:

• Teacher model $T$: a large, well-trained model with parameters $\theta_T$ (frozen during distillation)
• Student model $S$: a smaller model with parameters $\theta_S$ (trained during distillation)
• Training data $\mathcal{D} = \{(x_i, y_i)\}$: input sequences and ground-truth labels

The student is trained to minimize a combination of two losses:

$\mathcal{L} = \alpha \cdot \mathcal{L}_{\text{distill}} + (1 - \alpha) \cdot \mathcal{L}_{\text{hard}}$

where $\mathcal{L}_{\text{hard}}$ is the standard cross-entropy loss against ground-truth labels, and $\mathcal{L}_{\text{distill}}$ is the distillation loss that matches the student’s output to the teacher’s.

import torch
import torch.nn as nn
import torch.nn.functional as F

class DistillationLoss(nn.Module):
    """Combined distillation + hard label loss."""

    def __init__(self, temperature=2.0, alpha=0.5):
        super().__init__()
        self.temperature = temperature
        self.alpha = alpha

    def forward(self, student_logits, teacher_logits, targets):
        """
        Args:
            student_logits: [batch, seq_len, vocab_size]
            teacher_logits: [batch, seq_len, vocab_size]
            targets: [batch, seq_len] (ground-truth token IDs)
        """
        # Hard label loss: standard cross-entropy
        hard_loss = F.cross_entropy(
            student_logits.view(-1, student_logits.size(-1)),
            targets.view(-1),
            ignore_index=-100,
        )

        # Soft label loss: KL divergence with temperature
        soft_loss = self.kl_divergence_loss(student_logits, teacher_logits)

        return self.alpha * soft_loss + (1 - self.alpha) * hard_loss

    def kl_divergence_loss(self, student_logits, teacher_logits):
        """KL divergence between teacher and student soft distributions."""
        T = self.temperature

        # Soften distributions with temperature
        student_probs = F.log_softmax(student_logits / T, dim=-1)
        teacher_probs = F.softmax(teacher_logits / T, dim=-1)

        # KL(teacher || student) — note: teacher is the target distribution
        kl_loss = F.kl_div(
            student_probs.view(-1, student_probs.size(-1)),
            teacher_probs.view(-1, teacher_probs.size(-1)),
            reduction='batchmean',
        )

        # Scale by T^2 to make gradients comparable across temperatures
        return kl_loss * (T * T)

1.2 Why Soft Labels Contain More Information

Consider a vocabulary of $V = 32000$ tokens. A hard label is a one-hot vector with 1 bit of useful information per position (which token is correct). A soft label is a full probability distribution over 32000 tokens, containing approximately $-\sum_i p_i \log p_i$ bits of information (the entropy of the teacher’s distribution).

For a typical LLM, the teacher’s output entropy is 3-8 bits per token, compared to 0 bits from the one-hot (excluding the single correct token). The soft label provides 3-8x more training signal per example.

More precisely, the teacher’s distribution encodes:

1. Semantic similarity: tokens with similar meanings get similar probabilities
2. Syntactic constraints: grammatically valid continuations get higher probability
3. Uncertainty: ambiguous contexts produce flatter distributions, signaling genuine uncertainty rather than model failure

def analyze_teacher_output(teacher_logits, targets, tokenizer, top_k=10):
    """Examine what information is in the teacher's soft distribution."""
    probs = F.softmax(teacher_logits[0, -1], dim=-1)  # Last position
    target_token = targets[0, -1].item()

    # Top-k tokens and their probabilities
    topk_probs, topk_indices = probs.topk(top_k)

    print(f"Target token: '{tokenizer.decode([target_token])}' "
          f"(prob: {probs[target_token]:.4f})")
    print(f"\nTop-{top_k} teacher predictions:")
    for i, (prob, idx) in enumerate(zip(topk_probs, topk_indices)):
        token = tokenizer.decode([idx.item()])
        marker = " <-- target" if idx.item() == target_token else ""
        print(f"  {i+1}. '{token}': {prob:.4f}{marker}")

    # Entropy of teacher distribution
    entropy = -(probs * probs.clamp(min=1e-10).log()).sum()
    print(f"\nTeacher entropy: {entropy:.2f} nats ({entropy/0.693:.2f} bits)")
    print(f"Hard label entropy: 0 bits (one-hot)")
    print(f"Information gain from soft labels: {entropy/0.693:.2f} bits/token")

# Example output for "The capital of France is ___":
# Target token: 'Paris' (prob: 0.82)
# Top-10 teacher predictions:
#   1. 'Paris': 0.8200 <-- target
#   2. ' Paris': 0.0650
#   3. 'paris': 0.0180
#   4. 'Lyon': 0.0120
#   5. 'the': 0.0080
#   6. 'Pars': 0.0045
#   7. 'Par': 0.0032
#   8. 'located': 0.0028
#   9. 'Marseille': 0.0025
#   10. 'known': 0.0022
# Teacher entropy: 1.24 nats (1.79 bits)

The teacher’s distribution tells the student that “Paris” and ” Paris” (with leading space) are nearly interchangeable, that “Lyon” and “Marseille” are at least plausible (they are French cities), and that “table” or “running” are essentially impossible. A hard label says only “Paris is correct; everything else is equally wrong.”

2. Temperature Scaling

2.1 The Role of Temperature

The softmax function with temperature $T$ is:

$p_i = \frac{\exp(z_i / T)}{\sum_j \exp(z_j / T)}$

where $z_i$ are the logits (pre-softmax scores). Temperature controls the “softness” of the distribution:

• $T = 1$: standard softmax. The teacher’s natural distribution.
• $T > 1$: softer distribution. More probability mass on low-probability tokens. More information about the teacher’s relative preferences.
• $T \to \infty$: uniform distribution. All information is lost.
• $T \to 0$: hard distribution. Converges to one-hot on the argmax. Equivalent to hard labels.

def demonstrate_temperature(logits, temperatures):
    """Show how temperature affects the output distribution."""
    print(f"Raw logits: {logits.tolist()}")
    print()

    for T in temperatures:
        probs = F.softmax(logits / T, dim=-1)
        entropy = -(probs * probs.log()).sum().item()
        print(f"T={T:.1f}: probs={[f'{p:.4f}' for p in probs.tolist()]}, "
              f"entropy={entropy:.4f}")

# Example: 5-class problem
logits = torch.tensor([5.0, 3.0, 1.0, 0.5, -1.0])
demonstrate_temperature(logits, [0.5, 1.0, 2.0, 5.0, 10.0])
# T=0.5: probs=[0.9672, 0.0297, 0.0024, 0.0006, 0.0001], entropy=0.2064
# T=1.0: probs=[0.8360, 0.1131, 0.0153, 0.0093, 0.0021], entropy=0.6718
# T=2.0: probs=[0.5547, 0.2340, 0.0987, 0.0765, 0.0361], entropy=1.2845
# T=5.0: probs=[0.3222, 0.2434, 0.1838, 0.1670, 0.0836], entropy=1.5413
# T=10.0: probs=[0.2527, 0.2254, 0.2011, 0.1916, 0.1292], entropy=1.5873

2.2 Why T = 2-4 Works Best

At $T = 1$, the teacher’s distribution is often very peaked: the top token has 80-95% probability. Most of the information about relative token similarities is compressed into the remaining 5-20% of probability mass, spread across 32000 tokens. The gradients for low-probability tokens are tiny.

At $T = 2$, the distribution is softer. The top token might have 55% probability, and the next 10 tokens have meaningful probabilities (1-10% each). The student receives stronger gradient signals for these secondary tokens, learning the teacher’s preference structure more efficiently.

At $T = 5$, the distribution is too flat. The teacher’s fine-grained preferences are washed out. The student learns that many tokens are roughly equally plausible, which is not useful.

2.3 Optimal Temperature Selection

def find_optimal_temperature(
    teacher_model, student_model, val_dataloader,
    temperatures, alpha=0.5, device='cuda',
):
    """Search for the temperature that minimizes validation loss."""
    results = {}

    for T in temperatures:
        total_loss = 0
        n_batches = 0
        loss_fn = DistillationLoss(temperature=T, alpha=alpha)

        with torch.no_grad():
            for batch in val_dataloader:
                input_ids = batch['input_ids'].to(device)
                targets = batch['labels'].to(device)

                teacher_out = teacher_model(input_ids).logits
                student_out = student_model(input_ids).logits
                loss = loss_fn(student_out, teacher_out, targets)

                total_loss += loss.item()
                n_batches += 1

        avg_loss = total_loss / n_batches
        results[T] = avg_loss
        print(f"T={T:.1f}: val_loss={avg_loss:.4f}")

    best_T = min(results, key=results.get)
    print(f"\nBest temperature: {best_T}")
    return best_T

# Typical result:
# T=1.0: val_loss=3.245
# T=2.0: val_loss=3.102  <-- often best
# T=3.0: val_loss=3.118
# T=4.0: val_loss=3.156
# T=5.0: val_loss=3.231

3. Feature Distillation: Matching Intermediate Representations

3.1 Why Final-Layer Distillation Is Not Enough

Output-only distillation gives the student a single supervision signal: match the teacher’s final token distribution. But a 70B teacher with 80 layers computes 80 intermediate representations, each encoding progressively more abstract features. A 7B student with 32 layers must learn all these feature transformations in fewer steps. Feature distillation provides additional supervision by matching intermediate layer outputs.

3.2 Layer Mapping

The student has $L_S$ layers and the teacher has $L_T$ layers. We need a mapping $m(i) \to j$ that pairs student layer $i$ with teacher layer $j$. Common strategies:

Uniform mapping: match every $k$-th teacher layer, where $k = L_T / L_S$:

$m(i) = \text{round}(i \times L_T / L_S)$

Skip mapping: match only the last few layers (output layers carry the most task-relevant features).

Learned mapping: train a small projector network that maps student representations to teacher space, jointly optimized during distillation.

def uniform_layer_mapping(n_student_layers, n_teacher_layers):
    """Map student layers to teacher layers uniformly."""
    mapping = {}
    for s in range(n_student_layers):
        t = round(s * n_teacher_layers / n_student_layers)
        t = min(t, n_teacher_layers - 1)
        mapping[s] = t
    return mapping

# Example: 32-layer student, 80-layer teacher
mapping = uniform_layer_mapping(32, 80)
print(f"Student layer 0  -> Teacher layer {mapping[0]}")
print(f"Student layer 15 -> Teacher layer {mapping[15]}")
print(f"Student layer 31 -> Teacher layer {mapping[31]}")
# Student layer 0  -> Teacher layer 0
# Student layer 15 -> Teacher layer 38
# Student layer 31 -> Teacher layer 78

3.3 Projection Layers

The teacher’s hidden dimension ($d_T = 8192$ for 70B) differs from the student’s ($d_S = 4096$ for 7B). A learned linear projection aligns the dimensions:

class FeatureDistillationModule(nn.Module):
    """Match student intermediate representations to teacher's."""

    def __init__(self, student_dim, teacher_dim, n_pairs):
        super().__init__()
        # One projector per matched layer pair
        self.projectors = nn.ModuleList([
            nn.Linear(student_dim, teacher_dim, bias=False)
            for _ in range(n_pairs)
        ])

    def forward(self, student_hidden_states, teacher_hidden_states, layer_mapping):
        """
        Args:
            student_hidden_states: dict of {layer_idx: tensor [B, S, d_S]}
            teacher_hidden_states: dict of {layer_idx: tensor [B, S, d_T]}
            layer_mapping: dict of {student_layer: teacher_layer}
        Returns:
            Feature distillation loss (scalar)
        """
        total_loss = 0
        n_pairs = 0

        for pair_idx, (s_layer, t_layer) in enumerate(layer_mapping.items()):
            s_hidden = student_hidden_states[s_layer]  # [B, S, d_S]
            t_hidden = teacher_hidden_states[t_layer]  # [B, S, d_T]

            # Project student to teacher dimension
            s_projected = self.projectors[pair_idx](s_hidden)  # [B, S, d_T]

            # Normalize both representations (cosine similarity objective)
            s_norm = F.normalize(s_projected, dim=-1)
            t_norm = F.normalize(t_hidden.detach(), dim=-1)

            # MSE loss on normalized representations
            loss = F.mse_loss(s_norm, t_norm)
            total_loss += loss
            n_pairs += 1

        return total_loss / n_pairs

3.4 Attention Transfer

Beyond hidden states, we can match attention patterns. The teacher’s attention weights encode which tokens should attend to which — this structural information can accelerate student training:

class AttentionDistillation(nn.Module):
    """Match student attention patterns to teacher's."""

    def __init__(self, n_student_heads, n_teacher_heads, n_pairs):
        super().__init__()
        self.n_student_heads = n_student_heads
        self.n_teacher_heads = n_teacher_heads
        # If head counts differ, we average teacher heads in groups
        self.heads_per_group = n_teacher_heads // n_student_heads

    def forward(self, student_attentions, teacher_attentions, layer_mapping):
        """
        Args:
            student_attentions: dict of {layer: tensor [B, n_heads_S, S, S]}
            teacher_attentions: dict of {layer: tensor [B, n_heads_T, S, S]}
        """
        total_loss = 0
        n_pairs = 0

        for s_layer, t_layer in layer_mapping.items():
            s_attn = student_attentions[s_layer]  # [B, H_S, S, S]
            t_attn = teacher_attentions[t_layer]  # [B, H_T, S, S]

            # Average teacher heads in groups to match student head count
            # H_T=64 -> H_S=32 means average pairs of teacher heads
            t_attn_grouped = t_attn.view(
                t_attn.size(0),
                self.n_student_heads,
                self.heads_per_group,
                t_attn.size(2),
                t_attn.size(3),
            ).mean(dim=2)  # [B, H_S, S, S]

            # KL divergence on attention distributions (already normalized)
            loss = F.kl_div(
                s_attn.log().clamp(min=-100),
                t_attn_grouped.detach(),
                reduction='batchmean',
            )
            total_loss += loss
            n_pairs += 1

        return total_loss / n_pairs

4. Online Distillation

4.1 Concept

In standard (offline) distillation, the teacher is pre-trained and frozen. The student trains on the teacher’s fixed outputs. Online distillation removes this separation: teacher and student train simultaneously on the same data.

The advantage: the teacher’s distribution evolves during training, providing a moving target that can be more informative than a fixed one. The teacher starts uncertain (high entropy) and gradually becomes confident, naturally providing a curriculum from soft to hard labels.

4.2 Implementation

class OnlineDistillation:
    """Teacher and student train simultaneously."""

    def __init__(self, teacher_model, student_model, temperature=2.0,
                 alpha=0.5, teacher_lr=1e-4, student_lr=3e-4):
        self.teacher = teacher_model
        self.student = student_model
        self.temperature = temperature
        self.alpha = alpha

        self.teacher_optimizer = torch.optim.AdamW(
            teacher_model.parameters(), lr=teacher_lr
        )
        self.student_optimizer = torch.optim.AdamW(
            student_model.parameters(), lr=student_lr
        )
        self.distill_loss = DistillationLoss(temperature, alpha)

    def train_step(self, batch):
        """One step of online distillation."""
        input_ids = batch['input_ids']
        targets = batch['labels']

        # Step 1: Teacher forward (with gradients for teacher training)
        teacher_logits = self.teacher(input_ids).logits
        teacher_loss = F.cross_entropy(
            teacher_logits.view(-1, teacher_logits.size(-1)),
            targets.view(-1),
            ignore_index=-100,
        )

        # Step 2: Student forward
        student_logits = self.student(input_ids).logits
        student_loss = self.distill_loss(
            student_logits,
            teacher_logits.detach(),  # Stop gradient to teacher for student loss
            targets,
        )

        # Step 3: Update both
        self.teacher_optimizer.zero_grad()
        teacher_loss.backward()
        self.teacher_optimizer.step()

        self.student_optimizer.zero_grad()
        student_loss.backward()
        self.student_optimizer.step()

        return {
            'teacher_loss': teacher_loss.item(),
            'student_loss': student_loss.item(),
        }

4.3 Mutual Distillation

A variant where two models of similar size teach each other:

class MutualDistillation:
    """Two models teach each other (Deep Mutual Learning)."""

    def __init__(self, model_a, model_b, temperature=2.0):
        self.model_a = model_a
        self.model_b = model_b
        self.T = temperature
        self.opt_a = torch.optim.AdamW(model_a.parameters(), lr=3e-4)
        self.opt_b = torch.optim.AdamW(model_b.parameters(), lr=3e-4)

    def train_step(self, batch):
        input_ids = batch['input_ids']
        targets = batch['labels']

        # Both forward passes
        logits_a = self.model_a(input_ids).logits
        logits_b = self.model_b(input_ids).logits

        # Model A loss: hard labels + match Model B
        ce_a = F.cross_entropy(
            logits_a.view(-1, logits_a.size(-1)), targets.view(-1),
            ignore_index=-100,
        )
        kl_a = self._kl_loss(logits_a, logits_b.detach())
        loss_a = ce_a + kl_a

        # Model B loss: hard labels + match Model A
        ce_b = F.cross_entropy(
            logits_b.view(-1, logits_b.size(-1)), targets.view(-1),
            ignore_index=-100,
        )
        kl_b = self._kl_loss(logits_b, logits_a.detach())
        loss_b = ce_b + kl_b

        # Update both
        self.opt_a.zero_grad()
        loss_a.backward()
        self.opt_a.step()

        self.opt_b.zero_grad()
        loss_b.backward()
        self.opt_b.step()

        return {'loss_a': loss_a.item(), 'loss_b': loss_b.item()}

    def _kl_loss(self, student_logits, teacher_logits):
        T = self.T
        s = F.log_softmax(student_logits / T, dim=-1)
        t = F.softmax(teacher_logits / T, dim=-1)
        return F.kl_div(
            s.view(-1, s.size(-1)), t.view(-1, t.size(-1)),
            reduction='batchmean',
        ) * (T * T)

5. Self-Distillation

5.1 The Model Teaching Itself

Self-distillation uses the model’s own outputs as the teacher signal. This sounds circular, but it works because the teacher signal comes from a different context than the student’s training:

1. Temporal self-distillation: Use a past checkpoint (EMA or snapshot) as teacher
2. Multi-view self-distillation: The model generates outputs from augmented inputs; the student trains on the original input to match
3. Multi-token prediction self-distillation: DeepSeek V3’s approach — the model’s main head teaches auxiliary prediction heads

5.2 EMA Self-Distillation

class EMASelfDistillation:
    """Self-distillation using an exponential moving average of the model."""

    def __init__(self, model, ema_decay=0.999, temperature=2.0, alpha=0.3):
        self.model = model
        self.ema_model = self._create_ema(model)
        self.ema_decay = ema_decay
        self.temperature = temperature
        self.alpha = alpha
        self.optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4)
        self.distill_loss = DistillationLoss(temperature, alpha)

    def _create_ema(self, model):
        """Create EMA copy of model (no gradients)."""
        import copy
        ema = copy.deepcopy(model)
        for param in ema.parameters():
            param.requires_grad_(False)
        return ema

    @torch.no_grad()
    def _update_ema(self):
        """Update EMA parameters."""
        for ema_param, model_param in zip(
            self.ema_model.parameters(), self.model.parameters()
        ):
            ema_param.data.mul_(self.ema_decay).add_(
                model_param.data, alpha=1 - self.ema_decay
            )

    def train_step(self, batch):
        input_ids = batch['input_ids']
        targets = batch['labels']

        # Teacher: EMA model (no gradients)
        with torch.no_grad():
            teacher_logits = self.ema_model(input_ids).logits

        # Student: current model
        student_logits = self.model(input_ids).logits
        loss = self.distill_loss(student_logits, teacher_logits, targets)

        self.optimizer.zero_grad()
        loss.backward()
        self.optimizer.step()

        # Update EMA after each step
        self._update_ema()

        return loss.item()

5.3 Multi-Token Prediction Self-Distillation (DeepSeek V3)

DeepSeek V3 trains auxiliary prediction heads that predict tokens 2, 3, …, $k$ positions ahead. The main model’s next-token prediction head serves as the teacher for these auxiliary heads:

class MTPSelfDistillation(nn.Module):
    """Multi-Token Prediction with self-distillation (DeepSeek V3 style)."""

    def __init__(self, d_model, vocab_size, n_extra_heads=3):
        super().__init__()
        self.n_extra_heads = n_extra_heads

        # Main prediction head (standard next-token)
        self.main_head = nn.Linear(d_model, vocab_size, bias=False)

        # Auxiliary heads predict tokens 2, 3, ... positions ahead
        # Each has its own small transformer layer + projection
        self.aux_layers = nn.ModuleList([
            nn.TransformerEncoderLayer(
                d_model=d_model,
                nhead=16,
                dim_feedforward=d_model * 4,
                batch_first=True,
            )
            for _ in range(n_extra_heads)
        ])
        self.aux_heads = nn.ModuleList([
            nn.Linear(d_model, vocab_size, bias=False)
            for _ in range(n_extra_heads)
        ])

    def forward(self, hidden_states, targets):
        """
        Args:
            hidden_states: [B, S, d_model] from the main transformer
            targets: [B, S] token IDs
        """
        B, S, D = hidden_states.shape

        # Main head: predict next token (position i predicts token i+1)
        main_logits = self.main_head(hidden_states)  # [B, S, V]
        main_loss = F.cross_entropy(
            main_logits[:, :-1].reshape(-1, main_logits.size(-1)),
            targets[:, 1:].reshape(-1),
            ignore_index=-100,
        )

        # Auxiliary heads: predict tokens i+2, i+3, ...
        aux_loss = 0
        distill_loss = 0
        h = hidden_states
        for k, (aux_layer, aux_head) in enumerate(
            zip(self.aux_layers, self.aux_heads)
        ):
            h = aux_layer(h)
            aux_logits = aux_head(h)  # [B, S, V]

            # Offset for k-th auxiliary: predict token at position i+k+2
            offset = k + 2
            if S > offset:
                # Hard label loss for auxiliary head
                aux_ce = F.cross_entropy(
                    aux_logits[:, :-offset].reshape(-1, aux_logits.size(-1)),
                    targets[:, offset:].reshape(-1),
                    ignore_index=-100,
                )
                aux_loss += aux_ce

                # Self-distillation: match main head's distribution
                # Main head's prediction at position i+k+1 is the "teacher"
                # for auxiliary head's prediction at position i
                teacher_logits = main_logits[:, (offset-1):-(1)].detach()
                student_logits = aux_logits[:, :-(offset)]

                T = 2.0
                kl = F.kl_div(
                    F.log_softmax(student_logits / T, dim=-1).reshape(-1, aux_logits.size(-1)),
                    F.softmax(teacher_logits / T, dim=-1).reshape(-1, main_logits.size(-1)),
                    reduction='batchmean',
                ) * (T * T)
                distill_loss += kl

        total_loss = main_loss + 0.3 * aux_loss + 0.1 * distill_loss
        return total_loss, {
            'main_loss': main_loss.item(),
            'aux_loss': (aux_loss / self.n_extra_heads if isinstance(aux_loss, torch.Tensor) else 0),
            'distill_loss': (distill_loss / self.n_extra_heads if isinstance(distill_loss, torch.Tensor) else 0),
        }

6. Complete Distillation Training Loop

Here is a production-quality distillation implementation:

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
import math
import time

class DistillationTrainer:
    """Complete distillation training pipeline."""

    def __init__(
        self,
        teacher_model,
        student_model,
        train_dataloader,
        val_dataloader,
        temperature=2.0,
        alpha=0.5,
        feature_distill=False,
        feature_weight=0.1,
        learning_rate=3e-4,
        warmup_steps=1000,
        total_steps=100000,
        grad_clip=1.0,
        device='cuda',
    ):
        self.teacher = teacher_model.to(device).eval()
        self.student = student_model.to(device)
        self.train_dl = train_dataloader
        self.val_dl = val_dataloader
        self.device = device
        self.grad_clip = grad_clip
        self.total_steps = total_steps
        self.warmup_steps = warmup_steps

        # Freeze teacher
        for param in self.teacher.parameters():
            param.requires_grad_(False)

        # Losses
        self.distill_loss = DistillationLoss(temperature, alpha)

        # Feature distillation (optional)
        self.feature_distill = feature_distill
        self.feature_weight = feature_weight
        if feature_distill:
            self.feature_module = FeatureDistillationModule(
                student_dim=student_model.config.hidden_size,
                teacher_dim=teacher_model.config.hidden_size,
                n_pairs=min(student_model.config.num_hidden_layers, 8),
            ).to(device)
            self.layer_mapping = uniform_layer_mapping(
                min(student_model.config.num_hidden_layers, 8),
                teacher_model.config.num_hidden_layers,
            )

        # Optimizer
        params = list(self.student.parameters())
        if feature_distill:
            params += list(self.feature_module.parameters())
        self.optimizer = torch.optim.AdamW(params, lr=learning_rate, weight_decay=0.1)

    def get_lr(self, step):
        """Cosine schedule with warmup."""
        if step < self.warmup_steps:
            return self.optimizer.defaults['lr'] * step / self.warmup_steps
        progress = (step - self.warmup_steps) / (self.total_steps - self.warmup_steps)
        return self.optimizer.defaults['lr'] * 0.5 * (1 + math.cos(math.pi * progress))

    def train(self):
        """Main training loop."""
        step = 0
        best_val_loss = float('inf')
        log_interval = 100

        self.student.train()
        train_iter = iter(self.train_dl)

        while step < self.total_steps:
            # Get batch (restart dataloader if exhausted)
            try:
                batch = next(train_iter)
            except StopIteration:
                train_iter = iter(self.train_dl)
                batch = next(train_iter)

            # Update learning rate
            lr = self.get_lr(step)
            for pg in self.optimizer.param_groups:
                pg['lr'] = lr

            # Move to device
            input_ids = batch['input_ids'].to(self.device)
            targets = batch['labels'].to(self.device)
            attention_mask = batch.get('attention_mask', None)
            if attention_mask is not None:
                attention_mask = attention_mask.to(self.device)

            # Teacher forward (no gradients)
            with torch.no_grad():
                teacher_out = self.teacher(
                    input_ids,
                    attention_mask=attention_mask,
                    output_hidden_states=self.feature_distill,
                )

            # Student forward
            student_out = self.student(
                input_ids,
                attention_mask=attention_mask,
                output_hidden_states=self.feature_distill,
            )

            # Output distillation loss
            loss = self.distill_loss(
                student_out.logits, teacher_out.logits, targets
            )

            # Feature distillation loss (optional)
            if self.feature_distill:
                s_hidden = {
                    i: student_out.hidden_states[s_layer]
                    for i, s_layer in enumerate(self.layer_mapping.keys())
                }
                t_hidden = {
                    i: teacher_out.hidden_states[t_layer]
                    for i, t_layer in enumerate(self.layer_mapping.values())
                }
                feat_loss = self.feature_module(
                    s_hidden, t_hidden, dict(enumerate(range(len(self.layer_mapping))))
                )
                loss = loss + self.feature_weight * feat_loss

            # Backward
            self.optimizer.zero_grad()
            loss.backward()
            torch.nn.utils.clip_grad_norm_(self.student.parameters(), self.grad_clip)
            self.optimizer.step()

            # Logging
            if step % log_interval == 0:
                print(f"Step {step}/{self.total_steps}, "
                      f"loss={loss.item():.4f}, lr={lr:.2e}")

            # Validation
            if step % (log_interval * 10) == 0 and step > 0:
                val_loss = self.validate()
                if val_loss < best_val_loss:
                    best_val_loss = val_loss
                    torch.save(self.student.state_dict(), 'best_student.pt')
                    print(f"  New best val_loss: {val_loss:.4f}")
                self.student.train()

            step += 1

        return best_val_loss

    @torch.no_grad()
    def validate(self):
        """Compute validation loss."""
        self.student.eval()
        total_loss = 0
        n_batches = 0

        for batch in self.val_dl:
            input_ids = batch['input_ids'].to(self.device)
            targets = batch['labels'].to(self.device)
            attention_mask = batch.get('attention_mask', None)
            if attention_mask is not None:
                attention_mask = attention_mask.to(self.device)

            teacher_out = self.teacher(input_ids, attention_mask=attention_mask)
            student_out = self.student(input_ids, attention_mask=attention_mask)
            loss = self.distill_loss(student_out.logits, teacher_out.logits, targets)

            total_loss += loss.item()
            n_batches += 1

        avg_loss = total_loss / max(1, n_batches)
        print(f"  Validation loss: {avg_loss:.4f}")
        return avg_loss

6.1 Launching Distillation

from transformers import AutoModelForCausalLM, AutoTokenizer

def run_distillation():
    """End-to-end distillation: 70B teacher -> 7B student."""
    device = 'cuda'

    # Load teacher (frozen, FP16 for memory)
    teacher = AutoModelForCausalLM.from_pretrained(
        "meta-llama/Llama-3-70B",
        torch_dtype=torch.float16,
        device_map="auto",  # Spread across GPUs
    )

    # Load student (trainable, BF16 for training)
    student = AutoModelForCausalLM.from_pretrained(
        "meta-llama/Llama-3-8B",
        torch_dtype=torch.bfloat16,
    )

    tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3-8B")

    # Prepare data
    train_dl = create_dataloader(tokenizer, split='train', batch_size=4, seq_len=2048)
    val_dl = create_dataloader(tokenizer, split='validation', batch_size=4, seq_len=2048)

    # Run distillation
    trainer = DistillationTrainer(
        teacher_model=teacher,
        student_model=student,
        train_dataloader=train_dl,
        val_dataloader=val_dl,
        temperature=2.0,
        alpha=0.5,
        feature_distill=False,  # Output-only for 70B (memory constraint)
        learning_rate=3e-4,
        warmup_steps=1000,
        total_steps=50000,
        device=device,
    )

    best_loss = trainer.train()
    print(f"Distillation complete. Best val loss: {best_loss:.4f}")

    # Save distilled model
    student.save_pretrained("llama-8b-distilled-from-70b")

7. Distillation Quality Analysis

7.1 How Much Quality Does Distillation Retain?

The critical question: how close does the student get to the teacher? The answer depends on the capacity gap, the distillation data, and the task.

7.2 Where Distillation Fails

Distillation does not uniformly transfer all capabilities:

def analyze_distillation_by_task(teacher, student, eval_datasets):
    """Measure quality retention per task category."""
    results = {}

    for task_name, dataset in eval_datasets.items():
        teacher_score = evaluate(teacher, dataset)
        student_score = evaluate(student, dataset)
        retention = student_score / teacher_score * 100

        results[task_name] = {
            'teacher': teacher_score,
            'student': student_score,
            'retention': retention,
        }

    return results

# Typical results (7B distilled from 70B):
# Factual recall (TriviaQA):     teacher=78, student=58, retention=74%
# Reasoning (ARC-Challenge):     teacher=85, student=77, retention=91%
# Code generation (HumanEval):   teacher=67, student=54, retention=81%
# Instruction following (IFEval): teacher=72, student=68, retention=94%
# Math (GSM8K):                  teacher=84, student=71, retention=85%

The pattern: reasoning and instruction following transfer well (these are “skill” capabilities that depend on learned procedures). Factual recall transfers poorly (the student has fewer parameters to memorize facts). Code generation is intermediate (it requires both procedural knowledge and factual recall of APIs).

7.3 Data Requirements

Distillation is more data-efficient than pretraining because each example carries more information (soft labels). But it still requires significant data:

7.4 The Compute Trade-off

Distillation requires running both teacher and student for each batch. The teacher forward pass for a 70B model is approximately 140 TFLOP per 2048-token sequence. The student forward + backward for an 8B model is approximately 48 TFLOP. Total per step: 188 TFLOP.

Standard pretraining of the 8B model alone: 48 TFLOP per step. Distillation is 3.9x more expensive per step. But it achieves better quality in fewer steps (20B tokens for distillation vs 2T tokens for pretraining from scratch). Net compute: distillation uses approximately 20B/2T * 3.9 = 3.9% of the pretraining compute while retaining 85-93% of teacher quality.

def compute_distillation_flops(
    teacher_params_B, student_params_B, seq_len, n_tokens_B,
):
    """Estimate total FLOP for distillation."""
    # Teacher: forward only (no gradients)
    teacher_flops_per_token = 2 * teacher_params_B * 1e9  # 2 * params for forward
    # Student: forward + backward (3x forward)
    student_flops_per_token = 6 * student_params_B * 1e9  # 6 * params for fwd+bwd

    total_flops_per_token = teacher_flops_per_token + student_flops_per_token
    total_flops = total_flops_per_token * n_tokens_B * 1e9

    # Compare to pretraining student from scratch
    pretrain_flops_per_token = 6 * student_params_B * 1e9
    pretrain_tokens = 2e12  # Standard: 2T tokens
    pretrain_total = pretrain_flops_per_token * pretrain_tokens

    print(f"Distillation FLOP: {total_flops:.2e}")
    print(f"Pretraining FLOP:  {pretrain_total:.2e}")
    print(f"Distillation / Pretraining: {total_flops/pretrain_total:.1%}")

compute_distillation_flops(
    teacher_params_B=70,
    student_params_B=8,
    seq_len=2048,
    n_tokens_B=20,
)
# Distillation FLOP: 2.96e+21 (20B tokens * (2*70B + 6*8B) FLOP/token)
# Pretraining FLOP:  9.60e+22 (2T tokens * 6*8B FLOP/token)
# Distillation / Pretraining: 3.1%

8. Advanced Distillation Techniques

8.1 Task-Specific Distillation

Rather than distilling general language modeling ability, distill on data specific to your target task:

class TaskSpecificDistillation(DistillationTrainer):
    """Distillation focused on a specific task distribution."""

    def __init__(self, teacher, student, task_data, general_data,
                 task_weight=0.7, **kwargs):
        super().__init__(teacher, student, task_data, **kwargs)
        self.general_dl = iter(general_data)
        self.task_weight = task_weight

    def get_batch(self):
        """Mix task-specific and general data."""
        if torch.rand(1).item() < self.task_weight:
            return next(self.train_dl)
        else:
            try:
                return next(self.general_dl)
            except StopIteration:
                self.general_dl = iter(self.general_dl)
                return next(self.general_dl)

8.2 Progressive Distillation

Distill in stages through intermediate-sized models:

$70\text{B} \to 33\text{B} \to 13\text{B} \to 7\text{B}$

Each stage has a smaller capacity gap, making distillation more effective. The total compute is higher, but the final 7B model is typically 2-3% better than direct 70B-to-7B distillation:

def progressive_distillation(model_sizes, base_teacher_path, n_tokens_per_stage):
    """Multi-stage distillation through decreasing model sizes."""
    teacher_path = base_teacher_path
    results = []

    for i in range(len(model_sizes) - 1):
        teacher_size = model_sizes[i]
        student_size = model_sizes[i + 1]
        print(f"\nStage {i+1}: {teacher_size}B -> {student_size}B")

        teacher = load_model(teacher_path, teacher_size)
        student = create_model(student_size)

        trainer = DistillationTrainer(
            teacher_model=teacher,
            student_model=student,
            train_dataloader=create_dataloader(batch_size=4),
            val_dataloader=create_dataloader(split='val', batch_size=4),
            temperature=2.0,
            alpha=0.5,
            total_steps=int(n_tokens_per_stage / (4 * 2048)),
        )
        val_loss = trainer.train()

        # Save and use as next teacher
        student_path = f"distilled_{student_size}B_from_{teacher_size}B"
        student.save_pretrained(student_path)
        teacher_path = student_path
        results.append({'stage': f'{teacher_size}B->{student_size}B', 'val_loss': val_loss})

        # Free memory
        del teacher, student
        torch.cuda.empty_cache()

    return results

# Run: 70B -> 33B -> 13B -> 7B
results = progressive_distillation(
    model_sizes=[70, 33, 13, 7],
    base_teacher_path="meta-llama/Llama-3-70B",
    n_tokens_per_stage=10_000_000_000,  # 10B tokens per stage
)

8.3 Distillation with Quantized Teachers

Running a 70B FP16 teacher requires 140 GB of GPU memory. Quantizing the teacher to INT4 reduces this to 35 GB, making distillation feasible on fewer GPUs:

def load_quantized_teacher(model_name, bits=4):
    """Load teacher model with quantization for memory efficiency."""
    from transformers import BitsAndBytesConfig

    quantization_config = BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_compute_dtype=torch.bfloat16,
        bnb_4bit_use_double_quant=True,
        bnb_4bit_quant_type="nf4",
    )

    teacher = AutoModelForCausalLM.from_pretrained(
        model_name,
        quantization_config=quantization_config,
        device_map="auto",
    )

    # Freeze all parameters
    for param in teacher.parameters():
        param.requires_grad_(False)

    return teacher

# Memory comparison:
# FP16 70B: 140 GB (needs 2x A100 80GB)
# INT4 70B: 35 GB  (fits on 1x A100 80GB)
# Quality impact: teacher's soft labels degrade by ~0.5 PPL
# Student quality impact: negligible (within 0.1 PPL of FP16 teacher distillation)

9. Common Failure Modes and Fixes

9.1 Capacity Gap Too Large

When the student is too small relative to the teacher (e.g., distilling 70B into 0.5B), the student cannot represent the teacher’s distribution and distillation provides no benefit over standard training.

Fix: Use progressive distillation or choose a larger student.

9.2 Temperature Too High

Over-softened distributions provide weak gradients for the top predictions. The student learns that many tokens are roughly equally likely but fails to learn the sharp peaks needed for factual recall.

Fix: Use $T = 2$ as default. Validate on held-out data with different temperatures.

9.3 Alpha Imbalance

If $\alpha$ (soft label weight) is too high, the student ignores ground truth and can hallucinate more (it copies teacher mistakes). If $\alpha$ is too low, the student ignores teacher knowledge.

Fix: $\alpha = 0.5$ is a robust default. For factual tasks, reduce to $\alpha = 0.3$.

def diagnose_distillation(teacher, student, val_data, device='cuda'):
    """Diagnose common distillation failure modes."""
    metrics = {}

    with torch.no_grad():
        for batch in val_data:
            input_ids = batch['input_ids'].to(device)
            targets = batch['labels'].to(device)

            t_logits = teacher(input_ids).logits
            s_logits = student(input_ids).logits

            # Check 1: Are student and teacher distributions correlated?
            t_probs = F.softmax(t_logits[:, -1], dim=-1)
            s_probs = F.softmax(s_logits[:, -1], dim=-1)
            correlation = torch.corrcoef(
                torch.stack([t_probs.flatten(), s_probs.flatten()])
            )[0, 1]

            # Check 2: Top-1 agreement
            t_top1 = t_logits[:, -1].argmax(dim=-1)
            s_top1 = s_logits[:, -1].argmax(dim=-1)
            agreement = (t_top1 == s_top1).float().mean()

            # Check 3: KL divergence (should decrease during training)
            kl = F.kl_div(
                F.log_softmax(s_logits[:, -1], dim=-1),
                F.softmax(t_logits[:, -1], dim=-1),
                reduction='batchmean',
            )

            print(f"Correlation: {correlation:.4f} (should be > 0.8)")
            print(f"Top-1 agreement: {agreement:.4f} (should be > 0.5)")
            print(f"KL divergence: {kl:.4f} (should be < 1.0)")
            break

    return metrics

10. Summary

Knowledge distillation is one of the most practical techniques in the LLM deployment toolkit. The key numbers to remember:

1. A distilled 7B model retains 85-93% of a 70B teacher’s quality
2. Distillation uses 3-5% of the compute of pretraining from scratch
3. Temperature $T = 2$ and $\alpha = 0.5$ are robust defaults
4. Feature distillation adds 1-3% quality but doubles memory requirements
5. Progressive distillation (70B to 33B to 7B) adds 2-3% over direct distillation
6. Quantized teachers (INT4) work nearly as well as FP16, halving memory cost

The fundamental trade-off is simple: you pay the one-time cost of training a large teacher, then amortize that investment across many small student deployments. For serving at scale, this is almost always a good deal.

Reviewer Agent Validation

Challenge: Given a teacher model that outputs logits $z_T = [5.0, 3.0, 1.0]$ and a student that outputs logits $z_S = [4.0, 2.5, 1.5]$ for a 3-class problem with temperature $T = 2$, compute the KL divergence distillation loss (including the $T^2$ correction).

Step 1: Compute teacher soft probabilities at $T = 2$:

$p_T = \text{softmax}([5.0/2, 3.0/2, 1.0/2]) = \text{softmax}([2.5, 1.5, 0.5])$

$p_T = [\exp(2.5), \exp(1.5), \exp(0.5)] / Z$ where $Z = 12.1825 + 4.4817 + 1.6487 = 18.3129$

$p_T = [0.6652, 0.2447, 0.0901]$

Step 2: Compute student log soft probabilities at $T = 2$:

$\log q_S = \log\text{softmax}([4.0/2, 2.5/2, 1.5/2]) = \log\text{softmax}([2.0, 1.25, 0.75])$

$q_S = [\exp(2.0), \exp(1.25), \exp(0.75)] / Z$ where $Z = 7.3891 + 3.4903 + 2.1170 = 12.9964$

$q_S = [0.5685, 0.2685, 0.1629]$

$\log q_S = [-0.5647, -1.3147, -1.8147]$

Step 3: KL divergence = $\sum_i p_T(i) \cdot [\log p_T(i) - \log q_S(i)]$

$= 0.6652 \times [\ln(0.6652) - (-0.5647)] + 0.2447 \times [\ln(0.2447) - (-1.3147)] + 0.0901 \times [\ln(0.0901) - (-1.8147)]$

$= 0.6652 \times [-0.4076 + 0.5647] + 0.2447 \times [-1.4076 + 1.3147] + 0.0901 \times [-2.4076 + 1.8147]$

$= 0.6652 \times 0.1571 + 0.2447 \times (-0.0929) + 0.0901 \times (-0.5929)$

$= 0.1045 - 0.0227 - 0.0534 = 0.0284$

Step 4: Apply $T^2$ correction: $\mathcal{L}_{\text{distill}} = 0.0284 \times 2^2 = **0.1136**$.