Learnable Diagonal Induced Metric and the Curvature-Aware Variant

The induced metric with a fixed base $\gamma = I$ gives a uniform damping of the Euclidean step. The learnable scalar variant replaces the identity by $e^s I$ for a single learnable scalar $s$. The diagonal generalisation goes one step further: give every parameter its own learnable scale factor $e^{s_i}$. This is the most expressive diagonal extension of the original framework, and the variant with the highest peak accuracy on the project's 1000-trial MNIST MLP sweep.

This entry derives the metric, the online learning rule, the Riemannian Hessian of the resulting (now $\theta$-dependent) metric, the hierarchy among scalar and diagonal variants, the optimal curvature correction, the curvature-aware update rule, and the empirical results. The headline structural observation is that the learnable diagonal sits in the highest tier of the geodesic convexity hierarchy: it has enough degrees of freedom to flip eigenvalue signs of the Riemannian Hessian, which the scalar variant can do only when $\mathrm{tr}(H) > 0$ and the fixed metric cannot do at all. The headline empirical observation is that the variant's best-run accuracy is the highest on MNIST MLP but its median is well below Adam's, exposing a robustness gap that is the principal open problem.

1. The learnable diagonal metric

Parameterise the inverse base metric by

$$\gamma^{-1}_\phi(\theta) = \mathrm{diag}\bigl(e^{s_1(\theta)},\dots,e^{s_N(\theta)}\bigr),$$

where $s \in \mathbb{R}^N$ is a vector of per-parameter log-scale factors that evolves under its own learning rule (Section 4). Working in log space ensures positive-definiteness by construction. The induced metric becomes

$$g_{ij}(\theta) = e^{-s_i(\theta)}\,\delta_{ij} + \xi\,\ell_i(\theta)\,\ell_j(\theta), \qquad \ell = \nabla L.$$

Setting $s = 0$ recovers the original induced metric with $\gamma = I$. The base metric $\gamma$ is now $\theta$-dependent through $s(\theta)$, which makes the Christoffel symbols nonzero and produces a curvature correction in the Riemannian Hessian (Section 3).

Sign convention for $s$. This entry parameterises the inverse metric (equivalently, the preconditioner applied to the gradient): increasing $s_i$ inflates the preconditioner and increases the parameter step in direction $i$. The geodesic convexity entry parameterises the metric itself ($\gamma = \mathrm{diag}(e^{s_i})$), so its $s$ has the opposite sign: increasing $s_i$ in that entry's convention means more metric weight and a smaller step in direction $i$. The underlying object is the same; when porting a formula from there to here (or vice versa), substitute $s \to -s$.

2. How the metric enters the update

The preconditioned gradient is

$$\tilde g_{t,i} = e^{s_i}\,g_{t,i},$$

each gradient component scaled by its own factor. The metric trace, which determines the global damping factor of the Sherman-Morrison denominator, is

$$v_t = \xi\,\sum_i g_{t,i}\,\tilde g_{t,i} = \xi\,\sum_i e^{s_i}\,g_{t,i}^2.$$

This is where the global coupling enters: every parameter contributes to the single scalar $v_t$, which determines the global scaling factor $r_t = 1/(1 + v_t)$. A parameter with large $e^{s_i}$ and large gradient dominates the trace, and therefore influences the step size for every other parameter. Adam has no such mechanism: each Adam coordinate is scaled independently.

The full update is

$$\delta\theta_{t,i} = -\eta\,r_t\,e^{s_i}\,\hat m_{t,i} - \eta\,\lambda_{\mathrm{wd}}\,\theta_{t,i},$$

where $\hat m_t$ is the bias-corrected momentum and the last term is decoupled weight decay (AdamW-style; the Riemannian-correct coupled version is implementable but not what the production code uses).

3. The Riemannian Hessian with learnable metric

Metric field versus trajectory state. The Riemannian Hessian formula below assumes that $\gamma$ is an actual metric field on parameter space, i.e., that $s$ is a smooth function $s: \mathbb R^N \to \mathbb R^N$ with well-defined Jacobian $\partial_k s_i$. In the online optimizer, $s_t$ is trajectory state updated step by step from observables; the Jacobian $\partial_k s_i$ is not directly available, and there is no guarantee that the trajectory state is the value of some smooth global $s$ field at $\theta_t$. The pointwise-feasibility analysis of geodesic convexity §4.3 makes this explicit: the SDP/Gershgorin theorems are statements about pointwise first-jet feasibility, not about a global metric field. The bridge from theory to optimizer is therefore: the theory tells us what kind of $\sigma$ would help if it existed; the online rule produces a $\sigma$ from local observables that approximates this in spirit; whether the produced $\sigma$ is the Jacobian of any global $s$ is generally unknown and not strictly required for the optimizer to work.

When $\gamma$ depends on $\theta$, the Riemannian Hessian theorem for fixed metric generalises. The Christoffel bracket $\partial_i g_{jl} + \partial_j g_{il} - \partial_l g_{ij}$ decomposes additively into a $\gamma$-part (involving $\partial_k \gamma_{ij}$) and a $\xi$-part (involving $\xi(H_{ij} \ell_l + \ell_j H_{il})$ and similar). Each piece is treated separately under the Sherman-Morrison inverse $g^{kl} = \gamma^{kl} - \xi\,\ell^k \ell^l/(1 + \xi\,\lVert \ell\rVert_\gamma^2)$.

Define the curvature correction

$$C_{ij} = \Gamma^k_{ij}(\gamma)\,\ell_k, \qquad \Gamma^k_{ij}(\gamma) = \tfrac{1}{2}\,\gamma^{kl}\bigl(\partial_i \gamma_{jl} + \partial_j \gamma_{il} - \partial_l \gamma_{ij}\bigr),$$

where the Christoffel symbols are computed using $\gamma^{kl}$, not $g^{kl}$. A clean algebraic cancellation (verified symbolically in the geodesic convexity entry's Appendix A) gives

$$(\mathrm{Hess}_g L)_{ij} = \frac{H_{ij} - C_{ij}}{1 + \xi\,\lVert \ell\rVert_\gamma^2},$$

where $\lVert \ell\rVert_\gamma^2 = \ell^\top \gamma^{-1} \ell$. The formula is exact, no approximation.

The key structural fact: $C$ depends on $\nabla \gamma$ and $\ell$ but not on $H$. Therefore $C$ can have eigenvalue structure completely different from $H$ and can, in principle, cancel negative eigenvalues, making the corrected $H - C$ positive semidefinite. This is the mechanism that allows the learnable variants to do something the fixed metric cannot: convert saddle directions into convex ones.

The denominator $1 + \xi\,\lVert \ell\rVert_\gamma^2$ is always a positive scalar, so it does not affect eigenvalue signs. The sign structure of the Riemannian Hessian is determined entirely by $H - C$.

4. The online learning rule for $s$

The metric trace serves as a surrogate for the metric's adaptive utility. The per-parameter gradient of the surrogate $v = \xi\,\sum_k e^{s_k} g_k^2$ with respect to $s_i$ is

$$\frac{\partial v}{\partial s_i} = \xi\,e^{s_i}\,g_{t,i}^2.$$

The interpretation. Parameters with large gradient magnitudes get their scale factors pushed upward, increasing their influence on the metric trace and their effective preconditioning strength. Parameters with small gradients have their factors driven toward zero (and thus $e^{s_i}$ toward unity).

The base online update mirrors the scalar case:

$$s_i \leftarrow s_i + \eta_{\mathrm{metric}}\,\xi\,e^{s_i}\,g_{t,i}^2 - \eta_{\mathrm{metric}}\,\lambda_{\mathrm{reg}}\,s_i.$$

Two numerical safeguards.

1. Mean-centering per leaf. $s_i \leftarrow s_i - \bar s$. Without this, the overall scale of $e^{s_i}$ is degenerate with $\xi$: doubling all $e^{s_i}$ and halving $\xi$ gives the same trace. Mean-centering removes the degree of freedom by keeping the geometric mean of the diagonal at $1$.

2. Clipping. $s_i \in [-c, c]$ with $c = 4$ by default. Bounds the condition number of the preconditioner at $e^{2c} \approx 3000$ and prevents the metric from collapsing or diverging in pathological cases.

The base rule has been the production version for most of the project's sweep results. Section 6 below shows that it is structurally suboptimal as a curvature corrector, motivating the curvature-aware variant of Section 7.

5. Comparison with Adam

Property	Learnable diagonal	Adam
Per-parameter scaling	$e^{s_i}$, learned by gradient ascent on $v$	$1/(\sqrt{\hat v_i} + \varepsilon)$, EMA of $g_i^2$
Global coupling	Yes: scalar $r_t$ from $\sum e^{s_i} g_i^2$ damps all parameters	None: each coordinate is independent
Adaptation mechanism	Explicit gradient ascent on a surrogate	Implicit via second-moment EMA
Memory	$O(N)$ for $s$ plus $O(1)$ for the metric trace	$O(N)$ for $\hat v$
Geometric interpretation	Riemannian gradient flow with learned metric	Coordinate-wise rescaling
Curvature correction	Yes (Section 3, nonzero $C$)	No (Adam is not a pullback metric)

The crucial distinction is global coupling. The per-parameter scales feed into a shared scalar $r_t$ that modulates the entire update, so the optimizer can detect when the overall landscape is steep and dampen accordingly, even when individual parameters look tame. The curvature correction $C$ from Section 3 is a second structural advantage: it allows the metric to bend convergence away from sharp regions, a mechanism Adam lacks entirely.

6. The hierarchy of curvature corrections

The geodesic convexity question (when can the metric convert a non-convex landscape into a Riemannian convex one?) splits cleanly across three levels of metric complexity. The geodesic convexity entry gives the full theorem; this section summarises the part directly relevant to the learnable diagonal.

Level 0: constant $\gamma$

$C = 0$ identically. The Riemannian Hessian is the Euclidean Hessian divided by a positive scalar. Eigenvalue signs are preserved. Saddle points cannot be fixed.

Level 1: scalar learnable, $\gamma = e^s I$

The Christoffel symbols are $\Gamma^k_{ij} = \tfrac{1}{2}(\partial_i s\,\delta_{jk} + \partial_j s\,\delta_{ik} - \partial_k s\,\delta_{ij})$. Contracting with $\ell_k$:

$$C_{ij} = \tfrac{1}{2}\bigl(\partial_i s\,\ell_j + \partial_j s\,\ell_i - (\nabla s \cdot \ell)\,\delta_{ij}\bigr).$$

The trace is

$$\mathrm{tr}(C) = \frac{2 - N}{2}\,(\nabla s \cdot \ell).$$

This is dimension-dependent.

• $N = 2$: $\mathrm{tr}(C) = 0$. The correction is trace-free. It can redistribute eigenvalues but not change the total. For an indefinite $H$ with $\mathrm{tr}(H) > 0$, the trace-free $C$ can shift one eigenvalue up and the other down by equal amounts, equalising them at $\mathrm{tr}(H)/2 > 0$ (a clean construction: the optimal $\nabla s$ solves a $2 \times 2$ linear system with determinant $\lVert \ell\rVert^2$, always solvable when $\ell \ne 0$). When $\mathrm{tr}(H) \le 0$, the trace-free constraint forces $\mathrm{tr}(H - C) \le 0$, making positive-definiteness impossible.

• $N > 2$: $\mathrm{tr}(C) = \frac{2 - N}{2}(\nabla s \cdot \ell)$ depends on $\nabla s \cdot \ell$, whose sign the learner controls. The scalar's $N$ degrees of freedom (the components of $\nabla s$) control an $N(N + 1)/2$-dimensional space of symmetric matrices, with the additional structural constraint that $C$ is rank-two plus a scalar multiple of the identity. Sign-flipping becomes increasingly constrained as $N$ grows.

Level 2: diagonal learnable, $\gamma = \mathrm{diag}(e^{s_1}, \dots, e^{s_N})$

The Christoffel symbols of a diagonal metric involve cross-terms with factors $e^{s_i - s_j}$. The exponential anisotropy is what unlocks the trace constraint:

• The trace of $C$ is not constrained to depend only on $\nabla s \cdot \ell$. The $N$ independent log-scales $s_i$ give $N$ independent control directions, and the exponential ratios $e^{s_i - s_j}$ allow asymmetric amplification.
• Numerically verified for $N = 2$ on multiple test functions: the diagonal learnable can flip eigenvalue signs on the symmetric saddle $L = x^2 - y^2$ (where $\mathrm{tr}(H) = 0$ and the scalar cannot succeed), the asymmetric saddle $x^2 - 3 y^2$, the monkey saddle $x^3 - 3 x y^2$, and the Rosenbrock function.
• The constructive Gershgorin argument in the geodesic convexity entry (Appendix C.2) shows that for any symmetric indefinite $H$ and any $\ell$ with all components nonzero, there exists a (diagonal) $\sigma$ such that $H - C \succ 0$. The Frobenius norm of the required $\sigma$ depends on $H$ and $\ell$ (it scales with $\max_i (R_i - H_{ii})/\min_i |\ell_i|$, where $R_i$ is the $i$-th off-diagonal row sum of $H$), so "Frobenius-bounded" is bounded for each instance but the bound is instance-dependent. The construction degenerates on coordinate axes where some $\ell_i = 0$.

Summary

Variant	$C$ structure	Sign-flip capability
Constant $\gamma$	$C = 0$	None
Scalar learnable	rank-two plus scalar $\cdot I$, trace constrained	$N = 2$: iff $\mathrm{tr}(H) > 0$; $N > 2$: increasingly constrained
Diagonal learnable	$N$ independent degrees of freedom, unrestricted trace	All tested $N = 2$ cases; arbitrary $N$ proven in geodesic convexity entry

Each level strictly contains the previous one (setting all $s_i$ equal in the diagonal case recovers the scalar). This hierarchy gives a geometric explanation for the empirical observation that the diagonal learnable outperforms the scalar learnable on best-run ceilings: it has the richest curvature-correction capability available within the framework's diagonal class.

7. The optimal $\nabla s$ and why the base rule misses it

The base online rule (Section 4) was designed to maximise adaptive step-size calibration via the metric trace, not to optimise the curvature correction $C$. As a result, the rule's equilibrium $\sigma_{ij} = \partial s_i / \partial \theta_j$ does not match the structure of the optimal correction.

The optimal $\sigma$ at the canonical saddle

For $L = x^2 - y^2$ at $(1, 1)$, $H = \mathrm{diag}(2, -2)$ and $\ell = (2, -2)$. With $\xi = 1$ and $\gamma = I$, the diagonal optimisation gives the following budget-constrained optima for $\lambda_{\min}(\mathrm{Hess}_g L)$.

Frobenius budget $R$ on $\sigma$	$\lambda_{\min}$	Sign-flip?
$0$	$-0.222$	no
$1.5$	$-0.009$	no
$2$	$+0.055$	yes
$5$	$+0.360$	yes
$10$	$+0.770$	yes

The sign-flip threshold is $R \approx 1.6$. The optimal $\sigma$ at large $R$ has a characteristic anti-correlation structure: $\sigma_{11} \ll 0$ (shrink the scale along the positive-curvature direction) and $\sigma_{22} \gg 0$ (grow the scale along the negative-curvature direction). The asymmetry between $H_{11}$ and $H_{22}$ is what produces the nonzero trace shift that the scalar case cannot achieve.

What the base rule produces

At the same saddle with the base online rule and $\lambda_s = 10$, the equilibrium reasoning is more subtle than a single linearisation because the trajectory state $s_t$ is not literally the gradient of a global potential $s(\theta)$ (see §3 and §10 for the metric-field-vs-trajectory-state distinction). For a fixed-point analysis at the canonical saddle, treat the gradient as constant locally ($\ell_1 = 2, \ell_2 = -2$); the equilibrium is $s_i^* = \xi\ell_i^2/(\lambda_s - \xi\ell_i^2)$, the same value for both coordinates (4/6 = 2/3) because $\ell_1^2 = \ell_2^2 = 4$. With both $s_i^*$ equal, the metric is uniform across coordinates, the Christoffel-correction symmetry forces $C_{11} + C_{22} = 0$, and $H - C$ keeps one positive and one negative eigenvalue. The diagonal degree of freedom is wasted at this symmetric saddle: the rule drives both scales identically because it responds to gradient magnitudes ($\ell_i^2$, equal at this saddle), not to curvature signs ($H_{ii}$, opposite). For asymmetric saddles where $|\ell_i H_{ii}|$ differs across coordinates, the rule produces an asymmetric $s^*$, but the produced $\sigma$ derived from the spatial variation of $s^*$ still does not match the optimal anti-correlation structure of Section 7.

The qualitative conclusion (base rule wastes the diagonal degree of freedom at the symmetric saddle, produces a $\lambda_{\min}$ that is suboptimal but not as bad as no correction) is robust; the exact value $-2/27$ quoted in the table below depends on a specific linearisation that is recorded in the accompanying notebook notebooks/optimal-correction.nb.

Quantitative gap

For the canonical saddle with $\xi = 1$.

Method	$\lambda_{\min}(\mathrm{Hess}_g L)$	Geodesically convex?
No correction (constant $\gamma$)	$-2/9 \approx -0.222$	no
Base online rule ($\lambda_s = 10$, linearised at $\gamma = I$)	$-2/27 \approx -0.074$	no
Optimal $\sigma$ (diagonal, $R = 5$)	$+0.360$	yes

The base rule reduces the negative eigenvalue but never eliminates it. The optimal achieves strict positive-definiteness. The pattern persists across all tested saddles (asymmetric, monkey, Rosenbrock): the base rule is suboptimal, and the optimal with sufficient budget achieves sign-flipping.

8. The curvature-aware update rule

The key insight from Section 7: the update must distinguish directions by their curvature sign, not just by gradient magnitude. The proposed curvature-aware rule, in this entry's "inverse-metric" convention ($\gamma^{-1} = \mathrm{diag}(e^{s_i})$, so increasing $s_i$ means larger step in direction $i$),

$$s_i \leftarrow s_i + \mu\Bigl[\xi\,e^{s_i}\,\ell_i^2 - \beta\,\mathrm{sign}(\hat H_{ii})\,e^{s_i}\,\lvert \ell_i\rvert^2\Bigr] - \mu\,\lambda_s\,s_i.$$

The first term ($\xi\,e^{s_i}\,\ell_i^2$) is the existing adaptive step-size calibration. The second term ($-\beta\,\mathrm{sign}(\hat H_{ii})\,e^{s_i}\,\lvert \ell_i\rvert^2$) creates the anti-correlation we want: it shrinks $s_i$ where curvature is positive (so the step in that direction is smaller, the Newton direction) and grows $s_i$ where curvature is negative (a larger step in the saddle-escape direction). Setting $\beta = 0$ recovers the base rule.

Sign-convention check. The "optimal $\sigma$" structure derived in Section 7 was stated in the form "$\sigma_{11} \ll 0$ in positive-curvature directions, $\sigma_{22} \gg 0$ in negative-curvature directions." That statement uses the geodesic-convexity convention where $\gamma_{ii} = e^{s_i}$ (so $\sigma_{ii} < 0$ means smaller metric weight, smaller step). In this entry's inverse-metric convention, the same physical behaviour ("smaller step in high-curvature directions") corresponds to decreasing $s_i$ when $H_{ii} > 0$. The curvature-aware rule above does decrease $s_i$ when $\mathrm{sign}(\hat H_{ii}) > 0$ (because of the minus sign on the $\beta$ term), so the sign convention is internally consistent. The Section 7 derivation, transposed to this convention, would read "the optimal $\sigma_{ii}$ (under the inverse-metric convention) is positive in positive-curvature directions and negative in negative-curvature directions," that is, $\sigma_{ii}$ co-correlates with $-\mathrm{sign}(H_{ii})$ in this convention.

Estimating $H_{ii}$

The diagonal Hessian entries are approximated by secant differences between consecutive gradient evaluations:

$$\hat H_{ii} \approx \frac{\ell_i^{(t)} - \ell_i^{(t-1)}}{\theta_i^{(t)} - \theta_i^{(t-1)}}.$$

The cost is one extra gradient snapshot ($O(N)$ memory) and $O(N)$ computation per step. No additional backward passes are needed.

Caveats on the secant estimator

1. Bias from off-diagonal Hessian. The secant formula estimates $(H \delta\theta)_i / \delta\theta_i$, not $H_{ii}$. When $H$ has significant off-diagonal entries and the step $\delta\theta$ is not axis-aligned, the estimate is biased. The $\mathrm{sign}$ operation provides partial robustness: only the sign matters, not the magnitude, so moderate bias is tolerable as long as it does not flip the sign. The Hutchinson-based estimator analysed in the derived optimizer has a stronger robustness property ((1+δ)/(1-δ) relative-error bound, independent of dimension and condition number), but it requires the relative error to stay bounded below 1; the secant estimator does not satisfy this hypothesis in practice (the derived-optimizer entry §5.4 shows the initial relative error is an order of magnitude above the theorem's domain), so the formal bound does not apply to secant. Sign-robustness is the only formal guarantee available for the secant version.

2. Near $\hat H_{ii} = 0$. The $\mathrm{sign}$ function is discontinuous and the secant estimate is noisy, so the curvature-sign estimate is unreliable in the transition region. Smoothing via $\tanh(\hat H_{ii}/\tau)$ instead of $\mathrm{sign}$ would help; the derived optimizer entry Section 5 discusses why the $\tanh$ smoothing is insufficient on its own to recover the Newton target.

3. Small steps. When $\theta_i^{(t)} - \theta_i^{(t-1)}$ is small, the secant denominator becomes ill-conditioned. The implementation uses a minimum step-size threshold and EMA smoothing of the estimate.

4. No convergence proof. The curvature-aware rule is motivated by the structure of the optimal $\sigma$ but has not been proven to converge to it. The empirical results in Sections 9 and 10 are the evidence base.

The fully principled Newton target

A more aggressive principled version replaces both terms with the Newton target

$$D_i = -\lambda\,\log\max(\hat H_{ii},\,\varepsilon),$$

which by the diagonal minimax theorem (see the derived optimizer entry) is the unique minimax-optimal log-diagonal preconditioner on diagonal quadratics. This is the basis of the principled diagonal Newton variant explored in the derived optimizer entry. Its empirical fate on real architectures is the cautionary tale of that entry: the principled rule is structurally misspecified against the Kronecker structure of neural-network Hessians, and full-batch ablation shows that more aggressive preconditioning yields worse MNIST accuracy. The lesson is that local diagonal optimality is not enough; the structural class matters more than the choice within the class.

9. Empirical results

9.1 1000-trial MNIST MLP sweep

Setup: 2-layer MLP (784, 64, 64, 10), batch size 1024, 200 epochs, 1000 Optuna trials per optimizer. Raw runs in results/mnist_mlp/<optimizer>/itr_4/. The rows below are the learnable-diagonal variants and the Adam, SGD, and Muon baselines, all from the same sweep.

Variant	Best	Median	Mean	Worst
sgd_learn_diag (base rule)	98.41%	86.57%	64.36%	2.24%
sgd_learn_diag_curv (curvature-aware)	98.39%	78.04%	60.88%	5.36%
sgd_learn_diag_curv_log	98.38%	57.89%	57.27%	4.11%
sgd_learn_diag_log	98.36%	84.63%	62.04%	5.30%
Adam	98.09%	94.73%	72.69%	9.80%
SGD	98.04%	78.19%	58.43%	2.57%
Muon	98.39%	93.73%	70.53%	9.29%

Two things to note.

1. Best-run ceiling. The base learnable diagonal finds the highest single-trial accuracy across all 28 swept optimizers (98.41%, narrowly ahead of Muon and the curvature-aware variants). This is a real ceiling signal: with sufficient hyperparameter search, the variant can reach an accuracy that Adam does not match.

2. Robustness gap. The median is $86.57\%$ for the base rule and $78.04\%$ for the curvature-aware variant. Adam's median is $94.73\%$. The learnable variants are substantially more hyperparameter-sensitive than Adam. The hyperparameter search range includes regions where the variants behave pathologically, which is the source of the long left tail.

The honest summary: the learnable diagonal raises the best-run ceiling but does not produce a robust improvement under random hyperparameter draws.

9.2 Implicit sharpness-aware optimisation

A controlled measurement of Hessian sharpness during training shows that the learnable diagonal reaches matched loss levels at substantially flatter minima than Adam.

Setup: 2-layer MLP (8-16-4, 212 parameters), synthetic 200-sample classification.

Loss threshold	SGD max eig	Adam max eig	Learn-Diag max eig	LD vs Adam
$\le 0.10$	0.7011	4.4479	0.6831	84.6% flatter
$\le 0.05$	0.7011	2.2586	0.6831	69.8% flatter
$\le 0.01$	0.3629	0.6943	0.2946	57.6% flatter

The likely mechanism is a combination of two components, both present in the base learnable diagonal variant (the variant used in this measurement; the curvature-aware variant was not the one tested). First, the damping factor $r = 1/(1 + \lvert v\rvert)$ shrinks the step in directions where $\sum_i e^{s_i} g_i^2$ is large, which empirically correlates with sharper regions, preventing entry into sharp basins. Second, the metric update $s_i \leftarrow s_i + \mu\xi\,e^{s_i}\,g_i^2$ amplifies weight on high-gradient-magnitude directions, which geometrically stretches the step along those directions and (combined with the damping) effectively compresses the step in high-curvature directions where gradients spike sharply. The net behaviour resembles what SAM (Foret et al., 2021) achieves via adversarial perturbation.

This attribution is a hypothesis: the experiment shows that the learnable diagonal reaches flatter minima at matched loss levels, but it does not isolate which mechanism (damping, metric update, or both, or something else) is responsible. The fully principled Riemannian-Hessian curvature correction $C$ from Section 3 is a third candidate mechanism that would apply once the curvature-aware variant is tested with sharpness measurement; that experiment has not been done.

The qualitative comparison.

Property	SAM	Learnable diagonal metric (base)
Mechanism	adversarial perturbation $\max_\delta L(\theta + \delta)$	candidate mechanisms: damping factor $r$ plus learned-diagonal scaling; curvature correction $C$ is a candidate not isolated in this experiment
Extra cost per step	$2\times$ (two forward-backward)	$\approx 1.2\times$ (metric update)
Theoretical backing	PAC-Bayes generalisation bound	Riemannian Hessian formula (Section 3) for the curvature-correction candidate; no formal backing for the damping mechanism
Explicit sharpness penalty	yes	no (implicit through geometry, mechanism not yet isolated)

This is the most plausible positive mechanism for why the learnable diagonal reaches higher peak accuracy than Adam on best-run comparisons, despite the analysis showing that its preconditioning direction is anti-Newton. The hypothesis is that the variant wins not because of better preconditioning (the preconditioning is actively wrong) but because of implicit sharpness regularisation that biases convergence toward flatter, better-generalising minima. Confirming this requires the controlled experiments listed in Section 10 (matched-cost comparison against explicit SAM, scaled-up sharpness measurement, and ablation of the two candidate mechanisms above).

Note on baseline values. The sharpness experiment above uses a 212-parameter network and is a small-scale controlled measurement, not part of the 1000-trial MNIST sweep. The "Adam 98.09%" baseline cited in Section 9.1 (from the 1000-trial Optuna sweep on the 55k-parameter MNIST MLP) is a different setting from the older "Adam 97.17%" baseline that appears in the derived optimizer's full-batch ablation (smaller network, full-batch gradients). The two are not directly comparable. The implicit-SAM result here uses neither; it is its own controlled measurement on a third (smaller) network.

9.3 Small multimodal test functions

A focused sweep on standard 2D test functions (raw runs in results/small_examples_*/, itr_4) isolates where the curvature-aware diagonal helps. On rugged, multimodal landscapes it is genuinely useful; on smooth ones Adam remains strongest.

Function	Best optimizer	Notable result
Rastrigin	sgd_learn_diag_curv	$491/1000$ trials converged vs Adam's $24/1000$
Ackley	sgd_learn_diag_curv	best value $3.6\times10^{-7}$; the log variant sgd_learn_diag_curv_log converged $184/5000$ vs $2/5000$ for the non-log version (earlier itr_3 sweep)
Beale / Himmelblau / Rosenbrock	Adam	strongest on both best value and median; curvature-aware variant converges often but is less robust

The supportable claim is narrow: the curvature-aware diagonal converges far more often than Adam on rugged multimodal functions (Rastrigin, and Ackley with the log embedding) and is competitive but not dominant on smooth functions. The older "log embedding solves Ackley when all baselines fail" framing is too broad and is not quoted as a general result.

10. What stays open

1. Closing the robustness gap. The median accuracy of the learnable variants is well below Adam's on MNIST MLP. The structural causes are not pinned down. Candidate explanations include (i) the curvature-aware secant estimator's noise on the long left tail of the hyperparameter distribution, (ii) the mean-centering operation interacting poorly with weight decay at certain LR-decay combinations, and (iii) the metric clip ($\lvert s_i\rvert \le 4$) being too tight in some regimes and the metric saturating. A diagnostic study (which hyperparameter draws produce the divergent trials, and what feature of the learning dynamics distinguishes them) has not been done.

2. Implicit-SAM at scale. The sharpness measurement in Section 9.2 was on a 212-parameter network. Direct sharpness measurement on a realistic-scale task (MNIST CNN, CIFAR-10 SmallCNN) compared against explicit SAM at matched computational cost would be the next test.

3. Curvature-aware variant vs base rule. Empirically the two variants have the same peak ($98.39\%$ vs $98.41\%$) and similar medians. The theoretical advantage of the curvature-aware rule (anti-correlation with $\hat H_{ii}$, motivated by Section 7) does not translate into a measurable empirical edge on this task. Whether it does on harder landscapes (Rastrigin and other multimodal small functions) is positive: the curvature-aware variant converged in $491/1000$ Rastrigin trials versus Adam's $24/1000$ (Section 9.3).

4. Coupled weight decay. The Riemannian-correct weight decay (passed through $G^{-1}$) is implementable but has not been swept. The off-diagonal metric entry notes the same open question for the off-diagonal variant.

5. Combination with off-diagonal. The learnable diagonal and off-diagonal generalisations are structurally orthogonal (different mechanisms for enriching the metric). A combined variant has not been registered or swept.

Appendix A. Derivation of the Christoffel correction for diagonal $\gamma$

Convention used in this appendix. This appendix derives the Christoffel symbols of the metric $\gamma$ itself, using the convention $\gamma_{ij} = e^{s_i}\,\delta_{ij}$ (the geodesic-convexity-entry convention). This is opposite to the main-text convention of Section 1, which parameterises the inverse metric as $\gamma^{-1} = \mathrm{diag}(e^{s_i})$. To apply any formula from this appendix to the main-text parameterisation, substitute $s_i \to -s_i$ throughout. The derivation is reproduced here in the geodesic-convexity convention because that is the natural setting for the Christoffel calculus; the main text's sign convention is the natural setting for the online optimizer (preconditioner-multiplies-gradient).

For $\gamma_{ij} = e^{s_i}\,\delta_{ij}$ (no sum), the inverse is $\gamma^{ij} = e^{-s_i}\,\delta_{ij}$ and the derivatives are $\partial_k \gamma_{ij} = e^{s_i}\,\partial_k s_i\,\delta_{ij}$. The Christoffel symbols of $\gamma$ are

$$\Gamma^k_{ij}(\gamma) = \tfrac{1}{2}\,\gamma^{kl}\bigl(\partial_i \gamma_{jl} + \partial_j \gamma_{il} - \partial_l \gamma_{ij}\bigr).$$

Case $i = j = k$: $\Gamma^i_{ii} = \tfrac{1}{2}\,e^{-s_i}(\partial_i \gamma_{ii} + \partial_i \gamma_{ii} - \partial_i \gamma_{ii}) = \tfrac{1}{2}\,\partial_i s_i$.

Case $i = j \ne k$: $\Gamma^k_{ii} = \tfrac{1}{2}\,e^{-s_k}(\partial_i \gamma_{ik} + \partial_i \gamma_{ik} - \partial_k \gamma_{ii}) = -\tfrac{1}{2}\,e^{s_i - s_k}\,\partial_k s_i$.

Case $i \ne j$, $k = i$: $\Gamma^i_{ij} = \tfrac{1}{2}\,e^{-s_i}(\partial_i \gamma_{ji} + \partial_j \gamma_{ii} - \partial_i \gamma_{ij}) = \tfrac{1}{2}\,\partial_j s_i$.

Case $i \ne j$, $k = j$: $\Gamma^j_{ij} = \tfrac{1}{2}\,\partial_i s_j$.

Case $i, j, k$ all distinct: $\Gamma^k_{ij} = 0$.

The exponential ratios $e^{s_i - s_k}$ in the off-diagonal terms are what generate the asymmetric amplification mentioned in Section 6: when $s_i \gg s_k$, the contribution of $\partial_k s_i$ to $\Gamma^k_{ii}$ is exponentially amplified. This is the structural ingredient that lets the diagonal learnable achieve sign-flipping that the scalar cannot.

The full $C_{ij} = \Gamma^k_{ij}\,\ell_k$ for the diagonal case is a sum over the five cases above, contracted with $\ell_k$. The diagonal entries (used in the trace identity discussion in the main text) come from cases (i) and (ii),

$$C_{ii} = \tfrac{1}{2}\,\partial_i s_i\,\ell_i \;-\; \tfrac{1}{2}\sum_{k \ne i} e^{s_i - s_k}\,\partial_k s_i\,\ell_k,$$

and the off-diagonal entries ($i \ne j$) come from cases (iii) and (iv),

$$C_{ij} = \tfrac{1}{2}\,\bigl(\partial_j s_i\,\ell_i + \partial_i s_j\,\ell_j\bigr).$$

The full case-by-case expression is reproduced in the geodesic convexity entry's Appendix C.

Appendix B. Trace formula for scalar $\gamma$

For $\gamma_{ij} = e^s\,\delta_{ij}$ with scalar $s$, the Christoffel symbols are

$$\Gamma^k_{ij}(\gamma) = \tfrac{1}{2}\bigl(\partial_i s\,\delta_{jk} + \partial_j s\,\delta_{ik} - \partial_k s\,\delta_{ij}\bigr).$$

Contracting with $\ell_k$:

$$C_{ij} = \tfrac{1}{2}\bigl(\partial_i s\,\ell_j + \partial_j s\,\ell_i - (\nabla s \cdot \ell)\,\delta_{ij}\bigr).$$

The trace is

$$\mathrm{tr}(C) = \sum_i C_{ii} = \tfrac{1}{2}\bigl(\nabla s \cdot \ell + \nabla s \cdot \ell - N\,(\nabla s \cdot \ell)\bigr) = \frac{2 - N}{2}\,(\nabla s \cdot \ell).$$

The factor $(2 - N)/2$ is the source of the dimension-dependent trace constraint discussed in Section 6. For $N = 2$ the correction is trace-free. For $N > 2$ the factor $(2 - N)/2$ is negative, so Cauchy-Schwarz gives $|\mathrm{tr}(C)| \le \frac{N - 2}{2}\,\lVert \nabla s\rVert\,\lVert \ell\rVert$ as an upper bound on the magnitude of the (sign-controllable) trace shift. A positive trace shift requires $(\nabla s)\cdot \ell < 0$; a negative shift requires $(\nabla s)\cdot \ell > 0$. The bound on magnitude constrains how much shift the scalar learnable can produce in either direction.

Appendix C. Why $r = 1/(1 + v)$ damps sharp directions

The damping factor $r_t = 1/(1 + v_t)$ with $v_t = \xi\,\sum_i e^{s_i}\,g_{t,i}^2 \ge 0$ has a geometric interpretation. Near (but not at) a sharp minimum, at a comparable displacement from the basin centre, the gradient is large in the high-curvature directions (the directions in which $H_{ii}$ is large), and the metric weight $e^{s_i}$ tracks the gradient magnitude (the metric update pushes $s_i$ up where $g_i$ is large). Both factors combine: $e^{s_i}\,g_{t,i}^2$ is large in directions that are both sharp and currently being traversed quickly. The sum produces a large $v_t$, which shrinks the step.

This is structurally what the Sherman-Morrison denominator on the loss embedding does for the rank-one metric: it bounds the step at $1/(2\sqrt{\xi})$ regardless of the parameter-space gradient. With a learnable diagonal underneath, the bound becomes direction-aware: the directions that contribute most to $v_t$ (sharp and active) are the directions that get most damped.

Combined with the candidate curvature correction $C$ from Section 3 (which would, if it were the active mechanism, bias the Riemannian Hessian to be more positive-definite), the qualitative picture is implicit sharpness-aware behaviour: the metric prevents entry into sharp basins (via $r_t$) and (under the curvature-correction mechanism) bends the effective curvature toward flatter regions (via $C$). The empirical sharpness measurement in Section 9.2 is consistent with these mechanisms but does not isolate them; the curvature-aware variant was not the one tested, so $C$ is a candidate mechanism rather than a confirmed one.

Appendix D. Attached materials

Scripts (scripts/):

• scripts/jax_learn_diag.py: JAX implementation of the base learnable diagonal optimizer.
• scripts/jax_learn_diag_curv.py: implementation with the curvature-aware update rule.
• scripts/mnist_mlp_sweep.py: Optuna sweep entry point for the MNIST MLP results.
• scripts/sharpness_comparison.py: the Section 9.2 sharpness measurement.
• scripts/make_figures.py: reproduces every figure shown in this entry (PDF and PNG written to figures/).

Notebooks (notebooks/):

• notebooks/optimal-correction.nb: symbolic derivation of the optimal $\sigma$ for sign-flipping, including the canonical-saddle worked example.
• notebooks/hierarchy-verification.nb: numerical verification that the diagonal learnable achieves sign-flipping where the scalar cannot, across the four test saddles.

References

• Foret, P., Kleiner, A., Mobahi, H., Neyshabur, B. (2021). Sharpness-aware minimization for efficiently improving generalization. ICLR.
• Harvey, T. R. (2025). The optimiser hidden in plain sight: training with the loss landscape's induced metric. arXiv:2509.03594.
• Kingma, D. P., Ba, J. (2014). Adam: A method for stochastic optimization. ICLR.
• Keskar, N. S., Mudigere, D., Nocedal, J., Smelyanskiy, M., Tang, P. T. P. (2017). On large-batch training for deep learning: generalization gap and sharp minima. ICLR.