[Paper] YuriiFormer：一套 Nesterov 加速的 Transformers

Source: arXiv - 2601.23236v1

（未提供需要翻译的文本。）

概述

论文 “YuriiFormer: A Suite of Nesterov‑Accelerated Transformers” 将 transformer 层重新构想为经典优化算法的迭代步骤。通过将 self‑attention 视为 interaction energy 的梯度、将 MLP 视为 potential energy 的梯度，作者展示了标准的 GPT‑style 模型本质上是在对一个复合目标执行 vanilla gradient descent。基于这一洞见，他们设计了一种 Nesterov‑accelerated 变体，保持相同的 attention 和 MLP 构件，但加入了动量项，从而在语言建模基准上实现了可观的性能提升。

关键贡献

• 变分重新解释变压器 – 将每一层形式化为在 token 嵌入上进行的优化例程的迭代。
• 基于能量的分解 – 将模型的计算拆分为交互能（由自注意力处理）和势能（由前馈 MLP 处理）。
• Lie–Trotter 分裂视角 – 表明常见的交替注意力‑MLP 模式对应于用于最小化组合能量的 Lie–Trotter（算子分裂）方案。
• Nesterov 加速的变压器设计 – 引入基于动量的更新，同时保留原始注意力/MLP “oracle”。
• 实证验证 – 加速架构（YuriiFormer）在 TinyStories 和 OpenWebText 上始终优于强大的 nanoGPT 基线，尽管参数数量相当。

方法论

1. 能量表述 – 作者定义了一个标量目标 ( \mathcal{L}(X) = \mathcal{E}{\text{int}}(X) + \mathcal{E}{\text{pot}}(X) )，其中 (X) 为 token 嵌入。

   • 交互能 (\mathcal{E}_{\text{int}}) 捕捉 token 之间的成对关系；它的梯度正是自注意力计算的内容。
   • 势能 (\mathcal{E}_{\text{pot}}) 编码每个 token 的变换；它的梯度对应 MLP 前馈块。

2. 算子分裂 – 通过应用 Lie–Trotter 分裂，一个 transformer 层可以写成：

   [ X^{(k+1/2)} = X^{(k)} - \eta \nabla \mathcal{E}_{\text{int}}(X^{(k)}) \quad\text{（注意力步骤）} ]

   [ X^{(k+1)} = X^{(k+1/2)} - \eta \nabla \mathcal{E}_{\text{pot}}(X^{(k+1/2)}) \quad\text{（MLP 步骤）} ]

   这正好对应 GPT 块的前向传播。

3. Nesterov 加速 – 作者在上述基础上加入动量项：

   [ Y^{(k)} = X^{(k)} + \beta_k (X^{(k)} - X^{(k-1)}) ]

   然后在 (Y^{(k)}) 而不是 (X^{(k)}) 上计算注意力和 MLP 的梯度。系数 (\beta_k) 按经典的 Nesterov 调度进行，保证在凸问题中的加速收敛。

4. 实现 – 不需要新的算子；直接复用已有的注意力和 MLP 模块。唯一的额外开销是保存前一次隐藏状态并进行一次轻量的线性组合。

5. 训练设置 – 实验使用 nanoGPT 代码库，在两个语料上训练约 10 M 参数的模型：TinyStories（合成短篇小说）和 10 M token 的 OpenWebText 切片。超参数（学习率、批大小等）在基线和加速实验中保持完全相同，以隔离 Nesterov 步的影响。

结果与发现

• 训练速度：额外的动量计算每步仅增加 < 2 % 的开销，在现代 GPU 上几乎可以忽略不计。
• 稳定性：加速模型在更少的 epoch（约 15 % 更少的更新次数）内收敛，同时保持相当的梯度范数，表明优化过程更平滑。
• 泛化能力：收益在两个截然不同的语料库上均有体现，暗示该方法并非特定于某一数据集。

实际意义

• Plug‑and‑play 加速 – 由于 YuriiFormer 复用了已有的 attention/MLP 核心，开发者只需添加几行代码（存储前一隐藏状态，加入动量混合），即可升级现有的 transformer 代码库。
• 性价比高的性能 – 对于在边缘设备上常用的中小规模模型（10‑100 M 参数），Nesterov 步骤能够在不增加模型体积的情况下显著提升性能，从而在相同硬件预算下获得更好的下游任务表现。
• 训练效率 – 更快的收敛意味着更少的 GPU 时长，这对计算资源有限的创业公司和研究团队具有吸引力。
• 设计框架 – 变分视角为其他受优化算法启发的改动打开了大门（例如 Adam‑style 预条件、适应性步长），这些都可以作为架构层面的“oracle”实现，而无需重新设计整个模型。
• 可解释性 – 将层解释为梯度步骤提供了更透明的思维模型，便于调试训练动态，可能有助于自动化架构搜索工具。

限制与未来工作

• 凸性假设 – 理论加速保证仅适用于凸目标函数，而 Transformer 的训练高度非凸；观察到的收益是经验性的，且该方法在非常大规模模型（≥ 1 B 参数）上可能无法直接扩展，需要进一步调优。
• 动量调度 – 论文使用了标准的 Nesterov 调度；自适应或学习得到的动量可能带来更大的提升，但未在本文中进行探索。
• 更广泛的基准 – 实验主要聚焦于语言建模；将该方法应用于视觉 Transformer、多模态模型或指令微调的大语言模型仍是一个待解的问题。
• 消融深度 – 虽然作者已将动量项单独剥离，但更深入的消融实验（例如，改变拆分顺序、与其他优化器技巧结合）可以进一步阐明哪些组件驱动了性能提升。

Bottom line: YuriiFormer 表明，借鉴经典的优化技巧——此处的 Nesterov 加速——可以以低成本、高影响的方式从现有 Transformer 架构中挤出额外的性能。对于希望在不进行大幅重构的情况下提升模型效率的开发者，本文提供了一个具体、可直接实现的方案，并为未来的架构创新提供了新的视角。

作者

• Aleksandr Zimin
• Yury Polyanskiy
• Philippe Rigollet

论文信息

• arXiv ID: 2601.23236v1
• 分类: cs.LG, cs.AI, math.OC, stat.ML
• 发表时间: 2026年1月30日
• PDF: 下载 PDF