Knowledge Distillation

# MiniLLM: Reverse KL Divergence for LLM Distillation

Based on arXiv 2306.08543 (2024) - MiniLLM: Knowledge Distillation of Large Language Models

## Overview

**Source**: https://arxiv.org/abs/2306.08543
**GitHub**: https://github.com/microsoft/LMOps/tree/main/minillm

MiniLLM replaces forward KLD with reverse KLD for knowledge distillation, achieving better performance on generative language models.

## Problem with Standard KLD

### Forward KL Divergence (Standard)

**Formula**: `KL(Student || Teacher)`

**Minimization behavior**: Mode-seeking
```
Student tries to match teacher's MEAN behavior
→ Student focuses on teacher's highest probability regions
→ Student ignores low-probability but valid generations
```

**Issue for generative models**: Limits diversity, student generates safe but boring outputs.

### Why Forward KL Fails for Generation

```python
# Teacher distribution (diverse)
teacher_probs = [0.3, 0.3, 0.2, 0.1, 0.1]  # Multiple valid options

# Forward KL minimization
# Student learns: [0.6, 0.3, 0.1, 0.0, 0.0]
# Problem: Ignores options 4-5 entirely (mode-seeking)
```

## MiniLLM Solution: Reverse KLD

### Reverse KL Divergence

**Formula**: `KL(Teacher || Student)`

**Minimization behavior**: Mode-covering
```
Student tries to COVER all teacher's modes
→ Student learns diverse generation
→ Student doesn't ignore any valid teacher outputs
```

### Mathematical Formulation

**Forward KL** (standard distillation):
```
L_forward = Σ p_student(x) log(p_student(x) / p_teacher(x))
          = E_{x~student} [log p_student(x) - log p_teacher(x)]
```

**Reverse KL** (MiniLLM):
```
L_reverse = Σ p_teacher(x) log(p_teacher(x) / p_student(x))
          = E_{x~teacher} [log p_teacher(x) - log p_student(x)]
```

**Key difference**: Expectation over teacher distribution vs student distribution.

## Implementation

### Reverse KLD Loss

```python
import torch
import torch.nn.functional as F

def reverse_kl_loss(student_logits, teacher_logits, temperature=1.0):
    """
    Reverse KL divergence: KL(Teacher || Student).

    Args:
        student_logits: Model predictions (batch, seq_len, vocab_size)
        teacher_logits: Teacher predictions (batch, seq_len, vocab_size)
        temperature: Softening parameter

    Returns:
        Reverse KL divergence loss
    """
    # Teacher distribution (target, detached)
    p_teacher = F.softmax(teacher_logits / temperature, dim=-1)
    p_teacher = p_teacher.detach()  # Don't backprop through teacher

    # Student distribution (learnable)
    log_p_student = F.log_softmax(student_logits / temperature, dim=-1)

    # Reverse KL: -Σ p_teacher * log p_student
    reverse_kl = -(p_teacher * log_p_student).sum(dim=-1).mean()

    # Temperature correction
    return reverse_kl * (temperature ** 2)
```

### Policy Gradient Optimization

**Challenge**: Reverse KL requires sampling from teacher.

**Solution**: Use policy gradient with teacher samples.

```python
def minillm_policy_gradient(student_model, teacher_model, prompt_batch):
    """
    MiniLLM training with policy gradient.

    Steps:
    1. Sample responses from teacher
    2. Compute reverse KL using those samples
    3. Optimize student to cover teacher's distribution
    """
    # 1. Generate from teacher (detached)
    with torch.no_grad():
        teacher_outputs = teacher_model.generate(
            prompt_batch,
            max_new_tokens=256,
            do_sample=True,
            temperature=1.0,
            return_dict_in_generate=True,
            output_scores=True
        )

        teacher_sequences = teacher_outputs.sequences
        teacher_scores = teacher_outputs.scores

    # 2. Student evaluates teacher's samples
    student_outputs = student_model(
        input_ids=teacher_sequences,
        labels=teacher_sequences
    )

    # 3. Policy gradient loss
    # Maximize student's likelihood on teacher's samples
    loss = -student_outputs.logits.mean()

    return loss
```

## Training Procedure

### Two-Stage