How Does a Diffusion Model Work, Part 2: Flow Matching and Diffusion

Diffusion models have become the foundation of modern generative AI, powering state-of-the-art systems for image generation, video synthesis, protein design, and more. In the first episode of this series, we introduced generative modeling as a sampling problem and developed the probabilistic framework needed to think about how models learn complex data distributions.

In this second episode, we move into the core mathematics behind flow matching and diffusion modeling, continuing through the MIT lecture notes An Introduction to Flow Matching and Diffusion Models by Peter Holderrieth and Ezra Erives.

We explore how a simple, easy-to-sample distribution such as Gaussian noise can be continuously transformed into a complicated data distribution, and how this idea leads naturally to both deterministic and stochastic approaches to generative modeling.

Topics covered include:

• The basic idea behind flow matching
• Transforming a noise distribution into a data distribution
• Probability paths and time-dependent distributions
• Vector fields and continuous-time generative processes
• Ordinary differential equations (ODEs) in generative modeling
• Stochastic differential equations (SDEs) and the role of injected noise
• The relationship between diffusion models and continuous flows
• Forward and reverse diffusion processes
• How trained diffusion and flow models are used during inference
• Numerical inference algorithms for solving the resulting ODEs and SDEs
• Why the number of inference steps affects speed and generation quality
• How these mathematical formulations connect to practical modern diffusion systems

The goal of this episode is to build intuition for what a diffusion model is actually doing mathematically: rather than simply “removing noise,” we can view generation as learning a dynamical process that transports probability from a simple starting distribution toward the complex distribution of real data.

This framework provides the foundation for later discussions of score matching, denoising objectives, probability flow ODEs, modern sampling algorithms, and the architectures used in state-of-the-art generative models.

Lecture Notes:
https://diffusion.csail.mit.edu/docs/lecture-notes.pdf

Website:
https://arkitekt-ai.com

Contact:
support@arkitekt-ai.com

Transcript

SPEAKER_00 0:11

Hello, and welcome to another episode of AI Research Today. It's been a bit since I've released an episode, so apologies. Just been extremely busy with our company Architect getting that accelerated, a lot of different things transpiring at the same time. So it's been hard to quarantine time. But appreciate you taking a listen. So again, I'm the host Aaron McClendon. I lead the company Architect AI. And we do a lot of SAS rip and replace for various companies that are spending lots of money on SAS, as well as AI ML build-outs for specific industries, manufacturing, recycling, steel, plastics, that sort of thing. So in the last episode, we left off at uh I had started going through a document that was an open source lecture from MIT, and they were going through diffusion modeling and sort of how that process works. So in the first episode, we covered generally what it means generation as sampling. Um so uh just to summarize a couple of the takeaways, generation, the task of generating samples from a probability distribution, having access to a data set of samples sampled from that distribution during training. Um, and then guided distribution is where we condition this probability distribution on a label y and then sample from that distribution. Um, and we're training on data set pairs that are like Z and Y, where Z would be an image and Y would be an image label. So our goal really behind diffusion model modeling, generative modeling is to find a model that returns samples from this data distribution that are aligned with the distribution after training. So the first chapter was sort of a high-level introduction about some of the terminology that presents itself when discussing diffusion modeling. And uh today I'll I'll move I'll wrap up the diffusion over modeling overview, um, and then we'll get back into the regular research paper discussion. Um so what I'll do in this episode is walk through um chapter two in this open source lecture notes on how diffusion modeling works. So we'll walk through a flow model and then we'll walk through a diffusion model. And for each of them, we'll talk about the mathematical underpinnings as well as the inference procedures behind these models. Now, the training procedures are um a different beast, those are covered in later chapters, but I think so equation heavy that it would be difficult to do it justice on an audio uh format like a podcast. So um we'll cover the diffusion and flow models, what they are, and then uh for further reading, I would just suggest that you consult the text. It's uh not too bad of a read, not too mathy, but mathy enough to be satisfying. Um and uh yeah, worth a look for sure. So, okay, let's get into it a bit. So flow matching and diffusion models. These are two different models, they're both generative models, and they both spawn from the realm of differential equations. So flow modeling, um, flow matching comes from ordinary differential equations, ODE, ODEs, and um diffusion models come from stochastic differential equations, SDEs. So you can imagine the difference, perhaps, if you had to take a guess, one being called ordinary and one being stochastic, that one can be solved exactly analytically, and the other requires some approximations due to non-deterministic processes or distributions. We'll see random variables in the equations. Um, so what we're gonna do is construct these ODEs and SDEs and their simulation. Um, then we'll talk about how they can be used or how a deep learning network can be used to parameterize one of these equations. Um, and that'll lead us to a flow and diffusion model. So we'll start with the simpler one, um, the ODEs. So a solution to an ODE is a trajectory, um, and an ODE is fine defined over a vector field. So a vector field would just mean that if you think about like a 3D uh surface like XYZ normal spatial coordinates, at any point in that uh space, I assign a vector that has a magnitude and direction that indicates um the velocity and space at that point. So the solution to the ODE is a trajectory, which is a function. Let's say um we generalize to d dimensions, but between time steps zero and one, um it gives us a point in this d-dimensional space or three three-dimensional space, sticking with the example. So it's sort of like a time-based equation. So for a certain point in time, we end up being located at a certain point in the vector field, and that vector field at that point has a certain velocity that's pointing a certain direction that's defined by the vector field itself. Um, and so the ODE imposes a condition on the trajectory that we want it to follow along the lines of the vector field. So the vector field is going to be kind of pointing or pushing us in certain directions based on the vectors that we are traveling past. And so the trajectory should um outline a trace that makes sense moving through that vector field. So if we think about our trajectory, which is just where we're located at a certain point in time, um, if we take the derivative of that, uh we should we should equal the um value of the vector field at that point. So the derivative of that um trajectory is going to be a velocity, um, and then we want that velocity to equal the same velocity vector specified by the by the vector field, subject to some initial conditions. So this gives us kind of like the set of uh equations that we need to solve using a differential equation. We're gonna start at time zero at some point x naught, uh, and then we can ask uh the question okay, at a different point in time at time t, uh, where is the trajectory located? And the solution to that is called a flow, which um harkens back to the flow matching or flow modeling uh title of the section. So the solution to this ODE is called a flow, and the flow specifies uh if we start at time t0 at point x0, where do we end up at time t? So it's like a time evolution equation. Um, and the uh time derivative of this flow is equal to the uh velocity field um evol with sorry, let me back up for a second. I'm trying to think how to explain this equation a little bit nicer than it's written. So if you start at a point x naught, the equation, the flow, will evolve that point through time. And so the rate of change of that flow evolution is is going to be equal to the velocity field. So that's the ODE for the flow equation. Um and we have the same initial conditions that at time t0, um, the flow equation uh will be equal to x0, the the starting point. Um so we can um kind of envision how a flow field um evolves, a velocity evolves through a velocity field over time. Um they have a nice little illustration um in the textbook, figure one, I suggest you take a look at. But the flow it's mapping the same dimensions, so rd to rd, and you know the the velocity field as well is is similar. Um but if you can imagine uh if you look at a point um and that point is being pushed a certain direction by this velocity field, uh that point will travel in a certain direction uh as indicated by that velocity field. And the flow is what describes that time evolution. So the flow is really uh what we're looking for as a solution to this um ordinary differential equation. Now, up to this point, uh again, this is not connected to data sampling or anything like this. Really, we're just talking about points and trajectories and that sort of thing. So if you're asking yourself where the data distributions and all that, um, we have not covered that yet. So, but we will get there. So uh one theorem they point out is that if uh u, the velocity field or uh yeah, the velocity field is continuously differentiable with a bounded derivative, then the solution to the ordinary differential equation has a unique solution given by the flow. Um they point out it's a different diffeomorphism, so it's continuously differentiable with a continuously differentiable inverse. Um, but the important part uh to take away from the theorem is that the flow solution exists and is unique. Uh, it does specify that the derivative must be bounded and continuously differentiable, but luckily we're parameterizing this equation by a neural network. So the derivative of this neural network um is going to almost always have a bounded derivative. So we shouldn't need to worry about this. All we need to know is that the solution um does exist. Okay, so they have a simple example. Um I will talk through it, it's all equations, so it'll probably be hard to um visualize too heavily, but I'll I'll do my best here because I think it's a nice little example. So if we have a simple vector field, so let's say that our vector field um is a function uh that has negative theta times x. Negative theta is just some constant, x is the location, and so at a specific location, like let's just say we're moving along a single dimension, uh, so at uh location two, your velocity vector is negative two theta, um, and the uh evolution is constant, negative, and um based on theta. So the solution uh to this system of equations, the ODE, is straightforward to figure out using you know standard Diffy diffy q solutions techniques. Um and the solution turns out to be uh decaying exponential function, so the exponent or e to the negative theta t times the initial starting point x0. So uh e to the negative theta t times x0 at time t0, the exponent just goes to one and we're left at x0. So it checks out. So at time t0, we're at the starting point. The derivative of that exponential is negative theta times the exponent of the negative theta t x0, which is just the flow field itself. And so it is indeed, in fact, the the value of the derivative of the flow function at x0 is equal to the value of the velocity fields um at time t given the flow field at time t that started at x0. Um so this is uh gonna converge to zero um exponentially. Uh they do have a figure in here illustrating the decay over time. It's exactly like what you would imagine it to look like, nothing surprising there. Um but okay, so now um we have this uh general solution technique um for a diff-eq series of diff-eqs to get the flow field um through this sort of strange, seemingly unrelated to AI trajectory sort of setup that we've described here. Um in general, as you can probably imagine, it's not possible to compute the flow field directly if the velocity field is not simple. And the velocity fields uh in in this case will be parameterized by a neural network and could be quite complex. And so uh we instead use numerical methods to solve these ODEs rather than the sort of analytics analytical method, direct solution method, like the example we just walked through. Um the simplest technique um they call out in the book is the Euler method. Uh, and in the Euler method, we start with X, the the position to be X0, and we update it with a very simple rule um where we have some very small step size, we'll call H, and we have our velocity fields um at a specific point. So um the new position x is going to be equal to the initial position, x0, plus a small step size taken in the direction of the velocity field. So it's just an iterative update equation. Um there are different values that we could select for H, uh, and you know, different ones will give you different convergence uh behaviors. Um and they also uh mention another method called I I'm uh at risk of pronouncing this wrong, Hoines method. Um it's a similar kind of thing though. Basically, your position at time t plus h is equal to time position at time t plus some incremental update. Um, it uses like an average uh update rule rather than a direct velocity addition like the Euler method, but uh either way, we're sort of taking very small baby incremental steps in the direction the vector field tells us to go. Okay, so now the interesting part, how do we connect all this to a machine learning generative model? So we're gonna do that by making the vector field a neural network vector field. So if we have a vector field, um, we're gonna parameterize it by some neural network that has parameters theta. So in the book they use the term u for the vector field. Um and now we're gonna start using u sub theta to indicate that this vector field is a value predicted by a neural network. So it's a parameterized function again, mapping the same um Cartesian or product space of RD by um some time 0 to t to a final time or final velocity vector in the same d dimensions. Uh, we just now have a parameterized version of this theta, and that gives us a bunch of different possible choices for neural network architectures to use here. So the goal originally was to generate samples um from this data distribution, and the sample should be random. Um, but the ODE itself is not random. The flow matching model we just talked about, it's fully deterministic. Um, so to inject some randomness, uh, we make the initial condition random. So we choose an initial distribution. Uh in most cases we choose uh this based on a standard Gaussian. Um, but the important thing is that it's easy to sample from, and so the way we proceed is we have this initial distribution, we sample some point from this distribution, and then using the ODE and the flow field, we can evolve this point through time t uh to time t equal to one, um, and then that gives us the solution to the point at the end of the line, which should match the ideal or required distribution p data. Um so we start off with a random distribution, and then we evolve that point through time, and it should land in the distribution that we've optimized the model to train. Now, note here that I'm just describing the inference and setup procedures. The actual way to get the neural network and vector field to push the point along the correct flow journey to arrive at the actual distribution of image data or conditioned image data is a full training uh procedure, which again uh is covered later in the text. But this at least gives you the the problem and at least the overview of the architecture behind a generative model like this. So uh because because of this underlying velocity field being uh or vector field being um parameterized by this neural network, the flow that it's based off that that is based on that vector field is then parameterized. Um and we have a very simple procedure for how to sample from this flow model. So this is the algorithm for sampling from a flow model using the Euler method. So we start off with a neural network vector field ut of data, and some number of steps we're going to traverse in. Uh ut is trained prior, so we're starting with a trained model, but we start off at time zero, we iterate by step size h, which is equal to one over n. We draw a sample from our random distribution uh p init, which is again uh could be a Gaussian. Um and then with that starting sample, x naught, for n steps, we're going to update uh x of time plus step size h is equal to the uh position at time t plus h times the value of the vector field at time t. So then we update our x, we update our t. So now t becomes t plus h, um, and we just do that through our full samples until we end up with a final sample x1. So that is the generative process behind a flow model with the Euler. Method. Not as scary as it sounds when you first uh hear the term flow matching, flow modeling. Uh sounds like it's going to be extremely complex, but it is not necessarily like a lot of things in machine learning. At some level, ultimately, it's just some linear algebra and some probability theory. So definitely doable. And I again encourage you to read the remainder of the text. We'll just have covered the first two chapters and the first two podcast episodes here. Um, but yeah, okay, so now we've covered the flow model, and we're gonna move on to the diffusion modeling. So a stochastic differential equation extends the deterministic trajectories we would get from an ODE with stochastic trajectories. Not shocking. Um, the stochastic process is given by a random variable, uh, which again is just a function that takes a value at some time. Um, and so every time we draw uh to get the position, it's going to generate a random trajectory over time. So um we will simulate the same process, but we could get different outcomes because the dynamics are assigned to this random variable. Um the basis of the diffusion modeling comes from Brownian motion, which is a stochastic process. I always think about um particles in a fluid, like if you have dirt floating around in glass, the dirt particles seem to kind of diffuse everywhere, but there are a lot of um applications of Brownian motion. It's an entire field of study with physics. And in fact, the person who invented the diffusion modeling for AI came from physics uh and like a brownian motion background. Um, I'm 99% sure. Feel free to correct me if if I'm hallucinating that. Uh so um you can think of Brownian motion as a continuous random walk, um, where we have some process that starts out with a value zero. Uh they use the um value w for the random variable, uh it's called a vinier process, I believe that's how it's pronounced. It's spelled like Wiener, but I'll say vinor to sound more sophisticated. Uh and um so that's the random variable that's going to be defining our trajectory. And there are a couple of requirements around this Brownian motion setup. One of them is that the uh variance uh is Gaussian distributed and increasing linearly in time. So if we look at the value of this random variable at time t, compare it with a previous time at time s, uh the distribution of that difference is going to be um distributed with mean zero and a variance proportional to t minus s. Um, and so it's just uh increasing the farther that um larger that constant gets, the further in time we get. Um and the increments are independent. So the samples at time t, t naught, t1, t2, t3, those are all independent um random variables, independent draws from this function itself. So uh we can simulate this Brownian motion with a step size h by starting with uh w equal to zero, and then in a similar sort of Euler update, um w at time t plus h is equal to w of time t plus the square root of h times some error term, um, which is normally distributed as well. Um so this is again um centralist study of processes like the Gaussian distribution. Um it's a has applications from finance, statistical physics. They list these in the texts. I'm I'm just reciting the examples they give, epidemiology. So there's a lot of um examples um in finance, they call out Brownian Moshing use the model price of a complex financial instrument. Um, and the paths are continuous. Um, so there's a lot of interesting mathematical and physics applications as well. Um, so worth uh worth a read if you're a math fan. So the idea of the stochastic differential equation is to extend the deterministic dynamics of the ODE by adding this Brownian motion dynamic. Um and since we're now in a stochastic realm, uh normal rules of calculus don't apply. Um, so we can't take the derivative like we were doing previously. So previously, in the in the ODE flow matching equation approach, um the derivative of the velocity or position uh was equal to the value of the velocity field at that position. So it was all expressed in terms of derivatives. And so what they do is reformulate that in terms of infinitesimal updates, um, and they do that by expanding the derivative. Uh it's a little loose, plays fast and loose with some of the math notation um regarding derivatives and limits and that sort of thing, but it's fine. Physics, physics math. Um, and uh they uh re take they get rid of the derivative and they define it in terms of uh infinitesimal steps, um, where we have uh the position at time t plus h is now equal to the position time t plus h uh the step size multiplied by the value of the velocity vector plus um a uh function that has some uh randomness out of it, but for which uh as the limit of h goes to zero, this function's value takes zero as well. Um so in an infinitesimal sense, at every time step we take a small direction in the direction of ut. Um and now what we're gonna do is add uh some contribution from Brownian motion. So we took the setup for the ordinary differential equation, we rewrote it in terms of these infinitesimal updates rather than derivatives, and now um since the new position is equal to the old position plus velocity field uh or vector field contribution, um, we're gonna add in some noise to make this process uh stochastic. So we add this stochastic term, which is um the difference between uh w at t plus h minus w of t times what they call a diffusion coefficient and some static error term um they call R sub t. Uh so this is our diffusion modeling setup. Uh it is written in terms of like dwt, um which is an informal uh notation for uh small contributions. Um but the value, the point important point now is that the value at time x plus t or sorry, the value at time x t is not fully determined by x naught, um, since the evolution is stochastic. We have this um random randomness contributing to the time evolution of the equation. So, like a nice smooth analytic solution like we had in the example of the uh flow matching setup is no longer applicable here. However, um the U is still continuously differentiable with a bounded derivative and has a unique stochastic process satisfying the equation. Um they point out that you could spend an entire calculus class on some of the stochastic calculus rules and theorems, um, but they kind of take the result at face value in order to get to the point. Um and they note that you can transition uh from an SDE to an ODE simply by giving it a vanishing diffusion constant. Um so um they then talk about simulating a stochastic differential equation. Um so you you would ask how would we estimate this uh equation's value? And they do this using the Euler Maruyama method, um, which is basically the same thing as the method we use to numerically solve the ODE. And we're doing this because we're gonna start again doing the same thing, parameterizing these function components by a trained or learned distribution, which is uh given by a neural network. So the Euler uh Maruyama method here is just our new position is equal to the old position plus uh step constant times the velocity field at the old position plus our new um error term, which has a diffusion coefficient, and our error, which is sampled from a Gaussian uh distribution, some Gaussian noise. Um okay, so our original goal was to convert a distribution p init into a complex distribution p data. So, like for ODEs, simulation of this SDE randomly initialized is how we're gonna start. And we can parameterize the vector field again with a neural network. Um, and so the sampling procedure or inference pass through the diffusion model is given a neural network u theta, some number of steps in, and a diffusion coefficient sigma, that's new here, related to the stochastic process. We're gonna start time t. Uh, we're gonna set a constant step size, we're gonna draw a random sample from our initial distribution, and then for some step number, we're going to draw an error sample from our Gaussian distribution, and then we're gonna update our position um uh in this trajectory. So new position equals old one, plus the constant step size times the value that the neural network predicts for this velocity field at the old time, plus the contribution scaled contribution from this error term. We do that uh in times, and then we get our final sample x1, and that is our diffusion model. So uh throughout the I'll I'll read the little summary they have here at the end. So throughout the document, diffusion model consists of a neural network U with parameters theta that parameterize a vector field and a fixed diffusion coefficient. So to obtain samples from our stochastic model, the procedure is as follows. Initialize a random distribution, a Gaussian, sample from it, simulate the stochastic equation from time zero to one, and then the goal is to make the point at x1, the final point, um, have the same match the distribution p data. And we can reduce the stochastic model to a flow model by setting our diffusion constant equal to zero. All right. So that brings us to the end of chapter two. Um, chapter three goes into flow matching, um, which is how to train a flow model. Um it's uh an algorithm that they present. Um, however, there's a lot of integrals and things like that that would make it um difficult to talk through in a podcast form. So I will leave it to you. Um, however, it is a good read, not too difficult. Chapter four, they go through score matching, um, which is uh algorithm for training diffusion models. Um and from there uh they get into um things like conditioning on a prompt. So here we're simply talking about creating an image, right? Like there's no notion of give me a human text, like give me a picture of a dog eating an ice cream cone. So that would be like conditioned prompt inference. Um, and so that's what they get into in chapter five. Um, and uh they also have some language modeling with diffusion chapter, the entire book uh text is about 60 pages. Um, they do have a reminder of probability theory theory uh in the back and a few other things, um, VAE discussions, that sort of thing. So um again, really worth the read um if you have some time. I suggest it. I will put a link to the book in the episode description. Um, and yeah, other than that, I think next time I do an episode, we'll get back into the pure research papers. Um, and uh yeah, I I've I have a few that are on the docket um I would like to discuss. So uh yeah, outside of that, I appreciate everyone who listened and hope you have a great rest of your week.