分散损失抵消小型语言模型中的嵌入凝聚

# 分散损失 (LM-Dispersion) Source: https://chenliu-1996.github.io/projects/LM-Dispersion/ ## 分散损失抵消嵌入凝聚并提升小语言模型的泛化能力 \*同等贡献 **ICML 2026** 语言模型中嵌入凝聚现象的演示。 ## 一句话总结 是什么让大语言模型优于小语言模型？数据？参数？几何可能也扮演着关键角色！ ## 什么是嵌入凝聚？ 语言模型的每个Transformer层都将每个输入token表示为高维嵌入空间中的一个向量。我们注意到，随着这些向量通过Transformer层，它们往往表现得像是被限制在一个狭窄的锥体内：通过成对余弦相似度衡量，它们指向越来越相似的方向。我们将这种几何现象称为**嵌入凝聚**。该现象具有以下特征： 特征 1 在小模型中比在大模型中更为严重（图2）。 特征 2 在控制混杂因素的设置下可复现（图3）。 特征 3 在模型初始化时出现，并通过预训练得到缓解（图4）。 特征 4 通过从大模型进行知识蒸馏无法解决（图5）。 ## 关于本文的五分钟简介 **本文提出了一种基于观察的语言模型训练改进方法。** 我们观察到一种几何现象，我们将其称为**嵌入凝聚**，即在小语言模型中，token嵌入坍缩到一个狭窄的锥形子空间中。随后，我们设计了一种称为分散损失（dispersion loss）的训练目标来抵消这种效应。 图1. 嵌入凝聚现象示意图。在预训练语言模型中，同一输入序列的所有token嵌入在经过多个Transformer层处理后，会凝聚到一个狭窄的锥体内。在同一模型家族中，小模型的这一现象远较大模型更为明显，这引出了我们在第3.3节中的假设。**特征1：模型越大，凝聚越少。** 在同一模型家族中，小模型表现出更严重的嵌入凝聚，token嵌入向近平行方向坍缩，而大模型则抵抗这种坍缩。 图2. 嵌入凝聚现象的定性和定量观察。**a.** 余弦相似度热力图显示，小模型（如`GPT2`、`Qwen3-0.6B`）容易发生凝聚，因为随着嵌入进入更深层，token余弦相似度变得越来越正。相比之下，大模型（如`GPT2-xl`、`Qwen3-32B`）对嵌入凝聚更具抵抗力。**b.** 使用Spearman相关和Kendall's Tau的量化结果显示了多个语言模型家族中一致的“模型越大，凝聚越少”趋势。更多结果见图S1。这种效应对于输入数据集的选择也相当稳健。 图S2. 嵌入凝聚效应在不同输入文本数据集上保持一致。结果显示在四个数据集上：**(a)**`wikitext`, **(b)**`pubmed_qa`, **(c)**`imdb`, **(d)**`squad`。**特征2：在控制混杂因素时可复现。** 为了将模型大小的影响与其他混杂因素分离，我们进行了控制实验，预训练类似GPT2的模型，仅改变MLP维度，同时保持所有其他组件固定，包括层数、嵌入维度、数据集和训练设置。观察到相同的现象。 图3. 在高控制实验中，我们复现了“模型越大，凝聚越少”的观察结果。我们预训练了四个不同大小的`GPT2`类模型，它们仅在MLP维度上不同，同时保持所有其他因素固定，包括层数、嵌入维度、数据集和训练配置。所得模型在嵌入凝聚方面表现出一致的趋势，定性结果见面板**a**，定量结果见面板**b**。面板**a**中添加了水平虚线以便于视觉比较。 **特征3：凝聚早期出现。** 嵌入凝聚现象在模型初始化时就已出现，并随着预训练逐渐缓解，而非加剧。 图4. 嵌入凝聚在模型初始化后立即被观察到。我们分析了`Olmo-3-1025-7B`的检查点，涵盖初始化、中间预训练阶段和最终基础模型。每个检查点标注了其训练阶段和训练token数量。 **特征4：蒸馏不是解决方案。** 从大模型进行知识蒸馏不会传递所需的嵌入凝聚抵抗性。 图5. 知识蒸馏并非嵌入凝聚的补救措施，定性结果见面板**a**，定量结果见面板**b**。**分散损失** 嵌入凝聚通过将token嵌入向量坍缩到狭窄锥体，降低了Transformer的表达能力，未充分利用表示空间。我们假设，通过在训练中分散嵌入，小模型可以获得更接近大模型的表示质量，从而缩小性能差距，而无需增加参数数量。 图6. 分散损失及其替代公式如何促进嵌入分散的示意图。**a.** 分散损失通过沿单位超球面散布所有向量对来实现均匀的角度分散。**b.** 去相关损失鼓励不同特征维度保持不相关。**c.** l2-排斥损失增加成对欧氏距离，同时范数正则化防止无界扩张。**d.** 正交化损失散布形成锐角的向量，而钝角保持不变。 我们的分散损失受到"Diffuse and Disperse (https://arxiv.org/abs/2506.09027)"论文的启发，并进行了实际修改。 表1. 我们的分散损失及其替代公式。与Diffuse and Disperse (https://arxiv.org/abs/2506.09027)的主要实现差异用teal和magenta突出显示。包含或排除对角线项产生相同的梯度，因此仅为外观差异。对于分散损失和l2-排斥，我们采用`log-sum-exp`技巧以保证数值稳定性，这与`log(mean(exp(·)))`仅相差一个加性常数。对于l2-排斥，我们包含一个范数正则化项以防止嵌入的无界扩张。对于正交化，距离边际固定为1/2，因为我们使用角度距离，其中1/2对应正交性，因此作为理想边际。 分散损失在中期训练和预训练期间抵消了嵌入凝聚效应。下方展示了定性结果，更多定量结果可在论文中找到。 图7. 分散损失抵消嵌入凝聚现象。**a.** 从凝聚的嵌入开始（灰色虚线框），使用默认损失进行中期训练效果有限（绿色框）。**b.** 相比之下，使用我们的分散损失作为正则化器进行中期训练显著缓解了嵌入凝聚（蓝色框）。 **结论** 大语言模型优于小语言模型，但可能不仅仅是因为它们有更多参数。这可以部分归因于它们如何在潜在表示中组织信息。我们期待未来在这一有趣方向上的努力。 ## 免责声明 如果你考虑复现这项工作或借用其中部分内容，以下是我的两点建议。 - **嵌入凝聚。** 这部分在我们的手中表现良好。在许多模型家族、输入数据集以及受控设置下都能一致观察到。趋势可能因模型家族而异，但我们关键观察中所包含的内容都不需要经过挑选的运行。我们不能保证所有模型家族都存在这一现象，但你当然可以在你喜欢的模型上尝试。 - **分散损失。** 请将这部分视为探索性内容。收益有限：改进非常微妙，没有正式的统计检验很难与噪声区分（而我们所做的统计检验非常基础，因为我们不是好的统计学家）。一位在数学推理方面更有经验的朋友在论文被接收后评论说，我们的中期训练方案不太标准，因为常见的做法是增强特定领域能力，而不是在wikitext上继续训练。我们的预训练实验也比较单薄，因为大规模运行成本高昂。如果你对这种方法感兴趣，我建议先在你的团队常规流程中在小规模上尝试，然后再投入大的训练预算。 ## 未来方向 我个人强调几个看似有意义的潜在方向。 - **更好的正则化器。** 分散损失是一个非常简单的直接解决方案，可能有优点也有缺点。一个更精心设计的抵消嵌入凝聚的方法可能会更有帮助。 - **超越预训练。** 追踪嵌入凝聚在训练后期阶段（如监督微调SFT和强化学习RL）的演变。目前尚不清楚凝聚是否会重新出现、稳定下来，还是与对齐目标产生不同的交互。 - **机制与因果性。** 确定嵌入凝聚的根本原因，并在凝聚与泛化等下游行为之间建立更强的因果联系。 - **更好的架构。** 设计天生抵抗凝聚的模型家族或模块，补充或替代基于损失的正则化。 - **更好的初始化。** 开发使模型从较不凝聚的状态开始的初始化方案，可能减轻训练目标抵消几何坍缩的负担。 ## 引用 `` @inproceedings{liu2026dispersion, title={Dispersion loss counteracts embedding condensation and improves generalization in small language models}, author={Liu, Chen and Sun, Xingzhi and Xiao, Xi and Van Tassel, Alexandre and Xu, Ke and Reimann, Kristof and Liao, Danqi and Gerstein, Mark and Wang, Tianyang and Wang, Xiao and Krishnaswamy, Smita}, booktitle={International Conference on Machine Learning}, year={2026}, organization={PMLR} } `` ## 致谢 1. 这项工作的最初动机来自论文“A mathematical perspective on Transformers (https://arxiv.org/abs/2312.10794)”。我们在2025年4月初观看了关于该论文的讲座 (https://www.youtube.com/watch?v=3McmEtA3t_0) 后开始了这个项目。我们对论文中的一个理论结果很感兴趣：如果我们无限堆叠Transformer层，所有token嵌入将聚集到同一点，我们好奇这种行为是否能在经验上观察到。这引出了我们论文中嵌入凝聚的关键观察。 2. 分散损失的设计很大程度上受到Runqian (https://raywang4.github.io/)和Kaiming (https://scholar.google.com/citations?user=DhtAFkwAAAAJ)的论文“Diffuse and Disperse: Image Generation with Representation Regularization (https://arxiv.org/abs/2506.09027)”的启发。他们的论文在我们完成嵌入凝聚的初步观察并思考缓解方法时正好出现。当我读到那篇论文时，我立即感到它高度相关，可以用作一个合理的解决方案。