Demystifying Manifold Constraints in LLM Pre-training

📄 arXiv: 2605.04418v1 📥 PDF

作者: Kang An, Jiaxiang Li, Donald Goldfarb, Shiqian Ma

分类: cs.LG, cs.AI, math.OC

发布日期: 2026-05-06

💡 一句话要点

提出MACRO优化器，揭示流形约束在LLM预训练中的作用，提升稳定性和性能。

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 大型语言模型 预训练 流形约束 黎曼优化 优化器

📋 核心要点

  1. LLM预训练依赖启发式稳定技术，但约束权重的具体机制和动机尚不明确。
  2. 提出MACRO优化器，通过流形约束独立地限制激活尺度和稳定旋转平衡。
  3. 实验表明，MACRO在大型LLM架构上实现了极具竞争力的性能，并保证了理论收敛性。

📝 摘要（中文）

大型语言模型（LLM）预训练的成功在很大程度上依赖于启发式稳定技术，如显式归一化层和权重衰减。虽然最近显式约束权重的约束优化方法可能提高数值稳定性和性能，但添加约束的机制和动机仍然难以捉摸。本文系统地揭示了显式流形约束在LLM预训练中的作用。通过引入Msign对齐约束黎曼优化器（MACRO）——一个可证明收敛的单循环优化框架——我们的研究将权重正则化启发式方法与RMS归一化和解耦权重衰减等相互作用机制分离开来。理论分析和全面的经验评估表明，流形约束独立地限制了前向激活尺度并强制执行稳定的旋转平衡，从而涵盖了这些启发式机制的作用。对大型LLM架构的评估表明，MACRO在严格保持精确黎曼优化的理论保证的同时，实现了极具竞争力的性能。

🔬 方法详解

问题定义：大型语言模型预训练过程中，数值不稳定性和性能优化是一个关键问题。现有的方法，如显式归一化层和权重衰减，虽然有效，但缺乏理论支撑，且其内在机制尚不明确。因此，需要一种更系统和理论化的方法来解决LLM预训练中的稳定性和性能问题。

核心思路：论文的核心思路是通过引入流形约束来显式地控制模型权重，从而实现数值稳定性和性能提升。流形约束能够限制前向激活尺度，并强制执行稳定的旋转平衡，从而替代或补充现有的启发式稳定技术。通过将权重限制在特定的流形上，可以避免权重爆炸或梯度消失等问题，从而提高训练的稳定性和效率。

技术框架：论文提出了Msign-Aligned Constrained Riemannian Optimizer (MACRO)，这是一个单循环优化框架，它结合了流形约束和黎曼优化。该框架首先将权重限制在特定的流形上，然后使用黎曼优化算法在该流形上进行优化。MACRO框架包括以下几个主要步骤：1) 定义流形约束；2) 计算黎曼梯度；3) 使用Msign对齐策略更新权重；4) 投影权重到流形上。

关键创新：论文的关键创新在于将流形约束引入到LLM预训练中，并提出了MACRO优化器。MACRO优化器不仅具有理论上的收敛性保证，而且在实践中也表现出良好的性能。此外，论文还揭示了流形约束在LLM预训练中的作用，即独立地限制前向激活尺度和强制执行稳定的旋转平衡。

关键设计：MACRO优化器的关键设计包括：1) 流形约束的选择：论文中使用了球形流形约束，即限制权重的范数为1。2) Msign对齐策略：该策略用于确保权重更新方向与黎曼梯度方向一致，从而提高收敛速度。3) 投影操作：在每次权重更新后，需要将权重投影回流形上，以满足流形约束。论文中使用了简单的投影操作，即将权重除以其范数。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，MACRO优化器在大型LLM架构上实现了极具竞争力的性能，并且在训练过程中表现出更好的稳定性。具体来说，MACRO在保持理论收敛性的前提下，能够达到与现有启发式方法相当甚至更好的性能，证明了流形约束在LLM预训练中的有效性。

🎯 应用场景

该研究成果可应用于各种大型语言模型的预训练，尤其是在模型规模较大、训练数据量较多的情况下，能够有效提高训练的稳定性和效率，并最终提升模型的性能。此外，该研究对于理解LLM预训练的内在机制具有重要的理论价值，可以指导未来LLM架构的设计和优化。

📄 摘要（原文）

The empirical success of large language model (LLM) pre-training relies heavily on heuristic stabilization techniques, such as explicit normalization layers and weight decay. While recent constrained optimization approaches that explicitly restrict weights may improve numerical stability and performance, the mechanism and motivation for adding constraints still remain elusive. This paper systematically demystifies the role of explicit manifold constraints in LLM pre-training. By introducing the Msign-Aligned Constrained Riemannian Optimizer (MACRO)-a provably convergent, single-loop optimization framework-our study disentangles weight regularization heuristics from interacting mechanisms like RMS normalization and decoupled weight decay. Theoretical analyses and comprehensive empirical evaluations reveal that manifold constraints independently bound forward activation scales and enforce stable rotational equilibrium, thereby subsuming the roles of these heuristic mechanisms. Evaluations on large-scale LLM architectures demonstrate that MACRO achieves highly competitive performance while rigorously preserving the theoretical guarantees of exact Riemannian optimization.