MIT 6.S184, Flow Matching and Diffusion Models

MIT 6.S184, Flow Matching and Diffusion Models.

represent the generated objects as vectors: $z \in \mathbb{R}^d$.

---

data distribution: how likely are we to find this picture on the Internet?

how good an image is $\approx$ how likely it is under the data distribution

---

generation: sampling from the data distribution

data distribution $p_{\text{data}}$: distribution of generated objects
unknown probability density $p_{\text{data}}(z)$: given $z$, output probability of $z$ in data distr

$z \sim p_{\text{data}} \rightarrow z = \text{an image of a dog}$

---

dataset: finite number of samples from the data distribution

$z_1, \cdots, z_N \sim p_{\text{data}}$

---

conditional generation allows us to condition on prompts

data distribution $p_{\text{data}}$, fixed prompt, e.g., "dog"
conditional data distribution $p_{\text{data}} (\cdot \mid y)$

---

a generative model converts samples from a initial distr (e.g., Gaussian) into samples from the data distr:

---

Conditional Probability Paths

The basic premise of conditional flow matching: describe a conditional probability path $p_t(x|z)$, so that $p_1(x|z) = \delta_z(x)$, and $p_0(z) = p_{\text{simple}}$ (e.g., a Gaussian), and $p_t(x|z)$ interpolates continuously between $p_0(x|z)$ and $p_1(x|z)$. Such a conditional path can be seen as corresponding to some corruption process which (in reverse time) drives the point $z$ at $t=1$ to be distribution as $p_0(x|z)$ at time $t=0$. Such a corruption process is given by the ODE:

$dX_t = u_t^{\text{ref}}(X_t|z)\,dt,\quad \quad X_0 \sim p_{\text{simple}}.$

The drift $u_t^{\text{ref}}(X_t|z)$ is referred to as the conditional vector field. By averaging $u_t^{\text{ref}}(x|z)$ over all such choices of $z$, we obtain the marginal vector field $u_t^{\text{ref}}(x)$. Flow matching proposes to exploit the fact that the marginal probability path $p_t(x)$ generated by the marginal vector field $u_t^{\text{ref}}(x)$, bridges $p_{\text{simple}}$ to $p_{\text{data}}$. Since the conditional vector field $u_t^{\text{ref}}(x|z)$ is often analytically available, we may implicitly regress against the unknown marginal vector field $u_t^{\text{ref}}(x)$ by explicitly regressing against the conditional vector field $u_t^{\text{ref}}(x|z)$.

Gaussian Conditional Probability Paths

A Gaussian conditional probability path is used to transform a simple source $p_{\text{simple}} = \mathcal{N}(0, I_d)$ into a Gaussian mixture $p_{\text{data}}$. Later, we'll experiment with more exciting distributions. Recall that a Gaussian conditional probability path is given by

$p_t(x|z) = \mathcal{N}(x; \alpha_t z, \beta_t^2 I_d), \quad\quad\quad p_{\text{simple}} = \mathcal{N}(0, I_d),$

where $\alpha_t: [0,1] \to \mathbb{R}$ and $\beta_t: [0,1] \to \mathbb{R}$ are monotonic, continuously differentiable functions satisfying $\alpha_1 = \beta_0 = 1$ and $\alpha_0 = \beta_1 = 0$. In other words, this implies that $p_1(x|z) = \delta_z$ and $p_0(x|z) = \mathcal{N}(0, I_d)$ is a unit Gaussian. Before we dive into things, let's look at $p_{\text{simple}}$ and $p_{\text{data}}$.

In this section, we'll use

$\alpha_t = t \quad \quad \text{and} \quad \quad \beta_t = \sqrt{1-t}.$

It is not hard to check that both functions are continuously differentiable on $[0,1)$, and monotonic, that $\alpha_1 = \beta_0 = 1$, and that $\alpha_0 = \beta_1 = 0$.

Sample the Conditioning Variable

$z \sim p_{\text{data}}$

Sample from the Conditional Distribution

Sample from the conditional distribution $p_t(x|z) = \mathcal{N}(x;\alpha_t z,\beta_t^2 I_d)$.

A random variable $X \sim \mathcal{N}(\mu, \sigma^2 I_d)$ can be obtained via $X = \mu + \sigma Y$, where $Y \sim \mathcal{N}(0, I_d)$, so our sample is $x = \alpha_t z + \beta_t X_0$.

An image of the conditional probability path: from the central red cloud to the final red dot.

Conditional Vector Field

A conditional flow model is defined as $\psi_t^{\text{target}} (x | z) = \alpha_t z + \beta_t x$. Taking $\frac{d}{dt} \psi_t^{\text{target}} (x | z)$ gives the conditional vector field $u_t(x|z)$:

$u_t(x|z) = \left(\dot{\alpha}_t-\frac{\dot{\beta}_t}{\beta_t}\alpha_t\right)z+\frac{\dot{\beta}_t}{\beta_t}x.$

Visualize the conditional trajectories corresponding to the ODE

$d X_t = u_t(X_t | z)dt, \quad \quad X_0 = x_0 \sim p_{\text{simple}}.$

Conditional Score

$\nabla_x \frac12 \|u\|^2 = u^T \nabla_x u$

$\dot a = - \Big( a^T\nabla_x b + \nabla_x f + \underbrace{ u^T\nabla_x u + a^T\sigma \nabla_x u }_{extra terms} \Big)$

Compute the conditional distr $\nabla_x \log p_t(x|z)$ of $p_t(x|z)$, the conditional score of $p_t(x|z) = \mathcal{N}(x; \alpha_t z, \beta_t^2 I_d)$:

$\nabla_x \log p_t(x|z) = \nabla_x \log \mathcal{N}(x;\alpha_t z,\beta_t^2 I_d) = \frac{\alpha_t z - x}{\beta_t^2}.$

To check for correctness, verify that samples from the conditional SDE match the samples drawn analytically from the conditional probability path (left and middle plot are the same).

Note: You may notice that strange things happen for large (or even not-so-large) values of $\sigma_t$. Plugging in $\nabla_x \log p_t(x|z) = \frac{\alpha_t z - x}{\beta_t^2}$ into

$d X_t = \left[u_t(X_t|z) + \frac{1}{2}\sigma_t^2 \nabla_x \log p_t(X_t|z) \right]dt + \sigma_t \, dW_t$

yields

$d X_t = \left[u_t(X_t|z) + \frac{1}{2}\sigma_t^2 \left(\frac{\alpha_t z - X_t}{\beta_t^2}\right) \right]dt + \sigma_t\, dW_t.$

When $t \to 1$, $\beta_t \to 0$, so that the second term of the drift explodes (and this explosion scales quadratically with $\sigma_t$). With a finite number of simulation steps, we cannot accurately simulate this explosion and thus encounter numerical issues. In practice, this is usually circumvented by setting, e.g., $\sigma_t = \beta_t$, so that the exploding effect is canceled out by a gradually decreasing noise level.

Visualize the conditional trajectories corresponding to the SDE

$d X_t = \left[u_t(X_t|z) + \frac{1}{2}\sigma_t^2 \nabla_x \log p_t(X_t|z) \right]dt + \sigma_t\, dW_t, \quad \quad X_0 = x_0 \sim p_{\text{simple}},$

obtained by adding Langevin dynamics to the original ODE.

Flow Matching and Score Matching with Gaussian Conditional Probability Paths

Flow Matching with Gaussian Conditional Probability Paths

Our goal is to learn the marginal vector field $u_t(x)$ given by

$u_t^{\text{ref}}(x) = \mathbb{E}_{z \sim p_t(z|x)}\left[u_t^{\text{ref}}(x|z)\right].$

Unfortunately, we don't actually know what $u_t^{\text{ref}}(x)$ is. We therefore approximate $u_t^{\text{ref}}(x)$ with a neural network $u_t^{\theta}(x)$, and exploit the identity

$u_t^{\text{ref}}(x) = \text{argmin}_{u_t(x)} \,\,\mathbb{E}_{z \sim p_t(z|x)} \left[\lVert u_t(x) - u_t^{\text{ref}}(x|z)\rVert^2\right]$

to obtain the conditional flow matching objective,

$\mathcal{L}_{\text{CFM}}(\theta) = \,\,\mathbb{E}_{\textcolor{blue}{t \in \mathcal{U}[0,1), z \sim p(z), x \sim p_t(x|z)}} \textcolor{green}{\lVert u_t^{\theta}(x) - u_t^{\text{ref}}(x|z)\rVert^2}$

To model $u_t^{\theta}(x)$, we'll use a simple MLP. This network will take in both $x$ and $t$, and will return the learned vector field $u_t^{\theta}(x)$. Using a Monte-Carlo estimate of the form

$\frac{1}{N}\sum_{i=1}^N \textcolor{green}{\lVert u_{t_i}^{\theta}(x_i) - u_{t_i}^{\text{ref}}(x_i|z_i)\rVert^2}, \quad \quad \quad \forall i\in[1, \dots, N]: \textcolor{blue}{\,z_i \sim p_{\text{data}},\, t_i \sim \mathcal{U}[0,1),\, x_i \sim p_t(\cdot | z_i)}.$

Here, $N$ is the batch size.

Note that sampling uses conditional probability $p_t(x | z)$ and loss uses conditional vector field $u_t^{\text{ref}}(x|z)$ from before.

Score Matching with Gaussian Conditional Probability Paths

We have thus far used flow matching to train a model $u_t^{\theta}(x) \approx u_t^{\text{ref}}$ so that

$d X_t = u_t^{\theta}(X_t) dt$

approximately passes through the desired marginal probability path $p_t(x)$. For a diffusion process, we augment the reference marginal vector field $u_t^{\text{ref}}(x)$ with Langevin dynamics to add stochasticity while preserving the marginals, viz.,

$d X_t = \left[\textcolor{blue}{u_t^{\text{ref}}(x)} + \frac{1}{2}\sigma^2 \nabla \log p_t(x)\right] dt + \sigma d W_t.$

Substituting our learned approximation $u_t^{\theta}(x) \approx u_t^{\text{ref}}$ therefore yields

$dX_t = \left[\textcolor{blue}{u_t^{\theta}(x)} + \frac{1}{2}\sigma^2 \nabla \log p_t(x)\right] dt + \sigma d W_t.$

But what's the marginal score $\nabla \log p_t(x)$? In conditional score, we computed the conditional score $\nabla \log p_t(x|z)$ of the Gaussian probability path. In the same way that we learned an approximation $u_t^{\theta}(x) \approx u_t^{\text{ref}}$, we'd like to be able to learn a similar approximation $s_t^{\theta}(x) \approx \nabla \log p_t(x)$. Given the identity

$\nabla \log p_t(x) = \mathbb{E}_{z \sim p_t(z|x)}\left[\nabla \log p_t(x|z) \right].$

It then immediately follows that

$\nabla \log p_t(x) = \text{argmin}_{s_t(x)} \,\,\mathbb{E}_{z \sim p(z), x \sim p_t(x|z)} \left[\lVert s_t(x) - \nabla \log p_t(x|z)\rVert^2\right].$

We thus obtain the conditional score matching loss

$\mathcal{L}_{\text{CSM}}(\theta) \triangleq \mathbb{E}_{t \sim \mathcal{U}[0,1), z \sim p(z), x \sim p_t(x|z)} \left[\lVert s_t^{\theta}(x) - \nabla \log p_t(x|z)\rVert^2\right].$

Here, we will parameterize $s_t^{\theta}(x): \mathbb{R}^2 \to \mathbb{R}^2$ as a simple MLP, just like $u_t^{\theta}(x)$.

After training, the loss should converge, but not to zero b/c $\mathcal{L}_{\text{SM}}(\theta) = \mathcal{L}_{\text{CSM}}(\theta)$ and $\nabla_\theta \mathcal{L}_{\text{SM}}(\theta) = \nabla_\theta \mathcal{L}_{\text{CSM}}(\theta)$.

The diffusion process is defined as an SDE with Langevin dynamics:

$dX_t = \left[u_t^{\theta}(x) + \frac{1}{2}\sigma^2 s_t^{\theta}(x)\right] dt + \sigma d W_t.$

Deriving the Marginal Score from the Marginal Flow

For Gaussian probability path $p_t (x | z) = \mathcal{N}(\alpha_t z, \beta^2 I_d)$, can express vector field in terms of score:

$u_t^{\text{ref}}(x) = a_t x + b_t\nabla \log p_t^{\text{ref}}(x).$

where $(a_t, b_t) = \left(\frac{\dot{\alpha}_t}{\alpha_t}, \beta_t^2 \frac{\dot{\alpha}_t}{\alpha_t} - \dot{\beta}_t \beta_t\right)$. Rearranging yields

$\nabla \log p_t^{\text{ref}}(x) = \frac{u_t^{\text{ref}}(x) - a_t x}{b_t}.$

Therefore, we may instead exploit the fact that we have already trained $u_t^{\theta}(x)$, to parameterize $s_t^{\theta}(x)$ via

$\tilde{s}_t^{\theta}(x) = \frac{u_t^{\theta}(x) - a_t x}{b_t} = \frac{\alpha_t u_t^{\theta}(x) - \dot{\alpha}_t x}{\beta_t^2 \dot{\alpha}_t - \alpha_t \dot{\beta}_t \beta_t},$

so long as $\beta_t^2 \dot{\alpha}_t - \alpha_t \dot{\beta}_t \beta_t \neq 0$ (which is true for $t \in [0,1)$ by monotonicity). Here, we differentiate $\tilde{s}_t^{\theta}(x)$ paramterized via $u_t^{\theta}(x)$ from $s_t^{\theta}(x)$ learned indepedently using score matching. Plugging in $\alpha_t = t$ and $\beta_t = \sqrt{1-t}$, we find that

$\beta_t^2 \dot{\alpha}_t - \alpha_t \dot{\beta}_t \beta_t = \begin{cases} 1 - \frac{t}{2} & \text{if}\,\,t\in [0,1)\\0 & \text{if}\,\,{t=1}. \end{cases}.$

In the following visualization, we'll circumvent the issue at $t=1$ by taking $t=1 - \varepsilon$ in place of $t=1$, for small $\varepsilon \approx 0$.

Now, let's compare our learned marginal score $s_t^{\theta}(x)$ (using MLP) to our flow-parameterized score (score from VF). We'll do so by plotting the vector fields across time and space.

Note: The two score parameterizations will probably look a bit different, but should generally point in the same direction, especially around modes.
The two score fields would only match if both the learned flow and vector field match the true marginal score, but they are only approximations.

Flow Matching Between Arbitrary Distributions with a Linear Probability Path

In this section, we consider an alternative conditional probability path, the linear conditional probability path, which can be constructed as follows. Given a source distribution $p_{\text{simple}}$ and a data distribution $p_{\text{data}}$, for a fixed $z$ we may consider the interpolant

$X_t = (1-t) X_0 + tz$

where $X_0 \sim p_{\text{simple}}$ is a random variable. We may then define $p_t(x|z)$ so that $X_t \sim p_t(x|z)$. Then $p_0(x|z) = p_{\text{simple}}(x)$ and $p_1(x|z) = \delta_z(x)$. It is also not difficult to show that the conditional vector field is given by

$u_t^{\text{ref}}(x|z) = \frac{z - x}{1-t} \quad \text{for } t \in [0,1).$

We make two observations about the linear conditional probability path:

1. Unlike in the Gaussian probability path, we do not have a closed form for the conditional score $\nabla \log p_t(x|z)$.
2. There is no constraint that $p_{\text{simple}}$ be Gaussian, which we will exploit in from circles to checkerboard to construct flows between arbitrary choices of $p_{\text{simple}}$ and $p_{\text{data}}$.

First, let's visualize some more complicated choices of $p_{\text{data}}$.

The following shows linear probability paths of conditional (GT), conditional vector field (ODE), marginal:

Flow Matching with Linear Probability Paths

Train a flow matching model using the linear conditional probability path. The loss should converge, but not necessarily to zero.

Bridging Between Arbitrary Source and Target

For the linear probability path, there is no need for $p_{\text{simple}}$ to be a Gaussian. Set it to another distribution, e.g., circles:

From the output plots, the learned flow is able to transport samples from $p_{\text{simple}}$ to $p_{\text{data}}$ for several choices of source and target distributions. The match is typically better when the source and target are geometrically similar, while more different pairs, such as circles to checkerboard, produce more distorted intermediate marginals and are harder to learn. This suggests that linear probability paths are very flexible, but difficult source-target pairs may require more training and a higher-capacity model.

---

guided: "Generate an image."
unguided: "Generate an image of a cat baking a cake."

vanilla guided sampling:

unguided vs. guided (condition on $y$)

---

Classifier Free Guidance

Guidance

For unconditional generation, we simply wanted to generate any digit, we would now like to be able to specify, or condition, on the identity of the digit we would like to generate. That is, we would like to be able to say "generate an image of the digit 8", rather than just "generate an image of a digit". We will henceforth refer to the digit we would like to generate as $x \in \mathbb{R}^{1 \times 32 \times 32}$, and the conditioning variable (in this case, a label), as $y \in \{0, 1, \dots, 9\}$. If we imagine fixing our choice of $y$, and take our data distribution as $p_{\text{simple}}(x|y)$, then we have recovered the unconditional generative problem, and we can construct a generative model using e.g., a conditional flow matching objective via

$\begin{aligned} \mathcal{L}_{\text{CFM}}^{\text{guided}}(\theta;y) &= \,\,\mathbb{E}_{\square} \lVert u_t^{\theta}(x|y) - u_t^{\text{ref}}(x|z)\rVert^2\\ \square &= z \sim p_{\text{data}}(z|y), x \sim p_t(x|z) \end{aligned}$

We may now then allow $y$ to vary by simply taking our conditional flow matching expectation to be over $y$ as well (rather than fixing $y$), and explicitly conditioning our learned approximation on $u_t^{\theta}(x|y)$ on the choice of $y$. We therefore obtain the the guided conditional flow matching objective

$\begin{aligned} \mathcal{L}_{\text{CFM}}(\theta) &= \,\,\mathbb{E}_{\square} \lVert u_t^{\theta}(x|y) - u_t^{\text{ref}}(x|z)\rVert^2\\ \square &= z,y \sim p_{\text{data}}(z,y), x \sim p_t(x|z) \end{aligned}$

Note that $(z,y) \sim p_{\text{simple}}(z,y)$ is obtained in practice by sampling an image $z$, and a label $y$, from our labelled (MNIST) dataset. This is all well and good, and we emphasize that if our goal was simply to sample from $p_{\text{data}}(x|y)$, our job would be done (at least in theory). In practice, one might argue that we care more about the perceptual quality of our images. To this end, we will a derive a procedure known as classifier-free guidance.

CFG-CDM: Classifier-Free Guidance

For the sake of intuition, we will develop guidance through the lense of Gaussian probability paths, although the final result might reasonably be applied to any probability path. Recall from the lecture that for $(a_t, b_t) = \left(\frac{\dot{\alpha}_t}{\alpha_t}, -\frac{\dot{\beta}_t \beta_t \alpha_t - \dot{\alpha}_t \beta_t^2}{\alpha_t}\right)$, we have

$u_t(x|y) = a_t x + b_t\nabla \log p_t(x|y).$

This identity allows us to relate the conditional marginal velocity $u_t(x|y)$ to the conditional score $\nabla \log p_t(x|y)$. However, notice that

$\nabla \log p_t(x|y) = \nabla \log \left(\frac{p_t(x)p_t(y|x)}{p_t(y)}\right) = \nabla \log p_t(x) + \nabla \log p_t(y|x),$

so that we may rewrite

$u_t(x|y) = a_t x + b_t(\nabla \log p_t(x) + \nabla \log p_t(y|x)) = u_t(x) + b_t \nabla \log p_t(y|x).$

An approximation of the term $\nabla \log p_t(y|x)$ could be considered as a sort of noisy classifier (and in fact this is the origin of classifier guidance, which we do not consider here). In practice, people have noticed that the conditioning seems to work better when we scale the contribution of this classifier term, yielding

$\tilde{u}_t(x|y) = u_t(x) + w b_t \nabla \log p_t(y|x)$

where $w > 1$ is known as the guidance scale. We may then plug in $b_t\log p_t(y|x) = u^{\text{target}}_t(x|y) - u^{\text{target}}_t(x)$ to obtain

$\begin{aligned} \tilde{u}_t(x|y) &= u_t(x) + w b_t \nabla \log p_t(y|x)\\ &= u_t(x) + w (u^{\text{target}}_t(x|y) - u^{\text{target}}_t(x))\\ &= (1-w) u_t(x) + w u_t(x|y). \end{aligned}$

The idea is thus to train both $u_t(x)$ as well as the conditional model $u_t(x|y)$, and then combine them at inference time to obtain $\tilde{u}_t(x|y)$. Our recipe will thus be:

1. Train $u_t^{\theta} \approx u_t(x)$ as well as the conditional model $u_t^{\theta}(x|y) \approx u_t(x|y)$ using conditional flow matching.
2. At inference time, sample using $\tilde{u}_t^{\theta}(x|y)$.

But why must we train two models? We can instead treat $u_t(x)$ as $u_t(x|y)$, where $y=\varnothing$ denotes the absence of conditioning. We may thus augment our label set with a new, additional $\varnothing$ label, so that $y \in \{0,1,\dots, 9, \varnothing\}$. This technique is known as classifier-free guidance (CFG). We thus arrive at

$\boxed{\tilde{u}_t(x|y) = (1-w) u_t(x | \varnothing) + w u_t(x|y)}.$

Training and CFG: We must now amend our conditional flow matching objective to account for the possibility of $y = \varnothing$. Of course, when we sample $(z,y)$ from MNIST, we will never obtain $y = \varnothing$, so we must introduce the possibliity of this artificially. To do so, we will define some hyperparameter $\eta$ to be the probability that we discard the original label $y$, and replace it with $\varnothing$. In practice, we might set $\varnothing = 10$, for example, as it is sufficient to distinguish it from the other digit identities. When we go and implement our model, we need ony be able to index into some embedding, such as via torch.nn.Embedding. We thus arrive at our CFG conditional flow matching training objective:

$\begin{aligned} \mathcal{L}_{\text{CFM}}(\theta) &= \,\,\mathbb{E}_{\square} \lVert u_t^{\theta}(x|y) - u_t^{\text{ref}}(x|z)\rVert^2\\ \square &= t \sim \mathcal{U}[0,1),\, z,y \sim p_{\text{data}}(z,y),\, x \sim p_t(x|z),\, \text{replace } y \text{ with } \varnothing \text{ with probability } \eta \end{aligned}$

In plain English, this objective reads:

1. Sample an image $z$ and a label $y$ from $p_{\text{data}}$ (here, MNIST).
2. With probability $\eta$, replace the label $y$ with the null label $\varnothing \triangleq 10$.
3. Sample $t$ from $\mathcal{U}[0,1)$.
4. Sample $x$ from the conditional probability path $p_t(x|z)$.
5. Regress $u_t^{\theta}(x|y)$ against $u_t^{\text{ref}}(x|z)$.

Build a Diffusion Transformer

For images as input, MLP will no longer suffice. We will use a diffusion transformer.

Fourier Time Encoder

First, we implement a Fourier time encoder, which maps a scalar time value $t \in [0,1]$ to

$t^{\text{emb}} = \begin{bmatrix} \cos(2\pi w_1 t) & \cdots & \cos(2\pi w_d t) & \sin(2\pi w_1 t) & \cdots & \sin(2\pi w_d t) \end{bmatrix}^T,$

where the weights $w_i \sim \mathcal{N}(0, 1)$ are drawn from a unit normal.

Patchifier

The patchifier takes in an image-valued tensor of shape b 1 32 32, and patchifies it to shape b (h / p * w / p) d, where d denotes the diffusion transformer hidden dimension, p denotes the patch size, and h and w denote the image height and width, respectively. It does this in two steps:

1. Applies a convolutional layer mapping the an input of shape b 1 32 32 to b d h/p h/w.
2. Rearranges from b d h/p h/w to b (h/p h/w) d (i.e., n = h/p * h/w tokens with dimension d).

Diffusion Transformer

Adaptive LayerNorm formula

For hidden state $x$ and condition $c$, adaptive layernorm:

$\mathrm{AdaLN}(x,c) = (1 + \gamma(c)) \odot \mathrm{Norm}{x} + \beta(c)$

$\gamma(c)$ = scale
$\beta(c)$ = bias/shift
$\odot$ = elementwise multiply

In DiT-style adaLN-Zero, there is also a residual gate $\alpha(c)$:

$x_{\text{out}} = x + \alpha(c) \odot F\big((1+\gamma(c))\odot \mathrm{Norm}(x)+\beta(c)\big)$

where $F$ is either the attention block, or the feed-forward block.

Translated to code,

attn_scale, attn_bias, attn_gate, ff_scale, ff_bias, ff_gate = c.chunk(6, dim=-1)

x = x + attn_gate * self.attn(modulate(self.norm1(x), attn_scale, attn_bias))
x = x + ff_gate * self.ff(modulate(self.norm2(x), ff_scale, ff_bias))

with

$\mathrm{modulate}(x,\text{scale},\text{bias}) = x \odot (1+\text{scale}) + \text{bias}$

Scale, Bias, Gate

scale/bias: change the input to the branch
gate: change how much the branch contributes

Because the transformer layer has two different residual branches, there are scale, bias, gate for both attention and feed-forward:

1. Attention branch
2. Feed-forward branch

Each branch plays a different role, so conditioning them separately is useful.

Scale + Bias

These control how the condition changes the normalized activations before the sublayer.

• attention may need one type of conditioning
• feed-forward may need another

Sharing one modulation for both would be less flexible.

Gate: controls how much of that branch is added back through the residual.

$x + \alpha_{\text{attn}}(c), \text{Attn}(\cdot) \quad,\quad x + \alpha_{\text{ff}}(c), \text{FF}(\cdot)$

• the model may want strong conditioning in attention but weak in MLP, or vice versa
• separate gates make the layer more expressive
• with zero init, gates start near 0, so the block initially behaves close to identity, which improves stability

Depatchifier

After the diffusion transformer, we need to convert from b n d back to b 1 h w.

One approach would be to

1. Start with some sort of normalization. nn.LayerNorm or nn.RMSNorm wil suffice.
2. Pass through some MLP to obtain b (h/p w/p) (f p p) (ie, map to dimension $d=fp^2$, for patch size $p$ and "final dimension" $f$.
3. Rearrange to b f h w.
4. Pass through final convolution to obtain output of shape b 1 h w.

Putting It All Together

Finally, implement our DiT-based guided flow model $u_t^\theta(x|y)$.

1. To embed the guiding input y, you'll need class labels. You can reuse your solution from MLPConditionalVectorField. In particular, we recommend using nn.Embedding(num_classes=11, embedding_dim=dim). Note we use 11 - not 10 - classes to account for the presence of the null label ($\varnothing = 10$).
2. Follow the DiT overview diagram: embed $t$ and $y$, add them together to obtain the guiding variable. Pass $x$ through the patchifier, and then pass the patchifier output and $t + y$ through the diffusion transformer. Finally, pass the DiT output through the depatchifier.

Visualize final result on MNIST: