PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Ionut-Vlad Modoranu, Philip Zmushko, Erik Schultheis, Mher Safaryan, Dan Alistarh

arXiv:2602.02016v2cs.LG

TL;DR

Shampoo’s expensive internal operations, especially inverse matrix-root computation, create a substantial computational bottleneck. DASH accelerates these operations through GPU-efficient block execution and alternative inverse-root solvers, achieving up to 5.6× faster optimizer steps while Newton-DB attains the lowest validation perplexity per iteration among tested methods.

  • Problem

    Shampoo’s inverse-root computations remain costly and difficult to parallelize on GPUs, limiting efficient use of non-diagonal second-order information.

  • Method

    DASH stacks preconditioner blocks for GPU execution and evaluates Newton-DB and Chebyshev approximations as alternative inverse matrix-root solvers.

  • Results

    Newton-DB achieves lower validation perplexity than Coupled-Newton and standard EVD in both Distributed Shampoo and DASH implementations.

  • Takeaways & Limitations

    For the evaluated GPU setting, NDB with Power-Iteration scaling minimizes validation loss, whereas CN with FP16 and Frobenius normalization minimizes optimizer runtime.

  • Takeaways & Limitations

    CN diverges with BF16 because of its lower precision, despite BF16 having the same range as FP32.

Abstract

from arXiv · show

Shampoo is one of the leading approximate second-order optimizers: a variant of it has won the MLCommons AlgoPerf competition, and it has been shown to produce models with lower activation outliers that are easier to compress. Yet, applying Shampoo currently comes at the cost of significant computational slowdown, due to its expensive internal operations. In this paper, we take a significant step to address this shortcoming by proposing \method (for \textbf{D}istributed \textbf{A}ccelerated \textbf{SH}ampoo), a faster implementation of Distributed Shampoo based on two main new techniques: First, we show that preconditioner blocks can be stacked into 3D tensors to significantly improve GPU utilization; second, we introduce the Newton-DB iteration and the Chebyshev polynomial approximations as novel and faster approaches for computing the inverse matrix roots required by Shampoo. Along with these algorithmic contributions, we provide a first in-depth analysis of how matrix scaling critically affects Shampoo convergence. On the practical side, our GPU-aware implementation achieves up to $5.6\times$ faster optimizer steps compared to the well-optimized Distributed Shampoo, while Newton-DB attains the lowest validation perplexity per iteration among all tested methods. Our code is available at https://github.com/IST-DASLab/DASH.

1. Introduction

Shampoo can outperform diagonal optimizers but remains costly because inverse matrix roots create substantial per-step overhead and GPU efficiency gaps. DASH addresses these bottlenecks through GPU-oriented blocking, faster inverse-root methods, improved scaling, and convergence analysis.

  • Motivation: Diagonal optimizers reduce memory complexity to O(mn), but fail to capture complex parameter correlations.Full-matrix AdaGrad has O(m2n2) memory complexity for an m×n layer.
  • Motivation: Shampoo has recently achieved leading wall-clock performance in AlgoPerf, while its per-step computation remains substantially more expensive than AdamW.Inverse matrix roots typically scale as Θ(n3) for an n × n preconditioner matrix, motivating infrequent updates every 10–100 steps.
  • Motivation: Distributed Shampoo lowers inverse-root complexity to O(Bn2) through B × B blocking, but its default EVD1 is difficult to parallelize on GPUs.Its matrix-multiplication-based Coupled-Newton alternative is more GPU-friendly but is not enabled by default, likely because of numerical-stability concerns.
  • DASH: DASH is a high-performance Distributed Shampoo implementation that targets the inverse matrix root bottleneck while preserving numerical precision and leveraging modern GPUs.It aims to make high-quality preconditioning practical for widespread use.
  • DASH: DASH introduces Newton-DB and Chebyshev polynomial methods, analyzes Newton-iteration scaling, and proposes multi-Power-Iteration to estimate spectral radius efficiently.Multi-Power-Iteration uses half precision, avoids local maxima, and enables Newton procedures to satisfy convergence criteria rapidly.

2. Preliminaries on Shampoo

Shampoo maintains exponential-moving-average left and right preconditioners and applies their inverse fourth roots to gradient updates. DASH inherits Distributed Shampoo features including grafting and load balancing, while adding lower-precision iterations, efficient kernels, and improved memory utilization.

  • Shampoo: Shampoo forms left and right preconditioners from exponential moving averages of G_tG_t^⊤ and G_t^⊤G_t, then applies their regularized inverse fourth roots to G_t.The simplified update uses L_t and R_t with regularization terms εI.
  • Distributed Shampoo features: DASH uses block-wise Adam grafting to rescale Shampoo’s update direction to Adam’s update norm and inherits Distributed Shampoo’s implementation features.Grafting preserves Shampoo’s direction while transferring Adam’s tuned learning-rate schedule.
  • Distributed Shampoo features: DASH greedily assigns layers to the least-loaded GPU after sorting them by descending parameter count, scattering parameters to avoid redundant computation.This follows the optimizer-state partitioning strategy in ZeRO.
  • Implementation improvements: Around 10% lower optimizer-step runtime is achieved by introducing CN in FP16 instead of FP32, with no validation-perplexity degradation.FP16 improves validation perplexity and runtime for CBSHV but causes numerical instabilities for NDB.
  • Implementation improvements: At 953M parameters on 8 GPUs, DASH uses 73 GB per GPU for higher-rank workers and 71 GB for lower-rank workers, versus 76 GB for Distributed Shampoo.The improvement comes from stacking blocks together with load balancing.

3. Inverse Root Methods

The section compares iterative inverse-root methods with EVD, introducing Newton-DB and analyzing how matrix scaling affects convergence. It also improves Power-Iteration through multiple starting vectors to obtain more reliable eigenvalue estimates.

  • Inverse-root methods: EVD accurately computes matrix powers but is difficult to parallelize on GPUs because its iterative eigensolver builds and re-orthogonalizes a Krylov subspace.
  • Newton-DB: Newton-DB computes inverse square roots iteratively with 3 matmuls per iteration and obtains inverse fourth roots by chaining two calls.The first iteration can save two matmuls by directly initializing its closed-form values.
  • Matrix scaling: 10−100×: Distributed Shampoo’s Frobenius-norm scaling exceeds λmax(A) by around this factor, pushing eigenvalues toward zero.The Frobenius norm is computationally cheaper than Power-Iteration in their implementation.
  • Matrix scaling: 5 to 15 steps: after Frobenius scaling changes λ = 1e-2 to λ = 2e-4, CN’s convergence requirement increases from 5 steps to 15 steps.The example assumes the Frobenius norm is 50× larger than the largest eigenvalue.
  • Matrix scaling: Around 1: CN requires its peak number of iterations in x ∈(0.3, 1), whereas NDB requires fewer steps when λmax(A) is accurately approximated.The paper uses this behavior to support Power-Iteration-based scaling and reports that it is stable for larger block sizes where Frobenius scaling causes NDB instabilities.
  • Power-Iteration: 16 or 32 starting vectors: multi-Power-Iteration estimates largest eigenvectors in parallel and selects the one with the largest Rayleigh quotient.This reduces the likelihood that Power-Iteration converges to an eigenvector that is not the largest.

4. Blocking Strategy

DASH accelerates Distributed Shampoo by stacking compatible preconditioner blocks into batched 3D tensors, enabling efficient inverse-root computations and GPU-level processing. This strategy reduces optimizer-step running time by up to 4−5× compared with Distributed Shampoo while reducing memory fragmentation and enabling lower-precision matrix multiplications.

  • Stacking strategy: 4−5× lower optimizer-step running time is achieved by DASH compared with Distributed Shampoo.The reported speedup follows from the efficient blocking strategy implemented in DASH.
  • Stacking strategy: DASH stacks compatible blocks into 3D tensors and applies batched operations for inverse-root procedures, improving GPU utilization.The approach modifies CN, NDB, and CBSHV to support batched matrix multiplications, while EVD already supports batching.
  • Stacking strategy: For B = 1024, V = 32 000, and E = 2048, the embedding gradient forms 62 full blocks and 2 remainder blocks.These blocks have shapes Gfull ∈(62, 1024, 1024) and Grest ∈(2, 256, 1024).
  • Stacking strategy: Matrices Lfull, Rfull, and Rrest share shape (1024, 1024) and are stacked into Sfull ∈(126, 1024, 1024).The remaining matrix is stored as Srest = Lrest ∈(2, 256, 256).
  • Stacking strategy: DASHGPU stacks matrices across all layers assigned to one GPU, whereas DASHLAYERWISE performs stacking per layer.The GPU-level variant is faster than the layerwise implementation.
  • Stacking benefits: Stacking avoids memory fragmentation and makes non-EVD matrix multiplications more efficient, supporting lower-precision formats such as float16.It also enables more accurate estimates of the largest eigenvalue used to scale iterative-procedure inputs.

5. Experimental Results

Experiments on a 953M-parameter Llama model show that DASH largely preserves Distributed Shampoo’s validation perplexity while substantially reducing optimizer-step runtime. Newton-DB with Power-Iteration scaling achieves the lowest reported validation perplexity, while CN provides the largest speedups but is sensitive to numerical precision.

  • Setting: The study pretrains a 953M-parameter Llama model for 9089 optimization steps across three seeds, measuring average validation perplexity and optimizer-step runtime.Forward and backward runtimes are held constant, and the selected learning rate is η* = 1e-3.
  • Overall Trends: DASH matches DIST in almost all stable configurations while reducing optimizer-step runtime by up to 4× in one-to-one comparisons and 5× across inverse-root methods.Some NDB configurations also achieve lower validation perplexity than EVD.
  • EVD: For EVD with B = 2048 and f = 1, DASH reduces runtime from 2200ms to 1747ms for LW and 1755ms for GPU, while maintaining nearly identical validation perplexity.The reported speedups are 1.26× and 1.25×, respectively.
  • CN: DASH’s CN-FP16 GPU configuration reaches 5.6× faster runtime than DIST-CN-FP32 at block size B = 1024, reducing 666ms to 119ms.CN variants match DIST’s validation perplexity of 11.87 under Frobenius normalization.
  • NDB: NDB with Power-Iteration scaling achieves validation perplexity 11.68, outperforming Frobenius scaling’s 11.76 and matching or outperforming EVD across tested block sizes.NDB did not converge in FP16 or BF16 experiments.
  • NDB: NDB Power-Iteration at B = 1024 reduces runtime from 740ms for DIST to 194ms for DASH-LW and 177ms for DASH-GPU, with speedups of 3.81× and 4.18×.Power-Iteration is inefficient for DIST when preconditioning blocks sequentially.

6. Related Work and Discussion … A.3. Numerical Precision.

The paper situates DASH among efforts to reduce Shampoo’s overhead and presents Chebyshev-based inverse-root approximation, including matrix scaling, optimized evaluation, and mixed-precision execution. These techniques target fragmented execution and numerical instability while improving computational efficiency.

  • 6. Related Work and Discussion: Prior work reduces Shampoo overhead through 4-bit quantization of eigenvectors or preconditioning matrices, while randomized SVD enables block low-rank decomposition for AdaGrad.These approaches include eigenvector quantization, Cholesky-based preconditioning-matrix quantization, and shared-basis randomized SVD.
  • 6. Related Work and Discussion: DASH identifies fragmented execution and numerical instability as primary practical bottlenecks and addresses them through numerical analysis and system design.Its numerical contributions include convergence analysis, different solvers, and multi-Power-Iteration.
  • A. Chebyshev polynomials: Chebyshev polynomials provide an orthogonal polynomial basis for approximating Shampoo’s inverse roots, with the first kind selected for better endpoint accuracy.The first-kind polynomials Tn(x) are defined recurrently, and the second kind uses T1(x) = 2x.
  • A.1. Chebyshev polynomials for real numbers: For real functions, a degree-d Chebyshev approximation uses coefficients ck and can be evaluated with Clenshaw’s algorithm after mapping inputs to the interval [−1, 1].The coefficients are fitted once for f(x), degree d, and N points; evaluation then requires at most d + 1 multiplications per optimization step.
  • A.1. Chebyshev polynomials for real numbers: The approximation x^-1/p is accurate only on the coefficient-fitting interval [a, b], whose mapping is implicitly embedded in the fitted coefficients.Inputs must be mapped to [−1, 1] before evaluation, while the coefficients are fitted for [a, b].
  • A.2. Chebyshev polynomials for matrices: For matrices, the method fits coefficients on [ϵ, 1 + ϵ], scales A so its eigenvalues lie in [0, 1], then maps it to [−1, 1] before estimating A^-1/p.Scaling can use the Frobenius norm or Power-Iteration; the described algorithm uses the Frobenius norm.
  • A.2. Chebyshev polynomials for matrices: Optimizing the matrix Clenshaw recurrence reduces a degree-d polynomial from d + 2 to d −1 matrix multiplications by eliminating redundant zero-matrix and final-step computations.The optimized recurrence initializes Bd and Bd−1 directly and stops at B1 before computing the final approximation.
  • A.3. Numerical Precision.: Mixed precision stores Bk in float32 while multiplying S and converted iterates in float16 with float32 accumulation, exploiting two times higher tensor-core throughput.The paper reports that float16 multiplication with float32 accumulation improves results compared with float32 execution.

A.4. Experiments. · B. Additional Improvements to Distributed Shampoo

The experiments preliminarily evaluate CBSHV within Distributed Shampoo and DASH-A, finding precision- and normalization-dependent validation perplexity and runtime behavior. The paper also introduces further observations about the Distributed Shampoo implementation, while withholding the small-scale CBSHV results from the main body because both methods have drawbacks.

  • A.4. Experiments.: CBSHV is preliminarily tested at degree d = 60 on Llama-373M with E = 1024, B = 1024, and a 2-million-token batch in DIST and DASH-A.DASH-A updates normalization layers using Adam.
  • A.4. Experiments.: 18.6 validation perplexity occurs for float32 DIST, versus 16.15 for float16 under both Frobenius and Power-Iteration normalizations.The authors report no rigorous explanation for this precision effect.
  • A.4. Experiments.: Under float32 DASH-A, both normalizations converge, while Power-Iteration has higher validation perplexity and runtime remains under 90 milliseconds per optimizer step.This runtime represents an improvement of a bit less than 2×.
  • A.4. Experiments.: 76 milliseconds is the float16 DASH-A runtime when scaling by Frobenius norm, which converges in the reported experiment.The result indicates that DASH supports lower-precision computation, although the passage is truncated before stating the full implication.
  • A.4. Experiments.: The small-scale CBSHV results are excluded from the main paper because both DIST and DASH-A have drawbacks, despite the technique’s potential for broader improvements and applications.The authors present the results to bring CBSHV to the community’s attention.
  • A.4. Experiments.: The experiments next announce preliminary CBSHV results for a larger model, with Table 3 covering Llama-953M and degrees d ∈ {40, 60, 100}.Both Table 2 and Table 3 update preconditioners every optimization step, f = 1, and report time in milliseconds; Table 2 covers Llama-373M.
  • B. Additional Improvements to Distributed Shampoo: Section B records several observations about the Distributed Shampoo implementation that arose during preparation of the work.The passage introduces additional implementation-related observations without reporting a specific result.

B.1. Regularization (dampening) for EVD

EVD-based inverse roots require regularizing every preconditioner block, because failure of any block to converge fails training. The Distributed Shampoo implementation can apply inconsistent regularization by adding ϵ during eigendecomposition and again during eigenvalue adjustment, producing nonuniform shifts.

  • Regularization requirement: EVD inverse-root computation requires regularization of every preconditioner block for convergence; otherwise, the entire training fails.The paper reproduces the symmetric eigendecomposition procedure used for computing matrix root inverses.
  • Regularization procedure: The implementation computes λ, Q ← eigh(L + ϵI_n), then applies subsequent eigenvalue-adjustment steps before forming L_inv ← Q diag(λ_new^-r) Q^T.The eigendecomposition and returned root-inverse formulas are stated explicitly in the reproduced procedure.
  • Regularization inconsistency: Subtracting ϵ after eigendecomposition is necessary because the eigenvalues already include that regularization; otherwise, some entries receive 2ϵ while others receive less than ϵ.The paper identifies this as an inconsistent regularization caused by adding ϵ again during step 3.
  • Resulting eigenvalue shifts: The implementation increases λ_min by 2ϵ and each other eigenvalue λ_i≠λ_min by ϵ − λ_min, rather than increasing all eigenvalues uniformly by ϵ.This summary applies the paper’s stated behavior to the eigenvalues after the regularized eigendecomposition.

B.2. Our Dampening Heuristics.

The section proposes three dampening heuristics for Distributed Shampoo, centered on correcting the eigenspectrum and filtering or modifying small and negative eigenvalues before inverse-root computation. Experiments on two Llama models found no heuristic clearly superior because validation perplexities were similar, while the preferred Llama-953M frequency-1 setting was unaffordable to run.

  • Corrected spectrum: The corrected-spectrum procedure computes EVD on L + ϵI_n and then subtracts ϵ from the eigenvalues before inverse-root computation.This addresses inconsistencies in Distributed Shampoo regularization and avoids amplifying small-magnitude eigenvalues under the −1/p power.
  • Shifted-ReLU heuristic: Shifted-ReLU replaces each corrected eigenvalue with ReLU(λ −ϵ), retaining only significant positive entries and producing a rank-r_ϵ inverse.Here, r_ϵ counts the non-zero eigenvalues after shifted-ReLU.
  • ABS-based heuristic: The ABS-based heuristic applies absolute value to corrected eigenvalues and can optionally add ϵ to impose a lower bound.Unlike ReLU, this approach removes negative eigenvalues by making them positive.
  • Experiments: The study compares Shifted-ReLU, ABS-ADD, and SHMP—the original Distributed Shampoo heuristic—on Llama-373M with frequency 1 and Llama-953M with frequency 10.All methods are evaluated on DASH with EVD.
  • Results: Validation perplexities were similar, so no heuristic clearly stood out; the preferred Llama-953M frequency-1 experiment could not be run because compute was limited.The unavailable setting would have required a 17-hour run.

C. Standard Deviations for DASH

The reported standard deviations reflect experiments run with three seeds, selected because the runs were generally stable, particularly with Multi-PowerIteration.

  • Standard deviations for DASH: 3 seeds were used for only a few Table 1 experiments because the runs were quite stable.The paper reports the exact result obtained for each seed below.
  • Standard deviations for DASH: Runs were especially stable when using Multi-PowerIteration.

D. DASH for Computer Vision

DASH is evaluated on Vision Transformers by reshaping their 3D embedding layers into 2D, while comparing inverse-root methods on ImageNette and ImageWoof. NDB improves accuracy over CN, and Distributed Shampoo remains 10x slower than DASH despite being more accurate.

  • Scope and limitation: DASH is primarily designed for Transformer models with only 1D and 2D layers and has not been tested on other architectures.This limitation motivates converting higher-dimensional Vision Transformer layers into 2D.
  • Experimental setup: The experiments evaluate Tiny-5M and Small-21M Vision Transformers on ImageNette and ImageWoof.ImageNette contains 10 easy classes, whereas ImageWoof contains 10 difficult dog breeds.
  • Experimental setup: DASH converts each Vision Transformer’s 3D embedding layer from shape (E, 3, 16, 16) to 2D shape (E, 768).The final dimension merges the remaining dimensions because 768 = 3 * 16 * 16.
  • Inverse-root methods: When NDB is integrated into Distributed Shampoo, it computes power -1/2 for the first dimension and -1/4 for all remaining dimensions.For higher-dimensional layers, the original Shampoo procedure uses dimension-dependent inverse-root powers; DASH and Distributed Shampoo coincide for the first two dimensions.
  • Results: NDB is more accurate than CN in both DASH and Distributed Shampoo, with the largest accuracy advantage for Distributed Shampoo-NDB on ImageWoof/ViT-S.The results indicate that the inverse root used in the preconditioner does not impact the results in this converted 2D setting.
  • Results: 10x slower: Distributed Shampoo is slower than DASH despite being more accurate.These results show that NDB remains effective even when 3D and 4D tensors are converted to 2D for DASH.
Loading 2602.02016v2…