LM-Dispersion
The official code and project write-up for an ICML 2026 paper showing that smaller language models compress their token embeddings into a narrow region, and proposing a dispersion loss to counteract it.
LM-Dispersion is the author's repository for an ICML 2026 paper on a training-time effect the authors call embedding condensation, where token embeddings in smaller language models collapse into a narrow, cone-like subspace, reducing the effective representation space available to the model. Within a model family, smaller models show more of this collapse than larger ones, the effect appears early in training rather than growing during it, and knowledge distillation from a larger model does not on its own transfer resistance to it.
The repository implements a dispersion loss, inspired by a related paper called Diffuse and Disperse, that actively disperses embeddings during mid-training and pre-training to counteract condensation and narrow the performance gap between small and large models without adding parameters.
The project page has an illustrated walkthrough of the paper, and the official code is also hosted at the authors' lab repository.