Enter the Matrix

A matrix is little else than plain linear scaling but written in the most space and time saving way.

Take the equation \(y = a \cdot x\) which scales a number $x$ by a factor of $a$. What happens if we have multiple numbers and want to scale them all at once? For that we introduce little subscripts to indicate the first $x$ as $x_1$ and the second $x$ as $x_2$. Similarly, we denote the corresponding outputs as $y_1$ and $y_2$ and the scaling turning $x_1$ into $y_1$ as $a_1$ and the scaling turning $x_2$ into $y_2$ as $a_2$. So for example we would have $y_1 = a_1 \cdot x_1$ and $y_2 = a_2 \cdot x_2$ and so on. Stacking the $y_1$ and $y_2$ equations, we can it out as

In the equations above we’re clearly using $x_i$ as a single input to obtain a corresponding output $y_i$. Since Indian mathematicians came up with the concept of $0$ a long time ago, we can also theoretically add all the $x_i$’s to both equations by properly multiplying them by zero:

But now we for each $y_i$ have an expression that involves all the $x_i$’s, not just a single one. If we now want to distinguish that a particular $a_i$ is the scaling we use to scale a particular $x_j$, we need a way to indicate both what input $x_i$ is scaled how for each output $y_i$. So we have to denote each $a$ by it’s input and output indices $a_{\text{output} \leftarrow {\text{input}}}$. So we get

Writing arrows is cumbersome, especially for larger matrices, which is why mathematicians prefer the compact notation and we simply write $a_{ij}$ instead of $a_{i \leftarrow j}$ how input $x_j$ affects output $y_i$. Since this is just a sum, we can also write it as such:

Mathematicians are a lazy bunch, sorry, efficient bunch and they quickly become bored at writing out all these little subscripts $_1$ and $_2$ all the time, so they thought hard about how to write the two equations in a more compact form. The $y$’s, $a$’s and $x$’s only differ in their subscripts so we should be able to collect all the $y$’s, $a$’s and $x$’s into some sort of bundle up object that represents all of them at once, which leads us to the concept of vectors and matrices.

Let’s first collect the $y$’s and $x$’s and introduce introduce a collective object $\mathbf{x} = ( x_1 , x_2)$ and $\mathbf{y} = ( y_1 , y_2)$. This is straight forward as both $\mathbf{x}$ and $\mathbf{y}$ only have one dimension which is denoted by its single subscript. The $a$’s form a two-dimensional array since they have two subscripts which we call a matrix, and we denote it by

\(\mathbf{A} = \begin{bmatrix} a_{11} & a_{12} \\ a_{21} & a_{22} \end{bmatrix}\).

This brings us to the workhorse of linear algebra, the matrix-vector multiplication:

As a simple starting point, consider the diagonal matrix

But wait a second, if we squint really closely, we can rediscover the stacking we did with the vector $\mathbf{x}$ and $\mathbf{y}$ in the matrix $\mathbf{A}$ by interpreting the columns of the matrix as vectors as well:

so our matrix $\mathbf{A}$ can be seen as a collection of vectors as well!

With this representation of a matrix, we can now reinterpret matrix-vector multiplication as a combination of the matrix’s column vectors:

This now opens up a new interpretation as we’ve initially considered $\mathbf{y} = \mathbf{Ax}$ as the scaling of the input $\mathbf{x}$ by their respective $a$ coefficients. Now, we can also see it as forming the output $\mathbf{y}$ by combining the column vectors of $\mathbf{A}$, each scaled by the corresponding component of $\mathbf{x}$. Initially we thought about scaling $x$ with $a$, but now we’re flipping the script and scaling the column vectors in the matrix $\mathbf{A}$ by the components of $\mathbf{x}$.

The animation below highlights both of these approaches: We either interpret it as a scaling of the input $x$ by the values in $A$ through its column space. Alternatively, we interpret the linear transformation as a combination of the column space of $A$ through $x$.

The column space view shifts the analysis from the input $x$ onto the matrix $A$ itself, opening up some interesting options which we will see below. We thus turn the matrix $\mathbf A$ from a simple object of writing scaling $x$ compactly into a first class citizen that warrants its own analysis tools and distinct properties. One of those analysis tools is the rank of a matrix.

Rank of a Matrix

When inspecting the column space of a matrix

we can see that the resulting vector $\mathbf{y}$ is a linear combination of the columns of $\mathbf{A}$. Why? Because we’ve shifted from interpreting $a_{ij}$ as the linear coefficients to interpreting $x_i$ as the coefficient which scales the vectors in the column space of the matrix.

Consequently, we can ask what kind of vectors $\mathbf{y}$ we can generate from $\mathbf{A}$ for a given $x$. The answer to that is closely related to the rank of a matrix Let’s consider the matrix

which coincidentally can be interpreted as the basis vectors of a two dimensional coordinate system. A vector $\mathbf{x} = (x_1, x_2)$ scales the first column by $x_1$ and the second column by $x_2$

The column space of $\mathbf{A}$ consists of two different vectors as they both point into different directions. Essentially both vectors stay out of each other’s lanes and don’t interfere across dimension with each other. But what happens if we flip the entries in the second column of $\mathbf{A}$?

In this case regardless of what we use as input $\mathbf{x}$, the second dimension of our output vector $y_2$ will always be zero. This collapses the second dimension and in effect we can only generate a variable output in the first dimension. The important test is whether we have independent columns in the matrix $\mathbf{A}$ or whether the column space consists of collinear vectors. Collinear vectors simply mean that they all point into the same direction but can have different lengths.

To make it more interesting (as in not squashing things to zero) we can purposefully construct a column space where $\mathbf{a}_2$ is a multiple of $\mathbf{a}_1$, namely $\mathbf{a}_2=1/5 \cdot \mathbf{a}_1$.

Both columns have two nonzero components, but the second remains a multiple of the first. From the column space view of a matrix, we’re essentially creating an output vector $\mathbf{y}$ as a linear combination of vectors that all point into the same direction. They therefore provide only one independent direction, and every output remains on the line to which the columns point, just with possibly different lengths. We say that the matrix $\mathbf{A}$ has rank 1 even if it has more than 1 vectors in its column space, as it provides only one independent direction in its column space.

In the widget below, you can see that we can move the output vector $\mathbf{y}$ by adjusting the components of the input vector $\mathbf{x}$. In the $\text{rank}[A]=1$ case on the right, the output vector $\mathbf{y}$ is constrained to move along a single line, reflecting the fact that the column space of $\mathbf{A}$ has only one independent direction. If all vectors which we scale with $x_i$ point in the same direction, we’re strictly constrained to the subspace that are spanned by those collinear vectors as our basis system doesn’t allow for anything else.

Matrix Inverse

The rank discussion gives us a natural way to understand why some matrices have an inverse and others do not. The rank-one transformation can be viewed as a many-to-one map. So far we’ve covered how we can compute $\mathbf y = \mathbf A \mathbf x$ given some input $\mathbf x$. But now we might be interested to obtain the $\mathbf x$ that created $\mathbf y$ in the first place. So the natural thing to ask is “What is the inverse matrix $\mathbf A^{-1}$ with which we can compute $\mathbf x = \mathbf A^{-1} \mathbf y$?”. Is there anything we should pay attention to in the properties of the matrix $\mathbf A$ before we tackle this idea?

Recall our rank-one matrix

whose second column is a multiple of its first. Its matrix-vector multiplication can be written as

Although we multiplied a two-dimensional input, $x_1$ and $x_2$, into the matrix, the output is determined by the direction of the column space of the matrix scaled by a linear combination of the input components. Whatever two values $x_1$ and $x_2$ we choose, the direction of the resulting output vector $\mathbf{y}$ remains the same and the input only appears for determining the length of the vector. The entire two-dimensional input plane is squashed onto the one-dimensional line spanned by the vector $\begin{bmatrix}2.5 & 1\end{bmatrix}^{\mathsf T}$. Said differently, the $x_1 + \frac{1}{5}x_2$ is a scaling factor but which does not affect the direction of $\begin{bmatrix}2.5 & 1\end{bmatrix}^{\mathsf T}$ but only its length. The problem is many different inputs can lead to the same output, making it impossible to uniquely determine the original input from the output alone. In other words, the matrix $\mathbf{A}$ has lost information about the input $\mathbf{x}$ that cannot be recovered from the output.

This creates a kind of mathematical trapdoor: many different inputs fall onto the same output. For example, consider the two inputs

They are clearly different, but both produce exactly the same output:

Now imagine that we are only given the output $\mathbf{y}=(2.5, 1)$ and are asked to recover the input. Was the input $\mathbf{x}^{(1)}$, $\mathbf{x}^{(2)}$, or one of infinitely many other possibilities? The matrix has discarded the information required to choose a unique answer for $\mathbf x$.

We can see the lost information by subtracting the two inputs to obtain the null space of the matrix $\mathbf{A}$. The null space are all the vectors which get mapped to the zero vector by the matrix.

The direction $(-1, 5)$ belongs to the null space of $\mathbf{A}$. Moving the input in this direction does not move the output at all:

for any value of $t$. All inputs along this line are indistinguishable after applying $\mathbf{A}$. So any vector $\mathbf{x}$ that is a scaled version of $(-1, 5)$ will be mapped to the zero vector by $\mathbf{A}$.

An inverse matrix is supposed to undo a matrix transformation. If

then we would like to apply another matrix $\mathbf{A}^{-1}$ that recovers the original input. This is nothing else than the scalar case of $y =ax$ and getting $x = y / a$:

Doing something and then undoing it should be equivalent to doing nothing, so an inverse must satisfy

But our rank-one matrix cannot have such an inverse. If it did, applying $\mathbf{A}^{-1}$ to the collision above would give

even though the two inputs are different. This is again because multiple inputs $\mathbf{x}^{(1)}$ or $\mathbf{x}^{(2)}$ can be mapped to the same output $\mathbf{y}$ by a rank-deficient matrix.

Below is a little animation in which you can drag and move a vector $\mathbf{x}$ and see how it is mapped by the matrix $\mathbf{A}$ to it’s output $\mathbf{y}$. In the left animation, there is a line of inputs that all get mapped to the same output. Since all points on that line get mapped to the same point, we can not uniquely recreate the input vice-versa.

This gives us the essential condition for invertibility:

A square matrix is invertible exactly when it has full rank.

Determinants of Matrices

The previous sections asked whether a matrix preserves all directions or squashes some of them together. The determinant packages the same geometric story into a single number: it tells us how a matrix scales area.

We begin with the simplest case, a diagonal matrix:

The first column stays on the horizontal axis and the second stays on the vertical axis. Consequently, the unit square becomes a rectangle with side lengths $\lvert a_{11}\rvert$ and $\lvert a_{22}\rvert$, and therefore with area $\lvert a_{11}\rvert\lvert a_{22}\rvert=\lvert a_{11}a_{22}\rvert$. The corresponding signed area is

Drag the two column vectors below. They are constrained to their respective axes, making it visible that each diagonal entry scales one side independently and that their product scales the area.

Now we let the columns point in arbitrary directions:

The transformed unit square is still spanned by the two columns, but its right angles can now lean into a general parallelogram. Its signed area is

This is a deceptively simple equation which you will read in most text books. In order to understand how this simple subtractions comes to be, play around with the widget below which shows how we subtract the triangles from the parallelogram to obtain the term above. The determinant can be calculated by computing the outer bounding box and subtracting all triangles between the outer most bounding box and the actual parallelogram.

The magnitude $\lvert\det(\mathbf{A})\rvert$ is the area-scaling factor. A determinant of $2$ doubles every area, while a determinant of $1/2$ halves it. The sign keeps track of orientation: a positive determinant preserves the ordering of the two directions, while a negative determinant flips it, like turning a sheet of paper over.

This zero-area case is exactly the rank collapse from before. When the columns are collinear (try it out by making both vectors of the parallelogram identical in their direction), the matrix maps the plane onto a line, so $\det(\mathbf{A})=0$ and the lost direction makes the matrix non-invertible. When $\det(\mathbf{A})\neq 0$, the columns span the plane and the matrix is invertible. In higher dimensions the same idea remains: the determinant is the signed scaling factor for $n$-dimensional volume.

The determinant is associative in the sense if we want to know how a matrix $\mathbf{A}$ and a matrix $\mathbf{B}$ change a vector, we can chain the operations and compute each determinant separately, and then multiply them.

This makes intuitively sense because applying two transformations in sequence should scale areas by the product of their individual scaling factors.

Trace of a Matrix

The trace of a square matrix is the sum of its diagonal entries:

For a two-dimensional matrix, this is simply

Recall that $a_{ij}$ describes how input component $x_j$ contributes to output component $y_i$. The diagonal entries are therefore the contributions that remain in the same coordinate: $a_{11}$ maps the first input direction back onto the first output direction, and $a_{22}$ does the same for the second. The off-diagonal entries mix one direction into another and do not contribute to the trace.

At first, adding the diagonal entries may appear arbitrary. Its geometric meaning becomes visible when we apply a very small transformation. Consider the matrix $\mathbf{I}+\delta\mathbf{A}$, where $\delta$ is a small number:

The identity matrix leaves an area unchanged. The additional term $\delta\mathbf{A}$ applies a small deformation. To see how this deformation changes area, we calculate the determinant:

Because $\delta$ is small, $\delta^2$ is much smaller. If we focus only on the first-order change and ignore the $\delta^2$ term, we obtain

The trace therefore measures the immediate rate at which a small transformation expands or contracts area. A positive trace means local expansion, a negative trace means local contraction, and a zero trace means that there is no first-order change in area.

The trace also has a useful cyclic property. Although matrix multiplication generally does not commute, the trace operator actually does $\operatorname{Tr}(\mathbf{B}\mathbf{C})=\operatorname{Tr}(\mathbf{C}\mathbf{B})$. Writing out the diagonal entries makes the reason visible:

The entries $b_{ij}$ and $c_{ji}$ are ordinary numbers, so their multiplication commutes, and we can reorder the finite sums. The only thing we’ve really done is to switch the order of $c_{ji}$ and $b_{ji}$ while also switching out the sum operators. This does not mean that $\mathbf{B}\mathbf{C}=\mathbf{C}\mathbf{B}$, only that their traces agree. For three or more factors, this property allows us to rotate their order cyclically:

We cannot generally swap arbitrary factors, but we can move the first factor to the back.

The “Eigenheiten” of Matrices

What happens to individual directions under the transformation $\mathbf{A}$? Is there a nonzero vector whose direction survives without being turned, changing only in length or orientation?

Such a vector satisfies

We call $\mathbf{v}$ an eigenvector and $\lambda$ its eigenvalue. To be a bit on the nose: An eigenvector only changes its length via its corresponding eigenvalue and not its direction. Rearranging gives

Now remember our null space conversation: a vector lies in the null space of a matrix if and only if the matrix maps it to the zero vector. For a particular eigenvalue $\lambda$, the corresponding eigenvector $\mathbf{v}$ satisfies $(\mathbf{A}-\lambda\mathbf{I})\mathbf{v}=\mathbf{0}$.

Above we stated that the eigenvector $\mathbf{v}$ is nonzero. If it were nonzero, we would obtain the trivial solution $\mathbf{A} \mathbf{0}=\lambda\mathbf{0}$, which is a trivial eigenvector and more importantly would hold for any $\mathbf{A}$ and for any $\lambda$ at all times. So it’s kind of pointless, to be honest. A nontrivial null space means that some nonzero direction with arbitrary length is mapped to zero. Moving an input along this direction does not change its output:

Different inputs along the direction of $\mathbf v$ therefore produce the same output, so no inverse can uniquely recover the original input. Hence, $\mathbf A-\lambda\mathbf I$ is not invertible. A non-invertible matrix has a determinant of zero, which gives us the characteristic equation for finding eigenvalues:

For a two-dimensional matrix, we can write this equation out component by component:

This equation allows us to find the eigenvalues. Once $\lambda$ is known, the corresponding eigenvectors are the nonzero vectors in $\ker(\mathbf{A}-\lambda\mathbf{I})$.

Geometrically, an eigenvector marks a line preserved by the transformation, while its eigenvalue tells us whether that line is stretched, compressed, flipped, or collapsed. These directions are the natural axes of the transformation.

Drag $\mathbf{x}$ below and hunt for a direction that the matrix does not turn.

We can make the relationship between eigenvectors and the column space precise using our rank-one matrix from before:

Its characteristic equation is

so the eigenvalues are $\lambda=2.7$ and $\lambda=0$. Substituting $\lambda=2.7$ into $(\mathbf{A}-\lambda\mathbf{I})\mathbf{v}=\mathbf{0}$ gives

This is just a set of linear equations which we can solve to find the eigenvector $\mathbf{v} = [2.5 \; 1]^\mathsf{T}$ corresponding to $\lambda=2.7$. Substituting $\lambda=0$ similarly gives multiples of $\mathbf{v}=(-1,5)^\mathsf{T}$.

Decompositions of a Matrix

The eigenvector relationship actually gives us a neat entry into a matrix decomposition. We have the eigenvectors and eigenvalues relationship for each individual eigenvector, and by stacking them into matrices, we can express the entire decomposition compactly.

For a single eigenvector, we have

If $\mathbf{A}$ has enough linearly independent eigenvectors, we can stack them as the columns of a matrix

Each column of $\mathbf{V}$ interacts with the matrix $\mathbf{A}$ independently of the other eigenvectors. That’s just the property of a matrix matrix multiply. Conversely, on the right hand side with the eigenvalues $\lambda_i$, we have to scale each of the eigenvectors correspondingly by its eigenvalue:

Multiplying $\mathbf{V}$ by this diagonal matrix scales each column by its corresponding eigenvalue:

We have therefore bundled all the individual eigenvector equations into

Since we’re assuming a full rank matrix, the columns of $\mathbf{V}$ are linearly independent and $\mathbf{V}$ is invertible. Multiplying both sides from the right by $\mathbf{V}^{-1}$ gives

This is the eigendecomposition. Read from right to left, $\mathbf{V}^{-1}$ changes into eigenvector coordinates, $\boldsymbol{\Lambda}$ scales each eigendirection independently, and $\mathbf{V}$ changes back. A matrix with enough independent eigenvectors for this construction is called diagonalizable.

The singular value decomposition extends the same intuition to every real matrix:

Here, orthogonal changes of coordinates surround independent scalings. The number of nonzero singular values in $\boldsymbol{\Sigma}$ is exactly the rank of the matrix.

Matrix Exponentials

A matrix can describe not only one transformation, but also a transformation applied continuously. Consider the linear differential equation

Over a tiny interval $\Delta t$, the state changes approximately as

Dividing a duration $t$ into $n$ tiny updates gives the matrix analogue of repeated scalar compounding:

Expanding this limit defines the matrix exponential:

Thus, $\mathbf{A}$ specifies the infinitesimal change, while $e^{t\mathbf{A}}$ accumulates it into the finite transformation after time $t$.

Previously we’ve encountered the eigendecomposition of a matrix in the form of $\mathbf{A}=\mathbf{V}\boldsymbol{\Lambda}\mathbf{V}^{-1}$. Let’s plug that eigendecomposition into the first three components that we have in the infinite sum of the matrix exponential:

Since $\mathbf V^{-1} \mathbf V = \mathbf I$ by definition of the inverse the intermediate terms between teh $\mathbf \Lambda$ disappear and we get

and we can ever so slightly rearrange the $\mathbf V$ and $\mathbf V^{-1}$ out to isolate the exponential terms:

so each eigendirection evolves independently by $e^{t\lambda_i}$. The power-series definition remains valid even when $\mathbf{A}$ is not diagonalizable.

We can now connect the matrix exponential to two familiar quantities: the determinant and the trace. For a diagonalizable matrix, set $t=1$ and take the determinant. Since determinants multiply and are thus commutative, the change of coordinates and its inverse cancel:

The exponential of a diagonal matrix simply exponentiates each diagonal entry. Its determinant therefore multiplies the exponentials of the eigenvalues:

What is the sum of the eigenvalues? It is exactly the trace of $\mathbf{A}$. Using the cyclic property from the trace section, we move $\mathbf{V}$ from the front to the back of $\mathbf{V}\boldsymbol{\Lambda}\mathbf{V}^{-1}$. This brings $\mathbf{V}^{-1}$ and $\mathbf{V}$ together, where they cancel:

Putting these steps together gives

Although we used an eigendecomposition to derive it, this identity holds for every square matrix, including those that are not diagonalizable.

Einstein Notation: You don’t have to be Einstein to take derivatives of matrices

For one component of a matrix-vector multiplication, we write

Every output component follows the same pattern: fix the output index $i$ and sum over all input indices $j$,

Here, $i$ is a free index: it selects which component of $\mathbf{y}$ we are computing. The index $j$ is a repeated, or dummy, index: it runs over all input components and is summed away. Since a repeated index means “sum over this index,” we can drop the summation sign and write

This is the Einstein summation convention: less notation, but exactly the same matrix-vector multiplication.

Now let’s use it to take a derivative. Consider a product of three matrices of compatible sizes,

and suppose we want to differentiate an output entry $Y_{ij}$ with respect to an entry $B_{pq}$ in the middle matrix. We keep $\mathbf{A}$ and $\mathbf{C}$ fixed and treat the entries of $\mathbf{B}$ as independent variables. The matrix product may look complicated, but writing out one output component reveals nothing more than a sum:

The indices $i$ and $j$ select the output entry, while the repeated indices $k$ and $l$ are summed over. For two-dimensional matrices, we can write all four terms explicitly:

Despite all the matrix notation, this is just a linear expression in the entries of $\mathbf{B}$. Let’s differentiate with respect to $B_{12}$. The terms containing $B_{11}$, $B_{21}$, and $B_{22}$ do not depend on $B_{12}$, so their derivatives are zero. In the remaining term, $B_{12}$ differentiates to $1$, while the surrounding factors stay:

For an arbitrary entry $B_{pq}$, the same idea can be written compactly using the Kronecker delta:

An entry $B_{kl}$ depends on $B_{pq}$ only when both its row and its column match:

Substituting this into the component-wise derivative gives

The deltas select $k=p$ and $l=q$, so only the coefficient of $B_{pq}$ survives. In Einstein notation, we can skip the summation signs altogether:

So the trick is not to wrestle with the whole matrix expression at once. Write one output component, identify which terms contain the entry we are differentiating, and keep its coefficient. The chosen entry differentiates to $1$; all terms that do not contain it disappear.

Let’s try three slightly more involved examples. All matrices below are real and have compatible dimensions, and only $\mathbf{B}$ varies. The matrix gradient collects the entry-wise derivatives into a matrix with the same shape as $\mathbf{B}$.

1. The trace of three matrices

How do we differentiate with respect to the middle entry $B_{pq}$, and how do we collect the result into the matrix gradient? Remember that the trace sums the diagonal entries of the output matrix.

2. A matrix that appears twice

Find the derivative with respect to $B_{pq}$ and the matrix gradient. This time $\mathbf{B}$ appears twice, so both occurrences contribute to the derivative. Do not assume that $\mathbf{A}$ is symmetric.

3. A squared error through three matrices

Suppose we want $\mathbf{A}\mathbf{B}\mathbf{C}$ to match a fixed target matrix $\mathbf{D}$. Consider the loss

where the squared Frobenius norm is the sum of the squares of all entries. Find the derivative with respect to $B_{pq}$ and the matrix gradient. The outer square introduces a chain rule, but the derivative of the inner matrix product is one we already know.

There has been a tremendous interest in doing autoresearch with LLM agents. The core idea is that since LLM’s have intelligence (the amount for real world problem solving is still being debated), have infinite stamina and have access to the abundance of information on the web including all recent research, then they could automate parts of the research pipeline that we all do. Evidently, the writing was on the wall and there is a whole swath of start ups that try to capitalize on this idea with the idea of an AI scientist (Periodic Labs, FutureHouse etc).

LLM’s rely on the iterative filling of their context (to have context at all) and you can add a RAG system, but then you’re still dependent on what’s in your context. So whatever a LLM does to solve a problem is highly influenced by what’s in it’s context. Of course there is a lot of baseline intelligence in LLM’s, but the context is still a huge factor in what it can do.

But research is not a linear process in the same way that we linearly and monotonically fill up a context. Research consist of many rabbit holes, dead ends and resets. This resetting is something that LLM’s are not ideally placed to do due to the their dependence on the context when trying to find solutions to problems by trial and error.

That begs the question: can we wrap a harness around a LLM to make it more like a human researcher that’s able to deal with setbacks, rabbit hole failures and resetting the context when needed?

Graduate student descent

There’s an old joke in machine learning that the real optimization algorithm behind most papers isn’t Adam or SGD. It’s graduate student descent: a PhD student sits in front of a screen, semi-informed, and tweaks things. Bump the learning rate. Swap out a loss term. Try a slightly different architecture. Rerun. Squint at the validation number. Try again.

Most of this is busywork — hyperparameter fiddling, “what if I change this one thing,” waiting for the run to finish. Only a handful of the tries actually matter, and the student usually can’t tell in advance which ones, which is exactly why they have to keep trying.

Having done a PhD in machine learing I’m all too familiar with this process. I have spent countless hours doing this exact loop, it’s unfortunately necessary to get a good result, but it’s also mind numbing and not very fun when you haven’t found the magic combination yet. I doubt we’ll be able to do away with graduate student descent as machine learning is intrinsically a latent model problem where try to explain/model/predict data and generalize from it. If you’ve done machine learning research you’ll know th Bermuda triangle of hope between model architecture, parameter optimization and data quirks.

Your graduate student of choice would

• proposes a change, based on a hunch about what might help,
• evaluates it by running a training job and reading off a number,
• and then decides what to try next — usually by leaning into whatever’s been working, but every so often taking a flyer on a direction they’ve barely touched.

It’s a semi-informed search over the tree of possible experiments, steered by intuition and metrics. Intuition is hard to come by if you haven’t done a lot of research before, but LLM agents offer boundless stamina and will never stop trying stuff out. This has been epitomized in the Ralph loop where a semi-intelligent actor just hammers the problem over and over again. Or it’s your new grad student is still optimistic about the next experiment.

Branching and Resetting

The other modus operandi we usually face in research is that we often have to reset our context and start over. That’s the humiliating act of realizing that your ideas have failed and we need to figuratively and literally clean the white board.

Here’s my biggest gripe with using LLM’s for autoresearch: They’re hard at resetting their context and starting over in a principled way. The original idea of autoresearch by Karpathy was to use git to reset to the previous state if something failed but it still left the context unchanged. So fundamentally, while the code was reset the context wasn’t. It essentially set the autoresearch agent on a linear trajectory of git states with no way of completely changing track and pursuing a different line of research.

But git allows branching such that we could try different ideas from the same root node. We can try experiments and log the code state in the git tree and then branch off to try a different idea. Importantly, this allows us to keep a ledge or diary of sorts of the many experiments which can develop in a multitude of different ways. If we want to revisit a prior experiment, we simply have to checkout the git commit and we have the exact code state that was used to run that experiment. That allows us to jump around in the tree and pursue different lines of research.

Monte Carlo Tree Search

So up to now we’ve concluded that LLM are useful testing out ideas with unwavering enthusiasm and that git is a perfect way to log the experiments and keep track of the many different lines of research. The last missing piece of the puzzle is to combine the two in a principled way such that we can explore the tree of experiments in a way that balances exploration and exploitation.

A long while ago in 2017, I gave a talk on AlphaGo and the Monte Carlo Tree Search algorithm that they used as a fast rollout policy to explore the state space tree of a game of Go. Mind you, Go is a game with a huge state space and the number of possible moves is astronomical. So they used a combination of a neural network to evaluate the board state and a Monte Carlo Tree Search to explore the tree of possible moves.

In the Monte Carlo Tree Search algorithm we work fundamentally on a tree structure where each node is a state and each edge is a possible action that can be taken from that state.

MCTS is a way to explore the tree of possible states in a principled way that balances exploration and exploitation. In the case of AlphaGo we simply counted the number of times a node was visited and the number of wins that were achieved from that node.

As the saying goes “An image is worth a thousand words” so let’s look at the MCTS algorithm in a visual way:

← Prev

1 / 6Game tree

Next →

Initially we start out with a state tree with all the moves with their wins and visits.

• Select: We use a formula to compare all non-terminal child nodes in the tree and decide to which child node to go first. In this simple case we simply compare the number of wins in the subtree to all visits of the subtree.
• Expand: Reaching a terminal leaf node, we chose an action on how to change the state. In Go we would place a stone on the game board.
• Simulate: We then simulate the game from the new state to see the outcome. This could be a cheap and fast rollout policy until we win or loose.
• Backpropagate: We log the win (or loss) and backpropagate the ratio up the tree, automatically increasing the number of wins of any node between the leaf node and the root. This increases the percentage of wins all of those recursive subtrees.
• Action: With one win or loss recorded, we can choose the best action based on the updated statistics. We use a policy that does a bit of exploration and chooses the middle subtree.

Below is an animation of a MCTS search for a simple analytical case to visualize the search character. The reward is a simple function of a binary action whether to add a 0 or a 1 to a string of length 5. The reward is the number of 1’s in the string. The search starts at the root node and explores the tree of possible strings until it finds the optimal string of all 1’s. After having found the gren and dark green combinations the algorithms keeps on exploring the lesser valuable paths to make sure that it doesn’t miss a better solution.

Translating MCTS to Autoresearch

MCTS gives us a principled way to explore a tree of possible states all originating from a common root node and searching the tree for the best state.

In order to get the performance of a particular node, we define a validation script that isn’t allowed to change and returns a single metric that gauges the performance of the state. This metric is then used to compare the performance of different nodes in the tree.

MCTS was originally formulated for games which have a binary outcome: win or loss. In machine learning research we often have a continuous outcome (loss) and we want to minimize that loss. The simple switch is to set an upper and lower (preferabbly zero but negative loglikelihoods can go lower) bound and then just rescaling the validated performance in those bounds to a [0,1] range.

The selection is a formula that compares the performance of all child nodes and chooses the best one to explore further. A particular node can have children which are already validated or (pre-generated) proposals which still require an expansion.

An expansion is the code change that we did that changes the code base of a particular node in the tree. This could be a change in the model architecture, a change in the loss function, a change in the optimizer or a change in the data set.

After each expansion we run the validation script to see how good our change performed. This is the simulation step in MCTS. Of course our proposal can be so bad that we run immediatley into NaN’s or other problems. In that case we simply ignore the proposal and mark it as failed. Failed nodes have no repercussion and are simply skipped in the selection step and represent a dead end.

Once the simulation/validation ran, we backpropagate the [0,1] value up the path between the leaf node and the root node, updating the performance of all nodes in between.

Here’s an interactive example with random best values. You can play it a couple of times ot see how

Notice how the search keeps coming back to the best branch — but never only that branch. Every so often it wanders off to try something it has mostly ignored. That tension, between milking what already works and checking on what you’ve neglected, is the whole game. Hold onto it; in a moment we’ll see the exact formula that produces it.

Keeping things simple

The idea is all nice but the devil’s in the detail as we all know.

For the implementation I opted for a some very hard boundaries.

• Simple: The whole setup should consist of a single functional file and maybe an additional utils file
• A Single Contract: There is no try/except (looking at you coding agents), defaults, if/else and catch mechanism. There is a single way to run this simple set up, if you don't it will fail.
• YAGNI (Recycling git): Git provides everything we need via inheritance, branching, checking out commits in a super convenient package. We will use as much of git as possible.
• Extendable Base Case: With the three paradigms above in a single functional python script, it can easily be extended for more complex scenarios to your liking.

Wait, how do we inject “intelligence” into the whole setup? It’s simple: we use the Python SDK for Copilot/Claude. With everything we’ve talked about so far, every step of the way is a simple prompt to the LLM with a very specific task and a very specific context. Essentially, we have a very clear idea of what we want to do and we can prompt the LLM to do it for us. The thing is that we can actually isolate the LLM to whatever it’s supposed to do and not let it wander off into the weeds of the code base. Since everything runs on git, we can run git commands manually from within the research loop, and start a coding session with the LLM to implement a particular proposal. After the part with the intelligence is done, we can switch back to normal programming and run the validation script from within our loop. Since the agents get a fresh context and a clear cut, localized and fence off task to do, we manage everything else around them with a mini harness.

Basically, the MCTS algorithm, the validation script, the git plumbing, the proposals calls, the implementation calls are all simply python commands run from a single script. The agents might screw up, but we have the harness around them to make sure that they don’t screw up the whole system and can simply pursue a different line of research through the MCTS formalism.

The Plumbing under the hood

Once we’ve expanded the tree to a new node, it gets validated and we run a proposer loop that attaches five distinct proposals in text format to the node. Each proposal has a promise score and a rationale for why it might work.

Then we run a selection algorithm on the node and its children to decide which child node to explore next. The most basic form would be to run PUCTS (Predictor + Upper Confidence bound applied to Trees) which is a simple formula that balances exploration and exploitation:

Here, $Q(s,a)$ is the average value of the child node $a$ from state $s$, $N(s)$ is the number of times state $s$ has been visited (includes the subtree attached to that node), and $N(s,a)$ is the number of times action $a$ has been taken from state $s$. The prior $\pi(a \mid s)$ is the promise score of the proposal that created this child node, and $c$ is a constant that controls how much we favor exploration over exploitation.

It looks like a mouthful, but it’s just those two instincts sitting side by side:

• $Q(s,a)$ is how good this looks — the best score seen so far down this branch (and for a brand-new idea, we optimistically borrow the parent’s score, so it isn’t punished for being untried).
• The second term, $U$, is how under-explored it is — large when we’ve barely visited this branch and when our gut feeling about it, the prior $\pi$, was high. It shrinks a little every time we come back.
• $c = 0.5$ is just a knob for how adventurous we are overall.

The lovely part is how the balance shifts on its own. Early on, $U$ dominates and the search fans out, sampling lots of directions. As a branch piles up visits, its $U$ collapses and its hard-won $Q$ takes over. Exploration quietly hands off to exploitation, with no schedule to tune by hand.

With a small probability $\varepsilon$ we can also flat-sample a node and try a proposal from there, which is a simple way to keep the search from getting stuck in a local optimum. So essentially from time to time, we’ll look at all $N(s)$ and pick a node at random to explore, instead of always picking the best $Q(s,a) + U$. That’s a little bit of jitter to make the search jump out of a local optimum and explore other branches which might be blocked due to a series of bad validations.

What gets stapled to each commit. Earlier I waved vaguely at git “notes.” Concretely, each live commit carries a small JSON note recording its score, the proposal that created it, and the proposals still waiting to be tried from here:

{
  "metrics": {"loss": 0.0623},
  "winner":  {"plan": "CODE. Edit losses.py…", "promise": 0.72, "rationale": "…"},
  "open":    [{"plan": "HYPERPARAM…", "promise": 0.45, "rationale": "…"}],
  "state":   "evaluated"
}

The nice thing about this is that the entire tree is self-contained in git. Git does all the tracking. The metrics are the result of the validation script. I also opted for the script to parse in the entire tree into a nested tree data structure and do all the calculations in memory, then write the notes back to git at the end of each iteration. That way, if the machine dies mid-run, the tree is still consistent and we can pick up where we left off. The requires us to parse the whole git tree once per iteration, but the tree is small enough that it doesn’t matter. Most of our time will be spent anyways on GPU compute.

winner is the proposal that produced this commit (it lives on the child, not the parent); open is the queue of ideas not yet tried from here. The entire tree is rebuilt from these notes on every wake-up, so there’s no second copy of the state to fall out of sync. Each run namespaces its notes under an ID like YYYYMMDD_HHMMSS_µs-happy-otter-abc12345, so several searches can branch off the same starting commit without colliding.

The verify contract. After each commit the orchestrator runs your evaluation command and reads back a JSON file that has to contain at least {"loss": <number>}. To stop the agents from “improving” the metric by quietly rewriting the scorer, the verify script is hash-locked on the first run: its checksum is pinned, and if it ever changes, the run aborts on the spot.

The three states a node can be in. Not every experiment succeeds, and the tree has to cope. A node is one of:

evaluated — verify succeeded and a loss was recorded. The search can keep descending from here.

failed — verify errored or timed out. The search can still descend from here, since a fix might live one commit below, but the node is scored pessimistically so the search routes around it without giving up on the branch.

terminal — a genuine dead end, either declared by the verifier or because every one of its children is itself terminal. The search will not descend here again.

That last rule is worth dwelling on, because terminal deadness bubbles upward. The moment a node has no ideas left in its queue and every one of its children has gone terminal, it turns terminal too — so the exhaustion of an entire subtree propagates quietly back toward the root, one dead node at a time, and the search never wastes another thought on it.

One full iteration, end to end:

1. Rebuild the tree from git notes (pure reads)
2. Backpropagate once — post-order walk filling n, subtree_value, cascading terminal states
3. With prob. $\varepsilon$: flat-sample a node and append one proposal → go to 7
   Otherwise: PUCT-descend to pick (node, proposal_index)
4. Checkout that node’s commit
5. Implement: run the implementer agent, verify the diff, commit the child
6. Evaluate: run the verify script, record loss and state
7. Seed: run the proposer at the new node to populate its open list
8. Save: write notes for child and parent

The order is deliberately crash-proof: if the machine dies anywhere in the middle, the tree on disk is still consistent — either the new commit’s note exists and we carry on, or it doesn’t and the parent simply re-tries that idea next time. The entire footprint on disk is a couple of git refs, a checksum, and a log file. No database, no cache, nothing to corrupt.

Every time you start a new run, a random string is created with the root hash in it. This serves as a marker that we’ll only consider any git notes that are marked under this ID. It also automatically denotes our root such that we can keep running multiple runs from the same root in parallel and from different machines without any collisions. This also enables the script to run iteratively such that you can generate new runs on the fly based off previous autoresearch runs. If you know git, you don’t have to learn new concepts.

The whole thing is a single Python file of around 1700 lines, standard library only. Drop it into a git repo, point it at your task and your evaluation command, and let it climb overnight.

The thing I keep coming back to is what got automated and what stubbornly didn’t. The busywork — the “let me just bump this and see,” the overnight grind of running and re-running and squinting at numbers — turns out to be genuinely mechanizable. A tree, a metric, and a formula will happily do all of it while you sleep. But the taste — knowing which direction is worth a hunch in the first place, recognizing that a boring-looking number is secretly the interesting result.

So we’ve automated graduate student descent. We have not, it seems, automated the graduate student.

You can find the code here.

Fokker-Planck Equation

The Fokker-Planck equation describes the time evolution of the probability density function $p(\mathbf{x}, t)$ of a stochastic process:

where:

• $\mu(x, t)$ is the drift coefficient (deterministic force)
• $\sigma^2(x, t)$ is the diffusion coefficient (related to the noise strength)
• $\nabla$ is the gradient operator
• $\Delta = \nabla^2$ is the Laplacian operator

What this essential equation tells us is how the probability density of a system’s state changes over time due to both deterministic forces (drift) and random fluctuations (diffusion). The first term on the right-hand side represents the effect of the drift, while the second term accounts for the diffusion. If we compute $\partial_t p(x,t)$ for a given $p(x,t)$, we can determine how the probability distribution evolves over time, which is crucial for understanding the dynamics of stochastic systems. The nice thing is that normalization is automatically baked into this equation. The total probability is conserved, meaning that if we integrate $p(x, t)$ over all possible states $x$, it will always equal 1 for all time $t$.

The product of distributions

The product of two distributions $p(x)$ and $q(x)$ is defined as:

Why does the integral in the denominator appear? It ensures that the resulting distribution is properly normalized, meaning that the total probability integrates to 1. The product $p(x) q(x)$ represents the unnormalized joint distribution and therefore has the correct shape, and dividing by the integral $\int p(x’) q(x’) dx’$ normalizes it to become a valid probability distribution.

Reward Tilted SDE

We consider the product of the original distribution $q_t(x)$ and an exponential of a reward function $r_t(x)$, scaled by a factor $\beta_t$:

The big question is: what is the Fokker Planck equation of $p_t(x)$ from which we can easily read off the SDE if the process evolves according to the original process dynamics of $q_t(x)$? Said differently, what’s the probability of a sample $x$ under the tilted $p_t$ if it evolves according to the stochastic differential equation of $q_t$?

So we need

where $\mu$ and $\sigma_p^2$ are the drift and diffusion of the SDE that generates $q_t$.

In order to obtain the FPE for the tilted distribution $p_t(x)$, we will exploit the FPE in log space, which offers some nice algebraic advantages. To realize that we seek to obtain the log FPE, namely not $\partial_t p_t(x)$, but

FPE in log space

Let us first recall the log-derivative trick which we will use extensively in the following calculations. For any function (or distribution) $q_t(x)$, we have the following identities:

This allows us to substitute all the gradient calculations in the FPE of our original process $q_t$ with the log derivatives,

The term $\partial_t \log q_t$ is log transformed Fokker Planck equation and governs the evolution of the original distribution $q_t$ in log space.

How $q_t$ (or $\log q_t$) evolves over time is determined by the drift $\mu_q$ and the diffusion $\sigma_q^2$ of the original process. We’re able to model that with the FPE of the original process $q_t$, but here we’re actually interested in the dynamics of $p_t$. So the question is whether we can somehow massage the term

into the form of the FPE from which we can then read off the drift and diffusion of the SDE the tilted process $p_t$.

In fact, we can express the change $\log q_t$ in terms of the sought-after process $p_t$ by rearranging the terms,

The derivatives of $\log q_t$ are

Substituting these into the log FPE of $q_t$ gives us

And what do we spy with our little eyes? The first four terms are actually the log FPE of $p_t$ itself!

We we plug in $\partial_t \log q_t$ into $\partial_t \log p_t(x)$, we get

In order to go back to the non-log FPE, we just have to multiply both sides by $p_t$ and apply the log derivative identities from above in reverse,

The last thing is how we actually compute $\partial_t \log Z$. We can compute $\partial_t \log Z$ by differentiating the definition of $Z$. For notational simplicity, let us define $\tilde{p} = \int q_t \exp(\beta_t r(x)) dx$

It should immediately be obvious that the term $\partial_t \log q_t + \partial_t \beta_t \ r(x)$ is the same term into which we’ve inserted the log derivatives. So we expand the $\log q_t$ in the same way as we did above, and then multiply it with the $p_t$ in the integral, we get

which is literally just repeating the same steps that we did above to transform the log FPE to the FPE.

One of the key properties of the Fokker-Planck equation is that it conserves probability, meaning that the total probability integrates to 1 for all time $t$. So if we integrate all the outflows and inflows of probability across the entire state space, they should balance out to zero, ensuring that the total probability remains constant. If at any time step $t$, the change in probability $\partial_t$ over the entire space $x$ would not be zero, it would imply that probability is either being created or destroyed, which violates the fundamental principle of probability conservation. So if from one $t$ to the next $t+dt$, all the changes in probability across the state space would be $-0.1$, it would imply that the total probability has decreased by 0.1 which would render our distribution $p_t$ invalid as it would no longer integrate to 1.

But in fact, we can show this analytically by integrating the FPE of $p_t$ over the entire state space. For the left hand side, we have

where we’ve used Leibniz rule where we can interchange differentiation and integration for different variables. For the right hand side of the FPE, we have

Using integration by parts,

we can exploit that fact that for any reasonable probability distribution $p_t$, the probability density at the boundaries of the state space should be zero, meaning that $\lim_{x \to \infty} p_t(x) = 0$ and $\lim_{x \to -\infty} p_t(x) = 0$.

Thus we have,

where we’ve used the fact integration and derivatives cancel each other out while preserving the boundary evaluation conditions of $p_t$ at $\pm \infty$.

We can also play the integration by parts game one more time and show that

All of this let’s us rewrite the FPE of our tilted distribution $p_t$ as

Is it in fact quite interesting that the change of probability equation that we’ve derived consists of two different parts: the original FP equation of $p_t$ and an additional term that captures the effect of the reward tilt on the probability distribution.

The first term consists only of the original drift $\mu_q$ and diffusion $\sigma_q^2$ of the original process $q_t$, which implies that we can model the original SDE as is. It is only the second term that introduces the effect of the reward tilt, modifying the probability distribution according to the function $g(x)$. Every term in the second part $g(x)$ contains in fact the reward function $r(x)$, whether it be $r(x)$ directly or its gradients $\nabla r(x)$ and $\Delta r(x)$, which means that the reward function is the only source of change in the probability distribution $p_t$.

The Feynman-Kac Equation

The Feynman-Kac equation is a fundamental result that links partial differential equations (PDEs) to expectations of stochastic processes. For a function $p$ satisfying certain boundary conditions, it defines a PDE.

For a SDE for a particle $x_t$ defined as

We can define two key operators: The first one is the the generator that acts on the test function and describes its instantaneous rate of change of the test function $\theta$ along the paths of the SDE. We’ve seen this in the Feynman Kac equation where we used it to define

where the FK equation is defined as $\partial_t \phi +\mathcal{L}_t[\phi] + g \ \phi = 0$.

The adjoint generator acts on the probability distribution $p$ and describes how the probability distribution evolves over time under the dynamics of the SDE. It is defined as

Both generators are linked through the identity

Now, we’re interested in some test function on $x$ that we name $\phi(x, t)$ and we’re interested in what the distribution of $\phi(x,0)$ and the end of our reverse diffusion process is. We have defined the probability evolution of $p$ to be governed by the Feynman-Kac PDE,

This tells us how the probability distribution $p_t$ evolves over time under the influence of both the original dynamics of the SDE (captured by the first two terms) and the reward tilt (captured by the last term).

Connecting it to the Feynman Kac equation

In a previous blog post on the Feynman-Kac equation, we saw that for very particular PDE that are defined in the state space of the SDE, we can express the solution of the PDE as an expectation over the trajectories of the SDE. Namely, we have a PDE of the form

subject to the terminal condition $u(x,T) = \Phi_T(x)$, then the solution of this PDE is given by the Feynman-Kac formula. A straightforward option would be to evolve this very particular PDE backward in time form $T$ to an earlier $t$ and then read off the solution at $t=0$. This would give us $u(x,t)$.

If the dynamics of $x$ are governed by the Ito drift-diffusion process defined above, then the solution of the PDE is given by the Feynman-Kac formula

where the expectation is taken over all trajectories of the SDE that start at $x$ at time $t$ and evolve according to the dynamics of the SDE until time $T$ with values $x_T$.

We now have two differential equations

The samples $x$ evolve according to the $dx = \mu_t dt + \sigma_t dW_t$ and induce the stochastic process $q_t$ but we’re interested in the distribution $p_t(x) = \frac{1}{Z} q_t(x) \exp(\beta r(x))$ and want to know how likely it is to sample $x$ from $p_t$ if we evolve according to the SDE of $q_t$. The function $g$ consists almost exclusively of terms that contain the reward function $r(x)$, which nicely shows that any change in the probability distribution $p_t$ is solely due to the reward tilt encapsulated in $g$.

Let’s assume a test function $u(x,t)$ with some terminal condition $u(x,T)$ and we evolve it backward in time according to the Feynman-Kac PDE over the trajectories of the SDE defined by $\mu$ and $\sigma$. As we’re evolving some terminal function backwards in time, this is obviously very closely linked to the Kolmogorov backward equation.

The question the FKC correctors try to answer is: Given my samples $x$ generated from $q_t$ but tilted towards some reward function $r(x)$, what is the expectation of $u(x,T)$ if I were to evolve further from $t$ to $T$ according to the original SDE of $q_t$?

We’re interested in the expectation

In the equation above, $q$ generates the samples $x$ according to the SDE, $r(x)$ tilts our distribution towards higher rewards such that our samples $x$ should exhibit certain characteristics favoured by $r(x)$ and then we want to know the expected value of $u(x,T)$ under this tilted distribution.

Now we can ask ourselves how $\int u(x,T) \ p_T(x) dx$ evolves over time. To make things a tad easier, we’ll for now just consider the unnormalized distribution $\tilde{p}_t(x) = q_t(x) \exp(\beta_t r(x))$ and then we can later on divide by (some estimate of) $Z_t$ to get the normalized distribution $p_t$. Fortunately, every component in $\partial_t p_t$ is a function of $p_t$, so we extract $1/Z_t$ out of every component to get

The proof for this is in the FKC paper, but it is quite intuitive if you think about it. The original FPE of $p_t$ is linear in $p_t$, so if we replace $p_t$ with $\tilde{p}_t / Z_t$, we can just pull out the $1/Z_t$ from every term and get the FPE of $\tilde{p}_t$. The authors show the consistency of this in the appendix.

Thus we’re want to know

behaves. Using Leibniz rule, we can interchange differentiation and integration for different variables, which gives us

The integration by parts substitution is a pure algebraic busy work and I skipped it here. It’s just applying the product rule many times to fully separate every term, then applying integration by parts to every term by pulling in $u$ and recombining the terms.

Integrating the dynamics from $0$ to $T$ gives us

All that remains is substituing back the Feynman-Kac formula for $u(x,0)$ and substituting $\tilde{p}_t = p_t Z_t$, which gives us

Since at $t=0$, we have $p_0(x) = q_0(x) \exp(\beta_0 r(x)) / Z_0$ with $\beta_0 = 0$ we get $p_0(x) = q_0(x) / Z_0$ with $Z_0 = 1$, we can simplify the right hand side to get

This is pure importance sampling on a trajectory level, where we sample trajectories from the original SDE, calculate $g(x_s, s)$ along the trajectory and simply integrating it in log space as a simple sum over the time steps. In the SMC blog post we saw that we can use the weights from SMC to estimate $\hat{Z}_T = 1/K \sum_k^K \exp(\int_0^T g(x,s) ds)$, which gives us a practical way to compute the expectation of $u(x,T)$ under the tilted distribution $p_T$ by sampling trajectories from the original SDE and weighting them according to the reward tilt encapsulated in $g$.

Why we need it

Suppose you want to estimate $\mathbb{E}_p[f(X)]$ for some target distribution $p$. In a perfect world, you’d draw $N$ i.i.d. samples from $p$ and compute the sample mean. The variance of that estimator drops as $1/N$ — that’s what you learn in ML 101.

If we can draw i.i.d. samples from $p$, then there’s no real problem in computing the expectation. But alas, we can’t always draw i.i.d. samples from $p$.

But sampling from $p$ is not always possible. This can be for a variety of reasons, but one common reason is that the distribution $p$ is complex and high-dimensional and its normalization constant is not known. Oftentimes, we only know $p$ up to a normalizing constant $p = \tilde{p} / Z$. Without knowing $Z$, we can’t assign the correct probability to a point $x$ under $p(x)$. For the same value of $\tilde{p}(x)$, the probability of $x$ under $p(x)$ is different if $Z$ is different.

The mathematical version of this predicament shows up everywhere. Consider the workhorse of modern machine learning: Bayesian posterior inference. You observe data $\mathcal{D}$ and want to reason about parameters $\theta$ under the posterior

The numerator — likelihood times prior — you can evaluate pointwise, no problem. But the denominator $p(\mathcal{D})$, the marginal likelihood, requires the probability of your entire dataset $\mathcal{D}$ under the model.

So either you somehow magically know the probability of every sample in your dataset $\mathcal{D}$ under the model, or you need to estimate it by defining every possible value of $\theta$ and evaluating the likelihood and prior for each of them. Every possible value of $\theta$ for continuous parameters is an infinite number of values, so you’d need an infinite amount of computational resources to evaluate the likelihood and prior for each of them.

For anything beyond toy conjugate models, this integral is intractable: there is no closed-form solution and the dimensionality of $\theta$ makes brute-force numerical integration hopeless. So you’re stuck with a distribution you can evaluate (up to a constant) but can’t sample from directly. You know the shape of the landscape, but you have no way to generate points that are distributed according to it.

This is where sampling schemes come in. But in practice, the samples these schemes produce are almost never i.i.d. from the target:

• MCMC produces samples that are serially correlated (each sample depends on the previous one).
• Importance sampling draws from a proposal $q \neq p$, then corrects with weights — but those weights can be wildly uneven.
• Particle filters resample with replacement, creating duplicate particles (see Annealed Importance Sampling for a related staged-sampling scheme).

All of these schemes are ways to sample from $p$ even if you can’t sample from $p$ directly. In all these cases, your $N$ raw samples carry less information than $N$ truly independent samples from the OG $p$ would. You need a single number that answers: “how many i.i.d. samples is my collection of samples I proposed actually worth?” That number is the ESS.

Without it, you’re flying blind — you can’t meaningfully compare algorithms, diagnose poor mixing, set stopping criteria, or calibrate your confidence in Monte Carlo estimates.

Let’s get the equations rolling

Importance Sampling

We want to compute:

We can’t sample from $p$ directly, so we draw $X_1, \dots, X_N \overset{\text{i.i.d.}}{\sim} q$ from a proposal distribution $q$, and use the importance sampling identity:

where $w(x) = p(x)/q(x)$ are the (unnormalized) importance weights.

Below you can see the importance weights $w(x) = p(x)/q(x)$ for a Gaussian mixture model from the cover. If we sample from $q$ we have to take into consideration that we’re not sampling from $p$ directly. The correction for sampling from $q$ but wanting to estimate the expectation under $p$ is given by the importance weights $w(x) = p(x)/q(x)$. For all examples in the visualization below, $p(x) < q(x)$, such that the importance weights will all be smaller than 1. These weights $w(x) < 1$, will thus downweight the contribution of the samples from $q$ in the expectation estimator,

whereas parts where $p(x) > q(x)$, the importance weights will be greater than 1, thus upweighting the contribution of the samples from $q$ in the expectation estimator.

Deriving the self-normalized estimator

The problem with standard importance sampling

If we know $p$ exactly (including its normalizing constant), life is simple. The standard IS estimator is:

This is unbiased:

But in practice — especially in Bayesian inference — we only know $p$ up to a normalizing constant $p = \tilde{p} / Z$. While we’re interested in estimating $w = p / q$, we’re left holding the bag with its unnormalized cousin where $Z$ is unknown,

which makes us arrive again at the same spot we started from that we don’t know $Z$ in the first place. But with unnormalized weights $\tilde{w}$, there is a trick we can pull to bootstrap ourselves out of this seemingly dead end.

The key observation is that the normalizing constant $Z$ is itself is an integral over the whole space $x$ of $\tilde{p}(x)$. We can therefore use importance sampling to evaluate the integral as in

Given that we’ve generated the samples $x \sim q$ proportionally to the distribution $q$ by definition, the expectation reduces to a simple average over the samples $x$,

Going back to our original objective we want to identify

Whereas the importance estimator \(\hat{\mu}_{\text{IS}}\) assumes access to the target distribution $p$ including its normalization constant $Z$, using the unnormalized weights \(\tilde{w}\) to estimate \(\hat{Z}\) turns this estimator it to the self-normalized estimator \(\hat{\mu}_{\text{SN}}\).

We can simplify the equation for $\hat{\mu}_{\text{SN}}$ even further by observing that the sum in the denominator is self contained (as in sum over $j$), so we can write

where the normalized weights $\bar{w}$ are defined as

How much does our estimator actually vary?

The self-normalized estimator $\hat{\mu}_{\text{SN}}$ is a ratio of two random quantities. This makes its variance analysis more subtle than the simple i.i.d. case.

The key question is: how much does our estimate vary from one set of samples to the next? If we draw $N$ samples from $q$ repeatedly, how much will $\hat{\mu}_{\text{SN}}$ fluctuate around the true mean $\mu$?

Here’s a (somewhat) intuitive example: Let’s assume we want an estimate of the average height of people in Europe by sending $N$ random people a text message. Imagine you picking 100 random cell phone numbers and blasting off a text saying “Yo, how tall are you?”. These text message could arrive on a Greek island in the sun, in the eternal winter darkness of a Finnish village or in the land of the giants, Holland.

Problem is, you don’t know how tall nor how representative the person reading the text is for the average height of people in Europe. Ideally, you want to representatively sample from all countries in Europe from all height groups: the small people in Holland, the tall and small people in Portugal and the average tall people in Greece.

The variance of your little survey depends on two factors:

1. The intrinsically different heights in Europe that exists regardless of who you pick. Every country has differently tall people.
2. Who you send the text to in Europe from your 100 allowed text messages.

Particularly the second source of variability is tricky in our case. Did you send the text message it by accident to all Dutch men measuring 2.14m? Then your average is like 2.05m. Hardly realistic. Or maybe you send it to all Greek grandmas towering at a mighty 1.55m. Also hardly representative.

In more mathy terms we would say that the variance of $\hat{\mu}_{\text{SN}}$ depends on two factors:

1. The intrinsic variability of $f(X)$ under the target distribution $p$ captured by $\sigma_f^2 = \text{Var}_p[f(X)]$. (the natural distribution of heights in Europe)
2. The mismatch between proposal $q$ and target $p$, reflected in the weight variability. ($q$ being the proposed people you send your height text message to)

When $q$ is a poor approximation to $p$, a few samples will have very large weights while most have tiny weights. For simplicites sake assume that $q$ is a uniform $q=1/100$. The probability of picking the tiny Greek grandmas or the super tall male Dutchies is very small as their heights are at the far ends of the height distribution in Europe. The average height is around 1.75m, so picking a height of 1.55m or 2.14m is very unlikely.

It’s important to note that we’re self-normalizing the weights $\tilde{w}_i$ with all our samples, so we have \(\bar{w}_i = \tilde{w} / \sum_{j=1}^{100} \tilde{w}_j\).

If we ask 98 Greek grandmas and towering Dutchies and only 2 slightly less extreme samples, the contribution of those extreme 98 weights will be negligible. Remember that in order to compute the average weight we’ll compute,

But lo and behold, due to $p$ being very small, the weights of the 98 unrepresentative samples have a negligible contribution to the overall value. In effect, we’re trying to estimate the average European height with two samples. Not ideal when we had assumed that we used 100 samples.

This weight concentration in those two meaningful samples means we’re effectively using far fewer than $N$ samples, inflating the estimator variance. The effective sample size (ESS) quantifies exactly how many “equivalent i.i.d. samples” our weighted samples represent. In our case it’s two, not a hundred.

But how do we quantify this effective sample size?

The Variance of our Estimator

Let’s denote the number of effective samples as $N_{\text{eff}}$.

If we had $N$ i.i.d. samples from $p$, the variance of the sample mean would be:

where $\sigma_f^2 = \text{Var}_p[f(X)]$.

But alas, we don’t have access to samples from $p$ because if we had, I wouldn’t have to write this blog post explaining ESS.

What we do have is our estimator $\hat{\mu}_{\text{SN}}$,

The variance is easy to set up, but actually a bit tricky to compute analytically as we’re dealing with a ratio of random variables

The Delta Method (or “Everybody’s Favourite Tool: Linearization with Taylor Expansions”)

Imagine you have two random variables, say $X$ and $Y$, and you care about the ratio $g(X, Y) = X/Y$. You know the variances and covariance of $X$ and $Y$, but what’s the variance of the ratio? That’s awkward because the ratio is nonlinear — variances don’t just “pass through” nonlinear functions.

The delta method’s trick is simple. All we gonna do is locally linearize the function via a first-order Taylor expansion around the means, then compute the variance of that linear approximation. Since variance behaves nicely under linear transforms, you get a clean closed-form result (albeit an only locally valid result due to the linearization).

Let $g(X, Y) = X / Y$.

Let’s start with the one-dimensional case to build intuition. The Taylor expansion of a smooth function $g(X)$ around a point $\mu_X$ is:

Now let’s extend this to two dimensions. The exact formula is actually bit complicated for the infinite sum but the practically interesting parts are

Only using the first order yields the truncated version in the form of

For the delta method, we keep only the first-order term (linearization):

Taking the variance of both sides (and noting that $\mu_X , \mu_Y$ is a constant so $\mathbb{V}[\mu_X] = \mathbb{V}[\mu_Y] = 0$), we have

where the last term $\mathbb{C}[X, Y]$ is the covariance between $X$ and $Y$ as they might or might not be correlated.

The required partial derivatives are:

Substituting:

Naturally, this approximation degrades when $\mu_Y$ is close to zero since we’re effectively dividing by something close to zero or when the distributions have heavy tails. In ML contexts (e.g., variance of a metric ratio in A/B testing), if our samples are large enough for CLT to kick in, this works extremely well.

For our case, we would have

Linearizing the Ratio in the Variance Estimator

For the linearization to be applicable we need the expected values and the variances of the numerator and denominator in our self-normalizing importance sampler.

First up we need to compute the mean and the variance of numerator and denominator

All of the expectations and variances above are with respect to a random variable $X$ and it’s functions $\tilde{w}(X)$ or $f(X)$. The expectation and variance operator are both linear for a sum of random variables, so we can generically write

For the covariance calculation we’ve exploited the fact that over $N$ samples, we assume that only samples with the same index are correlated, whereas samples with different indices are independent.

Since the terms in $\mathbb{V}[X/Y]$ are all variances, we have to simply add a $1/N$ factor to each term, if we choose to work with the sum instead of the mean. In any case, since the $1/N$ factor can be inserted in both the numerator and denominator, it won’t change the final result anyway.

Putting all of these terms together we get

The equation above is an approximation of the variance of the self-normalized estimator. If we sample the samples from a proposal distribution $q$ that is different from the unnormalized target distribution $\tilde{p}$ resulting in the weights $\tilde{w}$.

But what if we wouldn’t sample from $q$ but from the true target distribution $p$ resulting in the weights $w$? This implies that we can really generate $N_p$ samples from $p$ directly, and thus our samples will be i.i.d. from $p$, not suffering from the approximation error of $q$ to $p$. Let’s calculate the variance of the self-normalized estimator if we were to sample from $p$ directly,

And thus we can see that the variance is actually decreasing with the number of samples $N_p$ from $p$ directly as we would expect from a standard i.i.d. estimator. If $N_p$ is large, we consequently decrease the inherent variance of the estimator by a factor of $N_p$.

But again, we can only sample from $q$, not from $p$. But we can pose a clever question: How many samples $N_q$ from $q$ would we need to get the same variance as $N_p$ samples from $p$? At the very beginning we said that our correlated samples from $q$ carry less information than $N_p$ truly independent samples would. So the ESS metric asks the question how many samples $N_q$ from $q$ are worth $N_p$ perfect, i.i.d, super-duper-excellent samples from $p$?

In order to solidify this, we’ll set the two variance estimators equal to each other, one with $N_p$ samples and one with $N_q$ samples. \(\begin{align*} \mathbb{V}_p\left[\frac{\sum \tilde{w} \ f}{\sum \tilde{w}}\right] &= \mathbb{V}_q\left[\frac{\sum \tilde{w} \ f}{\sum \tilde{w}}\right] \\ \frac{1}{N_p} \mathbb{E}_p\left[ \ \left( f - \mu_f\right)^2\right] &= \frac{1}{N_q} \frac{\mathbb{E}_q\left[ \tilde{w}^2 \ \left( f - \mu_f\right)^2\right]}{\mathbb{E}_q[\tilde{w}]^2} \\ \end{align*}\)

In order to arrive at the ESS formula, we’ll need to do a final approximation by assuming that the weights $\tilde{w}^2$ actually aren’t too correlated with $(f - \mu_f)^2$ which is in practice actually somewhat true,

The samples $N_q$ are then the effective number of samples from $q$ that are worth $N_p$ samples from $p$. The ratio $\frac{\mathbb{E}_q\left[ \tilde{w}^2 \right] }{\mathbb{E}_q[\tilde{w}]^2}$ can intuitively be interpreted as a discount factor. You can quite literally think of the discount factor like a conversion rate between the number of samples from $q$ and the number of samples from $p$ that are worth the same. This is really akin to conversion factors between currencies. Take $N_q$ Zimbabwean dollars and try to convert them to $N_p$ US dollars. Since Zimbabwean dollars aren’t worth anything, you got like 0.000001 Dollar per Zimbabwean dollar. So you need an inordinate amount of Zimbabwean dollars to get a single $N_p=1$ US dollar.

We can also express the effective sampling size in terms of normalized weights $\tilde{w}$ instead of $\tilde{w}^2$. Remember that

This gives us an easy way to calculate the ESS purely from the normalized weights $\tilde{w}_i$.

MCMC

In the importance sampling case, the problem was that we couldn’t sample from $p$ and had to use a proposal $q$, leading to weights that could be wildly uneven. MCMC flips the script entirely: we are sampling from $p$ (at least asymptotically), but our samples are serially correlated because each sample depends on the previous one.

Think of it this way: if importance sampling suffers from “some samples mattering way more than others”, MCMC suffers from “consecutive samples being almost the same thing”. You’re walking through the mountainous landscape we described at the very beginning, but you can only take small steps from where you currently are. If you’re standing on a peak, your next few hundred positions will all be near that same peak — you haven’t explored the valleys yet. Your chain might eventually visit the entire landscape, but locally, your samples are redundant.

So the question is the same one we’ve been asking all along: how many i.i.d. samples is our correlated chain actually worth?

Let’s set up the math. Denote the sample mean of our MCMC chain as

where $X_1, X_2, \ldots, X_N$ are the states of our Markov chain targeting $p$.

We want $\mathbb{V}[\bar{f}_N]$. Let’s expand this directly from the definition:

Now here’s where the correlation makes life interesting. What we’ve done is run a MCMC chain which always takes the last value $X_{t-1}$ as the starting point for the next step $X_t$. So for all intents and purposes, the samples $X_t$ are correlated. The variance of a sum is not just the sum of the variances when the summands are correlated. We have to account for all the cross-terms:

This is a double sum — every pair of samples $(t, s)$ contributes a covariance term. The diagonal terms ($t = s$) give us the familiar sum of variances, while the off-diagonal terms ($t \neq s$) capture the correlations between different samples. For i.i.d. samples, all the off-diagonal terms vanish because $\mathbb{C}[f(X_t), f(X_s)] = 0$ for $t \neq s$, and we’d recover

But for MCMC, they don’t vanish. Consecutive samples are correlated, and that correlation bleeds into the variance.

The saving grace is stationarity: if you run the chain long enough after which the chain is said to have “converged”, the covariance between $f(X_t)$ and $f(X_s)$ depends only on the lag $k = |t - s|$, not on the absolute positions $t$ and $s$. With stationarity, this allows us define the autocovariance function:

Note that $c_0 = \sigma_f^2$ (the marginal variance of $f$ under $p$, the same $\sigma_f^2$ we’ve been carrying around since the beginning) and $c_k = c_{-k}$ by symmetry — the covariance doesn’t care about the direction of time.

In the double sum $\sum_{t=1}^N \sum_{s=1}^N$, we need to count how many pairs $(t, s)$ produce each lag $k = t - s$. This is a straightforward combinatorial exercise:

• Lag $k = 0$: pairs $(1,1), (2,2), \ldots, (N,N)$ → exactly $N$ terms
• Lag $k = 1$: pairs $(2,1), (3,2), \ldots, (N, N-1)$ → exactly $N-1$ terms
• Lag $k = -1$: pairs $(1,2), (2,3), \ldots, (N-1, N)$ → exactly $N-1$ terms
• In general, lag $k$: exactly $N - |k|$ terms, for $k = -(N-1), \ldots, N-1$

This allows us to simplify the equation as for any fixed $k=t-s$, there are $N - |k|$ pairs $(t, s)$ that produce this lag with a fixed value of $c_k$ for the stationary/converged chain. So the double sum collapses to

Using the symmetry $c_k = c_{-k}$, we fold the negative lags into the positive ones:

Plugging this back into our variance expression:

Defining the autocorrelation function $\rho_k = c_k / c_0$ (the normalized version of the autocovariance), we arrive at

For i.i.d. samples, $\rho_k = 0$ for all $k \geq 1$, and the bracket collapses to $1$, recovering the familiar $\sigma_f^2/N$ we know and love.

The Bartlett weights $(1 - k/N)$ are worth pausing on. They arise purely from counting: at lag $k$, there are only $N - k$ pairs of samples separated by exactly $k$ steps. The further apart two samples are in the chain, the fewer such pairs exist. For large $N$, these weights are approximately $1$ for any fixed lag — the edge effects wash out.

Thus for large $N$, the Bartlett weights $(1 - k/N) \to 1$ for any fixed $k$. If the autocorrelations decay fast enough that $\sum_{k=1}^\infty \|\rho_k\| < \infty$ — which holds for any geometrically ergodic chain, i.e. pretty much every well-behaved MCMC sampler you’ll encounter in practice — we can take the limit:

where we’ve defined the integrated autocorrelation time (IAT):

The name “integrated” comes from viewing this as the discrete analogue of integrating the continuous autocorrelation function. It collapses the entire memory structure of the chain into a single number.

What does $\tau_f$ tell us?

• $\tau_f = 1$: no autocorrelation at all, we’re back in i.i.d. land because all $\rho_k = 0$ for $k \geq 1$
• $\tau_f = 10$: each sample carries only $1/10$-th the information of an independent draw
• $\tau_f = 100$: the chain is mixing like molasses — you need 100 steps to get one effective independent sample

One subtlety worth flagging: notice the subscript $f$ on $\tau_f$. Unlike the importance sampling ESS we derived above — which is purely a function of the weights and therefore the same no matter what $f$ you’re estimating — the MCMC IAT is function-dependent. A chain might mix quickly for the mean of $f$ but slowly for its variance, or zip along in one direction of parameter space while crawling in another. This is a fundamental difference between the two flavors of ESS.

We now play exactly the same game as in the importance sampling case. We ask: how many i.i.d. samples $N_{\text{eff}}$ from $p$ would give us the same variance as our $N$ correlated MCMC samples?

The i.i.d. baseline is, as before:

Setting this equal to the MCMC variance:

The $\sigma_f^2$ cancels — though this is a bit misleading, since $\tau_f$ itself depends on $f$, so the ESS is still function-dependent unlike the IS case:

The logic is beautifully parallel to what we did for importance sampling. There, the “discount factor” was the weight variability $\mathbb{E}_q[\tilde{w}^2] / \mathbb{E}_q[\tilde{w}]^2$. Here, the discount factor is the integrated autocorrelation time $\tau_f$. Both tell you how many of your $N$ samples you’re actually using.

Going back to our currency analogy: if importance sampling was converting Zimbabwean dollars to US dollars (with the exchange rate determined by weight concentration), then MCMC is like converting a time series of correlated price observations to independent ones (with the exchange rate determined by how sticky the prices are). Either way, you’re asking: what’s my real purchasing power?

Diffusion Models as Sequential Monte Carlo

We’ve now built up the entire ESS toolkit: for importance sampling, it measures weight concentration; for MCMC, it measures autocorrelation-induced information loss. These two stories converge in one of the most consequential generative modeling frameworks of the last few years — diffusion models — where the denoising process can be interpreted as a sequential Monte Carlo sampler.

But before we get to diffusion models, let’s set up the general SMC framework. The whole point of SMC is to do inference in state-space models (also known as hidden Markov models in the discrete case), where we have:

• A sequence of hidden states $x_0, x_1, \ldots, x_T$ that we can’t observe directly
• A prior over the initial state: $p(x_0)$
• A transition kernel $p(x_t | x_{t-1})$ describing how the hidden state evolves over time
• An emission model (also called the observation model) $p(y_t | x_t)$ describing how the observed data $y_t$ is generated from the hidden state

The joint distribution over the full trajectory and observations factorizes as

The goal is to compute the filtering distribution $p(x_t | y_{1:t})$ — what do we believe about the current hidden state, given everything we’ve observed so far? — or, more ambitiously, the smoothing distribution $p(x_{0:T} | y_{1:T})$ over the entire trajectory.

The emission probability $p(y_t | x_t)$ is the likelihood of the observed data $y_t$ given the hidden state $x_t$. In effect, it works as a critic of the trajectory $x_{0:t}$ that was sampled so far. At each step (but also maybe only at other steps) it evaluates the probability $p(y_t | x_{0:t})$ of the observed data $y_t$ given the trajectory $x_{0:t}$. If you did a lot of steps that moved your samples $x_t$ into some part of space where the data $y_t$ is very unlikely, your weight will be low, and thus that particular sample path $x_{0:t}$ will be highly discouraged.

The trouble is that the corresponding analytical posteriors are intractable in general. The exact recursion involves integrating over all possible previous states at every step, and this integral doesn’t have a closed form except in special cases (linear-Gaussian dynamics give you the Kalman filter; finite state spaces give you the forward-backward algorithm).

Sequential Monte Carlo approximates these intractable distributions using a weighted set of particles — sample trajectories ${x_{0:t}^{(i)}, w_t^{(i)}}_{i=1}^N$ that represent an empirical approximation to the posterior. The algorithm proceeds sequentially:

1. Propagate: Move each particle forward in time by sampling from the transition kernel $x_t^{(i)} \sim p(x_t | x_{t-1}^{(i)})$
2. Weight: Assign each particle an importance weight based on the emission model $\tilde{w}_t^{(i)} = p(y_t | x_t^{(i)})$
3. Normalize: Compute $\bar{w}_t^{(i)} = \tilde{w}_t^{(i)} / \sum_j \tilde{w}_t^{(j)}$
4. Evaluate: Compute the ESS $\hat{N}_{\text{eff}} = 1 / \sum_i (\bar{w}_t^{(i)})^2$
5. Resample: If the ESS drops below a threshold, resample particles according to the weights — duplicate the promising ones, discard the dead weight — and reset all weights to $1/N$

Essentially, you keep sampling trajectories and you keep using the “critic” to evaluate them by checking what the probability is that the trajectory would have produced the data $y_t$ that you observed. This information is encapsulated in the weights $w_t^{(i)}$ that you assign to each trajectory. If the weights become too low, this implies that your trajectories have become more and more unlikely for the “critic” $p(y_t | x_{0:t})$. What do you do? You look at your trajectories and you decide to keep only the ones that are still likely, and you discard the ones that are not. The normalized weights $\bar{w}_t^{(i)}$ can then be used as the probabilities for a categorical distribution that you use to sample new trajectories.

Now — what does any of this have to do with diffusion models?

For diffusion models, we have the generative reverse-time SDE:

where $W_t$ is a reverse-time Wiener process and $\nabla_{x_t} \log p_t(x_t)$ is the score function — the gradient of the log-density at time $t$. This is the quantity a diffusion model learns: a neural network $s_\theta(x_t, t) \approx \nabla_{x_t} \log p_t(x_t)$.

To actually generate samples, we discretize the reverse SDE with an Euler-Maruyama step of size $\Delta t$:

This discretization is our transition kernel — reading it off as a Gaussian:

This is the transition kernel that enters the SMC framework. Unconditional generation is then straightforward: sample $x_T \sim \mathcal{N}(0, I)$, apply $p_\theta(x_{t-\Delta t} | x_t)$ for each step backwards, and out pops a sample $x_0 \approx p_{\text{data}}$. This is just a Markov chain — no weights, no resampling, no ESS to speak of.

The one-step denoiser: Tweedie’s formula

Before we get to SMC, we need one more ingredient: the “critic” $p(y_t | x_t)$. At any noise level $t$, the diffusion model implicitly gives us a one-step denoiser — a direct estimate of what the clean data $x_0$ looks like, given only the current noisy state $x_t$.

This comes from Tweedie’s formula. Given the forward marginal $q(x_t | x_0) = \mathcal{N}(x_t;\, \alpha(t)\, x_0,\, \sigma^2(t)\, I)$, the posterior mean of $x_0$ given $x_t$ is:

Since the diffusion model learns the score $s_\theta(x_t, t) \approx \nabla_{x_t} \log p_t(x_t)$, we get $\hat{x}_0(x_t)$ essentially for free — it’s just a different way of reading off the network output.

Think of $\hat{x}_0(x_t)$ as the model’s best guess of the clean image, peeking through the noise at step $t$. Early in the denoising process (large $t$, lots of noise), this guess is blurry and uncertain. Late in the process (small $t$, little noise), it’s sharp and confident.

We’ll call this one-step denoiser prediction the “observation” that the current noisy state $x_t$ produces — and denote the observation model as $p(y_t | x_t)$, where $y_t$ encodes what we can infer about the data from $x_t$.

Conditional generation: enter the state-space model

Now suppose we don’t just want any sample from $p_{\text{data}}$ — we want a sample that satisfies some condition $y$. Maybe we want an image of a golden retriever, or we want to inpaint a missing region, or we want a molecule with certain properties.

What we want is to sample from the conditional reverse process:

The problem? The likelihood $p(y | x_0)$ only tells us about the final clean sample $x_0$. We don’t get any signal about whether our denoising trajectory is heading in the right direction until the very end. By then, it’s too late — all our computational budget is spent.

This is where the one-step denoiser saves the day. Instead of waiting until $x_0$ to evaluate $p(y | x_0)$, we can get an approximate per-step likelihood using the Tweedie estimate:

We’re peeking at the model’s current best guess of the clean output and asking: “does this look like it satisfies the condition?”

And now we have all the pieces to write down a state-space model — exactly the structure we introduced in the general SMC framework above:

• Hidden states $x_T, x_{T-\Delta t}, \ldots, x_0$ — the noisy latents at each denoising step
• Prior $p(x_T) = \mathcal{N}(0, I)$ — pure noise at the start of the reverse process
• Transition kernel $p_\theta(x_{t-\Delta t} \| x_t)$ — the Euler-Maruyama discretization of the reverse SDE
• Emission model $p(y_t | x_t) \approx p(y | \hat{x}_0(x_t))$ — how consistent is the current state with the desired condition, as estimated through the one-step denoiser

This is exactly the structure we need for a particle filter.

Combining everything we get the following algorithm:

Initialize: Draw $N$ particles from the prior (pure noise):

For each denoising step $t = T, T-\Delta t, \ldots, \Delta t$:

1. Propagate each particle through the transition kernel (one reverse SDE step):

2. Compute the unnormalized importance weights via the observation model:

The weight of each particle reflects how well its one-step denoiser prediction $\hat{x}_0(x_t^{(i)})$ matches the desired condition $y$.

3. Normalize the weights:

Look familiar? These are exactly the self-normalized weights $\bar{w}_i = \tilde{w}_i / \sum_j \tilde{w}_j$ we spent the first half of this post deriving.

4. Compute the ESS using the formula we derived:

5. Resample if necessary: if $\hat{N}_{\text{eff}}$ drops below some threshold (a common choice is $N/2$), resample the particles according to the weights — duplicate the high-weight particles and discard the low-weight ones. After resampling, reset all weights to $1/N$ and continue denoising from the refreshed particle cloud.

Why ESS matters here: the weight degeneracy problem

And here’s where our entire ESS derivation pays off.

Remember the Dutch giants and Greek grandmas? In the diffusion context, the particles are like our text message recipients. Each particle $x_t^{(i)}$ is a different “hypothesis” for what the clean image looks like. Some particles, by chance, will be heading toward images that satisfy the condition $y$ — samples that score highly under the observation model. Others will be heading toward something completely different: samples that bear no resemblance to the desired condition, effectively dead weight in the particle cloud.

The weights $\tilde{w}_t^{(i)} = p(y | \hat{x}_0(x_t^{(i)}))$ quantify this mismatch: particles whose one-step denoiser prediction aligns well with the condition $y$ get large weights, while those drifting toward irrelevant regions of the sample space get negligible weights. If only 2 out of 100 particles are heading in the right direction, the ESS will be approximately 2 — just like our Dutch-and-Greek height survey.

Without resampling, the particle cloud degenerates: a handful of particles carry all the weight, and you’re effectively running the denoising process with just those few trajectories. The estimator variance explodes, and the quality of your conditional samples plummets.

Resampling — triggered when the ESS drops below a threshold — is the remedy. It prunes the dead-weight particles and replicates the promising ones, refreshing the particle cloud. But the price is that resampled particles are no longer independent (they share common ancestors), introducing the autocorrelation that we analyzed in the MCMC section. So we’re back to our old friend: correlated samples carrying less information than independent ones.

This is the fundamental tension of SMC: resampling fights weight degeneracy but creates sample correlation. The ESS, in both its importance-sampling and MCMC flavors, is the diagnostic that navigates this tradeoff.

Deriving the incremental weights

So far we have the transition probabilities $p_\theta(x_{t-\Delta t} | x_t)$ — the reverse SDE that propagates our particles from noise toward data. But transitions alone just tell us how particles move; they say nothing about which particles are worth keeping. For that, we need the emission probabilities $p(y_t | x_t)$.

The idea is simple: we define a function that takes in a particle $x_t$ and returns a scalar value evaluating how well that particle’s trajectory aligns with the desired observation $y$. In the diffusion setting, we use the one-step denoiser to peek ahead: $p(y_t | x_t) \approx p(y | \hat{x}_0(x_t))$. This is our critic — it scores each particle at time $t$ by asking “if I were to denoise this particle all the way to $x_0$ right now, how well would the result match the condition $y$?”. In the literature, this is often called the “potential function” and denoted by the function $G_t(x_t)$.

The whole point of tracking importance weights is to keep score of which samples are good and which are bad, so that we can resample accordingly — duplicate the promising trajectories and prune the hopeless ones. And for that resampling decision to be well-informed, we need the most up-to-date emission probability at every time step $t$, not just at the very end.

We can achieve exactly this with a telescoping product. Define the per-step potential function $G_t(x_t)$:

This is the ratio of the current emission score to the previous one — how much our assessment of the particle has changed in a single denoising step. Remember, we’re going from $T \to 0$, so the previous step is $t + \Delta t$. If the particle is making progress toward satisfying the condition, $G_t > 1$ and the weight increases. If it’s drifting away, $G_t < 1$ and the weight decreases. The incremental unnormalized importance weight at step $t$ is then simply

Now watch what happens when we accumulate these incremental updates. The denominators and numerators of consecutive steps cancel each other:

since $\hat{x}_0(x_0) = x_0$ (fully denoised) and $p(y | \hat{x}_0(x_T))$ is essentially constant — pure noise carries no information about $y$. The telescoping gives us the best of both worlds: at any point in time $t$, the accumulated weight reflects the most up-to-date emission probability $p(y | \hat{x}_0(x_t))$, while the recursive update via $G_t$ means we only ever need to compute a single ratio per step (instead of keeping track of the entire trajectory).

The big picture

Let’s step back and admire the view from the top of this particular mountain. What we’ve shown is that conditional diffusion generation is, at its core, a particle filter:

• The prior $p(x_T) = \mathcal{N}(0, I)$ is where our particles start: pure isotropic noise.
• The transition kernel $p_\theta(x_{t-\Delta t} | x_t)$ — the discretized reverse SDE — propagates particles toward cleaner states.
• The emission model $p(\hat{x}_0 | x_t)$ evaluates each particle’s one-step denoiser prediction against the desired condition.
• The importance weights $\tilde{w}_t$ measure how well each particle’s denoised prediction matches $y$.
• The ESS — computed as $1/\sum_i \bar{w}_i^2$ — tells us how many particles are effectively contributing.
• Resampling prunes bad hypotheses and replicates promising ones when the ESS signals weight degeneracy.

Unconditional generation is the degenerate special case where $p(\hat{x}_0 | x_t) = \text{const}$ — there is no condition to satisfy, all weights are equal, the ESS equals $N$ at every step, and no resampling is ever triggered. You’re just running $N$ independent Markov chains in parallel. The moment you introduce any form of guidance — classifier guidance, classifier-free guidance, inpainting constraints, compositional objectives — you’ve entered SMC territory, and the ESS machinery we derived throughout this post becomes indispensable.

Recently there has been quite some interest in the idea of Feynman-Kac Steering.

For a given diffusion model, we can generate samples by iteratively removing noise, transforming a sample of pure noise into a sample of the target distribution. At each step during the denoising process we can let the model estimate what the fully denoised sample would look like. This can happen either through the score identity or is even more readily available if the model was trained on a denoising loss. These fully denoised samples are then evaluated with a criterion. This criterion is then “pulled back” into the denoising process to gauge whether the sample is going to be sufficiently good if fully denoised. Feynman-Kac steering then implements a resampling step where all the samples are evaluated with the criterion in parallel and the batch is resampled from the noisy samples proportional to each samples fully denoised criterion. This is essentially particle filtering with extra steps with stochastic PDE theory wrapped around it.

So let’s dissect the theory.

Ito Derivative

At the core of diffusion model is a stochastic process $X_t$ that evolves according to a stochastic differential equation (SDE). The SDE in question is typically of the form: \(dX_t = \mu(X_t, t) dt + \sigma(X_t, t) dW_t\) where $\mu$ is the drift term, $\sigma$ is the diffusion term, and $W_t$ is a Wiener process (or Brownian motion).

What happens to a function $f(X_t, t)$ as it evolves over time? The Ito derivative gives us a way to compute this by extending the classical Taylor expansion to the stochastic process realm:

An important factor that we’ll encounter in a second again is that we won’t consider effects of higher order as A) they get divided by ever decreasing factors and B) they are negligible for small $\Delta t^p$ where $p>2$.

We then get the Ito derivative:

For an infinitesimally small time increment $\Delta t$, we assume that the difference becomes a continuous $dt$ and we can plug in our SDE:

The next step to observe is that $dW_t = \epsilon \sqrt{dt}$ and that in the limit for a very small $dt$ (think of $dt=10^{-5}$), any effect of a term with $dt$ raised to any power larger than $1$ will diminish even faster (think of $dt^{1.5}=(10^{-5})^{1.5} = 10^{-7.5}$) and thus become negligible. This allows us to drop any term where $dt^p$ where $p>1$. Isn’t math convenient?

We thus get

Rearranging a bit we have

So the infinitesimal change on the right hand side, properly denoted by $dt$ and $dW_t$, would equate the difference in the function $f(X_{t+\Delta t}, t+ \Delta t)$ and $f(X_t, t)$. Given a step of $dt$ in time, the change in the function $f(X_t, t)$ is given by

Backward Kolmogorov Equation

For a derivation of the Kolmogorov backward equation via the Kramers-Moyal expansion, see Fokker, Planck & Kolmogorov Revisited; it is also the central object used in Doob’s h-transform.

Ito’s lemma gives us a way to compute the infinitesimal change in a function of a stochastic process. Since our process is stochastic, for any realized $x_t$, if we run the stochastic process again, we will get a different $X_T$ where $T= t + \Delta t$. All the stochasticity that we accrue over time $\Delta t$ will make the function $f(X_T, T)$ a random variable.

As a sort of cognitive crutch, you can think of $f(X_T)$ as your Temple Run gold coin counter. You always start from the same starting point $x_t$ and you always do the same 5 moves in the time $\Delta t$. If the coin positions stay fixed and you do the same five moves every time you run, you’ll always get the same score $f(x_T)$ at the end. This is what the deterministic part models. But now the stochastic part keeps on changing the gold coin positions. Even for a fixed score $f(x_t)$ and deterministic moves/dynamics, you final gold coin score $f(x_T)$ will vary. Sometimes you’ll get a lot of gold coins, sometimes you’ll get none. So essentially, your gold coin score $f(X_T)$ now is a random variable.

The next question we have to ask to naturally arrive at the Kolmogorov backward equation is whether there is a function $f_T(x_T)$ that describes the expected value of $f_T(X_T)$ for the current state $x_t$ given that the dynamics are stochastic!

So we’d be interested in the equation

where the expectation how likely a state $X_T$ is given the current state $x_t$ is given by the stochastic process $p(X_T | x_t)$. So $f_T(x_T)$ gives you an estimate of how much payout at the end you’ll get midway through your Temple Run.

Figure: Visual intuition for Feynman-Kac. The stochastic process $X_t$ evolves under random noise, while the expected value of a terminal criterion $f_T(X_T)$ is “pulled back” to earlier times via the backward PDE. Particle trajectories (red) and expectation (blue) illustrate the connection between stochastic paths and deterministic backward evolution.

But the definition $f_T(x_T)$ is inherently forward looking as it relies on the solving the SDE forward in time. Can we also reverse time in a way to compute earlier values of $f(x_t,t)$ starting from the terminal values $f_T(X_T)$?

It turns out that there is a partial differential equation hidden beneath all this stochasticity that allows us to compute the expected value of $f(x_t, t)$ backwards in time.

To show this we’ll start out with Ito’s lemma above but this time integrate it all the way from $t$ to $T$.

(where I should have used the integration variable $s$ or $\tau$ instead of $t$ but I didn’t want to introduce unnecessary notational noise ;-) ). More importantly, this equation tells us that if we start with some intermediate value $f(X_t, t)$ and integrate the infinitesimal changes in the function $f(X_t, t)$ from $t$ to $T$, we will end up with the final value of the function $f_T(X_T)$ which is our terminal payout.

And now comes the important step …

… drumroll 🥁 …

… we take the expectation. 👍

The important thing to note is that the expectation of a Wiener process, however he might be scaled over any time, is zero.

While $X_t$ is a random variable, its stochasticity arises completely from the input of the Wiener process $dW_t$. If we eliminate the influence of the Wiener process with the expectation $\mathbb{E}$ at every step, we will essentially obtain a deterministic ODE. With this intuition, the expectation essentially has no influence on the remaining terms and we can write

The equation above can only hold if the term we’re integrating

An important observation is that the integral above is a definite integral over the arbitrary time interval $[t, T]$. This means that the integrand must be zero for all $t$ in the interval $[t, T]$ as we could choose $t$ arbitrarily close to $T$ and the integral would still have to be zero. We can go backward in time from the terminal time $T$ to $T-dt$ and the integral has to be zero, i.e. $\int_{T-dt}^T \ldots dt = 0$. If that integral is zero, then the integral $\int_{T-2dt}^{T-dt} \ldots dt = 0$ also has to be zero in order for the whole integral $\int_{T-2dt}^T \ldots dt = 0$ to hold. This means that the integrand must be zero for all $t$ in the interval $[t, T]$. This gives us the backward Kolmogorov equation:

with some terminal condition $f_T(X_T)$. This result is quite remarkable as it allows us to compute the expected value of a function of a stochastic process at an earlier time $t$ given the terminal condition $f_T(X_T)$ with a PDE backward in time instead of solving a SDE forward in time.

While it looks deceptively similar to the Ito derivative, the backward Kolmogorov equation is a PDE that describes how the expected value of a function of a stochastic process evolves backward in time.

In practical terms, this means that if we have a function $f_T(X_T)$ at the terminal time $T$ with the stochastic dynamics of $dX_t$, we can compute the expected value of this function at an earlier time $t$ by solving the backward Kolmogorov equation.

For example, we could redefine the function $f_T(X_T)$ as the probability of $p(x_T)$ at the terminal time $T$. We could then solve the Kolmogorov backward equation in the form of a PDE backwards in time to obtain the expected value of the probability at an earlier time $t$, which we would denote as $p(X_T | X_t)$.

This equation would evolve backward the probability density $p(X_T | X_t)$ from the terminal time $T$ to the earlier time $t$. At some earlier point $t$, we could evaluate the how likely the state $X_T$ would be given the earlier state $X_t$.

A Stochastic Sidenote

For some function $f(x_t, t)$ and a terminal value $f_T(x_T)$, we derived the following identity above

But now with our previously gained knowledge we can ascertain that the deterministic integrand is always zero, so we obtain the following equation

which is A) a martingale and B) a random variable due to the Ito integral. We can make this more precise by observing as before that the expectation is zero because $\mathbb{E}[dW_t] = 0$. Also, we can quite easily compute the variance of this random variable by computing

By definition a Wiener process $W_t$ has the property that $\mathbb{E}[dW_t dW_{t’}] = \delta(t-t’) dt$ which says that the product of a Wiener process at two different times is zero unless the two times are the same. This means that we can rewrite the above equation as

because all the “cross terms” where $t \neq t’$ vanish due to the property of the Wiener process. This is also known as Ito Isometry. This gives us the mean and variance of the random variable $f_T(X_T) | X_t$:

Feynman-Kac Equation

Previously, we’ve observed how with a terminal condition at time $T$, we can compute the expected value of a function of a stochastic process at an earlier time $t$ by solving the backward Kolmogorov equation

where the dynamics of a stochastic process $X_t$ are given by the SDE \(dX_t = \mu(X_t, t) dt + \sigma(X_t, t) dW_t\) and the terminal condition is given by $f_T(X_T)$.

In their famous equation, Mark Kac and Richard Feynman asked “Yo, what if we add more terms to the Kolmogorov Backward Equation?” (essentially … please don’t quote me on this).

In order to align with the normal notation of the Feynman-Kac equation, we will be using $u$ instead of $f$. The FK equation then is posited as follows:

where $c(X_t, t)$ and $v(X_t, t)$ are some functions that we can choose.

We can interpret $c(X_t, t)$ as a “potential” term that modifies the expected value of the function $u(X_t, t)$ at time $t$ by interacting with the function $u(X_t, t)$. $v(X_t, t)$ can be interpreted as a “forcing” term that adds an additional contribution to the expected value of the function $u(X_t, t)$ at time $t$.

In order to make the notation a bit easier we’ll shorten the notation to

Solving this equation backwards in time from the terminal condition $u(X_T, T)$ gives us the expected value of the function $u(X_t, t)$ at an earlier time $t$ given the terminal condition $u(X_T, T)$.

So how do we solve this equation? Solving PDEs is fiendishly hard and no easy task.

Usually, there are a couple of ‘ansatz’s that we can use to solve PDE’s. For the FK equation we’ll chose the ansatz of defining a new function $y(x_t, t)$ as

where $u_s$ is the function $u(X_s, s)$ and $c_r$ is the function $c(X_r, r)$.

Differentiating this function with respect to $s$ and the product rule gives us

where $du$ is an Ito derivative but $u_s \ d e^{\int_t^s c_r dr}$ is not.

Both $u(s, X_s)$ and $\exp\left( \int_t^s c(X_r), dr \right)$ depend on the random path $X$, but only $u(s, X_s)$ is a function of the semimartingale $X_s$ at a point—so Itô’s lemma applies. $e^{\int_t^s c_r dr}$ is a functional of the entire path, not just $X_s$, and its exponent $\int_t^s c(X_r), dr$ is absolutely continuous (no Brownian term). You assume the values $X_r$ to be deterministic and the functional $\int c_r dr$ is evaluated on an “already” realized path of $X_r$’s. So the ordinary chain rule suffices: \(d e^{\int_t^s c_r dr} = e^{\int_t^s c_r dr} c_r ds\)

In short: use Itô’s lemma for functions of $X_s$; use the chain rule for pathwise functionals without stochastic terms.

Evaluating the Ito derivative we get

and we can compare that to our original FK equation

We can then proceed to plug our $-c \ u - v$ into our Ito derivative and get

Now we take the $du$ and plug that back into our definition of $dy$ and get

Integrating both sides from $s=t$ to $s=T$ gives us

and evaluating the left hand side gives us

where integrating any quantity a zero amount $\int_t^t \ldots dt$ is always equal to zero as nothing gets added to the integral.

We then ultimately have

Now, we’re going to pull off our favourite, slick stochastic process move and take the expectation of both sides which eliminates the martingale $\int dW_s$ term on the right hand side,

If we now set $c_r=0$ and $v_s =0$, we indeed recover OG Kolmogorov backward equation $u_t = \mathbb{E}[u_T]$.

Feynman-Kac Steering

Diffusion models are able to produce an estimate of a fully denoised sample $\hat{x}_0$ from a diffusive sample $x_t$ by either leveraging the score identity

or by using a denoising loss which directly aims at estimating the fully denoised sample $x_0$ from a noisy sample $x_t$.

In both cases, we can use the denoised sample $\hat{x}_0$ as an estimate of the fully denoised sample $x_0$. This denoised sample is then evaluated with a criterion $c(\hat{x}_0)$ that tells us how good the sample is.

The Feynman-Kac/Kolmogorov Backward equation give us a mathematical tool to “pull back” $c(\hat{x}_0)$ to a more diffused sample $x_t$. In essence it allows us to estimate a diffused version of the criterion $c(x_t)$ at an earlier time step $t$. We use the Feynman-Kac framework to estimate the expected value of the criterion at an earlier time step $t$. Theoretically, we could try to evaluate the full expectation $\mathbb{E}[c(x_t)] = \int c(x_t) p(x_t) dx_t$ but this is computationally expensive and not feasible in practice. We can instead filter or resample the batch, preferentially retaining samples with higher expected criterion values and discarding those predicted to perform poorly. FK steering then allows us to resample the batch of samples $x_t$ proportional to the expected value of the criterion $c(x_t)$. Since we’re dealing with a stochastic process, even samples in the mini batch which were resampled from the same noisy sample $x_t$ will have different expected values of the criterion $c(x_t)$ further down the reverse diffusion process.