Induced-Metric Optimizer: Learnable Scalar Metric

Extends the fixed induced-metric optimizers by making the inverse metric a learnable scalar parameter updated online during training.

Motivation

In the fixed metric variants, the base metric on parameter space is $\gamma I$ — a fixed identity scaling. The ratio $\xi/\gamma$ determines how strongly the optimizer responds to curvature, and it must be tuned as a hyperparameter. But the optimal ratio may change during training as the loss landscape evolves. A natural extension is to let the optimizer learn this ratio online.

The Learnable Inverse Metric

Instead of fixing $\gamma^{-1} = I$, parameterize the inverse metric as:

$$\gamma^{-1}_\phi = e^s \cdot I$$

where $s \in \mathbb{R}$ is a single scalar learned during training. Working in log-space ($s$ rather than $e^s$ directly) ensures the inverse metric is always positive definite by construction — no constraints needed. When $s = 0$, $\gamma^{-1} = I$ and we recover the fixed metric.

How the Metric Enters the Update

The inverse metric scales both the trace computation and the momentum:

Preconditioned gradient: $\tilde{g}_t = e^s \cdot g_t$. This is the gradient "seen through" the learned metric.

Metric trace: $v_t = \xi \cdot \langle g_t, \tilde{g}_t \rangle = \xi \cdot e^s \cdot \|g_t\|^2$. The trace now depends on $s$: a larger $e^s$ amplifies the trace, which increases $\hat{v}_t$, which decreases $r_t$. So the learnable metric directly controls how aggressively the induced metric dampens updates.

Preconditioned momentum: $\tilde{m}_t = e^s \cdot \hat{m}_t$. The momentum is also scaled by $e^s$ before the global $r_t$ is applied.

Full update: $\delta\theta_t = -\eta \cdot r_t \cdot \tilde{m}_t - \eta\lambda_{\text{wd}}\theta_t$, where $r_t = 1/(1 + |\hat{v}_t|)$ (plain) or $r_t = L_t/(L_t^2 + |\hat{v}_t|)$ (log-loss).

Online Metric Learning

The scalar $s$ is updated each step by ascending on a local surrogate objective — the trace $v(\phi)$ itself. The intuition: the optimizer wants $\gamma^{-1}$ to maximize the metric trace, because a larger trace means the metric is "seeing" more curvature and providing stronger preconditioning. Regularization prevents $s$ from diverging.

The gradient of the trace with respect to $s$ is:

$$\frac{\partial v}{\partial s} = \xi \cdot e^s \cdot \|g_t\|^2$$

The update rule for $s$ is:

$$s \leftarrow s + \eta_{\text{metric}} \cdot \frac{\partial v}{\partial s} - \eta_{\text{metric}} \cdot \lambda_{\text{reg}} \cdot s$$

The first term drives $s$ upward (stronger preconditioning), while the L2 regularization term $-\lambda_{\text{reg}} \cdot s$ pulls $s$ back toward zero (the identity metric). The balance between these determines the learned metric scale.

After the update, $s$ is clipped to $[-c, c]$ (default $c = 4$, so $e^s \in [e^{-4}, e^4] \approx [0.018, 54.6]$) to prevent extreme values during early training.

Variants

• Plain embedding — scaling $r_t = 1/(1 + |\hat{v}_t|)$
• Log-loss embedding — scaling $r_t = L_t/(L_t^2 + |\hat{v}_t|)$

Cost

O(1) learnable state (one scalar $s$) on top of the base optimizer state. The metric learning step adds negligible computation — one scalar gradient, one scalar update, one clip.