Dispersion Loss Counteracts Embedding Condensation and Improves Generalization in Small Language Models

TL;DR

This paper investigates a geometric bottleneck in Transformer representations called embedding condensation, which is more severe in small language models and manifests as token embeddings becoming nearly parallel across layers. The authors quantify layer-wise embedding alignment using pairwise cosine similarities and show that larger models exhibit embedding dispersion rather than condensation. To counteract condensation, they introduce dispersion loss, an auxiliary objective that promotes embedding dispersion with several variants (decorrelation, ℓ2-repel, orthogonalization). Across mid-training and full pre-training, dispersion-aware training yields consistent improvements on 10 benchmarks and notable gains on target tasks, demonstrating a principled, parameter-efficient path to improve generalization in small Transformers without increasing model size.

Abstract

Large language models (LLMs) achieve remarkable performance through ever-increasing parameter counts, but scaling incurs steep computational costs. To better understand LLM scaling, we study representational differences between LLMs and their smaller counterparts, with the goal of replicating the representational qualities of larger models in the smaller models. We observe a geometric phenomenon which we term $\textbf{embedding condensation}$, where token embeddings collapse into a narrow cone-like subspace in some language models. Through systematic analyses across multiple Transformer families, we show that small models such as $\texttt{GPT2}$ and $\texttt{Qwen3-0.6B}$ exhibit severe condensation, whereas the larger models such as $\texttt{GPT2-xl}$ and $\texttt{Qwen3-32B}$ are more resistant to this phenomenon. Additional observations show that embedding condensation is not reliably mitigated by knowledge distillation from larger models. To fight against it, we formulate a dispersion loss that explicitly encourages embedding dispersion during training. Experiments demonstrate that it mitigates condensation, recovers dispersion patterns seen in larger models, and yields performance gains across 10 benchmarks. We believe this work offers a principled path toward improving smaller Transformers without additional parameters.

Figures (7)

• Figure 1: Illustration of the embedding condensation phenomenon. In pre-trained language models, embeddings of all tokens from the same input sequence condense into a narrow cone after being processed by many Transformer layers. This phenomenon is substantially more pronounced in smaller models than in larger models within the same family, which led to our hypothesis in Section \ref{['sec:hypothesis']}.
• Figure 2: Qualitative and quantitative observations of the embedding condensation phenomenon. a. The cosine similarity heatmaps demonstrate that smaller models (e.g., GPT2, Qwen3-0.6B) are susceptible to condensation, since token cosine similarities become increasingly positive as the embeddings proceed to deeper layers. In contrast, larger models (e.g., GPT2-xl, Qwen3-32B) are more resistant to embedding condensation. b. Quantifications using Spearman correlation and Kendall's Tau demonstrate a consistent trend of "larger model, less condensation" across multiple families of language models. Additional results can be found in Figure \ref{['fig:supp_observation']}.
• Figure 3: Knowledge distillation is not a remedy to embedding condensation, shown qualitatively (panel a) and quantitatively (panel b).
• Figure 4: Embedding condensation is observed immediately after model initialization. We analyze checkpoints of Olmo-3-1025-7B spanning initialization, intermediate pre-training stages, and the final base model. Each checkpoint is annotated by its training stage and the number of training tokens.
• Figure 5: Illustration of how dispersion loss and its alternative formulations promote embedding dispersion. a. Dispersion loss enforces uniform angular dispersion by spreading out all pairs along the unit hypersphere. b. Decorrelation loss encourages different feature dimensions to remain uncorrelated. c.$\ell_2$-repel loss increases pairwise Euclidean distance, while the norm regularization prevents unbounded expansion. d. Orthogonalization loss spreads out vectors forming acute angles while leaving cobtuse ones unchanged.

Theorems & Definitions (6)

• Definition 1.1
• Definition 1.2
• Proposition 5.1
• proof
• Proposition 5.2
• proof