Soft-Label Distillation: Maximum Knowledge Transfer (and Its Memory Problem)

Soft-label distillation is the first and highest-fidelity of the three distillation techniques. The process: use a fixed, pre-trained Teacher LLM to generate softmax probabilities over the entire corpus — not just the teacher's single top prediction per token, but its FULL probability distribution across every token in the vocabulary, for every position. That same data is then passed through the untrained Student LLM to get its own softmax probabilities, and the student is trained to match the teacher's probabilities as closely as possible.

The reason this transfers the most knowledge of the three techniques is exactly because of what's in that full distribution: a teacher predicting 'cat' with 70% confidence but also assigning 20% to 'dog' and 5% to 'kitten' is communicating something genuinely useful about HOW it reasons — which alternatives it considered plausible and by how much — not just its single final answer. This visibility over the teacher's full reasoning (via its probability distribution) is what gives soft-label distillation its name and its strength.

But this technique has two real costs. First, you must have access to the teacher's weights to actually compute its output probability distribution — this rules out soft-label distillation for any closed-weight teacher model you can only query through an API. Second, even when you DO have weight access, there's a genuinely enormous memory problem: since you're generating softmax probabilities for EVERY input token over the ENTIRE vocabulary, the storage cost explodes with scale. With a 100k-token vocabulary and a 5-trillion-token training corpus, storing soft labels at float8 precision would require roughly 500 million GB of memory — a genuinely infeasible number for real training pipelines.

This memory problem is exactly what motivates the next technique, hard-label distillation — it deliberately sacrifices some of soft-label distillation's rich signal in exchange for a dramatically smaller memory footprint.

# Simplified soft-label distillation loss — matching the STUDENT's
# full probability distribution to the TEACHER's full distribution.
import torch
import torch.nn.functional as F

def soft_label_distillation_loss(
    student_logits: torch.Tensor, teacher_logits: torch.Tensor, temperature: float = 2.0,
) -> torch.Tensor:
    # Softened distributions reveal more about relative confidence
    # between tokens than a sharp, near-one-hot distribution would
    student_log_probs = F.log_softmax(student_logits / temperature, dim=-1)
    teacher_probs = F.softmax(teacher_logits / temperature, dim=-1)

    # KL divergence: how far the student's full distribution is from
    # the teacher's full distribution — not just matching the top token
    return F.kl_div(student_log_probs, teacher_probs, reduction="batchmean") * (temperature ** 2)

# Illustrating the memory problem with real numbers:
vocab_size = 100_000
corpus_tokens = 5_000_000_000_000  # 5 trillion
bytes_per_prob_fp8 = 1
total_gb = (vocab_size * corpus_tokens * bytes_per_prob_fp8) / (1024 ** 3)
print(f"Soft-label storage at fp8: ~{total_gb:,.0f} GB")  # ~500 million GB