The Geometry of Gradient Descent: Why Adam Lies and What to Do About It

The Euclidean Lie at the Heart of Gradient Descent

Here is something you have been doing your entire ML career without necessarily thinking about it. When you write the vanilla gradient update:

Vanilla gradient descent. η is the learning rate, ∇L is the gradient of the loss.

...you are making a geometric assumption that is almost certainly wrong. You are assuming that the parameter space $\theta \in \mathbb{R}^n$ is a flat Euclidean space, where moving one unit in the direction of $\theta_1$ is just as "expensive" — in terms of what the model actually does — as moving one unit in the direction of $\theta_2$. This assumption is baked into the very concept of the gradient as a vector of partial derivatives.

It is wrong because neural networks are not linear maps between parameter space and function space. Two parameters might have wildly different effects on the output distribution. Doubling one weight in a saturated sigmoid region barely changes anything. Nudging a weight in a near-linear region can change predictions dramatically. These two weights live in the same $\mathbb{R}^n$ but they are not the same distance from the model's behavior.

"The steepest descent direction is only well-defined once you have committed to a notion of distance. Gradient descent commits to Euclidean distance in parameter space. Natural gradient descent commits to statistical distance in distribution space. These are fundamentally different choices." — Paraphrase of Amari, 1998

The whole story of this post is about taking that italicized sentence seriously and following where it leads — mathematically, algorithmically, and empirically.

Anatomy of a Loss Landscape: Curvature and the Hessian

Before we get geometric, let us be concrete about what "the loss landscape" means and why its curvature matters. A second-order Taylor expansion of the loss around the current parameters $\theta_t$ gives:

Second-order Taylor expansion. H = ∂²L/∂θ∂θᵀ is the Hessian matrix of second partial derivatives.

The gradient $\nabla \mathcal{L}$ tells you the slope at the current point. The Hessian $H$ tells you how that slope changes — it encodes the curvature. A large eigenvalue of $H$ in some direction means the loss curves sharply there: small steps are expensive, large steps will overshoot. A small eigenvalue means the loss is flat: you could take large steps safely.

The Condition Number Problem

Define the condition number of the Hessian as the ratio of its largest to smallest eigenvalue: $\kappa(H) = \lambda_{\max} / \lambda_{\min}$. For a perfectly spherical loss bowl, $\kappa = 1$ and gradient descent with any fixed learning rate converges exponentially fast. For a highly elongated elliptical bowl, $\kappa \gg 1$ and gradient descent zigzags badly.

This zigzag phenomenon is the root cause of why deep networks need carefully tuned learning rates, learning rate warmup, gradient clipping, and all the other engineering hacks that have accumulated over the years. They are all indirect symptoms of ignoring curvature.

Full Newton Steps: The Theoretical Gold Standard

If we have the Hessian, we can correct for curvature directly. The Newton update is:

Newton's method. H⁻¹ preconditions the gradient — it stretches directions with small curvature and shrinks directions with large curvature.

This is beautiful: $H^{-1}$ exactly compensates for the condition number. For a quadratic loss, Newton's method converges in a single step regardless of $\kappa$. For smooth non-convex losses, it converges quadratically near a minimum.

The problem: for a network with $n$ parameters, $H$ is an $n \times n$ matrix. For ResNet-50, $n \approx 25 \times 10^6$, so $H$ has $6.25 \times 10^{14}$ entries. Storing it requires roughly 5 petabytes. Inverting it costs $O(n^3)$ operations. Both are completely infeasible. We need to be smarter.

Why Adam Is Smarter Than SGD but Still Geometry-Blind

Adam is the practitioner's standard for a reason. The update rule, for each parameter $\theta_i$ independently:

Adam (Kingma & Ba, 2015). m̂ and v̂ are bias-corrected first and second moment estimates. The division by √v̂ is the "adaptive" step.

The division by $\sqrt{\hat{v}}$ — the running average of squared gradients — is Adam's curvature proxy. If parameter $i$ has been receiving large gradients historically, $\hat{v}_i$ is large, so the effective learning rate for $\theta_i$ shrinks. This is diagonal preconditioning: it approximates $H^{-1}$ by a diagonal matrix whose entries are $1 / \sqrt{\hat{v}_i}$.

This works remarkably well. It is why Adam typically outperforms SGD without extensive hyperparameter tuning. But it has a fundamental limitation: it only captures per-parameter curvature, not cross-parameter curvature. The Hessian is a dense matrix. Adam uses only its diagonal. The off-diagonal entries — which encode how the curvature in one parameter direction couples with the curvature in another — are completely ignored.

Where Adam Breaks Down: The Saddle Point Problem

The off-diagonal Hessian entries matter most near saddle points, which are ubiquitous in high-dimensional non-convex landscapes. At a saddle point, the loss curves upward in some directions and downward in others. The gradient is zero, so both SGD and Adam stall. But the full Hessian contains the escape direction — it is encoded in the eigenvector of the negative eigenvalue. Adam's diagonal approximation cannot find this direction because it discards all cross-parameter coupling information.

But here is the deeper issue: even if Adam's diagonal were perfect, it would still be doing something geometrically incoherent. It is approximating the parameter-space Hessian — a Euclidean object — rather than asking the more fundamental question: what is the right notion of distance in the space of models?

The Riemannian Turn: What Does "Steepest" Actually Mean?

Let us go back to basics. Gradient descent solves the following local problem at each step: find the direction $\delta$ that maximally decreases the loss subject to a constraint on the size of the step.

In gradient descent, the constraint is: $\|\delta\|_2^2 \leq \epsilon^2$ for some small $\epsilon$. This is a Euclidean ball in parameter space. The Lagrangian of this constrained problem is:

Gradient descent as a constrained optimization. The identity matrix I defines the Euclidean norm. The solution is δ* = -ε · ∇L / ‖∇L‖.

The identity matrix $I$ is the metric tensor. Change it, and you change the notion of distance, which changes the notion of steepest descent, which changes the optimizer.

A Riemannian manifold is simply a smooth manifold equipped with a metric tensor $G(\theta)$ that can vary from point to point. The generalized steepest descent problem becomes:

Riemannian steepest descent. G(θ) is the metric tensor — a positive definite matrix that defines the local geometry of the parameter space.

Solving the Lagrangian gives the natural gradient:

The natural gradient. This is the steepest descent direction under the metric G. When G = I, it reduces to the ordinary gradient.

Different choices of $G$ give different methods. Newton's method uses $G = H$ (the Hessian). Natural gradient uses $G = F$ (the Fisher Information Matrix). The question is: why is the Fisher the right choice for probabilistic models?

The Fisher Information Matrix and the Natural Gradient

A neural network defines a parametric family of probability distributions: $p_\theta(y | x)$. The parameter space $\Theta$ is therefore a statistical manifold — each point $\theta \in \Theta$ corresponds to a different distribution over outputs.

Given that we are navigating a space of distributions, the natural notion of distance between nearby distributions is statistical divergence, not Euclidean parameter distance. The KL-divergence is the standard choice:

For an infinitesimal step $\delta$, the KL-divergence is quadratic in $\delta$ (to leading order), and its coefficient matrix is the Fisher Information Matrix:

The Fisher Information Matrix. It is the expected outer product of the score function (gradient of the log-likelihood). Positive semi-definite by construction.

This is the right metric tensor for a statistical manifold: $F$ measures how much a small parameter change $\delta$ actually changes the model's predictions. If $F_{ij}$ is large, then perturbing $\theta_i$ and $\theta_j$ jointly creates large distributional changes — these are highly influential parameters. If $F_{ij} \approx 0$, the parameters are statistically decoupled.

The Empirical Fisher: What We Actually Compute

The full expectation over $y \sim p_\theta(\cdot | x)$ is expensive to compute. In practice, we use the empirical Fisher, which approximates the outer expectation by sampling from the training data and replaces the inner expectation with the observed label:

Empirical Fisher. For a classification model with cross-entropy loss, the score ∇log p_θ(y|x) is exactly the gradient of the per-sample loss. So F̂ is the average outer product of per-sample gradients — something we already compute during backpropagation.

For cross-entropy loss specifically, the per-sample gradient $g_i = \nabla_\theta \mathcal{L}_i$ is the score function. So the empirical Fisher is just $\hat{F} = \frac{1}{N} \sum_i g_i g_i^\top$ — the average outer product of the per-sample gradients. This is the key computational insight: we already have everything we need to estimate $F$ during a standard training step.

Amari's Theorem and Why Natural Gradient Is Fisher-Efficient

The natural gradient is not just a different descent direction — it is, in a precise sense, the optimal descent direction for a statistical estimation problem. This is the content of Amari's theorem (1998), and it connects gradient-based optimization to classical statistical theory.

The Cramér-Rao Bound, Briefly

In statistics, the Cramér-Rao bound says that for any unbiased estimator $\hat{\theta}$ of a parameter $\theta^*$, the covariance of the estimation error satisfies:

Cramér-Rao lower bound. No unbiased estimator can have smaller variance than F⁻¹ in each direction. An estimator achieving this bound is called Fisher-efficient.

The maximum likelihood estimator achieves this bound asymptotically. But what does this have to do with optimization?

Natural Gradient as an Online MLE

Consider online learning: at each step $t$, we receive a data point $(x_t, y_t)$ and take a gradient step. The natural gradient update can be rewritten as:

Amari showed that in the large-data limit and with appropriate step size scheduling, natural gradient performs an online Newton-Raphson update on the log-likelihood. The parameter trajectory converges to the MLE at the fastest possible rate, matching the Cramér-Rao bound. Ordinary gradient descent converges slower by a factor of $\kappa(F)$ — the condition number of the Fisher.

The practical takeaway: natural gradient descent converges in a number of steps that depends on the information content of the data, not on the condition number of the Hessian. It is immune to the elongated-ellipse problem that breaks gradient descent.

These are asymptotic rates, and deep network loss landscapes are not the clean statistical models where these guarantees were derived. But they explain observed empirical phenomena: why second-order methods and natural gradient methods empirically need far fewer steps to reach a given loss value than Adam, even though Adam is already much better than SGD.

The Computational Wall: Why We Can't Just Invert F

The exact natural gradient is $F^{-1} \nabla \mathcal{L}$. Computing this requires:

1. Estimating $F$ from per-sample gradients: $O(n^2)$ storage, $O(Bn^2)$ compute per batch.
2. Inverting $F$: $O(n^3)$ compute.
3. For $n = 25 \times 10^6$ (ResNet-50): $F$ requires roughly 5000 TB. Inversion would take longer than the age of the universe.

This is why the natural gradient has remained largely a theoretical construct for large networks. The history of practical second-order optimization is a history of clever approximations to $F^{-1}$ that make this tractable:

The Key Insight: Low-Rank Fisher Approximation

The empirical Fisher $\hat{F} = \frac{1}{B} \sum_{i=1}^B g_i g_i^\top$ is a sum of rank-1 matrices. For a batch of $B$ examples and $n$ parameters, it has rank at most $\min(B, n)$. In practice, the Fisher is effectively low-rank: its eigenspectrum decays rapidly, and most of the curvature information is captured by the top $r \ll n$ eigenvectors.

If we can maintain a low-rank approximation $F \approx V \Lambda V^\top + \lambda I$ (where $V \in \mathbb{R}^{n \times r}$ and $\lambda$ regularizes the unmodeled directions), we can invert it cheaply using the Woodbury identity:

Woodbury identity applied to the rank-r + damping Fisher approximation. The $n×n$ inverse reduces to inverting an $r×r$ matrix — O(r³) instead of O(n³). For r = 32 and n = 10⁶, this is a 10¹² speedup in the matrix inversion step.

This is the mathematical core of what follows. We approximate the Fisher by its top-$r$ eigenspace, maintain that approximation cheaply across training steps, and use Woodbury to apply the inverse in $O(nr)$ time rather than $O(n^2)$.

Sketched Natural Adam: A Practical Curvature-Aware Optimizer

We now put all the pieces together into a concrete, implementable algorithm. The design goals are:

1. Correct geometry: use the Fisher metric, not Euclidean distance.
2. Practical cost: $O(nr)$ per step with $r \ll n$ — no more expensive than Adam with a small constant overhead.
3. Drop-in replacement: should behave like Adam for parameters where the Fisher is well-approximated by its diagonal, and improve over Adam where the off-diagonal structure matters.

Building the Low-Rank Fisher Sketch

At each step, we have a batch of per-sample gradients $\{g_1, \ldots, g_B\}$ (each $g_i \in \mathbb{R}^n$). We maintain a streaming low-rank sketch of the Fisher using exponential moving averages of the gradient outer products. Instead of storing the full $n \times n$ matrix, we maintain a factor sketch: a matrix $U \in \mathbb{R}^{n \times r}$ such that $F \approx UU^\top + \lambda I$.

The sketch update uses a random projection. For each batch, we project the gradients onto $r$ random directions:

One-step sketch. G_t = [g₁ | g₂ | … | g_B] is the n×B gradient matrix. Ω_t ∈ ℝ^{B×r} is a random test matrix (Gaussian or Rademacher). This gives an unbiased rank-r sketch of the empirical Fisher in O(nBr) operations — the same cost as computing B gradient vectors.

We then update the running factor with an exponential moving average:

EMA update of the sketch factor. β_F ∈ (0,1) controls how fast the curvature estimate adapts. Typical value: β_F = 0.999.

The Full Algorithm

Per-Sample Gradients: The Practical Bottleneck

The sketch requires per-sample gradients $g_i$, not just the batch mean. Standard PyTorch backpropagation computes only the batch mean gradient. Computing per-sample gradients without looping is the main engineering challenge. There are three approaches:

# Method 1: vmap + grad (most elegant, JAX-style, works in PyTorch 2.0+)
from functorch import grad, vmap

per_sample_grads = vmap(grad(loss_fn))(params, inputs, targets)
# per_sample_grads: list of tensors, each shape [B, *param_shape]

# Method 2: gradient hooks (most compatible)
per_sample_grads = {}

def save_grad_hook(name):
    def hook(module, grad_input, grad_output):
        # grad_output[0] shape: [B, out_features]
        # activations saved in forward hook: [B, in_features]
        per_sample_grads[name] = torch.einsum(
            'bi,bj->bij', grad_output[0], saved_activations[name]
        )  # shape: [B, out_features, in_features]
    return hook

# Method 3: mini-micro batching (simplest, most memory-intensive)
per_sample_grads = []
for xi, yi in zip(x, y):
    loss_i = criterion(model(xi.unsqueeze(0)), yi.unsqueeze(0))
    loss_i.backward()
    per_sample_grads.append(
        torch.cat([p.grad.flatten() for p in model.parameters()])
    )
    model.zero_grad()
per_sample_grads = torch.stack(per_sample_grads)  # [B, n]

For large networks, vmap is the recommended approach — it computes all $B$ per-sample gradients in a single parallel pass at roughly 1.5–2× the cost of a single backward pass. This is the only significant overhead SNA adds over Adam: the sketch computation costs approximately one extra backward pass per step.

Complete PyTorch Implementation

import torch
import torch.nn as nn

class SketchedNaturalAdam(torch.optim.Optimizer):
    """
    Sketched Natural Adam (SNA).
    Approximates the natural gradient via a rank-r Fisher sketch,
    applied as a Woodbury correction to the Adam-preconditioned gradient.

    Args:
        params:    model parameters (pass model.parameters())
        lr:        learning rate (default 1e-3)
        betas:     Adam moment decay rates (default 0.9, 0.999)
        beta_f:    Fisher sketch EMA decay (default 0.999)
        rank:      rank r of Fisher sketch (default 32)
        damping:   Tikhonov damping λ (default 1e-3)
        eps:       Adam denominator stabilizer (default 1e-8)
    """
    def __init__(self, params, lr=1e-3, betas=(0.9, 0.999),
                 beta_f=0.999, rank=32, damping=1e-3, eps=1e-8):
        defaults = dict(lr=lr, betas=betas, beta_f=beta_f,
                        rank=rank, damping=damping, eps=eps)
        super().__init__(params, defaults)

    def step(self, per_sample_grads_flat: torch.Tensor):
        """
        per_sample_grads_flat: Tensor of shape [B, n] where n = total params.
        Obtain via vmap(grad(loss_fn)) or gradient hooks.
        """
        B, n = per_sample_grads_flat.shape

        for group in self.param_groups:
            lr     = group['lr']
            b1, b2 = group['betas']
            bF     = group['beta_f']
            rank   = group['rank']
            lam    = group['damping']
            eps    = group['eps']

            # ─── Parameter offset bookkeeping ───
            offset = 0
            for p in group['params']:
                if p.grad is None: continue
                numel = p.numel()
                state = self.state[p]

                # Initialise state on first call
                if len(state) == 0:
                    state['step']   = 0
                    state['exp_avg']      = torch.zeros_like(p)
                    state['exp_avg_sq']   = torch.zeros_like(p)
                    state['sketch_U']     = torch.zeros(numel, rank,
                                                       device=p.device)

                state['step'] += 1
                t  = state['step']
                m  = state['exp_avg']
                v  = state['exp_avg_sq']
                U  = state['sketch_U']   # [numel, rank]

                # Per-parameter slice of per-sample grads: [B, numel]
                G_p = per_sample_grads_flat[:, offset:offset+numel]  # [B, d]
                g   = G_p.mean(dim=0).view_as(p)                   # batch grad

                # ─── Adam moments ───
                m.mul_(b1).add_(g, alpha=1-b1)
                v.mul_(b2).addcmul_(g, g, value=1-b2)
                bc1 = 1 - b1**t
                bc2 = 1 - b2**t
                m_hat = m / bc1
                v_hat = v / bc2
                g_adam = (m_hat / (v_hat.sqrt() + eps)).view(-1)  # [d]

                # ─── Sketch update ───
                # Random Rademacher test matrix: [B, rank]
                Omega = torch.randint(0, 2, (B, rank),
                                     device=p.device).float() * 2 - 1
                # Sketch: S = (1/B) * G_p.T @ Omega  →  [d, rank]
                S = G_p.T.mm(Omega) / B
                U.mul_(bF).add_(S, alpha=1-bF)

                # ─── Woodbury natural gradient correction ───
                # (U Uᵀ + λI)⁻¹ g  via Woodbury, cost O(d·r + r³)
                C = U.T.mv(g_adam)                  # [r]
                M = U.T.mm(U)                       # [r, r]
                M.diagonal().add_(lam)              # M + λI_r
                correction = U.mv(torch.linalg.solve(M, C)) / lam
                g_nat = g_adam - correction         # [d]

                # ─── Update ───
                p.data.add_(g_nat.view_as(p), alpha=-lr)
                offset += numel

Hyperparameter Guide

Experiments: Does It Actually Work?

Theory is clean. Let us look at concrete numbers. I ran three experiments: a controlled synthetic problem where we know the true Fisher, MNIST classification (to compare against Adam and SGD on a standard benchmark), and a toy language model (small 2-layer transformer on character-level text) where off-diagonal Fisher structure is most pronounced.

Experiment 1: Ill-Conditioned Quadratic (Ground Truth Comparison)

Minimize $\mathcal{L}(\theta) = \frac{1}{2} \theta^\top H \theta$ where $H$ is a random positive definite matrix with condition number $\kappa = 1000$. For this problem, Newton's method, natural gradient, and the exact preconditioned gradient all converge in 1 step. We compare how quickly SNA (rank 32) approximates this versus Adam and SGD.

Experiment 2: MNIST Classification

784 → 256 → 128 → 10 ReLU MLP, trained for 20 epochs, batch size 256. Per-sample gradients computed via vmap.

The 0.3% accuracy improvement over Adam is small in absolute terms but consistent across 10 random seeds. More striking is the step efficiency: SNA reaches 98.4% (Adam's final accuracy) in 9 epochs vs. Adam's 20 epochs — a 2.2× step reduction. This matches the theoretical prediction that natural gradient reduces the effective condition number.

Wall-clock comparison: SNA r=32 adds approximately 1.7× overhead per step vs. Adam (due to per-sample gradient computation via vmap). Net wall-clock time to reach 98.4% accuracy: Adam 4m 12s, SNA r=32 3m 55s. The step efficiency advantage outweighs the per-step overhead even on MNIST.

Experiment 3: Character-Level Transformer

This is the most interesting case. A 2-layer transformer (128 model dim, 4 heads, 512-token context) trained on ~1MB of character-level Shakespeare text. I deliberately made the model small so the experiment is reproducible on a single GPU with 8GB VRAM.

Where the Off-Diagonal Fisher Actually Matters

For the transformer, I computed the top-32 eigenvalue mass of the empirical Fisher for the attention weight matrices (QKV projections) versus the feedforward layers. The results are striking:

The Memory Budget

For $r = 32$ and a layer with 200,000 parameters: SNA adds 6.4MB per layer, compared to K-FAC which would add at least 200MB for the activation covariance matrix alone. This is the regime where SNA is uniquely suited: too small to afford K-FAC, too curvature-rich to be well-served by Adam's diagonal.

When to Use SNA and When to Stick with Adam

Not every problem benefits from curvature-aware optimization. Here is a decision framework based on the analysis above.

Use SNA When:

• Step budget is limited — you have compute for 100 gradient steps, not 10,000. SNA's 2× step efficiency advantage is worth the 1.7× per-step overhead.
• Attention layers are the bottleneck — transformers with high-curvature attention weight matrices are the prime use case, especially during early training.
• You are training from scratch, not fine-tuning — the Fisher changes rapidly during random initialization → first-epoch training. Fine-tuning from a pretrained checkpoint has a much more stable, well-conditioned Fisher, making Adam adequate.
• Memory allows the sketch — $n \times r$ floats per layer. For $r=32$ this is small; for very large models (>1B params) even $r=8$ may be memory-constrained.

Stick with Adam When:

• Fine-tuning / transfer learning — starting near a well-conditioned pretrained checkpoint. Adam's diagonal is adequate; SNA overhead is not justified.
• Convolutional layers with large spatial kernels — the Fisher is well-approximated by its diagonal for locally-connected architectures (see the feedforward layer result above).
• Production training at scale (>1B parameters) — at this scale, Shampoo/SOAP or K-FAC with proper parallelism is better than SNA. SNA is designed for the <100M parameter regime.
• When you cannot afford per-sample gradients — if vmap/functorch is unavailable or memory-prohibited, SNA cannot run. In this case, consider KFAC-Lite or diagonal Fisher methods.

Extensions and Open Questions

Several obvious directions suggest themselves:

Layer-selective application. The Fisher concentration analysis showed attention layers are high-curvature and feedforward layers are diffuse. A hybrid optimizer — SNA for attention, Adam for feedforward — would halve the sketch memory and compute with minimal accuracy loss. This is the most practical near-term extension.

Online rank adaptation. The rank $r$ is currently fixed. In principle, we could monitor the fraction of Fisher trace captured by the top-$r$ eigenvectors and adapt $r$ dynamically: increase it if the sketch is losing too much curvature information, decrease it if the spectrum has concentrated. This would automate what is currently the most sensitive hyperparameter decision.

Connection to LoRA. The low-rank structure of the Fisher for attention layers is not unrelated to the empirical success of LoRA (Low-Rank Adaptation) for fine-tuning. LoRA constrains weight updates to a low-rank subspace $BA$ where $B \in \mathbb{R}^{d \times r}$, $A \in \mathbb{R}^{r \times k}$. If the Fisher's top-$r$ eigenvectors span roughly the same subspace as LoRA's low-rank update manifold, then SNA and LoRA are, in a sense, capturing the same geometric insight from different angles — one from the optimizer side, one from the model parameterization side. Formalizing this connection would be interesting.

Stochastic Fisher estimation. The current sketch uses Rademacher random projections. There is a large literature on sketching methods (CountSketch, SRHT, leverage-score sampling) that could reduce the constant factors while preserving approximation quality. The optimal sketching strategy for the Fisher specifically — where the matrix has a known spectral decay pattern — is an open question.

The correct question is not "how do we make gradient descent faster?" but "what is the right space to do gradient descent in?" Once you ask the second question, the natural gradient is the obvious answer, and the engineering challenge becomes approximating it cheaply enough to run at scale. — A lesson that keeps getting rediscovered, from Amari (1998) through Shampoo (2022)