Dispersion loss counteracts embedding condensation and improves generalization in small language models

大语言模型为什么比小语言模型强？数据？参数？几何结构或许也颇为关键！

现在基于 Transformer 架构的大语言模型的每一层会把每个输入 token 表示为嵌入空间中的高维向量（通常每一层对应的空间维度均相等）。 随着向量逐层前传，我们发现这些向量常常像是被挤进一个狭窄圆锥里： token 表征两两之间指向的方向越来越趋同。我们把这种几何现象称为嵌入塌缩 (embedding condensation)。该现象具有以下特点：

为什么小语言模型比大语言模型弱？一个直观答案当然是参数少、数据少、算力少。 但如果我们打开模型内部，看它每一层的 token 表征，会不会发现一些更具体的差异？

我们最近从表征空间的几何结构出发，观察到一个很有意思的现象。我们称之为嵌入塌缩： 在小语言模型中，token 的表征往往会坍缩到狭窄的锥形子空间内。为了对抗这一现象，我们设计了一个损失函数 dispersion loss 。

特征一：模型越大，塌缩越弱
这个现象在小模型中尤其明显。 在同一个模型家族里，模型越小，token 表征越容易在深层发生这种 “方向挤压”。越大的模型通常对这种现象越有抵抗力。

我们换了四个文本数据集，都能复现这个现象。

特征二：控制变量后依然存在
为了排除干扰因素， 我们还训练了一组基于 GPT2 的变种。 我们特意只改变 MLP 宽度，保持其他所有变量完全一致， 包括层数、嵌入维度、训练数据和训练设置。 结果发现，“模型越大，塌缩越弱”的趋势仍然存在。

这说明嵌入塌缩不只是某个具体模型或某个训练设置下的偶然现象，而可能反映了小语言模型在表征空间利用上的一种系统性限制。

特征三：模型初始化时就已萌芽
有一个略微反直觉的发现是， 这个现象在模型初始化时就已经存在。模型在真正开始学习之前，它的表征空间几何结构就已经表现出这种 “挤在一起”的倾向。预训练会在一定程度上缓解、而非加剧这一现象。

特征四：知识蒸馏无法解决
有人可能会问： 既然大模型不怎么塌缩，那我们直接蒸馏大模型不就解决了吗？ 凭直觉而言，如果一个小模型向一个大模型学习，也许它的表征空间会变得更像大模型。 但实验显示，蒸馏虽然可以约束输出分布，却并不会自然修复中间层表征的几何结构。 换句话说，学生可以学到老师的答案，但不一定学到老师的内部空间组织方式。

解决方法：dispersion loss

基于这些观察，我们提出了一个简单的训练正则项，叫 dispersion loss。

它的想法很直接：如果 token 表征在角度上太挤，那我们就在训练时显式鼓励它们分散开。 这有点像把原本挤成一团的向量摊开，让模型使用更多的表征方向。

我们 dispersion loss 的设计很大程度上借鉴了 凯明大神和 Runqian 的一篇文章，但略作调整。对比如下。

实验结果

我们通过实验发现，这个简单的几何约束可以有效缓解嵌入塌缩。

在 mid-training 和 pre-training 中，使用 dispersion loss 能够给小语言模型带来一定的泛化提升。详情可见最新版文章。

总结

大模型的优势，可能不只是“更大”， 而是它更充分地利用了表示空间。而这也许可以通过训练目标来部分弥补。

若你打算复现本文或借鉴其中部分思路，以下是我们的一点个人建议。

• 嵌入塌缩: 这一部分在我们手中表现稳定：在多种模型家族、输入数据集以及控制实验设置下均可重复观察到。 趋势强弱会因模型家族而异，但关键观察均相当稳定，完全无需 cherry pick。 我们无法保证所有模型家族都会出现该现象，但你完全可以在自己关心的模型上试一试。
• dispersion loss: 这一部分的稳定性有待提高。在我们的实验中，它带来的效果比较有限： 提升幅度较小，若无正式统计检验很难与噪声区分(而我们做的检验也非常基础，毕竟我们不是统计专业出身)。 另外，一位在 Reasoning 方向上更有经验的朋友在论文录用后指出， 我们的 mid-training 配方并不符合主流规范：常见做法往往是在特定领域上继续强化能力，而不是在 wikitext 上做类似预训练的持续训练。 我们的预训练实验也由于训练成本的限制比较有限。 若你对该方法感兴趣，我们建议先在小规模试跑，再决定是否投入大额算力。

我个人觉得以下几方面可能值得关注。

• 更好的正则项。dispersion loss 简单直接，但很可能并非最优解。若能更精细地设计对抗嵌入塌缩的方法，或许会有很大提升空间。
• 预训练以外的世界。可以跟踪嵌入塌缩在监督微调 (SFT)、 强化学习 (RL) 等后续阶段如何演化： 塌缩是会恶化还是缓解、与下游任务是否会耦合等等，这些目前仍不清楚。
• 机制与因果。厘清嵌入塌缩的根因，并在塌缩与泛化等下游行为之间建立更强的因果证据。
• 更好的架构。设计在结构上就更能抵抗塌缩的架构或模块，与纯损失正则互为补充或替代。
• 更好的初始化。设计更优的初始化，使模型从一开始就处于不塌缩的状态。

1. 本工作最初受到论文 A mathematical perspective on Transformers 的启发。我们在 2025 年 4 月初观看了关于该论文的一场报告后启动项目。文中一个理论结果指出： 若无限堆叠 Transformer 层，所有 token 嵌入将聚到同一点； 我们对这一现象能否在实验中被验证很感兴趣，并沿着这个方向展开了关于嵌入塌缩的探索。
2. Dispersion loss 的设计在很大程度上受到 Runqian 与 Kaiming 的论文 Diffuse and Disperse: Image Generation with Representation Regularization 启发。该论文在我们刚刚完成嵌入塌缩的初步观察、并思考如何缓解现象之际，正好被上传到了arXiv。 我阅读后立刻意识到它与我们的研究方向高度相关，并借鉴了一些损失函数的设计。