Hard-Label Distillation: Solving the Memory Blowup

Hard-label distillation directly solves soft-label distillation's memory problem by changing what gets extracted from the teacher. Instead of capturing the teacher's full softmax probability distribution over the entire vocabulary, hard-label distillation uses a fixed, pre-trained Teacher LLM to just get its final one-hot output token — the single word or token the teacher actually produced, not its full internal probability spread across every alternative.

The training procedure otherwise looks similar to soft-label distillation: the untrained Student LLM processes the same data to get its own softmax probabilities, and the student is trained to match the teacher's output — but now that target is a single hard label per token position (effectively a one-hot vector) rather than a rich probability distribution across the whole vocabulary.

This is a dramatic memory improvement: instead of storing a full probability vector (size = vocabulary size) for every single token position in the training corpus, you only need to store one token ID per position — reducing the storage requirement by roughly the vocabulary size itself (a factor of 100,000x smaller in the earlier worked example). This is exactly what makes hard-label distillation practical at the scale modern LLM training operates at, where soft-label distillation's storage cost was simply infeasible.

The trade-off, as you'd expect, is some loss of signal: a one-hot hard label can't communicate 'the teacher was 70% confident in this token but also considered these two alternatives plausible' the way a full soft distribution can — you only know what the teacher's top choice was, not how confident it was or what else it considered. This is a real fidelity trade-off, but in practice it's a trade-off worth making given the memory numbers involved. DeepSeek used exactly this technique to distill DeepSeek-R1's knowledge into the smaller Qwen and Llama 3.1 models.

import torch
import torch.nn.functional as F

def hard_label_distillation_loss(student_logits: torch.Tensor, teacher_hard_labels: torch.Tensor) -> torch.Tensor:
    # teacher_hard_labels: just the teacher's argmax token IDs, not a
    # full probability distribution — standard cross-entropy against
    # a single "correct" target per position, exactly like normal
    # supervised training, except the "ground truth" IS the teacher's output
    return F.cross_entropy(student_logits, teacher_hard_labels)

def get_teacher_hard_labels(teacher_logits: torch.Tensor) -> torch.Tensor:
    return teacher_logits.argmax(dim=-1)  # just the single top token, not the full distribution

# Memory comparison: hard labels need only 1 token ID per position,
# vs. an entire probability vector per position for soft labels
vocab_size = 100_000
corpus_tokens = 5_000_000_000_000
hard_label_gb = (corpus_tokens * 4) / (1024 ** 3)  # 4 bytes per int32 token ID
print(f"Hard-label storage: ~{hard_label_gb:,.0f} GB")  # dramatically smaller than soft-label's ~500M GB