Clean：通过 Nystr\"om Sketching 以线性内存成本进行二阶 LLM 训练

arXiv:2610.04204v2 公告类型：替换 摘要：训练大型语言模型（LLM）涉及一个根本性权衡：Adam 等内存高效的优化器会丢弃参数间的曲率信息，而 SOAP 等完整曲率方法虽然能够加速收敛，却需要高得令人难以承受的内存成本。我们提出了 Clean，这是一种内存高效且利用完整曲率的优化器，旨在解决这一瓶颈。Clean 利用随机化 Nystr\"om 方法，准确近似 SOAP 中的左右预条件器，并将优化器的内存复杂度从关于模型维度的二次复杂度降至线性复杂度。随后，我们重新整合子空间外分量，以捕获低秩近似之外的曲率信息，在极低内存成本下保留丰富的曲率信息。我们进一步提出了 Q-Clean，这是一种积极压缩优化器状态的低精度变体。在预训练 LLaMA-1.3B 架构时，与 Muon 相比，Q-Clean 将优化器内存消耗降低了 \textbf{50\% 以上}，同时保持了强劲且具有竞争力的预测性能。值得注意的是，与标准 AdamW 相比，Clean 所需的优化器状态存储空间更小，同时仅用 \textbf{快 26\%} 的实际运行时间即可达到 AdamW 的最终性能。此外，我们的方法还独特地支持在单张 80GB GPU 上预训练一个 13B 参数模型，为大规模模型优化提供了一种可扩展、高效且易于使用的方案。