TL;DR. The KL view connects Muon and RMNP to AdaGrad through different choices of Kronecker structure. Our Shared-Shape AdaGrad (SS-AdaGrad) chooses the Kronecker structure that keeps both row scales and within-row correlations while reducing second-moment memory from $O(m^2n^2)$ to $O(m+n^2)$.
Notation. Let $G_t=\nabla_W\ell_t(W_t)\in\mathbb R^{m\times n}$. Write its $i$-th row as $r_{t,i}^{\top}$, where $r_{t,i}\in\mathbb R^n$. We flatten matrices by stacking their rows:
$$ g_t=\operatorname{vec}_r(G_t), \qquad w_t=\operatorname{vec}_r(W_t). $$
This convention gives:
$$ \operatorname{vec}_r(AGB^{\top})=(A\otimes B)\operatorname{vec}_r(G). $$
AdaGrad. Full-matrix AdaGrad stores an $mn\times mn$ second-moment matrix. This requires $O((mn)^2)$ memory.
<aside> 💡
Algorithm 1 (Full-matrix AdaGrad). Choose learning rates $\eta_t>0$. At each step, update
$$ C_t=\sum_{s=1}^{t}g_sg_s^{\top}, \qquad w_{t+1}=w_t-\eta_t C_t^{-1/2}g_t. $$
</aside>
The KL objective. Many methods approximate AdaGrad's second-moment matrix $C$ by a Kronecker product $Q=A\otimes B$, where $A\succ0$ and $B\succ0$. Storing these smaller factors reduces the memory cost. Yet, how should we choose them?
Lin et al. use KL divergence to measure the difference between two zero-mean Gaussian distributions with covariance matrices $C$ and $Q$. Minimizing this objective gives a way to choose the Kronecker factors:
$$ \mathcal K(C,Q) =D_{\mathrm{KL}}\!\left(\mathcal N(0,C)\,\|\,\mathcal N(0,Q)\right) =\frac12\left[\operatorname{tr}(Q^{-1}C)+\log\det Q-\log\det C-mn\right]. $$
Warmup: fix the right factor. Fixing the factor $B=I_n$, we can derive the closed form of $A$.
<aside> 💡
Theorem 1 (The one-sided KL approximation). The unique minimizer of $\mathcal K(C_t,A\otimes I_n)$ over $A\succ0$ is
$$ A_t^*=\frac1n\sum_{s=1}^{t}G_sG_s^{\top}. $$
The corresponding matrix update is
$$ W_{t+1}=W_t-\eta_t(A_t^*)^{-1/2}G_t. $$
</aside>
Proof.
Step 1: simplify the KL objective. With $Q=A\otimes I_n$, we have
$$ Q^{-1}=A^{-1}\otimes I_n, \qquad \log\det Q=n\log\det A. $$
The row-vectorization identity gives
$$ g_s^{\top}(A^{-1}\otimes I_n)g_s =\operatorname{tr}(A^{-1}G_sG_s^{\top}). $$
Summing over the gradient history gives
$$ \begin{aligned} \operatorname{tr}(Q^{-1}C_t) &=\sum_{s=1}^{t}\operatorname{tr}(A^{-1}G_sG_s^{\top})\\ &=n\operatorname{tr}(A^{-1}L), \qquad L=\frac1n\sum_{s=1}^{t}G_sG_s^{\top}. \end{aligned} $$
Terms that do not depend on $A$ do not affect the minimizer. The remaining objective is
$$ \Phi(A)=\frac n2\left[\operatorname{tr}(A^{-1}L)+\log\det A\right]. $$
Step 2: find the minimum. Set $P=A^{-1}$. Then
$$ \Phi(P)=\frac n2\left[\operatorname{tr}(PL)-\log\det P\right], \qquad \nabla_P\Phi=\frac n2(L-P^{-1}). $$
This function is strictly convex in $P\succ0$. Its derivative is zero at $P=L^{-1}$, so the unique solution is $A=L$. Finally, $(A\otimes I_n)^{-1/2}=A^{-1/2}\otimes I_n$ gives the stated matrix update.
End of proof.
Recovering Muon. Keeping only the current gradient recovers the update rule of idea Muon:
<aside> 💡
Algorithm 2 (Ideal Muon without momentum).
$$ U_t=(G_tG_t^{\top})^{-1/2}G_t, \qquad W_{t+1}=W_t-\eta_t U_t. $$
</aside>