DASH：通过批量化分块预条件子和高效逆根求解器加速 Shampoo 优化器 DASH: Faster Shampoo via Batched Block Preconditioning and Efficient Inverse-Root Solvers
通过分块堆叠与牛顿迭代改进，将 Shampoo 优化器速度提升 5.6 倍且不损失精度
前置知识
Shampoo 优化器
Shampoo 是一种近似二阶优化方法，由 Gupta 等人在 2018 年提出。它通过维护每个网络层的左预条件矩阵 $L_t \in \mathbb{R}^{m \times m}$ 和右预条件矩阵 $R_t \in \mathbb{R}^{n \times n}$ 来捕获参数间的相关性。具体来说，左预条件矩阵通过梯度外积 $G_t G_t^\top$ 的指数移动平均（EMA）累积，右预条件矩阵通过 $G_t^\top G_t$ 累积。参数更新时使用 $\theta_{t+1} = \theta_t - \eta_t \cdot L_t^{-1/4} G_t R_t^{-1/4}$，将全矩阵自适应方法的 $O(m^2 n^2)$ 复杂度降至 $O(m^2 + n^2)$。它在 MLCommons AlgoPerf 竞赛中获胜，并被证明能产生激活离群值更少、更容易量化的模型。
DASH 的全部工作就是加速 Shampoo 的内部计算，不理解 Shampoo 的基本机制就无法理解本文的优化目标和方法设计
矩阵逆根（Matrix Inverse Root）
给定对称正定矩阵 $A$ 和正整数 $p$，矩阵逆根 $A^{-1/p}$ 是满足 $X^p A = I$ 的矩阵 $X$。在 Shampoo 中，需要对预条件矩阵 $L_t$ 和 $R_t$ 分别计算 $L_t^{-1/4}$ 和 $R_t^{-1/4}$。标准方法是对 $A$ 做特征值分解（EVD）$A = Q \Lambda Q^\top$，然后 $A^{-1/p} = Q \Lambda^{-1/p} Q^\top$。但 EVD 的时间复杂度为 $\Theta(n^3)$ 且难以在 GPU 上并行化，是 Shampoo 的主要性能瓶颈。
逆根计算是 Shampoo 最昂贵的操作，本文提出的所有新迭代方法（NDB、CBSHV）都是为了替代 EVD 来加速这一步
Coupled-Newton（CN）迭代
CN 迭代是 Higham 在 2008 年提出的一种基于矩阵乘法的迭代方法，用于计算 $A^{-1/p}$。它维护两个矩阵序列 $X_k$ 和 $M_k$，其中 $X_k \to A^{-1/p}$ 且 $M_k \to I_n$。每次迭代需要 3 次矩阵乘法（$p=2$）或 4 次矩阵乘法（$p=4$），相比 EVD 的 $\Theta(n^3)$ 复杂度，CN 可以利用 GPU 的 Tensor Core 高效执行矩阵乘法。该方法被集成到 Distributed Shampoo 中作为 EVD 的替代方案，但默认未启用。
CN 是本文对比的基线方法，DASH 中引入的 Newton-DB 迭代正是要改进 CN 的收敛速度和数值稳定性
分布式分块预条件子（Distributed Block Preconditioner）
Distributed Shampoo（Shi et al., 2023）将预条件矩阵分割为 $B \times B$ 的块，将逆根计算的复杂度从 $\Theta(n^3)$ 降至 $\Theta(B^3)$（对每个块而言）。然而原始实现中这些块是按顺序逐个处理的，导致 GPU 利用率低下。分块还涉及负载均衡算法，将不同层分配给不同 GPU 处理，处理完后广播同步参数。
DASH 的核心架构创新就是将这些分块堆叠成 3D 张量进行批量并行处理，理解原始分块机制是理解 DASH 加速原理的前提
Frobenius 范数与谱范数
Frobenius 范数 $\|A\|_F = \sqrt{\sum_{i,j} a_{ij}^2}$ 是矩阵元素的平方和开根号，计算简单但通常是最大特征值 $\lambda_{\max}(A)$ 的宽松上界（在实际中可大 $10-100$ 倍）。谱范数 $\|A\|_2 = \lambda_{\max}(A)$ 是矩阵的算子范数。迭代方法计算逆根时需要先将矩阵缩放到合适的范围（如 $[0, 1]$），缩放方式的选择直接影响收敛速度。
本文发现用 Frobenius 范数缩放会导致小特征值被过度压缩，从而需要更多迭代步数，这是本文提出 multi-Power-Iteration 的理论动机
研究动机
Shampoo 是目前表现最好的近似二阶优化器之一，在 AlgoPerf 竞赛中以最短的墙钟时间达到目标训练性能获胜，且已被证明能产生泛化能力更好、对量化更鲁棒的模型。然而，使用 Shampoo 的代价是每个优化器步的计算开销巨大。具体来说，Shampoo 需要计算逆矩阵根 $A^{-1/4}$，这个操作对于 $n \times n$ 的预条件矩阵通常需要 $\Theta(n^3)$ 的时间复杂度。虽然 Distributed Shampoo 将预条件矩阵分割为 $B \times B$ 的块以降低单块复杂度到 $\Theta(B^3)$，但底层算法仍然存在严重的效率瓶颈：默认实现依赖 EVD（特征值分解），这是一个在 GPU 上难以并行化的操作，因为它需要构建 Krylov 子空间、迭代重正交化和三对角化。即使引入了基于矩阵乘法的 Coupled-Newton 迭代作为替代，由于对数值稳定性的顾虑，默认并未启用。更关键的是，现有实现按顺序逐个处理分块，无法充分利用 GPU 的大规模并行能力。Vyas 等人（2024）的研究表明，更频繁地更新预条件子直接带来更好的优化性能，但这在现有实现下意味着更高的计算代价。
本文的目标是本文的目标是显著降低 Shampoo 的计算开销，同时保持其数值精度，使高质量的二阶预条件优化在大规模模型训练中变得实用。具体而言，作者希望通过算法和系统的协同设计，将优化器步的运行时间降低到与 Diagonal 方法（如 AdamW）可接受的差距内。从实验结果来看，一个可量化的目标是在 Llama-953M 模型上将优化器步时间从 Distributed Shampoo 的 675ms（CN-FP32，$B=2048$）降低到 150ms 以下，同时保持甚至改善验证困惑度。
与已有工作不同的是，本文的独特切入角度在于从「系统-算法协同设计」的视角重新审视 Shampoo 的性能瓶颈。已有工作如 SOAP、4-bit Shampoo 等主要从算法层面（如特征基变换、量化）来降低开销，但忽视了一个关键观察：Distributed Shampoo 的分块存储结构导致了内存碎片化和顺序计算，这是系统层面的效率损失。DASH 抓住的核心被忽视的点是：同层和跨层的预条件矩阵块具有相同的形状（$B \times B$），可以堆叠成 3D 张量进行批量处理，从而充分利用 Tensor Core 的高吞吐量。此外，本文还首次系统分析了矩阵缩放方式对迭代收敛的影响，发现 Distributed Shampoo 使用的 Frobenius 范数缩放远非最优，提出了基于 Power-Iteration 的更精确缩放方法。
核心方法
DASH 的设计思路可以用一个类比来理解：想象你有一堆同样大小的乐高积木（$B \times B$ 的预条件矩阵块），原来的 Distributed Shampoo 是一块一块地按照说明书（EVD/CN）拼装，而 DASH 的做法是把这些积木放在传送带上，用机器人手臂（Tensor Core）一次性批量处理。技术路线分为两个层面：第一，系统层面——将分块堆叠为 3D 张量，用批量化矩阵乘法（bmm）替代逐块顺序处理；第二，算法层面——引入 Newton-DB 迭代和 Chebyshev 多项式逼近作为更快的逆根求解器，并提出 multi-Power-Iteration 来优化矩阵缩放。这两个层面的改进叠加，实现了对 Distributed Shampoo 的全面加速。
DASH 最核心的创新是两个互补的技术。第一个是分块堆叠策略：同一层的 $L$ 矩阵块、$R$ 矩阵块以及它们的逆根可以堆叠在一起，因为它们具有相同的 $(B, B)$ 形状。这不仅是简单的批处理，而是一个精心设计的内存管理方案——它消除了 Distributed Shampoo 中因用列表存储独立块而产生的内存碎片化，使矩阵乘法可以充分利用 Tensor Core 的高吞吐量。更进一步，作者发现同一 GPU 上所有层的块都可以跨层堆叠（DASH_GPU），实现更高程度的并行化。第二个核心创新是对 Newton 迭代收敛行为的深入分析：作者证明了 Frobenius 范数缩放会将小特征值过度压缩（例如一个本应 5 步收敛的 $\lambda = 10^{-2}$ 特征值被压缩到 $2 \times 10^{-4}$ 后需要 15 步才能收敛），并提出了 multi-Power-Iteration 方法来获得更精确的谱半径估计。这两个创新的结合使得 DASH 可以用更少的迭代次数、更低精度的浮点格式来完成逆根计算，从而大幅降低运行时间。
方法步骤详情
DASH 的方法步骤如下。首先，对于每个优化步，将梯度矩阵 $G \in \mathbb{R}^{m \times n}$ 分割为 $B \times B$ 的块，形成 3D 张量 $\text{block}(G) \in \mathbb{R}^{N \times B \times B}$，其中 $N = N_m \cdot N_n$ 是总块数。对于不足 $B$ 的残余块单独处理。接着，用批量化操作计算 $L$ 和 $R$ 的分块：$GG^\top$ 和 $G^\top G$ 通过张量转置和 bmm 操作高效完成。然后，将所有 $(B, B)$ 形状的 $L$ 块、$R$ 块堆叠为 3D 张量。对于逆根计算，先用 multi-Power-Iteration 估计每个块的谱半径，使用 16-32 个随机起始向量并行估计，选取最大 Rayleigh 商对应的特征向量，然后以 $2\lambda_{PI}$ 进行缩放。之后调用选定的逆根求解器（NDB、CN、CBSHV 或 EVD）对整个堆叠张量进行批量处理。NDB 迭代每次需要 3 次矩阵乘法，但通过跳过第一次迭代中的冗余计算可以省去 2 次。CN 迭代对于 $p=4$ 需要 4 次矩阵乘法。最后，将逆根结果用于计算预条件梯度 $U_t = L_t^{-1/4} G_t R_t^{-1/4}$，结合 Adam grafting 缩放学习率后更新参数。
技术新颖性
DASH 的技术新颖性体现在多个层面。与 Distributed Shampoo 相比，最本质的区别是「从顺序处理到批量并行处理」的范式转换——Distributed Shampoo 的分块只是降低了单次计算的复杂度，但并未改变执行模式，而 DASH 通过堆叠策略实现了真正的 GPU 友好并行化。与 Coupled-Newton 相比，Newton-DB 的引入本身就是新颖的：NDB 可以同时计算平方根和逆平方根（CN 只能计算逆根），且在使用精确的谱半径估计时，NDB 在 $x \in (0.3, 1)$ 区间需要的迭代次数显著少于 CN。此外，论文首次对 Newton 迭代中的矩阵缩放问题进行了系统性分析，发现 Frobenius 范数缩放是次优的，并用 multi-Power-Iteration 解决了 Power-Iteration 可能收敛到非最大特征向量的问题。Chebyshev 多项式逼近（CBSHV）虽然不是全新方法，但将其应用于 Shampoo 逆根计算并通过 Clenshaw 算法的优化将矩阵乘法次数从 $d+2$ 降至 $d-1$，也是有意义的工程创新。
实验结果
DASH 在 Llama-953M 模型（$E=2048$）上进行了全面的基准测试，使用 C4 数据集、Chinchilla-optimal token 数量（约 180 亿 tokens）、9089 个优化步。核心发现如下：（1）在所有稳定的配置中，DASH 的验证困惑度与 Distributed Shampoo 一致（差异约 $\pm 0.01$），同时显著降低运行时间。最突出的改进是 CN-FP16 配置在 $B=1024$ 时，GPU 版本的优化器步时间从 675ms（DIST-CN-FP32）降至 119ms，实现了 5.6 倍加速。（2）NDB（Newton-Denman-Beavers）迭代在验证困惑度上一致优于 CN，甚至在多个配置中超过 EVD。使用 Power-Iteration 归一化的 NDB 在 $B=1024$ 时达到最低的验证困惑度 11.68，而 EVD 为 11.72，CN 为 11.87。（3）Power-Iteration 归一化相比 Frobenius 范数归一化能显著改善 NDB 的数值稳定性和优化质量——在 $B=2048$ 时，Frobenius 归一化的 NDB 在 DIST 中全部种子失败，而 Power-Iteration 归一化稳定运行。（4）FP16 精度对 CN 无损（困惑度不变），但能降低约 10% 的运行时间。然而 NDB 在 FP16 下存在数值不稳定问题，需要进一步研究。（5）从整体训练时间来看，以 9000 步、前向 1000ms + 后向 3000ms 为例，使用 DASH-CN-FP16（138ms/步）替代 DIST-EVD-f10（355ms/步），总训练时间从 10 小时 53 分钟降至 10 小时 21 分钟，节省约 30 分钟（5%）。（6）视觉 Transformer 的初步实验同样支持 NDB 的优势：在 ImageNette/ImageWoof 上，DASH-NDB 相比 DASH-CN 测试准确率提升 0.8-6 个百分点，且运行时间仅为 9-15ms/步（DIST 为 104-145ms/步）。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| Llama-953M 验证困惑度（C4，B=1024） | Validation Perplexity (↓) | DASH-NDB-PIM-GPU: 11.68 | DIST-EVD: 11.80, DIST-CN-FP32: 11.87 | 比 EVD 低 0.12，比 CN 低 0.19 |
| Llama-953M 优化器步时间（B=1024, CN-FP16） | Opt Step Time (ms, ↓) | DASH-GPU: 119ms | DIST-CN-FP32: 666ms | 5.6× 加速 |
| Llama-953M 优化器步时间（B=2048, CN-FP32） | Opt Step Time (ms, ↓) | DASH-GPU: 207ms | DIST-CN-FP32: 675ms | 3.26× 加速 |
| Llama-953M 优化器步时间（B=1024, NDB-PIM） | Opt Step Time (ms, ↓) | DASH-GPU: 177ms | DIST-NDB-PIM: 740ms | 4.18× 加速 |
| ImageNette/ViT-Tiny 测试准确率 | Test Accuracy (%, ↑) | DASH-NDB: 82.45% | DASH-CN: 81.62%, DIST-CN: 81.05% | 比 DASH-CN +0.83pp，比 DIST-CN +1.40pp |
| ImageWoof/ViT-Small 测试准确率 | Test Accuracy (%, ↑) | DASH-NDB: 70.28% | DASH-CN: 64.14%, DIST-NDB: 72.08% | 比 DASH-CN +6.14pp，略低于 DIST-NDB (-1.80pp) |
局限与改进
论文存在以下局限性。首先，作者明确指出 NDB 在 FP16/BF16 精度下存在数值不稳定问题，无法收敛，这限制了 NDB 在更低精度下的应用。相比之下 CN 在 FP16 下运行良好，说明 NDB 的数值稳定性是一个需要单独解决的问题。其次，Chebyshev 多项式（CBSHV）方法在小模型（Llama-373M）的初步实验中表现不够稳定——Distributed Shampoo 的 FP32 版本困惑度高达 18.6（FP16 为 16.15），且在更大模型（Llama-953M）上，CBSHV 的困惑度（11.98-12.11）不如 CN（11.87）和 NDB（11.68），作者因此选择将其放在附录中作为初步探索。第三，实验仅在 Llama-953M 上进行，缺乏更大规模模型（如 7B、13B）和更多架构的验证。视觉 Transformer 的实验也仅是初步结果，且需要将 3D/4D 张量转换为 2D 来适配当前实现。第四，虽然优化器步时间大幅降低，但对整体训练时间的实际改善仅约 5%（953M 模型场景），因为前向和后向传播仍占主要时间。第五，论文未与 SOAP 等近年的其他 Shampoo 改进方法进行直接对比。
独立分析的弱点
从独立分析的角度，DASH 存在以下可改进之处。第一，NDB 在 FP16 下不收敛是一个严重的实际限制——在大规模训练中 FP16/BF16 是标配精度，如果 NDB 只能在 FP32 下工作，则其速度优势会被精度成本部分抵消。改进方向可以是引入误差反馈机制（如 Modoranu et al. 2023 所示）或混合精度策略，在 NDB 的关键步骤使用 FP32 累积而其余步骤使用 FP16。第二，multi-Power-Iteration 虽然有效，但引入了额外的超参数（起始向量数量 16-32），且在特征值分布特殊的情况下仍可能收敛到非最大特征向量。可以考虑结合 Lanczos 方法或随机 SVD 来获得更稳健的谱估计。第三，分块堆叠策略要求同一 GPU 上的块具有相同形状，但残余块（不足 $B$ 的部分）需要单独处理，这在层形状高度异构的模型中可能导致负载不均衡。第四，论文对 CBSHV 的探索不够深入，仅在小模型上做了初步实验，考虑到 Chebyshev 多项式可以用 $d-1$ 次矩阵乘法实现 $d+1$ 阶精度的理论优势，值得在更大模型上验证。
未来方向
作者提出了几个明确的未来方向。第一，动态求解器选择：根据预条件矩阵块的条件数，为每个块选择最优/最廉价的逆根求解器，例如条件数好的块用 CN，条件数差的用 NDB。第二，稳定 NDB 的低精度执行：结合随机方法或误差校正技术。第三，在更大模型和数据规模上验证 DASH，以及在张量并行训练场景中的应用。此外，基于本文的成果还可以延伸以下方向：将 DASH 的堆叠策略推广到 3D/4D 卷积层（当前限制在 2D），探索 NDB 与其他矩阵函数计算的结合（如矩阵对数、矩阵指数），以及将 multi-Power-Iteration 应用于其他需要谱估计的优化器（如 Muon）。
复现评估
DASH 的代码已在 GitHub 开源（https://github.com/IST-DASLab/DASH），这是复现的最大利好。实验使用的是公开的 C4 数据集和标准 Llama 架构，数据获取无门槛。在算力需求方面，953M 模型的完整实验使用 8 张 GPU，每组实验约 4-10 小时，3 个种子约 12-30 小时的 GPU 时间——这在学术环境中可行但在个人设备上较难。关键的超参数选择（学习率 $\eta^*=10^{-3}$ 经过网格搜索、块大小 $B \in \{1024, 2048\}$、预条件子更新频率 $f \in \{1, 10\}$）都有详细记录。唯一的不确定性是 NDB 在 FP16 下的数值不稳定性可能导致复现时需要额外调试，且 CBSHV 方法的实验仅在附录中初步展示。总体而言，复现难度为中等偏低——代码开源、数据公开、实验规模可控，但需要一定的 GPU 集群资源。
论文图表