Clean 用 Nyström 草图把二阶优化器内存降到线性，单张 80GB GPU 可预训练 13B 模型
相关研究与后续进展
核心概要
作者提出 Clean 与低精度变体 Q-Clean：用随机 Nyström 方法近似 SOAP 的左右预条件子，把优化器状态内存从随层维度二次增长降为线性，并把低秩核心之外的 Main、A、B、C 四个正交块重新整合进更新；在 C4 上预训练 LLaMA-350M 与 1.3B 时，Clean 困惑度接近 SOAP 而优化器内存小于 AdamW，Q-Clean 相比 Muon 减少超过 50% 优化器内存，Clean 达到 AdamW 最终性能的墙钟时间快 26%，并首次使 13B 模型可在单张 80GB H100 上预训练。
Figure 1: Measured memory consumption for pre-training a LLaMA-13B model on a single 80 GB H100 GPU (maximum sequence length of 1024, batch size of 1). All non-quantized matrices are stored in bfloat16. Activation + Other combines activation memory with an estimated optimizer workspace. Methods marked ∗ ran out of memory, so their optimizer state is estimated.
arXiv深度剖析
Clean 用随机 Nyström 近似替代 SOAP 中显式累积的完整梯度协方差，只保存薄因子，使预条件子内存随层维度线性而非二次增长。 此前 SOAP、Shampoo 及多种高效优化器至少有一个优化器状态随模型维度二次增长；Clean 是作者所述首个以线性内存完成 SOAP 式双侧预条件的二阶优化器。 论文给出内存复杂度对比表，并在 LLaMA-350M 与 1.3B 预训练中报告优化器内存：Clean 为 1.25G 与 4.48G，SOAP 为 7.64G 与 13.62G，AdamW 为 1.40G 与 5.11G。
Clean 不丢弃低秩核心之外的梯度，而是把更新拆成 Main、A、B、C 四个正交块，对 Main 用完整 Adam，对 A、B 用池化二阶矩，对 C 只用原始梯度与标量二阶矩，再投影回参数空间重组为满秩更新。 低秩投影类方法通常只保留核心子空间，Clean 显式重新整合被丢弃的残差块，并用投影感知的矩搬运处理基变化。 作者在 LLaMA-350M 上统计各块占梯度范数的比例，显示 Main 在几乎所有层占主导而 C 贡献较小，但非 Main 块在训练后期仍携带可观信号，据此设计分块矩状态。
Q-Clean 对优化器状态做低精度量化：一阶矩与投影矩阵用 INT8 分块非对称线性量化并配合随机舍入，主二阶矩用 bitsandbytes 非线性量化，A、B 的二阶矩保持全精度。 该方案把 Clean 的内存优势进一步放大，同时用随机舍入避免小更新在舍入中被抹掉。 在 LLaMA-1.3B 预训练中，Q-Clean 优化器内存为 1.28G，相比 Muon 的 2.80G 减少超过 50%，困惑度为 12.10，与全秩方法相当。
线性内存使 Clean 与 Q-Clean 成为所评估基线中唯一能在单张 80GB H100 上预训练 13B 参数模型的优化器，且 Clean 达到 AdamW 最终性能的墙钟时间快 26%。 这为在更易获得的硬件上训练更大基础模型提供了可行路径，而其他低秩二阶基线如 ALSO 在相同设置下仍超过 80GB。 论文报告 Clean 用 64 小时达到 AdamW 的 81 小时最终性能，并只消耗 77% 的 token 预算，Q-Clean 为 87%；13B 单卡结论来自附录 D 的内存分析。
启示与展望
该结果面向以 LLaMA 架构在 C4 上从零预训练的场景，覆盖 350M 与 1.3B 参数、序列长度 1024、有效批大小 4096，并按 Chinchilla 规则训练；13B 单卡结论来自附录 D 的内存分析而非完整预训练实验。方法适用于矩阵权重层，1-D 参数仍用 Adam 更新。对希望降低优化器状态内存、同时保留双侧曲率信息的研究者与工程团队，这提供了可复用的分块矩设计与量化配方。
1.3B 的学习率只做了缩短扫描并只续跑最佳配置，更大规模下的超参数敏感性仍需观察；分块矩的池化与基变化时的近似搬运带来多少误差，正文以图 3 与附录说明为主，未给出逐块误差上界；13B 单卡结论基于内存分析，尚无该规模的完整预训练困惑度；量化在更长训练或不同架构上的稳定性仍是开放问题。