Langevin Sampling

Introduction

As I was studying the score-matching approach to diffusion, I stumbled upon Langevin sampling. It is particularly useful when we need to sample from a distribution with an intractable normalizing constant. From an application perspective, it seems straightforward, but I couldn't understand how we got to that equation, so I looked up the physics motivation. These are my notes on this topic as I worked through chapter 1 of the book Nonequilibrium Statistical Mechanics by Robert Zwanzig. Tbh, I wish I had earned a degree in maths or physics as an undergraduate, but oh well... I'm sure there are a lot of typos, mistakes, and incorrect assumptions in these notes, but I hope they still show the general picture.

Visualization:

Code that generates this visualization (generated by ChatGPT): langevin_2d_video.py

Notes

We study the Brownian motion of a particle with mass , radius , and velocity in a fluid with viscosity . For simplicity, this note works in one spatial dimension, so , , and the Gaussian noise variables are scalars. The same ideas can be generalized to multiple dimensions by replacing them with vectors.

Stokes' law gives the drag force , so Newton's second law gives

where

Separating variables,

Integrating,

so

This means

This is correct for the deterministic drag-only model. However, it does not fully describe a particle in thermal equilibrium because we haven't accounted for the random forces exerted on the particle by the medium (collisions with all the molecules).

Let the random fluctuating force acting on the particle be . Then we have

Equivalently,

Let's study

which is the total impulse from the molecules.

Assume the molecules hit at a rate . Then the total number of collisions during the interval is

The total impulse over is

We assume that the molecular collisions are independent, including collisions in disjoint time intervals. Since their impulses come from random directions, we also assume

When

the Central Limit Theorem gives

Therefore,

Since

we have

Let's use Zwanzig's notation and set the strength to

Therefore,

The first two moments of are then

and

For Gaussian noise, the zero correlation at means that the fluctuating forces at two different times are independent, consistent with our assumption of independent molecular collisions.

The Dirac delta is not an ordinary function and does not have a pointwise value at . It is defined by how it behaves inside an integral: it vanishes away from and concentrates all of its unit mass there.

We do not measure impulse at a single time point, but rather over an interval. Define

Then

With a fixed , the inner integral is . Hence,

which is what we already know.

To summarize,

where is the total fluctuating impulse, Gaussian with

and

Let's derive the value of .

But before we do that, note that in a more general setting, the particle is also under the influence of other conservative forces such as gravity, electromagnetic forces, etc., whose total potential energy is , where is the particle's position.

Hence the total conservative force is

so the equation of motion becomes

For now, we can assume

i.e., there is no external conservative force.

This is because and are two different effects: one comes from external conservative forces, while the other comes from the thermal bath. We can isolate one to study the other.

An analogy is measuring the acceleration due to gravity from a falling object. Air resistance can affect the observed motion, so we perform the experiment in a vacuum to eliminate that effect. Removing air resistance does not change gravity itself. Similarly, we can temporarily set to remove the external conservative force while deriving the thermal-noise strength.

We have

or equivalently,

This is a linear, inhomogeneous first-order differential equation of the form

Without , it is easy to see that

where is a constant.

To use variation of constants, define

This does not restrict the form of , because is never zero and every can equivalently be written as

Without , would be constant. With present, differentiate to see how it varies:

Substituting

gives

Integrating from to ,

Since ,

Therefore,

Using this result for

we have

Let's look at . It has three terms.

Taking expectation,

hence this term also disappears.

Taking expectation,

Using

we obtain

The delta function evaluates the inner integral:

Evaluating,

Therefore,

As , the system reaches thermal equilibrium.

For one dimension,

Therefore,

which gives

This is the fluctuation-dissipation theorem.

Now we know

and

Let's get back to our equation,

Equivalently,

where

• is the momentum change,
• is the conservative force impulse,
• is the friction impulse,
• is the random molecular impulse.

To study the deterministic velocity relaxation, assume for the moment that there is no conservative force and no thermal noise. Then

Hence,

The relaxation time is defined as

After one relaxation time,

Let be the characteristic timescale over which the particle's position or the external conditions change. In the overdamped regime,

meaning the particle quickly adjusts to its new velocity due to friction from the surrounding medium.

We can rewrite

Due to the exponential decay of , the change is of the same order of magnitude as .

For example, in the ideal case from the previous derivation,

Thus,

and

so

(I honestly think this logic is overreaching, as the form of is not .)

Therefore,

which is the overdamped assumption.

This means that, in the equation

we can neglect the left-hand side and rearrange to obtain

Hence,

Now the Boltzmann rule states that, for a state with energy ,

The state of the particle is , and

Therefore,

Marginalizing over the velocity,

Factoring out the terms independent of ,

The integral does not depend on , so it can be absorbed into the normalization constant .

Therefore,

Since

taking the logarithm gives

Therefore,

Substituting this back into the equation,

We also know

Hence,

Let

Then

Define

The equation becomes

Let's perform a variable transformation. Define a characteristic length , and introduce the dimensionless time

where has units of , so is dimensionless. This is different from the physical relaxation time defined earlier.

Also define

which is also dimensionless.

Since

we have

Furthermore,

and

Substituting the transformed variables into the SDE,

Therefore,

Let's use the names and instead of and :

We can already see how this equation can generate a sampling trajectory for : the score pulls the trajectory toward regions of higher probability density and the modes of the distribution, while the Brownian noise introduces random motion that helps the trajectory explore the distribution and potentially move between different modes.

Let's discretize it first.

For a finite, very small time step

moves from to .

Also,

This is the Euler–Maruyama approximation. Over one sufficiently small time step, we approximate the score by its value at the beginning of the step, . This works when the score changes smoothly; otherwise, we need to choose a smaller . Hence,

When the target distribution is written as

where is a dimensionless model energy, equivalent to physical energy divided by , and is an intractable normalization constant, we have

Hence, the update rule simply replaces with :

This is how Langevin sampling is used in energy-based models.

This discretized method is called the unadjusted Langevin algorithm. With a fixed nonzero step size , it generally samples an approximation to rather than exactly , because discretization introduces a small bias. Using a smaller step size reduces this bias but requires more steps. If exact correction is needed, the Metropolis-adjusted Langevin algorithm adds an accept-or-reject step. I've yet to study that one, but let's stop here.