令牌几何

# 标记几何学
来源：https://arxiv.org/html/2607.01455
###### 摘要

语言模型在离散符号上学习连续程序，嵌入表和LM头作为两者之间的读/写接口。我们表明，该接口具有与密集隐藏权重不同的梯度几何，可以利用此几何来改进监督微调、强化学习和预训练的帕累托前沿，同时仅需数千字节的优化器状态。我们引入了*Ember*，一种针对嵌入和LM头矩阵的轻量级优化器，它使用O\(V\+D\)\\mathcal\{O\}\(V\+D\)显存，而不是Adam的O\(2VD\)\\mathcal\{O\}\(2VD\)，并且无需分片两个令牌表的优化器状态。我们提供了经验证据，表明Ember在批量大小和参数数量上都能有效扩展。我们表明，令牌的优化轨迹可以用简单的1D射线很好地描述，这与流行的观念（即神经网络参数在高度非凸的景观中导航）相反。我们为Transformer训练所需的优化器空间（意外地狭窄）提供了原则性的观点。最后，我们开源了我们的分布式Ember实现，该实现可以干净地集成到现有的ZeRO/FSDP设置中，以支持进一步的研究 (github.com/katop1234/ember (https://github.com/katop1234/ember))。

![[无标题图片]](https://arxiv.org/html/2607.01455v1/results/headline_figure.png)

图1：Ember优化器几乎不使用嵌入表和LM头的优化器状态，并且普遍匹配或改进了当前最先进技术。在稀疏梯度/低批量情况下（图1B），Ember大幅优于Adam；在微调和强化学习中（图1C、1D），它在与Adam的种内平行时达到了持平，同时在大多数记录的检查点表现更好。
## 1 引言

嵌入表和LM头矩阵，即“令牌接口”参数，目前使用Adam[1 (https://arxiv.org/html/2607.01455#bib.bib1)]进行优化。这与当前开源前沿报告（DeepSeek V4[14 (https://arxiv.org/html/2607.01455#bib.bib14)]）相符，该报告仅将线性层的优化器改为Muon[15 (https://arxiv.org/html/2607.01455#bib.bib15)]。值得注意的是，在文献中，几乎所有提出的Adam后继者都仅适用于线性层，并避开嵌入表。因此，改进是相关的，因为Adam需要大量的分布式训练工程，并且没有充分利用令牌接口的几何结构。在本文中，我们介绍*Ember*，一种针对此类Transformer参数的轻量级优化器，它解决了这些问题并取得了有竞争力的结果。我们从信息几何的一些初步知识开始：

费希尔度量。在正则性假设下，交叉熵损失具有一个显著性质：给定参数的二阶导数仅是其导数平方，

−Epθ\[∂2∂θi2log⁡pθ\(x\)\]=Epθ\[\(∂∂θilog⁡pθ\(x\)\)2\]\.\-\\mathbb\{E\}\_\{p\_\{\\theta\}\}\\left\[\\frac\{\\partial^\{2\}\}\{\\partial\\theta\_\{i\}^\{2\}\}\\log p\_\{\\theta\}\(x\)\\right\]=\\mathbb\{E\}\_\{p\_\{\\theta\}\}\\left\[\\left\(\\frac\{\\partial\}\{\\partial\\theta\_\{i\}\\}\\log p\_\{\\theta\}\(x\)\\right\)^\{2\}\\right\]\.这也被称为费希尔信息度量[8 (https://arxiv.org/html/2607.01455#bib.bib8),9 (https://arxiv.org/html/2607.01455#bib.bib9),10 (https://arxiv.org/html/2607.01455#bib.bib10)]，在最优值附近它与交叉熵损失的Hessian矩阵重合。虽然标准的优化基线是牛顿法，它使用二阶导数的逆来调节每个梯度，但标准的反向传播不能廉价地返回二阶导数。幸运的是，费希尔度量（或二阶导数）可以仅使用上述一阶导数的信息来廉价估计。

##### 交叉熵的规范几何。

费希尔度量特别优雅，因为它不仅作为近似有用，而且是Kullback-Leibler散度的精确二阶项：

DKL\(pθ∥pθ\+dθ\)=12dθ⊤F\(θ\)dθ\+O\(‖dθ‖3\)\.D\_\{\\mathrm\{KL\}\}\\left\(p\_\{\\theta\}\\;\\\|\\;p\_\{\\theta\+d\\theta\}\\right\)=\\frac\{1\}\{2\}d\\theta^\{\\top\}F\(\\theta\)d\\theta\+\\mathcal\{O\}\(\\\|d\\theta\\\|^\{3\}\)\.根据Chentsov定理[6 (https://arxiv.org/html/2607.01455#bib.bib6)]，对于KL散度，除了费希尔度量之外不存在其他规范局部几何。因此，对于交叉熵损失，超越SGD的优化必须利用该量；这已被确立为“自然梯度下降”[7 (https://arxiv.org/html/2607.01455#bib.bib7),2 (https://arxiv.org/html/2607.01455#bib.bib2),11 (https://arxiv.org/html/2607.01455#bib.bib11),12 (https://arxiv.org/html/2607.01455#bib.bib12)]。

有几个注意事项，很容易补救。首先，费希尔度量作为二阶校正，假设了精确的二次损失景观。其次，它在批量梯度下降下数值不稳定，因为在任何给定的训练步骤中，我们只观察到经验令牌频率p^i\\hat\{p\}\_\{i\}，而不是精确的pip\_\{i\}。

### 1.1 相关工作

SGD[16 (https://arxiv.org/html/2607.01455#bib.bib16)]是最简单的梯度下降形式，使用反向传播梯度乘以学习率来更新参数。Adam[1 (https://arxiv.org/html/2607.01455#bib.bib1)]在此基础上通过除以梯度RMS进行构建，并维护一阶和二阶矩的偏差校正EMA，每个令牌表占用O\(2VD\)\\mathcal\{O\}\(2VD\)内存。

Lion[17 (https://arxiv.org/html/2607.01455#bib.bib17)]表明仅对梯度符号使用动量具有竞争力。Muon[15 (https://arxiv.org/html/2607.01455#bib.bib15)]强制每个密集线性层的梯度矩阵的奇异值成为单位值，并被前沿模型使用。Adafactor[18 (https://arxiv.org/html/2607.01455#bib.bib18)]使用了与Ember类似的形式，它完全丢弃了Adam的一阶矩缓冲区，并对二阶矩进行分解，尽管有4个额外配置。我们注意到我们独立地收敛到了这个设置，并在附录6.5 (https://arxiv.org/html/2607.01455#S6.SS5)中进行了阐述，相信这进一步印证了我们的方法。

[19 (https://arxiv.org/html/2607.01455#bib.bib19)]声称线性层从根本上是对称的，因此一个好的优化器应该利用这一点，而不是像对待1D连接向量那样更新其参数。Modula[4 (https://arxiv.org/html/2607.01455#bib.bib4),13 (https://arxiv.org/html/2607.01455#bib.bib13)]是最早提出优化器统一理论的文献之一，该理论数学上建立在范数对偶性基础上。他们推导出行归一化对嵌入是最优的，这与本文的分析一致。

#### 1.1.1 平方根费希尔度量

在高层次上，我们认为这些优化器在SGD之上共享相同的基本原理。每篇论文使用了不同的术语，但我们发现最简洁的解释是“将每个梯度转换为z分数”（数学上，使用费希尔度量的平方根进行调节）。回想一下，费希尔度量只是平方梯度，当模型接近最优值时，它等于方差。因此，除以其平方根类似于除以梯度估计量的标准差。

有一种观点认为，Adam通过其分母捕捉到了这一度量，对于给定的参数梯度gg，它以\(g2\)−1/2\(g^\{2\}\)^\{\-1/2\}的量级进行调节。奇怪的是，可以证明Muon优化器与此类似，因为它用\(GTG\)−1/2\(G^\{T\}G\)^\{\-1/2\}来调节矩阵梯度GG。不难看出，这正好是除以G平方的矩阵类比量的平方根。我们声称，这种简单形式构成了有竞争力的Transformer优化器的集合。

在标准优化中，简单地用曲率的逆来调节梯度是一个众所周知的基线[21 (https://arxiv.org/html/2607.01455#bib.bib21)]。引人注目的是，我们认为现代深度学习只做了这一点，并进行了适度调整，例如应用平方根而不是完整的曲率校正，因为不能假设精确的二次景观，并且我们在每一步只观察到一个小批量。

令人惊讶的是，该领域的文献分散在各种线性代数推导或经验发现中，并没有以费希尔度量为主导，尽管它是一个中心数学对象，具有易于操作的特性。虽然这些工作似乎在传达相同的基本概念，但使用围绕费希尔度量的统一词汇（它被规范地证明是交叉熵损失的二阶项）可以加速该领域的研究。

## 2 方法

### 2.1 推导

我们从观察开始：Adam尽管结构相对简单，但在广泛的深度学习任务中表现极好。我们注意到Adam近似于每个参数的平方根经验费希尔度量，因为梯度的平方近似于曲率，这是本文中将使用的假设，因为Adam在经验上对所有参数类都表现良好。值得注意的是，由于二阶导数在技术上是对所有参数对进行平均的，我们注意到二阶矩EMA项在训练过程中会自动调整这一点，因为纠缠参数会接收到总体上较低的梯度。

首先，我们发现完全移除一阶矩EMA并用瞬时梯度替换它有利于性能（图11 (https://arxiv.org/html/2607.01455#S6.F11)），并在输入嵌入表和LM头中节省了V×DV\\times D的优化器状态。

我们观察到，行向梯度能量是令牌参与概率的无偏估计量，我们将其表示为pip\_\{i\}。直观上，这是因为反向传播在每次选择令牌时都会累积梯度，因此频率更高的令牌其梯度l2\\ell\_\{2\}范数与频率成正比。因此，由于行向费希尔度量与pi2p\_\{i\}^\{2\}成比例（附录6.8 (https://arxiv.org/html/2607.01455#S6.SS8)），曲率可以通过VV个行向平方梯度EMA而不是标准的V×DV\\times D缓冲区来很好地近似。根据标准优化器文献，我们应用偏差校正，并取该曲率度量的平方根，这具有一个很好的性质，即净梯度实际上是每一行中的z分数。我们的第一个主要发现是，这种切换节省了大量内存（≈D×\\approx D\\times）状态，并且与Adam持平（图2 (https://arxiv.org/html/2607.01455#S2.F2)）。

我们继续这一点，观察到这种方法虽然高效，但通过假设特征各向同性，为令牌中的每个参数分配相同的曲率，而这一点通常不成立。因此，为了校正这一点，我们使用一个廉价的O\(D\)\\mathcal\{O\}\(D\)列向因子（它在转置后与上述相同）的外积来估计逐元素平方梯度（图12 (https://arxiv.org/html/2607.01455#S6.F12)，附录6.6 (https://arxiv.org/html/2607.01455#S6.SS6)），并观察到它在基准测试中超越或匹配Adam的性能。由于词汇量很大，根据大数定律，列因子集中在单位值附近，因此其贡献很小；因此，在本文中，我们将分析重点放在Ember的仅行版本上。

包含列因子会使分母增加一个额外的gg单位，因此我们推导出，减少方差（6.7 (https://arxiv.org/html/2607.01455#S6.SS7)）并将单位与平方根费希尔度量匹配的数学上最优方式是将外积除以行和列缓冲区几何平均值（算法1 (https://arxiv.org/html/2607.01455#alg1)）。

本质上，我们丢弃了Adam的一阶矩，并将其二阶矩缓冲区替换为1D行向和列向因子的廉价外积（算法1 (https://arxiv.org/html/2607.01455#alg1)），该外积逐点复制了密集二阶矩（图12 (https://arxiv.org/html/2607.01455#S6.F12)），并发现它在替换嵌入表中的Adam时提供了最先进的结果。唯一的超参数是β2=0.999\\beta\_\{2\}=0.999，我们直接从原始Adam论文中借用。

算法1 Ember，我们提出的针对嵌入和LM头矩阵的优化器。Ember维护行和列的二阶矩估计，并形成轻量级的外积预处理器。本文中使用的默认设置为α=10−3\\alpha=10^\{\-3\}和β2=0.999\\beta\_\{2\}=0.999。1:

α\\alpha: 学习率

2:

β2∈\[0,1\)\\beta\_\{2\}\\in\[0,1\): 二阶矩估计的动量项

3:

Rt\(θ\)R\_\{t\}\(\\theta\): 时间步

tt的奖励

4:

θ0∈RV×D\\theta\_\{0\}\\in\\mathbb\{R\}^\{V\\times D\}: 嵌入或LM头矩阵的初始参数

5:

r0←0∈RVr\_\{0\}\\leftarrow\\mathbf\{0\}\\in\\mathbb\{R\}^\{V\}⊳\\triangleright初始化行二阶矩向量

6:

c0←0∈RDc\_\{0\}\\leftarrow\\mathbf\{0\}\\in\\mathbb\{R\}^\{D\}⊳\\triangleright初始化列二阶矩向量

7:

t←0t\\leftarrow 0⊳\\triangleright初始化时间步

8:while

θt\\theta\_\{t\}未收敛do

9:

t←t\+1t\\leftarrow t\+1
10:

gt←∇θRt\(θt−1\)g\_\{t\}\\leftarrow\\nabla\_\{\\theta\}R\_\{t\}\(\\theta\_\{t\-1\}\)⊳\\triangleright通过反向传播计算梯度

11:

rt←β2rt−1\+\(1−β2\)meanj⁡\(gt2\)r\_\{t\}\\leftarrow\\beta\_\{2\}r\_\{t\-1\}\+\(1\-\\beta\_\{2\}\)\\operatorname\{mean\}\_\{j\}\(g\_\{t\}^\{2\}\)⊳\\triangleright更新行二阶矩

12:

ct←β2ct−1\+\(1−β2\)meani⁡\(gt2\)c\_\{t\}\\leftarrow\\beta\_\{2\}c\_\{t\-1\}\+\(1\-\\beta\_\{2\}\)\\operatorname\{mean\}\_\{i\}\(g\_\{t\}^\{2\}\)⊳\\triangleright更新列二阶矩

13:

r^t←rt/\(1−β2t\)\\hat\{r\}\_\{t\}\\leftarrow r\_\{t\}/\(1\-\\beta\_\{2\}^\{t\}\)⊳\\triangleright偏差校正行估计

14:

c^t←ct/\(1−β2t\)\\hat\{c\}\_\{t\}\\leftarrow c\_\{t\}/\(1\-\\beta\_\{2\}^\{t\}\)⊳\\triangleright偏差校正列估计

15:

st←r^t ̄c^t ̄s\_\{t\}\\leftarrow\\sqrt\{\\overline\{\\hat\{r\}\_\{t\}\}\\,\\overline\{\\hat\{c\}\_\{t\}\}\}⊳\\triangleright几何均值归一化

16:

v~t←r^tc^t⊤/st\\tilde\{v\}\_\{t\}\\leftarrow\\hat\{r\}\_\{t\}\\hat\{c\}\_\{t\}^\{\\top\}/s\_\{t\}⊳\\triangleright形成分解预处理器

17:

θt←θt−1\+αgt/\(v~t\+10−8\)\\theta\_\{t\}\\leftarrow\\theta\_\{t\-1\}\+\\alpha\\,g\_\{t\}/\(\\sqrt\{\\tilde\{v\}\_\{t\}\}\+10^\{\-8\}\)⊳\\triangleright上升奖励RtR\_\{t\}

18:endwhile

19:return

θt\\theta\_\{t\}

图2：GPT-2 small / FineWeb在两种批量大小下的嵌入优化器消融实验。有趣的是，仅行版本达到了与规范Ember和Adam几乎持平的效果，表明大部分曲率来自参与概率，并且在批量较大时，这可以充分估计。
### 2.2 优化稀疏激活的参数

每个令牌行以不同的速率参与训练。令pip\_\{i\}表示令牌ii在训练分布下激活的频率，gig\_\{i\}其净梯度，Fi=Ez\[∇θilog⁡pθ\(z\)∇θilog⁡pθ\(z\)⊤\]F\_\{i\}=\\mathbb\{E\}\_\{z\}\[\\nabla\_\{\\theta\_\{i\}\\}\\log p\_\{\\theta\}\(z\)\\nabla\_\{\\theta\_\{i\}\\}\\log p\_\{\\theta\}\(z\)^\{\\top\}\]令牌费希尔度量。关于用pip\_\{i\}缩放令牌梯度存在几种数学上有效的观点，每种选择对应于嵌入表上的不同度量。

按pp缩放。一个基本的观点是让更新随参与度增长，因此gi∝pig\_\{i\}\\propto p\_\{i\}，因为直观上能量应该与我们对其确信程度成比例。然而，这可能导致死学习，因为令牌通常根据幂律被选择，并且会偏向于学习n-gram统计量[3 (https://arxiv.org/html/2607.01455#bib.bib3)]。