Dispersion loss counteracts embedding condensation in small language models
摘要
该文提出 embedding condensation 现象：相同序列的 token 嵌入在小语言模型的 Transformer 层中会塌陷为狭窄锥形子空间，现象在小模型中更明显（GPT2 vs GPT2-xl、Qwen3-0.6B vs Qwen3-32B）。通过控制变量预训练 GPT2-like 模型仅改变 MLP 维度等实验验证了 “更大模型更抗塌陷” 的趋势，现象在初始化时就存在，且不受输入数据集影响；知识蒸馏也无法传递抗塌陷能力。作者设计 dispersion loss（受 Diffuse and Disperse 启发，含对数和指数和等变式）作为正则化，在中 / 预训练阶段分散嵌入，能缓解塌陷并缩小小模型与大模型的性能差距。
荐读理由
这篇论文直接给出小模型 token 嵌入在初始化阶段就坍缩成窄锥体的几何事实，以及通过中训练加入 dispersion loss 能实质缓解该现象并缩小与大模型表征质量差距的量化对比（图7），让我的独立 AI 工程创业决策从“参数越多越好”转向“用简单正则器补齐小模型表征空间”
原文
This paper presents an observation-driven improvement on language model training.
We observe a geometric phenomenon which we term embedding condensation, where token embeddings collapse into a narrow cone-like subspace in smaller language models. We then design a training objective called dispersion loss to counteract the effect.
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 motivates our hypothesis in Section 3.3.
Feature 1: Larger model, less condensation. Within the same model family, smaller models exhibit more severe embedding condensation, with token embeddings collapsing toward near-parallel directions, while larger models resist this collapse.
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 S1.
````````
This effect is also quite robust to the choice of input datasets.
Figure S2. The embedding condensation effect is consistent regardless of the input text dataset. Results are shown for four datasets, namely (a) wikitext, (b) pubmed_qa, (c) imdb, and (d) squad.
********
Feature 2: Reproducible when controlling for confounders. To isolate the effect of model size from other confounding factors, we conduct a controlled experiment where we pre-train GPT2-like models, varying only the MLP dimension while keeping all other components fixed, including the number of layers, embedding dimension, dataset, and training settings. The same phenomenon is observed.
Figure 3. In a highly controlled experiment, we reproduced the observation of “larger model, less condensation”. We pre-trained four GPT2-like models of varying sizes that differ only in MLP dimension, while keeping all other factors fixed, including the number of layers, embedding dimension, dataset, and training configuration. The resulting models exhibit consistent trends in embedding condensation, shown qualitatively (panel a) and quantitatively (panel b). Horizontal dashed lines are added to panel a for easier visual comparison.
``************
Feature 3: Condensation occurs early on. The embedding condensation phenomenon emerges at model initialization and is gradually mitigated, not exacerbated, by pre-training.
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.
Feature 4: Distillation is not a solution. Knowledge distillation from a larger model does not transfer the desired resistance to embedding condensation.
Figure 5. Knowledge distillation is not a remedy to embedding condensation, shown qualitatively (panel a) and quantitatively (panel b).
Dispersion loss Embedding condensation reduces the expressivity of Transformers by collapsing token embedding vectors into narrow cones, under-utilizing the representation space. We hypothesize that by dispersing embeddings during training, smaller models can achieve representational qualities more similar to larger models, thus narrowing the performance gap without increasing the number of parameters.
Figure 6. 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. ℓ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 obtuse ones unchanged.
Our dispersion loss is inspired by the "Diffuse and Disperse" paper with practical modifications.
Table 1. Our dispersion loss and its alternative formulations. Main implementation differences from Diffuse and Disperse are highlighted in teal and magenta. Including or excluding diagonal terms yields identical gradients and is therefore cosmetic. For dispersion loss and ℓ2-repel, we adopt the log-sum-exp trick for numerical stability, which differs from log(mean(exp(·))) only by an additive constant. For ℓ2-repel, we include a norm regularization term to prevent unbounded expansion of embeddings. For Orthogonalization, the distance margin is fixed to 1⁄2 since we use angular distance, where 1⁄2 corresponds to orthogonality and thus serves as the ideal margin.
Dispersion loss counteracts the embedding condensation effect during mid-training and pre-training. A qualitative result is shown below, while more quantitative results can be found in the paper.
<img src="https://raw.githubusercontent.com/ChenLiu-1996/LM-Dispersion/main/assets/results_condensation_counteract.png" alt="" width="2942" height="892">
*Figure 7. Dispersion loss counteracts the embedding condensation phenomenon. a. Starting from condensed embeddings (gray dashed box), mid-training with the default loss has a limited impact (green box). b. In contrast, mid-training with our dispersion loss as a regularizer substantially mitigates embedding condensation (blue box).*
********
**Conclusion**
Larger language models are better than smaller language models, but might not merely because they have more parameters. It can be partially attributed to how they organize the information in the latent representations. We hope to see future efforts along this interesting direction.
这条对你有帮助吗？