Avoiding unsafe sets when training with Langevin Dynamics

Training a model with noisy gradient descent can be idealized as overdamped Langevin dynamics on the loss landscape, and a natural safety question is to bound the probability $\nu_t(\mathcal{A}_H) = \mathbb{P}(Q_t \in \mathcal{A}_H)$ that the trajectory lies in a designated failure region $\mathcal{A}_H$.

We study this for a smooth, strongly convex loss in $d$ dimensions and a failure region separated from the minimizer by an energy gap.

Three bounds emerge.

At the end of training, the equilibrium mass $\pi(\mathcal{A}_H)$ is exponentially small in $d$, with a complementary energy-barrier rate when the noise is small.

Along the trajectory, a shape-free bound $\nu_t(\mathcal{A}_H) \le \pi(\mathcal{A}_H)(1 + \sqrt{\chi_0^2/\pi(\mathcal{A}_H)}\,e^{-mt})$ shows that the in-set probability relaxes to (twice) the static value after a burn-in time of order $d$, using only the global spectral gap $m$ of the loss.

A worked Ornstein-Uhlenbeck example shows this burn-in is necessary: an angular slice of the equilibrium shell can transiently swell by a factor exponential in $d$, even though its equilibrium mass is tiny.

To rule such swelling out we introduce a local relaxation rate attached to the failure region, defined through the spectral measure of its centered indicator rather than a Dirichlet-form Rayleigh quotient.

For geometrically isolated regions this rate exceeds the global one, shrinking the burn-in proportionally, and combined with a maximum-principle ceiling it caps the trajectory probability uniformly in time.

The picture is that strong convexity sets how fast training relaxes, but the shape of the unsafe set decides whether the trajectory bulges through it on the way home.