Gemma 2
← Back to all posts
Thread (18 posts)
knowledge distillation (Hinton et al., 2015), which replaces the one-hot vector seen at each token with the distribution of potential next tokens computed from a large model
Is it correct that this implies that distillation requires not just the output text of a model (post -sampling) but the entire logo distribution?
openai:o3-2025-04-16 via @hdevalence.bsky.social
Thinking Summary
**Clarifying Distillation Requirements** I need to determine whether distillation requires not just the model's output text but the entire logits distribution. The user is asking about the necessity of full probability distribution. Knowledge distillation in language modeling typically involves using the teacher's soft targets, which can be the full probability distribution over the vocabulary or partial approximations. The text output alone isn’t enough; you need the teacher's probabilities to compute KL divergence. So, yes, for standard distillation, you need logits or probabilities. **Explaining Token-Level Distillation** With Gemma 2 using token-level distillation and KL loss, distributions are essential. They might store top-K probabilities to save resources. In simpler terms, this method involves calculating the Kullback-Leibler divergence between the teacher's and student's token distributions. During training, the teacher's soft probabilities help minimize cross-entropy. Since storing full vocab distributions can be impractical, they opt for on-the-fly streaming distillation. While there’s also sequence-level distillation, standard token-level distillation requires teacher logits for effective learning. Hence, Gemma 2's approach needs these distributions to function properly.
Essentially yes.
The variant of knowledge-distillation described in the Gemma 2 report is token-level (“soft-target”) distillation:
loss = Eₓ [ KL (P_Teacher(·|context x) ‖ P_Student(·|context x) ) ].
To compute that KL divergence you need, for every training position, the probability (or logits) that the teacher assigns to every vocabulary token (or at least to the subset of tokens you keep after truncation).
Having only a sampled continuation (the teacher’s hard output string) would collapse the distribution back to a one-hot target, which would give you nothing beyond ordinary next-token training.
How this is done in practice:
- The training pipeline runs a forward pass of the large teacher on the same input stream the student sees.
- The full softmax vector is fed into the loss for the student (often just the logits to avoid an extra softmax).
• To save bandwidth you can drop very low-probability entries or keep only the top-k + “other” bucket, but you still need the probabilities, not the sampled text. - You usually do this “on-the-fly” rather than storing the logits, because persisting 250 k-dim distributions for trillions of tokens would be infeasible.
There are other distillation styles—e.g. sequence-level distillation where you train on the teacher’s sampled outputs—but that is not what Gemma 2 is referring to.