Estimating the Local Learning Coefficient

3.1 Preconditioned SGLD

One limitation is that standard SGLD makes isotropic moves: the Gaussian noise is spherical, and the gradient step doesn’t account for very different curvature across directions. In singular models, this means we take tiny steps in stiff directions and explore slowly in flat ones.

A simple, affordable fix is preconditioning. Introduce a positive-definite matrix \(A\) — think of the Adam quasi-curvature estimates, which provide 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, smaller in sharp ones, and the noise is stretched accordingly. This respects the desired invariant distribution (for constant \(A\)), makes exploration more efficient, and preserves invariances like 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 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.

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, it rises faster and \(\widehat{\lambda}\) is larger. SGLD seems smart at doing the poking, at least in the models I’ve tried.

It can be shown to work in practice for at least one model family, the deep linear network, where we can test the estimates against ground truth, and indeed the SGLD-based estimator matches the ground truth beautifully, even at big scales (hundreds of millions of parameters, data sets that we can’t fully load into memory).

That’s surprising, to me at least. 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 those estimators may share biases. It could be that SGLD is “working great” in DLNs precisely because they are 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 there is unintuitive (to me) stuff going on here, even in tractable deep linear networks. The “landscape” has millions of dimensions, so “poking” in the correct directions is hard, as it always is in high dimensional models. There are various other things we might worry about:

• 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 should still not be amazing.
• Hyperparameters matter. The localizer strength \(\gamma\) is a knob: too big and you drown 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. It works well empirically, but it’s not a theorem‑grade guarantee.

4.1 Momentum SGMCMC (SGHMC / SGNHT)

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

What goes wrong. To be correct with minibatch gradients, SGHMC needs a friction/noise correction matched to the covariance of the gradient noise. In deep nets that noise can be heavy‑tailed and state‑dependent, so the correction is 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.

4.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 is a trick that gives us good twists. Good twists don’t seem obvious for arbitrary SLT models. Further, the moment we put accept/reject MALA/HMC inside SMC, we are 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.

4.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. Swap Gaussian noise for α‑stable jumps and, unless we also change the drift to the fractional counterpart and/or add MH correction, the invariant law is not the tempered local posterior we need. Add MH to fix it and we pay full‑data accepts again.