Pith. sign in

REVIEW 3 major objections 6 minor 2 cited by

Shampoo optimizer steps run up to 4.8× faster with batched block preconditioning and a Newton-based inverse-root solver.

Reviewed by Pith at T0; open to challenge. T0 means a machine referee read the full paper against a public rubric. the ladder, T0–T4 →

Challenge this review Re-run · record.json Download PDF Read on arXiv ↗

Batched 3D block preconditioning plus Newton–Denman–Beavers and Chebyshev inverse-root solvers cut Shampoo's per-step cost up to ~4.8× on a 953M-parameter Llama model without hurting perplexity.

T0 review reviewed 2026-08-03 challenge →

load-bearing objection Genuinely useful Shampoo engineering with honest reporting; speedups hold in the tested setting, but the scaling-convergence argument is heuristic and the end-to-end gain is modest. the 3 major comments →

arxiv 2602.02016 v2 pith:EC76XWZX submitted 2026-02-02 cs.LG

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

Ionut-Vlad Modoranu , Philip Zmushko , Erik Schultheis , Mher Safaryan , Dan Alistarh This is my paper
classification cs.LG
keywords Shampoo optimizersecond-order optimizationinverse matrix rootsNewton-Denman-Beavers iterationbatched block preconditioningpower iterationpreconditioner scalinglanguage model pretraining
verification ladder T0 review T1 audit T2 compute T3 formal T4 reserved

The pith

A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.

The reading

The paper aims to make Shampoo, a powerful second-order optimizer, cheap enough for everyday large-model training. It argues that Shampoo's main bottleneck — computing inverse matrix roots of preconditioner blocks — can be accelerated by stacking all blocks into a single 3D tensor and applying the Newton-Denman-Beavers (NDB) iteration in parallel with half-precision matmuls. It further shows that scaling each block by an estimate of its largest eigenvalue (via a robust multi-power-iteration) is critical: the usual Frobenius-norm scaling slows convergence and can cause divergence. In a 953M-parameter language-model pretraining run, the resulting DASH implementation matches or improves validation perplexity while reducing optimizer step time by up to 4.83× relative to a well-optimized Distributed Shampoo baseline.

Core claim

The central claim is that the practical cost of Shampoo preconditioning can be reduced by about fivefold without loss of quality, by replacing sequential block-by-block inverse-root computation with batched operations on stacked 3D tensors, and by replacing the default Eigen-Value Decomposition (EVD) solver with the Newton-Denman-Beavers (NDB) iteration scaled by a robust power-iteration estimate of the spectral radius. Under this scheme, NDB achieves the lowest validation perplexity of all tested solvers (11.68 vs 11.80 for EVD and 11.87 for Coupled-Newton at block size 1024), and the DASH-CN-FP16 configuration runs in 138 ms per optimizer step versus 666 ms for the baseline Dist-CN-FP32 —

What carries the argument

The argument rests on three pieces: (1) a blocking strategy that stacks all left and right preconditioner blocks of matching shape into a single 3D tensor, so inverse-root routines execute as batched matrix multiplications that exploit GPU tensor cores; (2) the Newton-Denman-Beavers (NDB) iteration, a matrix-inverse-square-root fixed-point scheme (paired once to get the inverse fourth root) that needs only matmuls and converges quickly when the preconditioner spectrum is scaled near 1; and (3) multi-Power-Iteration, which estimates the largest eigenvalue of every block simultaneously in FP16 using a pool of starting vectors, and scales each block by 2λPI so the iteration's convergence condit

Load-bearing premise

The entire speedup and quality claim depends on the multi-power-iteration estimate λPI underestimating each block's true largest eigenvalue by no more than a factor of two, for every block, across the whole training run, so that scaling by 2λPI keeps the spectrum inside the convergence interval for all NDB and CN iterations.

What would settle it

A direct test: run the same DASH configuration but force a single block to have a spectrum whose true λmax is more than 2× the estimated λPI (e.g., by artificially injecting a large outlier eigenvalue into one preconditioner block during training). If the optimizer's validation perplexity degrades sharply or the iteration diverges, the scaling assumption is the load-bearing point; if training survives, the method is more robust than the paper's convergence analysis suggests.

Watch this falsifier. Get emailed when new claim-graph text bears on it.

If this is right

  • If the speedups hold, Shampoo-style preconditioning becomes affordable at larger scale, and practitioners could update preconditioners more frequently (the paper uses update frequency f=1) without paying the previous runtime penalty.
  • The perplexity ordering NDB < EVD < CN suggests that the choice of inverse-root solver and matrix scaling, not just the optimizer family, affects final model quality; switching from Frobenius to power-iteration scaling alone improved NDB perplexity from 11.76 to 11.68.
  • Because DASH matches baseline perplexity while cutting step time in all stable configurations, the main obstacle to adopting Shampoo — its per-step cost — is reduced, potentially making second-order methods competitive with Adam-style optimizers on wall-clock grounds.
  • The reported 4.83× is an optimizer-step speedup; the paper estimates that on a 953M model this translates to roughly a 5% reduction in total training time (about 30 minutes over ~11 hours), meaning the end-to-end benefit is real but modest.

Where Pith is reading between the lines

These are editorial extensions of the paper, not claims the author makes directly.

  • The paper's own instability results (NDB with Frobenius scaling failing at block size 2048; NDB and CN diverging under BF16) suggest that the speedup is tied to the precise combination of FP16, power-iteration scaling, and block size; other precisions, block sizes, or model shapes may require adaptive solver or scaling selection.
  • The 'dynamic solver selection' the paper proposes in its discussion — estimating a block's condition number and picking the cheapest convergent solver — could be tested directly by instrumenting each block's eigenvalue distribution during training to verify that the λPI estimate stays within the required factor of two.
  • Because the whole evaluation rests on one 953M Llama architecture on C4, the approach's generality to much larger models, transformers with different layer shapes, or non-transformer architectures remains open; a reader should treat the speedup as demonstrated for this configuration rather than universal.
Share X Bluesky LinkedIn Reddit HN

Editorial analysis

A structured set of objections, weighed in public.

Desk editor's note, referee report, simulated authors' rebuttal, and a circularity audit.

Referee Report

3 major / 6 minor

Summary. The paper proposes DASH, a distributed Shampoo implementation that batches preconditioner blocks into 3D tensors for better GPU utilization and introduces Newton-DB (NDB) and Chebyshev-polynomial (CBSHV) inverse-root solvers, together with a multi-power-iteration (multi-PI) scaling scheme. The authors report up to 4.83× faster optimizer steps than the well-optimized Distributed Shampoo implementation in a Llama-953M/C4 setting, and claim that NDB with power-iteration scaling attains the lowest validation perplexity per iteration among EVD, CN, and NDB. The paper includes a numerical analysis of matrix scaling, a description of the blocking strategy, and an appendix with preliminary Chebyshev experiments.

Significance. If the claims hold, the paper is a useful step toward making Shampoo practical: the 3D block stacking and batched inverse-root solvers address a real systems bottleneck, and the comparison against the Distributed Shampoo baseline is concrete. The code release is a strength, and the empirical finding that NDB can match or beat EVD at lower cost is potentially valuable. However, the evidence base is narrow (one model, one dataset, one training horizon), and the central stability argument for the 2λPI scaling contains a logical gap. The contribution is therefore significant but conditional on additional analysis and evaluation.

major comments (3)
  1. [§3.4–§3.5] The claim that dividing each preconditioner block by 2λPI 'makes sure' the NDB/CN convergence conditions hold does not follow from the stated one-sided bound λPI < λmax. For NDB, convergence requires ||I − A/(2λPI)||₂ < 1, i.e. λmax/(2λPI) < 2, or λmax/λPI < 4. For CN with the paper's c = (1+p)^{-1/p}, the required interval is λmax/(2λPI) < 1, i.e. λmax/λPI < 2. A fixed-budget multi-power-iteration can underestimate λmax by an arbitrarily large factor, so no guarantee is provided. The failure cases in Table 1 and Appendix A.4 (NDB-FRO fails for DIST at B=2048; NDB diverges in FP16/BF16; CBSHV-PI fails in FP16) show the normalization is fragile in practice. Since the headline NDB-PI perplexity result depends on every block staying in the convergent regime under a fixed ~10-iteration budget, this gap is load-bearing. The authors should either certify a lower bound on λPI (e.g., λPI ≥ λmax/
  2. [§3.4, Fig. 2] The CN-vs-NDB convergence comparison is normalized in a way that disadvantages CN. Equations (1)–(3) define M0 = A/c^p; the text sets c = (1+p)^{-1/p} to give the same interval [0,1] for A. With this c, M0 = (1+p)A, so as an eigenvalue of A approaches 1, the corresponding eigenvalue of M0 approaches p+1, exactly the boundary where the CN coefficient C0 = (1+1/p)I − (1/p)M0 vanishes. The 'peak around 1' in Figure 2 is therefore an artifact of this particular c, not an intrinsic property of CN. A matched comparison with c=1 (A already scaled to λmax(A)<1, so M0=A) removes the peak and changes the iteration-count ordering. The claim that NDB is intrinsically faster than CN is not supported by this experiment; the conclusion must be conditioned on the choice of c or the analysis redone with a standard normalization.
  3. [§5, Table 1] All quantitative claims are based on a single pretraining configuration: Llama-953M on C4, 9089 steps, batch size 2M tokens, 8 GPUs, and the GPU model is not stated. No standard deviations are reported because 'converged runs are extremely stable,' and there is no second model, dataset, or alternative block-size comparison for the solver ordering. Optimizer-step time is highly hardware-dependent, and the reported 4.83× speedup is in ms/step on this one setup. The perplexity ordering of EVD/CN/NDB is also a 3-seed observation at one scale. To support the abstract's general claims, at least one additional model/dataset or a clear scope limitation is needed; if the paper is intended as a systems report for this specific configuration, the claims should be narrowed accordingly.
minor comments (6)
  1. [Abstract vs §5] The abstract states 'up to 5.6× faster optimizer steps,' while Table 1 and §5 report 4.83× as the maximum. Reconcile.
  2. [Table 1] The table is difficult to read: the two block-size groups are not visually separated, and the meaning of ✗ is not footnoted. Add explicit grouping and a legend.
  3. [§5] The GPU hardware is not specified. For a timing-centric paper, the experimental setup should state the GPU model and framework versions.
  4. [§5, B=2048 NDB-FRO] The text says 'the runs using Frobenius normalization failed for DIST across all seeds and therefore we skipped the results,' but Table 1 lists a DASH-FRO value (11.68). Clarify whether the failure was DIST-only and why the DASH result is reported.
  5. [§3.5] Multi-Power-Iteration pool size (16 or 32) is mentioned but not reported per experiment. Specify the pool size in the experimental setup or in Table 1.
  6. [Appendix A.4] The Chebyshev results are preliminary and mixed: PI/FP16 fails, and perplexity is above NDB. The abstract and introduction still list Chebyshev as a main contribution. Either strengthen the evaluation or demote the claim.

Circularity Check

0 steps flagged

No load-bearing circularity: the core claims are empirical comparisons against an external baseline, and the only self-citations are contextual.

full rationale

The paper's central claims—faster optimizer steps and lower validation perplexity—are empirical results benchmarked against the external Distributed Shampoo implementation of Shi et al. (2023), not derivations from premises that contain the conclusions. The inverse-root methods (NDB, CN, CBSHV) are explicitly attributed to standard sources (Higham 2008; Cody 1970; Clenshaw 1955), and the scaling analysis in Section 3.4 is a direct numerical experiment on scalar eigenvalues, not a fitted parameter being renamed as a prediction. The self-citations (Modoranu et al. 2023; Pascanu et al. 2025; Vlassis et al. 2025) appear only as background or as a suggested future direction and are not load-bearing for the speedup or perplexity claims. The possible objection that scaling by 2λPI does not mathematically guarantee NDB/CN convergence is a correctness or stability risk, not a circular reduction: no equation in the paper is equivalent to its input by construction, and no fitted quantity is later reported as an independent prediction. Accordingly, no circular step meets the evidentiary bar, so the steps list is empty.

Axiom & Free-Parameter Ledger

6 free parameters · 6 axioms · 0 invented entities

The paper contributes no new mathematical objects; its novelties are algorithmic and engineering variants (NDB-in-Shampoo, Chebyshev/Clenshaw for matrix roots, multi-Power-Iteration, 3D block stacking). The ledger therefore lists the training and solver hyperparameters the empirical claims rest on, together with the standard numerical-analysis premises (convergence conditions from Higham 2008) the root solvers inherit. The fixed iteration budget is the most consequential free choice: §3.4 shows per-eigenvalue error depends on how many iterations are spent, and the entire §5 comparison runs at a fixed budget.

free parameters (6)
  • Learning rate η* = 1e-3
    Selected by grid search over {1e-4,...,4e-3} for AdamW and reused for Shampoo via Adam grafting (§5 'Learning Rate'); all val-ppl numbers depend on this choice.
  • Block size B = 1024 (2048 in some runs)
    B is a hand-chosen design parameter (§5 'Block Size'); the paper recommends B=1024 after observing similar perplexity at lower runtime.
  • Preconditioner update frequency f = 1 (10 for EVD)
    The headline speedups are for f=1 configurations; f is a Shampoo hyperparameter that trades per-step cost against approximation quality (§5).
  • Fixed iteration budget for NDB/CN = ~10 (per Shampoo convention)
    Iterative root solvers run a fixed number of steps (§3.4: 'Since we use a fixed number of steps (e.g., 10)'); approximation error per eigenvalue then depends on eigenvalue magnitude, which is the crux of their analysis.
  • multi-Power-Iteration pool size = 16–32 starting vectors
    Chosen to use tensor cores efficiently (§3.5); the quality of the λPI estimate sets the scaling factor 2λPI and thereby the convergence margin.
  • EVD regularization ϵ = 1e-10 (example)
    Dampening for EVD (§Appendix B); the paper's corrected-spectrum heuristic changes EVD results relative to the official implementation.
axioms (6)
  • standard math Newton-DB computes A^{±1/2} for matrices satisfying ||I−A||₂ < 1
    Invoked in §3.3 (Eqs. 4-6) via Higham (2008) Eq. 6.35; the paper relies on the classical convergence condition without proof and chains two NDB calls to obtain A^{−1/4}.
  • standard math Coupled-Newton converges for spectra in [0, (p+1)c^p]; choosing c=(1+p)^{−1/p} gives the interval [0,1]
    §3.2 and §3.4 (Eqs. 1-3): the entire scaling discussion is built on this convergence interval.
  • standard math Power iteration with a pool of 16–32 random starts finds the dominant eigenvector; the Rayleigh quotient underestimates λmax
    §3.4–3.5: the 2λPI scaling factor is justified only if λPI ≤ λmax and close enough that the scaled spectrum stays inside the convergence regions.
  • domain assumption Tensor-core batched matmul and FP16 arithmetic preserve enough precision for correct optimizer behavior
    §2.4 and §4: the speedups rest on bmm/FP16; the paper itself shows BF16 diverges for CN and FP16/BF16 diverge for NDB, so the precision margin is thin.
  • domain assumption Adam grafting is required for numerically stable Shampoo, and the grafted learning rate transfers
    §2.3: 'Grafting is mandatory to have a numerically stable implementation for Shampoo'; all val-ppl results therefore inherit the AdamW LR tuned on the same grid.
  • ad hoc to paper One pretraining setup (Llama-953M, C4, Chinchilla-optimal 20 tokens/param, 9089 steps, 8 GPUs) is representative enough to rank root solvers
    §5 'Setting': the entire empirical contribution (speedup factors and perplexity ordering) is measured on this single configuration; the Discussion admits scale validation is future work.

reviewed 2026-08-03 · how reviews work

0 comments
Cite this review

Pith. "Pith review of DASH: Faster Shampoo via Batched Block Preconditioning and Efficient Inverse-Root Solvers." pith.science (2026). https://pith.science/paper/EC76XWZX

@misc{pith2026260202016,
  author       = {Pith},
  title        = {Pith review of: DASH: Faster Shampoo via Batched Block Preconditioning and Efficient Inverse-Root Solvers},
  year         = {2026},
  howpublished = {\url{https://pith.science/paper/EC76XWZX}},
  note         = {Machine review of arXiv:2602.02016}
}
Share X Bluesky LinkedIn Reddit HN
abstract

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.

Figures

Figures reproduced from arXiv: 2602.02016 by Dan Alistarh, Erik Schultheis, Ionut-Vlad Modoranu, Mher Safaryan, Philip Zmushko.

Figure 2
Figure 2. Figure 2: Number of steps required for NDB and CN to compute the square and inverse square roots of scalars between 0 and 1 (in linear scale) up to precision 10−10. We added a shift for NDB iterations to improve visibility on the y-axis. 0.0 0.2 0.4 0.6 0.8 1.0 eigen-value x 5 10 15 number of iterations iters for x (SR) via NDB iters for x 1 (ISR) via NDB iters for x 1 (ISR) via CN The required number of steps to ac… view at source ↗

discussion (0)

Sign in with ORCID, Apple, or X to comment. Anyone can read and Pith papers without signing in.

Forward citations

Cited by 2 Pith papers

Reviewed papers in the Pith corpus that reference this work. Sorted by Pith novelty score.

  1. Hierarchical Muon: Tiled Newton-Schulz Updates for Efficient Muon Optimization

    math.NA 2026-06 unverdicted novelty 7.0

    HiMuon partitions momentum-gradient matrices into T x T tiles, runs independent Newton-Schulz iterations on each tile, and reassembles the results, reducing leading cost to O(H W T K) while defining a local rather tha...

  2. SOAP, Muon, and Beyond: Pushing LLM Pretraining Scales

    cs.LG 2026-07 conditional novelty 6.0

    SOAP and Muon, stabilized by per-step QR eigenbasis updates and KL-Shampoo covariance accumulation, beat AdamW on large-batch LLM pretraining up to 100M-token batches.

Reference graph

Works this paper leans on

4 extracted references · cited by 2 Pith papers

  1. [1]

    Compute symmetric eigendecompositionλ, Q←eigh(L)whereλ∈R n andQ∈R n×n

  2. [2]

    Computeλ min ←min i λi

  3. [3]

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

    Computeλ new ←λ−min(λ min,0)1 +ϵ1. 15 DASH: Faster Shampoo via Batched Block Preconditioning and Efficient Inverse-Root Solvers

  4. [4]

    In the Distributed Shampoo implementation, the authors implement step 1 as λ, Q←eigh(L+ϵI n) and then proceed with steps 2, 3 and 4

    Form and return matrix root inverseL inv ←Qdiag(λ −r new)Q T . In the Distributed Shampoo implementation, the authors implement step 1 as λ, Q←eigh(L+ϵI n) and then proceed with steps 2, 3 and 4. Our observation is that the eigenvalues λ already contain the regularization ϵ and we state it should be subtracted from λ after the first step, otherwise in ste...

This paper was first reviewed by deepseek-v4-flash on August 3, 2026.