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

会议: ICML 2026
arXiv: 2602.00217
代码: https://github.com/KrishnaswamyLab/LM-Dispersion
领域: 模型压缩 / 表征学习 / 小模型训练
关键词: 嵌入坍缩, 分散损失, 小模型泛化, 知识蒸馏, GPT2 / Qwen3

一句话总结¶

本文系统观测到 "小语言模型的 token 嵌入会随深度坍缩到一个窄锥体"（embedding condensation）这个普遍现象——大模型反而不会——并设计了一个角度分散损失 \(\mathcal{L}_{\text{disp}}\) 直接逼嵌入散开，无须加参数就让 Qwen3 / GPT2 在 10 个 benchmark 上平均提升 3.3%。

研究背景与动机¶

领域现状：LLM 的能力随规模 scaling 持续提升，但训练 / 部署成本飞涨，迫切需要 "用小模型复现大模型的关键性质"。已有压缩路线主要是蒸馏、量化、剪枝，都侧重模仿大模型的输出分布。

现有痛点：作者从 representation geometry 视角发现，小模型（GPT2-small、Qwen3-0.6B）的 token 嵌入在深层会几乎全部对齐到同一方向，pairwise cosine similarity 趋近 1；大模型（GPT2-xl、Qwen3-32B）则维持嵌入分散。Geshkovski 2025 的理论也证明 Transformer 在层数趋无穷时嵌入会坍缩成一个点，但经验上没人系统验证过和性能的关系。

核心矛盾：嵌入坍缩意味着模型实际可用的 "表征方向" 越来越少，表达能力被几何上锁死。即使蒸馏从大教师学到 logit 分布，也无法把大模型的几何性质继承下来——因为蒸馏目标只约束输出，不约束中间嵌入。

本文目标：(1) 把 embedding condensation 现象量化测量，确认 "大模型抗坍缩" 是普适规律；(2) 验证蒸馏不能缓解；(3) 设计直接作用于几何的辅助损失，让小模型主动散开嵌入。

切入角度：既然大模型 "自动" 维持分散，那分散本身可能就是性能的瓶颈条件。与其堆参数让模型 "自然" 分散，不如显式加一个目标函数强制分散。

核心 idea：用基于角度的分散损失 \(\mathcal{L}_{\text{disp}} = \log \sum_{i \neq j} \exp(-\arccos(\cos\text{sim}(z_i, z_j)) / \pi\tau)\) 把所有 token 嵌入推向单位超球面上的均匀分布，零额外参数。

方法详解¶

整体框架¶

整篇方法围绕"先把病诊断清楚、再开一剂直接作用于几何的药"展开。诊断阶段量化嵌入坍缩——用 Spearman \(\rho\) 和 Kendall \(\tau\) 衡量 layer-wise 平均余弦相似度是否随深度单调爬升，爬得越陡说明坍缩越重。干预阶段把一个角度分散损失 \(\mathcal{L}_{\text{disp}}\) 当正则项挂到原训练目标上：\(\mathcal{L} = \mathcal{L}_{\text{train}} + \lambda_{\text{disp}} \cdot \mathcal{L}_{\text{disp}}\)，零额外参数，既能 retrofit 现成 checkpoint 也能从零预训练时塑形。

关键设计¶

1. 嵌入坍缩的量化诊断：用层级余弦相似度的秩相关测出"病"有多重

要对症下药，先得把"病"测准。做法是：把输入序列喂进模型，逐层 \(l\) 收集该层所有 token 的嵌入 \(z_i^{(l)}\)，算出全部 \(N^2\) 对的 pairwise 余弦相似度，再用该层均值 \(\mu^{(l)} = \frac{1}{N^2}\sum_{i,j}\cos\text{sim}(z_i^{(l)}, z_j^{(l)})\) 概括这一层的对齐程度；最后计算序列 \(\{\mu^{(l)}\}_{l=1}^L\) 与层号 \(\{l\}_{l=1}^L\) 之间的秩相关——Spearman \(\rho\) 与 Kendall \(\tau\)。\(\rho/\tau\) 越接近 \(+1\)，说明相似度随深度单调爬升、坍缩越重；接近 \(0\) 表示没有系统趋势；为负则说明嵌入随深度反而散开（dispersion）。之所以用秩相关而不是"末几层的平均余弦"，是因为秩相关只看单调趋势、不受绝对尺度和非线性失真影响，趋势比直接取平均更干净、跨模型可比。正是这把尺子让作者量化出"模型越大、坍缩越轻"的普适规律（并在只改 MLP 维度的 confounder 控制实验里复现），为后面的干预提供了可测量的靶子。

2. 角度分散损失：把所有 token 嵌入往单位超球面的均匀分布上推

针对的痛点是小模型深层嵌入几乎全对齐到同一方向、pairwise cosine 趋近 1，可用的表征方向被几何上锁死。做法是对每层每个 token pair \((z_i, z_j)\) 先把余弦相似度映射成角度距离 \(D(z_i, z_j) = \arccos(\cos\text{sim}(z_i, z_j)) / \pi \in [0, 1]\)，再用 log-sum-exp 聚合成 \(\mathcal{L}_{\text{disp}} = \log \sum_{i \neq j} \exp(-D(z_i, z_j)/\tau)\)——两个 token 越同向、\(D\) 越小、\(\exp\) 项越大，损失就越重地把它们推开；接近正交时 \(\exp\) 项几乎为零、不再施力。所有层的损失求和，每 batch 复杂度 \(\mathcal{O}(N^2 F)\)，可在 token 维度子采样减负。几个细节都是为稳定服务：用 \(\arccos\) 而非裸 cosine 是避免在 \(\pm 1\) 两端梯度饱和；log-sum-exp 比直接取 mean 更稳健，差一个加性常数也不影响梯度；显式排除对角项 \(i=j\) 防止 \(z_i\) 自相似爆梯度；选 angular 而非欧式距离，是因为坍缩本质是方向坍塌而非长度问题。

3. 三个备选公式：用消融隔离"角度均匀分散"到底比别的分散方式好在哪

dispersion 是个抽象诉求，可以有很多实现，作者另造三种和主损失对打。Decorrelation 最小化嵌入协方差矩阵的非对角元，从特征维度间去耦合、间接逼分散；\(\ell_2\)-repel 直接拉大 token 间欧式距离，但必须配一个 norm 正则 \(\lambda_{\text{norm}} \|\mathcal{Z}\|_2^2\)，否则模型会靠膨胀 norm 而非真正散开来作弊；Orthogonalization 用铰链式损失 \(\max(0, 1/2 - D(z_i, z_j))^2\)，只惩罚 \(D < 1/2\) 的锐角对、放任钝角对自由生长。这一组对比是为了说明"在角度空间均匀分散"比"在特征维度去相关"或"在欧氏空间硬排斥"都更直接、更有效，从侧面坐实主方法的选择。

4. 同一损失覆盖 mid-training 与 full pre-training 两种流程

要证明这剂药既能给老模型续命也能从娘胎里塑形，就得在两种真实训练场景里都验证。Mid-training 拿现成的 GPT2 / Qwen3 在 wikitext-103 上再续训 200M token，单张 A100 就能跑完，是低成本的 proof-of-concept 和超参扫描场所；full pre-training 让 Qwen3 在 C4 上从零训 156B token、动用 640 GPU，验证这个几何信号从一开始就能塑造出更好的表征结构、从根本上改变模型可用容量。两种场景都只是把 \(\lambda_{\text{disp}} \cdot \mathcal{L}_{\text{disp}}\) 加到 cross-entropy 上，在每个 forward 同时计算多层嵌入的 dispersion 并加权，pipeline 改动极小。

损失函数 / 训练策略¶

最终训练目标 \(\mathcal{L} = \mathcal{L}_{\text{CE}} + \lambda_{\text{disp}} \cdot \mathcal{L}_{\text{disp}}\)，温度 \(\tau\) 与权重 \(\lambda_{\text{disp}}\) 是仅有的两个主要超参，扫描结果在 Appendix。Mid-training 用 3 个 seed 报均值方差，full pre-training 只跑单 seed 但 token 量足够大、结果已经稳定。

实验关键数据¶

主实验¶

GPT2 mid-training（10 个 benchmark 平均）：

配置	训练成本	平均分↑	rank↓	显著性
GPT2 原版（无 mid-training）	—	34.35	6.1	p<0.0001
+ \(\mathcal{L}_{\text{CE}}\) only	1.122 A100h	34.95	6.2	p<0.01
+ noisy embedding	1.122	35.15	4.3	p<0.01
+ active forgetting	1.127	35.36	3.2	n.s.
+ Dispersion loss	1.13 (1.01×)	35.52+	最佳	—

Qwen3 full pre-training（156B token from scratch）：加 dispersion loss 平均提升 +1.17 分（3.3% 相对增益），所有 benchmark 上稳定提升。

消融实验¶

四种 dispersion 变体对比：

变体	平均分	备注
Decorrelation	35.1	间接，受 feature dim 影响
\(\ell_2\)-repel	35.0	需 norm 正则才稳定
Orthogonalization	35.2	只惩罚锐角
Dispersion (canonical)	35.5+	角度均匀分散，最优

规模对照（confounder-controlled）：从零训四个 GPT2-like 模型，只改 MLP 维度其余固定，larger MLP → 更小 condensation，验证 "大模型抗坍缩" 不是其他因素带来的伪相关。

关键发现¶

蒸馏救不了坍缩：Qwen2.5 系列蒸馏后嵌入几何与从零训几乎一样，因为 KD loss 只约束输出 logits，不规范中间表征——这成为该论文最直观的动机证据。
坍缩从初始化就有，但训练能缓解：Olmo-3-7B 的 checkpoint 显示 condensation 指标在初始化时正且大，随训练降到负值，说明 SGD 本身就在抵抗坍缩，分散损失只是加速 + 强化。
小模型增益更大：Qwen3-0.6B 提升最明显，Qwen3-32B 几乎无收益，符合 "大模型本来就分散" 的假说。
mid-training 即可生效：不需要重训，对已有 checkpoint 加 200M token + dispersion 即获明显收益，落地成本极低。
代价 < 1% 训练时间：1.13 vs 1.122 A100h，因为 \(N^2\) 配对可子采样。

亮点与洞察¶

"小模型的瓶颈是几何不是参数"：把性能差距归因于表征几何而非容量本身，是该论文最有想象力的命题——意味着可以在不加参数的前提下逼近大模型性能上限。
角度而非欧氏距离：用 \(\arccos\) 把 cosine 映成均匀的角度距离，避开 cosine 在两端的饱和，是工程上很关键的稳定化 trick。
理论 → 经验 → 干预 → 验证闭环：从 Geshkovski 2025 的理论坍缩定理出发，提供大规模经验证据，设计干预，再用 confounder 控制实验证实，整篇文章的论证链非常完整。
零参数增量：与剪枝、量化、LoRA 等需要改架构的方法不同，dispersion loss 只是训练时附加项，可与任何主流 LM 训练 pipeline 即插即用。

局限与展望¶

未扩展到对齐 / 推理任务：实验全在 zero/few-shot 通用 NLU，未测 reasoning / math / code 等需要更复杂表征几何的任务。
大模型几乎无增益：32B 上没看到提升，分散假说是否对大模型成立未确认；也许大模型的瓶颈在别处。
token 子采样的代价：实际跑大模型时 \(N^2\) 还是要子采样，子采样如何影响收敛和最终性能没仔细消融。
未与 anti-collapse 技术（如 SimSiam stop-gradient、BarlowTwins）比较：自监督表征学习里有大量去坍缩方法，应横向比对。
未解释为何 \(\lambda_{\text{disp}}\) 不需要复杂调度：直觉上训练早期需要更强分散、后期可减弱，但论文用固定权重就 work，这是个有趣的现象但未深入分析。
架构敏感性：在 RMSNorm（无 affine LN）、不同位置编码下 dispersion 是否仍生效未测。

评分¶

新颖性: ⭐⭐⭐⭐ 把理论坍缩定理转化为可训练的辅助损失，并配套提出 4 种变体，属于 "已知现象 + 干净干预" 类新颖性，不是石破天惊但完整可信。
实验充分度: ⭐⭐⭐⭐⭐ mid-training + full pre-training（156B token / 640 GPU）双重验证，10 个 benchmark + 3 seed + confounder 控制 + 4 种 dispersion 变体消融，少见的扎实。
写作质量: ⭐⭐⭐⭐⭐ 论证链 "理论 → 经验观测 → 反例（蒸馏无效）→ 干预 → 验证" 非常顺，图表清晰，把抽象的几何现象讲得像故事一样。
价值: ⭐⭐⭐⭐ 几乎零成本即插即用，对小模型训练社区直接受益；只是大模型上无收益、未测推理 / 代码任务，长期影响力有上限。