Purifying Shampoo: Investigating Shampoo's Heuristics by Decomposing its Preconditioner

Extracted source text · arXiv 2506.03595

\documentclass{article}




\usepackage[preprint]{neurips_2025}








\usepackage[utf8]{inputenc} %
\usepackage[T1]{fontenc}    %
\usepackage[%
  colorlinks = true,
  citecolor  = Maroon,
  linkcolor  = MidnightBlue,
  urlcolor   = Maroon,
  unicode,
  linktocpage,
]{hyperref}       %
\usepackage{url}            %
\usepackage{booktabs}       %
\usepackage{amsfonts}       %
\usepackage{nicefrac}       %
\usepackage{microtype}      %
\usepackage[dvipsnames]{xcolor}         %

\input{math_commands.tex}


\usepackage[most]{tcolorbox}
\usepackage{graphicx}
\usepackage{subcaption}
\usepackage{cleveref}
\usepackage{makecell}
\usepackage{xspace}
\usepackage{algorithm}
\usepackage{algpseudocode}
\usepackage{multirow}
\usepackage{threeparttable}  %

\definecolor{tue1}{rgb}{0.55294118, 0.17647059, 0.22352941}
\definecolor{tue2}{rgb}{0.21568627, 0.25490196, 0.29019608}
\definecolor{tue3}{rgb}{0.68235294, 0.62352941, 0.42745098}
\definecolor{tue4}{rgb}{0.        , 0.41176471, 0.66666667}
\definecolor{tue5}{rgb}{0.68627451, 0.70196078, 0.71764706}
\definecolor{tue6}{rgb}{0.49019608, 0.64705882, 0.29411765}
\definecolor{tue7}{rgb}{0.78431373, 0.31372549, 0.23529412}
\definecolor{tue8}{rgb}{0.68627451, 0.43137255, 0.58823529}
\definecolor{tue9}{rgb}{0.70588235, 0.62745098, 0.58823529}
\definecolor{tue10}{rgb}{0.56862745, 0.41176471, 0.2745098}

\usepackage{amsthm}
\newtheorem{proposition}{Proposition}
\newtheorem{corollary}{Corollary}
\newtheorem{lemma}{Lemma}
\newtheorem{remark}{Remark}
\newtheorem{definition}{Definition}
\newtheorem{observation}{Observation}
\newtheorem{claim}{Claim}
\Crefname{claim}{Claim}{Claims}

\newcommand*{\ie}{i.e.\@\xspace}
\newcommand*{\eg}{e.g.\@\xspace}
\newcommand*{\cf}{c.f.\@\xspace}
\newcommand*{\wrt}{w.r.t.\@\xspace}
\newcommand*{\iid}{i.i.d.\@\xspace}
\newcommand*{\psd}{p.s.d.\@\xspace}
\newcommand*{\pd}{p.d.\@\xspace}

\usepackage{pifont}
\newcommand{\cmark}{\ding{51}}
\newcommand{\xmark}{\ding{55}}

\title{Purifying Shampoo: Investigating Shampoo's Heuristics by Decomposing its Preconditioner}


\author{%
  Runa Eschenhagen\thanks{Work performed while an intern and visiting researcher at FAIR, Meta Platforms.} \\
  Department of Engineering\\
  University of Cambridge\\
  \texttt{re393@cam.ac.uk} \\
  \And
  Aaron Defazio \\
  Fundamental AI Research \\
  Meta Platforms, Inc. \\
  \texttt{adefazio@meta.com} \\
  \AND
  Tsung-Hsien Lee \\
  AI and Systems Co-Design \\
  Meta Platforms, Inc. \\
  \texttt{zong@meta.com} \\
  \And
  Richard E. Turner \\
  Department of Engineering\\
  University of Cambridge\\
  The Alan Turing Institute \\
  \texttt{ret26@cam.ac.uk} \\
  \And
  Hao-Jun Michael Shi \\
  AI and Systems Co-Design \\
  Meta Platforms, Inc. \\
  \texttt{hjmshi@meta.com} \\
}


\begin{document}


\maketitle


\begin{abstract}
    The recent success of Shampoo in the AlgoPerf contest has sparked renewed interest in Kronecker-factorization-based optimization algorithms for training neural networks.
    Despite its success, Shampoo relies heavily on several heuristics such as learning rate grafting and stale preconditioning to achieve performance at-scale.
    These heuristics increase algorithmic complexity, necessitate further hyperparameter tuning, and lack theoretical justification.
    This paper investigates these heuristics from the angle of Frobenius norm approximation to full-matrix Adam and decouples the preconditioner's eigenvalues and eigenbasis updates.
    We show that grafting from Adam mitigates the staleness and mis-scaling of the preconditioner's \emph{eigenvalues} and how correcting the eigenvalues directly can eliminate the need for learning rate grafting.
    To manage the error induced by infrequent \emph{eigenbasis} computations, we propose an adaptive criterion for determining the eigenbasis computation frequency motivated by terminating a warm-started QR algorithm.
    This criterion decouples the update frequency of different preconditioner matrices and enables us to investigate the impact of approximation error on convergence.
    These practical techniques offer a principled angle towards removing Shampoo's heuristics and developing improved Kronecker-factorization-based training algorithms.
\end{abstract}

\setcounter{footnote}{0}  %

\section{Introduction}
\label{sec:introduction}
\input{contents/01_introduction}

\section{Background}
\label{sec:background}
\input{contents/02_background}

\section{Learning rate grafting compensates for Shampoo's eigenvalues}
\label{sec:grafting}
\input{contents/03_learning_rate_grafting}

\section{Controlling the approximation error induced by stale eigenbases}
\label{sec:adaptive-eigenbases}
\input{contents/05_adaptive_eigenbasis_computation_frequency}

\section{Discussion and conclusion}
\label{sec:conclusion}
\input{contents/06_conclusion}

\section*{Acknowledgments}
\input{contents/acknowledgements}

\bibliography{references}
\bibliographystyle{iclr_bibstyle}


\newpage
\appendix
\input{contents/07_appendix}



\end{document}

\typeout{get arXiv to do 4 passes: Label(s) may have changed. Rerun}


\usepackage{amsmath,amsfonts,bm}


\newcommand{\vectorized}{\mathrm{vec}}
\newcommand{\mat}{\mathrm{mat}}
\newcommand{\reshape}{\mathrm{reshape}}
\newcommand{\diag}{\mathrm{diag}}
\newcommand{\tr}{\mathrm{tr}}

\usepackage{amssymb}
\makeatletter
\newcommand*{\transpose}{\bgroup\@ifstar{\mathpalette\@transpose{\mkern-3.5mu}\egroup}{\mathpalette\@transpose\relax\egroup}}
\newcommand*{\@transpose}[2]{\setbox0=\hbox{\m@th$#1#2\intercal$}\raise\dp0\box0}
\makeatother


\newcommand{\figleft}{{\em (Left)}}
\newcommand{\figcenter}{{\em (Center)}}
\newcommand{\figright}{{\em (Right)}}
\newcommand{\figtop}{{\em (Top)}}
\newcommand{\figbottom}{{\em (Bottom)}}
\newcommand{\captiona}{{\em (a)}}
\newcommand{\captionb}{{\em (b)}}
\newcommand{\captionc}{{\em (c)}}
\newcommand{\captiond}{{\em (d)}}

\newcommand{\newterm}[1]{{\bf #1}}


\def\figref#1{figure~\ref{#1}}
\def\Figref#1{Figure~\ref{#1}}
\def\twofigref#1#2{figures \ref{#1} and \ref{#2}}
\def\quadfigref#1#2#3#4{figures \ref{#1}, \ref{#2}, \ref{#3} and \ref{#4}}
\def\secref#1{section~\ref{#1}}
\def\Secref#1{Section~\ref{#1}}
\def\twosecrefs#1#2{sections \ref{#1} and \ref{#2}}
\def\secrefs#1#2#3{sections \ref{#1}, \ref{#2} and \ref{#3}}
\def\eqref#1{equation~\ref{#1}}
\def\Eqref#1{Equation~\ref{#1}}
\def\plaineqref#1{\ref{#1}}
\def\chapref#1{chapter~\ref{#1}}
\def\Chapref#1{Chapter~\ref{#1}}
\def\rangechapref#1#2{chapters\ref{#1}--\ref{#2}}
\def\algref#1{algorithm~\ref{#1}}
\def\Algref#1{Algorithm~\ref{#1}}
\def\twoalgref#1#2{algorithms \ref{#1} and \ref{#2}}
\def\Twoalgref#1#2{Algorithms \ref{#1} and \ref{#2}}
\def\partref#1{part~\ref{#1}}
\def\Partref#1{Part~\ref{#1}}
\def\twopartref#1#2{parts \ref{#1} and \ref{#2}}

\def\ceil#1{\lceil #1 \rceil}
\def\floor#1{\lfloor #1 \rfloor}
\def\1{\bm{1}}
\newcommand{\train}{\mathcal{D}}
\newcommand{\valid}{\mathcal{D_{\mathrm{valid}}}}
\newcommand{\test}{\mathcal{D_{\mathrm{test}}}}

\def\eps{{\epsilon}}


\def\reta{{\textnormal{$\eta$}}}
\def\ra{{\textnormal{a}}}
\def\rb{{\textnormal{b}}}
\def\rc{{\textnormal{c}}}
\def\rd{{\textnormal{d}}}
\def\re{{\textnormal{e}}}
\def\rf{{\textnormal{f}}}
\def\rg{{\textnormal{g}}}
\def\rh{{\textnormal{h}}}
\def\ri{{\textnormal{i}}}
\def\rj{{\textnormal{j}}}
\def\rk{{\textnormal{k}}}
\def\rl{{\textnormal{l}}}
\def\rn{{\textnormal{n}}}
\def\ro{{\textnormal{o}}}
\def\rp{{\textnormal{p}}}
\def\rq{{\textnormal{q}}}
\def\rr{{\textnormal{r}}}
\def\rs{{\textnormal{s}}}
\def\rt{{\textnormal{t}}}
\def\ru{{\textnormal{u}}}
\def\rv{{\textnormal{v}}}
\def\rw{{\textnormal{w}}}
\def\rx{{\textnormal{x}}}
\def\ry{{\textnormal{y}}}
\def\rz{{\textnormal{z}}}

\def\rvepsilon{{\mathbf{\epsilon}}}
\def\rvtheta{{\mathbf{\theta}}}
\def\rva{{\mathbf{a}}}
\def\rvb{{\mathbf{b}}}
\def\rvc{{\mathbf{c}}}
\def\rvd{{\mathbf{d}}}
\def\rve{{\mathbf{e}}}
\def\rvf{{\mathbf{f}}}
\def\rvg{{\mathbf{g}}}
\def\rvh{{\mathbf{h}}}
\def\rvu{{\mathbf{i}}}
\def\rvj{{\mathbf{j}}}
\def\rvk{{\mathbf{k}}}
\def\rvl{{\mathbf{l}}}
\def\rvm{{\mathbf{m}}}
\def\rvn{{\mathbf{n}}}
\def\rvo{{\mathbf{o}}}
\def\rvp{{\mathbf{p}}}
\def\rvq{{\mathbf{q}}}
\def\rvr{{\mathbf{r}}}
\def\rvs{{\mathbf{s}}}
\def\rvt{{\mathbf{t}}}
\def\rvu{{\mathbf{u}}}
\def\rvv{{\mathbf{v}}}
\def\rvw{{\mathbf{w}}}
\def\rvx{{\mathbf{x}}}
\def\rvy{{\mathbf{y}}}
\def\rvz{{\mathbf{z}}}

\def\erva{{\textnormal{a}}}
\def\ervb{{\textnormal{b}}}
\def\ervc{{\textnormal{c}}}
\def\ervd{{\textnormal{d}}}
\def\erve{{\textnormal{e}}}
\def\ervf{{\textnormal{f}}}
\def\ervg{{\textnormal{g}}}
\def\ervh{{\textnormal{h}}}
\def\ervi{{\textnormal{i}}}
\def\ervj{{\textnormal{j}}}
\def\ervk{{\textnormal{k}}}
\def\ervl{{\textnormal{l}}}
\def\ervm{{\textnormal{m}}}
\def\ervn{{\textnormal{n}}}
\def\ervo{{\textnormal{o}}}
\def\ervp{{\textnormal{p}}}
\def\ervq{{\textnormal{q}}}
\def\ervr{{\textnormal{r}}}
\def\ervs{{\textnormal{s}}}
\def\ervt{{\textnormal{t}}}
\def\ervu{{\textnormal{u}}}
\def\ervv{{\textnormal{v}}}
\def\ervw{{\textnormal{w}}}
\def\ervx{{\textnormal{x}}}
\def\ervy{{\textnormal{y}}}
\def\ervz{{\textnormal{z}}}

\def\rmA{{\mathbf{A}}}
\def\rmB{{\mathbf{B}}}
\def\rmC{{\mathbf{C}}}
\def\rmD{{\mathbf{D}}}
\def\rmE{{\mathbf{E}}}
\def\rmF{{\mathbf{F}}}
\def\rmG{{\mathbf{G}}}
\def\rmH{{\mathbf{H}}}
\def\rmI{{\mathbf{I}}}
\def\rmJ{{\mathbf{J}}}
\def\rmK{{\mathbf{K}}}
\def\rmL{{\mathbf{L}}}
\def\rmM{{\mathbf{M}}}
\def\rmN{{\mathbf{N}}}
\def\rmO{{\mathbf{O}}}
\def\rmP{{\mathbf{P}}}
\def\rmQ{{\mathbf{Q}}}
\def\rmR{{\mathbf{R}}}
\def\rmS{{\mathbf{S}}}
\def\rmT{{\mathbf{T}}}
\def\rmU{{\mathbf{U}}}
\def\rmV{{\mathbf{V}}}
\def\rmW{{\mathbf{W}}}
\def\rmX{{\mathbf{X}}}
\def\rmY{{\mathbf{Y}}}
\def\rmZ{{\mathbf{Z}}}

\def\ermA{{\textnormal{A}}}
\def\ermB{{\textnormal{B}}}
\def\ermC{{\textnormal{C}}}
\def\ermD{{\textnormal{D}}}
\def\ermE{{\textnormal{E}}}
\def\ermF{{\textnormal{F}}}
\def\ermG{{\textnormal{G}}}
\def\ermH{{\textnormal{H}}}
\def\ermI{{\textnormal{I}}}
\def\ermJ{{\textnormal{J}}}
\def\ermK{{\textnormal{K}}}
\def\ermL{{\textnormal{L}}}
\def\ermM{{\textnormal{M}}}
\def\ermN{{\textnormal{N}}}
\def\ermO{{\textnormal{O}}}
\def\ermP{{\textnormal{P}}}
\def\ermQ{{\textnormal{Q}}}
\def\ermR{{\textnormal{R}}}
\def\ermS{{\textnormal{S}}}
\def\ermT{{\textnormal{T}}}
\def\ermU{{\textnormal{U}}}
\def\ermV{{\textnormal{V}}}
\def\ermW{{\textnormal{W}}}
\def\ermX{{\textnormal{X}}}
\def\ermY{{\textnormal{Y}}}
\def\ermZ{{\textnormal{Z}}}

\def\vzero{{\bm{0}}}
\def\vone{{\bm{1}}}
\def\vmu{{\bm{\mu}}}
\def\vtheta{{\bm{\theta}}}
\def\va{{\bm{a}}}
\def\vb{{\bm{b}}}
\def\vc{{\bm{c}}}
\def\vd{{\bm{d}}}
\def\ve{{\bm{e}}}
\def\vf{{\bm{f}}}
\def\vg{{\bm{g}}}
\def\vh{{\bm{h}}}
\def\vi{{\bm{i}}}
\def\vj{{\bm{j}}}
\def\vk{{\bm{k}}}
\def\vl{{\bm{l}}}
\def\vm{{\bm{m}}}
\def\vn{{\bm{n}}}
\def\vo{{\bm{o}}}
\def\vp{{\bm{p}}}
\def\vq{{\bm{q}}}
\def\vr{{\bm{r}}}
\def\vs{{\bm{s}}}
\def\vt{{\bm{t}}}
\def\vu{{\bm{u}}}
\def\vv{{\bm{v}}}
\def\vw{{\bm{w}}}
\def\vx{{\bm{x}}}
\def\vy{{\bm{y}}}
\def\vz{{\bm{z}}}

\def\evalpha{{\alpha}}
\def\evbeta{{\beta}}
\def\evepsilon{{\epsilon}}
\def\evlambda{{\lambda}}
\def\evomega{{\omega}}
\def\evmu{{\mu}}
\def\evpsi{{\psi}}
\def\evsigma{{\sigma}}
\def\evtheta{{\theta}}
\def\eva{{a}}
\def\evb{{b}}
\def\evc{{c}}
\def\evd{{d}}
\def\eve{{e}}
\def\evf{{f}}
\def\evg{{g}}
\def\evh{{h}}
\def\evi{{i}}
\def\evj{{j}}
\def\evk{{k}}
\def\evl{{l}}
\def\evm{{m}}
\def\evn{{n}}
\def\evo{{o}}
\def\evp{{p}}
\def\evq{{q}}
\def\evr{{r}}
\def\evs{{s}}
\def\evt{{t}}
\def\evu{{u}}
\def\evv{{v}}
\def\evw{{w}}
\def\evx{{x}}
\def\evy{{y}}
\def\evz{{z}}

\def\mA{{\bm{A}}}
\def\mB{{\bm{B}}}
\def\mC{{\bm{C}}}
\def\mD{{\bm{D}}}
\def\mE{{\bm{E}}}
\def\mF{{\bm{F}}}
\def\mG{{\bm{G}}}
\def\mH{{\bm{H}}}
\def\mI{{\bm{I}}}
\def\mJ{{\bm{J}}}
\def\mK{{\bm{K}}}
\def\mL{{\bm{L}}}
\def\mM{{\bm{M}}}
\def\mN{{\bm{N}}}
\def\mO{{\bm{O}}}
\def\mP{{\bm{P}}}
\def\mQ{{\bm{Q}}}
\def\mR{{\bm{R}}}
\def\mS{{\bm{S}}}
\def\mT{{\bm{T}}}
\def\mU{{\bm{U}}}
\def\mV{{\bm{V}}}
\def\mW{{\bm{W}}}
\def\mX{{\bm{X}}}
\def\mY{{\bm{Y}}}
\def\mZ{{\bm{Z}}}
\def\mBeta{{\bm{\beta}}}
\def\mPhi{{\bm{\Phi}}}
\def\mLambda{{\bm{\Lambda}}}
\def\mSigma{{\bm{\Sigma}}}

\DeclareMathAlphabet{\mathsfit}{\encodingdefault}{\sfdefault}{m}{sl}
\SetMathAlphabet{\mathsfit}{bold}{\encodingdefault}{\sfdefault}{bx}{n}
\newcommand{\tens}[1]{\bm{\mathsfit{#1}}}
\def\tA{{\tens{A}}}
\def\tB{{\tens{B}}}
\def\tC{{\tens{C}}}
\def\tD{{\tens{D}}}
\def\tE{{\tens{E}}}
\def\tF{{\tens{F}}}
\def\tG{{\tens{G}}}
\def\tH{{\tens{H}}}
\def\tI{{\tens{I}}}
\def\tJ{{\tens{J}}}
\def\tK{{\tens{K}}}
\def\tL{{\tens{L}}}
\def\tM{{\tens{M}}}
\def\tN{{\tens{N}}}
\def\tO{{\tens{O}}}
\def\tP{{\tens{P}}}
\def\tQ{{\tens{Q}}}
\def\tR{{\tens{R}}}
\def\tS{{\tens{S}}}
\def\tT{{\tens{T}}}
\def\tU{{\tens{U}}}
\def\tV{{\tens{V}}}
\def\tW{{\tens{W}}}
\def\tX{{\tens{X}}}
\def\tY{{\tens{Y}}}
\def\tZ{{\tens{Z}}}


\def\gA{{\mathcal{A}}}
\def\gB{{\mathcal{B}}}
\def\gC{{\mathcal{C}}}
\def\gD{{\mathcal{D}}}
\def\gE{{\mathcal{E}}}
\def\gF{{\mathcal{F}}}
\def\gG{{\mathcal{G}}}
\def\gH{{\mathcal{H}}}
\def\gI{{\mathcal{I}}}
\def\gJ{{\mathcal{J}}}
\def\gK{{\mathcal{K}}}
\def\gL{{\mathcal{L}}}
\def\gM{{\mathcal{M}}}
\def\gN{{\mathcal{N}}}
\def\gO{{\mathcal{O}}}
\def\gP{{\mathcal{P}}}
\def\gQ{{\mathcal{Q}}}
\def\gR{{\mathcal{R}}}
\def\gS{{\mathcal{S}}}
\def\gT{{\mathcal{T}}}
\def\gU{{\mathcal{U}}}
\def\gV{{\mathcal{V}}}
\def\gW{{\mathcal{W}}}
\def\gX{{\mathcal{X}}}
\def\gY{{\mathcal{Y}}}
\def\gZ{{\mathcal{Z}}}

\def\sA{{\mathbb{A}}}
\def\sB{{\mathbb{B}}}
\def\sC{{\mathbb{C}}}
\def\sD{{\mathbb{D}}}
\def\sF{{\mathbb{F}}}
\def\sG{{\mathbb{G}}}
\def\sH{{\mathbb{H}}}
\def\sI{{\mathbb{I}}}
\def\sJ{{\mathbb{J}}}
\def\sK{{\mathbb{K}}}
\def\sL{{\mathbb{L}}}
\def\sM{{\mathbb{M}}}
\def\sN{{\mathbb{N}}}
\def\sO{{\mathbb{O}}}
\def\sP{{\mathbb{P}}}
\def\sQ{{\mathbb{Q}}}
\def\sR{{\mathbb{R}}}
\def\sS{{\mathbb{S}}}
\def\sT{{\mathbb{T}}}
\def\sU{{\mathbb{U}}}
\def\sV{{\mathbb{V}}}
\def\sW{{\mathbb{W}}}
\def\sX{{\mathbb{X}}}
\def\sY{{\mathbb{Y}}}
\def\sZ{{\mathbb{Z}}}

\def\emLambda{{\Lambda}}
\def\emA{{A}}
\def\emB{{B}}
\def\emC{{C}}
\def\emD{{D}}
\def\emE{{E}}
\def\emF{{F}}
\def\emG{{G}}
\def\emH{{H}}
\def\emI{{I}}
\def\emJ{{J}}
\def\emK{{K}}
\def\emL{{L}}
\def\emM{{M}}
\def\emN{{N}}
\def\emO{{O}}
\def\emP{{P}}
\def\emQ{{Q}}
\def\emR{{R}}
\def\emS{{S}}
\def\emT{{T}}
\def\emU{{U}}
\def\emV{{V}}
\def\emW{{W}}
\def\emX{{X}}
\def\emY{{Y}}
\def\emZ{{Z}}
\def\emSigma{{\Sigma}}

\newcommand{\etens}[1]{\mathsfit{#1}}
\def\etLambda{{\etens{\Lambda}}}
\def\etA{{\etens{A}}}
\def\etB{{\etens{B}}}
\def\etC{{\etens{C}}}
\def\etD{{\etens{D}}}
\def\etE{{\etens{E}}}
\def\etF{{\etens{F}}}
\def\etG{{\etens{G}}}
\def\etH{{\etens{H}}}
\def\etI{{\etens{I}}}
\def\etJ{{\etens{J}}}
\def\etK{{\etens{K}}}
\def\etL{{\etens{L}}}
\def\etM{{\etens{M}}}
\def\etN{{\etens{N}}}
\def\etO{{\etens{O}}}
\def\etP{{\etens{P}}}
\def\etQ{{\etens{Q}}}
\def\etR{{\etens{R}}}
\def\etS{{\etens{S}}}
\def\etT{{\etens{T}}}
\def\etU{{\etens{U}}}
\def\etV{{\etens{V}}}
\def\etW{{\etens{W}}}
\def\etX{{\etens{X}}}
\def\etY{{\etens{Y}}}
\def\etZ{{\etens{Z}}}

\newcommand{\pdata}{p_{\rm{data}}}
\newcommand{\ptrain}{\hat{p}_{\rm{data}}}
\newcommand{\Ptrain}{\hat{P}_{\rm{data}}}
\newcommand{\pmodel}{p_{\rm{model}}}
\newcommand{\Pmodel}{P_{\rm{model}}}
\newcommand{\ptildemodel}{\tilde{p}_{\rm{model}}}
\newcommand{\pencode}{p_{\rm{encoder}}}
\newcommand{\pdecode}{p_{\rm{decoder}}}
\newcommand{\precons}{p_{\rm{reconstruct}}}

\newcommand{\laplace}{\mathrm{Laplace}} %

\newcommand{\E}{\mathbb{E}}
\newcommand{\Ls}{\mathcal{L}}
\newcommand{\R}{\mathbb{R}}
\newcommand{\emp}{\tilde{p}}
\newcommand{\lr}{\alpha}
\newcommand{\reg}{\lambda}
\newcommand{\rect}{\mathrm{rectifier}}
\newcommand{\softmax}{\mathrm{softmax}}
\newcommand{\sigmoid}{\sigma}
\newcommand{\softplus}{\zeta}
\newcommand{\KL}{D_{\mathrm{KL}}}
\newcommand{\Var}{\mathrm{Var}}
\newcommand{\standarderror}{\mathrm{SE}}
\newcommand{\Cov}{\mathrm{Cov}}
\newcommand{\normlzero}{L^0}
\newcommand{\normlone}{L^1}
\newcommand{\normltwo}{L^2}
\newcommand{\normlp}{L^p}
\newcommand{\normmax}{L^\infty}

\newcommand{\parents}{Pa} %

\DeclareMathOperator*{\argmax}{arg\,max}
\DeclareMathOperator*{\argmin}{arg\,min}

\DeclareMathOperator{\sign}{sign}
\DeclareMathOperator{\Tr}{Tr}
\let\ab\allowbreak

Structured non-diagonal, and especially Kronecker-factored, preconditioned stochastic gradient algorithms for training neural networks have been extensively studied in the machine learning community \citep{heskes2000natural,martens2010deep,martens2015optimizing,li2018preconditioned,gupta2018shampoo}.
Despite promising results, diagonal methods like Adam and its variants have remained the de facto methods for training neural networks over the past decade \citep{duchi2011adaptive,kingma2014adam}.
Recently, a distributed implementation of the Shampoo algorithm \citep{gupta2018shampoo,anil2020scalable,shi2023distributed} won the external tuning track of the AlgoPerf neural network training algorithm competition \citep{dahl2023benchmarking,kasimbeg2025accelerating}.\footnote{The \textit{external tuning} track requires a submission to specify a hyperparameter search space, in contrast to the hyperparameter-free \textit{self-tuning} track.}
This has spurred renewed interest in non-diagonal neural network training algorithms, including new methods like Muon \citep{jordan2024muon,bernstein2025deriving}
and SOAP \citep{vyas2025soap}.

The winning Shampoo submission to the AlgoPerf competition depended on multiple crucial heuristics on top of the Shampoo algorithm \citep{anil2020scalable,shi2023distributed}.
In particular, the update of each layer is re-scaled by the update magnitude of a reference optimizer, also known as \textit{learning rate grafting} \citep{agarwal2020disentangling}; see \Cref{fig:stale-instability} for an illustration of the importance of the heuristic. 
Additionally, to amortize the computational overhead, the root-inverse of the preconditioner is only periodically re-computed every $100$ steps, resulting in the use of a \textit{stale} preconditioner.
However, both heuristics lack theoretical justification and have not been thoroughly understood.

In this paper, we investigate the role of these two heuristics by decoupling the updates of the preconditioner's eigenvalues and eigenbasis. 
In \Cref{sec:grafting}, we empirically demonstrate that Shampoo requires learning rate grafting in order to address the staleness and mis-scaling of the preconditioner's \textit{eigenvalues}. Correcting Shampoo's eigenvalues at every step, \textit{à la} SOAP \citep{vyas2025soap}, removes the need for grafting.
We make this intuition more mathematically precise by comparing bounds on the magnitude of the Shampoo and full-matrix Adam updates.
Given the importance of controlling the approximation error of the eigenvalues, we turn to the choice of the \textit{eigenbasis} update frequency in \Cref{sec:adaptive-eigenbases}.
Motivated by a termination criterion for a QR algorithm that bounds the relative error of the Kronecker factor approximation induced by the current eigenbasis \citep{golub2013matrix}, we show how to adaptively determine the eigenbasis update frequency.
We empirically demonstrate that the approximation error of the Kronecker factors evolves and impacts convergence depending on the stage of training and the properties of the parameters.
Using an adaptive criterion may increase training efficiency in cases where more frequent eigenbasis computations notably improve convergence in steps.

\begin{figure}[t]
    \centering
    \includegraphics[width=\textwidth]{figures/stale_shampoo_instability.pdf}
    \caption{Shampoo with stale preconditioner (updating the root inverse matrices every $F=100$ steps) without grafting for different choices of the learning rate $\alpha$ and $\epsilon$ on Imagewoof. All tested hyperparameter combinations underperform AdamW and, by extension, Shampoo with Adam grafting.}
    \label{fig:stale-instability}
\end{figure}

We consider neural network training as a standard stochastic optimization problem, where we seek to minimize the expected loss function
\begin{equation}
\min_{\vtheta \in \R^d} \gL(\vtheta) = \E_{(\vx, \vy) \sim p_\gD(\vx, \vy)} \big[ \ell(f_\vtheta(\vx), \vy)  \big],
\end{equation}
where $f_\vtheta: \R^N \rightarrow{\R^c}$ is the neural network prediction function with parameters $\vtheta \in \R^d$ flattened and concatenated into a vector, $p_\gD(\vx, \vy)$ is the joint data distribution from which inputs $\vx \in \R^N$ and targets $\vy \in \R^c$ are sampled, and $\ell: \R^c \times \R^c \rightarrow{\R}$ is the loss function. The neural network is commonly trained using preconditioned stochastic gradient methods that update the parameters at each iteration $t$ by 
\begin{equation}
\label{eq:PGD}
    \vtheta_t = \vtheta_{t-1} - \alpha_t \mC_t^{-p} \vg_t = \vtheta_{t - 1} - \alpha_t \mQ_{\mC_t} \mLambda_{\mC_t}^{-p} \mQ_{\mC_t}^\transpose \vg_t,
\end{equation}
where $\vg_t = \nabla_{\vtheta} \ell(f_{\vtheta_t}(\vx_t), \vy_t) \in \R^d$ is the stochastic (mini-batch) gradient with respect to sample $(\vx_t, \vy_t) \sim p_\gD(\vx, \vy)$, $\alpha_t > 0$ is the step size, and $\mC_t \in \R^{d \times d}$ is a symmetric positive-definite \textit{preconditioner} or \textit{scaling} matrix.
The second equality in \Cref{eq:PGD} expresses the update using the eigendecomposition of the preconditioner matrix $\mC_t = \mQ_{\mC_t} \mLambda_{\mC_t} \mQ_{\mC_t}^\transpose$, where the eigenbasis matrix $\mQ_t \in \R^{d \times d}$ is orthogonal and eigenvalue matrix $\mLambda_{\mC_t} \in \R^{d \times d}$ is diagonal.

We focus on $\mC_t$ as an accumulation of gradient outer products, i.e., $\bar{\mA}_t = \sum_{s = 1}^t \vg_s \vg_s^\transpose$  or $\mA_t = \beta_2 \mA_{t-1} + (1 - \beta_2) \vg_s \vg_s^\transpose$, which correspond to full-matrix Adagrad and RMSprop/Adam, respectively \citep{duchi2011adaptive,kingma2014adam}.
We can also consider alternatives like the Fisher information matrix \citep{amari1998natural}; see \Cref{sec:fisher-connection}.
When $\mC_t$ is diagonal, this scaling recovers simplified versions of popular algorithms like AdaGrad \citep{duchi2011adaptive} or Adam \citep{kingma2014adam}, ignoring additional modifications like exponential moving averages, momentum, bias corrections, etc. However, when $\mC_t$ is non-diagonal, these preconditioned methods require storing and possibly inverting dense $d \times d$ matrices, which is prohibitive in practice.

To mitigate this high cost, two structured approximations are commonly applied: (1) a layer-wise block-diagonal approximation, where each block captures pairwise correlations within each layer; (2) a Kronecker product approximation that takes a $m n \times mn$ matrix block for each layer and re-writes it as a Kronecker product of two smaller matrices of size $m \times m$ and $n \times n$, and has convenient properties for computing the inverse-root matrix-vector product. Notable examples include Kronecker-Factored Approximate Curvature (K-FAC) \citep{heskes2000natural,martens2015optimizing,grosse2016kronecker,martens2018kroneckerfactored,eschenhagen2023kfac}, Shampoo \citep{gupta2018shampoo,anil2020scalable,shi2023distributed}, and their variants \citep{george2018ekfac,gao2020trace,ren2021tensor,feinberg2023sketchy,duvvuri2024caspr,lin2024structured,lin2024remove}. Alternative approaches, such as Hessian-free or inexact Newton methods, avoid explicitly storing the preconditioner matrix, \eg, by leveraging Hessian-vector products \citep{martens2010deep,li2018preconditioned,pooladzandi2024curvature}.

\subsection{Shampoo}
\label{sec:shampoo}

The Shampoo preconditioner was originally developed as a Kronecker-factorized upper bound to full-matrix AdaGrad in its regret analysis \citep{gupta2018shampoo}. 
For a weight matrix $\mW_t \in \R^{m \times n}$ with stochastic gradient $\mG_t \in \R^{m \times n}$ where $\vg_t = \vectorized(\mG_t)$, Shampoo stores symmetric positive semi-definite matrices that approximate an \textit{idealized} preconditioner with $\mL_t \approx \E[\mG_t \mG_t^\transpose] \in \R^{m \times m}$ and $\mR_t \approx \E[\mG_t^\transpose \mG_t] \in \R^{n \times n}$, and updates the preconditioner and weight matrix by 
\begin{align}
\label{eq:shampoo}
    \mL_t & = \beta_2 \mL_{t-1} + (1 - \beta_2) \mG_t \mG_t^\transpose, \\
    \mR_t & = \beta_2 \mR_{t-1} + (1 - \beta_2) \mG_t^\transpose \mG_t, \\
    \mW_t & = \mW_{t-1} - \alpha_t \mL_t^{-\frac{1}{4}} \mG_t \mR_t^{-\frac{1}{4}}
\end{align}
given $\beta_2 \in [0, 1)$.\footnote{We drop the subscript in the expectation, taken with respect to $(\vx_t, \vy_t) \sim p_\gD(\vx, \vy)$.}\footnote{A pseudo-inverse that ignores the null space of the matrix or perturbing the matrix by a regularization term $\epsilon I$ for $\epsilon > 0$ can be used to handle the symmetric positive semi-definite case.} Note that the original Shampoo algorithm uses a sum to accumulate over the gradient outer products, although the exponential moving average form as presented here is more commonly used in practice. Therefore, we are interested in approximating the preconditioner of full-matrix Adam given by $\mA_t = \beta_2 \mA_{t - 1} + (1 - \beta_2) \vg_t \vg_t^\transpose \approx \E[\vg_t \vg_t^\transpose]$. 

The search direction or update in matrix form is given as $\mU_t^\mathrm{Shampoo} = -\mL_t^{-\frac{1}{4}} \mG_t \mR^{-\frac{1}{4}}$. Equivalently, we can set $\mC_t^\mathrm{Shampoo} = (\mR_t \otimes \mL_t)^{\frac{1}{2}}$ and $p=1/2$ in \Cref{eq:PGD}.
This update rule can also be generalized to higher-order weight tensors.
We will also consider \textit{Shampoo$^2$} defined as $\mC_t^\mathrm{Shampoo^2} = \mR_t \otimes \mL_t$, which is a tighter approximation to full-matrix Adagrad \citep{morwani2025new}.


\subsection{Eigenvalue-corrected Shampoo and SOAP}
\label{sec:eigenvalue-corrected}

By leveraging a Kronecker product approximation, Shampoo's preconditioner is restricted to a Kronecker product structure for both its eigenvectors and eigenvalues. Specifically, if we have a matrix $\mC_t = \mR_t \otimes \mL_t$ with eigendecompositions $\mL_t = \mQ_{\mL_t} \mLambda_{\mL_t} \mQ_{\mL_t}^\transpose$ and $\mR_t = \mQ_{\mR_t} \mLambda_{\mR_t} \mQ_{\mR_t}^\transpose$, then $\mC_t$ has eigendecomposition $\mC_t = (\mQ_{\mR_t} \otimes \mQ_{\mL_t}) (\mLambda_{\mR_t} \otimes \mLambda_{\mL_t}) (\mQ_{\mR_t} \otimes \mQ_{\mL_t})^\transpose$. Preconditioning $-\vg_t$ with $\mC_t^{-p}$ in its matrix form is equivalent to the matrix transformation:
\begin{equation}
    \label{eq:kronecker-factored-update}
    \mU_t^\mathrm{Shampoo} = -\mL_t^{-\frac{1}{4}} \mG_t \mR_t^{-\frac{1}{4}} = -\mQ_{\mL_t} \mLambda_{\mL_t}^{-\frac{1}{4}} (\mQ_{\mL_t}^\transpose \mG_t \mQ_{\mR_t}) \mLambda_{\mR_t}^{-\frac{1}{4}} \mQ_{\mR_t}^\transpose.
\end{equation}
This can be interpreted as applying a Kronecker-product-based coordinate-wise scaling to a gradient with changed basis, \ie, $\tilde{\mG}_t = \mQ_{\mL_t}^\transpose \mG_t \mQ_{\mR_t}$. After scaling, the transformed gradient is converted back to its original basis. 

Instead of restricting the eigenvalues to be a Kronecker product, \citet{liu2018rotate} and \citet{george2018ekfac} have proposed \textit{correcting} the eigenvalues by decoupling the scaling from the basis in K-FAC. This involves accumulating a separate scaling matrix $\mD_t \approx \E\big[\tilde{\mG}_t^{\odot 2} \big] \in \R^{m \times n}$ and using it in lieu of $\mLambda_{\mR_t} \otimes \mLambda_{\mL_t}$.\footnote{$\mX^{\odot 2}$ denotes the element-wise square.}
\citet{anil2020scalable} remarked that the same correction can also be applied to Shampoo.
More recently, \citet{vyas2025soap} presented promising empirical results for language models using an instance of eigenvalue-corrected Shampoo (EShampoo) called SOAP, which updates the scaling by $\mD_t = \beta_2 \mD_{t - 1} + (1 - \beta_2) \tilde{\mG}_t^{\odot 2}$.
Since then, SOAP has also been shown to perform well for physics-informed neural networks \citep{wang2025gradient} and to reduce outlier features in transformers, which potentially improves quantization \citep{he2024understanding}. See \Cref{sec:optimal}
for a discussion of the optimal choice of $\mD_t$ in Frobenius norm and how it relates to SOAP.


Originally motivated by decoupling the update magnitude from its direction to account for different implicit learning rate schedules induced by different optimizers, \citet{agarwal2020disentangling} proposed a technique called \textit{learning rate grafting} that combines the layer-wise update of an optimizer with the layer-wise update magnitude from a different optimizer. More precisely in Shampoo, given the Shampoo update $\mU_t^\mathrm{Shampoo}$ and the grafting method update $\mU_t^\mathrm{Grafting}$ (typically Adam) for a single layer with weight $\mW_t$, the Shampoo per-layer update is re-scaled as
\begin{equation}
\label{eq:grafting}
    \mW_{t  + 1} = \mW_t + \alpha_t \frac{||\mU_t^\mathrm{Grafting}||_F}{||\mU_t^\mathrm{Shampoo}||_F} \; \mU_t^\mathrm{Shampoo}.
\end{equation}
A complete description of Shampoo with Adam grafting is given in \Cref{app:algorithms-pseudocode}, \Cref{alg:shampoo-adam-grafting}.

As mentioned in \Cref{sec:shampoo}, \citet{shi2023distributed} has claimed that learning rate grafting from a base optimizer like Adam has been crucial to Shampoo's performance; see also \citet{anil2020scalable,kasimbeg2025accelerating}.
We illustrate this in \Cref{fig:stale-instability} by showing that independent of the hyperparameter setting, Shampoo without grafting (and updating the root inverse matrices every $F = 100$ steps) underperforms AdamW, with many hyperparameter settings diverging.
\citet{anil2020scalable} claims that grafting is used to account for differences in magnitude of the eigenspectrum of the Kronecker factors for different layers and the infrequent updates of the Kronecker factors and their inverse roots (see \citet{anil2020scalable}, Appendix G).
Despite its practical necessity in Shampoo, grafting's algorithmic role has not been thoroughly investigated.

Here, we focus on the case of Adam grafting, which was used in conjunction with Shampoo ($F=100$). This configuration was used in the winning submission to the external tuning track of the AlgoPerf competition \citep{kasimbeg2025accelerating}.
Since grafting re-scales the update using the Frobenius norm, we compare the Frobenius norm of the updates of full-matrix and diagonal Adam, Shampoo, and EShampoo.
Note that the magnitude of the eigendecomposed update in \Cref{eq:PGD} with Kronecker-factored $\mQ_{\mC_t} = \mQ_{\mR_t} \otimes \mQ_{\mL_t}$ is determined by the norm of the stochastic gradient $\mG_t$ as well as the eigenvalues of $\mC_t$, which is formalized in the following lemma.
\begin{lemma}
    \label{lem:bound-kronecker-update}
    Let $\mU = \mQ_{\mL} (\mD^{\odot-p} \odot (\mQ_{\mL}^\transpose \mG \mQ_{\mR})) \mQ_{\mR}^\transpose \in \R^{m \times n}$ be the generalized eigendecomposed Kronecker-factored update given by orthogonal matrices $\mQ_{\mL} \in \R^{m \times m}$, $\mQ_{\mR} \in \R^{n \times n}$, and dense scaling matrix $\mD \in \R^{m \times n}$. Then we have:
    \begin{equation}
        (\max_{i, j} \mD_{i, j})^{-p} || \mG ||_F \leq || \mU ||_F \leq (\min_{i, j} \mD_{i, j})^{-p} || \mG ||_F.
    \end{equation}
\end{lemma}
\Cref{lem:bound-kronecker-update} generalizes multiple algorithmic cases. Specifically, if $\mQ_{\mL} = \mI \in \R^{m \times m}$, $\mQ_{\mR} = \mI \in \R^{n \times n}$, and $\mD = \E\big[\mG^{\odot 2}]$ with $p = 1/2$, $\mU$ gives the idealized Adam update. If $\mD = \mLambda_\mR \otimes \mLambda_\mL$ with $\mL = \E[\mG \mG^\transpose] = \mQ_{\mL} \mLambda_{\mL} \mQ_{\mL}^\transpose$ and $\mR = \E[\mG^\transpose \mG] = \mQ_{\mR} \mLambda_{\mR} \mQ_{\mR}^\transpose$, then $\mU$ corresponds to the idealized Shampoo update. Selecting $\mD = \E\big[(\mQ_L^\transpose \mG \mQ_R)^{\odot 2}\big]$ instead yields the idealized EShampoo update.
\begin{proposition}
    \label{prop:magnitude}
    Assume that $\E[\vg\vg^\transpose]$ is symmetric positive definite. The magnitude of the idealized updates for full-matrix Adam, diagonal Adam, and EShampoo are all bounded by the power of the extreme eigenvalues of full-matrix Adam: 
    \begin{equation}
        \label{eq:update-magnitude-bound}
        \lambda_{\max}(\E[\vg \vg^\transpose])^{-p} \|\mG\|_F \leq \|\mU\|_F \leq \lambda_{\min}(\E[\vg \vg^\transpose])^{-p} \|\mG\|_F,
    \end{equation}
    for all $p > 0$. However, under the simplifying assumption that $\E[\mG] = \vzero$ and $\mG_{i, j}$ is independent from $\mG_{k, l}$ for $(i, j) \neq (k, l)$ and has bounded second moment, $\lambda_{\min}(\E [\vg \vg^\transpose]) \leq \E [\mG_{i, j}^2] \leq \lambda_{\max}(\E [\vg \vg^\transpose])$ and Shampoo has dimension-dependent bounds:
    \begin{equation}
        \label{eq:dimension-dependent-magnitude}
        m^{-p/2} n^{-p/2} \lambda_{\max}(\E[\vg \vg^\transpose])^{-p} \|\mG\|_F \leq \|\mU\|_F \leq m^{-p/2} n^{-p/2} \lambda_{\min}(\E[\vg \vg^\transpose])^{-p} \|\mG\|_F.
    \end{equation}
\end{proposition}
    

    
This highlights the first potential issue of Shampoo without grafting: the use of stale eigenvalues can result in update magnitudes that lie outside of the bounds of full-matrix Adam. Notably, the basis does not influence the update magnitude.
\Cref{prop:magnitude} also suggests that, even in simple settings, Shampoo's eigenspectrum is mis-scaled relative to Adam's and EShampoo's, due to the dependency on the dimensions of the preconditioner. While this additional scaling can be absorbed into the learning rate when only handling a single matrix, it is potentially problematic when one needs to handle multiple parameters simultaneously.

This leads us to the following hypothesis:
\begin{tcolorbox}[colframe=tue2, colback=tue2!3, title={The role of learning rate grafting in Shampoo},bottom=0mm,top=0mm,middle=0mm]
Learning rate grafting compensates for the staleness and scaling of Shampoo's eigenvalues.
\end{tcolorbox}
From the Frobenius norm approximation perspective, we know that $\mC_t^{\textrm{Shampoo}^2}$ is a tighter approximation to full-matrix Adam compared to $\mC_t^\textrm{Shampoo}$ \citep{morwani2025new}. This directly addresses the mismatch of the eigenvalues' exponent in \Cref{eq:dimension-dependent-magnitude}.
Additionally, we can rescale the preconditioner by $S^{-1} = \tr(\mR_t)^{-1} = \tr(\mL_t)^{-1}$, which ensures exactness in the case where full-matrix Adam $\mA_t$ is defined by a single Kronecker product.
The scaling $S^{-1}$ has been previously introduced for Tensor Normal Training \citep[TNT]{ren2021tensor} to approximate the Fisher and is also discussed in \citet{morwani2025new}. Naturally, Shampoo$^2$ with the trace scaling $S^{-1}$ will have the right scaling when full-matrix Adam can be decomposed into a single Kronecker product.


\begin{figure}
    \centering
    \includegraphics[width=\textwidth]{figures/grafting.pdf}
    \caption{Training results with different Shampoo variants and eigendecomposition frequencies $F$ on the Imagewoof dataset. Shampoo with eigenvalue correction achieves a better training loss compared to Shampoo with Adam grafting, and the optimal learning rate for Adam transfers to both variants.}
    \label{fig:lr-transfer}
\end{figure}

Based on our observations here, we can make several predictions:
\begin{enumerate}
    \item By updating Shampoo's preconditioner at every iteration, we can address the issue of stale eigenvalues. In this scenario, we anticipate that Shampoo without grafting will perform well when the scalings across different layers are relatively consistent, but may struggle when these scalings exhibit substantial discrepancies. This discrepancy is more likely to occur in deep neural networks with highly varying parameter shapes.
    \item By using Shampoo$^2$ with the trace scaling $S^{-1} \mC_t^{\textrm{Shampoo}^2}$ instead of Shampoo $\mC_t^\textrm{Shampoo}$ and updating at every iteration, we anticipate that the issues of staleness and scaling may be partially addressed, assuming that full-matrix Adam can nearly decompose into a single Kronecker product. In that case, Shampoo$^2$ could match Shampoo with grafting.
    \item An eigenvalue correction, as used in SOAP, that is updated every iteration also addresses the issues with staleness and scaling, and should match Shampoo with grafting.
\end{enumerate}





\textbf{Empirical validation.} To test our predictions, we ablate different settings of Shampoo on the Imagewoof dataset with vision transformer (ViT) and ConvNeXt V2 models. 
We plot the final training loss after $100$ epochs against the learning rate.
Following the specification of the AlgoPerf Shampoo submission, we update the preconditioner every $100$ steps when grafting the learning rate from Adam, and the eigenbasis when using the eigenvalue correction.
Our results are presented in Figure \ref{fig:lr-transfer}.
We observe that all three predictions are confirmed. Specifically, Shampoo$^2$ scaled by $S^{-1}$ or eigenvalue correction is able to match or even surpass Shampoo with grafting on both problems.
Shampoo without grafting can match the final loss for the Imagewoof ViT workload, but not for the ConvNeXt V2 model.

\textbf{Hyperparameter transfer.} The optimal learning rate of Adam transfers to Shampoo with grafting and EShampoo.
However, it is not completely clear whether the optimal learning rate transfers to Shampoo$^2$ scaled by $S^{-1}$. Learning rate transfer only makes sense when jointly transferring $\epsilon$ since it affects the effective step size, but Shampoo does not directly add $\epsilon$ to the root of the eigenvalues, as done in Adam and other diagonal-scaling-based optimizers.
While we have not tested this, the optimal learning rate may potentially transfer when matching the effective $\epsilon$.












\textbf{Computational and memory costs.} While updating Shampoo every iteration incurs a prohibitive computational overhead, the eigenvalue correction is only slightly more expensive than Adam grafting since it requires more matrix multiplications.
EShampoo requires the same memory overhead to Shampoo as Adam grafting, specifically requiring an additional $d$-dimensional buffer for the second moment accumulation. One could further reduce the memory overhead by leveraging memory-efficient versions of Adam like Adam-mini, which only requires a single scalar for each layer \citep{zhang2025adammini}, and could be applied to both grafting as well as the eigenvalue correction.

\textbf{Application to other Kronecker-factored methods.} The issue of stale eigenvalues extends to all sparsely-updated Kronecker-factored preconditioners such as K-FAC variants and TNT. Although we are unaware of any work that combines K-FAC or TNT with grafting, it is common practice to apply the Adam update to gradients preconditioned with K-FAC \citep{pauloski2021kaisa,osawa2023asdl,osawa2023pipefisher,eschenhagen2023kfac}, which we hypothesize could serve a similar function as grafting.

\begin{figure}
    \centering
    \includegraphics[width=\textwidth]{figures/adaptive_eigenbasis_ablation.pdf}
    \caption{(\textbf{left}) On the Imagewoof ViT problem, setting the maximum number of iterations $I < 10$ with threshold $\tau = 0.2$ for the adaptive QR algorithm leads to significant increase in wall-clock time compared to using adaptive \texttt{eigh}. Even with $I = 10$, adaptive \texttt{eigh} is faster in terms of wall-clock time. The default SOAP setting achieves worse final loss and is also slightly slower.
    (\textbf{right}) Using the adaptive criterion of the QR algorithm to determine when to skip the computation of an eigendecomposition (\texttt{eigh}) improves efficiency by 20\% in wall-clock time compared to the AlgoPerf setting of updating every 100 iterations on Imagewoof ViT.}
    \label{fig:adaptive-eigenbasis}
\end{figure}

While we have seen the importance of frequently updating the eigenvalues of the preconditioner in \Cref{sec:grafting}, we are still unable to predict how much performance we lose by only periodically computing the eigenbasis. The frequency of the eigenbasis computation is often tuned by hand and must be chosen in a problem-dependent way that balances the time of each iteration against the method's per-iteration convergence rate. Shampoo currently employs a fixed update frequency for all matrices, which does not distinguish between: (1) the stage of training, \ie, the beginning and the end of training; (2) the types of layers, \ie, the first and last layer of the model; and (3) the different Kronecker factors that precondition different dimensions for a single layer, \ie for $\mL$ and $\mR$ for the gradient matrix of a linear layer. Furthermore, it does not consider the amount of compute spent on each eigendecomposition, which depends on the parameter shapes.


\subsection{Can we adaptively control the approximation error induced by the stale eigenbasis?}

In the SOAP algorithm, the eigenbasis is computed via an eigendecomposition for the first iteration, and then is subsequently refined by a single step of power iteration followed by a QR decomposition, \cf Algorithm 4 in \citet{vyas2025soap}.
The authors attribute the application of this method to \citet{wang2024bitshampoo}, Appendix B.\footnote{This approach was referred to as randomized SVD in \citet{wang2024bitshampoo}.}
The method corresponds to a single iteration of a warm-started simultaneous iteration, an extension of the power iteration for iteratively computing eigendecompositions \citep{golub2013matrix}.
However, this approach may still yield a poor approximation to the true eigenbasis, especially when it changes rapidly. This depends on the choice of $\beta_2$ and the dynamics of the stochastic gradient.

Instead, we can control the error induced by the eigenbasis by allowing for multiple iterations of a \textit{warm-started QR algorithm or iteration} to be performed until a relative approximation-error-based convergence criterion is satisfied. 
There are three major benefits to this approach: (1) each Kronecker factor matrix can be treated independently from others, which decouples the frequency used across different layers, factors, and stages of training; (2) evaluating the criterion with the last computed eigenbasis prior to the first QR algorithm enables us to \textit{skip} eigenbasis updates adaptively; (3) we can control the amount of error we permit in the eigenbasis computation through the choice of the hyperparameter $\tau \in [0, 1)$, which quantifies its \textit{inexactness}.


Consider a single Kronecker factor $\mL_t \in \R^{m \times m}$ without loss of generality. We want to solve inexactly for an orthogonal matrix $\mQ_{\mL_t}$ such that we obtain an approximation to the eigendecomposition $\mL_t = \mQ_{\mL_t} \mLambda_{\mL_t} \mQ_{\mL_t}^\transpose$, with the diagonal eigenvalues matrix $\mLambda_{\mL_t}$. If the previous approximate eigenbasis $\hat{\mQ}_{\mL_{t - 1}}$ is a good approximation to the current eigenbasis $\mQ_{\mL_t}$, then we expect its approximation of the eigenvalues induced by the stale eigenbasis $\hat{\mLambda}_{\mL_t} =\hat{\mQ}_{\mL_{t - 1}}^\transpose \mL_t \hat{\mQ}_{\mL_{t - 1}}$ to be nearly diagonal. We can leverage this observation to define a unified criterion for skipping the eigenbasis computation and terminating the warm-started QR algorithm.
Diagonalizing the approximate eigenvalues $\hat{\mLambda}_{\mL_t}$ produces an approximation to $\mL_t$ induced by the stale eigenbasis, specifically, $\hat{\mL}_t = \hat{\mQ}_{\mL_{t - 1}} \diag(\hat{\mLambda}_{\mL_t}) \hat{\mQ}_{\mL_{t - 1}}^\transpose$. Bounding the relative error of this approximation is equivalent to estimating the relative error in the Frobenius norm of the off-diagonal entries of $\hat{\mLambda}_{\mL_t}$, which is cheaper to compute:
\begin{tcolorbox}[colframe=tue2, colback=tue2!3, title={Adaptive eigenbasis computation frequency criterion},bottom=0mm,top=0mm,middle=0mm]
\begin{equation}
\label{eq:criterion}
    \frac{||\mL_t - \hat{\mL}_t||_F}{||\mL_t||_F} = \frac{||\hat{\mQ}_{\mL_{t - 1}}^\transpose \mL_t \hat{\mQ}_{\mL_{t - 1}} - \diag(\hat{\mLambda}_{\mL_t})||_F}{||\hat{\mQ}_{\mL_{t - 1}}^\transpose \mL_t \hat{\mQ}_{\mL_{t - 1}}||_F} = \frac{||\hat{\mLambda}_{\mL_t} - \diag(\hat{\mLambda}_{\mL_t})||_F}{||\hat{\mLambda}_{\mL_t}||_F} \leq \tau.
\end{equation}
\end{tcolorbox}
This condition provides a guarantee of the quality of the eigendecomposition approximation. If the condition is satisfied, we set $\hat{\mQ}_{\mL_t} = \hat{\mQ}_{\mL_{t - 1}}$. Otherwise, we can run the QR algorithm until the condition is satisfied to refine the eigenbasis and set this to $\hat{\mQ}_{\mL_t}$. A complete pseudocode of the algorithm is provided in \Cref{app:algorithms-pseudocode}, \Cref{alg:adaptive-qr-iteration}.
Note that evaluating the approximate eigenvalues $\hat{\mLambda}_{\mL_t}$ at the first step requires two matrix multiplications, which is significantly cheaper than computing a step of the QR algorithm. 



\subsection{Does this adaptivity translate to efficiency gains?}

\begin{table}
  \caption{Results on a subset of the AlgoPerf workloads. We show the mean and standard error of the steps/time to the targets across the runs that reach them. See \Cref{app:additional-experiments}, \Cref{tab:full-algoperf-results} for more results.}
  \label{tab:algoperf-results}
  \centering
  \begin{tabular}{cllll}
    \toprule
    Workload & Shampoo Variant & Hits Target & Steps & Time [min] \\
    \midrule
    \multirow{3}{4em}{FastMRI} & Adam grafting ($F=100$) & $4/5$ & $4301 \pm 109$ & $13.96 \pm 0.44$ \\
    & $\mC^\mathrm{EShampoo}$ ($F=100$) & $5/5$ & $2536 \pm 66$ & $10.44 \pm 0.21$ \\
    & $\mC^\mathrm{EShampoo}$ ($\tau=0.1$, $F=50$) & $5/5$ & $2468 \pm 145$ & $10.81 \pm 0.72$ \\
    \midrule
    \multirow{3}{4em}{ImageNet ViT} & Adam grafting ($F=100$) & $1/1$ & $79907$ & $894.27$ \\
    & $\mC^\mathrm{EShampoo}$ ($F=100$) & $1/1$ & $76226$ & $894.85$ \\
    & $\mC^\mathrm{EShampoo}$ ($\tau=0.1$, $F=50$) & $1/1$ & $77459$ & $935.89$ \\
    \midrule
    \multirow{3}{4em}{OGBG} & Adam grafting ($F=100$) & $2/5$ & $12574 \pm 708$ & $39.20 \pm 1.88$ \\
    & $\mC^\mathrm{EShampoo}$ ($F=100$) & $3/5$ & $8320 \pm 1203$ & $33.02 \pm 4.05$ \\
    & $\mC^\mathrm{EShampoo}$ ($\tau=0.1$, $F=50$) & $5/5$ & $7117 \pm 328$ & $27.55 \pm 3.49$ \\
    \bottomrule
  \end{tabular}
\end{table}

The practical efficiency gains of adaptively controlling the approximation error depend on multiple factors.
First, the efficiency of the QR algorithm depends heavily on the relative cost of a single QR iteration versus an eigendecomposition, as well as the convergence rate of the QR algorithm, specifically, the number of iterations required to reach a given threshold $\tau$.
Second, evaluating the criterion adds a constant, albeit small, overhead, independent of the actual frequency of eigendecompositions.


Interestingly, we found that choosing a smaller maximum number of iterations $I < 10$ in the QR algorithm leads to drastic slowdown in wall-clock time since the total number of QR iterations is larger than the case where we run with a larger budget of $I \geq 10$ maximum iterations.
However, even with a larger number of maximum iterations, the QR algorithm wound up slightly more expensive than computing eigendecompositions with \texttt{torch.linalg.eigh} (short: \texttt{eigh}) whenever \Cref{eq:criterion} does not hold.
The default configuration of SOAP, which computes a single step of the warm-started simultaneous iteration every 10 steps, was also slightly slower in wall-clock time compared to adaptively computing eigendecompositions and reaches a worse final loss.
This result conflicts with the observation in Figure 7 in \citet{vyas2025soap} that the SOAP setting matches \texttt{eigh} when using the same fixed frequency -- likely due to the considered language modeling task.
See \Cref{fig:adaptive-eigenbasis} (left).

Therefore, we focus on leveraging the criterion for deciding when to call \texttt{eigh}. To potentially bound the total possible number of eigendecompositions performed over the course of training, we evaluate the criterion at a fixed frequency $F$ as opposed to every step. We ablate different choices of the choice of $F$ and $\tau$ in \Cref{fig:adaptive-eigenbasis} (right). We compare against the baseline setting of the winning AlgoPerf submission, which re-computes the eigendecomposition of all factor matrices every $100$ steps on the Imagewoof ViT problem. We observe that this setting without adaptivity requires $20\%$ more wall-clock time compared to an adaptive setting ($\tau=0.1$, $F=50$).

To test whether the results generalize to other problems, we consider a subset of the AlgoPerf workloads and compare the winning Shampoo submission with Adam grafting to: (1) EShampoo with the same fixed eigenbasis computation frequency of $F=100$; and (2) EShampoo with $\tau = 0.1$ and $F = 50$;
see \Cref{tab:algoperf-results}.
First, as predicted by \Cref{sec:grafting}, EShampoo with the fixed eigenbasis update frequency matches or outperforms Shampoo with Adam grafting in steps and wall-clock time, using the same hyperparameters.
Second, using the adaptive criterion with $\tau=0.1$ and $F = 50$ matches the fixed frequency $F=100$ for the FastMRI and OGBG workloads, and does slightly worse for ImageNet ViT.
We also present results with other fixed and adaptive frequencies in \Cref{app:additional-experiments}, \Cref{tab:full-algoperf-results}.
Overall, we find that the performance on the considered problems is quite robust to the eigenbasis computation frequency (\cf \Cref{app:additional-experiments}, \Cref{tab:full-algoperf-results}).
We expect that when Shampoo's convergence benefits from more frequent eigenbasis updates, adaptively determining the frequency can result in higher efficiency.



\subsection{What patterns of adaptivity emerge?}
\label{sec:patterns}

\begin{figure}
    \centering
    \includegraphics[width=\textwidth]{figures/patterns_of_adaptivity.pdf}
    \caption{We show the mean with standard error across preconditioners corresponding to the labels in the legends, for $\tau=0.01$ and $F=1$ on Imagewoof ViT. The eigenbases for biases and layer normalization parameters are changing faster than for weight matrices and linear layers, respectively.}
    \label{fig:adaptivity-patterns}
\end{figure}

The criterion in \Cref{eq:criterion} also provides a means of measuring how rapidly the preconditioner's eigenbasis is changing.
We examine how many eigendecomposition computations are performed across training steps broken down across different parameter and layer types in our Imagewoof ViT setup. We set $\tau=0.01$ and check the condition every iteration ($F = 1$); see \Cref{fig:adaptivity-patterns}.

Generally, we observe that more frequent eigendecomposition computations are performed early in training.
When comparing weight and bias parameters, we find that the eigenbases of the preconditioners corresponding to the biases evolve more rapidly than for weight matrices. Moreover, the eigenbases for layer normalization parameters change faster than for linear layer parameters, with on average almost $4 \times$ total computations. 
Finally, we also compare the number of eigendecompositions across the two Kronecker factors $\mL_t$ and $\mR_t$ for weight matrices, and $\mA_t$ for biases and layer normalization parameters.\footnote{Full-matrix Adam is also used for weight matrices $\mW_t \in \R^{m \times n}$ with $mn \leq $ \texttt{max\_preconditioner\_dim}, which is a hyperparameter to control the largest possible dimension of the preconditioner.} 
We observe more frequent updating for biases and layer normalization parameters, but it is unclear whether the underlying mechanism is related to the use of full-matrix Adam or some other properties about the training dynamics of the biases and layer normalization parameters. We also find that the eigenbasis is updated more frequently for $\mL_t$ compared to $\mR_t$.

\subsection{How does the error induced by the eigenbasis affect convergence?}

While we have investigated the effect of directly controlling the error induced by the preconditioner's stale eigenbasis, we do not care about this error itself but its impact on convergence.
In the Imagewoof ViT setting, we find that the approximation quality (determined by $\tau$) during early iterations is dramatically more important than during later iterations; see \Cref{fig:error_impact} (left).
Setting $\tau=0.8$ for the first $90\%$ and $\tau=0.01$ for the last $10\%$ of iterations leads to an identical final loss as $\tau=0.8$ for all of training ($\approx 0.09-0.1$).
In contrast, setting $\tau=0.01$ for the first $10\%$ and $\tau=0.8$ for the remaining $90\%$ of iterations leads to a significantly lower final loss ($\approx 0.04$).
In fact, setting $\tau=0.99$ for all of training, which effectively computes the eigenbases for the first iteration and then freezes them for the rest of training still significantly outperforms AdamW ($\approx 0.12$ vs. $0.3$).\footnote{For $5$ out of $172$ Kronecker factors the eigenbasis is computed twice instead of just once.}

\begin{figure}[t]
    \centering
    \includegraphics[width=\linewidth]{figures/error_impact_on_convergence.pdf}
    \caption{(\textbf{left}) The error in the eigenbases is dramatically more important for early iterations. A single eigenbasis computation at the first iteration ($\tau=0.99$) is sufficient to outperform AdamW. (\textbf{right}) The difference between the convergence behavior of AdamW and EShampoo on this problem can be exclusively attributed to the eigenbases corresponding to 2D parameters.}
    \label{fig:error_impact}
\end{figure}

Inspired by the discrepancy in the rate of change of the eigenbases corresponding to 1D and 2D parameters (\cf \Cref{fig:adaptivity-patterns}, right), we compare an Imagewoof ViT run with $\tau=0.01$ for 2D parameters and no change of basis, \ie Adam, for 1D parameters and vice versa.
Perhaps surprisingly, using Adam for 1D parameters does not affect the convergence behavior and final loss at all.
Running Adam for 2D parameters and changing the basis according to $\tau=0.01$ for 1D parameters closely matches Adam's convergence and final loss; see \Cref{fig:error_impact} (right).

Considering that the eigenbasis corresponding to 1D parameters is changing much faster than for others and about $45\%$ of the preconditioner matrices correspond to $1$D parameters, this means that a large fraction of all eigendecompositions which are computed to respect a fixed threshold $\tau$ do not impact convergence at all.
For example, in this setting, using $\tau = 0.01$ for all parameters results in a $2.9\times$ increased runtime compared to only using EShampoo for 2D and Adam for 1D parameters, while achieving the same final loss.
In practice, SOAP already uses Adam for 1D parameters, although purely motivated by reducing the computational and memory overhead \citep{vyas2025soap}.



In this paper, we find that frequently updating the eigenvalues while periodically updating the eigenbasis of Shampoo's preconditioner provides a principled and practical approach for eliminating grafting.
It remains unclear whether the same correction directly applies to the AdaGrad summation, or if a different basis-aware eigenvalue correction is needed (\cf \Cref{sec:basis-aware-derivation}).
We also show how to control the approximation error induced by the stale eigenbasis, which can improve efficiency.
The remaining obstacles to determining the eigenbasis computation frequency independently of the problem are understanding how approximation error affects convergence and integrating systems considerations, such as batched kernel efficiency, into the algorithmic design.
Further exploration is needed to understand how these trade-offs evolve with scale, \eg for large language models.

While our work focuses on Shampoo, the ideas are not limited to the AdaGrad family, and can be adapted to other methods such as K-FAC and TNT.
Turning to Shampoo's theoretical analysis, a key open question is how to integrate approximation quality into Shampoo's regret bound (\cf \Cref{sec:optimal}).
Finally, we assume in this work that full-matrix Adam is the ideal algorithm to approximate; however, there are alternative explanations for Shampoo's effectiveness \citep{carlson2015stochastic,benzing2022gradient,bernstein2024oldoptimizernewnorm,pethick2025training,zhang2025concurrence}.










We thank Anna Cai, Parameswaran Raman, and Ke Sang for their in-depth review of the paper.
We are grateful for insightful discussions with Ganesh Ajjanagadde, Rohan Anil, Anna Cai, Rong Jin, Jihao Andreas Lin, Bruno Mlodozeniec, Vinay Rao, Isaac Reid, and Xingyu (Alice) Yang.
We also acknowledge managerial support from Yuchen Hao, Guna Lakshminarayanan, Maxim Naumov, Sandeep Parab, Chunqiang Tang, and Lin Xiao.
Runa Eschenhagen is supported by ARM, the Cambridge Trust, and the Qualcomm Innovation Fellowship.
Richard E. Turner is supported by Google, Amazon, ARM, Improbable and an EPSRC Prosperity Partnership (EP/T005386/1) and the EPSRC Probabilistic AI Hub (EP/Y028783/1).

\begin{table}
  \caption{Shampoo variants and their properties.}
  \label{tab:grafting-shampoo-properties}
  \centering
  \begin{tabular}{lllll}
    \toprule
    \cmidrule(r){1-2}
    \makecell{Shampoo \\ variant} & \makecell{Justified from \\ approximation \\ perspective} & \makecell{Matches grafting \\ in steps} & \makecell{Matches grafting \\ in compute \& memory} & \makecell{Learning rate \\ transfer} \\
    \midrule
    grafting (Adam) & \xmark & \cmark &  \cmark & \cmark  \\
    $\mC^\mathrm{Shampoo}$ & \xmark & \xmark & \xmark & \xmark \\
    $S^{-1} \mC^{\mathrm{Shampoo}^2}$ & \cmark & \cmark & \xmark & ? \\
    $\mC^\mathrm{EShampoo}$ & \cmark & \cmark & \cmark & \cmark \\
    \bottomrule
  \end{tabular}
\end{table}

\section{Connecting full-matrix Adagrad to the Fisher}
\label{sec:fisher-connection}

We can also consider a more general class of preconditioned stochastic gradient methods with other choices of $\mC_t$ and $p > 0$ 
that update the parameters at each iteration $t$ by \Cref{eq:PGD}.
For example, one could choose $\mC_t$ as approximations of the Hessian $\nabla^2 \gL(\vtheta)$ via subsampling or secant approximation with $p = 1$, which yields the class of \textit{stochastic Newton} or \textit{quasi-Newton} methods \citep{keskar2016adaqn,bollapragada2018progressive,berahas2020quasi,goldfarb2020practical}.
Another common choice is the Fisher information matrix $\mF_t = \E_{\hat{\vy} \sim f_{\vtheta}(\vx), \vx \sim p_\gD(\vx)} \big[\nabla_{\vtheta} \ell(f_{\vtheta}(\vx), \hat{\vy}) \nabla_{\vtheta} \ell(f_{\vtheta}(\vx), \hat{\vy})^\transpose\big]$, which yields the \textit{natural gradient} or, for common loss functions, \textit{generalized Gauss-Newton method} with $p = 1$ \citep{amari1998natural,martens2014new}.
A summary of different methods is provided in \citet{bottou2018optimization}.

By instead summing over per-sample gradient outer products, $\hat{\mA}_t$ can be connected to the empirical Fisher. In general, the empirical Fisher is not expected to be a good approximation of the Fisher information matrix \citep{kunstner2019limitations}.
However, by replacing the accumulation with an expectation over the conditional distribution given by the model $f_{\vtheta}$, full-matrix AdaGrad can be made equivalent to the Fisher. This equivalence reveals a natural extension of Shampoo to approximate the Fisher, called Tensor Normal Training \citep[TNT]{ren2021tensor}, and is closely related to the K-FAC preconditioner \citep{anil2020scalable}. In fact, both approximations are exactly equivalent to the (block-diagonal) Fisher for simple cases such as deep linear networks (also with weight sharing) and mean squared error loss \citep{bernacchia2018exact,eschenhagen2023kfac,morwani2025new}.
The modifications in TNT deviate from the original Shampoo update \citep{gupta2018shampoo} because it was motivated by upper bounds to full-matrix AdaGrad in its non-smooth, nonconvex regret analysis, rather than this Fisher approximation perspective.

\section{Algorithms pseudocode}
\label{app:algorithms-pseudocode}

In this section, we provide the pseudocode for all algorithms, including eigenvalue-corrected Shampoo (\Cref{alg:eigenvalue-correction}), Shampoo with Adam grafting (\Cref{alg:shampoo-adam-grafting}), and the adaptive warm-started QR algorithm (\Cref{alg:adaptive-qr-iteration}). \Cref{alg:eigenvalue-correction} and \Cref{alg:shampoo-adam-grafting} present simplified versions of the algorithm ignoring common modifications like momentum or an exponential moving average over the stochastic gradient, bias corrections, and weight decay.  

Note that different instances of eigenvalue-corrected Shampoo can be employed by changing the exponential moving average in \Cref{eq:ema-shampoo}, eigenbasis computation in \Cref{eq:eigenbasis-computation}, or eigenvalue correction update in \Cref{eq:eigenvalue-correction}.
For example, SOAP delays the update of the eigenbasis until after the update step, and approximates the eigenvalue correction (\Cref{eq:eigenvalue-correction}) by 
\begin{equation}
    \label{eq:soap-update}
    \mD_t = \beta_2 \mD_{t - 1} + (1 - \beta_2) \tilde{\mG}_t^{\odot 2}
\end{equation}
and uses a single iteration of the warm-started simultaneous iteration. 

\begin{algorithm}
    \caption{Idealized eigenvalue-corrected Shampoo (EShampoo) pseudocode}
    \label{alg:eigenvalue-correction}
    \begin{algorithmic}[1]
        \Require Parameter $\mW_1 \in \R^{m \times n}$, learning rate $\alpha_t > 0$, $\epsilon > 0$.
        \State Initialize $\mL_0 = \vzero \in \R^{m \times m}$, $\mR_0 = \vzero \in \R^{n \times n}$, and $\mD = \vzero \in \R^{m \times n}$.
        \For{t = 1, ..., T}
            \State Compute (mini-batch) stochastic gradient: $\mG_t = \nabla_{\theta} \ell(f_{\vtheta_t}(\vx), \vy)$.
            \State Update factor matrices: 
            \begin{equation}
                \label{eq:ema-shampoo}
                \mL_t = \beta_2 \mL_{t - 1} + (1 - \beta_2) \mG_t \mG_t^\transpose, ~~~ \mR_t = \beta_2 \mR_{t - 1} + (1 - \beta_2) \mG_t^\transpose \mG_t.
            \end{equation}
            \State Compute or approximate orthonormal eigenbasis of the factor matrices:
            \begin{equation}
                \label{eq:eigenbasis-computation}
                \begin{aligned}
                    \mLambda_{\mL_t}, \mQ_{\mL_t} & = \mathrm{eig}(\mL_t), & \mLambda_{\mR_t}, \mQ_{\mR_t} & = \mathrm{eig}(\mR_t).
                \end{aligned}
            \end{equation}
            \State Transform gradient basis: $\tilde{\mG}_t = \mQ_{\mL_t}^\transpose \mG_t \mQ_{\mR_t}$.
            \State Compute or update eigenvalue correction: 
            \begin{equation}
                \label{eq:eigenvalue-correction}
                \mD_t^* = \argmin_{\mD \in \R^{m \times n}} \|\mA_t - (\mQ_{\mR_t} \otimes \mQ_{\mL_t}) \diag(\vectorized(\mD)) (\mQ_{\mR_t} \otimes \mQ_{\mL_t})\|_F.
            \end{equation}
            \State Compute update: $\mW_{t + 1} = \mW_t - \alpha_t \mQ_{\mL_t} (\tilde{\mG}_t / (\sqrt{\mD_t^*} + \epsilon \vone \vone^\transpose )) \mQ_{\mR_t}$.
        \EndFor
    \end{algorithmic}
\end{algorithm}

\begin{algorithm}
    \caption{Shampoo with Adam grafting pseudocode}
    \label{alg:shampoo-adam-grafting}
    \begin{algorithmic}[1]
        \Require Parameter $\mW_1 \in \R^{m \times n}$, learning rate $\alpha_t > 0$, exponential moving average constant $\beta_2 \in (0, 1)$, $\epsilon > 0$.
        \State Initialize $\mL_0 = \vzero \in \R^{m \times m}$, $\mR_0 = \vzero \in \R^{n \times n}$, and $\mD = \vzero \in \R^{m \times n}$.
        \For{t = 1, ..., T}
            \State Compute (mini-batch) stochastic gradient: $\mG_t = \nabla_{\theta} \ell(f_{\vtheta_t}(\vx), \vy)$.
            \State Update factor matrices:
            \begin{equation*}
                \mL_t = \beta_2 \mL_{t - 1} + (1 - \beta_2) \mG_t \mG_t^\transpose, ~~~ 
                \mR_t = \beta_2 \mR_{t - 1} + (1 - \beta_2) \mG_t^\transpose \mG_t.
            \end{equation*}
            \State Update Adam grafting state: $\mD_t = \beta_2 \mD_{t - 1} + (1 - \beta_2) \mG_t^{\odot 2}$.
            \State Compute matrix root inverse of the factor matrices:
            \begin{equation*}
                \begin{aligned}
                    \mL_t^{-1/4} & = \mathrm{rootinv}(\mL_t), & \mR_t^{-1/4} & = \mathrm{rootinv}(\mR_t).
                \end{aligned}
            \end{equation*}
            \State Compute update: $\mW_{t + 1} = \mW_t - \alpha_t \frac{\|- \mG_t / (\sqrt{\mD_t} + \epsilon \vone \vone^T)\|_F}{\| - \mL_t^{-1/4} \mG \mR_t^{-1/4}\|_F} \mL_t^{-1/4} \mG_t \mR_t^{-1/4}$.
        \EndFor
    \end{algorithmic}
\end{algorithm}

\begin{algorithm}
    \caption{Warm-started QR iteration with relative error termination criterion.}
    \label{alg:adaptive-qr-iteration}
    \begin{algorithmic}[1]
        \Require Matrix $\mL = \mL_t$, previous eigenbasis $\hat{\mQ} = \hat{\mQ}_{\mL_{t - 1}}$, relative tolerance $\tau \in [0, 1)$, maximum number of iterations $I$.
        \State $i \leftarrow 0$
        \State $\hat{\mLambda} \leftarrow \hat{\mQ}^\transpose \mL \hat{\mQ}$
        \While{$\|\hat{\mLambda} - \diag(\hat{\mLambda})\|_F \leq \tau \|\hat{\mLambda}\|_F$ and $i < I$}
        \State $\mQ, \mR \leftarrow \texttt{QR}(\hat{\mLambda})$
        \State $\hat{\mLambda} \leftarrow \mR \mQ$
        \State $\hat{\mQ} \leftarrow \hat{\mQ} \mQ$
        \State $i \leftarrow i + 1$
        \EndWhile \\
        \Return $\hat{\mQ}$, $\hat{\mLambda}$
    \end{algorithmic}
\end{algorithm} 

\section{Proofs}
\label{sec:proofs}

We provide complete proofs for all lemmas, propositions, and corollaries in the paper.
\begin{lemma}
    \label{lem:bound-kronecker-update-appendix}
    Let $\mU = \mQ_{\mL} (\mD^{\odot-p} \odot (\mQ_{\mL}^\transpose \mG \mQ_{\mR})) \mQ_{\mR}^\transpose \in \R^{m \times n}$ be the generalized eigendecomposed Kronecker-factored update given by orthogonal matrices $\mQ_{\mL} \in \R^{m \times m}$, $\mQ_{\mR} \in \R^{n \times n}$, and dense scaling matrix $\mD \in \R^{m \times n}$. Then we have:
    \begin{equation}
        (\max_{i, j} \mD_{i, j})^{-p} || \mG ||_F \leq || \mU ||_F \leq (\min_{i, j} \mD_{i, j})^{-p} || \mG ||_F.
    \end{equation}
\end{lemma}
\begin{proof}
    Since the Frobenius norm of a matrix is invariant to orthogonal transformations and the entries of $\mD$ are bounded, i.e., $\min_{i, j} \mD_{i, j} \leq \mD_{i, j} \leq \max_{i, j} \mD_{i, j}$, we can show that
    \begin{align*}
        \|\mU\|_F & = \|\mQ_{\mL} (\mD^{\odot-p} \odot (\mQ_{\mL}^\transpose \mG \mQ_{\mR})) \mQ_{\mR}^\transpose\|_F \\
        & = \|\mD^{\odot-p} \odot (\mQ_{\mL}^\transpose \mG \mQ_{\mR}))\|_F \\
        & \in [(\max_{i, j} \mD_{i, j})^{-p}, (\min_{i, j} \mD_{i, j})^{-p}] \cdot \|\mQ^\transpose_{\mL} \mG \mQ_{\mR}\|_F \\
        & \in [(\max_{i, j} \mD_{i, j})^{-p}, (\min_{i, j} \mD_{i, j})^{-p}] \cdot \|\mG\|_F,
    \end{align*}
    as desired.
\end{proof}

\begin{proposition}
    \label{prop:magnitude-appendix}
    Assume that $\E[\vg\vg^\transpose]$ is symmetric positive definite. The magnitude of the updates for full-matrix Adam, diagonal Adam, and eigenvalue-corrected Shampoo are all bounded by the power of the extreme eigenvalues of full-matrix Adam: 
    \begin{equation}
        \label{eq:update-magnitude-bound-appendix}
        \lambda_{\max}(\E[\vg \vg^\transpose])^{-p} \|\mG\|_F \leq \|\mU\|_F \leq \lambda_{\min}(\E[\vg \vg^\transpose])^{-p} \|\mG\|_F,
    \end{equation}
    for all $p > 0$. However, under the simplifying assumption that $\E[\mG] = \vzero$ and $\mG_{i, j}$ is independent from $\mG_{k, l}$ for $(i, j) \neq (k, l)$ and has bounded second moment, $\lambda_{\min}(\mathbb{E}[\vg \vg^\transpose]) \leq \mathbb{E}[\mG_{i, j}^2] \leq \lambda_{\max}(\mathbb{E}[\vg \vg^\transpose])$ and Shampoo has dimension-dependent bounds:
    \begin{equation}
        \label{eq:dimension-dependent-magnitude-appendix}
        m^{-p/2} n^{-p/2} \lambda_{\max}(\E[\vg \vg^\transpose])^{-p} \|\mG\|_F \leq \|\mU\|_F \leq m^{-p/2} n^{-p/2} \lambda_{\min}(\E[\vg \vg^\transpose])^{-p} \|\mG\|_F.
    \end{equation}
\end{proposition}
\begin{proof}    
    Note that \Cref{eq:update-magnitude-bound-appendix} holds for full-matrix Adam since $\|\mU\|_F = \|\vu\|_2 = \|\E[\vg \vg^\transpose]^{-p} \vg\|_2 \in [\lambda_{\max}(\E[\vg \vg^\transpose])^{-p}, \lambda_{\min}(\E[\vg \vg^\transpose])^{-p}] \cdot \|\vg\|_2$, where $\vu = \vectorized(\mU)$ and $\vg = \vectorized(\mG)$. 
    
    For diagonal Adam and eigenvalue-corrected Shampoo, it is sufficient to show that $\lambda_{\min}(\E[\vg \vg^\transpose]) \leq \mD_{i, j} \leq \lambda_{\max}(\E[\vg \vg^\transpose])$ for all $i, j$ and apply Lemma \ref{lem:bound-kronecker-update-appendix}. To see this, note that $\mD_{i, j}$ can be represented by a Rayleigh quotient, i.e., 
    \begin{align*}
        \mD_{i, j} & = \mathbb{E}[\mG_{i, j}^2] = \ve_{i, j}^\transpose \E[\vg \vg^\transpose] \ve_{i, j} & & \mbox{(Adam)} \\
        \mD_{i, j} & = \E[(\mQ_L^\transpose \mG \mQ_R)_{i, j}^2] = \ve_{i, j}^\transpose (\mQ_R \otimes \mQ_L)^\transpose \E[\vg \vg^T] (\mQ_R \otimes \mQ_L) \ve_{i, j} & & \mbox{(EShampoo)}
    \end{align*}
    where $\ve_{i, j} = \vectorized(\mE_{i, j}) \in \R^{m n}$ and $\mE_{i, j} \in \R^{m \times n}$ with
    \begin{equation*}
        (\mE_{i, j})_{k, l} = \begin{cases}
            1 & \mbox{ if } (k, l) = (i, j) \\
            0 & \mbox { otherwise}.
        \end{cases}
    \end{equation*}
    Since $\|\ve_{i, j}\|_2 = \|(\mQ_R \otimes \mQ_L) \ve_{i, j}\|_2 = 1$, the desired bound by the extreme eigenvalues follows from the Courant–Fischer–Weyl min-max theorem.

    To prove \Cref{eq:dimension-dependent-magnitude-appendix}, observe that under independence between components of the gradient $\mG_{i, j}$, the preconditioner for full-matrix Adam $\E[\vg \vg^\transpose]$ is a diagonal matrix whose diagonal entries consist of $\E[\mG_{i, j}^2]$ for all $i, j$. Hence, $\E[\mG_{i, j}^2]$ is bounded by the minimum and maximum eigenvalues, i.e., $\E[\mG_{i, j}^2] \in \big[\lambda_{\min}(\E[\vg \vg^\transpose]), \lambda_{\max}(\E[\vg \vg^\transpose])\big]$. 

    Due to independence, $\E[\mG \mG^\transpose ]$ and $\E[\mG^\transpose \mG]$ are also diagonal. Expanding their diagonal entries gives 
    \begin{align*}
        (\E[\mG \mG^\transpose])_{i, i} & = \sum_{j = 1}^{n} \E[\mG_{i, j}^2] \in n \cdot \big[\lambda_{\min}(\E[\vg \vg^\transpose]), \lambda_{\max}(\E[\vg \vg^\transpose]) \big] &  \forall ~ i = 1, ..., m, \\
        (\E[\mG^\transpose \mG])_{j, j} & = \sum_{i = 1}^{m} \E[\mG_{i, j}^2] \in m \cdot \big[\lambda_{\min}(\E[\vg \vg^\transpose]), \lambda_{\max}(\E[\vg \vg^\transpose]) \big] & \forall ~ j = 1, ..., n.
        \end{align*} 
    Therefore, the Shampoo preconditioner $(\E[\mG \mG^\transpose] \otimes \E[\mG^\transpose \mG])^{1/2}$ is also diagonal with eigenvalues lying in the interval $\big[m^{1/2} n^{1/2} \lambda_{\min}(\E[\vg \vg^\transpose]), m^{1/2} n^{1/2} \lambda_{\max}(\E[\vg \vg^\transpose])\big]$. The desired result follows.
\end{proof}

Note that more general bounds can be derived by relaxing the assumption $\E[\mG] = \vzero$, although the bounds are more complex and are not more conceptually informative.

\section{On the gap between the optimal and practical eigenvalue correction}
\label{sec:optimal-practical-gap}

Eigenvalue-corrected Shampoo shares similarities with memory-efficient optimizers such as GaLore, which apply Adam in a low-dimensional subspace spanned by the largest singular vectors of the gradient matrix \citep{zhao2024galore, su2025galore}. However, GaLore relies on the assumption that the gradient resides in a low-rank subspace that evolves gradually, allowing the same optimizer state to be used even as the subspace is updated. A recent method called LDAdam proposed by \citet{robert2024ldadam} describes a projection-aware method that corrects the scaling matrix through both a projection-aware update and generalized error feedback mechanism when transitioning between subspaces to address this issue. We will see that the naive eigenvalue correction as used in SOAP suffers from similar limitations.

\subsection{Optimal eigenvalue correction in Frobenius norm}
\label{sec:optimal}

Note that we can determine the optimal eigenvalue correction by minimizing its Frobenius norm approximation to full-matrix AdaGrad or Adam:
\begin{proposition}\label{prop:optimal}
    Given symmetric matrix $\mC \in \R^{d \times d}$ and orthogonal matrix $\mQ \in \R^{d \times d}$, the optimal eigenvalue correction $\mD^*$ that minimizes the Frobenius norm distance is given by:
    \begin{equation}
        \mD^* := \diag(\mQ^\transpose \mC \mQ) \in \arg\min_{\mD \in \mathbb{D}^d} || \mC - \mQ \mD \mQ^\transpose ||_\mathrm{F}.\footnote{$\mathbb{D}^{d}$ denotes the set of $d \times d$ diagonal matrices.}
    \end{equation}
    The exact expression for $\mD^*$ depends on the form of accumulation used for computing $\mC$; note that the damping term is omitted here since it does not change the optimal solution. $T$ denotes the current iteration.
    \begin{enumerate}
        \item (Idealized) If $\mC = \E\left[ \vg \vg^\transpose \right]$, then $\mD^* = \E\left[ \diag(\mQ^\transpose \vg)^{\odot 2} \right]$ \citep{george2018ekfac}.
        \item (AdaGrad) If $\mC = \sum_{t=1}^T \vg_t \vg_t^\transpose$, then $\mD_T^* = \sum_{t=1}^T \diag(\mQ_T^\transpose \vg_t)^{\odot 2}$. 
        \item (Adam) If $\mC = (1 - \beta_2) \sum_{t=1}^T \beta_2^{T-t} \vg_t \vg_t^\transpose$, then $\mD_T^* = (1 - \beta_2) \sum_{t=1}^T \beta_2^{T-t} \diag(\mQ_T^\transpose \vg_t)^{\odot 2}$. Note that $\mC$ can be generated recursively via an exponential moving average $\mC_T = \beta_2 \mC_{T-1} + (1 - \beta_2) \vg_T \vg_T^\transpose$ with $\mC_0 = \vzero \in \R^{mn \times mn}$.
    \end{enumerate}
\end{proposition}

\begin{proof} (Informal.)
    Since $\mQ$ is orthogonal, $\|\mC - \mQ \mD \mQ^\transpose\|_F = \|\mQ^\transpose \mC \mQ - \mD\|_F$, with $\mD$ diagonal. Therefore, the optimal solution has the form $\mD^* = \diag(\mQ^\transpose \mC \mQ)$. Each case then follows by observing that $\mC$ is the expectation or weighted sum of $\vg \vg^\transpose$, and passing $\mQ$ into the sum or expectation.
\end{proof}

While $\mD^*$ is optimal in this sense, it does not guarantee anything regarding the similarity of the (root) inverse of the approximation.
Using \Cref{prop:optimal}, we can establish the following corollary.
\begin{corollary}
    The optimal eigenvalue correction yields a tighter Frobenius norm approximation than Shampoo and CASPR, i.e., 
    \begin{equation*}
        || \mC - \mQ \mD^* \mQ^\transpose ||_\mathrm{F} \leq \min \{ || \mC - \mC^\mathrm{Shampoo} ||_\mathrm{F}, || \mC - \mC^\mathrm{CASPR} ||_\mathrm{F} \},
    \end{equation*}
    where $\mC^\mathrm{CASPR}$ is the preconditioner used by the CASPR algorithm \citep{duvvuri2024caspr}.
\end{corollary}
\begin{proof}(Informal.)
    The first inequality trivially follows from \Cref{prop:optimal} since by definition the eigenvectors of Shampoo are equal to $\mQ = (\mQ_{\mR} \otimes \mQ_{\mL})^\transpose$. The second inequality follows from \Cref{prop:optimal} together with Lemma 3.4 in \citet{duvvuri2024caspr}, which states that the eigenvectors of $\mC^\mathrm{Shampoo}$ and $\mC^\mathrm{CASPR}$ are identical.
\end{proof}

This means that, in the sense of the Frobenius norm approximation, using the optimal eigenvalue correction yields a tighter approximation to the full-matrix quantity compared to Shampoo or CASPR. 
It remains an open theoretical question whether a tighter Frobenius norm approximation can yield a tighter regret bound compared to Shampoo and CASPR. For example, one can obtain a tighter regret bound than Shampoo by only considering one-sided Shampoo, which is not generally a better approximation to full-matrix AdaGrad \citep{xie2025structured,an2025asgo}.

\subsection{Basis-aware eigenvalue correction}
\label{sec:basis-aware-derivation}

Despite its optimality, computing $\mD^*$ for cases 2 and 3 in \Cref{prop:optimal} is not feasible in practice since we would have to store and transform prior gradients from \textit{all previous iterations} with $\mQ_T$, or have access to the full-matrix quantity $\mC_T$.
For simplicity, we will focus on case 3, although the results generalize without loss of generality.\footnote{To address case 2, simply drop $1 - \beta_2$ and $\beta_2^{T-t}$.}
In contrast to the optimal eigenvalue correction, which uses a fixed basis matrix $\mQ_T$ based on the most recent statistics $\mL_T$ and $\mR_T$, SOAP uses potentially different basis matrices $\mQ_t$ based on previous statistics $\mL_t$ and $\mR_t$ available at every step $t = 1, ..., T$:
\begin{equation}
\label{eq:change-coordinate-system-naive}
    \mD^* = (1 - \beta_2) \sum_{t=1}^T \beta_2^{T-t} \diag( \mQ_T^\transpose \vg_t )^{\odot 2} \approx (1 - \beta_2) \sum_{t=1}^T \beta_2^{T-t} \diag( \mQ_{\textcolor{red}{t}}^\transpose \vg_t )^{\odot 2} =: \hat{\mD}_T.
\end{equation}
Intuitively, this means that the eigenvalue correction is  accumulated inconsistently \textit{across different coordinate systems}. This can lead to a mismatch when preconditioning $\mG_T$, which is only transformed to the \textit{current} coordinate system determined by $\mQ_T$.
When the basis remains approximately constant, \ie, $\mQ_T \approx \mQ_t$ for all $t$, \Cref{eq:change-coordinate-system-naive} can be a tight approximation and $\mD^* \approx \mD_T$.

Note that the naive eigenvalue correction may be approximately correct, i.e. the basis is approximately constant, when the basis is updated infrequently, but this can be a poor approximation overall if $\mQ$ does not approximate the changing eigenbasis of $\mC$, thereby increasing the approximation error.
The approximation potentially becomes milder for case 3 because the terms in the sum in \Cref{eq:change-coordinate-system-naive} are down-weighted by $\beta_2^{T-t}$ through the exponential moving average. Therefore, depending on $\beta_2$, the contribution from the eigenvalue correction statistic in previous coordinate systems might be negligible. However, this is not the case for case 2, where all terms are weighted equally.

\begin{figure}
    \centering
    \includegraphics[width=\textwidth]{figures/beta2_beta3_ablation.pdf}
    \caption{Decoupling the exponential moving average for the eigenbasis ($\beta_2$) and eigenvalues ($\beta_3$) for different eigendecomposition frequencies ($F$). We can observe a remarkable invariance to the choice of $\beta_2$ and $\beta_3$. Interestingly, the correction for the change of bases seems to \textit{hurt} performance overall and especially for low $\beta_2$, high $\beta_3$, and low $F$ values -- a pattern that we might expect \textit{without} the correction.}
    \label{fig:beta2-beta3}
\end{figure}

In order to address the theory-practice gap between the naive eigenvalue correction (used in SOAP) and the optimal correction in Frobenius norm, we first observe that
\begin{equation}
\begin{split}
    \mD^* & = (1 - \beta_2) \sum_{t=1}^T \beta_2^{T-t} \diag( \mQ^\transpose \vg_t )^{\odot 2} \\
    & = (1 - \beta_2) \sum_{t=1}^T \beta_2^{T-t} \diag( \mQ_T^\transpose \vg_t \vg_t^\transpose \mQ_T) \\
    & = \diag\big( \mQ_T^\transpose \big( (1 - \beta_2) \sum_{t=1}^T \beta_2^{T-t} \vg_t \vg_t^\transpose \big) \mQ_T \big) \\
    & = \diag\big( \mQ_T^\transpose \mC_T \mQ_T \big) =: \diag\big( \hat{\mC}_T \big),
\end{split}
\end{equation}
with $\hat{\mC}_T = \mQ_T^\transpose \mC_T \mQ_T$. Since $\mQ_T$ is orthogonal, we can recursively write
\begin{equation}
\begin{split}
    \mD^* = \diag\big( \hat{\mC}_T \big) &= \diag\big( \mQ_T^\transpose \big( \beta_2 \mC_{T-1} + (1 - \beta_2) \vg_T \vg_T^\transpose \big) \mQ_T \big) \\
    &= \diag\big( \beta_2 \mQ_T^\transpose \mC_{T-1} \mQ_T + (1 - \beta_2) \mQ_T^\transpose \vg_t \vg_t^\transpose \mQ_T  \big) \\
    &= \diag\big( \beta_2 \mQ_T^\transpose \mQ_{T-1} \hat{\mC}_{T - 1}  \mQ_{T-1}^\transpose \mQ_T + (1 - \beta_2) \mQ_T^\transpose \vg_T \vg_T^\transpose \mQ_T  \big) \\
    &= \beta_2 \diag\big( \mR_{T, T-1} \hat{\mC}_{T - 1}  \mR_{T, T-1}^\transpose \big) + (1 - \beta_2) \diag\big( \mQ_T^\transpose \vg_T  \big)^{\odot 2},
\end{split}
\end{equation}
where $\mR_{T, T-1} := \mQ_T^\transpose \mQ_{T-1}$ is the transition matrix between different bases.

While this is the exact expression for our desired quantity, it requires keeping track of the full matrix $\hat{\mC}_{T-1}$ to compute $\diag\big( \mR_{T, T-1} \hat{\mC}_{T - 1} \mR_{T, T-1}^\transpose \big)$, which is not tractable.
Since we explicitly construct $\mQ_{T-1}$ to be close to the best Kronecker-factored basis for $\mC_{T-1}$ through the choice of the Shampoo preconditioner with $\mL_{T-1}$ and $\mR_{T-1}$, we make the additional assumption that $\hat{\mC}_{T - 1}$ is approximately diagonal, \ie, $\hat{\mC}_{T - 1} \approx  \diag(\hat{\mC}_{T - 1}) = \diag(\vv_{T-1}^\textrm{corrected})$.

Substituting this back into the recursive equation above, we have that
\begin{equation}
\begin{split}
    \mD^* = \diag\big( \hat{\mC}_T \big) &\approx \beta_2 \diag\big( \mR_{T, T-1} \diag(\vv_{T-1}^\textrm{corrected}) \mR_{T, T-1}^\transpose \big) + (1 - \beta_2) \diag\big( \mQ_T^\transpose \vg_T \big)^{\odot 2} \\
    &= \beta_2 \diag\big( \mR_{T, T-1}^{\odot 2} \vv_{T-1}^\textrm{corrected} \big) + (1 - \beta_2) \diag\big( \mQ_T^\transpose \vg_T \big)^{\odot 2} \\
    &= \diag\big(  \vv_T^\textrm{corrected} \big).
\end{split}
\end{equation}
This gives a recursive update rule that, in contrast to the naive solution, accounts for the changes of bases between iterations.
This also recovers the exact solution when $\mQ_T = \mQ_t$ for all $t$ since $\mR_{t, t-1} = \mathbf{I}$.
This correction for the changing basis has a similar motivation to and resembles parts of the LDAdam algorithm \citep{robert2024ldadam}.

We design an experiment to empirically test whether the implicit approximation in \Cref{eq:change-coordinate-system-naive} helps in practice, and find that interestingly, the correction appears to hurt performance in the settings where we would expect it to help, see \Cref{fig:beta2-beta3}.
A satisfying explanation of this phenomenon remains an open question.

\section{Additional experimental details and results}
\label{app:additional-experiments}

For all experiments we used $1\times$ NVIDIA A100 80GB GPU per run, with the exception of the ImageNet ViT experiments, for which we used $4\times$ NVIDIA A100 80GB GPUs per run.
All of the experiments were conducted on an internal compute cluster and we estimate that the total required compute was around $1440$ GPU hours or $60$ days.
We ran additional exploratory experiments not reported in this paper which increases the total compute cost of the project.

\subsection{Imagewoof experimental details}
\label{sec:imagewoof-details}

We use the Imagewoof dataset.\footnote{The Imagewoof dataset is available at \url{https://github.com/fastai/imagenette}.}
All models are trained with cross entropy loss for $100$ epochs, using a learning rate schedule consisting of a linear warmup for 353 steps followed by cosine decay. We use a batch size of $128$, randomized cropping and horizontal flips as data augmentation, and the default settings for $\beta_1 = 0.9$ and $\beta_2 = 0.999$.

For the vision transformer (ViT) model, we use SimpleViT \citep{beyer2022better} with patch size $16$, $6$ heads, a depth of $12$ layers, an MLP dimension of $1536$, dimension of $384$, gradient clipping with threshold $1$, and weight decay of $10^{-4}$.
For the ConvNeXt V2 architecture \citep{woo2023convnextv2} we use weight decay of $0.05$ and drop paths with rate $0.1$.

For \Cref{fig:lr-transfer}, we fix $\epsilon = 10^{-10}$ and sweep the following learning rates $\alpha$:\footnote{We fixed $\epsilon$ based on a wide sweep over $\alpha$ and $\epsilon$ for AdamW. We also use this $\epsilon$ when grafting from Adam.}
\begin{itemize}
    \item Vision transformer
    \begin{itemize}
        \item AdamW: $\alpha \in \{10^{-4}, 3 \cdot 10^{-4}, 6 \cdot 10^{-4}, 10^{-3}, 3 \cdot 10^{-3}, 6 \cdot 10^{-3}, 10^{-2} \}$
       \item Shampoo with grafting: $\alpha \in \{ 10^{-4}, 3 \cdot 10^{-4}, 10^{-3}, 3 \cdot 10^{-3} \}$
        \item Shampoo: $\alpha \in \{3 \cdot 10^{-3}, 6 \cdot 10^{-3}, 10^{-2}, 3 \cdot 10^{-3} \}$
        \item Shampoo$^2$ with trace scaling: $\alpha \in \{ 10^{-4}, 3 \cdot 10^{-4}, 3 \cdot 10^{-3} \}$
        \item EShampoo: $\alpha \in \{ 10^{-4}, 3 \cdot 10^{-4}, 10^{-3}, 3 \cdot 10^{-3} \}$
    \end{itemize}
    \item ConvNeXt V2
    \begin{itemize}
        \item AdamW: $\alpha \in \{10^{-4}, 3 \cdot 10^{-4}, 6 \cdot 10^{-4}, 10^{-3}, 3 \cdot 10^{-3}, 6 \cdot 10^{-3}, 10^{-2} \}$
        \item Shampoo with grafting: $\alpha \in \{10^{-4}, 3 \cdot 10^{-4}, 10^{-3}, 3 \cdot 10^{-3}, 10^{-2} \}$
        \item Shampoo: $\alpha \in \{ 10^{-4}, 3 \cdot 10^{-4}, 10^{-3}, 3 \cdot 10^{-3}\}$
        \item Shampoo$^2$ with trace scaling: $\alpha \in \{ 10^{-4}, 3 \cdot 10^{-4}, 10^{-3} \}$
        \item EShampoo: $\alpha \in \{10^{-4}, 3 \cdot 10^{-4}, 10^{-3}, 3 \cdot 10^{-3}, 10^{-2} \}$ 
    \end{itemize}
\end{itemize}
For EShampoo in \Cref{fig:adaptive-eigenbasis}, \Cref{fig:adaptivity-patterns}, and \Cref{fig:error_impact} we use the learning rate $\alpha = 6 \cdot 10^{-4}$ and $\epsilon = 10^{-10}$.





\subsection{AlgoPerf workloads}

We follow the standard AlgoPerf setup and consider wall-clock time to pre-specified validation metric targets.
See \citet{dahl2023benchmarking} and \citet{kasimbeg2025accelerating} for more details on the AlgoPerf benchmark.\footnote{The benchmark code is available at \url{https://github.com/mlcommons/algorithmic-efficiency/}.}
The FastMRI dataset can be attributed to \citet{knoll2020fastmri,zbontar2019fastmri}, the ImageNet dataset to \citet{krizhevsky2012imagenet}, and the OGBG dataset to \citet{hu2021opengraphbenchmark}.

We choose this specific subset of the workloads because (1) we want to include a larger scale vision transformer (ImageNet ViT), an architecture we use in the small-scale experiments, (2) the FastMRI and OGBG workloads share the same hyperparameter settings with the Imagenet ViT workload, hence excluding them as a confounding factor, and (3) the $\beta_2$ used for these workloads is the smallest among all hyperparameter settings of the winning Shampoo submission, resulting in the fastest moving average and thereby potentially faster changing eigenbases.

We run the winning Shampoo submission with Adam grafting, $F=100$, and the best hyperparameter setting for each workload.\footnote{The submission is available at \url{https://github.com/mlcommons/submissions_algorithms/tree/main/previous_leaderboards/algoperf_v05/submissions/external_tuning/shampoo_submission}.}
For EShampoo, we use the same best-performing hyperparameter setting from Shampoo but turn off learning rate grafting from Adam, and modify $F$ and $\tau$.

\begin{table}
  \caption{Results for all considered settings on a subset of the AlgoPerf workloads. We show the mean and standard error of the steps and time to the targets across the runs that actually hit the targets.}
  \label{tab:full-algoperf-results}
  \centering
  \begin{tabular}{cllll}
    \toprule
    Workload & Shampoo Variant & Hits Targets & Steps & Time [min] \\
    \midrule
    \multirow{4}{4em}{FastMRI} & Adam grafting ($F=100$) & $4/5$ & $4301 \pm 109$ & $13.96 \pm 0.44$ \\
    & $\mC^\mathrm{EShampoo}$ ($F=100$) & $5/5$ & $2536 \pm 66$ & $10.44 \pm 0.21$ \\
    & $\mC^\mathrm{EShampoo}$ ($F=50$) & $5/5$ & $2578 \pm 86$ & $10.86 \pm 0.27$ \\
    & $\mC^\mathrm{EShampoo}$ ($F=10$) & $5/5$ & $2311 \pm 73$ & $14.93 \pm 1.97$ \\
    & $\mC^\mathrm{EShampoo}$ ($F=1$) & $5/5$ & $2101 \pm 31$ & $35.34 \pm 0.70$ \\
    & $\mC^\mathrm{EShampoo}$ ($\tau=0.1$, $F=100$) & $5/5$ & $2553 \pm 154$ & $16.33 \pm 2.10$\footnotemark \\
    & $\mC^\mathrm{EShampoo}$ ($\tau=0.1$, $F=50$) & $5/5$ & $2468 \pm 145$ & $10.81 \pm 0.72$ \\
    & $\mC^\mathrm{EShampoo}$ ($\tau=0.1$, $F=10$) & $5/5$ & $2420 \pm 92$ & $10.76 \pm 0.73$ \\
    & $\mC^\mathrm{EShampoo}$ ($\tau=0.1$, $F=1$) & $5/5$ & $2367 \pm 96$ & $10.93 \pm 0.42$ \\
    & $\mC^\mathrm{EShampoo}$ ($\tau=0.01$, $F=1$) & $5/5$ & $2208 \pm 41$ & $27.43 \pm 0.49$ \\
    \midrule
    \multirow{5}{4em}{ImageNet ViT} & Adam grafting ($F=100$) & $1/1$ & $79907$ & $894.27$ \\
    & $\mC^\mathrm{EShampoo}$ ($F=100$) & $1/1$ & $76226$ & $894.85$ \\
    & $\mC^\mathrm{EShampoo}$ ($F=10$) & $1/1$ & $73237$ & $1160.53$ \\
    & $\mC^\mathrm{EShampoo}$ ($\tau=0.1$, $F=100$) & $1/1$ & $74010$ & $852.66$ \\
    & $\mC^\mathrm{EShampoo}$ ($\tau=0.1$, $F=50$) & $1/1$ & $77459$ & $935.89$ \\
    & $\mC^\mathrm{EShampoo}$ ($\tau=0.01$, $F=10$) & $1/1$ & $75841$ & $1188.76$ \\
    \midrule
    \multirow{4}{4em}{OGBG} & Adam grafting ($F=100$) & $2/5$ & $12574 \pm 708$ & $39.20 \pm 1.88$ \\
    & $\mC^\mathrm{EShampoo}$ ($F=100$) & $3/5$ & $8320 \pm 1203$ & $33.02 \pm 4.05$ \\
    & $\mC^\mathrm{EShampoo}$ ($F=50$) & $5/5$ & $7173 \pm 443$ & $26.17 \pm 1.31$ \\
    & $\mC^\mathrm{EShampoo}$ ($F=10$) & $3/5$ & $6645 \pm 357$ & $37.55 \pm 1.74$ \\
    & $\mC^\mathrm{EShampoo}$ ($F=1$)\footnotemark & $-$ & $-$ & $-$ \\
    & $\mC^\mathrm{EShampoo}$ ($\tau=0.1$, $F=100$) & $4/5$ & $8047 \pm 369$ & $27.60 \pm 1.15$ \\
    & $\mC^\mathrm{EShampoo}$ ($\tau=0.1$, $F=50$) & $5/5$ & $7117 \pm 328$ & $27.55 \pm 3.49$ \\
    & $\mC^\mathrm{EShampoo}$ ($\tau=0.1$, $F=10$) & $5/5$ & $7151 \pm 416$ & $29.11 \pm 1.98$ \\
    & $\mC^\mathrm{EShampoo}$ ($\tau=0.1$, $F=1$) & $5/5$ & $6758 \pm 273$ & $34.16 \pm 1.65$ \\
    & $\mC^\mathrm{EShampoo}$ ($\tau=0.01$, $F=10$) & $2/5$ & $7234 \pm 361$ & $39.15 \pm 2.01$ \\
    \bottomrule
  \end{tabular}
\end{table}
\footnotetext{This wall-clock time statistic seems to be negatively impacted by either an issue with the AlgoPerf code or the hardware setup.}
\footnotetext{Runs failed due to \url{https://github.com/mlcommons/algorithmic-efficiency/issues/866}.}

\begin{figure}
    \label{fig:fastmri}
    \centering
    \includegraphics[width=\textwidth]{figures/fastmri.pdf}
    \caption{FastMRI AlgoPerf workload.}
\end{figure}

\begin{figure}
    \label{fig:imagenet-vit}
    \centering
    \includegraphics[width=\textwidth]{figures/imagenet_vit.pdf}
    \caption{ImageNet ViT AlgoPerf workload.}
\end{figure}

\begin{figure}
    \label{fig:ogbg}
    \centering
    \includegraphics[width=\textwidth]{figures/ogbg.pdf}
    \caption{OGBG AlgoPerf workload.}
\end{figure}








\begin{thebibliography}{61}
\providecommand{\natexlab}[1]{#1}
\providecommand{\url}[1]{\texttt{#1}}
\expandafter\ifx\csname urlstyle\endcsname\relax
  \providecommand{\doi}[1]{doi: #1}\else
  \providecommand{\doi}{doi: \begingroup \urlstyle{rm}\Url}\fi

\bibitem[Agarwal et~al.(2020)Agarwal, Anil, Hazan, Koren, and Zhang]{agarwal2020disentangling}
Naman Agarwal, Rohan Anil, Elad Hazan, Tomer Koren, and Cyril Zhang.
\newblock Disentangling adaptive gradient methods from learning rates.
\newblock arXiv2002.11803, 2020.

\bibitem[Amari(1998)]{amari1998natural}
Shun{-}Ichi Amari.
\newblock Natural gradient works efficiently in learning.
\newblock \emph{Neural computation}, 10\penalty0 (2), 1998.

\bibitem[An et~al.(2025)An, Liu, Pan, Ma, Goldfarb, and Zhang]{an2025asgo}
Kang An, Yuxing Liu, Rui Pan, Shiqian Ma, Donald Goldfarb, and Tong Zhang.
\newblock {ASGO}: Adaptive structured gradient optimization.
\newblock {arXiv} 2503.20762, 2025.

\bibitem[Anil et~al.(2020)Anil, Gupta, Koren, Regan, and Singer]{anil2020scalable}
Rohan Anil, Vineet Gupta, Tomer Koren, Kevin Regan, and Yoram Singer.
\newblock Scalable second order optimization for deep learning.
\newblock ar{X}iv 2002.09018, 2020.

\bibitem[Benzing(2022)]{benzing2022gradient}
Frederik Benzing.
\newblock Gradient descent on neurons and its link to approximate second-order optimization.
\newblock In \emph{{ICML}}, 2022.

\bibitem[Berahas et~al.(2020)Berahas, Jahani, Richt{\'a}rik, and Tak{\'a}{\v{c}}]{berahas2020quasi}
Albert~S Berahas, Majid Jahani, Peter Richt{\'a}rik, and Martin Tak{\'a}{\v{c}}.
\newblock {Quasi-Newton} methods for deep learning: Forget the past, just sample.
\newblock arXiv 1901.09997, 2020.

\bibitem[Bernacchia et~al.(2018)Bernacchia, Lengyel, and Hennequin]{bernacchia2018exact}
Alberto Bernacchia, Mate Lengyel, and Guillaume Hennequin.
\newblock Exact natural gradient in deep linear networks and its application to the nonlinear case.
\newblock In \emph{{N}eur{IPS}}, 2018.

\bibitem[Bernstein(2025)]{bernstein2025deriving}
Jeremy Bernstein.
\newblock Deriving {Muon}, 2025.
\newblock URL \url{https://jeremybernste.in/writing/deriving-muon}.

\bibitem[Bernstein \& Newhouse(2024)Bernstein and Newhouse]{bernstein2024oldoptimizernewnorm}
Jeremy Bernstein and Laker Newhouse.
\newblock Old optimizer, new norm: An anthology.
\newblock arXiv 2409.20325, 2024.

\bibitem[Beyer et~al.(2022)Beyer, Zhai, and Kolesnikov]{beyer2022better}
Lucas Beyer, Xiaohua Zhai, and Alexander Kolesnikov.
\newblock Better plain {ViT} baselines for {ImageNet-1k}.
\newblock arXiv 2205.01580, 2022.

\bibitem[Bollapragada et~al.(2018)Bollapragada, Nocedal, Mudigere, Shi, and Tang]{bollapragada2018progressive}
Raghu Bollapragada, Jorge Nocedal, Dheevatsa Mudigere, Hao-Jun Shi, and Ping Tak~Peter Tang.
\newblock A progressive batching {L-BFGS} method for machine learning.
\newblock In \emph{International Conference on Machine Learning}, pp.\  620--629. PMLR, 2018.

\bibitem[Bottou et~al.(2018)Bottou, Curtis, and Nocedal]{bottou2018optimization}
L{\'e}on Bottou, Frank~E Curtis, and Jorge Nocedal.
\newblock Optimization methods for large-scale machine learning.
\newblock \emph{SIAM review}, 60\penalty0 (2):\penalty0 223--311, 2018.

\bibitem[Carlson et~al.(2015)Carlson, Cevher, and Carin]{carlson2015stochastic}
David Carlson, Volkan Cevher, and Lawrence Carin.
\newblock Stochastic spectral descent for restricted {Boltzmann} machines.
\newblock In \emph{{ICML}}, 2015.

\bibitem[Dahl et~al.(2023)Dahl, Schneider, Nado, Agarwal, Sastry, Hennig, Medapati, Eschenhagen, Kasimbeg, Suo, Bae, Gilmer, Peirson, Khan, Anil, Rabbat, Krishnan, Snider, Amid, Chen, Maddison, Vasudev, Badura, Garg, and Mattson]{dahl2023benchmarking}
George~E. Dahl, Frank Schneider, Zachary Nado, Naman Agarwal, Chandramouli~Shama Sastry, Philipp Hennig, Sourabh Medapati, Runa Eschenhagen, Priya Kasimbeg, Daniel Suo, Juhan Bae, Justin Gilmer, Abel~L. Peirson, Bilal Khan, Rohan Anil, Mike Rabbat, Shankar Krishnan, Daniel Snider, Ehsan Amid, Kongtao Chen, Chris~J. Maddison, Rakshith Vasudev, Michal Badura, Ankush Garg, and Peter Mattson.
\newblock Benchmarking neural network training algorithms.
\newblock ar{X}iv 2306.07179, 2023.

\bibitem[Duchi et~al.(2011)Duchi, Hazan, and Singer]{duchi2011adaptive}
John Duchi, Elad Hazan, and Yoram Singer.
\newblock Adaptive subgradient methods for online learning and stochastic optimization.
\newblock \emph{{JMLR}}, 12\penalty0 (61), 2011.

\bibitem[Duvvuri et~al.(2024)Duvvuri, Devvrit, Anil, Hsieh, and Dhillon]{duvvuri2024caspr}
Sai~Surya Duvvuri, Fnu Devvrit, Rohan Anil, Cho-Jui Hsieh, and Inderjit~S Dhillon.
\newblock {CASPR}: Combining axes preconditioners through {K}ronecker approximation for deep learning.
\newblock In \emph{ICLR}, 2024.

\bibitem[Eschenhagen et~al.(2023)Eschenhagen, Immer, Turner, Schneider, and Hennig]{eschenhagen2023kfac}
Runa Eschenhagen, Alexander Immer, Richard~E. Turner, Frank Schneider, and Philipp Hennig.
\newblock {K}ronecker-{F}actored {A}pproximate {C}urvature for modern neural network architectures.
\newblock In \emph{NeurIPS}, 2023.

\bibitem[Feinberg et~al.(2023)Feinberg, Chen, Sun, Anil, and Hazan]{feinberg2023sketchy}
Vladimir Feinberg, Xinyi Chen, Y.~Jennifer Sun, Rohan Anil, and Elad Hazan.
\newblock Sketchy: Memory-efficient adaptive regularization with frequent directions.
\newblock In \emph{{N}eur{IPS}}, 2023.

\bibitem[Gao et~al.(2020)Gao, Liu, Huang, Wang, Wang, Xu, and Yu]{gao2020trace}
Kai-Xin Gao, Xiaolei Liu, Zheng-Hai Huang, Min Wang, Zidong Wang, Dachuan Xu, and F.~Yu.
\newblock A trace-restricted {K}ronecker-factored approximation to natural gradient.
\newblock In \emph{AAAI Conference on Artificial Intelligence}, 2020.

\bibitem[George et~al.(2018)George, Laurent, Bouthillier, Ballas, and Vincent]{george2018ekfac}
Thomas George, César Laurent, Xavier Bouthillier, Nicolas Ballas, and Pascal Vincent.
\newblock Fast approximate natural gradient descent in a {K}ronecker-factored eigenbasis.
\newblock In \emph{{N}eur{IPS}}, 2018.

\bibitem[Goldfarb et~al.(2020)Goldfarb, Ren, and Bahamou]{goldfarb2020practical}
Donald Goldfarb, Yi~Ren, and Achraf Bahamou.
\newblock Practical {quasi-Newton} methods for training deep neural networks.
\newblock \emph{Advances in Neural Information Processing Systems}, 33:\penalty0 2386--2396, 2020.

\bibitem[Golub \& Van~Loan(2013)Golub and Van~Loan]{golub2013matrix}
Gene~H Golub and Charles~F Van~Loan.
\newblock \emph{Matrix computations}.
\newblock JHU press, 2013.

\bibitem[Grosse \& Martens(2016)Grosse and Martens]{grosse2016kronecker}
Roger~B. Grosse and James Martens.
\newblock A {K}ronecker-factored approximate {F}isher matrix for convolution layers.
\newblock In \emph{{ICML}}, 2016.

\bibitem[Gupta et~al.(2018)Gupta, Koren, and Singer]{gupta2018shampoo}
Vineet Gupta, Tomer Koren, and Yoram Singer.
\newblock Shampoo: Preconditioned stochastic tensor optimization.
\newblock In \emph{{ICML}}, 2018.

\bibitem[He et~al.(2024)He, Noci, Paliotta, Schlag, and Hofmann]{he2024understanding}
Bobby He, Lorenzo Noci, Daniele Paliotta, Imanol Schlag, and Thomas Hofmann.
\newblock Understanding and minimising outlier features in neural network training.
\newblock In \emph{NeurIPS}, 2024.

\bibitem[Heskes(2000)]{heskes2000natural}
Tom Heskes.
\newblock On ``natural'' learning and pruning in multilayered perceptrons.
\newblock \emph{Neural Computation}, 12\penalty0 (4), 2000.

\bibitem[Hu et~al.(2021)Hu, Fey, Zitnik, Dong, Ren, Liu, Catasta, and Leskovec]{hu2021opengraphbenchmark}
Weihua Hu, Matthias Fey, Marinka Zitnik, Yuxiao Dong, Hongyu Ren, Bowen Liu, Michele Catasta, and Jure Leskovec.
\newblock Open graph benchmark: Datasets for machine learning on graphs.
\newblock arXiv 2005.00687, 2021.

\bibitem[Jordan et~al.(2024)Jordan, Jin, Boza, You, Cesista, Newhouse, and Bernstein]{jordan2024muon}
Keller Jordan, Yuchen Jin, Vlado Boza, Jiacheng You, Franz Cesista, Laker Newhouse, and Jeremy Bernstein.
\newblock Muon: An optimizer for hidden layers in neural networks, 2024.
\newblock URL \url{https://kellerjordan.github.io/posts/muon/}.

\bibitem[Kasimbeg et~al.(2025)Kasimbeg, Schneider, Eschenhagen, Bae, Sastry, Saroufim, Feng, Wright, Yang, Nado, Medapati, Hennig, Rabbat, and Dahl]{kasimbeg2025accelerating}
Priya Kasimbeg, Frank Schneider, Runa Eschenhagen, Juhan Bae, Chandramouli~Shama Sastry, Mark Saroufim, Boyuan Feng, Less Wright, Edward~Z. Yang, Zachary Nado, Sourabh Medapati, Philipp Hennig, Michael Rabbat, and George~E. Dahl.
\newblock Accelerating neural network training: An analysis of the {A}lgo{P}erf competition.
\newblock In \emph{ICLR}, 2025.

\bibitem[Keskar \& Berahas(2016)Keskar and Berahas]{keskar2016adaqn}
Nitish~Shirish Keskar and Albert~S Berahas.
\newblock adaqn: An adaptive {quasi-Newton} algorithm for training {RNNs}.
\newblock In \emph{Machine Learning and Knowledge Discovery in Databases: European Conference, ECML PKDD 2016, Riva del Garda, Italy, September 19-23, 2016, Proceedings, Part I 16}, pp.\  1--16. Springer, 2016.

\bibitem[Kingma \& Ba(2015)Kingma and Ba]{kingma2014adam}
Diederik~P Kingma and Jimmy Ba.
\newblock Adam: A method for stochastic optimization.
\newblock In \emph{ICLR}, 2015.

\bibitem[Knoll et~al.(2020)Knoll, Zbontar, Sriram, Muckley, Bruno, Defazio, Parente, Geras, Katsnelson, Chandarana, Zhang, Drozdzalv, Romero, Rabbat, Vincent, Pinkerton, Wang, Yakubova, Owens, Zitnick, Recht, Sodickson, and Lui]{knoll2020fastmri}
Florian Knoll, Jure Zbontar, Anuroop Sriram, Matthew~J. Muckley, Mary Bruno, Aaron Defazio, Marc Parente, Krzysztof~J. Geras, Joe Katsnelson, Hersh Chandarana, Zizhao Zhang, Michal Drozdzalv, Adriana Romero, Michael Rabbat, Pascal Vincent, James Pinkerton, Duo Wang, Nafissa Yakubova, Erich Owens, C.~Lawrence Zitnick, Michael~P. Recht, Daniel~K. Sodickson, and Yvonne~W. Lui.
\newblock fast{MRI}: A publicly available raw k-space and {DICOM} dataset of knee images for accelerated {MR} image reconstruction using machine learning.
\newblock \emph{Radiology: Artificial Intelligence}, 2\penalty0 (1):\penalty0 e190007, 2020.
\newblock \doi{10.1148/ryai.2020190007}.
\newblock PMID: 32076662.

\bibitem[Krizhevsky et~al.(2012)Krizhevsky, Sutskever, and Hinton]{krizhevsky2012imagenet}
Alex Krizhevsky, Ilya Sutskever, and Geoffrey~E Hinton.
\newblock {ImageNet} classification with deep convolutional neural networks.
\newblock In \emph{NIPS}, 2012.

\bibitem[Kunstner et~al.(2019)Kunstner, Balles, and Hennig]{kunstner2019limitations}
Frederik Kunstner, Lukas Balles, and Philipp Hennig.
\newblock Limitations of the empirical {F}isher approximation for natural gradient descent.
\newblock In \emph{NeurIPS}, 2019.

\bibitem[Li(2018)]{li2018preconditioned}
Xi-Lin Li.
\newblock Preconditioned stochastic gradient descent.
\newblock \emph{IEEE Transactions on Neural Networks and Learning Systems}, 29\penalty0 (5), 2018.

\bibitem[Lin et~al.(2024{\natexlab{a}})Lin, Dangel, Eschenhagen, Bae, Turner, and Makhzani]{lin2024remove}
Wu~Lin, Felix Dangel, Runa Eschenhagen, Juhan Bae, Richard~E. Turner, and Alireza Makhzani.
\newblock Can we remove the square-root in adaptive gradient methods? a second-order perspective.
\newblock In \emph{ICML}, 2024{\natexlab{a}}.

\bibitem[Lin et~al.(2024{\natexlab{b}})Lin, Dangel, Eschenhagen, Neklyudov, Kristiadi, Turner, and Makhzani]{lin2024structured}
Wu~Lin, Felix Dangel, Runa Eschenhagen, Kirill Neklyudov, Agustinus Kristiadi, Richard~E. Turner, and Alireza Makhzani.
\newblock Structured inverse-free natural gradient: Memory-efficient \& numerically-stable {KFAC} for large neural nets.
\newblock In \emph{ICML}, 2024{\natexlab{b}}.

\bibitem[Liu et~al.(2018)Liu, Masana, Herranz, van~de Weijer, L{\'o}pez, and Bagdanov]{liu2018rotate}
Xialei Liu, Marc Masana, Luis Herranz, Joost van~de Weijer, Antonio~M. L{\'o}pez, and Andrew~D. Bagdanov.
\newblock Rotate your networks: Better weight consolidation and less catastrophic forgetting.
\newblock \emph{{ICPR}}, 2018.

\bibitem[Martens(2010)]{martens2010deep}
James Martens.
\newblock Deep learning via {Hessian}-free optimization.
\newblock In \emph{{ICML}}, 2010.

\bibitem[Martens(2014)]{martens2014new}
James Martens.
\newblock New insights and perspectives on the natural gradient method.
\newblock \emph{JMLR}, 21\penalty0 (146), 2014.

\bibitem[Martens \& Grosse(2015)Martens and Grosse]{martens2015optimizing}
James Martens and Roger Grosse.
\newblock Optimizing neural networks with {K}ronecker-factored approximate curvature.
\newblock In \emph{ICML}, 2015.

\bibitem[Martens et~al.(2018)Martens, Ba, and Johnson]{martens2018kroneckerfactored}
James Martens, Jimmy Ba, and Matt Johnson.
\newblock {K}ronecker-factored curvature approximations for recurrent neural networks.
\newblock In \emph{{ICLR}}, 2018.

\bibitem[Morwani et~al.(2025)Morwani, Shapira, Vyas, Malach, Kakade, and Janson]{morwani2025new}
Depen Morwani, Itai Shapira, Nikhil Vyas, Eran Malach, Sham Kakade, and Lucas Janson.
\newblock A new perspective on {S}hampoo's preconditioner.
\newblock In \emph{ICLR}, 2025.

\bibitem[Osawa et~al.(2023{\natexlab{a}})Osawa, Ishikawa, Yokota, Li, and Hoefler]{osawa2023asdl}
Kazuki Osawa, Satoki Ishikawa, Rio Yokota, Shigang Li, and Torsten Hoefler.
\newblock {ASDL}: A unified interface for gradient preconditioning in {P}y{T}orch.
\newblock ar{X}iv 2305.04684, 2023{\natexlab{a}}.

\bibitem[Osawa et~al.(2023{\natexlab{b}})Osawa, Li, and Hoefler]{osawa2023pipefisher}
Kazuki Osawa, Shigang Li, and Torsten Hoefler.
\newblock {P}ipe{F}isher: Efficient training of large language models using pipelining and {F}isher information matrices.
\newblock In Dawn Song, Michael Carbin, and Tianqi Chen (eds.), \emph{{MLS}ys}, 2023{\natexlab{b}}.

\bibitem[Pauloski et~al.(2021)Pauloski, Huang, Huang, Venkataraman, Chard, Foster, and Zhang]{pauloski2021kaisa}
J.~Gregory Pauloski, Qi~Huang, Lei Huang, Shivaram Venkataraman, Kyle Chard, Ian~T. Foster, and Zhao Zhang.
\newblock {KAISA:} an adaptive second-order optimizer framework for deep neural networks.
\newblock In \emph{International Conference for High Performance Computing, Networking, Storage and Analysis ({SC}21)}, 2021.

\bibitem[Pethick et~al.(2025)Pethick, Xie, Antonakopoulos, Zhu, Silveti-Falls, and Cevher]{pethick2025training}
Thomas Pethick, Wanyun Xie, Kimon Antonakopoulos, Zhenyu Zhu, Antonio Silveti-Falls, and Volkan Cevher.
\newblock Training deep learning models with norm-constrained {LMO}s.
\newblock ar{X}iv 2502.07529, 2025.

\bibitem[Pooladzandi \& Li(2024)Pooladzandi and Li]{pooladzandi2024curvature}
Omead Pooladzandi and Xi-Lin Li.
\newblock Curvature-informed {SGD} via general purpose lie-group preconditioners.
\newblock arXiv 2402.04553, 2024.

\bibitem[Ren \& Goldfarb(2021)Ren and Goldfarb]{ren2021tensor}
Yi~Ren and Donald Goldfarb.
\newblock Tensor normal training for deep learning models.
\newblock In \emph{{N}eur{IPS}}, 2021.

\bibitem[Robert et~al.(2024)Robert, Safaryan, Modoranu, and Alistarh]{robert2024ldadam}
Thomas Robert, Mher Safaryan, Ionut-Vlad Modoranu, and Dan Alistarh.
\newblock {LDA}dam: Adaptive optimization from low-dimensional gradient statistics.
\newblock arXiv 2410.16103, 2024.

\bibitem[Shi et~al.(2023)Shi, Lee, Iwasaki, Gallego-Posada, Li, Rangadurai, Mudigere, and Rabbat]{shi2023distributed}
Hao-Jun~Michael Shi, Tsung-Hsien Lee, Shintaro Iwasaki, Jose Gallego-Posada, Zhijing Li, Kaushik Rangadurai, Dheevatsa Mudigere, and Michael Rabbat.
\newblock A distributed data-parallel {PyTorch} implementation of the {Distributed Shampoo} optimizer for training neural networks at-scale.
\newblock arXiv 2309.06497, 2023.

\bibitem[Su et~al.(2025)Su, Gu, Xu, Tian, and Zhao]{su2025galore}
DiJia Su, Andrew Gu, Jane Xu, Yuandong Tian, and Jiawei Zhao.
\newblock {GaLore} 2: Large-scale {LLM} pre-training by gradient low-rank projection.
\newblock arXiv 2504.20437, 2025.

\bibitem[Vyas et~al.(2025)Vyas, Morwani, Zhao, Kwun, Shapira, Brandfonbrener, Janson, and Kakade]{vyas2025soap}
Nikhil Vyas, Depen Morwani, Rosie Zhao, Mujin Kwun, Itai Shapira, David Brandfonbrener, Lucas Janson, and Sham Kakade.
\newblock {SOAP}: Improving and stabilizing {S}hampoo using {A}dam.
\newblock In \emph{ICLR}, 2025.

\bibitem[Wang et~al.(2025)Wang, Bhartari, Li, and Perdikaris]{wang2025gradient}
Sifan Wang, Ananyae~Kumar Bhartari, Bowen Li, and Paris Perdikaris.
\newblock Gradient alignment in physics-informed neural networks: A second-order optimization perspective.
\newblock {arXiv} 2502.00604, 2025.

\bibitem[Wang et~al.(2024)Wang, Zhou, Li, and Huang]{wang2024bitshampoo}
Sike Wang, Pan Zhou, Jia Li, and Hua Huang.
\newblock 4-bit {S}hampoo for memory-efficient network training.
\newblock In \emph{NeurIPS}, 2024.

\bibitem[Woo et~al.(2023)Woo, Debnath, Hu, Chen, Liu, Kweon, and Xie]{woo2023convnextv2}
Sanghyun Woo, Shoubhik Debnath, Ronghang Hu, Xinlei Chen, Zhuang Liu, In~So Kweon, and Saining Xie.
\newblock {ConvNeXt V2}: Co-designing and scaling {ConvNets} with masked autoencoders.
\newblock arXiv 2301.00808, 2023.

\bibitem[Xie et~al.(2025)Xie, Wang, Reddi, Kumar, and Li]{xie2025structured}
Shuo Xie, Tianhao Wang, Sashank Reddi, Sanjiv Kumar, and Zhiyuan Li.
\newblock Structured preconditioners in adaptive optimization: A unified analysis.
\newblock {arXiv} 2503.10537, 2025.

\bibitem[Zbontar et~al.(2019)Zbontar, Knoll, Sriram, Murrell, Huang, Muckley, Defazio, Stern, Johnson, Bruno, Parente, Geras, Katsnelson, Chandarana, Zhang, Drozdzal, Romero, Rabbat, Vincent, Yakubova, Pinkerton, Wang, Owens, Zitnick, Recht, Sodickson, and Lui]{zbontar2019fastmri}
Jure Zbontar, Florian Knoll, Anuroop Sriram, Tullie Murrell, Zhengnan Huang, Matthew~J. Muckley, Aaron Defazio, Ruben Stern, Patricia Johnson, Mary Bruno, Marc Parente, Krzysztof~J. Geras, Joe Katsnelson, Hersh Chandarana, Zizhao Zhang, Michal Drozdzal, Adriana Romero, Michael Rabbat, Pascal Vincent, Nafissa Yakubova, James Pinkerton, Duo Wang, Erich Owens, C.~Lawrence Zitnick, Michael~P. Recht, Daniel~K. Sodickson, and Yvonne~W. Lui.
\newblock fast{MRI}: An open dataset and benchmarks for accelerated {MRI}.
\newblock arXiv 1811.08839, 2019.

\bibitem[Zhang et~al.(2025{\natexlab{a}})Zhang, Moniri, Nagwekar, Rahman, Xue, Hassani, and Matni]{zhang2025concurrence}
Thomas~T. Zhang, Behrad Moniri, Ansh Nagwekar, Faraz Rahman, Anton Xue, Hamed Hassani, and Nikolai Matni.
\newblock On the concurrence of layer-wise preconditioning methods and provable feature learning.
\newblock In \emph{{ICML}}, 2025{\natexlab{a}}.

\bibitem[Zhang et~al.(2025{\natexlab{b}})Zhang, Chen, Li, Ding, Wu, Kingma, Ye, Luo, and Sun]{zhang2025adammini}
Yushun Zhang, Congliang Chen, Ziniu Li, Tian Ding, Chenwei Wu, Diederik~P. Kingma, Yinyu Ye, Zhi-Quan Luo, and Ruoyu Sun.
\newblock Adam-mini: Use fewer learning rates to gain more.
\newblock In \emph{ICLR}, 2025{\natexlab{b}}.

\bibitem[Zhao et~al.(2024)Zhao, Zhang, Chen, Wang, Anandkumar, and Tian]{zhao2024galore}
Jiawei Zhao, Zhenyu Zhang, Beidi Chen, Zhangyang Wang, Anima Anandkumar, and Yuandong Tian.
\newblock {GaLore}: Memory-efficient {LLM} training by gradient low-rank projection.
\newblock arXiv 2403.03507, 2024.

\end{thebibliography}