Induced Metric from Loss Embedding

The induced metric is a Riemannian metric on parameter space, obtained by embedding $\theta \mapsto (\theta, L(\theta))$ into a $(N{+}1)$-dimensional ambient space with a block-diagonal metric and pulling back. It is the mathematical foundation of the Induced-Metric Optimiser (IMO), introduced by Harvey (2025), and serves as the base construction that every variant in the project extends: scalar and diagonal learnable metrics, off-diagonal rank-two perturbations, the Newton-target derived optimizer, and the pullback unification that recovers Gauss-Newton, Fisher, natural gradient, and KFAC as special cases.

1. Setup

Let $\theta \in \mathbb{R}^N$ be the parameters of a model with loss $L : \mathbb{R}^N \to \mathbb{R}$.

1.1 The loss embedding

Define

$$\varphi : \mathbb{R}^N \to \mathbb{R}^{N+1}, \qquad \varphi(\theta) = (\theta^1, \dots, \theta^N,\, L(\theta)).$$

The image of $\varphi$ is the graph of the loss, that is, the loss landscape visualised as a literal hypersurface sitting in one dimension above parameter space.

1.2 The ambient metric

Equip $\mathbb{R}^{N+1}$ with the block-diagonal metric

$$h = \mathrm{diag}(\gamma\, I_N,\, \xi),$$

where $\gamma > 0$ scales the parameter directions and $\xi > 0$ couples the loss direction. Setting $\xi = 0$ collapses the loss direction and recovers flat Euclidean geometry on parameter space.

2. The induced metric

The Jacobian of $\varphi$ is

$$\frac{\partial \varphi^\alpha}{\partial \theta^i} = \begin{pmatrix} I_N \\ \nabla L^\top \end{pmatrix},$$

and the pullback of $h$ is

$$G_{ij} = h_{\alpha\beta}\,\frac{\partial \varphi^\alpha}{\partial \theta^i}\,\frac{\partial \varphi^\beta}{\partial \theta^j} = \gamma\, \delta_{ij} + \xi\, (\partial_i L)(\partial_j L),$$

that is, in matrix form,

$$\boxed{\; G(\theta) = \gamma\, I + \xi\, \nabla L\, \nabla L^\top.\;}$$

$G$ is a rank-one perturbation of $\gamma I$ in the gradient direction. This single structural fact controls every downstream property.

3. Inverse and the scaling factor

The Sherman-Morrison formula with $A = \gamma I$ and $u = v = \sqrt{\xi}\, \nabla L$ gives

$$\boxed{\; G^{-1} = \frac{1}{\gamma}\left(I - \frac{\xi\, \nabla L\, \nabla L^\top}{\gamma + \xi\, \lVert \nabla L\rVert^2}\right). \;}$$

Applied to the gradient,

$$G^{-1}\, \nabla L = \frac{\nabla L}{\gamma + \xi\, \lVert \nabla L\rVert^2}.$$

The Riemannian gradient direction is parallel to $\nabla L$, and the only effect of the metric on the gradient flow is a scalar damping factor

$$r(\theta) = \frac{1}{\gamma + \xi\, \lVert \nabla L(\theta)\rVert^2}.$$

Both $G$ and $G^{-1}$ act on vectors in $O(N)$ time, never materialising the dense matrix.

4. Geometric interpretation

Distances under $G$ are distances measured along the loss-landscape surface, not in flat parameter space. In steep regions ($\lVert \nabla L\rVert$ large) the embedding is nearly vertical, so a small parameter step corresponds to a large surface distance; the metric inflates the cost of moving in the gradient direction, and the optimiser takes a smaller step. In flat regions ($\lVert \nabla L\rVert \approx 0$) the embedding is nearly horizontal, $G \approx \gamma I$, and the optimiser behaves like ordinary SGD.

This is a smoothed variant of gradient clipping: instead of hard-clipping $\lVert \nabla L\rVert$ at a fixed threshold, the induced metric continuously dampens through the denominator $\gamma + \xi\, \lVert \nabla L\rVert^2$. The two regimes meet smoothly at $\lVert \nabla L\rVert \sim \sqrt{\gamma/\xi}$.

5. Riemannian gradient flow

The continuous-time gradient descent under $G$ is

$$\dot\theta = -G^{-1}\nabla L = -\frac{\nabla L}{\gamma + \xi\, \lVert \nabla L\rVert^2}.$$

The trajectory direction is identical to plain gradient flow; only the time parameterisation differs. The flow is a time-rescaled gradient flow, with the rescaling factor $r(\theta)$ depending on the current gradient norm.

A consequence: every critical point of $L$ is a critical point of the induced-metric flow, and vice versa. The induced metric does not change where the optimiser converges, only the speed at which it traverses each region of parameter space. Modifications that do change the critical-point geometry (sign-flipping saddles, basin enlargement, curvature corrections) require richer ambient metrics; see the geodesic convexity entry for the analysis and the learnable diagonal entry for the implementation.

6. Discrete update rule

Let $g_t = \nabla L(\theta_t)$. The practical update rule used in the reference implementation is:

1. Momentum. $m_t = \beta_m\, m_{t-1} + (1 - \beta_m)\, g_t$, bias-corrected $\hat m_t = m_t/(1 - \beta_m^t)$.
2. Metric trace (EMA). $v_t = \beta_v\, v_{t-1} + (1 - \beta_v)\, \xi\, \lVert g_t\rVert^2$, bias-corrected $\hat v_t = v_t/(1 - \beta_v^t)$.
3. Scaling factor. $r_t = 1/(1 + \hat v_t)$ (absorbing $\gamma$ into the learning rate).
4. Update. $\theta_{t+1} = \theta_t - \eta\, r_t\, \hat m_t - \eta\, \lambda_{\mathrm{wd}}\, \theta_t$.

The scaling factor is a scalar: it rescales the entire momentum vector uniformly, unlike Adam which rescales each coordinate independently. This is what makes the cost $O(N)$ per step on top of the gradient.

Reference JAX/Flax implementation: scripts/induced_metric_optimizer.py.

7. Hyperparameters

Parameter	Role	Typical range
$\gamma$	base metric on parameter space	absorbed into $\xi$ (set to $1$)
$\xi$	coupling to the loss dimension	$10^{-2}$ to $10^{1}$
$\beta_m$	momentum decay	$0.9$
$\beta_v$	metric EMA decay	$0.8$ to $0.99$
$\eta$	learning rate	task-dependent
$\lambda_{\mathrm{wd}}$	weight decay	$0$ to $10^{-2}$

8. Variants

The rank-one induced metric in this entry is the base construction. The project explores several extensions, each tracked in its own entry.

• Loss-embedding variants (custom_sgd_log, custom_sgd_rms): change $\xi\, \lVert g_t\rVert^2$ to $\xi\, L(\theta_t)\, \lVert g_t\rVert^2$ or use per-parameter RMS denominators. Cheap rewrites, no new geometric content.
• Scalar and diagonal learnable metrics: replace $\gamma I$ with $\gamma(\theta) = \mathrm{diag}(e^{s_1}, \dots, e^{s_N})$, with the $s_i$ themselves learned. Adds $N$ parameters and a coupled $(\theta, s)$ dynamics.
• Off-diagonal rank-two extension: replaces the rank-one $\nabla L\, \nabla L^\top$ perturbation with the symmetric pair $\xi(a\, a^\top + b\, b^\top)/2$ for observables $a, b \in \{g, m, \theta\}$. Adds curvature coupling without per-parameter state.
• Newton-target rule: derives the (anti-)Newton direction of the learnable-diagonal variant by solving a minimax problem, and explains why the naive variational rule is anti-Newton.
• Output-embedding generalisation: replaces the scalar loss $L$ with a vector residual $f(\theta) = \Sigma^{1/2} J(\theta)$ and recovers Gauss-Newton, Fisher, natural gradient, and KFAC under specific choices of $\Sigma$ and Kronecker structure on $J^\top J$.

The master minimax framework tracks which of these variants admit minimax-optimality theorems that are non-tautological. Each variant's benchmark standing lives in its own entry's empirical section.

9. Empirical results

The base construction has no learnable state: the metric trace is a fixed function of the gradient (Section 6). Two registered optimizers realise it, sgd_metric (the plain rank-one denominator) and sgd_rms (a per-parameter RMS denominator), and both are evaluated against Adam, SGD, and Muon under the project's standard sweeps.

9.1 MNIST MLP (1000-trial Optuna sweep)

Setup: 2-layer MLP (784, 64, 64, 10), batch size 1024, 200 epochs, 1000 trials per optimizer. Raw runs in results/mnist_mlp/<optimizer>/itr_4/.

Optimizer	Best	Median	Mean	Worst
sgd_rms (fixed RMS denominator)	98.30%	94.51%	81.91%	5.79%
sgd_metric (plain rank-one)	98.34%	80.80%	61.94%	3.12%
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%

The notable result is sgd_rms: the simplest fixed-denominator variant is the most robust member of the induced-metric family, with a median ($94.51\%$) essentially matching Adam ($94.73\%$) and a mean above all three baselines. Unlike the learnable variants — which buy a higher best-run ceiling at the cost of median robustness — the fixed denominator does not open a robustness gap.

9.2 Shakespeare MiniGPT

Setup: small transformer, 250 epochs, 200 trials per optimizer; metric is minimum validation perplexity (lower is better). Raw runs in results/shakespeare_minigpt/.

Optimizer	Best	Median
sgd_rms	4.293	12.945
Adam	4.365	12.145
Muon	4.441	13.425
SGD	5.737	28.862

sgd_rms found the best single run, narrowly ahead of Adam; Adam holds the better median. This is a ceiling/stability signal for the simple denominator, not a robust win.

10. The Adam-IMO bolt-on: a registered negative

A falsifiable check: bolt the smooth induced-metric denominator onto Adam by using $\gamma = \mathrm{diag}(\sqrt{\hat v_t} + \varepsilon)$ as the base metric in the Sherman-Morrison update,

$$u_t = \frac{\hat m_t}{\sqrt{\hat v_t} + \varepsilon},\qquad c_t = \frac{1}{1 + \xi\,\sum_i g_{t,i}^2/(\sqrt{\hat v_{t,i}} + \varepsilon)},\qquad \theta \leftarrow \theta - \eta\,c_t\,u_t.$$

Setting $\xi = 0$ recovers Adam exactly. The bolt-on is theoretically unmotivated, since Adam is not a pullback metric (Section 7 of the master minimax framework), but it was registered with pre-specified gating criteria as an honest engineering test.

Setup: CIFAR-10 SmallCNN (~90k params), 30 epochs, batch 128, cosine LR schedule, three seeds $\times$ five learning rates, $\xi \in \{0.01, 0.1, 1.0\}$. Raw runs in results/cifar10_simple_cnn/.

Configuration	Mean test acc	$\Delta$ vs Adam
Adam	82.16%	baseline
Adam-IMO, $\xi = 0.01$	81.80%	$-0.36$ pp
Adam-IMO, $\xi = 0.1$	66.39%	$-15.8$ pp
Adam-IMO, $\xi = 1.0$	49.88%	$-32.3$ pp

All three pre-registered gating criteria fail: the best Adam-IMO mean is $0.36$ pp below Adam (not the $\ge 0.3$ pp above that was required); Adam never diverges in the swept range, so there is no stability envelope to widen; and seed variance is $2.7\times$ Adam's. The denominator acts as a near-constant global learning-rate rescaler ($c_t \approx 0.09$ to $0.17$ throughout training) rather than a stability mechanism, and the win mode the experiment targeted (stability widening at high LR) is unavailable because Adam does not diverge on this task. The bolt-on is a dead end at the tested scale and $\xi$ values; the only untested regime that could still help is a high-LR sweep where Adam itself diverges.

11. What stays open

1. Robustness versus ceiling. The fixed denominator (sgd_rms) is robust but does not raise the ceiling; the learnable variants raise the ceiling but lose robustness. No single induced-metric variant currently dominates Adam on both best-run and median across tasks.
2. The bolt-on's high-LR regime. Whether the smooth denominator helps when bolted onto an optimizer that does diverge at the tested learning rates (the only regime where geometric clipping can demonstrably help) has not been run.
3. Scaling. The benchmarks here are MNIST MLP, a small transformer, and a SmallCNN. The fixed denominator's behaviour at larger scale is untested.

Appendix: Attached materials

Notebooks (notebooks/):

• notebooks/derivations.nb: symbolic verification of the pullback construction, the Sherman-Morrison inverse, $G G^{-1} = I$, the rank-one spectral structure, and the small- and large-$\lVert \nabla L\rVert$ asymptotics of $r$.

Scripts (scripts/):

• scripts/induced_metric_optimizer.py: reference JAX implementation of the Section 6 update rule, with optional momentum and weight decay.
• scripts/make_figures.py: reproduces every figure shown in this entry (PDF and PNG written to figures/).

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.