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

背景

近年来，大型语言模型（LLM）在各种任务中表现出色，而小型语言模型（small LM）的性能往往存在明显差距。人们通常将这种差距归因于参数数量或数据规模，但一项来自2026年ICML的研究指出，几何特性可能也扮演着关键角色。该研究发现，在Transformer层中，词元（token）的嵌入向量会逐渐坍缩到一个狭窄的锥形子空间中，这种现象被称为「嵌入凝聚」（embedding condensation），且在小模型中更为严重。研究团队设计了一种名为「分散损失」（dispersion loss）的训练目标，旨在对抗这种几何坍缩，从而在不增加参数的情况下缩小模型之间的性能差距。

核心内容

语言模型的每一层Transformer都将每个输入词元表示为一个高维嵌入空间中的向量。当这些向量逐层传递时，它们的行为类似于被限制在一个狭窄的锥体内：两两之间的余弦相似度越来越高，即方向趋于一致。研究者将这一几何现象称为嵌入凝聚。

该现象具有以下四个特征（原文中的Feature 1–4）：

1. 更大模型，凝聚更少。在同一模型家族内部，小型模型表现出更严重的嵌入凝聚，词元嵌入几乎变成平行方向；而大型模型则能抵抗这种坍缩。该效应在不同输入数据集上均稳健存在。
2. 控制混淆变量后仍可复现。为了排除模型大小的混淆效应，研究者进行了一项受控实验：预训练一系列GPT2风格的模型，仅改变MLP维度，其他所有组件（层数、嵌入维度、数据集、训练设置）均保持一致。结果观察到同样的现象。
3. 凝聚在早期出现。嵌入凝聚在模型初始化时就已存在，并且预训练过程会逐渐缓解它，而非加剧。
4. 知识蒸馏并非解决方案。从更大的模型进行知识蒸馏，并不能将抵抗嵌入凝聚的能力传递给小模型。

基于上述观察，研究者提出了分散损失（dispersion loss）作为训练目标，其灵感来自《Diffuse and Disperse》论文，并做了实用修改。该损失的目标是在训练过程中分散嵌入向量，使小模型获得更接近大模型的表征质量，从而缩小性能差距，且不增加参数数量。分散损失在预训练和中期训练（mid-training）中均可有效对抗嵌入凝聚效应。论文提供了定性结果，更多定量结果见原文。

结论指出：大型语言模型优于小型语言模型，可能不仅仅是因为参数更多，部分原因在于它们如何在潜在表征中组织信息。研究者希望未来能沿着这一有趣方向继续探索。

此外，作者在论文接受后给出了一些实用建议：

• 嵌入凝聚：这一现象在多种模型家族和输入数据集上均可一致观察到，且无需刻意挑选运行结果。不能保证所有模型家族都存在该现象，但读者可以尝试。
• 分散损失：这一部分更具探索性。带来的提升较微妙，需要正式统计检验才能与噪声区分。此外，论文的中期训练配方可能不够标准，预训练实验也因成本限制而较单薄。建议在投入大规模预算前，先在团队正常的小规模流程中试用。

作者还列出了几个有潜力的研究方向：

• 更好的正则化：分散损失只是简单直接的方案，可能有优缺点，需要更精心设计的方法。
• 超越预训练：跟踪嵌入凝聚在后续阶段（如SFT、RL）中的演化，看是否重新出现、稳定或与对齐目标交互。
• 机制与因果：确定嵌入凝聚的根本原因，并建立与下游行为（如泛化）的因果联系。
• 更好的架构：设计天生抵抗凝聚的模型家族或模块，以补充或替代基于损失的正则化。
• 更好的初始化：开发使模型从更少凝聚状态开始的初始化方案，减少训练目标对抗几何坍缩的负担。

该工作的最初动机源于论文《A mathematical perspective on Transformers》中的一个理论结果：如果无限堆叠Transformer层，所有词元嵌入将聚集到同一点。研究人员于2025年4月初观看了该论文的演讲后，开始好奇这一行为是否能在经验中观察到，从而导致了关键观察。分散损失的设计则主要受《Diffuse and Disperse: Image Generation with Representation Regularization》一文的启发。

关键要点

• 嵌入凝聚是小型语言模型中普遍存在的几何现象：词元嵌入在高维空间中向狭窄锥体方向坍缩，余弦相似度增大。
• 该现象在更小的模型中更严重，在大模型中较轻，且对输入数据集不敏感。
• 通过控制MLP维度（其他因素不变）的受控实验，确认了模型大小是导致该现象的原因。
• 凝聚在模型初始化时即出现，预训练过程会缓解而非加剧这一现象。
• 从大模型进行知识蒸馏并不能将抵抗凝聚的能力传递给小模型。
• 分散损失是一种针对性的训练目标，灵感来自图像生成领域的正则化方法，可在中期训练和预训练中有效对抗嵌入凝聚。
• 分散损失带来的性能提升相对微妙，需要统计检验才能确认，且其实践配方可能未达到标准，建议在小规模实验中先验证。
• 该工作揭示了大型语言模型的优势可能部分源于其几何表征的疏密组织，而非仅参数数量。
• 未来方向包括设计更好的正则化、探索后续训练阶段的影响、建立因果机制、开发固有抗凝聚架构以及改进初始化方法。

意义与影响

这项研究从一个新颖的几何视角解释了大型语言模型为何优于小型语言模型，挑战了「纯粹参数决定论」。它指出了潜在表征空间的组织方式（即嵌入是否坍缩）对模型表达能力的关键作用。分散损失作为一种轻量级、无额外参数的训练约束，为提升小模型性能提供了一条可行路径，尤其适合计算资源受限的场景。此外，该工作为后续正则化方法、初始化方案以及架构设计提供了明确的目标：抵抗嵌入凝聚。未来若能建立嵌入凝聚与泛化性能之间的更强因果联系，将有助于指导更高效的模型训练策略，甚至可能推动对Transformer内部表征理论的更深入理解。研究团队的开源实验细节和坦诚的局限性说明（如提升的微妙性、配方非标准）也为学界提供了宝贵的可复现性参考。