Technical Summary: DASH – Faster Shampoo via Batched Block Preconditioning and Efficient Inverse-Root Solvers

Problem Statement

Shampoo is a leading approximate second-order optimizer known for producing models with lower activation outliers and improved compressibility. However, its widespread adoption is hindered by significant computational overhead, primarily due to the expensive operation of computing inverse matrix roots ($A^{-1/4}$) for preconditioning matrices. While Distributed Shampoo attempts to mitigate this by splitting preconditioners into $B \times B$ blocks to reduce complexity from $O(n^3)$ to $O(Bn^2)$, existing implementations suffer from critical efficiency gaps:

1. Sequential Processing: Standard implementations process blocks sequentially, failing to fully utilize modern GPU parallelism.
2. Inefficient Solvers: The default reliance on Eigen-Value Decomposition (EVD) is difficult to parallelize on GPUs. Alternative iterative methods like Coupled-Newton (CN) are often disabled by default due to concerns over numerical stability and convergence speed.
3. Suboptimal Scaling: Current methods often scale matrices using the Frobenius norm, which is computationally cheap but can push eigenvalues toward zero, requiring significantly more iterations to converge to a target precision.

Methodology

The paper proposes DASH (Distributed Accelerated SHampoo), a high-performance implementation designed to bridge the gap between theoretical efficiency and hardware capabilities. The methodology rests on two primary pillars:

1. Batched Block Preconditioning (System-Level Optimization)

Instead of processing preconditioner blocks sequentially, DASH stacks independent blocks into 3D tensors.

• 3D Tensor Stacking: Gradient blocks $G_{ij}$ are stacked into a matrix of shape $(N, B, B)$, where $N$ is the total number of blocks.
• Parallel Execution: This allows the optimizer to apply inverse root procedures (EVD, CN, or new iterative methods) to the entire batch of blocks simultaneously using batched matrix multiplication (bmm) primitives.
• GPU Utilization: This architectural change maximizes the throughput of GPU tensor cores, reducing memory fragmentation and enabling efficient use of half-precision (FP16) arithmetic for certain solvers.

2. Efficient Inverse-Root Solvers (Algorithmic Optimization)

The paper introduces and analyzes advanced linear algebraic approaches to replace or augment standard EVD and CN:

• Newton-Denman-Beavers (NDB) Iteration: A novel approach for Shampoo that computes both square and inverse square roots. The authors optimize the first iteration to eliminate redundant matrix multiplications, reducing the overhead to 3 matmuls per step (effectively fewer in practice).
• Chebyshev Polynomial Approximations: Utilizing Clenshaw's algorithm to approximate matrix powers, offering an alternative iterative solver.
• Multi-Power-Iteration for Scaling: To address convergence issues, the paper replaces the standard Frobenius norm scaling with a Multi-Power-Iteration method.
  • Mechanism: It estimates the spectral radius ($\lambda_{max}$) using a pool of starting vectors (up to 32) processed in parallel via Power-Iteration.
  • Benefit: This provides a more accurate scaling factor than the Frobenius norm, ensuring eigenvalues remain in a regime where iterative solvers converge rapidly (closer to 1), thereby reducing the number of required iterations and improving numerical stability.

Key Contributions

1. DASH Implementation: A GPU-efficient framework that replaces sequential loops with parallel 3D tensor operations, achieving up to 5.6× faster optimizer steps compared to well-optimized Distributed Shampoo.
2. Novel Solvers: The introduction of the Newton-DB (NDB) iteration and Chebyshev polynomial approximations as viable, faster alternatives to EVD and CN. The paper demonstrates that NDB achieves lower validation perplexity than both Coupled-Newton and standard EVD.
3. Convergence Analysis: A detailed behavioral analysis showing that standard Frobenius norm scaling is suboptimal for iterative methods, often leading to slower convergence. The paper provides intuitive explanations for the distinct convergence behaviors of CN versus NDB.
4. Multi-Power-Iteration: An efficient, half-precision implementation of Power-Iteration that robustly estimates spectral radii, enabling optimal matrix scaling that satisfies convergence criteria for NDB and CN with fewer steps.

Experimental Results

The authors evaluated DASH on a Llama model with 953M parameters (embedding size 2048, sequence length 1024) trained on the C4 dataset.

• Speed:
  • DASH reduces optimizer step runtime by 4× to 5.6× compared to Distributed Shampoo, depending on the inverse root method and precision used.
  • The most significant speedup (5.6×) was observed when comparing Distributed Shampoo's EVD (FP32) against DASH's CN (FP16) with GPU-level stacking.
  • Even with frequent preconditioner updates ($f=1$), DASH maintains substantial speed advantages.
• Accuracy (Validation Perplexity):
  • DASH matches Distributed Shampoo in almost all stable configurations.
  • NDB consistently yields lower validation perplexity than CN and often outperforms EVD.
  • Using Power-Iteration scaling with NDB resulted in the lowest validation perplexity (11.68) compared to Frobenius scaling (11.76), validating the theoretical analysis.
• Precision:
  • The implementation supports FP16 for CN and EVD without degrading perplexity, further reducing runtime.
  • NDB in FP16/BF16 was found to be numerically unstable in their experiments, though the authors note this as an area for future work.

Significance and Claims

The paper claims that DASH significantly reduces the computational barrier to using second-order optimizers like Shampoo, making high-quality preconditioning practical for widespread use in large-scale deep learning.

• Bridging Theory and Hardware: By aligning algorithmic design (iterative solvers, matrix scaling) with hardware capabilities (tensor cores, batched operations), DASH allows researchers to investigate optimizer properties without prohibitive runtimes.
• Numerical Stability vs. Speed: The work challenges the assumption that iterative methods are inherently less stable or slower than EVD. By introducing better scaling (Multi-Power-Iteration) and optimized iterations (NDB), the authors demonstrate that iterative methods can achieve superior convergence and lower perplexity while being significantly faster.
• Future Directions: The authors suggest that the superior performance of Newton-DB opens avenues for dynamic solver selection based on block condition numbers and the potential for stabilizing NDB in lower-precision formats using stochastic or error-correction methods.

The paper concludes that the primary bottleneck in Shampoo was not mathematical complexity, but fragmented execution and numerical instability, both of which DASH addresses effectively.