提出无偏低秩梯度估计器，让大模型训练更稳定高效。
PLUMAGE: Probabilistic Low rank Unbiased Min Variance Gradient Estimator for Efficient Large Model Training
- 基于概率方法实现无偏低秩梯度估计，避免传统方法的偏差和方差问题。
- 在相似内存开销下，相比GaLoRE降低33%预训练损失、28%GLUE训练损失。
- 无需额外超参数，解决优化器状态错位问题，适合资源受限的大模型训练。
加速器内存与网络约束已成为训练数十亿参数大语言模型的主要瓶颈。现有低秩梯度估计方法如GaLoRE和FLORA通过将权重梯度投影到秩为r的子空间来压缩梯度和优化器张量，实现消费级硬件上的大模型训练。然而，这些方法存在偏差或高估计算法方差的问题。此外，当投影更新时，基于一阶与二阶矩估计的优化器状态会因子空间变化而失配，导致训练不稳定性。我们提出PLUMAGE：概率性低秩无偏最小方差梯度估计器。PLUMAGE可直接替换现有方法，仅需指定秩r与更新间隔，无需新增超参数。同时，解决了优化器状态失配问题，防止异常权重更新，提升训练稳定性。实验表明，相较于全秩优化，在多个模型上平均降低33%的预训练评估损失，且在相同计算与内存开销下，平均降低GLUE基准的训练损失28%。
原文摘要 · Abstract (English)
Accelerator memory and networking constraints have emerged as dominant bottlenecks when training large language models LLMs with billions of parameters. Existing low rank gradient estimators such as GaLoRE and FLORA compress gradients and optimizer tensors by projecting weight gradients onto a rank r subspace, enabling LLM training on consumer hardware. Yet, these methods are either biased or subject to high estimator variance. Moreover, the optimizer state based on the first and second moments estimates expressed in the previous subspace becomes misaligned whenever the projection is updated, leading to instabilities during training. We propose PLUMAGE: Probabilistic Low rank Unbiased Minimum vAriance Gradient Estimator. PLUMAGE is a drop in replacement for existing low rank gradient estimators. It does not introduce new hyperparameters beyond the chosen rank r and the update interval. In addition, we resolve optimizer state misalignment issues to prevent spurious weight updates and enhance training stability. We empirically demonstrate that PLUMAGE shrinks the full rank optimization's gap over the pre training evaluation loss by 33% on average across models and the average training loss across the GLUE benchmark by 28% within a similar computational and memory footprint as GaloRE.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。