🤖 machine learning

Riemannian Gradient Descent for Low-Rank Architectures

本文研究了针对深度学习中秩分解矩阵参数的十种算法设计的黎曼梯度下降，但发现尽管进行了调优，这些方法在应用于小型语言模型的多头注意力机制时，并未能得出明显优于 AdamW 基准线的结论。

原作者： Nicholas Knight

发布于 2026-06-02

📖 1 分钟阅读☕ 轻松阅读

原作者： Nicholas Knight

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

核心理念：在山峦间航行 vs. 开车行驶

想象你正试图在一个广阔且大雾弥漫的山谷中寻找最低点（这代表着训练一个犯错更少的机器学习模型）。通常，计算机科学家使用一种被称为 AdamW 的标准方法，这就像是在驾驶一辆拥有非常智能 GPS 的汽车。GPS 会告诉你“下坡”的方向，然后你向前迈进一步。如果你遇到了颠簸，汽车的悬挂系统（动量）会帮助你保持移动。这种方法效果很好，但它将引擎的每个部件都视为独立、互不相关的部分。

这篇论文提出了一个不同的问题：如果我们把引擎部件视为一个单一的、相互连接的系统呢？

在现代 AI 中，许多大脑的部分（模型）是使用“低秩”（low-rank）矩阵构建的。不要把这些看作是一个巨大的、实心的数据块，而要看作是一张折叠的纸。你可以通过仅仅知道折痕的坐标（两个较小的矩阵 $A$ 和 $B$ ）来描述整张纸的形状。

作者认为，标准方法将折痕视为相互独立的，这就像是在试图抚平一张皱巴巴的纸，却只顾着拉扯纸角，而没有意识到纸本身是连在一起的。因此，作者提议使用 黎曼梯度下降法 (Riemannian Gradient Descent)。

类比：

标准方法（欧几里得几何）： 想象你在平坦的地面上行走。你朝着 GPS 指引的方向迈出一步。如果你走在了一张皱巴巴的纸上，你可能会不小心踩到纸的边缘之外，导致你必须跳回纸面上。
黎曼方法： 想象你是一名在浪尖上冲浪的冲浪者。你被约束在波浪的表面上。你不仅仅是在行走，而是在波浪的曲面上滑行。这种方法确保你永远不会离开“低秩矩阵”的形状，使数学逻辑保持严谨且符合理论。

实验：测试新的冲浪板

作者基于这种“乘浪”的思想构建了十种不同的“冲浪板”（算法）。他在一个小型语言模型（一个学习预测句子下一个词的“大脑”）上进行了测试。

他测试了两种主要的波浪类型：

定秩波浪 (Fixed-Rank Waves)： 仅仅保持纸张以特定的紧凑度进行折叠。
部分等距波浪 (Partial Isometry Waves)： 一个更严格的规则，要求纸张必须是完美的正交状态（就像一个完美的直角），这就像是一种更僵硬、更有结构的波浪。

他还测试了一个“网格 (Grid)”版本，即模型的多个部分共享相同的折痕（就像一群人共同拉着同一根绳子）。

结果：一个“好”但不“卓越”的结局

以下是剥离了宣传辞令后的诚实结论：

它有效，但并非万能灵药： 新方法成功地训练了模型。它们没有崩溃，并且学会了任务。这证明了数学逻辑是成立的，代码也是可以运行的。
没有明显的胜利： 在仔细调整了新“冲浪板”的“速度”（学习率）后，它们并没有一致地击败标准的 AdamW 汽车。
- 在某些情况下，它们表现得稍好一些。
- 在另一些情况下，它们表现得稍差一些。
- 总的来说，这种差异微小到可以归因于随机噪声（就像微风改变了落叶的路径）。
代价： 新方法在计算上更加昂贵。这就像是在驾驶一辆需要特殊充电站的高科技电动车，而标准方法则是一辆可靠的燃油车，随处可用。由于新车并没有让你更快或更好地到达目的地，作者得出结论：它尚未准备好取代作为日常选择的标准方法。

给普通读者的核心要点

理论很优美： 将矩阵参数视为一个弯曲的曲面（流形）而非平坦网格的想法，在数学上是优雅且逻辑自洽的。它尊重了数字之间隐藏的关系。
实践很复杂： 虽然理论上说“你应该能更快到达”，但现实中的深度学习（由于其噪声、海量数据和复杂的架构）意味着传统的、更简单的方法（AdamW）目前仍然是王者。
未来的潜力： 作者保持乐观。仅仅因为一个新引擎在小赛道上没有赢得比赛，并不意味着它不会在巨大的高速公路上获胜。作者暗示，随着模型规模的扩大（扩展化），这些几何方法或许最终会展现出它们的真正威力。

本论文并未声称的内容

它没有声称这种方法能治愈疾病或解决气候变化。
它没有声称这就是 AI 的未来。
它没有声称该方法更快或更便宜（事实上，它承认该方法可能更慢、更复杂）。

总而言之： 作者构建了一种复杂的、几何化的方式来训练 AI 模型，这种方式尊重了数据的内部结构。它运行得非常完美，但就目前而言，传统的、简单的方法（AdamW）仍然是完成任务的最优选择。这种新方法是一个充满希望的“长线投资”，或许在未来我们构建更大规模的模型时，它会带来回报。

技术摘要：针对低秩架构的黎曼梯度下降法

问题陈述
在现代深度学习中，模型参数通常是受低秩约束的矩阵，以降低计算成本和参数量（例如在多头注意力机制中， $W = AB^T$ ）。标准的训练过程使用欧几里得梯度独立更新因子 $A$ 和 $B$ ，忽略了秩为 $r$ 的矩阵集的几何结构。本文研究了将秩约束矩阵 $W$ 视为光滑子流形上的一个点并应用黎曼优化，是否能产生更好的收敛性或解。核心挑战在于因子化 $W=AB^T$ 并不唯一（在 $A \to AS, B \to BS^{-T}$ 下具有不变性），且秩为 $r$ 的矩阵集合是一个流形而非向量空间。作者旨在仅通过修改优化器来限制改动，而不改变模型架构，同时保持与标准方法相当的单次迭代复杂度（ $O(mr^2 + nr^2)$ ）。

方法论
作者提出了一个为秩分解参数定制的黎曼梯度下降框架。通用的迭代过程包括：

从因子梯度 $G_A$ 和 $G_B$ 中恢复环境欧几里得梯度 $G_W$ 。
将此梯度投影到特定流形的切空间上，以获得黎曼梯度 $G_W^R$ 。
更新动量并归一化步长。
沿测地线（使用收缩映射/retraction）移动到下一个迭代点。
将动量向量进行平行移动至新的切空间。

该论文将此模板实例化为十种特定的算法设计，这些设计根据几何结构和约束进行了分类：

固定秩矩阵的几何结构 ( $R_{r}^{m \times n}$ ):
- 嵌入几何 (Embedded Geometry)： 使用环境 Frobenius 度量。通过对 $W + \Xi$ 进行秩- $r$ 截断 SVD 来执行收缩映射。平行移动使用度量投影。
- 商几何 (Quotient Geometry)： 将流形视为因子 $(A, B)$ 乘积空间的商。度量在“上方”的因子上定义，从而在商空间上诱导出一个度量。收缩映射在因子上是加性的；传输涉及水平投影。
秩- $r$ 部分等距矩阵的几何结构 ( $P_{r}^{m \times n}$ ):
- 此处， $W$ 被约束为其非零奇异值均为 1（等价于 $A$ 和 $B$ 是列正交的）。
- 嵌入几何： 限制 Frobenius 度量。收缩映射使用带有特定正交约束的截断 SVD。
- 商几何： 使用 Stiefel 流形之积 ( $St_{m,r} \times St_{n,r}$ ) 作为“上方”空间。收缩映射使用 QR 分解或极分解。
- 规范几何 (Canonical Geometry)： 使用 Stiefel 因子的规范度量，从而诱导出一个乘积度量。
权重共享（网格流形）：
- 将上述方法扩展到因子在块行和块列之间共享的架构（例如，分组查询注意力 GQA）。
- 对于固定秩，作者展示了通过将因子堆叠成更大的矩阵，可以使用单体算法，前提是需进行块内秩保持检查。
- 对于部分等距，简单的堆叠由于正交约束而失效。作者开发了更复杂的算法，包括一种针对嵌入几何情况的“混合”方法，即将其提升到商几何进行收缩映射和传输，然后再投影回原空间。

核心贡献

系统性探索： 本文对十种不同的黎曼优化变体进行了全面的综述和实现，涵盖了固定秩和部分等距约束，以及权重共享场景。
实现细节： 文中详细说明了每种几何结构所需的特定收缩映射（SVD、QR、Polar）和向量传输，解决了因子分解的不唯一性和保持流形约束（秩和正交性）的问题。
实验评估： 这些方法被应用于小规模语言模型（6 层，512 嵌入维度）中的 Q、K、V 矩阵，并使用了 FineWeb 数据集。

实验结果

性能： 在调整学习率后，黎曼方法并未得出超越 AdamW 基准线的结论。在多头注意力（MHA）和分组查询注意力（GQA）设置中，损失值的差异通常在随机种子的方差范围内。
超参数： 度量归一化提供了轻微的改进（0.01–0.05 nats）。重动量（ $\nu=7/8$ ）通常会对性能产生负面影响，或者相比基准线带来的收益微乎其微。
稳定性： 部分运行出现了损失值激增（例如在第 6660 步左右），但随后得以恢复，这通常与训练集中的结构化数据相关。
复杂度： 作者承认虽然这些方法在数学上是严谨的，但它们的单次迭代计算复杂度高于 AdamW，这使得它们在所进行的特定小规模实验中吸引力不足。

意义与主张
论文对其新颖性和直接影响保持了谦逊的态度：

重新发现： 作者明确指出，任何表现出的新颖性“几乎可以肯定是一种重新发现”，大量借鉴了已有的教科书（Absil 等，Boumal）和先前的文献。
启发式动机： 该工作受启发于这样一个假设：更好地考虑参数关系（通过流形几何）可能会带来更快的收敛或更好的解，尽管这一点在实验中并未得到证实。
未来潜力： 作者认为，在小规模模型中未能表现出优越性并不排除其在大规模场景下的价值。他们认为该方法在数学上是严谨的，且“成本-精度权衡”可能会随着更大规模的模型或不同的架构（如基于张量的模型：MLA、MoLAE）而发生有利的变化。
局限性： 研究承认存在显著局限性，包括模型规模较小、缺乏学习率退火，以及仅使用了一个未经过超参数调优的基准线（AdamW）。

总之，本文是对将黎曼优化应用于低秩深度学习组件的一次严谨的技术演示，它提供了一套算法库和基准，为未来的研究奠定了基础，同时也承认目前的实验结果尚未足以证明其可以取代 AdamW 等标准优化器。