Match flows, not scores

While a plethora of high-quality text-to-image diffusion models (Ramesh et al. 2022; Rombach et al. 2022; Podell et al. 2024) emerged in the last few years, credit mostly goes to the tremendous engineering efforts put into them. The fundamental theory behind them however, remained largly untouched – well, untill recently. The traditional “markov-chain” (Ho, Jain, and Abbeel 2020) or “SDE” (Song et al. 2021) perspective is now being replaced by an arguably simpler and flexible alternative. Under the new approach, we ditch the usual notion of a stochastic mapping as the model and adopt a deterministic mapping instead. Turns out that such model, although existed, never took off due the absense of a scalable learning algorithm. In this article, we provide a relatively easy and visually guided tour of “Flow Matching”, followed by ideas like “path straightening” & “reflow”. It is worth mentioning that this very idea powered the most recent version of Stable Diffusion (i.e. SD3 by Esser et al. (2024)).

Some prerequisites

Note that this article relies on the fact that you are familiar with the core ideas of Diffusion Models. So if you haven’t already, please visit my previous blogs (Blog #1, Blog #2 & Blog #3) on Diffusion/Score models.

Every scalable generative model follows a “sampling first” philosophy, i.e. they are defined in terms of a (learnable) mapping (say $F_{\theta }$), which maps (known) pure noise $z$ to desired data $x$

$$\begin{array}{}\text{(1)}& x=F_{\theta }(ϵ),\text{ where }ϵ\sim \mathcal{N}(0,I)\end{array}$$

The learning objective is crafted in a way that the $F_{\theta }$ induces a model distribution on $x$, i.e. $p_{\theta }(x)$ that matches a given $q_{data}(x)$ as closely as possible. That is, of course, from the samples of $q_{data}(x)$ only. Different generative model families learn (parameters of) the mapping in different ways.

Brief overview of Diffusion Models

The generative process

Although there are mutiple formalisms to describe the underlying theory of Diffusion Models, the one that gained traction recently is the “Differential Equations” view, mostly due to Song et al. (2021). Under this formalism, Diffusion’s generative mapping (in Equation 1) can be realized by integrating a differential equation in time $t$. Song et al. (2021) also showed that there are two equivalent generative processes, one stochastic and another deterministic, that leads to exact same noise-to-data mapping.

Stochastic Generative Process

$$\begin{array}{}\text{(2)}& d{x}_t=[f(t)x_t-g^2(t)\underset{\approx s_{\theta }(x_t,t)}{\underbrace{{\mathrm{\nabla }}_{x_t}\log p_t(x_t)}}]dt+g(t)d\overline{w}\end{array}$$ The hyper-parameters $f(t)$ and $g(t)$ are scalar functions of time.

Deterministic Generative Process

$$\begin{array}{}\text{(3)}& d{x}_t=[f(t)x_t-0.5\cdot g^2(t)\underset{\approx s_{\theta }(x_t,t)}{\underbrace{{\mathrm{\nabla }}_{x_t}\log p_t(x_t)}}]dt\cancel{+g(t)d\overline{w}}\end{array}$$ The hyper-parameters $f(t)$ and $g(t)$ are scalar functions of time.

where $p_t(x_t)$ is a noisy version of the data density¹ induced by a (known & fixed) forward process over time $t$

¹ Knowns as the forward process marginal

Forward SDE

$$d{x}_t=f(t)x_{t}dt+g(t)dw$$

Practical sampling equation

$$\begin{array}{}\text{(4)}& x_t=\alpha (t)x_0+\sigma (t)ϵ\end{array}$$

Conditional density

$$p_t(x_t|x_0)=\mathcal{N}(x_t;\alpha (t)x_0,\sigma (t{)}^2\mathrm{I})$$

Conversion between $f,g$ and $\alpha ,\sigma$

The relationship between $\{f(t),g(t)\}$ and $\{\alpha (t),\sigma (t)\}$ is as follows.

$$\begin{array}{}\text{(5)}& f(t)=\frac{d}{dt}\log \alpha (t),\text{ and }g(t{)}^2=-2\sigma (t{)}^2\frac{d}{dt}\log \frac{\alpha (t)}{\sigma (t)}\end{array}$$

Despite being able sample $x_t$ using Equation 4, the quantity ${\mathrm{\nabla }}_{x_t}\log p_t(x)$² in Equation 2 or Equation 3 is not analytically computable and therefore learned using a neural function.

² The ‘score’ of the marginal

The learning objective

A parametric neural function $s_{\theta }(x_t,t)$ is regressed using a different regression target than the ‘true’ target

Practical learning objective

$$\begin{array}{}\text{(6)}& \mathcal{L}(\theta )={\mathbb{E}}_{t,x_0\sim q_{data}(x),x_t\sim p_t(x_t|x_0)}[||s_{\theta }(x_t,t)-{\mathrm{\nabla }}_{x_t}\log p_t(x_t|x_0)|{|}^2]\end{array}$$ This objective is known as “Denoising Score Matching” (Vincent 2011).

True learning objective

$$\begin{array}{}\text{(7)}& \mathcal{L}(\theta )={\mathbb{E}}_{t,x_0\sim q_{data}(x),x_t\sim p_t(x_t|x_0)}[||s_{\theta }(x_t,t)-{\mathrm{\nabla }}_{x_t}\log p_t(x_t)|{|}^2]\end{array}$$ This is the true objective that is not directly computable.

It was proved initially by Vincent (2011) and later re-established by Song et al. (2021) that using ${\mathrm{\nabla }}_{x_t}\log p_t(x_t|x_0)$³ (in Equation 6) as an alternative target for regression can still lead to an ubiased estimate of the true ${\mathrm{\nabla }}_{x_t}\log p_t(x_t)$. The figure below show graphically, the quantity (red arrows) our parametric score model regresses against at an arbitrary timestep $t$.

³ Please note that ${\mathrm{\nabla }}_{x_t}\log p_t(x_t|x_0)=-\frac{x_t-x_0}{{\sigma }_t^2}=-\frac{ϵ}{{\sigma }_t}$

The parametric model regresses against the vector that points towards the original data point.

Matching flows, not scores

The idea of re-interpreting reverse diffusion as “flow”⁴ stems from a holistic observation. Note that given Equation 3, an ODE dynamics is guaranteed to exists which induces a deterministic mapping from noise to any given data distribution. In diffusion model’s framework, we were only learning a part of the dynamics – not the dynamics itself.

⁴ This term comes from Continuous Normalizing Flows (CNFs)

$$\begin{array}{}\text{(8)}& d{x}_t=[\underset{\text{.. why not learn this ?}}{\underbrace{f(t)x_t-0.5\cdot g^2(t)\stackrel{\text{Instead of this ..}}{\overbrace{s_{\theta }(x_t,t)}}}}]dt\end{array}$$

A deterministic model

Following the observation above, we can assume a generative model realized by a deterministic ODE simulation, whose parametric dynamics subsumes $f(t),g(t)$ and the parametric scores $s_{\theta }(x_t,t)$

$$\begin{array}{}\text{(9)}& d{x}_t=v_{\theta }(x_t,t)dt,\text{ where }{x}_1:=ϵ\sim \mathcal{N}(0,I)\end{array}$$

Turns out, models like Equation 9 have already been investigated in generative modelling literature under the name of Continuous Normalizing Flows (Chen et al. 2018). However, these models never made it to the scalable realm due to their “simulation based”⁵ learning objective. The dynamics is often called “velocity”⁶ or “velocity field” and denoted with $v$.

⁵ One must integrate or simulate the ODE during training.

⁶ It is a time derivative of position.

⁷ .. or as it is now called, the ‘Flow Matching’ loss

Upon inspecting the pair of Equation 8 and Equation 6, it is not particularly hard to sense the existance of an equivalent ‘velocity matching’ loss⁷ for the new flow model in Equation 9.

An equivalent objective

To see that exact form of the flow matching loss, simply try recreating the ODE dynamics in Equation 8 within Equation 6 by appending some extra terms that cancel out

$$\begin{array}{rl}\mathcal{L}(\theta )={\mathbb{E}}_{t,x_0\sim q_{data}(x),x_t\sim p_t(x_t|x_0)}& [||\frac{2}{g^2(t)}\{(f(t)x_t-\frac{1}{2}{g}^2(t)s_{\theta }(x_t,t))\\ & {-(f(t)x_t-\frac{1}{2}{g}^2(t){\mathrm{\nabla }}_{x_t}\log p_t(x_t|x_0))\}||}^2]\end{array}$$

The expression within the first set of parantheses ( .. ) is equivalent to what now call the parametric velocity/flow or $v_{\theta }(x_t,t)$. The expression in second set of parantheses is a proxy regression target (let’s call it $v(x_t,t)$), equivalent to ${\mathrm{\nabla }}_{x_t}\log p_t(x_t|x_0)$ in score matching (see Equation 6). With the help of $\{f,g\}\rightleftarrows \{\alpha ,\sigma \}$ conversion (see Equation 5), it’s relatively easy to show that $v(x_t,t)$ can be written as the time-derivative of the forward sampling process

Reveal Proof

$$\begin{array}{rl}v(x_t,t)& \triangleq f(t)x_t-\frac{1}{2}{g}^2(t){\mathrm{\nabla }}_{x_t}\log p_t(x_t|x_0)\\ & =f(t)(\alpha (t)x_0+\sigma (t)ϵ)-[-\sigma (t)\frac{g(t{)}^2}{2\sigma (t)}\frac{ϵ}{\sigma (t)}]\\ & =\frac{\dot{\alpha }(t)}{\alpha (t)}(\alpha (t)x_0+\sigma (t)ϵ)-[\sigma (t)(\frac{\dot{\alpha }(t)}{\alpha (t)}-\frac{\dot{\sigma }(t)}{\sigma (t)})ϵ]\\ & =\dot{\alpha }(t)x_0+\cancel{\frac{\dot{\alpha }(t)}{\alpha (t)}\sigma (t)ϵ}-\cancel{\sigma (t)\frac{\dot{\alpha }(t)}{\alpha (t)}ϵ}+\dot{\sigma }(t)ϵ\end{array}$$

$$v(x_t,t)=\dot{\alpha }(t)x_0+\dot{\sigma }(t)ϵ={\dot{x}}_t$$

To summarize, the following is the general form of flow matching loss

$$\begin{array}{}\text{(10)}& {\mathcal{L}}_{FM}(\theta )={\mathbb{E}}_{t,x_0\sim q_{data}(x),x_t\sim p_t(x_t|x_0)}[||v_{\theta }(x_t,t)-(\underset{{\dot{x}}_t}{\underbrace{\dot{\alpha }(t)x_0+\dot{\sigma }(t)ϵ}})|{|}^2]\end{array}$$

In practice, as proposed by many (Lipman et al. 2022; Liu, Gong, and Liu 2022), we discard the weightning term just like Diffusion Model’s simple loss popularized by Ho, Jain, and Abbeel (2020). We may think of ${\dot{x}}_t$ as a stochastic velocity induced by the forward process. The model, when minimized with Equation 10, tries to mimic the stochastic velocity, but without having access to $x_0$.

Be careful with time direction

While the learning objective regresses against ${\dot{x}}_t$, we need $-{\dot{x}}_t$ for sampling. The extra negative is induced automatically while simulating Equation 9 in reverse time ($dt$ being negative). It is therefore equivalent to think the regression target to be $-{\dot{x}}_t$.

The parametric model regresses against the instantaneous stochastic velocity at any point $x_t$ on the stochastic path.

The minima & its interpretation

The loss in Equation 10 can be shown to be equivalent⁸ to

⁸ Their gradients are equal, but the losses are not.

$$\begin{array}{}\text{(11)}& {\mathcal{L}}_{FM}(\theta )={\mathbb{E}}_{t,x_t\sim p_t(x_t)}[||v_{\theta }(x_t,t)-{\mathbb{E}}_{x_0\sim p(x_0|x_t)}[{\dot{x}}_t]|{|}^2]\end{array}$$

which implies that the loss reaches its minima⁹ when the model perfectly learns

⁹ This is a typical MMSE estimator.

$$v_{\ast }(x_t,t)\triangleq {\mathbb{E}}_{x_0\sim p(x_0|x_t)}\left[{\dot{x}}_t\right]$$

Hence, $v_{\theta }$ can be regarded as the variational approximation of the posterior velocity $v_{\ast }$. Note that $v_{\ast }$ is non-causal, i.e. it has access to the true data distribution $q_{data}(x_0)$. On the other hand, the model $v_{\theta }$ must be causal. Hence, the learning process “causalizes” (Liu, Gong, and Liu 2022) the stochastic path.

The forward stochastic paths $x_t$ are also overlapping, while the model is a function¹⁰. The expectation ${\mathbb{E}}_{x_0\sim p(x_0|x_t)}[\cdot ]$ averages the stochastic velocity over all possible real data points $x_0|x_t$, leading to smooth velocity fields learned by the model.

¹⁰ cannot have multiple value at one given point

The optimally learned generative process (Equation 9) can therefore be expressed as

$$d{x}_t=v_{\ast }(x_t,t)dt,\text{ where }{x}_1:=ϵ\sim \mathcal{N}(0,I)$$

Straightening & ReFlow

$t$-independent stochastic velocity

With the general theory understood, it is now easy to concieve the idea of straight flows. It simply refers to the following special case

$$\alpha (t)=1-t,\sigma (t)=t$$

which implies the forward process and its time derivative to be

$$\begin{array}{rl}{x}_t& =(1-t)\cdot x_0+t\cdot x_1\\ ⟹{\dot{x}}_t& =-x_0+x_1\end{array}$$

What is important is the stochastic velocity is independent of time¹¹ and the $x_t$ trajectory itself is a stright line.

¹¹ Not “constant” – it still depends on data $x_0$ and noise sample $ϵ$, which are random variables.

This however, does not mean that the learned model will produce straight paths – it only means we’re supervising the model to follow a path as straight as possible. An analogous illustration was provided in Liu, Gong, and Liu (2022) which is a bit more descriptive.

The stochastic straight paths (left) used for training had crossovers. However, the learned model (right) resolves them. (Taken from Liu, Gong, and Liu (2022))

This learning problem (Equation 11) effectively turns an independent stochastic ‘coupling’ $q_{data}(x_0)\mathcal{N}(x_1;0,I)$ into a deterministic coupling $p(z_0,z_1)$ with some dependence. This prcoess $(x_0,x_1)\to (z_0,z_1)$ has been termed (by Liu, Gong, and Liu (2022)) as “Rectification”. Samples from this deterministic coupling can be drawn by drawing noise sample $z_1:=ϵ\sim \mathcal{N}(ϵ;0,I)$ and then simulating the flow in Equation 9 with the learned model

$$z_0=z_1+{\int }_{t=1}^{t=0}{v}_{\theta }(z_t,t)dt.$$

It can be proved that the samples $(z_0,z_1)$ are, in average, closer to each other than that of $(x_0,x_1)$.

Reveal Proof

$$\begin{array}{rl}{\mathbb{E}}_{p(z_0,z_1)}[||z_0-z_1||]& ={\mathbb{E}}_{p(z_0,z_1)}[||{\int }_1^0{v}_{\ast }(z_t,t)dt||]\\ & \le {\mathbb{E}}_{p(z_0,z_1)}[{\int }_1^0||v_{\ast }(z_t,t)||dt]\\ & ={\mathbb{E}}_{p(x_0,x_1)}[{\int }_1^0||v_{\ast }(x_t,t)||dt]\\ & ={\mathbb{E}}_{p(x_0,x_1)}[{\int }_1^0||\mathbb{E}[{\dot{x}}_t|x_t]||dt]\\ & \le {\mathbb{E}}_{p(x_0,x_1)}[{\int }_1^0\mathbb{E}[||{\dot{x}}_t|||x_t]dt]\\ & ={\int }_1^0{\mathbb{E}}_{p(x_0,x_1)}[\mathbb{E}[||x_0-x_1|||x_t]]dt\\ & ={\mathbb{E}}_{p(x_0,x_1)}[||x_0-x_1||]\end{array}$$

The proof uses the following

1. The fact that $||\cdot ||$ is a convex cost function.
2. Convex functions can be exchanged with $\mathbb{E}$ according to Jensen’s inequality.
3. $z_t$ and $x_t$ has the same marginal and can be exchanges as a random variable.
4. Assumes our model learns the perfect $v_{\ast }$.
5. Law of iterated expectation.

Please see section 3.2 of Liu, Gong, and Liu (2022) for more details on the proof.

Reflow

The process of rectification however, does not guarantee (as you can see in the above figure) the new coupling to have straight paths between each pair. Liu, Gong, and Liu (2022) suggested the “Reflow” procedure, which is nothing but learning a new model using the samples of $p(z_0,z_1)$.

$${\mathcal{L}}_{FM}(\varphi )={\mathbb{E}}_{t,(z_0,ϵ)\sim p(z_0,ϵ)}[||v_{\varphi }(z_t,t)-(-z_0+ϵ)|{|}^2]$$

The ‘reflow’-ed coupling $(z_2,ϵ)$ is guaranteed to have paths straighter than $(z_1,ϵ)$. One can, in fact, repeat this procedure as many times as they want. Another figure from Liu, Gong, and Liu (2022) perfectly demonstrates successive reflows

Successive reflows produce more and more straight paths.

---

In this article, we talked about Flow Matching, Rectification and Reflow – some of the emerging new ideas in Diffusion Model literature. Specifically, we looked into the theoretical definitions and justifications behind the ideas. Despite having an appealing outlook, some researchers are skeptical of it being a special case of good old Diffusion Models. Whatever the case maybe, it did deliver one of the best text-to-image model so far (Esser et al. 2024), pehaps with a bit of clever engineering, which is a topic of another day.

References

Chen, Tian Qi, Yulia Rubanova, Jesse Bettencourt, and David Duvenaud. 2018. “Neural Ordinary Differential Equations.” In NeurIPS.

Esser, Patrick, Sumith Kulal, Andreas Blattmann, Rahim Entezari, Jonas Müller, Harry Saini, Yam Levi, et al. 2024. “Scaling Rectified Flow Transformers for High-Resolution Image Synthesis.” arXiv Preprint arXiv:2403.03206.

Ho, Jonathan, Ajay Jain, and Pieter Abbeel. 2020. “Denoising Diffusion Probabilistic Models.” In NeurIPS.

Lipman, Yaron, Ricky TQ Chen, Heli Ben-Hamu, Maximilian Nickel, and Matthew Le. 2022. “Flow Matching for Generative Modeling.” In The Eleventh International Conference on Learning Representations.

Liu, Xingchao, Chengyue Gong, and Qiang Liu. 2022. “Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow.” arXiv Preprint arXiv:2209.03003.

Podell, Dustin, Zion English, Kyle Lacey, Andreas Blattmann, Tim Dockhorn, Jonas Müller, Joe Penna, and Robin Rombach. 2024. “SDXL: Improving Latent Diffusion Models for High-Resolution Image Synthesis.” In The Twelfth International Conference on Learning Representations.

Ramesh, Aditya, Prafulla Dhariwal, Alex Nichol, Casey Chu, and Mark Chen. 2022. “Hierarchical Text-Conditional Image Generation with Clip Latents.” arXiv Preprint arXiv:2204.06125.

Rombach, Robin, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer. 2022. “High-Resolution Image Synthesis with Latent Diffusion Models.” In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, 10684–95.

Song, Yang, Diederik P. Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. 2021. “Score-Based Generative Modeling Through Stochastic Differential Equations.” In ICLR.

Vincent, Pascal. 2011. “A Connection Between Score Matching and Denoising Autoencoders.” Neural Computation.