MoARa: Module-Aware Rank Allocation and Structure-Preserving Decomposition for Low-Rank LLM Pre-training

📄 arXiv: 2609.15037v1 📥 PDF

作者: Keunyoung Kim, Nojun Kwak

分类: cs.LG, cs.CL

发布日期: 2026-09-14

备注: Accepted to EMNLP 2026 Main Conference. 30 pages, 16 figures

💡 一句话要点

提出MoARa以优化低秩LLM预训练效率

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

关键词: 低秩预训练 大规模语言模型 梯度投影 模块感知 效率优化 深度学习

📋 核心要点

  1. 现有低秩梯度投影方法在优化器状态内存成本上有所改善，但在达到目标质量的步骤和时间上仍存在显著不足。
  2. 论文提出的MoARa方法通过模块感知的投影秩分配和块级幅度-方向分解，优化了梯度投影的效率和效果。
  3. 实验结果显示，MoARa在多个Transformer架构上显著减少了训练步骤和时间，提升了预训练效率。

📝 摘要（中文）

低秩梯度投影可以减少大规模语言模型（LLM）预训练中的优化器状态内存成本，但达到目标质量所需的步骤和实际时间仍有提升空间。现有方法的不足在于，投影秩预算在不同Transformer模块间均匀分配，且原始梯度的投影同时削弱了其幅度和方向。本文提出MoARa，结合基于静态分析的模块感知投影秩分配与块级幅度-方向分解，默认块大小设置在注意力头维度附近。在涵盖Llama、Qwen和DeepSeek等五种Transformer架构的实验中，GaLore结合MoARa在Llama 2 7B上以37%的步骤和34%的时间减少达到了标准GaLore的最终困惑度，且在标准图编译下仅有0.2%的峰值保留内存开销。

🔬 方法详解

问题定义：本文旨在解决现有低秩梯度投影方法在大规模语言模型预训练中效率不足的问题。现有方法在投影秩预算分配上存在均匀化的问题，导致不同模块的投影敏感性未得到合理利用，同时原始梯度的投影削弱了其幅度和方向。

核心思路：论文提出的MoARa方法通过静态分析实现模块感知的投影秩分配，针对不同模块的特性进行优化。同时，采用块级幅度-方向分解技术，分别处理梯度的幅度和方向，以提高投影效果。

技术框架：MoARa的整体架构包括两个主要模块：模块感知投影秩分配和块级幅度-方向分解。首先，通过静态分析评估各模块的投影敏感性，动态分配投影秩；然后，针对每个块进行幅度和方向的独立处理，优化梯度投影。

关键创新：MoARa的核心创新在于模块感知的投影秩分配与块级幅度-方向分解的结合。这一设计使得不同模块的特性得以充分利用，显著提升了预训练的效率。与现有方法相比，MoARa在处理梯度时不再简单地进行均匀投影，而是根据模块的需求进行优化。

关键设计：在MoARa中，默认块大小设置在注意力头维度附近，以便更好地适应Transformer架构的特性。此外，论文中还对投影秩的动态分配进行了详细的参数设置，以确保在不同模型和任务中均能获得最佳效果。通过这些设计，MoARa在内存开销和训练效率上实现了显著的优化。

🖼️ 关键图片

fig_0
fig_1

📊 实验亮点

实验结果表明，结合MoARa的GaLore在Llama 2 7B上以37%的步骤和34%的时间减少达到了标准GaLore的最终困惑度。在六种低秩预训练方法中，模块感知的秩分配在所有方法上均实现了方向性一致的步骤减少，最高可达41.7%的步骤减少和37.1%的时间减少。

🎯 应用场景

该研究的潜在应用领域包括大规模语言模型的预训练和优化，尤其是在资源受限的环境中。通过提高训练效率，MoARa能够加速模型开发周期，降低计算成本，促进更广泛的AI应用。此外，未来可能在其他深度学习任务中借鉴该方法，进一步提升模型训练的效率和效果。

📄 摘要（原文）

Low-rank gradient projection reduces the optimizer-state memory cost of large language model (LLM) pretraining, but the steps and wall-clock time needed to reach a target quality remain a meaningful axis for improvement. We attribute this to two design choices in existing methods: the projection-rank budget is allocated uniformly across Transformer modules with heterogeneous projection sensitivity, and projecting a raw gradient attenuates its magnitude and direction jointly. We propose MoARa, which combines a static profiling-based module-aware projection-rank allocation with a block-wise magnitude-direction decomposition; the default block size is set in the neighborhood of the attention head dimension. Across five Transformer architectures spanning Llama, Qwen, and DeepSeek at 300M to 7B scales, GaLore with MoARa reaches standard GaLore's final perplexity in 37% fewer steps and 34% less wall-clock time on Llama 2 7B, with only 0.2% peak reserved memory overhead under standard graph compilation. Across the six low-rank pretraining methods we evaluate, module-aware rank allocation alone delivers directionally consistent step reductions on all six. On compatible hosts, the two-component design reaches up to 41.7% step reduction and 37.1% wall-clock reduction.