Flow and Diffusion Models: The Mathematical Perspective

Modern generative AI relies on sophisticated mathematical frameworks to transform random noise into complex, structured data like images, text, and music. Two of the most powerful approaches in this domain are flow models and diffusion models. In this post, we'll explore the elegant mathematics that makes these models work, starting with flow models.

Flow models represent a fascinating approach to generative modeling where we define vector fields that transform simple distributions (like Gaussian noise) into complex data distributions (like images or music). Unlike their more famous cousins, diffusion models, flow models are fundamentally deterministic - given the same starting point, they'll always produce exactly the same result.

Someone gave me this analogy for flow models, which isn't perfect but is a good starting point: Imagine pouring a bucket of water (representing a simple probability distribution) onto a carefully sculpted landscape. The water flows along paths determined by the terrain (that we will later define as a vector field), eventually forming streams that lead towards our target probability distribution. Flow models work in a similar way - they define "probability currents" that guide simple distributions into complex ones through continuous transformations.

Just like how the slope of terrain determines which way water flows at any point, the vector field in a flow model tells each point in our probability distribution exactly which way to move.

Now let's translate this intuition into precise mathematics. At first, you might think of flow models as simple updates where $u_t(x_t)=x_{t+1}$, meaning the vector field directly gives the next position. However, this isn't quite right because flow models operate in continuous time, not discrete steps.

We can define a vector field as a function:

$u:{\mathbb{R}}^d\times [0,1]\to {\mathbb{R}}^d,(x,t)\mapsto u_t(x)$(1)copytex

Let's break down what this means:

• The vector field $u$ takes two inputs: a position $x$ in $d$-dimensional space and a time $t$
• It outputs a vector that tells us which direction to move and how fast at that specific point and time
• The time parameter $t$ ranges from 0 to 1, representing the complete transformation from start to finish

The notation $u_t(x)$ emphasizes that our vector field can change over time - the directions might be different at the beginning (t=0) compared to the end (t=1) of our transformation.

A trajectory is the path that a single point follows through our vector field over time. Mathematically, we define it as:

$X:[0,1]\to {\mathbb{R}}^d,t\mapsto X_t$(2)copytex

It's simply saying:

• $X$ is a function that maps each time $t$ between 0 and 1 (where 0 is the start of our full noise distribution and 1 is the end of our fully transformed distribution such as an image) to a position $X_t$ in our space
• It tells us where a particular point is at any time during the transformation
• Every point follows its own unique trajectory through the space

Note that $X_0=x_0$ specifies the initial position of our point sampled from the noise distribution, which serves as the starting condition.

For example, in a 2D vector field, we can think of our starting point as a single point in 2D space, and our trajectory as the path that point follows through time. A point starting at (1,3) at t=0 might move to (2,4) at t=0.5, then to (3,5) at t=1. However, a different point starting at (1,4) at t=0 would follow a completely different trajectory because it has a different initial condition. This is why specifying $X_0=x_0$ is crucial - it determines which unique trajectory we follow through the time-varying vector field.

The relationship between our vector field and trajectories is formalized through an ordinary differential equation:

$\frac{d}{dt}{X}_t=u_t(X_t)$(3)copytex

with the initial condition $X_0=x_0$.

This equation is the heart of flow models. Breaking it down:

• $\frac{d}{dt}{X}_t$ is the rate of change of position with respect to time (the velocity)
• $u_t(X_t)$ is the vector field evaluated at the current position and time
• The equation says "the velocity (basically the rate of change of position with respect to time) of a point equals the vector field at that point" - in other words, the vector field directly controls how each point moves

While trajectories describe the path of individual points, the flow function captures the entire transformation all at once:

$\psi :{\mathbb{R}}^d\times [0,1]\to {\mathbb{R}}^d,(x,t)\mapsto {\psi }_t(x)$(4)copytex

The flow function ${\psi }_t(x)$ tells us where an initial point $x$ will be at time $t$. It satisfies two key properties:

1. Identity at time zero: ${\psi }_0(x_0)=x_0$ (at the start, points are at their initial positions)
2. Governed by the vector field: $\frac{d}{dt}{\psi }_t(x_0)=u_t({\psi }_t(x_0))$ (the flow evolves according to our ODE)

The flow function is powerful because it encapsulates the complete transformation from our simple distribution to our complex target distribution. If we know ${\psi }_1(x)$ for all $x$, we know exactly how our initial distribution transforms into our final distribution.

Flow models possess several crucial mathematical properties that make them uniquely powerful for generative modeling:

Given the same starting point $x_0$, the trajectory will always follow the exact same path through the vector field. This comes from a fundamental property of ODEs: they have unique solutions given an initial condition.

This determinism contrasts with diffusion models, which involve stochastic noise and don't guarantee the same output given the same input.

Flow models use invertible transformations, meaning we can trace the path backward from the target distribution to the source distribution. This bidirectional property is essential for both training and sampling.

Mathematically, this means we can define an inverse flow ${\psi }_t^{-1}$ that takes us from time $t$ back to time 0.

The transformation occurs continuously over the time interval $[0,1]$, providing smooth interpolation between the source and target distributions.

This continuity allows us to "watch" the transformation process at any intermediate time step, which can be helpful for understanding and debugging these models.

While the mathematical theory of flow models operates in continuous time, practical implementation requires discrete numerical methods.

The key insight comes from the definition of a derivative. Recall that:

$\frac{d}{dt}{X}_t=\underset{h\to 0}{lim}\frac{X_{t+h}-X_t}{h}$(5)copytex

This is the formal definition of a derivative - the limit of the difference quotient as the time step $h$ approaches zero. For small but non-zero values of $h$, we can approximate this as:

$\frac{X_{t+h}-X_t}{h}\approx u_t(X_t)+R_t(h)$(6)copytex

Here, $R_t(h)$ is a remainder term that captures the approximation error when $h$ is not exactly zero. This term is crucial for understanding the accuracy of our numerical methods, but is often overlooked in simplified presentations.

Rearranging the approximation by multiplying both sides by $h$, we obtain:

$X_{t+h}=X_t+h{u}_t(X_t)+h{R}_t(h)$(7)copytex

This equation provides the discrete equivalent of our continuous ODE. It tells us how to compute the next position $X_{t+h}$ given the current position $X_t$ and the vector field $u_t$.

If we ignore the remainder term $h{R}_t(h)$, we obtain the Euler method for numerically solving ODEs:

$X_{t+h}\approx X_t+h{u}_t(X_t)$(8)copytex

This approximation method works by:

1. Taking the current position $X_t$
2. Evaluating the vector field to get the direction $u_t(X_t)$
3. Scaling this direction by the step size $h$
4. Adding it to the current position to estimate the next position

Despite its simplicity, the Euler method works remarkably well for many applications, especially when the step size is sufficiently small.

If you've ever heard of Taylor series: The numerical approximation methods used in flow models have a deep connection to Taylor series expansion. The Taylor series provides a way to approximate functions using polynomial terms:

$f(x+h)=f(x)+h\cdot f^{\prime }(x)+\frac{h^2}{2!}\cdot f^{″}(x)+\frac{h^3}{3!}\cdot f^{‴}(x)+\dots$(9)copytex

We basically calculate the rate of change of a given function at a specific point and use it to approximate the function at a nearby point. If the function was infinitely differentiable, we could use the Taylor series to exactly represent the function by infinitely many polynomial terms.

In our discrete formulation:

$X_{t+h}=X_t+h\cdot u_t(X_t)+h\cdot R_t(h)$(10)copytex

This corresponds exactly to a first-order Taylor approximation where:

• $X_t$ represents the current value (analogous to $f(x)$)
• $h\cdot u_t(X_t)$ represents the first derivative term (analogous to $h\cdot f^{\prime }(x)$)
• $h\cdot R_t(h)$ contains all the higher-order terms we're neglecting

This connection explains why numerical methods for solving ODEs work: they're essentially using truncated Taylor series to approximate the continuous evolution of the system. There are actually more sophisticated methods that use more terms of the Taylor series to get a more accurate approximation, but the Euler method is the simplest one and works surprisingly well.

Let's connect all these mathematical pieces to understand how flow models work in practice:

1. Initialization: We start by sampling points from a simple distribution (typically Gaussian noise) at time $t=0$.

2. Vector Field Definition: We define (or more commonly, learn) a vector field $u_t(x)$ that will guide our points from the simple distribution to the complex distribution. This vector field will basically be our model that we train.

3. Numerical Integration: We use numerical methods like Euler's method to approximate the ODE and push our samples forward through time:

   $X_{t+h}\approx X_t+h\cdot u_t(X_t)$

4. Final Distribution: By time $t=1$, our initially simple distribution has been transformed into the target data distribution through this continuous, deterministic flow.

For training, we typically use the inverse direction: starting with real data samples, we learn a vector field that pushes them back to a simple distribution (like a Gaussian). Once trained, we can sample from the simple distribution and push forward to generate new data samples. But this will be a topic for another post.

---

Now that we understand the deterministic approach of flow models, let's explore their more chaotic cousins: diffusion models. While flow models give us the same output every time for the same input, diffusion models embrace randomness to create more flexible and diverse generative processes.

The key insight behind diffusion models is simple yet powerful: real-world processes aren't perfectly predictable. By adding controlled randomness to our transformations, we can model more complex and varied data distributions.

Remember our flow model ODE:

$\frac{d}{dt}{X}_t=u_t(X_t)$(11)copytex

Diffusion models extend this by adding a stochastic component, transforming our ordinary differential equation (ODE) into a stochastic differential equation (SDE):

$d{X}_t=u_t(X_t)dt+{\sigma }_{t}d{W}_t$(12)copytex

This single equation captures the essence of diffusion models. Let's break it down:

• $u_t(X_t)dt$: The deterministic drift - the predictable part (same as flow models)
• ${\sigma }_{t}d{W}_t$: The stochastic diffusion - the random part that makes diffusion models special

The $dt$ term represents an infinitesimal time step, while $d{W}_t$ represents an increment of something called Brownian motion, which we'll explain shortly. ${\sigma }_t$ is a scalar that controls the amount of noise added at each time step - note that if this is 0, we get back our original ODE of theflow model.

The first term $u_t(X_t)dt$ works exactly like in flow models. It provides a consistent direction based on the current position and time. If we removed the stochastic part entirely (set ${\sigma }_t=0$), we'd get back our original flow model.

The second term ${\sigma }_{t}d{W}_t$ introduces controlled randomness:

• ${\sigma }_t$: The diffusion coefficient - controls HOW MUCH randomness to add at time t
• $d{W}_t$: A Brownian motion increment - provides the random direction

Think of it like walking: the deterministic part tells you the general direction to head (toward your destination), while the stochastic part adds random stumbles and course corrections that make your path more varied.

Brownian motion is the mathematical foundation for the randomness in diffusion models. It's named after the botanist Robert Brown, who observed the random movement of particles suspended in fluid.

For any two time points $s<t$, Brownian motion satisfies:

$W_t-W_s\sim \mathcal{N}(0,(t-s)I_d)$(13)copytex

This means the change in Brownian motion between times s and t follows a normal distribution with:

• Mean of 0 (no bias in any direction)
• Variance proportional to the time difference $(t-s)$
• Note how each path is different

The crucial insight is that variance scales with time:

• Small time steps → small random movements
• Large time steps → large random movements
• The standard deviation scales as $\sqrt{t-s}$, not $(t-s)$

This scaling property is what makes Brownian motion perfect for diffusion models - it ensures that our random process behaves consistently regardless of how we discretize time.

Just like flow models need numerical methods to implement continuous ODEs, diffusion models need the Euler-Maruyama method to handle SDEs.

For a small time step $h$, we discretize our SDE as:

$X_{t+h}=X_t+h{u}_t(X_t)+{\sigma }_t(W_{t+h}-W_t)$(14)copytex

The Brownian increment $W_{t+h}-W_t$ has the property:

$W_{t+h}-W_t\sim \mathcal{N}(0,h{I}_d)$(15)copytex

In practice, we implement this by sampling from a standard normal distribution and scaling appropriately:

$X_{t+h}=X_t+h{u}_t(X_t)+{\sigma }_t\sqrt{h}ϵ$(16)copytex

where $ϵ\sim \mathcal{N}(0,I_d)$ is sampled fresh at each step.

You might wonder: why $\sqrt{h}$ and not just $h$?

This comes from a fundamental property of variance scaling. When you multiply a random variable by a constant $c$, its variance gets multiplied by $c^2$.

Since we want our Brownian increment to have variance $h$, and we're starting with a standard normal (variance 1), we need to multiply by $\sqrt{h}$ because $(\sqrt{h}{)}^2=h$.

Unlike flow models, running a diffusion model twice with the same input will produce different results due to the random components. This makes them excellent for generating diverse outputs.

The diffusion coefficient ${\sigma }_t$ allows fine-grained control over randomness:

• ${\sigma }_t=0$: Reduces to a flow model approach (pure determinism)
• Large ${\sigma }_t$: High randomness and diversity
• Small ${\sigma }_t$: Low randomness, more predictable

The stochastic component allows diffusion models to "explore" the probability space more thoroughly than deterministic flow models. This exploration capability often leads to better coverage of complex data distributions.

Here's the complete algorithm for sampling from a diffusion model:

Require: Vector field $u_t$, number of steps $n$, diffusion coefficient ${\sigma }_t$

1. Set $t=0$
2. Set step size $h=\frac{1}{n}$
3. Set $X_0=x_0$ (initial condition)
4. For $i=1,\dots ,n-1$ do:
   • Draw $ϵ\sim \mathcal{N}(0,I_d)$
   • $X_{t+h}=X_t+h{u}_t(X_t)+{\sigma }_t\sqrt{h}ϵ$
   • Update $t\leftarrow t+h$
5. Return $X_0,X_h,X_{2h},\dots ,X_1$

The key insight is that we're taking the deterministic Euler step from flow models and adding a carefully scaled random component at each iteration.

Diffusion models have become incredibly successful in generative AI because they strike a perfect balance between structure and randomness. The deterministic component ensures the generated samples follow the learned data distribution, while the stochastic component allows for exploration and diversity.

This mathematical framework explains why diffusion models can generate such diverse and high-quality samples - they're not just following a single deterministic path (trajectory) like flow models, but exploring a whole family of possible paths guided by both learned structure and controlled randomness.

So which approach reigns supreme? The answer isn't straightforward. Diffusion models currently dominate the commercial landscape, powering everything from Stable Diffusion to DALL-E, and have proven their reliability in real-world applications. However, recent advances in flow matching show that flow models can achieve superior performance in controlled studies, offering faster training and more efficient sampling paths. The most exciting development is that cutting-edge models like Stable Diffusion 3 and OpenAI's Sora are combining both approaches, using diffusion transformers with flow matching techniques. Rather than one approach replacing the other, the future of generative AI lies in hybrid methods that make use of the deterministic approach of flow models alongside the proven diversity and robustness of diffusion models.