Akshat Valse

Stochastic foundations of machine learning · Iowa State University

research

Research

Three projects, ordered by how central they are to one question: what can be proved about the stochastic objects machine learning already runs on. The answer is usually less than one would hope, and more than gets checked.

Advisor
David Herzog
Department
Mathematics
Status
Ongoing

Bears on
SGLD · annealed Langevin · diffusion correctors

Langevin dynamics: sampling, discretization bias, transport coefficients

The overdamped diffusion leaves invariant, so discretizing it gives a sampler. That sampler is not a niche method. Stochastic gradient Langevin dynamics is its Euler–Maruyama discretization with a minibatch gradient; the inner loop of annealed Langevin sampling is the same chain at a fixed noise level; the corrector step of a predictor–corrector diffusion sampler is the same chain against a learned score. In all three the Metropolis correction is dropped, usually without comment, and the invariant law that results is .

I measure that gap where it is computable rather than bounding it. On an anisotropic Gaussian the unadjusted chain is a coordinatewise AR(1) whose stationary variance solves exactly,

so a measured bias that disagrees indicts the sampler, not the theory. On a two-component mixture the crossing rate follows Kramers' law with fitted slopes and against the predicted . Together the two results are the argument for noise schedules. Bias is polynomial in and can be halved by halving the step; mixing is exponential in the barrier, and no step size buys back an exponential.

The second thread is transport coefficients: how fast a sampler actually covers a rough landscape, as against how fast it would move on a flat one.

Fig. 3. Effective diffusivity in the cosine potential, β = 1, 1024 independent replicas, Euler–Maruyama at h = 0.01. The dashed line is the Lifson–Jackson value Deff = 1/I₀(1)² = 0.6238603604…, computed from the Bessel series and cross-checked against quadrature. Nothing here is fitted, and the trace is unsmoothed: the residual at t = 1000 is the transient, not Monte Carlo error, and the band is ±1 s.e. across four disjoint replica groups.

In one periodic dimension this is not an asymptotic bound but an identity. Lifson–Jackson gives

the averages taken over one period. For both are , so at , where , the answer is , and the simulation either finds it or is wrong. Any Langevin sampler, including the inner loop of a score-based one, can be pointed at this potential and required to return that number. One that gets the invariant measure right and the diffusivity wrong has a bug in its time discretization. Derivation and check.

Advisor
Ruoyu Wu
Department
Mathematics
Status
Ongoing

Bears on
PAC-Bayes · mean-field limits of training · SGD escape

Large-deviation asymptotics for mean-field interacting systems

Large deviations sits underneath a good deal of learning theory. Sanov's theorem and the Donsker–Varadhan variational formula for relative entropy are the machinery behind PAC-Bayes generalization bounds and variational mutual-information estimators. Freidlin–Wentzell small-noise theory is how escape times from saddles and shallow basins are analyzed for SGD and for Langevin-based optimizers. And the empirical-measure LDP for a mean-field interacting particle system is structurally the same construction that sits one level above the mean-field limit of wide-network training.

I work on rate functionals and the variational characterization of rare events for the empirical measure of such a system, currently in the setting of load balancing. The object is a large-deviation principle on path space; the rate functional then supplies the change of measure for an importance-sampling estimator of rare events, benchmarked against the analytic rate function by exact and Monte Carlo numerics. The application domain is queueing, but the technique carries well past it.

The results are held back until the preprint is out. What the figure below shows instead is the mechanism the whole method rests on, on the one target where the answer is known to machine precision.

Fig. 4. The mechanism, on a target with a known answer. Estimating P(Z > a) for Z ~ 𝒩(0,1) by crude Monte Carlo (hollow) and by importance sampling under the tilted law 𝒩(a, 1) with weight exp(−aX + a²/2) (filled), against the exact Φ(−a), dashed. Both carry ±1 s.e. Watch crude Monte Carlo degrade rather than fail cleanly: its markers drift off the curve, go bright as the relative error climbs, survive on a single chance hit two decades out of place with an error bar spanning the axis, and then stop appearing. The tilted estimator tracks the exact curve five decades further down. The readout reports where each method last held a relative standard error under 50%. The reference comes from an erf series and the Laplace continued fraction, agreeing with SciPy to 7×10⁻¹⁵.

Naive Monte Carlo does not merely get slow in this regime. It stops being an estimator: it returns zero occurrences and no information about the magnitude. Tilting the sampling law so the rare event becomes typical, then reweighting by the likelihood ratio, recovers an answer, and the tilt has to come from somewhere principled, because a badly chosen change of measure produces an estimator with infinite variance that fails silently. In the Gaussian case the right tilt is a mean shift and can be written down. In the mean-field model it comes from the large-deviation rate function, and checking whether the asymptotics have actually taken hold at the system sizes anyone simulates is the part of the project I find most worth doing.

Advisor
Jue Yan
Department
Mathematics
Status
Ongoing

Bears on
learned operators · neural surrogates · rollout stability

Stability of neural approximations to evolution equations

A cell-average neural network solver learns a one-step finite-volume operator

a learned surrogate for a numerical flux. Conservation comes for free: the same network is evaluated at both faces of a cell, so the update telescopes whatever the weights turn out to be. Stability does not. Networks trained identically up to the random seed, reaching losses you cannot tell apart, behave completely differently under long rollouts, and since the losses are indistinguishable the loss cannot tell you which one you have. It is the autoregressive-surrogate failure mode, in the simplest setting where classical numerical analysis can still say something about it.

A trained solver is still a numerical scheme: its weights are frozen once training ends, so what you have is a fixed deterministic update rule that happened to be fitted rather than derived. Numerical analysis has a long history of asking whether update rules of that kind are stable. I am adapting von Neumann analysis to this setting, deriving the linearised amplification factor of the learned time-stepping scheme and characterizing how perturbations propagate through it, toward a condition on the weights alone that can be checked before committing to an expensive rollout.

preprints

Preprints will be posted here as they become available.