Figure 1. The projection of a curve to the loss landscape from the higher dimensional space. When accounting for the higher dimensional space, the distance travelled is significantly further in highly curved regions.
This repository contains the experiments and implementation for induced metric optimisers. These are preconditioning-based optimisers that utilise the induced metric, which is often illustrated in graphical visualisations of loss landscapes. Specifically, the approach involves pulling back the metric from a higher-dimensional space onto the loss landscape. This can be interpreted as a smoothed variant of gradient clipping.
Implementations in both JAX and PyTorch are included, along with a file demonstrating the implementations are equivilent. All experiments were performed in JAX.
For further details, see the paper: https://arxiv.org/abs/2509.03594
Figure 2. The parameter update vs the gradient, where is a scale that depends on the local curvature.
