Co-Distillation: Training Teacher and Student Together

Both soft-label and hard-label distillation share an assumption: the teacher is already fully trained and fixed before distillation begins — you extract signal from a finished teacher to train a fresh student. Co-distillation breaks that assumption entirely: it starts with an UNTRAINED Teacher LLM and an UNTRAINED Student LLM, and trains them SIMULTANEOUSLY rather than sequentially.

The process: for each batch, generate softmax probabilities from both models over that batch. The teacher is trained as usual, directly on the hard (ground-truth) labels for the task — normal supervised training, no distillation involved for the teacher's own objective. The student is trained to match its softmax probabilities to those of the teacher, exactly like soft-label distillation's matching objective, except the teacher's own probabilities are themselves evolving over the course of training rather than being fixed from the start.

There's an obvious bootstrapping problem here worth being explicit about: during the initial stages of training, the teacher hasn't learned much yet, so its soft labels aren't accurate — training the student to match a still-learning, still-inaccurate teacher would be counterproductive early on. Co-distillation's fix is to train the student using BOTH the teacher's soft labels AND the ground-truth hard labels simultaneously, rather than relying purely on the teacher's (initially poor) guidance. As training progresses and the teacher genuinely improves, its soft-label signal becomes increasingly useful and the student benefits more from it.

Llama 4 used exactly this technique to train Llama 4 Scout and Maverick from Llama 4 Behemoth — notably, this differs from the DeepSeek-R1-to-Qwen/Llama distillation example (hard-label, post-training, from an already-finished teacher). Co-distillation is specifically a pre-training-stage technique, since it requires both models to genuinely be trained together from the start, which isn't meaningful to do with a teacher that's already finished training.

import torch
import torch.nn.functional as F

def co_distillation_step(
    teacher_logits: torch.Tensor, student_logits: torch.Tensor,
    hard_labels: torch.Tensor, student_soft_weight: float = 0.5,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Both models train on the SAME batch simultaneously — the
    teacher's own objective is normal supervised learning; the
    student blends the (still-improving) teacher's soft signal with
    ground-truth hard labels, to avoid over-trusting an early, weak teacher."""
    # Teacher trains normally, directly on ground truth
    teacher_loss = F.cross_entropy(teacher_logits, hard_labels)

    # Student blends soft-label matching with hard-label supervision
    soft_loss = F.kl_div(
        F.log_softmax(student_logits, dim=-1),
        F.softmax(teacher_logits.detach(), dim=-1),  # detach — student doesn't backprop into teacher
        reduction="batchmean",
    )
    hard_loss = F.cross_entropy(student_logits, hard_labels)
    student_loss = student_soft_weight * soft_loss + (1 - student_soft_weight) * hard_loss

    return teacher_loss, student_loss