Better Embeddings with Coupled Adam¶

会议: ACL 2025
arXiv: 2502.08441
代码: 无
领域: LLM效率 / 优化器
关键词: 词嵌入各向异性, Adam优化器, Coupled Adam, 表示退化, LLM训练

一句话总结¶

从理论上证明 Adam 优化器的逐 token 二阶矩是导致 LLM 词嵌入各向异性（均值偏移）的根因，提出 Coupled Adam——对嵌入层的二阶矩取词汇平均——消除了各向异性问题，并在大规模实验中提升了嵌入质量和下游性能。

领域现状：LLM 训练出的词嵌入表现出各向异性——所有嵌入聚集在远离原点的小型子空间中，限制了语义表达力。这是一个被广泛观察但根因不明的现象。
现有痛点：Biś et al. (2021) 发现各向异性主要由嵌入均值偏移导致，但将原因归于"common enemy effect"（非目标 token 的梯度都指向 \(-h\) 方向），这个解释不够精确——它忽略了梯度被预测概率缩放的因素。
核心矛盾：如果嵌入梯度之和恒为零（论文证明了这一点），那为什么均值还会偏移？到底是梯度本身还是优化算法的问题？
本文要解决什么？ 找到嵌入各向异性的真正根因，并提供一个简单有效的修复方案。
切入角度：数学分析嵌入更新向量在 SGD 和 Adam 下的求和性质——SGD 下更新之和恒为零（均值不变），Adam 下由于逐 token 的二阶矩不同导致更新之和非零。
核心 idea 一句话：Adam 的二阶矩为每个 token 引入了不同的有效学习率，破坏了梯度求和恒零的性质，导致均值偏移；将二阶矩在词汇维度取平均（Coupled Adam）即可修复。

理论分析：证明 SGD 下嵌入均值不变 vs Adam 下嵌入均值会偏移
方法提出：Coupled Adam——仅对嵌入层修改，将二阶矩 \(\hat{v}_i\) 替换为词汇级平均 \(\hat{\nu} = \frac{1}{V}\sum_i \hat{v}_i\)
实验验证：小规模（125M-760M）+ 大规模（1.3B-2.6B）

理论推导——梯度之和恒零:
做什么：证明对于标准语言模型头，所有嵌入梯度之和恒为零 \(\sum_i g_i = 0\)
核心思路：由 softmax 概率之和 \(\sum_i p_i = 1\) 直接推出。这意味着如果优化器的更新向量与梯度成正比，均值不会变化
设计动机：建立了理论基础——问题不在梯度，而在优化器如何处理梯度
SGD vs Adam 的关键差异:
SGD 更新：\(u_i = -\eta \cdot g_i\)，更新之和 \(= -\eta \sum g_i = 0\) → 均值不变
Adam 更新：\(u_i = -\eta_i \cdot \hat{m}_i\)，其中 \(\eta_i = \eta / (\sqrt{\hat{v}_i} + \epsilon)\) 是 token 依赖的有效学习率。由于 \(\eta_i\) 对不同 token 不同，加权求和 \(\sum \eta_i \hat{m}_i \neq 0\) → 均值偏移
关键发现：\(\mathbb{E}[\hat{v}_i] \propto \tilde{p}_i\)（二阶矩期望正比于 unigram 概率），实验验证 \(R^2=0.85\)
Coupled Adam（核心创新）:
做什么：仅对嵌入参数，将 Adam 的二阶矩替换为词汇级平均
核心思路：\(\hat{\nu}^{(\tau)} = \frac{1}{V}\sum_{i=1}^V \hat{v}_i^{(\tau)}\)，然后 \(\hat{v}_i^{(\tau)} \leftarrow \hat{\nu}^{(\tau)}\) 对所有 \(i\)
效果：有效学习率变为 token 无关的，更新之和恢复为零，均值保持不变
实现：仅需 2 行代码修改，对非嵌入参数不做任何改变

数据/模型	方法	Test Loss	Iso↑	\(\\|\mu\\|^r\)↓	词相似 \(\bar{r}\)↑
20B/125M	Standard	3.03	0.10	0.82	5
20B/125M	Coupled	2.97	0.83	0.03	57
20B/355M	Standard	2.79	0.25	0.82	5
20B/355M	Coupled	2.75	0.95	0.02	57
20B/760M	Standard	2.68	0.28	0.73	3
20B/760M	Coupled	2.65	0.94	0.01	58

模型/数据	方法	Loss↓	Acc↑	Iso↑
1.3B/26B	Standard	2.46	0.397	0.22
1.3B/26B	Coupled	2.44	0.399	0.94
2.6B/210B	Standard	2.17	0.425	0.22
2.6B/210B	Coupled	2.16	0.428	0.97

Coupled Adam 将各向同性从 0.1-0.28 提升到 0.83-0.97，均值偏移比从 0.63-0.82 降到 0.01-0.03——几乎完美修复
词语义相似度评测（SimLex999 等）从个位数提升到 55-58——嵌入质量质的飞跃
在足够大的数据集上（20B+），Coupled Adam 同时改善了 perplexity 和下游准确率
在小数据集上效果中性或略差——说明各向同性带来的收益需要足够训练量才能体现
2.6B 模型 210B tokens 仍持续受益，说明方法在大规模场景下有效