Off-Diagonal Induced Metric

The induced metric is derived from a particular ambient inner product on $\mathbb{R}^{N+1}$: block-diagonal, $h = \mathrm{diag}(\gamma I_N, \xi)$. Parameter space and the loss dimension are orthogonal in the ambient view. There is no geometric reason for this choice, and dropping the orthogonality assumption yields a strictly more expressive family of optimizers.

This entry derives the metric, the inverse via Woodbury, the resulting update rule, the combinatorial space of design choices, and the empirical results on MNIST MLP. The headline empirical finding is that the off-diagonal variant beats Adam on best-run accuracy but is dominated by the learnable diagonal on both best and median, suggesting that per-parameter adaptation matters more than off-diagonal coupling at the tested scale.

1. The off-diagonal ambient metric

Replace the block-diagonal ambient metric with

$$h = \begin{pmatrix} \gamma I_N & \xi\, a \\ \xi\, b^\top & \xi \end{pmatrix},$$

where $a, b \in \mathbb{R}^N$ are vectors that introduce off-diagonal coupling between parameter space and the loss dimension. The factor $\xi$ scales the entire loss-dimension block. When $a = b = 0$, $h = \mathrm{diag}(\gamma I_N, \xi)$, the original ambient metric, and we recover the standard induced metric.

The embedding is unchanged: $\varphi(\theta) = (\theta, L(\theta))$ with Jacobian $J = \bigl(I_N;\, \nabla L^\top\bigr)^\top$. The pullback metric is

$$G_{ij} = \gamma\,\delta_{ij} + \xi\bigl(a_i\, \partial_j L + \partial_i L\, b_j + \partial_i L\, \partial_j L\bigr),$$

which in matrix form, writing $\ell = \nabla L$, is

$$G = \gamma I + \xi\bigl(a\,\ell^\top + \ell\, b^\top + \ell\,\ell^\top\bigr).$$

This is a rank-two perturbation of $\gamma I$ in general, rank-one when $a = 0$ or $b = 0$, and rank-zero (returning to $\gamma I$) when $a = b = 0$ and $\ell = 0$. The structure $a\,\ell^\top + \ell\, b^\top$ is a sum of outer products with a shared vector $\ell$, which keeps the inverse tractable.

2. The inverse via Woodbury

Decompose $G = \gamma I + \xi\, U V^\top$ with

$$U = \begin{pmatrix} a & \ell \end{pmatrix} \in \mathbb{R}^{N \times 2}, \qquad V = \begin{pmatrix} \ell & b + \ell \end{pmatrix} \in \mathbb{R}^{N \times 2}.$$

The Woodbury identity gives

$$G^{-1} = \frac{1}{\gamma}\!\left(I - \frac{\xi}{\gamma}\, U\,\bigl(I_2 + \tfrac{\xi}{\gamma}\, V^\top U\bigr)^{-1}\, V^\top\right).$$

The inner matrix $M = I_2 + \tfrac{\xi}{\gamma}\, V^\top U$ is $2 \times 2$. Writing $\alpha = \xi/\gamma$ for brevity,

$$M = \begin{pmatrix} 1 + \alpha\,\langle \ell, a\rangle & \alpha\,\lVert \ell\rVert^2 \\ \alpha\,\langle b + \ell, a\rangle & 1 + \alpha\,\langle b + \ell, \ell\rangle \end{pmatrix}.$$

The inverse uses Cramer's rule: compute $\det(M) = M_{11} M_{22} - M_{12} M_{21}$ and four dot products ($\langle \ell, a\rangle$, $\lVert \ell\rVert^2$, $\langle b + \ell, a\rangle$, $\langle b + \ell, \ell\rangle$). When $\lvert\det(M)\rvert < 10^{-8}$ the metric is near-singular and the implementation falls back to identity scaling to avoid numerical instability.

The total cost per step is dominated by the dot products and the matrix-vector products $U^\top\,\text{(vector)}$ and $V^\top\,\text{(vector)}$, all of which are $O(N)$. No matrix factorisation of $G$ itself is required.

Special cases.

• $a = 0$: $G = \gamma I + \xi\,\ell\,(b + \ell)^\top$ is rank-one. The Sherman-Morrison identity applies directly.
• $b = 0$: $G = \gamma I + \xi\,(a + \ell)\,\ell^\top$ is rank-one. Sherman-Morrison applies directly.
• Both nonzero: full rank-two Woodbury with Cramer's rule.

3. The combinatorial design space

The three vectors in the off-diagonal metric ($\ell$, $a$, $b$) can each be set to one of several observables available at each step. The naming convention uses the pair $(a, b)$.

Choice	Vector	Interpretation
zero	$0$	drops that off-diagonal term
gradient	$g_t$	current minibatch gradient
momentum	$\hat m_t$	bias-corrected momentum (smoothed gradient)
parameters	$\theta_t$	current parameter values

Three "modes" are exposed per vector ($a$, $b$, and the base direction $\ell$), giving $4^3 = 64$ combinations in principle. The implementation registers the most physically motivated subset, with $b$ defaulting to "same as $a$" to reduce the combinatorial explosion. The remaining sweep dimensions are $a \in \{\text{gradient}, \text{momentum}, \text{zero}\}$ with $b = a$, plus the base $\ell$ (typically the gradient).

4. Why the choices matter: the role of $a$ vs $b$, and what kind of object $G$ is

The expression

$$G = \gamma I + \xi\bigl(a\,\ell^\top + \ell\, b^\top + \ell\,\ell^\top\bigr)$$

is not symmetric unless $a = b$. Direct check: the antisymmetric part is

$$G_{ij} - G_{ji} = \xi\bigl[(a_i - b_i)\,\ell_j - (a_j - b_j)\,\ell_i\bigr],$$

which is nonzero whenever $a - b$ is not parallel to $\ell$. A concrete counterexample: with $a = (1, 0)$, $b = (0, 1)$, $\ell = (1, 1)$, $\gamma = \xi = 1$, the formula gives $G = \bigl(\begin{smallmatrix} 3 & 3 \\ 1 & 3 \end{smallmatrix}\bigr)$, which is not symmetric.

Two consequences need to be acknowledged explicitly, because they bear on how the construction should be interpreted.

(i) For $a = b$, the construction is a pullback bilinear form that is symmetric, but symmetry alone is not enough. A Riemannian metric also requires positive-definiteness. The ambient $h = \begin{pmatrix}\gamma I & \xi a \\ \xi a^\top & \xi\end{pmatrix}$ is symmetric when $a = b$. It is positive-definite iff $\gamma > 0$, $\xi > 0$, and the Schur complement is positive: $\xi - \xi^2\,a^\top (\gamma I)^{-1} a > 0$, which under $\gamma, \xi > 0$ reduces to $\xi \cdot \lVert a\rVert^2 < \gamma$ (in the scalar-$\gamma$ case). When this fails the ambient bilinear form is indefinite, and the pullback $G$ can be indefinite as well. The Christoffel symbols, geodesics, and the geodesic convexity machinery apply when all three conditions hold.

(ii) For $a \ne b$, the construction is a non-symmetric bilinear form rather than a metric. The "update rule" $\delta\theta = -\eta\,G^{-1} \nabla L$ is still well defined whenever $G$ is invertible, but the resulting flow is not a Riemannian gradient flow. It is a preconditioned gradient descent with a non-symmetric preconditioner. There is no Riemannian Hessian, no notion of geodesic convexity, and no second-fundamental-form interpretation.

This is the honest framing. The empirical asymmetry between $(a, b) = (\mathrm{momentum}, \mathrm{gradient})$ and $(a, b) = (\mathrm{gradient}, \mathrm{momentum})$ reported in Section 6 is not a feature of the geometry; it is a feature of the non-symmetric preconditioner. When swapping $a$ and $b$ the preconditioner $G$ is replaced by its transpose, which produces a different update rule. The transpose has the same eigenvalues as $G$ but acts on the gradient differently in a way that depends on the gradient's specific direction.

Empirically (Section 6), the asymmetry is large enough that the $a \ne b$ variants are worth keeping as engineering tools. Their theoretical status is more constrained than the symmetric case, and any claims that invoke the geodesic-convexity framework must restrict to $a = b$.

5. Three structural identities

5.1 $a = b = \ell$ recovers the rank-one metric

Setting $a = b = \ell$ (both off-diagonal vectors equal to the base direction) gives

$$G = \gamma I + \xi\,(\ell\,\ell^\top + \ell\,\ell^\top + \ell\,\ell^\top) = \gamma I + 3\xi\,\ell\,\ell^\top.$$

This is the original induced metric, with $\xi$ rescaled to $3\xi$. The rank-one metric is therefore recovered at two points in the $(a, b)$ design space: trivially at $a = b = 0$ (with $\xi$), and along the aligned diagonal $a = b = \ell$ (with $3\xi$).

5.2 Coupled routing of an externally-added regulariser

Given that one has added an L2 regulariser to the loss, the continuous-time Riemannian flow routes the resulting decay step through $G^{-1}$: the decay term $\theta^i$ is contravariant (parameter-direction) while the metric $G_{ij}$ is covariant (cotangent direction), and the natural flow contracts them, giving a step $-\eta\lambda G^{-1}\theta$. This is a consequence of index placement, not a design choice — given the regulariser. The metric itself does not generate the decay; it determines how an externally-added decay is routed.

The current implementations use decoupled weight decay (AdamW-style) for practical reasons: the decay is added separately, outside $G^{-1}$. The Riemannian-correct version is implementable but has not been swept. Whether coupled weight decay would help in practice is an open question; the small-batch implicit-regularisation literature on AdamW suggests the answer is non-obvious and benchmark-dependent.

This section is about routing of an externally-added regulariser. It is logically distinct from §5.3 below, which is about a $\theta$-direction term that the metric itself introduces into the update — with no regulariser added — when the params mode is selected for $a$. The two phenomena have opposite signs and are easy to conflate.

5.3 The params mode introduces structural anti-decay

When $a$ is set to $\theta$ (the params mode of Appendix B), the metric explicitly depends on $\theta$, and the update $\delta\theta = -\eta G^{-1}\ell$ acquires a $\theta$-direction component for free. The structural fact, from the Woodbury form $G = \gamma I + \xi\,U V^\top$ with $U = (a,\ell)$ and $V = (\ell, b+\ell)$:

$$G^{-1}\ell \in \mathrm{span}(a, \ell).$$

The output is a linear combination of the columns of $U$. Therefore a $\theta$-direction component in the update appears iff $a = \theta$. Setting $b = \theta$ alone puts $\theta$ into $V$, where it only multiplies things and never appears as an output direction. The role of $a$ and $b$ is asymmetric in a sharper way than §4's antisymmetric-part argument captures.

When $a = \theta$, the coefficient of $\theta$ in $\delta\theta$ works out to

$$+\eta\,\frac{\alpha\,\lVert\ell\rVert^2}{\gamma\,\det M},$$

where $M = I_2 + (\xi/\gamma)\,V^\top U$ is the inner $2\times 2$ matrix from §2 and $\alpha$ is the corresponding coefficient of $\ell$ in $G^{-1}\ell$. In the positive-definite regime ($\det M > 0$, which is exactly the condition under which $G$ is a Riemannian metric), this coefficient is strictly positive: the update adds a positive multiple of $\theta$, producing weight growth, not decay. The sign can only flip outside the PD regime, where the construction is no longer a Riemannian metric.

Combined with AdamW-style decoupled decay $-\eta\lambda\theta$, the effective weight-decay coefficient becomes

$$\lambda_{\mathrm{eff}} = \lambda - \frac{\alpha\,\lVert\ell\rVert^2}{\gamma\,\det M}.$$

Across the sweep bounds in Appendix A ($\lambda \le 10^{-1}$, $\xi/\gamma$ reaching $10^5$), the structural anti-decay term typically dominates the configured decay, so the $a = \theta$ variants effectively run with net weight growth. The empirical signature is documented in §6.1.

The detailed derivation lives in notebooks/params-mode-analysis.nb.

6. Results on MNIST MLP

Setup: 2-layer MLP (784, 64, 64, 10), 200 epochs, batch size 1024, 5 trials per configuration with Optuna. The same sweep ranges as the project MNIST MLP sweep (Appendix A).

(Provenance: the 5-trial figures in the table below come from a one-off local run recorded in notes/offdiag-benchmark-results.md. The 1000-trial sweep numbers cited in §6.1 are from the project MNIST MLP sweep, results/mnist_mlp/.)

Configuration $(a, b)$	Best	Mean	Notes
$a = \mathrm{momentum},\, b = \mathrm{gradient}$	96.63%	94.03%	most consistent
$a = \mathrm{gradient},\, b = \mathrm{momentum}$	96.43%	86.20%	1 of 5 trials diverged (57.26%)
$a = b = \mathrm{loss\;gradient}$	91.42%	91.42%	only 1 trial
sgd_learn_diag (best baseline)	97.99%	n/a	best peak overall
sgd_metric (rank-one)	97.98%	n/a	close second
Adam	95.89%	95.2%	most consistent baseline

Three observations.

1. Off-diagonal beats Adam on peak in this small probe, and the 1000-trial sweep agrees on best but not on robustness. In this 5-trial sweep, both $(m, g)$ and $(g, m)$ achieve higher peak accuracy than Adam ($96.63\%$ and $96.43\%$ vs $95.89\%$). The improvement is modest ($0.54$ to $0.74$ percentage points). The genuinely rank-two 1000-trial sweep (§6.1) gives both orderings a best of $98.27\%$, slightly above Adam's $98.09\%$ on the same sweep; the means collapse to $54.79\%$ for $(m, g)$ and $48.74\%$ for $(g, m)$, well below Adam's $72.69\%$. The ceiling claim survives the larger sweep; the robustness claim does not.

2. The asymmetry is large. $(a = \mathrm{momentum}, b = \mathrm{gradient})$ achieves a mean of $94.03\%$ across 5 trials. The reversed $(a = \mathrm{gradient}, b = \mathrm{momentum})$ has a mean of $86.20\%$ because one trial of five diverged to $57.26\%$. The other four trials of the reversed configuration are in the $90\%$ to $96\%$ range, comparable to the forward configuration. The conclusion is that the choice of which observable plays the "direction" role and which plays the "cogradient" role matters substantially for training stability.

3. Off-diagonal is dominated by the learnable diagonal at this scale. The best learnable diagonal run ($97.99\%$) is $1.4$ percentage points above the best off-diagonal run ($96.63\%$). The off-diagonal variant adds rank-two correction to the ambient metric but does not give each parameter its own scale factor. The result suggests that per-parameter adaptation matters more than off-diagonal coupling for this task, although a fair direct comparison would require matching the sweep budget (5 trials for off-diagonal versus 1000 for learn-diag).

6.1 The 1000-trial sweep

The project 1000-trial MNIST MLP sweep (results/mnist_mlp/, itr_4, sweep ranges as in Appendix A) covers every registered off-diagonal variant. The three rows that bear on the rank-two construction:

Variant	$(a, b)$	Best	Median	Mean	Worst
sgd_offdiag_m_l	$(\mathrm{momentum}, \mathrm{grad})$	98.27%	57.77%	54.79%	4.78%
sgd_offdiag_l_m	$(\mathrm{grad}, \mathrm{momentum})$	98.27%	40.34%	48.74%	4.99%
sgd_offdiag_l_l*	$(\mathrm{grad}, \mathrm{grad})$	98.37%	41.48%	49.72%	4.42%

Baselines in the same sweep: Adam best $98.09\%$ / median $94.73\%$ / mean $72.69\%$; SGD $98.04\%$ / $78.19\%$ / $58.43\%$; Muon $98.39\%$ / $93.73\%$ / $70.53\%$.

Three structural reads.

1. The $(m,g) / (g,m)$ asymmetry persists at scale. Both orderings hit the same best ($98.27\%$), but the median and mean differ by ${\approx}17$ and ${\approx}6$ percentage points, with $(g, m)$ producing the longer left tail. This matches the qualitative signal from the 5-trial probe in Section 6 and shows the asymmetry is not a small-sample artefact.

2. sgd_offdiag_l_l is not a genuinely off-diagonal experiment (the * row). Per §5.1, $a = b = \ell$ collapses the construction to the rank-one induced metric with $\xi \to 3\xi$. Its 1000-trial best of $98.37\%$ measures the rank-one metric with a particular $\xi$ scaling, not the rank-two mechanism, and should not be cited as evidence for the latter.

3. The params-mode collapse asymmetry matches §5.3. Aggregating over the full set of registered variants, the divergent-trial rate (final accuracy $< 10\%$) separates cleanly by where $\theta$ appears in $(a, b)$:

   Group	Trials	Collapse rate
   $a = \theta$ (4 variants)	4000	35.6%
   $b = \theta$, $a \ne \theta$ (3 variants)	3000	19.4%
   no $\theta$ (8 variants)	8000	19.9%

   The $a = \theta$ group collapses at nearly double the rate of the no-$\theta$ control; the $b = \theta$ group is statistically indistinguishable from it. This is the exact signature predicted by §5.3: $a = \theta$ activates the structural anti-decay term (whose magnitude scales with $\lVert g\rVert^2$ and so amplifies divergent trials), $b = \theta$ alone does not.

The ceiling story for the genuinely rank-two variants holds in the larger sweep ($98.27\%$ vs Adam $98.09\%$), but the robustness gap is in the same direction as the other learnable-metric variants across the sweep: high best, low median, high hyperparameter sensitivity.

7. Position within the design space

The off-diagonal generalisation sits at one end of a hierarchy of extensions to the original rank-one induced metric.

Variant	Ambient metric structure	Pullback rank	Parameter state
Original (rank-one)	$\mathrm{diag}(\gamma I, \xi)$	rank-one perturbation	none
Scalar learnable	$\mathrm{diag}(e^s I, \xi)$	rank-one perturbation	$1$ (scalar $s$)
Diagonal learnable	$\mathrm{diag}(\mathrm{diag}(e^{s_i}), \xi)$	rank-one perturbation	$N$ (vector $s$)
Off-diagonal (this entry)	$\bigl[\begin{smallmatrix}\gamma I & \xi a \\ \xi b^\top & \xi\end{smallmatrix}\bigr]$	rank-two perturbation	$0$ (vectors $a, b$ are observables)
Output embedding (KFAC)	$\mathrm{diag}(I_N, \xi\, \Sigma)$, with $C$-dim output	rank-$C$ perturbation ($J^\top \Sigma J$)	none

The off-diagonal variant occupies a niche: it raises the perturbation rank from one to two without introducing learnable per-parameter state, by using observables ($g_t$, $\hat m_t$, $\theta_t$) as the off-diagonal vectors. This is structurally distinct from the learnable-diagonal direction, which keeps the perturbation rank at one but adds $N$ learnable scalars.

A head-to-head between "more rank" (off-diagonal) and "more state" (learnable diagonal) at matched 1000-trial budget on MNIST MLP: the state-rich variants win on best (sgd_learn_diag $98.41\%$ vs the rank-two off-diagonal best of $98.27\%$) and dominate on median ($86.57\%$ vs $57.77\%$ for $(m,g)$). Whether off-diagonal becomes competitive at larger model sizes where the per-parameter overhead of learnable diagonal becomes prohibitive is open.

8. Connection to the broader curvature framework

Restricting to the symmetric case $a = b$ (Section 4), the off-diagonal construction is a genuine pullback metric, and the geodesic convexity machinery applies. Two further structural facts.

First, $a$ and $b$ are observables (functions of the gradient, momentum, or parameters), not learnable parameters with their own update rule. The metric is therefore parameter-dependent through $a(\theta)$, $b(\theta)$, and $\ell(\theta)$, but not through any auxiliary state. Working out the Christoffel symbols for this $\theta$-dependence is more involved than the loss-embedding case (where only $\ell$ depends on $\theta$): the bracket carries contributions from $\partial_k a$ and $\partial_k b$ as well as $\partial_k \ell$. The full derivation has not been carried out in the geodesic convexity entry; its Appendix B treats only the fixed off-diagonal case (constant $a$, $b$) for the symmetric subset.

Second, in the constant-$a = b$ case (and assuming the Schur condition above so that the construction is actually a Riemannian metric), the off-diagonal correction contributes terms proportional to $H$ in the Riemannian Hessian. The eigenvalue-sign structure of $H$ is preserved up to a possible global flip if the proportionality scalar happens to be negative (a global flip exchanges minima and maxima but does not turn an indefinite Hessian into a definite one). In that case the off-diagonal metric is in the same theoretical class as the rank-one fixed metric: it cannot convert a non-convex landscape into a geodesically convex one.

The variable-$a = a(\theta)$ case is open. Off-diagonal vectors that themselves depend on $\theta$ in principle add Christoffel-symbol corrections that could in principle change the eigenvalue structure. Whether they do is a calculation not yet completed.

The learnable diagonal variant, by contrast, sits in Level 2 of the geodesic convexity hierarchy and provably flips eigenvalue signs. This is one structural reason why the learnable diagonal dominates the off-diagonal variant empirically at the MNIST MLP scale: a richer and proven class of curvature corrections is available.

9. Implementation notes

The canonical implementation in repo/optimisers/jax_offdiag.py (custom_sgd_offdiag) handles the rank-two Woodbury inverse with a $2 \times 2$ Cramer's-rule solve and a near-singular fallback, and applies AdamW-style decoupled weight decay outside $G^{-1}$.

The sweep treats $\gamma$, $\xi$, learning rate, momentum, and weight decay as free hyperparameters (Appendix A), along with the categorical base_mode and use_momentum_for_update. Earlier prototypes fixed $\gamma = 1$ on the grounds that the effective metric strength is $\xi/\gamma$, but the registered sweep keeps both free.

Sweep coverage caveat for the params variants. Per §5.3, when $a = \theta$ the metric itself contributes a $\theta$-direction term to the update with effective coefficient $-\alpha\lVert\ell\rVert^2/(\gamma\det M)$ (anti-decay). Across the swept volume in Appendix A ($\lambda \le 10^{-1}$, $\xi/\gamma$ reaching $10^5$), this structural term typically dominates the configured weight_decay, so the net weight-decay coefficient $\lambda_{\mathrm{eff}}$ is usually negative for the sgd_offdiag_theta_* variants. The sweep therefore does not meaningfully exercise a net-decay regime for these four variants; reaching one would require either widening the weight_decay upper bound or restricting $(\xi, \gamma)$ ranges.

10. What stays open

1. Asymmetry mechanism. The $a$-vs-$b$ asymmetry has two distinguishable pieces. For variants involving the params mode, §5.3 gives a clean mechanistic explanation: $a = \theta$ activates a structural anti-decay term scaling with $\lVert g\rVert^2$, $b = \theta$ alone activates nothing, and the §6.1 collapse-rate signature confirms this. For variants between $(\mathrm{momentum}, \mathrm{gradient})$ and $(\mathrm{gradient}, \mathrm{momentum})$, the divergence between configurations is empirically large but not mechanistically pinned down. The natural hypothesis is that the variance properties of $a\,\ell^\top$ vs $\ell\, b^\top$ differ when $a$ is the momentum (low variance) versus the gradient (high variance), but this has not been tested with a controlled-variance experiment.

2. Coupled weight decay. The Riemannian-correct decay (passed through $G^{-1}$) has not been swept against the AdamW-style decoupled decay. A direct test would clarify whether the structural coupling is empirically beneficial.

3. Scaling to larger models. The MNIST MLP results suggest that learnable diagonal dominates off-diagonal at this scale. Whether off-diagonal becomes competitive at scales where the $N$-dimensional learnable state of sgd_learn_diag becomes a memory or stability bottleneck is open.

4. Combination with learnable diagonal. The off-diagonal mechanism and the learnable diagonal mechanism are structurally orthogonal: one raises the ambient-metric rank, the other adds learnable per-parameter state. A combined variant has not been registered. The closest registered variant is sgd_offdiag_l_l (off-diagonal with $a = b = \mathrm{loss\,gradient}$), which does not include the learnable diagonal.

Appendix A. Sweep configuration details

The MNIST MLP 1000-trial sweep uses the unified hyperparameter bounds from repo/parameters/optimizer_registry.py (single source of truth for both the Optuna and W&B backends). The same bounds apply to every registered off-diagonal variant; there are no per-task overrides. The ranges are confirmed against the raw itr_4 configs.

Parameter	Range	Scale
learning_rate	$[10^{-6}, 10^{1}]$	log-uniform
momentum	$[0, 0.99]$	uniform
xi ($\xi$)	$[10^{-3}, 10^{1}]$	log-uniform
gamma ($\gamma$)	$[10^{-4}, 10^{1}]$	log-uniform
weight_decay	$[10^{-6}, 10^{-1}]$	log-uniform
base_mode	$\{\mathrm{grad}, \mathrm{momentum}\}$	categorical
use_momentum_for_update	$\{\mathrm{True}, \mathrm{False}\}$	categorical

The $(a, b)$ pair is fixed per variant rather than swept inside an Optuna trial: e.g. sgd_offdiag_m_l is $(a, b) = (\mathrm{momentum}, \mathrm{grad})$, sgd_offdiag_l_m is $(\mathrm{grad}, \mathrm{momentum})$, sgd_offdiag_l_l is $(\mathrm{grad}, \mathrm{grad})$. The full set of 16 registered $(a, b)$ pairs (each over $\{0, \mathrm{grad}, \mathrm{momentum}, \mathrm{params}\}^2$ minus the redundant $(0, 0)$) is enumerated in repo/parameters/optimizer_registry.py; each is a separate Optuna study with the bounds above.

Appendix B. Mode semantics

The four modes for the off-diagonal vector $a$ (and $b$) correspond to the following physical observables.

Mode	Vector	Captures
momentum	$\hat m_t$	smoothed gradient direction
grad	$g_t$	current minibatch gradient direction
params	$\theta_t$	current parameter values
zero	$0$	drops the $a\,\ell^\top$ term

With $b = a$ (the default in the sweep), the metric becomes

$$G = \gamma I + \xi\,(a\,\ell^\top + \ell\, a^\top + \ell\,\ell^\top).$$

The symmetric case $a = b = \ell$ collapses to the rank-one metric with $\xi \to 3\xi$ (Section 5.1). The case $a = 0$ recovers the rank-one metric exactly. The case $a \in \{\mathrm{momentum}, \mathrm{grad}\}$ with $b = a$ is the substantive new symmetric variant. The params mode is qualitatively different: when $a = \theta$, the metric depends explicitly on the parameter, and the update acquires a $\theta$-direction component for free (a structural anti-decay; see §5.3).

Appendix C. Connection to nondiag-preconditioner literature

The off-diagonal generalisation is related to but distinct from the broader "nondiag-preconditioner" literature on neural-network Hessians. The classical observation (see the scenic-route note nondiag-preconditioner-scenic-route.md in the project repository) is that neural-network Hessians have substantial off-diagonal mass, so any diagonal preconditioner is structurally limited. The off-diagonal metric here is a different construction: it adds off-diagonal terms to the ambient metric on $\mathbb{R}^{N+1}$, which pull back to the rank-two correction on parameter space. It does not address the off-diagonal entries of the Hessian directly; it constructs a different preconditioner.

The structural answer to off-diagonal Hessian mass is the Kronecker factorisation discussed in the pullback unification entry, which targets per-layer Kronecker structure $H \approx A \otimes B$. The off-diagonal variant here targets a different signal (correlations between parameter directions and the loss covector), and the two are complementary in principle.

Appendix D. Attached materials

The canonical optimiser and sweep code live in the project repository, not in this entry. This entry attaches only the figure-reproduction script and the symbolic-derivation notebooks.

Scripts (scripts/):

• scripts/make_figures.py: reproduces every figure shown in this entry (PDF and PNG written to figures/). Reads the project-level results/mnist_mlp/sgd_offdiag_{m_l,l_m}/itr_4/run_*.csv.

Notebooks (notebooks/):

• notebooks/off-diagonal-derivations.nb: Mathematica derivation and numerical sanity-checks for §§1-5.1 (pullback, Woodbury inverse, Sherman-Morrison limits, antisymmetric part, $a = b = \ell \to 3\xi$ collapse).
• notebooks/params-mode-analysis.nb: Mathematica derivation backing §5.3 — master Woodbury formula, the $G^{-1}\ell \in \mathrm{span}(a,\ell)$ structural fact, numerical confirmation across all 15 non-trivial $(a,b)$ mode pairs, closed-form $\theta$-coefficient in $\delta\theta$, sign analysis in and out of the PD regime, interaction with AdamW decoupled decay, and the 1000-trial collapse-rate summary feeding §6.1.

Canonical sources (in the project repository, not duplicated here):

• repo/optimisers/jax_offdiag.py — custom_sgd_offdiag (the implementation referenced in §9, with decoupled weight decay).
• repo/parameters/optimizer_registry.py — unified hyperparameter bounds and the enumeration of registered $(a, b)$ variants (source of truth for Appendix A).
• repo/parameters/sweep_mnist_mlp.py, sweep_cifar.sbatch, sweep_shake.sbatch — the Optuna / W&B sweep drivers.

References

• Harvey, T. R. (2025). The optimiser hidden in plain sight: training with the loss landscape's induced metric. arXiv:2509.03594.
• Sherman, J., Morrison, W. J. (1950). Adjustment of an inverse matrix corresponding to a change in one element of a given matrix. Annals of Mathematical Statistics 21, 124-127.
• Woodbury, M. A. (1950). Inverting modified matrices. Memorandum Report 42, Statistical Research Group, Princeton.