Pullback Unification: Gauss-Newton, Fisher, and KFAC from Embedding Geometry

The induced metric framework starts from a single construction: pick an embedding of parameter space into a higher-dimensional ambient space, give the ambient space an inner product, and pull it back to a metric on the parameters. The choice of embedding fixes the optimizer.

The original work embedded $\theta$ via the loss, $\phi(\theta) = (\theta, \sqrt{\xi}\,L(\theta))$, which gave a rank-one update of the identity and produced smooth gradient clipping. This entry shows what happens for a different choice: embed via the network output instead of the loss. That embedding reproduces the regularized Gauss-Newton method, Fisher information, the natural gradient, and (under one further approximation) KFAC. The framework is not introducing a new optimizer here. It is providing a single geometric origin for several optimizers that are usually derived separately.

This entry is also the structural-level sequel to the derived diagonal optimizer. That optimizer was clean on diagonal quadratics but failed on real networks because layer Hessians are Kronecker-structured rather than diagonal (its Section 11). The Kronecker minimax theorem of Section 7 is the same minimax argument one structural level up: it lands on KFAC where the diagonal version landed on Newton.

1. The output embedding

Let $f(\theta) \in \mathbb{R}^C$ denote the network output (logits, class probabilities, regression targets, whatever the loss looks at). Define the embedding

$$\phi(\theta) = \bigl(\theta_1,\dots,\theta_N,\;\sqrt{\xi}\,f_1(\theta),\dots,\sqrt{\xi}\,f_C(\theta)\bigr) \;\in\; \mathbb{R}^{N+C}.$$

The differential is

$$D\phi = \begin{pmatrix} I_N \\ \sqrt{\xi}\, J \end{pmatrix}, \qquad J = \frac{\partial f}{\partial \theta} \in \mathbb{R}^{C \times N}.$$

Equip $\mathbb{R}^{N+C}$ with the standard inner product. The pullback metric is

$$G = (D\phi)^\top (D\phi) = I_N + \xi\, J^\top J.$$

This generalises the loss-embedding case, where the rank-one term $\xi\,\nabla L \nabla L^\top$ is replaced by the rank-$C$ term $\xi\,J^\top J$.

2. Inverting the metric via Woodbury

Direct inversion of $G$ is $O(N^3)$, which is prohibitive when $N$ is the parameter count. The Woodbury identity reduces it to a $C \times C$ solve:

$$G^{-1} = \bigl(I_N + \xi\, J^\top J\bigr)^{-1} = I_N - \xi\, J^\top \bigl(I_C + \xi\, J J^\top\bigr)^{-1} J.$$

The core inversion is over a $C \times C$ matrix. For ten-class classification, that is a $10 \times 10$ system, regardless of how many parameters the model has.

The preconditioned gradient is

$$G^{-1} \nabla_\theta L = \bigl(I_N + \xi\, J^\top J\bigr)^{-1} J^\top \nabla_f L,$$

where we used the chain rule $\nabla_\theta L = J^\top \nabla_f L$. The cost per step is: one ordinary backward pass to form $\nabla_f L$ (differentiating the loss with respect to outputs); one vector-Jacobian product to form $J^\top \nabla_f L$ (the parameter-space gradient); $C$ Jacobian-vector products to form the rows of $J J^\top$ explicitly (or $C$ JVPs if $J$ is needed); and the $C \times C$ solve.

3. Spectral structure and step-size bound

Let $J = U \Sigma V^\top$ be the SVD of $J$, with singular values $\sigma_1 \ge \dots \ge \sigma_{\min(C,N)} \ge 0$. The eigenvalues of $G$ are

$$\lambda_k = 1 + \xi\,\sigma_k^2,$$

with multiplicity $N - \min(C,N)$ for the value $1$ corresponding to the kernel of $J$. The condition number is therefore

$$\kappa(G) = \frac{1 + \xi\,\sigma_1^2}{1} = 1 + \xi\,\sigma_1^2.$$

This is finite even when $J$ is rank-deficient, which is the usual case (most parameters cannot independently change every output). Unregularised Gauss-Newton would have $\kappa(J^\top J) = \infty$ on the same problem.

The step-size bound generalises the rank-one loss-embedding bound. Writing $h(\sigma) = \sigma/(1 + \xi\,\sigma^2)$, calculus gives $\max_\sigma h(\sigma) = 1/(2\sqrt{\xi})$, attained at $\sigma = 1/\sqrt{\xi}$. Therefore

$$\bigl\lVert G^{-1} \nabla_\theta L\bigr\rVert \;\le\; \frac{\lVert \nabla_f L \rVert}{2\sqrt{\xi}}.$$

The bound holds regardless of $\lVert \nabla_\theta L\rVert$, so the optimizer cannot take an arbitrarily large step even if a single minibatch produces a very large parameter-space gradient.

4. The weighted embedding and exact Gauss-Newton

For a general loss with output Hessian $\Sigma(\theta) = \partial^2 L/\partial f^2$, the Gauss-Newton Hessian is

$$H_{\mathrm{GN}} = J^\top \Sigma\, J.$$

The unweighted embedding gives $G = I + \xi\, J^\top J$, which corresponds to $\Sigma = I_C$. For mean-squared error this is exact (the loss Hessian in output space is the identity), but for general losses we need to weight the embedding.

Equip output space with the local Mahalanobis metric induced by $\Sigma$ (rather than the Euclidean inner product), evaluated at the current iterate with $\Sigma$ treated as a stop-gradient constant for the purpose of the pullback. The induced metric is then

$$G_w = I_N + \xi\, J^\top \Sigma\, J = I_N + \xi\, H_{\mathrm{GN}}.$$

(Equivalently, replace $f$ by $\Sigma^{1/2} f$ in the embedding only if $\Sigma$ is treated as $\theta$-independent at the point of pullback; differentiating through the $\theta$-dependence of $\Sigma$ would produce extra terms not present in $H_{\mathrm{GN}}$. The construction is therefore a metric-at-a-point rather than a single embedding valid globally.)

This is the regularised Gauss-Newton matrix with Tikhonov parameter $\lambda = 1/\xi$. The identity term provides automatic regularisation on rank-deficient $J$, which is necessary because raw $H_{\mathrm{GN}}$ is singular when the model is overparameterised.

In the limit $\xi \to \infty$ (no regularisation), $G_w \to \xi\,H_{\mathrm{GN}}$ and the update reduces to standard Gauss-Newton. The framework therefore interpolates continuously between SGD ($\xi = 0$, $G = I$) and Gauss-Newton ($\xi \to \infty$, $G \propto H_{\mathrm{GN}}$).

5. Recovery of Fisher and natural gradient

For an exponential-family likelihood $p(y \mid f(\theta))$, the loss is the negative log-likelihood $L = -\log p$. The Bartlett identity states

$$F(\theta) = \mathbb{E}_y\bigl[\nabla_\theta L\, \nabla_\theta L^\top\bigr] = \mathbb{E}_y\bigl[H_{\mathrm{GN}}\bigr],$$

where $F$ is the Fisher information matrix. Choosing $\Sigma$ in the weighted embedding to be the Fisher information in output space gives

$$G_F = I_N + \xi\, F(\theta),$$

which is the regularised natural-gradient preconditioner (Amari 1998). In the $\xi \to \infty$ limit this is pure natural gradient.

This is not a derivation of natural gradient, which already has a standard one (Amari 1998), but it places natural gradient inside the same construction as Gauss-Newton and the loss-embedding clip: all three are pullbacks of a Euclidean ambient metric under different choices of embedding.

6. The Kronecker approximation and KFAC

For a fully-connected layer $W \in \mathbb{R}^{m \times n}$, the Gauss-Newton Hessian admits a Kronecker structure

$$H_{\mathrm{GN}}^{\mathrm{layer}} = A \otimes B,$$

where $A = X^\top X / |\mathcal{B}|$ is the empirical second moment of the layer input (averaged over batch size $|\mathcal{B}|$) and $B = G_o^\top G_o / |\mathcal{B}|$ is the empirical second moment of the per-sample output-gradients (writing $G_o$ for the matrix of output-gradients to avoid collision with the pullback metric $G$ used elsewhere in this entry). This factorisation is exact per sample for the GN Hessian of a single dense layer with no nonlinearity between $W$ and the loss: the per-sample GN Hessian is the Kronecker product of the per-sample input outer product and the per-sample output-Hessian. The batch-averaged version is exact when the per-sample output-Hessian $\Sigma$ is the same across samples (true for MSE), and is an approximation when $\Sigma$ varies sample-to-sample (true for cross-entropy, where $\Sigma$ depends on the per-sample probabilities). KFAC adopts the approximation by averaging $A$ and $B$ separately rather than averaging the per-sample Kronecker products. The approximation is good when the per-layer activations and gradients are weakly correlated across the batch and is the source of the "block-diagonal Fisher" approximation discussed in Martens and Grosse.

A Kronecker-factored metric approximation is any matrix of the form $M = P \otimes Q$ with $P \in \mathrm{SPD}(m)$ and $Q \in \mathrm{SPD}(n)$. The applied preconditioner is $M^{-1} = P^{-1} \otimes Q^{-1}$, and the corresponding update on the layer weight (using the standard vectorisation $\mathrm{vec}(\delta W) = M^{-1}\,\mathrm{vec}(\nabla_W L)$) is

$$\delta W = P^{-1}\, \nabla_W L\, Q^{-1},$$

where the left/right multiplication reflects the standard $y = W x$ orientation with $P$ acting on the output index (left) and $Q$ on the input index (right). The question is which $P, Q$ to choose. The KFAC paper of Martens and Grosse answered this by setting $P = A$ (output-gradient factor, acting on the left) and $Q = B$... wait, this contradicts the input/output assignment used elsewhere in the entry. Convention used in this entry: $A$ is the input second moment ($X^\top X / |\mathcal B|$, acting on the right via $Q^{-1}$), and $B$ is the output-gradient second moment ($G_o^\top G_o / |\mathcal B|$, acting on the left via $P^{-1}$). The applied preconditioner is $\delta W = B^{-1}\, \nabla_W L\, A^{-1}$. (Different KFAC papers use different left/right assignments; we adopt this one for consistency with the Pullback Unification convention $H_{\mathrm{GN}} = A \otimes B$ with $A$ on the input side.) The induced-metric framework gives an alternative answer for which factors to use from minimax optimality of the convergence rate.

7. The Kronecker minimax theorem

The theorem isolates the algorithmic structure of KFAC from the Fisher-approximation argument that originally motivated it.

Theorem. Let $H = A \otimes B$ with $A \in \mathrm{SPD}(m)$ and $B \in \mathrm{SPD}(n)$, and let $\{\lambda_i\} = \mathrm{spec}(P^{-1} A)$ and $\{\mu_j\} = \mathrm{spec}(Q^{-1} B)$ (positive by the SPD hypothesis on $A, B, P, Q$). Among Kronecker-factored preconditioners $M = P \otimes Q$ with $P, Q$ SPD and the volume constraint $\det(P)^n \det(Q)^m = 1$, the minimax contraction factor

$$\min_{P,Q}\,\min_\alpha\,\max_{i,j}\,\bigl|1 - \alpha\,\lambda_i \mu_j\bigr|$$

(the smallest worst-case contraction factor over choices of $P, Q$ and learning rate $\alpha$) is zero, attained uniquely at $P^* \propto A$, $Q^* \propto B$, with $\alpha^* = 1/(\lambda \mu)$ where $\lambda, \mu$ are the (then-equal) eigenvalues of $P^{*-1}A, Q^{*-1}B$.

The proof has four steps.

Step 1. Eigenvalues factorise. $M^{-1} H = (P^{-1} A) \otimes (Q^{-1} B)$, so $\mathrm{spec}(M^{-1} H) = \{\lambda_i \mu_j\}_{i,j}$ where $\{\lambda_i\}$ and $\{\mu_j\}$ are the spectra of $P^{-1} A$ and $Q^{-1} B$ respectively.

Step 2. Product constraint. Setting $p_{ij} = \alpha\,\lambda_i \mu_j$,

$$\prod_{i,j} p_{ij} = \alpha^{mn}\,\bigl(\textstyle\prod_i \lambda_i\bigr)^n\,\bigl(\textstyle\prod_j \mu_j\bigr)^m = \alpha^{mn}\,\frac{\det(A)^n \det(B)^m}{\det(P)^n \det(Q)^m}.$$

Under the volume constraint, $\det(P)^n \det(Q)^m = 1$, so $\prod p_{ij} = \alpha^{mn} \det(A)^n \det(B)^m$. This is a constant independent of the choice of $P, Q$.

Step 3. Equalisation by endpoint argument. With $\prod_{ij} p_{ij}$ fixed (Step 2), suppose $\max_{ij}|1 - p_{ij}|$ is attained simultaneously at the largest $p_{ij}$ (call it $p^+ > 1$) and the smallest (call it $p^- < 1$). Any spread that decreases $p^-$ further or increases $p^+$ further (while preserving the product) strictly increases $\max|1 - p_{ij}|$, because $1 - p^-$ grows when $p^- \to 0$ and $p^+ - 1$ grows when $p^+ \to \infty$. Conversely, contracting the spread (moving every $p_{ij}$ toward the geometric mean while preserving the product) weakly decreases the maximum. The minimum is therefore attained at $p_{ij} = c^{1/(mn)}$ for all $(i, j)$. (The proof is by the same endpoint-equalisation argument used in the diagonal minimax theorem; strict convexity of $x \mapsto |1 - e^x|$ holds only on $\{x > 0\}$, so the cleaner argument is via endpoints rather than via direct second-derivative convexity.)

Step 4. Rank-one factorisation. The constraint $\lambda_i \mu_j = c$ for all $(i,j)$ with $\lambda_i, \mu_j > 0$ forces all $\lambda_i$ equal and all $\mu_j$ equal. Fix $j = 1$: then $\lambda_i = c/\mu_1$ is independent of $i$. Fix $i = 1$: then $\mu_j = c/\lambda_1$ is independent of $j$. So $P^{-1} A = \lambda I_m$ and $Q^{-1} B = \mu I_n$, giving $P = A/\lambda \propto A$ and $Q = B/\mu \propto B$. With $\alpha^* = 1/(\lambda \mu)$, all contraction factors $|1 - \alpha^* \lambda_i \mu_j| = 0$.

The normalising constants are determined by the volume constraint: $P^* = \det(A)^{-1/m} A$ and $Q^* = \det(B)^{-1/n} B$.

The structural parallel with the diagonal minimax theorem (the IMO-84 result that $s_i^* = -\log H_{ii} + \overline{\log H}$ is the unique minimax-optimal log-diagonal scaling) is exact: product constraint, then equalisation. The Kronecker case adds Step 4 (the equalisation forces rank-one factorisation), which is what makes the two factors separately determined.

8. Geometric clipping for the layered update

The pullback metric $G = M + \xi\,g g^\top$ with $M = P^* \otimes Q^*$ admits a Sherman-Morrison inverse:

$$\delta\theta = -\frac{M^{-1} g}{1 + \xi\, g^\top M^{-1} g}.$$

Per layer this reads

$$\delta W = -\frac{A^{-1}\, \nabla_W L\, B^{-1}}{1 + \xi\, s}, \qquad s = \sum_{\ell} \mathrm{tr}\bigl(\nabla_{W_\ell} L^\top\, A_\ell^{-1}\, \nabla_{W_\ell} L\, B_\ell^{-1}\bigr),$$

where the sum runs over layers. The numerator is KFAC. The denominator $1 + \xi\,s$ is the geometric clipping factor inherited from the embedding.

The step-size bound from Section 3 specialises to a bound in the $M$-norm: $\lVert \delta\theta\rVert_M \le 1/(2\sqrt{\xi})$, where $\lVert v\rVert_M^2 = v^\top M v$. The Euclidean norm is then bounded by $\lVert\delta\theta\rVert_2 \le \lVert\delta\theta\rVert_M / \sqrt{\lambda_{\min}(M)}$, which depends on the conditioning of the Kronecker base. For the identity-base loss embedding ($M = I$) the two norms coincide and recover the rank-one bound. Near a minimum, $\nabla L \to 0$ and $s \to 0$, so the denominator approaches $1$ automatically and the clip turns itself off. Far from a minimum, $s$ can be large and the denominator suppresses the step.

9. Why this is not a new optimizer

The base preconditioner $A^{-1} \nabla_W L\, B^{-1}$ is identical to KFAC. The framework reproduces KFAC from a different derivation (minimax over the Kronecker class instead of Fisher approximation) and adds the geometric denominator on top. The contribution is the unified derivation and the clipping, not a new direction of update.

The original KFAC controls step magnitude through a damping term: replace $A$ by $A + \lambda_A I$ and $B$ by $B + \lambda_B I$. This is effective but changes the relative eigenvalue scaling, which changes the preconditioning direction on a generic gradient (adding $\lambda I$ leaves the eigenvectors of a symmetric matrix unchanged, but it compresses the relative weighting between large and small eigenvalues, so the resulting update direction differs). The geometric denominator only scales magnitude, leaving the direction $A^{-1} \nabla_W L\, B^{-1}$ unchanged.

The comparison table.

Property	KFAC	Induced metric with Kronecker base
Preconditioner direction	$A^{-1} \nabla W\, B^{-1}$	$A^{-1} \nabla W\, B^{-1}$ (same)
Magnitude control	external damping, gradient clip	denominator $1/(1+\xi s)$
Step-size bound	none	$\lVert \delta\theta\rVert_M \le 1/(2\sqrt{\xi})$ in the metric norm (proven)
Damping	$\lambda I$ added to factors, distorts direction	$\xi$ controls magnitude only
Near optimum	risk of over-damping if $\lambda$ large	denominator $\to 1$ automatically
Far from optimum	risk of divergence	smooth clip prevents it

10. The unification table

Different embeddings recover several optimizers as special cases of the same construction. Each row is a choice of $\phi$ and the resulting metric.

Embedding $\phi(\theta)$	Induced metric $G$	Resulting optimizer	Exactness
$(\theta,\sqrt{\xi}\,L)$	$I + \xi\,\nabla L\,\nabla L^\top$	Loss-embedding clip	exact
$(\theta,\sqrt{\xi}\,f)$	$I + \xi\,J^\top J$	Regularised GN (MSE)	exact (MSE only)
$(\theta,\sqrt{\xi}\,\Sigma^{1/2} f)$	$I + \xi\,J^\top \Sigma J = I + \xi\,H_{\mathrm{GN}}$	Regularised Gauss-Newton (any loss)	exact
Same, $\xi \to \infty$	$\xi\,H_{\mathrm{GN}}$	Gauss-Newton	exact in limit
Same, $\Sigma =$ Fisher	$I + \xi\,F$	Regularised natural gradient	exact
Layered, $J^\top \Sigma J \approx A \otimes B$	$I + \xi\,(A \otimes B)$	KFAC	exact under Kronecker approx
Layered, then add the SM denominator	$A \otimes B + \xi\,g g^\top$	KFAC + geometric clipping	the contribution here

The geometric interpretation: equipping output space with the Mahalanobis metric from the loss Hessian $\Sigma$ pulls back to the natural-gradient preconditioner on parameter space. The Kronecker approximation pushes the construction down to the per-layer factors $A$ and $B$. The Sherman-Morrison denominator adds smooth clipping at the bottom.

11. Computational cost

For an $N$-parameter model with $C$-dimensional output.

Method	Core size	Cost per step
Woodbury (output embedding)	$C \times C$	$O(N C^2 + C^3)$ plus $C$ JVPs for full Jacobian
Full Gauss-Newton	$N \times N$	$O(N^3)$
KFAC	per-layer $A, B$	$O(N d)$ where $d$ is layer width

For $N = 10^6$ and $C = 10$, the Woodbury form inverts a $10 \times 10$ matrix per step. KFAC inverts per-layer factors whose dimension is the layer width; for typical hidden widths of $64$ to $1024$, this is also tractable.

12. Empirical: KFAC plus geometric clipping on MNIST

A targeted ablation on MNIST MLP (784-64-64-10), full batch on 10k samples, 1000 gradient steps, damping $10^{-4}$, momentum $0.9$. The step is

$$\mathrm{step} = \frac{\mathrm{KFAC\_step}}{1 + \xi\,\lVert \mathrm{KFAC\_step}\rVert_M^2}.$$

Stability table over $\xi$.

$\xi$	Stable lr range	Step-norm bound $1/(2\sqrt{\xi})$
$0$ (plain KFAC)	$[0.01, 0.05]$	$\infty$
$0.01$	$[0.01, 0.1]$	$5.00$
$0.1$	$[0.01, 0.1]$	$1.58$
$1.0$	$[0.01, 1.0]$	$0.50$
$10.0$	$[0.01, 1.0]$	$0.16$

Plain KFAC diverges at $\mathrm{lr} \ge 0.1$. With $\xi \ge 1.0$ the entire swept range up to $\mathrm{lr} = 1.0$ is stable. The stable range widens by a factor of $20$.

Accuracy at $1000$ steps.

Configuration	Best accuracy
Adam, $\mathrm{lr} = 10^{-3}$	$93.3\%$
KFAC + clip, $\xi = 1.0$, $\mathrm{lr} = 0.01$	$83.8\%$
KFAC + clip, $\xi = 10$, $\mathrm{lr} = 0.05$	$81.9\%$
Plain KFAC, $\mathrm{lr} = 0.05$	$75.9\%$

Geometric clipping improves over plain KFAC by about $8$ percentage points but does not reach Adam. The accuracy ceiling appears to be set by the KFAC prototype's factor estimation and damping, not by the clipping mechanism. The clip behaves as a stability mechanism rather than an accuracy enhancer, which is consistent with the prediction from the loss-embedding analysis that the geometric denominator is transient near the optimum.

Clipping dynamics, at $\xi = 1.0$ and $\mathrm{lr} = 0.01$.

• Early (steps $0$ to $50$): clip factor about $0.03$. The step is reduced by a factor of $30$.
• Late (steps $500$ and beyond): clip factor $1.0$. The denominator turns off as the preconditioned gradient norm $g^\top M^{-1} g \to 0$.

For aggressive learning rates ($0.5$, $1.0$), clipping stays active throughout training, which is exactly the role for which the bound was designed.

13. Summary of contributions in this entry

What is reproduced from existing literature:

• Gauss-Newton, regularised Gauss-Newton, Fisher information, natural gradient, and KFAC are all standard optimizers with their own derivations.
• The unification of these as pullback metrics under different embeddings is an information-geometric observation that has appeared (in pieces) in the literature.

What the induced-metric framework adds:

1. A single embedding-based derivation that exhibits all of the optimizers above as instances of one construction (embedding plus pullback) rather than four separate constructions. This is exposition rather than a new theorem; its value is pedagogical and organisational.
2. An alternative derivation of the KFAC factor identity from minimax over the Kronecker class (Section 7), independent of the Fisher-approximation derivation in Martens and Grosse. The minimax instance is tautological in the sense made precise in the master minimax framework (the preconditioner class is rich enough to invert any Kronecker Hessian, so minimax forces the inverse). The contribution is the algorithmic form ($P \propto A$, $Q \propto B$ with the volume gauge fix), not a novel optimality property of KFAC as an optimizer.
3. Geometric clipping as an addition to KFAC (Section 8). The Sherman-Morrison denominator gives a step-size bound $\lVert\delta\theta\rVert \le 1/(2\sqrt{\xi})$ that is proven from the geometry, applies to any embedding in the family, and turns off automatically near a minimum. This is a substantive addition: standard KFAC controls step size through damping ($\lambda I$ added to factors), which distorts the preconditioning direction; the geometric denominator scales magnitude only.
4. A continuous interpolation between SGD, regularised GN, and pure GN through the scalar $\xi$. Useful as a conceptual handle.

Contribution 3 is the only one that adds a new algorithmic mechanism. The other three reorganise existing material, which is valuable for exposition but should not be advertised as algorithmic novelty.

14. Tautology check

The diagonal minimax theorem (s-targets) and the Kronecker minimax theorem (KFAC) are both tautological in the sense made precise in the master minimax framework: the preconditioner class is rich enough to invert every Hessian in the problem class, so the optimum is forced to be the per-instance inverse. This is not a defect of the theorems but a property of the problem-class choice. They still give algorithmic derivations of known optimizers from the embedding framework, and that derivation is what the unification table makes explicit. The non-tautological version of the same minimax (with diagonal preconditioner and rank-one off-diagonal Hessian) is the next step in the program.

15. What stays open

1. Empirical KFAC vs Adam. The asymptotic accuracy of the prototype with geometric clipping does not match Adam on MNIST. Likely causes are damping calibration, factor staleness, and lack of momentum tuning, not the clipping mechanism. A production-quality KFAC implementation with EMA factors and per-block damping is needed before the comparison is meaningful.

2. Empirical evaluation on larger architectures. The Kronecker structure becomes a Kronecker approximation as we move to convolutions, attention, and normalisation layers. Each of these has its own factorisation (see the Shampoo and SOAP literature), and the framework needs to be tested on those.

3. Beyond Kronecker. The minimax theorem in Section 7 is specific to the Kronecker class. The general non-tautological case (preconditioner class poorer than the Hessian class) is what the master minimax framework is set up to address.

Appendix A. Woodbury derivation

Start from $(I + \xi\, J^\top J) X = Y$ and pre-multiply by $J$: $J X + \xi\, J J^\top J X = J Y$, so $(I + \xi\, J J^\top)(J X) = J Y$ and $J X = (I + \xi\, J J^\top)^{-1} J Y$. Substituting back, $X = Y - \xi\, J^\top J X = Y - \xi\, J^\top (I + \xi\, J J^\top)^{-1} J Y$. Therefore

$$\bigl(I + \xi\, J^\top J\bigr)^{-1} = I - \xi\, J^\top \bigl(I + \xi\, J J^\top\bigr)^{-1} J,$$

which is the Woodbury identity for this matrix.

Appendix B. Bartlett identity and Fisher recovery

For an exponential-family likelihood $p(y\mid \eta)$ with natural parameter $\eta = f(\theta)$, the score is $\nabla_\theta \log p = J^\top \nabla_\eta \log p$. The Fisher information is

$$F(\theta) = \mathbb{E}\bigl[\nabla_\theta \log p\,\nabla_\theta \log p^\top\bigr] = J^\top\,\mathbb{E}\bigl[\nabla_\eta \log p\,\nabla_\eta \log p^\top\bigr]\,J = J^\top I(\eta) J,$$

where $I(\eta)$ is the Fisher information in output space. For exponential families $I(\eta) = -\mathbb{E}[\nabla^2_\eta \log p] = \mathbb{E}[\Sigma]$ where $\Sigma = -\nabla^2_\eta \log p = \partial^2 L/\partial f^2$ is the output Hessian of the negative log-likelihood. So $F = \mathbb{E}[H_{\mathrm{GN}}]$, and setting $\Sigma$ to its expectation recovers the natural gradient.

For softmax cross-entropy with $C$ classes, $\Sigma = \mathrm{diag}(p) - p p^\top$ has rank $C - 1$ (the probability simplex constraint removes one degree of freedom). The weighted embedding effectively adds $C - 1$ ambient dimensions, matching the intrinsic dimensionality of output space.

Appendix C. Frobenius distance from $H^{-1}$ to the Kronecker class

When the per-layer GN Hessian is not exactly Kronecker, the KFAC factor identity is the Frobenius projection of the Hessian $H$ onto the Kronecker class $\{P \otimes Q\}$, not of its inverse. The optimal projection is

$$\bigl(P^*, Q^*\bigr) = \arg\min_{P,Q}\,\lVert H - P \otimes Q\rVert_F.$$

The applied preconditioner is then $(P^*)^{-1} \otimes (Q^*)^{-1}$, which is a Kronecker-factored approximation of $H^{-1}$ inherited from the projection on $H$ (not the Frobenius-best Kronecker approximation of $H^{-1}$, which would be a different optimisation). For the Hessian factorisation $H \approx A \otimes B$ used by KFAC, this projection is approximate, and the residual

$$\bigl\lVert H_{\mathrm{GN}} - A \otimes B\bigr\rVert_F$$

is the structural error of the Kronecker assumption. Empirically this is small for the GN Hessian of a single dense layer with a smooth nonlinearity and large batches, and larger for layers with strongly batched-correlated activations or gradients (early convolutional layers, attention output projections). Direct measurement of this residual at scale on real architectures has not been carried out in the current project; the open empirical question is which layers in real models have small residual vs. which are dominated by it.

Appendix D. Spectral structure for rank-deficient $J$

If $\mathrm{rank}(J) = r < \min(C, N)$, the SVD $J = U \Sigma V^\top$ has $r$ positive singular values $\sigma_1, \dots, \sigma_r$ and $\min(C, N) - r$ zero singular values. The eigenvalues of $G$ are

$$\underbrace{1 + \xi\,\sigma_1^2,\,\dots,\,1 + \xi\,\sigma_r^2}_{r \text{ values}},\quad \underbrace{1, 1, \dots, 1}_{N - r \text{ values}}.$$

The $N - r$ unit eigenvalues correspond to directions in parameter space that do not affect the output. In those directions the metric reduces to the identity (no preconditioning, no regularisation), so the gradient flows as if no embedding were applied. The construction therefore does not over-regularise the null space of $J$, which is the usual failure mode of naive Gauss-Newton on overparameterised models.

Appendix E. Attached materials

Notebooks (notebooks/):

• notebooks/output-embedding-gn.nb: symbolic derivation of $G = I + \xi\,J^\top \Sigma J$, Woodbury inverse, spectral analysis, and the unification table.
• notebooks/kronecker-minimax-optimality.nb: full proof of the Kronecker minimax theorem with worked $2 \times 2 \otimes 2 \times 2$ and $3 \times 3 \otimes 4 \times 4$ examples, numerical verification that $\kappa \to 1$.
• notebooks/kronecker-condition-analysis.nb: measurement of per-layer condition numbers under diagonal vs Kronecker preconditioning on MNIST MLP and MiniGPT, showing the $10^7$ to $10^{15}$ reductions.

Scripts (scripts/):

• scripts/kronecker_geometric_clip_benchmark.py: the MNIST MLP benchmark in Section 12.
• scripts/kronecker_condition_analysis.py: the per-layer condition number measurements.
• scripts/make_figures.py: reproduces every figure shown in this entry (PDF and PNG written to figures/).

References

• Amari, S. (1998). Natural gradient works efficiently in learning. Neural Computation 10, 251-276.
• Bartlett, M. S. (1953). Approximate confidence intervals. II. More than one unknown parameter. Biometrika 40, 306-317.
• Gupta, V., Koren, T., Singer, Y. (2018). Shampoo: Preconditioned stochastic tensor optimization. ICML.
• Martens, J., Grosse, R. (2015). Optimizing neural networks with Kronecker-factored approximate curvature. ICML.
• Vyas, N., et al. (2024). SOAP: Improving and stabilizing Shampoo using Adam. arXiv:2409.11321.
• Woodbury, M. A. (1950). Inverting modified matrices. Memorandum Report 42, Statistical Research Group, Princeton.