The physics behind diffusion models

Key Takeaways

In 2015, researchers made an unexpected connection between machine learning and non-equilibrium thermodynamics, the branch of physics that studies how ink moves in water or how smoke disperses in air. Somehow these laws of physics led to fake Will Smith eating fake spaghetti. First in the uncanny valley, but then very convincingly. The bridge between the physics of motion and generative AI are diffusion models. Now, in machine learning, there are a lot of metaphors and a lot of questionable references to adjacent fields like neuroscience. But diffusion models are different. They truly are based on the same mathematical framework as physical diffusion. Going into this, my question was, why do we need to bring in such a foreign concept? Is it a solution in search of a problem? Well, spoiler alert, it's not. But before we get started, let me introduce you to the main character in this video. Julie. Julie, touch. Good girl. You'll see her a bunch. All right, let's get started. Think of data as a landscape shaped by probability. Hills represent structured images that are meaningful to the human eye. These are high probability regions. Valleys represent pure noise. These are regions with almost zero probability. Of course, in practice, this landscape has more than just two data axes, x0 and x1. In raw pixel space, each pixel channel would have its own axis. And in latent embedding space, it's common to have over 16,000 dimensions, but good luck drawing that. Instead, we'll visualize two data dimensions plus a third one for the probability. If we had a model of this probability landscape P of X, generating data would mean starting from a launch point and navigating the terrain in search of peaks. But of course, we don't know the shape of the landscape. That's the crux of the problem. Initially, all we have is a training set with real images, a few sticks in the ground, if you will. One way to formulate our learning task would be to infer the landscape that I just hid to build a complete map of the world covered by the training data. But that's hard and also doesn't generalize to uncharted territories outside of the training set. Just like knowing every street in New York City won't help you navigate San Francisco. So what we truly want is a compass, a model that has general navigational abilities. When placed in a specific position, it should make a local decision, taking one step in the direction of higher probability compared to where it's currently standing. While following these local nudges, the larger structure reveals itself gradually through this exploration. This gradual reveal of the landscape is basically a time variant probability distribution. Instead of working with a fixed P of X, we model P index T of X. Formally, this is a conditional probability. The probability of data X at a fixed time T. This notion is rare in traditional machine learning. This is why we turn to physics. In physics, time dependent fields are everywhere. When dropping ink in water, concentration changes over time and space. When heat travels through the environment, temperature evolves in time and space. What connects these real world examples is diffusion, the tendency of particles to move from high to low concentration, which changes their positional distribution over time. We want the same thing for a probability distribution, a P that depends both on X and T. Here's how diffusion fits in the lifetime of a model. We'll consider three stages. Training data generation, model training, and model inference, also known as sampling for generative tasks. Diffusion comes in before and after model training. The forward direction from structure to chaos produces training data and the reverse direction from chaos back to structure is what enables us to find new data. Let's start with stage one training data generation. If we place imaginary training particles in the vicinity of real images and then let them diffuse, we can model the landscape evolution based on the spread of the particles. At time zero, when particles are heavily concentrated around real data, the landscape reflects the target distribution. That's the distribution of the real data. As time passes and particles spread, so do the hills. Now, why would we want spreading hills? Well, that's a way of working around the sparity of the landscape. Since real images cover a tiny portion of the pixel space, there's a lot of flat land. When the ball is in the middle of nowhere, its local compass will be useless. There's no gradient to climb. Spreading the hills when the ball is far away extends the existing structure throughout the terrain so that the ball can make more informed steps. And as the ball progresses, we want to sharpen back the hills so that it can find the real data points. So hopefully you now see why diffusing particles starting from real images outward is a way of creating supervised training data that will model the landscape in a timed dependent way. This process is called forward diffusion. A forward diffusion path takes a training particle from real data to noise. Any isolated point in this path together with the direction that we came from can become a training instance. In fact, we'll arbitrarily choose a single point from this path to train on. We really don't want the model to memorize this specific path or this specific part of the landscape. We only want general directional abilities. Okay. So, how can we simulate this forward motion? How do particles actually move in physical diffusion systems? In 1827, botonist Robert Brown noticed that Poland particles suspended in water moved chaotically. Initially, he thought that Poland was alive, but ruled that out after seeing the same behavior in dust. It turned out the jittery motion came from countless collisions with surrounding water molecules, a phenomenon now called Brownian motion. In 1905, Einstein gave the first quantitative theory. He showed that while each particle's path is unpredictable, a pattern starts to emerge when looking at the entire cloud of particles, it tends to spread out as time goes by. Specifically, the mean square displacement of particles grows linearly with time. This smooth deterministic spreading at macro scale is what we know today as diffusion. In 1923, Norbert Viner gave a precise mathematical model for Brownian motion. It's now called the Vener process. He described the motion of a single particle as a stochastic process, meaning an entire family of random variables. If you freeze time at some instant t, the possible positions of the particle are encoded by the variable wt, which is a gausian with mean zero and variance t. This is a powerful formulation because it connects the physics of motion with probability theory. In the next few sections, we'll go through some equations, but I promise I'm not including them just to look smart on the internet. I'm intentionally cherrypicking the math that actually helps build an intuition. But don't let it scare you. I'm not a mathematician nor a physicist. So if I can follow, you can follow, too. And if you're a coward and click away, I'll see it in the analytics and I'll come and find you. In this section, we'll deduce the most general mathematical expression for forward diffusion, which various papers instantiate in various ways. We'll start from the most basic formula you learned in middle school. Distance X equals speed V multiplied by time T. This describes the motion of a single particle with constant velocity. To express displacement between two positions X1 and X2, we can rewrite the equation in terms of delta X and delta T. When displacements are infinite decimally small, this turns into a differential equation in terms of differentials dx and dt. We can now generalize the constant speed v into a velocity function f that can be position and time dependent. Here I'm toying with the sinosoidal that will slow down the ball right in the middle. Next we'll add brownian motion using the vener process discussed earlier. Since DW is a stochastic term, this is our first stochastic differential equation or SDE. Make sure to remember this abbreviation because we'll use it throughout the rest of the video. Just to be pedantic, X and W are both functions of T. In the physics literature, the T index is often emitted for brevity. There's one final generalization to make, namely control the amount of jitter with a function g that can also be position and time variant. Here I'm choosing another sinosoidal that will make the ball jitter more right in the middle. So this final expression is known as the forward sde. It describes the motion of a single particle forward in time. The first term which is deterministic is known as the drift term and the second is known as the diffusion term. Various papers make various choices for functions f and g. One of the most influential papers in this space is denoising diffusion probabilistic models known as ddpm. This is how DDPM chooses the drift and diffusion terms. It's not very intuitive at first glance. Let's build it step by step. The simplest choice would be pure Brownian motion applied independently to each pixel. To the naked eye, this looks reasonable. Julie seems to be completely erased. But was she really? If I repeat this noising process with 20 different seeds and average the final frames, she starts to reemerge. That's because simple brownian motion preserves the statistical mean. In two dimensions, this effect is obvious. Particles that start diffusing from a fixed point will spread evenly around it. Of course, their variance grows with time, but the cloud will forever be centered around the initial position. Now, why is that a problem? It's because it introduces a discrepancy between training and inference. During inference, we need a home base from where to start the sampling journey. We often plan the flag at the origin and then pick a random starting point around it. This is equivalent to setting the prior distribution to a gausian with mean zero and variance one in all dimensions. From here, we'll gradually find our way towards structure guided by the learned model. The training data needs to mirror these paths from origin outwards. But currently our training paths only cover local neighborhoods. This is why we need a drift term minus x that actively pulls the clouds away from their original position and towards the origin. And finally, there's a time dependent beta term that shows up both in the drift and diffusion terms. It's called a noise schedule, and it controls how fast noise is added at each moment in time. Let's look at a linear noise schedule where beta increases with time. This leads to paths diverging towards the origin and covering more ground than before. During training, the schedule can also control where the particles linger more. Too close to the real data, we won't generalize well. Too far from the real data, there's little structure to learn. The sweet spot is somewhere in the middle. The noise schedule is critical because as we saw earlier, it also controls the shape of the induced landscape and therefore how easily the ball can find its way to real data during sampling. That's why it's the subject of intense research. With a recipe for generating forward diffusion paths, we can now train a model to follow them in reverse and thus learn how to navigate the landscape. Remember, the train model acts as a local compass. More formally, for a fixed particle position or input image, it outputs the local gradient of the probability landscape which points in the direction of real data. We take the logarithm because it's mathematically convenient but also because it gives a scale invariant compass. It always points towards regions of higher probability but without being distorted by how small the probability itself is. In this video, I won't cover the more traditional machine learning aspects of this model, including the architecture or the training objective. Instead, we'll just treat the model as a black box that operates as a GPS, and we'll focus on the diffusion process instead. So, given such a model, how can we leverage it to build an entire reverse path from noise back to structure? At first, gradient ascent might come to mind. simply follow the gradient one step at a time. The problem is a gradient ascent is an optimization algorithm. It seeks to find the one best solution, the mode of the distribution. But image generation is fuzzier than that. There's no one single best picture of a dog. Sometimes we might want to climb this hill, sometimes the other. In other words, we need a sampling algorithm, one that chooses hills randomly proportional to their probability. This is where diffusion comes in again. In 1982, Australian electrical engineer Brian Anderson published a very unintuitive result. Diffusion is reversible. Well, at least on paper. Just like there's a forward stochastic differential equation that helps simulate the motion from data to noise, Anderson derived a reverse SDE for the opposite motion. Now this doesn't mean that physical diffusion is reversible in practice. We don't have enough fine grain control over the moving particles but in the virtual world we do. The reverse SDE might seem scary at first but the structure should look familiar. Just like the forward SDE there's a deterministic drift term and a stochastic diffusion term. Unsurprisingly, the reverse SDE depends on the same functions f and g. But what's most interesting is that the drift term depends on the gradient of the log probability. That's exactly what our model predicts. So what we have here is basically a more principled way of following the compass of the model compared to naive gradient ascent. The reverse SDE completes the theoretical framework for generating images and other types of data. It's a complete mathematical recipe ready for implementation. So let's talk about how we can actually implement this in practice. When generating training data, we start from a real image X0. Remember, we don't want the entire path. We'll pick an arbitrary time, call it tilda, between zero and capital T. The only particle position or image we need to actualize is xt tilda. For any forward SDE, there's a closed form solution, a formula that expresses xt tilda in terms of x0 and can be easily translated to code. However, at inference time, we need to generate an entire path xt from noise back to data. t is not a single point in time anymore. It ranges from zero to capital T. Of course, in practice, we'll have to discretise this path into a series of steps. This is where samplers come in. A sampler is just a numerical scheme for carrying out these discrete steps. Open source tools like automatic 111 give you a long list of samplers. Each is just a different numerical solver. The fact that the reverse equation is stochastic makes discretization inefficient since noise is injected at every infinite decimal step. The motion is jittery and big steps give poor approximations. If only we could drop the stochastic term. Then the reverse path would be smooth and we could approximate it with far fewer steps. It turns out we actually can without losing mathematical rigor. So our landscape is a timevariant probability distribution PT of X. The reverse SDE gives us one recipe for producing paths XD that follow this distribution, but nobody said it's the only recipe. In 2021, Stanford researcher Yang Song and colleagues showed that there's an equivalent deterministic recipe. This is what the new recipe looks like. Compared to the stochastic equation, the major difference is that it's missing the diffusion term, making it an ordinary differential equation or OD. Solving this OD produces paths whose probabilities at each time t match those from the SDE. In other words, both recipes transport probability mass in the exact same way. So in our end to end recipe, we can just replace the reverse SDE with an OD. Instead of a thousand sampling steps, we might get away with just a few dozen. That's a 20x speed up, 20 times fewer passes through the model. Modern diffusion systems often combine both perspectives. They can use SDEEs when diversity is desired and ODEs when speed or controllability matters. In practice, models like stable diffusion 3 lean heavily on OD based samplers. So this was the big picture. Physical diffusion provides a mathematical framework for time variant probabilities. There are many aspects that I glossed over including model training architectures and conditioning on text prompts. But I really wanted to decouple those and get to the core of the connection with physics without getting distracted by implementation details. Now, diffusion models are still a nent field. That's why in the literature, you'll find various perspectives that are different on the surface, but are slowly starting to converge. If you want to follow along, I mostly covered Yang Song's work from Stanford, but I'll include my entire reading list on Patreon for free. The next big question is how does language fit into this formulation? The equations of motion assume continuous space, but language is inherently discreet. It's made of tokens. But that's a different story for another time. For now, thanks for putting up with the math, and I'll see you next time.