arXiv:2502.05790cs.LG2025-02NeurIPS被引 6

提出重要性采样方法,突破大模型训练中低秩优化的固定子空间瓶颈。

Breaking the Frozen Subspace: Importance Sampling for Low-Rank Optimization in LLM Pretraining

  • 通过重要性采样动态选择低秩梯度更新方向,避免子空间僵化。
  • 在LLM预训练中显著提升性能,收敛性有理论保证。
  • 适合追求高效低显存大模型训练的研究者与工程师。

低秩优化已成为实现大语言模型(LLMs)内存高效训练的有前景方法。现有方法通常将梯度投影到低秩子空间,降低优化器状态存储开销。核心挑战在于选择合适子空间以确保有效优化轨迹。多数方法选择主成分子空间以保留梯度信息,这看似最优。然而我们发现,实践中主成分子空间在预训练过程中基本不变,导致权重更新受限于相似子空间。本文提出一种用于LLM预训练的低秩优化重要性采样方法,具有可证明的收敛性保证,而主成分方法不具备此性质。实验表明,该方法在大语言模型预训练任务中显著优于先前方法。

原文摘要 · Abstract (English)

Low-rank optimization has emerged as a promising approach to enabling memory-efficient training of large language models (LLMs). Existing low-rank optimization methods typically project gradients onto a low-rank subspace, reducing the memory cost of storing optimizer states. A key challenge in these methods is selecting suitable subspaces to ensure an effective optimization trajectory. Most existing approaches select the dominant subspace to preserve gradient information, as this intuitively provides the best approximation. However, we find that in practice, the dominant subspace stops changing during pretraining, thereby constraining weight updates to similar subspaces. In this paper, we propose importance sampling for low-rank optimization in LLM pretraining with a provable convergence guarantee, which the dominant subspace approach does not have. Empirically, we demonstrate that our method significantly outperforms previous methods in LLM pretraining tasks.

低秩优化大模型训练重要性采样

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。