Potts model

Consider the Potts model with colors on the complete graph with vertices. A configuration is with . Define the energy At inverse temperature , the Boltzmann distribution is . On the complete graph, the only relevant macroscopic variable is the empirical proportions vector

so that . A short computation rewrites the energy in terms of : up to an additive constant irrelevant for Gibbs weights. The number of configurations with a given is , where is the Shannon entropy. Putting energy and entropy together, typical samples from concentrate near minimizers of the mean-field free-energy functional:

constrained to the probability simplex. Everything that follows is geometry: how the minima of evolve as the inverse temperature varies.

Local minima: Two features make the analysis almost trivial. First, permuting colors leaves unchanged. Second, a stationary point under the simplex constraint satisfies a Lagrange-multiplier condition for . Hence each coordinate must solve the same scalar equation. From this, one finds that all local minima are necessarily of the following two types:

• Disordered point (uniform): always a stationary point.

• Ordered points (one dominant color, the rest equal): for some dominant proportion . There are such points by choosing which color is dominant.

The remaining work is algebra: determine, as a function of , when ordered stationary points exist, and which of the stationary points are local minima versus saddles. The details are routine and add little insight.

The phase diagram in

Assume . Then the model exhibits a first-order transition, and metastability appears on both sides. There are two key inverse temperatures:

• Spinodal threshold : the ordered stationary points appear as local minima.
• Coexistence threshold : the disordered minimum and the ordered minimum have equal free energy.

For the normalization above, the coexistence point is

The geometry of splits into three regimes.

Zonal spherical harmonics

Understanding the eigenfunctions of the spherical Laplacian is a central task. These eigenfunctions form the building blocks of harmonic analysis on spheres.

A particularly important situation arises when the function of interest depends only on the angular separation from a fixed axis.

Such zonal functions retain full rotational symmetry around that axis and reduce the spherical Laplacian to a one-dimensional operator. Its eigenfunctions turn out to be polynomials on the interval with a very specific weight. These polynomials are the Gegenbauer polynomials, which generalize several classical families of orthogonal polynomials, including the Legendre and Chebyshev polynomials.

Reminder on differential operators

Many problems require the ability to compute gradients, divergences, and Laplacians in arbitrary coordinate systems. Before specializing to spherical coordinates in higher dimensions, it is useful to recall the geometric meaning of these operators and the minimal formalism needed to manipulate them. Suppose that a point in space is described by a collection of coordinates If we make a small change , the physical displacement has a length denoted by . In an orthogonal coordinate system, each coordinate direction comes with a scale factor such that a change corresponds to a “physical” displacement of length along the unit vector . These functions encode how the coordinate system stretches or compresses distances along each axis. Once the scale factors are known, the rest of the differential operators follow directly.

The gradient: For a scalar function , the change produced by varying only is approximately . The physical distance traveled in this move is . The rate of increase of per unit distance in the direction is therefore . This is the th component of the gradient. Summing over all directions gives

The divergence describes the net outflow of a vector field from an infinitesimal volume around a point. Consider a vector field . Place a tiny coordinate-aligned box at a point. Its physical edge lengths are , so its infinitesimal volume is: The flux of through the pair of faces orthogonal to is multiplied by the physical area of those faces: . The divergence is the total outward flux divided by the infinitesimal volume:

This expression is simply “flux in direction ” minus “flux in direction ” across the opposite face, normalized by the physical volume.

The Laplacian is the divergence of the gradient. It describes how a scalar function curves around a point since it measures the net outflow of the gradient field. It also encodes how compares to its local averages over small spheres. In orthogonal coordinates, inserting the expression for the gradient into the divergence formula yields:

This is the expression we will apply to spherical coordinates in the next section.

Masked Discrete Diffusion

We consider a finite state space where is a special masked state and correspond to token values in a vocabulary of size . On the time interval , we define a continuous-time Markov chain with initial distribution and time-dependent infinitesimal rate matrix so that for any ,

If the total jump rate out of state is , then

Bayes’ rule implies that the time-reversal of this Markov chain is itself Markov, with infinitesimal rate matrix satisfying where is the marginal distribution of at time . We have:

We are interested in modeling a Markov chain that progressively masks the initial value into the masked state as time goes from to . Transitions are only allowed from any token to the masked state , and once in the process remains there. Thus, outside the diagonal, the only nonzero entries of are . As it will be useful later, we denote by the jump time to and we assume almost surely, so that with probability one, and that has a continuous distribution. In words: the process starts at some token value and at a random time jumps to the masked state , where it remains until time .

Extension to Sequences

We are interested in modeling sequences comprised of discrete tokens, eg: binary images, genomic sequences, chemical compounds, protein sequences, etc… Each token takes value in . We denote by the data distribution over such sequences. For this purpose, we consider independent copies of the above Markov chain, one per coordinate: each with rate matrix as defined previously. At time , the process reaches the fully masked sequence with probability one. Denote by the jump time of coordinate . Since the jump times are almost surely distinct, the infinitesimal rate matrix of the joint process is nonzero only when a single coordinate changes. If differ by a single coordinate , we have

As before, the time-reversal has infinitesimal rate matrix satisfying

where is the marginal distribution of . Since and differ at coordinate only, for to be non-zero, necessarily and . Let be the set of unmasked coordinates in . To observe configuration at time , the masked coordinates must have and the unmasked ones :

Similarly, and since differs from only at coordinate :

This shows that the time-reversal rate matrix becomes

with time dependent scalar . To simulate the reverse process that progressively unmasks a fully masked sequence, one only needs to model the conditional distribution of the data distribution . This is precisely the prediction task of masked language models such as BERT, which estimate token probabilities conditioned on visible context.

Training

To train the denoising model, Equation 1 shows that it is natural to parametrize the conditional distribution

for all sets with . Once done, one can define the rate matrix of the time-reversal process as:

If one denotes by the law of the forward noising process started from , and by the law of the time-reversal process started from the fully masked sequence at time and with learned denoising model , one can train the model by minimizing

Consider a trajectory of the forward noising process. The jump at time of the -th coordinate is denoted by . For simplifying the notations, we denote the reverse jump by . The log-likelihood ratio between the two processes is easily shown to be:

where and are the total jump rates of the forward and reverse processes respectively. Since in fact does not depend on , minimizing the KL divergence is equivalent to minimizing:

where is the set of unmasked coordinates at time . It is more convenient to rewrite this quantity as an integral over time so that one can sample a time uniformly in during training. With the Dirac delta function , we can rewrite this expectation as:

where so that . For training, it suffices to sample uniformly in , then choose , then sample the noised configuration according to the forward process, and finally obtain an unbiased estimate of the loss. The term is large for small , counter-balancing the fact that the reconstruction task is much easier when only a few tokens are masked. For standard denoising diffusion model, there is a similar “signal-to-noise” weighting term that balances the easy and hard denoising tasks.

Conclusion

Discrete diffusion models with one-way masking are mathematically almost identical to masked language models. Hence their similar behavior and performance on text generation tasks are not coincidental. The ideas summarized in these notes were developed in a very interesting stream of papers, including (Ou et al. 2024), (Sahoo et al. 2024), (Shi et al. 2024) and a number of more recent works. One of the potential drawbacks of such masked discrete diffusion models is that the support of the noising distribution is typically strictly smaller, and indeed often much smaller, than the whole state space. This means that when the denoising model is not perfect and wanders outside the support of the noising distribution, one can quickly end up in regions never seen during training. This can leads to poor sample quality and unstable behavior. Other discrete diffusion models such as the ones reaching the uniform distribution over all tokens at time are not as badly affected by this issue, although they do suffer from other important computational and modeling challenges. Exciting research directions remain to be explored in this area!

Static Shrödinger Bridge Problem

A standard way to address these issues is to add an entropic regularization term to the objective function. The resulting problem is known as the Schrödinger bridge problem and can be formulated as follows. Consider a reference joint distribution over and find the coupling that minimizes the Kullback-Leibler divergence to while matching the marginals and :

Remark (invariance under separable rescaling). Only the “interaction structure” of matters: if one replaces by with positive , then the optimal coupling can still be written in the same form below, with the factors absorbed into the potentials. Equivalently, the solution depends on only through its kernel up to left/right diagonal scaling.

A common choice for the reference measure is a distribution of the form (optionally times separable factors). This choice encourages the coupling to put more mass on pairs that are close according to the distance , and the resulting optimization problem can be rewritten (up to an additive constant) as:

where is the entropy of the coupling . Note that since the marginals of are fixed, it is also equivalent (up to constants) to replacing the entropy term by the Kullback-Leibler divergence to the independent coupling ,

Writing the Lagrange dual formulation of the problem Equation 2 provides fast algorithms to solve it such as the iterative proportional fitting procedure (IPFP), also known as the Sinkhorn-Knopp algorithm in the optimal transport literature (Cuturi 2013). Importantly, it also almost immediately shows that the optimal coupling has a very simple form,

These potentials must satisfy the marginal constraints, equivalently the Schrödinger/Sinkhorn scaling equations:

Dynamic Schrödinger Bridge Problem

Now, suppose that the reference distribution is given as the two-time marginal of a Markov process with transition kernels and path measure on trajectories .

The dynamic Schrödinger bridge problem consists in finding a new path measure on trajectories such that the starting and ending marginals are and , while minimizing the Kullback-Leibler divergence to the reference path distribution :

When the reference is Markov, the chain rule for KL (disintegration with respect to ) shows this is equivalent to the static Schrödinger bridge problem for the two-time marginal of : one first solves for the optimal endpoint coupling , and then fills in intermediate times using the reference bridges :

1. Sample the endpoints .
   
2. Sample the intermediate points according to the conditional law of the reference process given the endpoints, i.e. .

Continuous-Time Stochastic Processes:

A typical setting is when the reference process is given as the solution of a stochastic differential equation of the form

started from some initial distribution (one may take it to be without loss of generality, since any mismatch can be absorbed into the endpoint tilt below). The solution to the dynamic Schrödinger bridge problem is given by the twisted path distribution:

for endpoint potentials and satisfying the marginal constraints at and .

The marginal density at intermediate time factorizes as:

where is the time- marginal density of the reference process, and where the time-dependent Schrödinger potentials are defined by:

(with and , up to normalization). Naturally, the dynamics of the Schrödinger bridge process can be described as a new stochastic differential equation obtained by applying a Doob h-transform to the reference SDE,

where, as explained in these previous notes, the function is the harmonic extension of the terminal potential defined above.

Discrete-Time Markov Chains:

It is often useful to state the Schrödinger bridge dynamics in a discrete setting. Let the reference process be a Markov chain with one-step kernels , i.e. Let be the terminal potential at time and define the backward (harmonic) potentials defined by and recursively as:

Then the Schrödinger bridge is Markov and its forward transition kernels are given by the discrete Doob -transform: This is the discrete-time counterpart of the continuous-time drift correction .

In general, Schrödinger bridge problems are difficult problems since the endpoint potentials and need to be found such that the marginal constraints are satisfied. These recent years have seen the development of many numerical methods to solve this problem approximately, especially in the machine learning community, eg: (Shi et al. 2023); (De Bortoli et al. 2021); (Vargas et al. 2021); (Chen, Liu, and Theodorou 2021).

Some natural examples

One can then asks oneself what are some natural entire functions that vanish at some predetermined set of zeros. For example, what functions vanish on all the integers? Such a function cannot be of order less than one since otherwise it could be writtem as , but this product does not converge. Any entire function of order that has a simple zero at each integer is of the form

Checking the Taylor expansion of sine at zero, one finds that the celebrated formula:

and a similar interesting example by taking the derivatives of . Now, what about a function that vanishes on all the negative integers? Again, such a function needs to be of order at least one. Hadamard’s factorization theorem then tells us that such a function is, up to a multiplicative constant, of the form

Naturally, since such a function vanishes on all the negative integers, one of the first things one would like to try is to look at and relate it to itself. Since vanishes on , one knows that it can be expressed as so that is almost the same as . One can then do some algebra to choose the constant so that . One finds that the correct choice is the Euler-Mascheroni constant,

This gives the final expression for as

Furthermore, since as , it follows that , from which the identity gives that . In other words, is an analytic continuation on the whole complex plane of the inverse of the Gamma function,

This generalizes the usual definition valid for to the whole complex plane.

The connection to the sine function is also almost immediate. Indeed, the function vanishes on all the integers and the infinite product immediately shows that . But since , one obtains that , i.e.

This is the celebrated Euler reflection formula, from which it also follows that .

Zeta function

Let’s conclude these notes by deriving the analytic continuation of the Riemann zeta function , originnaly defined for , to the whole complex plane. I have always found this proof very elegant. A common trick that is used in many places is to do a change of variable in the definition of the Gamma function to obtains that

It is valid for . This allows one to obtain express the zeta function as:

From the previous discussion, one knows that is an entire function of order , and the integral converges for except for ; in other words, this already gives a meromorphic continuation of to the domain with a simple pole at . Since things get nasty near , the standard approach consist in splitting the integral into two parts. The part defines an entire function and one only needs to take care of the integral . This can be done by expressing , where are the Bernoulli numbers and integrating each term, but the way Bernhard Riemann did it is way more fun. The idea is to note use the boring but instead introduce the way more interesting function

where , defined for . Note that the function decreases exponentially rapidly to as . In other words, instead of Equation 1 one can just as easily write

This slight change of parametrization allows to write the zeta function as

One has not gained much doing this since there is indeed still an issue at where diverges. However, the Jacobi theta function enjoys some interesting symmetries. Crucially, the Poisson summation formula applied to the Gaussian function gives that satisfies modular inversion symmetry:

This means that splitting in Equation 2, using the change of variable to map to and finally use the modular inversion symmetry of the theta function leads after standard algebra:

First, one can note that since decreases exponentially quickly to as , the integral above defines and entire function. This means that the expression above defines a meromorphic continuation of to the whole complex plane with a simple pole at . There is no pole at since the simple zero of takes care of it and gives the value . This also shows that the function inherits from a simple zero at all the negative even integers, i.e.  for . And there are indeed a few other zeros, as the plot below shows… and they seem to be located on the critical line …

What is remarkable is that the term inside the curly brackets of Equation 3 is symmetric in and , i.e. symmetric with respect to the vertical line in the complex plane. This means that the function satisfies the functional equation

Fisher-Rao metric

Suppose we want to define a distance on the space of probability densities on . A natural but naive approach is to use an -type distance:

where are densities with respect to the Lebesgue measure. However, this definition has several shortcomings. For instance, if we change the base measure to for some positive density , and define the distance as

we obtain a different value. Perhaps more troubling, the distance is not invariant under reparametrizations. Let be a diffeomorphism of , and set . Then the transformed densities become , where is the Jacobian determinant. In general,

so the distance depends on the choice of coordinates. That is, measuring in Cartesian or polar coordinates yields different results—an undesirable feature. Ideally, we seek a distance that is invariant under reparametrizations and changes of base measure, such as the total variation distance,

One potential drawback of the total variation distance is that it is not differentiable, which can make it difficult to use in optimization problems. An alternative is to consider -divergences, defined as

where is a convex function with . These divergences are differentiable and invariant under reparametrizations and changes of base measure, although they are not symmetric and thus not true distances. Locally, however, all -divergences are equivalent, as a second-order Taylor expansion shows:

This means that all these divergences describe the same local geometry, defined by the Fisher-Rao information metric. Furthermore, it is relatively straightforward to derive the global geometry induced by the Fisher information metric. Consider the mapping , which maps a density to an element of the unit sphere of . Since

for , and we have just seen that any -divergence is locally equivalent to the Fisher-Rao metric, it follows that the geometry induced by the Fisher-Rao information information metric is the same as the geometry induced by the -norm on the unit sphere of . This implies that the geodesic distance between two densities and is given (up to an irrelevant constant) by the -geodesic distance between the points and . In other words, the geodesic distance, i.e. the Fisher-Rao distance, is given by

The geodesic path is a great circle, , where for . This shows, for example, that the Fisher-Rao geodesic between two Gaussian densities is composed of densities that are Gaussian mixtures; i.e., the geodesic is not composed of Gaussian densities in general. In other words, probability mass is not transported along the geodesic but reshaped, unlike the Wasserstein metric, which describes transport of probability mass. Note in passing that the Hellinger distance, defined as

is just just a slightly rescaled version of the Fisher-Rao distance since they are related by a deterministic function, i.e. . In this sense, the Hellinger distance is equivalent to the Fisher-Rao distance, and both describe a “correct” way to measure distances between probability densities for many applications.

Gradient flow

What do gradient flows look like in this Fisher-Rao geometry? For example, for a given distribution , the gradient flow of under the Wasserstein metric is given by the transport equation:

which describes the Langevin dynamics of the process , as informally described in these notes. So what does the gradient flow of look like in the Fisher-Rao geometry?

To answer this, one can consider the square-root mapping , express everything in terms of , compute the -gradient on the unit sphere (which is straightforward), and finally map back to the density using the inverse mapping . One readily finds that the gradient flow is described by:

This is quite intuitive: the flow tries to increase in regions where and decrease in regions where . Discretizing this flow can naturally be done using sampling-based methods. If is a system of weighted particles approximating , following the Fisher-Rao gradient flow corresponds to updating the weights as

for a small time-step. Indeed, it is because , where , can be discretized by updating the weights as . This is very much related to the resampling step in sequential Monte Carlo methods, and the recent article (Crucinio and Pathiraja 2025) make these connections explicit.

Table of Contents

• Linear Systems
• Adjoint Method
• PDE Inverse Problems
• Controlled Diffusions

The adjoint method is, at its core, the same idea as backpropagation or reverse-mode automatic differentiation. In practice, though, it’s often helpful to understand how it works under the hood. Basic implementations of backprop can be sub-optimal or impractical (eg. memory intensive), especially in settings like PDE-constrained optimization or stochastic optimal control.

Nystrom approximation

Before describing the simple and most important case of additive Gaussian noise, let’s give a brief reminder on the Nystrom approximation. The distribution of is Gaussian with mean and covariance . This means that is distributed as the unconditional distribution . In particular:

where is the so-called Nystrom approximation of the covariance matrix based on the inducing variable . This shows that the Nystrom approximation simply consists in ignoring the conditional variance term , and is thus an underestimate of the covariance matrix . Furthermore, if is very informative of , then is small and the Nystrom approximation is accurate.

Observation with additive Gaussian noise

The case of additive Gaussian noise is particularly simple. Assume that

where the noise terms are independent. Since , algebra gives that

Using that and the matrix inversion lemma it quickly follows that optimal variational distribution is with

Indeed, these are approximations of the exact condition moments,

where the Nystrom approximation is used instead. One then finds that . With the optimal variational distribution , Equation 4 gives that the free energy is:

Furthermore, note that exact likelihood of the observations is so that the free energy can be expressed as

for pseudo-likelihood . This shows that is given by:

The term is just the sum of the conditional variances and can be thought of as a regularization term,

As the number of inducing variables increases, the pseudo-likelihood becomes more accurate , the conditional variances shrink to zero, and the KL divergence approaches zero.

The animation at the start of this note illustrates the effect of optimizing the location of the inducing points with a very simple gradient descent. A few experiments show that it is worth being careful with the initial choice of inducing points. Inducing points chosen very far from the data essentially remain fixed during the optimization (ie. the gradient is very small). Initializing with k-means++ clustering of the data points seems to be a robust strategy and give an almost optimal choice of inducing points.

The problems and notations

Recall that we are trying to estimate the connective constant of self-avoiding walks (SAW) in the 2D lattice . If denotes the number of SAWs of length , we have the following asymptotic behavior:

for some unknown constants , , and . The main objective of the assignment is to estimate , which can also be expressed as the limit of as . As of today, the best known estimate is , which required several tens of thousand hours of CPU time to compute. To estimate , one must approximate the number of SAWs of length starting at the origin for large values of if one hopes to get a good estimate.

Consider a sequence of distinct vertices in with and for all , i.e., a walk of length . For notational convenience, let us introduce the function that returns one if is a correct walk of length , and zero otherwise. In particular, this function returns zero if two consecutive vertices are the same, or if the walk does not start at zero. Similarly, introduce the function that returns one if is a SAW of length . One can define two important probability mass functions:

They describe the uniform distributions on all the walks of length and all the SAWs of length , respectively.

Importance sampling

One can approximate with naive Monte Carlo by estimating the proportion of walks that are SAWs,

This is an absolute disaster since the proportion of SAWs among all walks is extremely small. One can do significantly better using importance sampling. For this, consider a proposal distribution that starts at the origin and continues by choosing uniformly among the four neighbors of the last vertex that have not been visited yet. If there are no unvisited neighbors, the walk continues by standing still until length is reached: the resulting path is not even a valid walk, so as well as . The probability mass function of the proposal distribution is easy to compute, so estimating with importance sampling is straightforward. This is usually called the Rosenbluth method (Rosenbluth and Rosenbluth 1955). [Note to students: make it much clearer in your report that the Rosenbluth method is just importance sampling. Do note that even the “rejected” walks have to be taken into account!]

As one can see, the quality quickly deteriorates as increases. This is because the number of accepted walks is very small, and, among them, the importance weights are highly unequal.
[Note to students: you should explain this much more clearly, and possibly explore this more quantitatively. The reason it is failing is not only that the number of accepted walks is small]

Recursive formulation

We have just seen that importance sampling will not be able to estimate for large values of . This makes accurate estimates of difficult to obtain this way.

To make progress, one can exploit the recursive structure of the problem. Let us define the concatenation of two walks. Given a first walk and a second walk , one can define a new walk of length by starting at the origin, following the increments of the first walk, then the increments of the second. The concatenation of two SAWs is not always a SAW. However, it is not hard to prove the following. Define as the probability that, when sampling SAWs and independently and uniformly at random, their concatenation is still a SAW. Then:

[Note to students: it is OK for you to use this fact. It’s even better if you can prove it, but not absolutely necessary.]

Assuming one can generate SAWs of length uniformly at random ( a problem that will be discussed later), we can estimate in several ways:

1. For small values of , the number of SAWs is known exactly (e.g., , ). Suppose one can generate SAWs of length . One can then estimate empirically. Since , it follows that, for fixed and , Using the fact that , one can then estimate from the estimate of .

2. Alternatively, one can estimate for large recursively. For example, starting from , estimate to compute , then use to compute , and so on. Using this method and about hours of CPU time (see below for details) with SAWs of lengths , I obtained .

Generating SAWs

The previous discussion shows that, once we know how to generate uniform SAWs, we can estimate relatively easily. One of the most common methods is the pivot algorithm: see here for a nice visualization. The principle is simple: given a SAW, randomly select a pivot site and apply a symmetry operation (like rotation or reflection) to one part of the walk. If the resulting walk remains self-avoiding, accept it; otherwise, reject it. Repeating this process generates diverse, approximately uniform SAWs.
[Note to students: explain this much more clearly if you decide to use it]

In short, the pivot algorithm updates a SAW by applying a symmetry operation to a subpath. Given a SAW , one can obtain another SAW by applying to it the pivot algorithm a (large) number of times. To obtain a nearly independent SAW of length starting from , one typically need to apply about pivot steps. While it can be slow for large , it is far more efficient than naive importance sampling.
[Note to students: efficiently implementing the pivot algorithm is non-trivial, but LLM assistants can help a lot, and are actually quite useful for code optimization]

Sequential Monte Carlo

To estimate for large , one can use Sequential Monte Carlo (SMC). The idea is to grow a population of SAWs in parallel and estimate by recursively estimating the ratios . Suppose you have SAWs of length . Try to extend each SAW by choosing a neighbor of the last vertex that has not been visited yet. This is a form of importance sampling, giving new walks of length with associated weights (some of them being non-valid walks!). Then, resample times from this weighted set to get new SAWs of length (with possible duplicates). Apply the pivot algorithm to eliminate these duplicates and generate more diverse SAWs.
[Note to students: if you decide to use SMC, explain it much more clearly. It’s not entirely straightforward to understand or implement, but it is one of the most powerful and versatile Monte Carlo methods to this day. A good investment of your time if you decide to understand SMC]

Improving the estimation of

Suppose you have estimates of for various . Since

you can fit a linear regression to estimate , , and . I tried this approach using a naive and non-optimized SMC implementation with and , running for 10 hours on a free (and bad) online CPU, and obtained .
[Note to students: can you do much better?]

Running long simulations

The best known estimate of required several tens of thousands of CPU hours. While writing these notes, I was able to run simulations easily and for free using deepNote: it was my first time using it, and it was very user friendly. This allowed me to run simulations for 8 hours on a (free but slow) CPU without issue. Launch simulations in the evening and let them run overnight. [Note to students: for the more motivated ones, you can try writing GPU-friendly code to run simulations, possibly on Google Colab]

Some References:

• The original papers by Jarzynski (Jarzynski 1997) and Crooks (Crooks 1999).
• The book (Stoltz, Rousset, et al. 2010) is excellent!
• The two papers that prompted these notes: (Vargas et al. 2024) and (Albergo and Vanden-Eijnden 2024).

Readings:

• The introduction paper (Lei et al. 2018) is really good
• I am really curious about (Gibbs, Cherian, and Candès 2023) and it’s next on my reading list

Change of measure

Consider a diffusion in given by:

for an initial distribution and for drift and volatility functions and . On the time interval , this defines a probability on the path-space . For two functions and , consider the probability distribution defined as

where denotes the normalizing constant

The distribution places more probability mass on trajectories such that is large. As described in these notes on Doob h-transforms, the path distribution can be described by a diffusion process with dynamics

The control function is of the gradient form

and the function is described by the conditional expectation,

The expression is intuitive; to describe the tilted measure that places more probability mass on trajectories such that is large, the optimal control should point towards promising states, i.e. states such that the expected “reward-to-go” quantity is large.

Variational Formulation

To obtain a variational description of the optimal control function , it suffices to express it as the solution of an optimization problem of the form

for an appropriately chosen distance. Here is the probability distribution of the controlled diffusion with dynamics

for some control function and initial distribution . Note that we have that . The KL-divergence is natural (pseudo) distance since it elegantly deals with the intractable constant and the ratio is easy to compute. Girsanov Theorem gives that

From this expression, one can readily write . Minimizing over the control and the initial distribution shows that the optimal control is:

Minimizing this loss attempts to find a control that drives the quantity large while keeping the control effort small. Equivalently, this can be expressed as a maximization problem,

Note that since , the optimal control is such that for any trajectory we have:

Stochastic Control

In the previous section, there was nothing special about the starting point and the time horizon . This means that the same derivation gives the solution to the following stochastic optimal control problem. Consider the reward-to-go function (a.k.a. value function) defined as

where is the controlled diffusion Equation 3. We have that

This shows that optimal control can also be expressed as

The expression is intuitive: since we are trying to maximize the reward-to-go function, the optimal control should be in the direction of the gradient of the reward-to-go function.

Hamilton-Jacobi-Bellman

Finally, let us mention that one can easily derive the Hamilton-Jacobi-Bellman equation for the reward-to-go function . We have

with . For , we have

where is the generator of the uncontrolled diffusion. Since is a simple quadratic function, the supremum over the control can be computed in closed form,

as we already knew from Equation 5. This implies that the reward-to-go function satisfies the HJB equation

with terminal condition . Another route to derive Equation 6 is to simply use the fact that ; since the Feynman-Kac gives that the function satisfies , the conclusion follows from a few lines of algebra by starting writing , expanding and expressing everything back in terms of . The term naturally arises when expressing the diffusion term as a function of the second derivative of ; it is the idea of the standard Cole-Hopf transformation.

Finally, Ito’s lemma and Equation 6 give that for , the optimally controlled diffusion satisfies:

Since and , writing this expression in between time and gives the formula Equation 4 for the log-normalizing constant .

Importance Sampling on path-space

Consider a functional on path-space; a typical example is

Suppose that we would like to evaluate the expectation of under the measure . Naive Monte-Carlo (MC) would require sampling trajectories from Equation 2 and computing the average of on these trajectories. To reduce the variance of this naive MC estimator, one can also use importance sampling by sampling trajectories from the measure and compute the average

with weights given by the Radon-Nikodym derivative

Choosing the optimal “control” function that minimizes the variance of the estimator is not entirely straightforward, although this previous note already gives the answer. More on this in another note.

General diffusions

Consider a diffusion process

with generator and initial distribution . We are interested in describing the dynamics of the “conditioned” process given by the probability distribution defined in Equation 4. Algebra applied to Equation 6 then shows that

where the function is described in Equation 5. Since , this reveals that the probability distribution describes a diffusion process with dynamics

The additional drift term is involves a “control” with

Note that the initial distribution of the conditioned process is

Unfortunately, apart from a few straightforward cases such as a Brownian motion or an Ornstein-Uhlenbeck process, the function is generally intractable. However, there are indeed several numerical methods available to approximate it effectively.

Brownian bridge

What about a Brownian motion in conditioned to hit the state at time , i.e. a Brownian bridge? In that case, the function is given by

for some irrelevant normalization constant that only depends on . Plugging this into Equation 7 gives that the conditioned Brownian motion has dynamics

The additional drift term is intuitive: it points in the direction of and gets increasingly large as .

Positive Brownian motion

What about a scalar Brownian conditioned to stay positive at all times? Let us consider and let us condition first on the event that the Brownian motion stays positive within and later consider the limit . The function reads

This can easily be calculated with the reflection principle. It equals

for a standard Gaussian . Plugging this into Equation 7 gives that the additional drift term is

as . This shows that a Brownian motion conditioned to stay positive at all times has a upward drift of size ,

Incidentally, it is the dynamics of a Bessel process of dimension , i.e. the law of the modulus of a three-dimensional Brownian motion. More generally, if one conditions a Brownian motion to stay within a closed domain , the conditioned dynamics exhibit a repulsive drift term of size about near the boundary of the domain, as described below.

Brownian motion staying in a domain

What about a Brownian motion conditioned to stay within a domain forever? As before, consider an time horizon and define the function as

One can see that the function satisfies the PDE

and equals zero on the boundary of the domain. Furthermore as for all . Consider the eigenfunctions of the negative Laplacian with Dirichlet boundary conditions on . Recall that is a positive operator with a discrete spectrum of non-negative eigenvalues. The eigenfunction corresponding to the smallest eigenvalue is the principal eigenfunction and it is standard that it is a positive function within the domain , as a “slight” generalization of the Perron-Frobenius in linear algebra shows it. Expanding in the basis of eigenfunctions gives that

Since we are interested in the regime , it holds that

This shows that the conditioned Brownian motion has a drift term expressed in terms of the principal eigenfunction of the Laplacian:

For example, if for a 1D Brownian motion, the principal eigenfunction is . This shows that there is a upward drift of size near and a downward drift of size near .

Manifold Random Walk Metropolis-Hastings

Assume that is the current position of the MCMC chain. To generate a proposal that will eventually be accepted or rejected, one can proceed very similarly to the standard RWM algorithm with Gaussian perturbations with variance . First, generate a vector from a centred Gaussian distribution with covariance on the tangent space to at . To do so, it suffices for example to generate a standard Gaussian vector in and orthogonal-project it onto . Indeed, one cannot simply define the proposal as since it would not necessarily lie on . Instead, one projects back to . To do so, one needs to define the direction used for the projection and the manifold RWM algorithm uses , for reasons that will become clear later. In other words, the proposal is obtained by seeking a vector such that .

If one calls the Jacobian matrix of at , i.e. the matrix whose rows are the gradients of the components of , this projection operation boils down to finding a vector such that

Note that Equation 1 is a non-linear equation in that can have no solution, one solution or many solutions – this can seem like a fundamental roadblock to the design of a valid MCMC algorithm, but we will see that it is not! Before discussing in slightly more details the resolution of Equation 1, assume that a standard root-finding algorithm takes the pair as input and attempts to produces the projection ,

The algorithm will either converge to one of the possible solutions or fail. If the algorithm fails to converge, one can simply reject the proposal and set and set . If the algorithm converges, this defines a valid proposal . To ensure reversibility, and it is one of the main novelty of the article (Zappa, Holmes-Cerfon, and Goodman 2018), one needs to verify that the reverse proposal is possible.

To do so, note that the only possibility for the reverse move to happen is if where

The uniqueness follows from the decomposition . The reverse move is consequently possible if and only if the following reversibility check condition is satisfied,

This reversibility check is necessary as it is not guaranteed that the root-finding algorithm started from converges at all, or converges to in the case when there are several solutions. If Equation 2 is not satisfied, the proposal is rejected and one sets . On the other hand, if Equation 2 is satisfied, the proposal is accepted with the usual Metropolis-Hastings probability

where denotes the Gaussian density on the tangent space The above description defines a valid MCMC algorithm on that is reversible with respect to the target distribution .

Projection onto the manifold

As described above, the main difficulty is to solve the non-linear equation Equation 1 describing the projection of the proposal back to the manifold . The projection is along the space spanned by the columns of , i.e. find a vector such that

One can use a standard Newton’s method to solve this equation started from . Setting for notational convenience , this boils down to iterating

As described in (Barth et al. 1995), it can sometimes be computationally advantageous to use a quasi-Newton method and use instead

with fixed positive definite matrix since one can then pre-compute a decomposition of and use it to solve the linear systems at each iterations. In some recent and related work (Au, Graham, and Thiery 2022), we observed that the standard Newton method performed well in the settings we considered and there was most of the time no computational advantage to using a quasi-Newton method. In practice, the main computational bottleneck is to compute the Jacobian matrix , although it is problem-dependent and some structure can typically be exploited. In practice, only a relatively small number of iterations are performed and the root-finding algorithm is stopped as soon as is below a certain threshold. If the stepsize is small, i.e. , it is typically the case that the Newton’s method will converge to a solution in only a very small number of iterations – indeed, Newton’s method is quadratically convergent when close to a solution.

In the figure above, I have implemented the RWM algorithm above described to sample from the uniform distribution supported on a double torus described by the constraint function given as

The figure shows chains ran in parallel, which is straightforward to implement in practice with JAX (Bradbury et al. 2018). All the chains are initialized from the same position so that one can visualize the evolution of the density of particles.

One can for example monitor the usual expected squared jump distance

and maximize it to tune the RWM step-size; it would probably make slightly more sense to monitor the squared geodesic distances instead the naive squared norm , but that’s way to much hassle and would probably make only a negligible difference. In the figure above, I have plotted the expected squared jump distance as a function of the acceptance rate for different step-sizes. It is interesting to see a pattern extremely similar to the one observed in the standard RWM algorithm (Roberts and Rosenthal 2001): in this double torus example, the optimal acceptance rate is around . Note that since the target distribution is uniform, the rate of acceptance is only very slightly lower than the proportion of successful reversibility checks.

Hamiltonian Monte Carlo (HMC) on manifolds

While the Random Walk Metropolis-Hastings algorithm is interesting, exploiting gradient information is often necessary to design efficient MCMC samplers. Consider a single iteration of a standard Hamiltonian Monte Carlo (HMC) sampler targeting a density on . The method proceeds by simulating from a dynamics that is reversible with respect to an extended target density on defined as

for a user-defined mass parameter . In general, the mass parameter is a positive definite matrix but generalizing this to manifolds is slightly less useful in practice. For a time-discretization step and a current position , the method proceeds by generating a proposal defined as

This proposal is accepted with probability . Indeed, in standard implementation, several leapfrog steps are performed instead of a single one. One can also choose to perform a single leapfrog step as above and only do a partial refreshment of the momentum after each leapfrog step – this may be more efficient or easier to implement when running a large number of HMC chains in parallel on a GPU for example. To adapt the HMC algorithm to sample from a density supported on a manifold , one can proceed similarly to the RWM algorithm by interleaving additional projection steps. These projections are needed to ensure that the momentum vectors remain in the right tangent spaces and the position vectors remain on the manifold ,

As in the RWM algorithm, reversibility checks need to be performed to ensure that the overall algorithm is reversible with respect to the target distribution . The resulting algorithm for generating a proposal reads as follows:

If any of the projection operations fail, the proposal is rejected. If no failure occurs, a reversibility check is performed by running the algorithm backward starting from . If the reversibility check is successful, the proposal is accepted with the usual Metropolis-Hastings probability .

The article (Lelièvre, Rousset, and Stoltz 2019) provides a detailed description of several of these ideas along with detailed analysis and extensions.

Same, but without involution

in some situations, the requirement for to be an involution can seem cumbersome. What if we consider the more general situation of a smooth bijection and its inverse ? In that case, one can directly apply what has been described in the previous section: it suffices to consider an extended state-space obtained by including an index and the involution defined as

This allows one to define a Markov kernel that lets the distribution invariant. Things can even start to get a bit more interesting if a deterministic “flip” is applied after each application of the Markov kernel above describe: doing so avoids immediately coming back to in the event the move is accepted. There are indeed quite a few papers exploiting this type of ideas.

A mixture of deterministic transformations?

To conclude these notes, here is a small riddle whose answer I do not have. One can check that for any , the function is an involution of the real line. This means that for any target density on the real line, one can build the associated Markov kernel defined as

for an acceptance probability described as above,

Finally, choose a values and consider the mixture of Markov kernels

The Markov kernel lets the distribution invariant since each Markov kernel does, but it is not clear at all (to me) under what conditions the associated MCMC algorithm does converge to . One can empirically check that if is very small, things can break down quite easily. On the other, for large, the mixuture of Markov kernels empirically seems to behave as if it were ergodic with respect to .

For values chosen at random, the illustration aboves shows the empirical distribution of the associated Markov chain ran for iterations and targeting the standard Gaussian distribution : the fit seems almost perfect.

Homogenization

Consider the presence of an additional intermediate time scale , with the same assumption that for any fixed the process is ergodic with respect to the probability distribution . The same reasoning as in the averaging case shows that averaging the term is relatively straightforward and has the exact same expression: it suffices to average under . This means that one can study instead

with, informally, . The new interesting phenomenon is coming from the intermediate time scale . Contrarily to the averaging phenomenon of the previous section that was only relying on a Law of Large Numbers, dealing with the intermediate time-scale requires exploiting a CLT and quantifying the rate of mixing of the fast process Note that since , for the dynamics to not explode one needs the centering condition:

Because of the centering condition*, the term will contribute an additional noise term in the effective dynamics of the slow process. To describe this additional noise term, assume an ergodic central limit theorem (CLT) for the fast process : for a test function with zero expectation under we have:

for asymptotic variance . For a time increment and assuming we have

The second integral term is an averaging term that can be treated easily. Approximating the process by , the first integral on the RHS of Equation 5 can be approximated as

After a time-rescaling, one can readily see that the first term is described by the CLT of Equation 4,

The second term is further approximated as

the second equality coming from the time-rescaling . The process mixes on scale so that the term inside bracket converges to its expectation. Setting , one obtains

In conclusion, the fast-slow system

can be described in the regime by the effective dynamics

for two independent Brownian motions and . The volatility terms comes from the CLT and the drift term comes from the self-interaction term:

For the drift function, the scaling may look a bit surprising at first sight as one may expect instead. Note that since the process mixes on a time scale and the centering condition holds, the expectation goes to zero as soon as . This means that only the subset of really matters in that double integral, hence the normalization factor.

Closed form solution & Poisson equation:

The drift and volatility terms and quantify the mixing properties of the fast process . While formulas Equation 6 are intuitive, they can be difficult to deal with if one needs the exact expressions of the drift and volatility functions. Instead, they can also be expressed in terms of the solution to an appropriate Poisson equations.

where the function is solution to the Poisson equation

for all and is the generator of the fast process . The last equality in Equation 7 follows from the integral representation of the Poisson equation. Similarly, and also as explained here, the asymptotic variance term can also be expressed in terms of the function ,

Example: integrated OU process

Consider a slow process obtained by integrating an OU process,

where is just a fixed time-scaling parameter. The fast OU process mixes on time scales of order and has a standard Gaussian distribution as invariant distribution. Homogenization gives that in the regime , the slow process can be approximated as

since the asymptotic variance is

where is the autocorrelation function of the fast OU process, as explained here. The fact that the effective diffusion is (twice) the integrated autocorrelation of the fast process is an example of Green-Kubo relations.

Example: Overdamped Langevin Dynamics

This example does not exactly fall within the homogenization result described in the previous section, but almost. Consider a potential and the slow-fast dynamics:

For any fixed value of , the fast OU-dynamics

converges to a Gaussian distribution with mean and unit variance. The same arguments as the previous section immediately give that, starting from , we have

The terms comes from the OU asymptotic variance. this shows that the slow process converges as to the overdamped Langevin dynamics

Example: Stratonovich Corrections

Consider a function and the slow-fast system

where is a fast OU process mixing on scales of order and with standard centred Gaussian invariant distribution .The discussion leading to Equation 8 suggest that the term can be heuristically be thought of as , which would imply that the effective dynamics for the slow-process is

We will see that this heuristic is wrong! In order to obtain the effective dynamics of the slow process as , since the generator of the fast-OU reads , one can solve the Poisson equation to obtain that . One already knows that . The drift term is given by

Putting everything together gives that the effective slow dynamics reads

where denotes Stratonovich integration.

Readings

The book (Pavliotis and Stuart 2008) is beautiful, and I quite like the section on multiscale expansion in (Weinan 2011). For proving this type of results with the “martingale problem” approach (Stroock and Varadhan 1997), the lectures (Papanicolaou 1977) are nicely done.