Low-Mask-Only Distillation in UltraBERT
Pruning makes a large encoder compact, but distillation can behave differently at different masking ratios; UltraBERT applies teacher-logit distillation only to low-masked examples and keeps standard masked-language-model training for high-masked examples.
Read the paper ↗What it is
- Selective teacher guidance A compact encoder learns from a large teacher on lightly masked text, but uses ordinary masked-word training on heavily masked text.
- Prune, then distill The encoder is made smaller by removing width, then trained with guidance from the larger model; the process can be repeated for a still smaller encoder.
Why it matters
Teacher guidance is not automatically useful for every kind of training example. Separating the objectives lets UltraBERT use distillation where the reported comparison supports it, without applying it to high-masked examples that showed no benefit in that comparison.
Before getting into the setup, here is the basic idea: make a smaller encoder by pruning a larger one, then let the large model guide the student only on lightly masked examples. When many words are hidden, the student instead does ordinary masked-word prediction. The interesting move is not just compression; it is deciding when the teacher is useful.
-
A smaller model still needs a learning signal: pruning removes parts of the larger encoder, so distillation can help the resulting student recover useful behavior. The canonical example here is the 1.3B encoder pruned to 380M. The smaller model is not trained in isolation: it is distilled from its larger parent. More generally, pruning shrinks width while distillation supplies teacher logits as guidance. But does that guidance help under every masking condition?
-
Masking changes how informative the teacher can be: compare a sentence with only a few masked words to one with half its tokens masked. In the first case, the model has substantial context; in the second, the prediction problem is much harder and the teacher's signal may be less useful. The reported comparison is at 15% and 50% masking: at 15%, distillation improves GLUE from 90.2 to 90.8 and leaves BEIR at 46.0; at 50%, the distilled model's BEIR score falls from 47.2 to 46.5 against the from-scratch model. The suggested reason is that the teacher signal at 50% could be trivial, though that remains an interpretation. So why not give the teacher's advice only where it appears useful?
-
Route the objectives by masking regime: UltraBERT mixes 5%-masked and 50%-masked sequences in a 4:1 ratio and gives each ratio its own language-model head. The low-mask examples receive KL divergence loss from the teacher; the high-mask examples receive the standard masked-language-model objective. In everyday terms, one notebook page is compared with the teacher's answer, while the other is learned from the masked-word task alone. That analogy breaks because these are separate heads and losses, not literal pages or a human decision. The source describes the loss routing but gives no displayed loss equation here. How does the compact model get built before this selective training?
The pruning itself is a cascade: prune the 1.3B encoder to 380M and distill, then prune that model to 150M and distill. Width is reduced while keeping the 24 layers; activations are used to rank channels, attention heads, and feed-forward neurons for retention. My take is that the useful lesson is the routing rule, not a claim that teacher guidance always helps: the observed benefit depends on which examples receive it.
Read the deep dive
Why People Wanted It
A smaller encoder can inherit useful behavior from a larger one, but the teacher's guidance is not equally helpful under every masking regime. In the reported comparison, distillation at 15% masking improved GLUE for the pruned model while leaving BEIR unchanged; at 50% masking, the distilled model did worse than its from-scratch counterpart on BEIR. So the real question was not simply whether to distill, but where the teacher's logits should be trusted.
The Core Move
Prune the large encoder into a compact student, then use teacher-logit distillation only on low-masked examples; train high-masked examples with the ordinary masked-language-model objective. The source gives this as a training rule, not as a displayed loss formula: low-mask examples receive KL divergence loss, and high-mask examples receive standard MLM loss.
Why Distillation Depends on the Masking Ratio
Now the version for readers who want more of the machinery: the key distinction is the information available to the teacher at each masking ratio. The comparison described here uses 15% and 50% masking: after pruning the 1.3B encoder to 380M, each student is distilled from its 1.3B parent, with the corresponding masking ratio. At 15%, the distilled model's GLUE score rises from 90.2 to 90.8 while BEIR stays at 46.0; at 50%, BEIR falls from 47.2 to 46.5 relative to the from-scratch model. The proposed interpretation is that the teacher signal at 50% masking may be trivial, but the result is stated as a possible explanation, not a settled cause. If the teacher's guidance can be weak at high masking, why apply it there by default?
Why UltraBERT Keeps Distillation on Low-Masked Examples
Here is the part that gets technical: UltraBERT separates the two training signals by example type. Its batches mix 5%-masked and 50%-masked sequences at a 4:1 ratio, with a separate language-model head for each masking ratio. The 5% head receives KL divergence loss from the 1.3B teacher; the 50% head receives standard MLM loss. The source specifies these objectives and routing, but does not provide a displayed loss equation here. This is not distillation everywhere with a different weight: high-masked examples are explicitly trained without the teacher-logit objective. How does the student become smaller before this split training begins?
Why Pruning Is a Cascade
Now one level deeper: the compact models are produced by repeated width pruning and distillation, while retaining all 24 layers. The 1.3B encoder is pruned to 380M and distilled; that 380M model is then pruned to 150M and distilled. The stated pruning procedure scores hidden channels, attention heads, and feed-forward neurons from activations, then keeps the top-scoring units. The training description says both compact models are distilled from the 1.3B teacher, while the appendix describes the intermediate cascade; the essential mechanism is pruning to reduce width, followed by teacher guidance. How much capability survives this sequence is a tradeoff, not something pruning alone guarantees.
Minimal Causal Code
This plain PyTorch sketch makes the routing visible. It represents the described objectives; it is not a complete reproduction of the training system.
for tokens, mask_ratio in batches:
student_logits = student(tokens, mask_ratio=mask_ratio)
if mask_ratio == 0.05:
with torch.no_grad():
teacher_logits = teacher(tokens, mask_ratio=mask_ratio)
loss = kl_divergence(student_logits, teacher_logits)
else:
loss = masked_language_model_loss(student_logits, tokens)
loss.backward()
optimizer.step()
optimizer.zero_grad()
What You Gain, and What You Lose
You retain teacher guidance on low-masked examples while avoiding teacher-logit distillation on the high-masked examples where the reported comparison did not show a BEIR gain. The cost is a more specialized training setup: pruning, two masking regimes, separate heads, and different objectives routed by masking ratio. The reported 15% result is not a universal win across tasks, and the high-mask degradation is not assigned a definitive cause. Distillation is therefore selective supervision, not a guarantee that a smaller student will improve everywhere.
UltraBERT does not distill every example: it reserves teacher logits for low masking and lets high-masked examples learn through standard MLM.