Annotations Mitigate Post-Training Mode Collapse

Post-training via supervised fine-tuning improves instruction-following, but often induces semantic mode collapse by biasing models toward low-entropy fine-tuning data at the expense of the high-entropy pretraining distribution. Crucially, we find this trade-off worsens with scale. To close this semantic diversity gap, we propose annotation-anchored training, a principled method that enables models to adopt the preference-following behaviors of post-training without sacrificing the inherent diversity of pretraining. Our approach is simple: we pretrain on documents paired with semantic annotations, inducing a rich annotation distribution that reflects the full breadth of pretraining data, and we preserve this distribution during post-training. This lets us sample diverse annotations at inference time and use them as anchors to guide generation, effectively transferring pretraining's semantic richness into post-trained models. We find that models trained with annotation-anchored training can attain 6x less diversity collapse than models trained with SFT, and improve with scale.

Modern AI language models learn from enormous amounts of text, which lets them write about almost any topic in many different styles. To turn them into helpful assistants, developers fine-tune them on curated examples of good responses. This extra step has a hidden cost: the models become repetitive, offering the same few kinds of answers and losing much of their original variety—and, surprisingly, this gets worse for larger models. We show the problem comes from what fine-tuning rewards, which pushes the model to copy the narrow range of its training examples. To fix this, we introduce annotation-anchored training. During the model's initial training, we attach short descriptive labels to each document, teaching it a wide range of possible directions a response could take. During fine-tuning, we let the model improve its answers while keeping that range of directions unchanged. When the model is used, it first picks a direction and then writes a response for it. This restores diversity—six times less collapse than standard fine-tuning—while keeping the assistant just as helpful, and the benefit grows with model size. More broadly, it offers a recipe for improving an AI system without erasing the valuable variety it learned early on.