Estimating the Local Learning Coefficient

3.1 Preconditioned SGLD

One limitation of standard SGLD is that it makes isotropic moves: the Gaussian noise is spherical and the gradient step ignores differences in curvature across directions. In singular models, this makes us take tiny steps in stiff directions and explore slowly in flat ones.

A simple, affordable amelioration is to precondition. We introduce a positive-definite matrix \(A\) — think of the Adam quasi-curvature estimates, which give a diagonal approximation to \(A\). The update becomes:

\[ w_{t+1} = w_t - \tfrac{\epsilon}{2}\,A\big(\beta^*\,\widehat{\nabla} L_n(w_t) + \gamma(w_t-w^*)\big)\;+\;\xi_t,\quad \xi_t \sim \mathcal N(0,\,\epsilon A). \]

Now both the drift and the noise are scaled by \(A\). Intuitively, steps are larger in flat directions and smaller in sharp ones, with the noise scaled accordingly. This respects the desired invariant distribution (for constant \(A\)), makes exploration more efficient and preserves invariances such as weight rescaling in ReLU nets.

On paper, unadjusted Langevin with minibatches should make us nervous. In practice, SGLD tends to work because the LLC setup stacks the deck in our favour:

• It’s local. The quadratic \(\tfrac{\gamma}{2}\|w-w^*\|^2\) keeps the chain near \(w^*\). We don’t need to cross modes or map the whole loss landscape.
• It’s hot. \(\beta^*=1/\log n\) is small for big \(n\). Energy barriers are softened; local exploration is easier.
• It’s cheap. Minibatch gradients mean each step costs about an SGD step, so long chains are feasible.
• It’s tolerant. We only need \(\mathbb{E}[L_n(w)]\) to be accurate under the local law. We don’t require immaculate global mixing.
• It may be approximating nonparametric variational Bayes in some useful sense. See Hoffman and Ma (2020).

Intuitively, I think of the estimator as “poke the basin around \(w^*\) and see how quickly the loss rises.” If the basin is broad (singular/degenerate), the average loss near \(w^*\) barely increases and \(\widehat{\lambda}\) is small; if it’s sharp, the loss rises faster and \(\widehat{\lambda}\) is larger. SGLD seems smart at doing the poking, at least in the models I’ve tried.

We can show it works in practice for at least one model family: the deep linear network. There, we can test the estimates against ground truth. Indeed, the SGLD-based estimator matches the ground truth beautifully, even at big scales (hundreds of millions of parameters and datasets we can’t fully load into memory).

That’s surprising to me, and impressive. I think we have to be careful not to oversell it. DLNs are special: their singularities are structured in a way we can analyze, and the loss surface has symmetries that play nicely with the Gaussian driving noise.

In contrast, in nonlinear nets—with ReLU, GeLU, and friends—we don’t have an analytic ground truth for \(\lambda(w^*)\). All we can do is compare different estimators (SGLD vs MALA vs diagnostics), but they may share biases. It could be that SGLD is “working great” in DLNs precisely because they’re unusually well‑behaved, and that in truly nonlinear models the same algorithm is only approximately right.

So, while the DLN results give me confidence that the estimator isn’t nonsense, they shouldn’t be read as a general proof of correctness in deep nonlinear nets. This gap—we can’t validate LLC estimates in the models we actually care about—is, to me, concerning.

3.2 SGLD shortcomings

Note that there’s some unintuitive (to me) stuff going on here, even in tractable deep linear networks. The “landscape” has millions of dimensions, so “poking” in the right directions is hard, as it always is in high-dimensional models. There are various other things we might worry about:

• We aren’t actually using a great preconditioner. In practice we take Adam-style diagonal estimates, which assume our directions of interest are well captured by axis-aligned transformations; that seems plausible but far from certain. More sophisticated preconditioners (K-FAC, Shampoo, Muon) might help, but they’re more expensive and complicated.
• No universal guarantees in singular models. The standard SGLD convergence conditions (vanishing step sizes, Lyapunov tails) often fail for deep linear and modern nets. There are even clean examples where unadjusted Langevin diverges on light‑tailed targets.
• Discretization bias. With constant step sizes (what we actually use), we need strong diagnostics to keep bias small. In practice we rely on a MALA‑style acceptance‑probability diagnostic (not a real accept step) to tune the step size, and on long burn‑in until the minibatch‑loss trace flattens.
• Anisotropy hurts. Flat/stiff directions force tiny steps. Preconditioning (fixed diagonal, Fisher‑like) helps, but… is that enough? We know that in interesting models the local curvature matrix is not positive definite, so things are still not great.
• Hyperparameters matter. The localizer strength \(\gamma\) is a knob: too big and we drown out the likelihood geometry; too small and the chain can drift, even producing negative \(\widehat{\lambda}\).
• Finite‑\(n\) wrinkles. We evaluate at \(w^*\), which was trained on the same data we use inside the estimator. This works well empirically, but it’s not a theorem‑grade guarantee.

5.1 Momentum SGMCMC (SGHMC / SGNHT)

Why I hoped it’d help. Momentum should cruise along flat manifolds and mix faster.

What goes wrong. For minibatch gradients, SGHMC requires a friction/noise correction matched to the gradient‑noise covariance. In deep nets, that noise can be heavy‑tailed and state‑dependent. The correction then becomes ill‑defined or expensive to estimate, and mis‑tuning yields a different stationary law. Net effect: we risk biasing the expectation that defines \(\widehat{\lambda}\), while adding tricky hyperparameters.

5.2 Twisted / optimally‑guided SMC

Why I hoped it’d help. SMC gives ESS diagnostics, tempering ladders, and MH‑corrected move steps (MALA/HMC) that avoid unadjusted‑Langevin pathologies.

What goes wrong. In very high dimension, we need excellent “twists” to avoid particle collapse. Twisted SMC works when there’s a trick that gives us good twists. Good twists aren’t obvious for arbitrary SLT models. Further, the moment we put accept/reject MALA/HMC inside SMC, we’re back to near full‑batch costs per particle. We can do it for small models (and it’s a good sanity check), but it’s not a scalable replacement.

5.3 Non‑Gaussian / Lévy noise (“spikier” SGLD)

Why I hoped it’d help. Heavy‑tailed jumps might traverse singular valleys faster.

What goes wrong. If we swap Gaussian noise for α‑stable jumps, then unless we also change the drift to the fractional counterpart and/or add an MH correction, the invariant law won’t be the tempered local posterior we need. If we add an MH step to fix it, we’re back to paying for full‑data accepts.