Flow matching and diffusion have emerged as the state-of-the-art techniques for generation of continuous data modalities, in particular images, audio and video [6, 7]. Their strong generation performance, training stability and flexibility motivated their extension to discrete data, including texts, graphs, or molecule sequences [1, 2, 5]. This blog post is an introduction to flow matching for discrete data. A second part TODO link is dedicated to discrete diffusion.

The first section is an introduction to discrete data generation, and discrete flow matching is introduced in Section 2.

Generating discrete data

Discrete data

Autoregressive generative models have being especially successful for generation of sequences of discrete elements. For instance, language models generate sequences of tokens (words or subwords) from a predefined discrete set (the vocabulary).

More precisely, in this blog post, we present flow matching/diffusion for sequences of discrete data:

• sequences $\mathbf{x} = x^1 \dots x^S$ of length $S \in \mathbb{N}$,
• where each element $x^i$ belongs to a predefined discrete/finite set called the vocabulary $\mathcal{V}$: $x^i \in \mathcal{V}$.

Sequence x = x¹x²x³x⁴ of length 4

x¹A

x²A

x³B

x⁴B

Vocabulary of 3 elements 𝒱 = {A, B, m}

 A

 B

 m

Elements of the vocabulary are called tokens. For example, for text data, tokens could be the letters: $\mathcal{V} = \{a, b, c, \dots, z \}$, while DNA sequences would correspond to $\mathcal{V}=\{A, T, C, G\}$. In the vocabulary, a special token can be added, the mask token $m$.

Discrete data generation process

Discrete generative modelling on Sudoku. The goal is to generate valid Sudoku grids, which are $9 \times 9$ grids filled with numbers from $1$ to $9$, such that each number appears exactly once in each row, column. The vocabulary is the set of numbers from $1$ to $9$: $\mathcal{V} = \{1, 2, \ldots, 9, m\}$, and the sequence length is number of squares to fill $S= 9 \times 9 = 81$.

Discrete diffusion/flow matching on Sudoku. Discrete flow matching starts from a fully masked grid: each square is initially filled with the mask token $m$. The model then iteratively fills this grid with non-mask tokens $\{1, 2, \ldots, 9\}$. For example, one can iteratively unmask the grid one square after one square, and finally reach a filled Sudoku grid.

QBE: todo do a sort of gif?

Generating a Sudoku grid

GIF mode

Generate new grid

t = 0.00

Diffusion

random order

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

81Masked

0Generated

Autoregressive

left to right

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

m

81Masked

0Generated

Masked [M]

Generated

Section 2 provides the details on to generate discrete data with discrete flow matching.

Discrete flow matching

Transporting distributions

Discrete flow matching iteratively transforms fully masked samples/images into samples/images from the target distribution. At the distribution level, discrete flow matching can be seen as a transport of a simple source probability distribution $p_0$ (e.g., $x_0 \sim p_0$ fully masked) to a target distribution $p_1$, the data distribution we want to generate.

QBE: introduce later?

x₀ ∼ p₀

mmmm

mmmm

mmmm

mmmm

x_t ∼ p_t

mAmm

Ammm

mBmm

mAAB

x₁ ∼ p₁

AAAB

ABBB

BBBB

AAAB

MaskedUniform

t = 0.50▶ PlayResample

The source distribution is typically chosen to be the mask distribution ($p_0 = \delta_{m\ldots m}$): all tokens in the starting sequence are equal to the mask token $m$; $\delta$ is the Dirac delta distribution.

QBE: schema transport distributions?

QBE: todo footnote uniform - The uniform distribution ($p_0 = p_{\text{unif}}$): All tokens in the starting sequence are sampled uniformly across the vocabulary $\mathcal{V}$.

QBE: decide what to do with this figure. Not sure it is cristal clear.

Flow matching: from continuous-space discrete-space data

Reminder on flow matching for continuous data

See the former blog post for a comprehensive explanation. In continuous-space flow matching, generated samples are obtained by solving the following Ordinary Differential Equation (ODE)

$$\begin{equation} \htmlData{equation-label=eq:flow-ode}{\htmlId{eq:flow-ode}{}} \begin{cases} x(0) = x_0 \\ \partial_t x(t) = u_t(x(t)), \quad \forall t \in [0,1] \end{cases} \end{equation}$$

which is driven by the so-called velocity field $u: \mathbb{R}^d \times [0, 1] \to \mathbb{R}^d$.

From continuous-space to discrete-space

QBE: comparison Algorithm Euler continuous vs discrete? + pixel vs sequence

QBE: discussion “shape” of u

QBE todo rephrase. As solving an ODE on a discrete space is not possible (since at each time $t$, the current state is $x_t$ should be a discrete sequence), the generation process in discrete flow matching is different.

It is defined as

$$\begin{equation} \htmlData{equation-label=eq:euler-sampling}{\htmlId{eq:euler-sampling}{}} \mathbf{x}_{t + h} \sim \text{Cat}\left(\mathbf{e}_{\mathbf{x}_t} + h u_t(\mathbf{x}_t)\right ) + o(h), \end{equation}$$

where $u : \mathcal{S} \times [0, 1] \to \mathbb{R}^S$ is the velocity field, $\mathbf{e}_{\mathbf{x}} \in \mathbb{R}^{|\mathcal{S}|}$ is the one-hot vector corresponding to state $\mathbf{x} \in \mathcal{S}$, and $\text{Cat}$ denotes the categorical distribution. The velocity $u_t(\mathbf{x})$ can be seen as the instantaneous rate at which the random variable $X_t$ transitions from state $\mathbf{x}$ to state $\mathbf{y}$ at time $t$. Below is a direct side by side analogy between the continuous and discrete flow matching generation processes.

Formal definition of the velocity $u_t$.

Definition 1 (velocity $u_t$).

Let $u : \mathcal{S} \times [0, 1] \to \mathbb{R}^S$, let $(X_t)_{t\in[0,1]}$ be random variables taking values in the space of sequences $\mathcal{S}$, and having laws $(p_t)_{t \in [0, 1]}$. Denote $u_t := u(\cdot, t)$. In the discrete case, we say that $u$ generates the probability path $(p_t)_{t \in [0, 1]}$ if for every sequences $\mathbf{x}, \mathbf{y} \in \mathcal{S}$,

$$\begin{equation} \htmlData{equation-label=eq:prob-velocity}{\htmlId{eq:prob-velocity}{}} \mathbb{P}(X_{t+h} = \mathbf{y} | X_t = \mathbf{x}) = \mathbb{I}(\mathbf{x} = \mathbf{y}) + h u_{t}(\mathbf{x})_{\mathbf{y}} + o(h) \end{equation}$$

If we index the state space as $\mathcal{S} = \{\mathbf{y}_1, \dots, \mathbf{y}_{|\mathcal{S}|}\}$ and choose $k$ such that $\mathbf{y}_k = \mathbf{x}$ (the current state), then the full transition-probability vector reads

$$\begin{equation} \htmlData{equation-label=eq:transition-vector}{\htmlId{eq:transition-vector}{}} \underbrace{\mathbb{P}(X_{t+h} = \cdot | X_t=\mathbf{x})}_{\in \mathbb{R}^{|\mathcal{S}|}} = \begin{bmatrix} h u_t(\mathbf{x})_{\mathbf{y}_1} \\ \vdots \\ h u_t(\mathbf{x})_{\mathbf{y}_{k-1}} \\ 1 + h u_t(\mathbf{x})_{\mathbf{y}_k} \\ h u_t(\mathbf{x})_{\mathbf{y}_{k+1}} \\ \vdots \\ h u_t(\mathbf{x})_{\mathbf{y}_{|\mathcal{S}|}} \end{bmatrix} + o(h). \end{equation}$$

These two equations mean that as $t$ increases to $t+h$, the variable $X_t$ stays in its current state $\mathbf{x}$ with probability $1 + h u_t(\mathbf{x})_{\mathbf{x}}$, or moves to any other state $\mathbf{y} \neq \mathbf{x}$ with probability $h u_t(\mathbf{x})_{\mathbf{y}}$.

A fundamental point is that unlike continuous FM, where sampling is deterministic once $x_0$ has been sampled from $p_0$, sampling in discrete flow matching is inherently stochastic at each step. This explains why we are allowed to take $p_0$ equal to a Dirac $\delta_{m\ldots m}$ and still produce different final outputs, which would not be possible with a deterministic sampling.

Justification

Positivity guarantees nonnegative transition probabilities:

$$\begin{equation} \htmlData{equation-label=eq:velocity-positivity-proof}{\htmlId{eq:velocity-positivity-proof}{}} \mathbb{P}(X_{t+h}=\mathbf{y}| X_t=\mathbf{x}) = h\,u_t(\mathbf{x})_{\mathbf{y}} + o(h) \ge 0, \qquad \mathbf{y}\neq\mathbf{x}. \end{equation}$$

Normalization guarantees total mass conservation:

$$\begin{equation} \htmlData{equation-label=eq:velocity-normalization-proof}{\htmlId{eq:velocity-normalization-proof}{}} \sum_{\mathbf{y}\in\mathcal{S}} \mathbb{P}(X_{t+h}=\mathbf{y}| X_t=\mathbf{x}) = 1 + h\sum_{\mathbf{y}\in\mathcal{S}}u_t(\mathbf{x})_{\mathbf{y}} + o(h) = 1 + o(h). \end{equation}$$

Fundamental equation linking $p_t$ and $u_t$: Kolmogorov equation

In continuous flow matching, we have the guarantee that a velocity field $u_t$ generates a probability path $(p_t)_{t \in [0, 1]}$ if the pair $(p_t, u_t)$ satisfies the continuity equation $\partial_t p_t + \nabla \cdot (u_t p_t) = 0$. The equivalent result for discrete data relies on the so-called Kolmogorov Equation:

Proposition 3 (Kolmogorov equation).

The velocity $u_t$ generates $(p_t)$ in the sense of Eq. (?) if and only if the pair $(u_t, p_t)$ satisfies the Kolmogorov equation:

$$\begin{equation} \htmlData{equation-label=eq:kolmogorov-equation}{\htmlId{eq:kolmogorov-equation}{}} \forall \mathbf{y} \in \mathcal{S}, \quad \frac{\mathrm{d}}{\mathrm{d}t} p_t(\mathbf{y}) = \sum_{\mathbf{x} \in \mathcal{S}} u_{t}(\mathbf{x})_{\mathbf{y}} \, p_t(\mathbf{x}) . \end{equation}$$

Proof

We can derive this directly by expanding the marginal probability at $t+h$ using the law of total probability and our velocity definition:

$$\begin{align} p_{t+h}(\mathbf{y}) &:= \mathbb{P}(X_{t+h} = \mathbf{y}) \notag \\ &= \sum_{\mathbf{x} \in \mathcal{S}} \mathbb{P}(X_{t+h} = \mathbf{y} | X_t = \mathbf{x})\, \mathbb{P}(X_t = \mathbf{x}) \htmlData{equation-label=eq:kolmogorov-proof-2}{\htmlId{eq:kolmogorov-proof-2}{}} \\ &= \sum_{\mathbf{x} \in \mathcal{S}} \left(\mathbb{I}(\mathbf{y} = \mathbf{x}) + h\,u_t(\mathbf{x})_{\mathbf{y}} + o(h)\right) p_t(\mathbf{x}) \htmlData{equation-label=eq:kolmogorov-proof-3}{\htmlId{eq:kolmogorov-proof-3}{}} \\ &= p_t(\mathbf{y}) + h \sum_{\mathbf{x} \in \mathcal{S}} u_t(\mathbf{x})_{\mathbf{y}}\, p_t(\mathbf{x}) + o(h) \htmlData{equation-label=eq:kolmogorov-proof-4}{\htmlId{eq:kolmogorov-proof-4}{}} \end{align}$$

Rearranging the terms and taking the limit as $h \to 0$ gives the desired result.

For the reverse implication (if the equation holds, then $u$ generates $p_t$), let $\tilde{p}_t$ denote the probability path that the velocity field $u_t$ generates, starting from the initial distribution $\tilde{p}_0 = p_0$. By the forward implication proven above, $\tilde{p}_t$ satisfies the Kolmogorov equation. By hypothesis, the given probability path $p_t$ also satisfies this exact same linear ordinary differential equation (ODE). Because both $p_t$ and $\tilde{p}_t$ are solutions to the same first-order linear ODEs and share the identical initial condition ($p_0 = \tilde{p}_0$), the Picard-Lindelöf uniqueness theorem guarantees that $p_t(\mathbf{y}) = \tilde{p}_t(\mathbf{y})$ for all times $t$ and all states $\mathbf{y} \in \mathcal{S}$. Therefore, the vector field $u_t$ generates the probability path $p_t$.

Learning probability velocities taking values in the full sequence space $\mathcal{S}=\mathcal{V}^S$ is cumbersome, as the cardinality of the latter is $|\mathcal{V}|^S$, e.g. $10^{81}$ for the sudoku example! To make computations tractable, discrete flow matching makes a factorization assumption: each transition is applied to every token of the sequence independently, according to a token-wise velocity $u_t^i : \mathcal{S} \times [0, 1] \to \mathbb{R}^{|\mathcal{V}|}$. The quantity $u_t^i(\mathbf{x}_t)_v$ then predicts the rate of probability change from the current token $x_t^i$ toward the token $v$, that still depends on the whole noisy sequence $\mathbf{x}_t$. The proposition in the toogle formalizes this assumption and its consequences.

Proposition

Proposition 4 (Factorized velocity).

Let $(X_t)$ be random variables of laws $p_t$, and assume that the transition probabilities factorize over tokens, namely for $\mathbf{x}, \mathbf{y} \in \mathcal{S}$:

$$\begin{align} \mathbb{P}(X_{t+h} = \mathbf{y} | X_t = \mathbf{x}) &:= \mathbb{P}(X_{t+h}^1 = y_1, X_{t+h}^2 = y_2, \dots, X_{t+h}^S = y_S | X_t = \mathbf{x}) \htmlData{equation-label=eq:factorized-path-1}{\htmlId{eq:factorized-path-1}{}} \\ &= \prod_{i=1}^S \mathbb{P}(X_{t+h}^i = y^i | X_t = \mathbf{x}) \htmlData{equation-label=eq:factorized-path-2}{\htmlId{eq:factorized-path-2}{}} \end{align}$$

Suppose that there is a velocity $u_t$ generating $(p_t)$. Then, the velocity writes as a sum of $S$ tokenwise probability velocities $u_t^i$, from $\mathcal{S}$ to $\mathcal{V}$:

$$\begin{equation} \htmlData{equation-label=eq:factorized-velocity}{\htmlId{eq:factorized-velocity}{}} u_t(\mathbf{x})_{\mathbf{y}} = \sum_{i=1}^S \mathbb{I}(\mathbf{y}^{\neq i} = \mathbf{x}^{\neq i}) u_t^i(\mathbf{x})_{y^i} \end{equation}$$

where $\mathbf{x}^{\neq i} \in \mathcal{V}^{S-1}$ denotes the sequence except $\mathbf{x}$ except its $i$-th element $x^i$.

Proof

We first formally define the token-wise velocity $u_t^i$:

$$\begin{align} \mathbb{P}(X_{t+h}^i = y^i | X_t = \mathbf{x}) &= \sum_{\mathbf{z} \in \mathcal{S}} \mathbb{P}(X_{t+h} = \mathbf{z} | X_t = \mathbf{x}) \mathbb{I}(z^i = y^i) \\ &= \sum_{\mathbf{z} \in \mathcal{S}} \left[\mathbb{I}(\mathbf{x} = \mathbf{z}) + h u_t(\mathbf{x})_{\mathbf{z}} + o(h)\right] \mathbb{I}(z^i = y^i) \\ &= \mathbb{I}(x^i = y^i) + h \underbrace{\sum_{\mathbf{z} \in \mathcal{S}} u_t(\mathbf{x})_{\mathbf{z}}\mathbb{I}(z^i = y^i)}_{=: u_t^i(\mathbf{x})_{y^i}} + o(h) \\ &= \mathbb{I}(x^i = y^i) + h u_t^i(\mathbf{x})_{y^i} + o(h). \htmlData{equation-label=eq:token-transition}{\htmlId{eq:token-transition}{}} \end{align}$$

Substituting this into our factorized joint probability, we get:

$$\begin{align} \mathbb{P}(X_{t+h} = \mathbf{y} | X_t = \mathbf{x}) &= \prod_{i=1}^S \left[ \mathbb{I}(y^i = x^i) + h u_t^i(\mathbf{x})_{y^i} + o(h) \right] \\ &\hspace{-8em}{}= \prod_{i=1}^S \mathbb{I}(y^i = x^i) + h \sum_{i=1}^S \left[ \mathbb{I}(\mathbf{y}^{\neq i} = \mathbf{x}^{\neq i}) u_t^i(\mathbf{x})_{y^i} \right] + o(h) \\ &\hspace{-8em}{}= \mathbb{I}(\mathbf{y} = \mathbf{x}) + h \sum_{i=1}^S \mathbb{I}(\mathbf{y}^{\neq i} = \mathbf{x}^{\neq i}) u_t^i(\mathbf{x})_{y^i} + o(h) \htmlData{equation-label=eq:factorized-path-2-3}{\htmlId{eq:factorized-path-2-3}{}} \end{align}$$

where $\mathbb{I}(\mathbf{y}^{\neq i} = \mathbf{x}^{\neq i})$ equals 1 if $\mathbf{y}$ and $\mathbf{x}$ differ at most at position $i$, and we used $\prod_{i=1}^S \mathbb{I}(y^i = x^i) = \mathbb{I}(\mathbf{y} = \mathbf{x})$.

By comparing this result with the definition of the velocity (Eq. (?)), we directly identify:

$$\begin{equation} \htmlData{equation-label=eq:factorized-velocity-2}{\htmlId{eq:factorized-velocity-2}{}} u_t(\mathbf{x})_{\mathbf{y}} = \sum_{i=1}^S \mathbb{I}(\mathbf{y}^{\neq i} = \mathbf{x}^{\neq i}) u_t^i(\mathbf{x})_{y^i} \end{equation}$$

Example 5.

To understand Eq. (?), take the vocabulary $\mathcal{V}=\{A,B\}$, sequence length $S=3$, and current sequence $\mathbf{x}=ABA$. We represent the vector $u_t(\mathbf{x})$ in $\mathbb{R}^8$. From Eq. (?), the sequence can only transition to 4 possible values (its current state, or state differing by a single token), and the velocity writes as:

ABA

State $\mathbf{y}$	Velocity $u_t(\mathbf{x})_{\mathbf{y}}$
AAA	$u_t^{\color{#dc2626}{2}}(\mathbf{x})_{\color{#dc2626}{A}}$
AAB	$0$
ABA	$u_t^{\color{#15803d}{1}}(\mathbf{x})_{\color{#15803d}{A}}+u_t^{\color{#dc2626}{2}}(\mathbf{x})_{\color{#dc2626}{B}}+u_t^{\color{#2563eb}{3}}(\mathbf{x})_{\color{#2563eb}{A}}$
ABB	$u_t^{\color{#2563eb}{3}}(\mathbf{x})_{\color{#2563eb}{B}}$
BAA	$0$
BAB	$0$
BBA	$u_t^{\color{#15803d}{1}}(\mathbf{x})_{\color{#15803d}{B}}$
BBB	$0$

While $u_t$ outputs a vector in $\mathbb{R}^8$, $u_t^1, u_t^2$ and $u_t^3$ output vectors in $\mathbb{R}^2$ (as $|\mathcal{V}| = 2$).

This modeling assumption has several important consequences:

• for infinitesimally small time steps, a sequence $\mathbf{x}$ can only transition to states $\mathbf{y}$ differing by at most one token (even if we draw all tokens simultaneously, the probability of two or more tokens changing is $o(h)$)
• learning-wise, it implies a substantial computational gain: instead of learning one transition object $u$ on the full sequence space $\mathcal{S}=\mathcal{V}^S$ (exponential in $S$), we can learn $S$ local transition objects $u^i$, each on the vocabulary space $\mathcal{V}$ (linear in $S$)
• at each sampling step, we draw each token independently following:

$$\begin{equation} \htmlData{equation-label=eq:token-sampling}{\htmlId{eq:token-sampling}{}} x_{t+h}^i \sim \text{Cat}\left(\mathbf{e}_{x_t^i} + h u_t^i(\mathbf{x}_t)\right) \end{equation}$$

The following algorithms and figure illustrates the analogy between one Euler step in continuous and discrete flow matching, and makes explicit what this token-wise velocity represents and how to use it to sample sequences.

SamplesStates

Continuous — an image x(t) ∈ R^d

x(t)

x(t+h) = x(t) + h u_t(x(t))deterministic

x(t+h)

Discrete — a sequence x_t ∈ S (here V = {A, B, m}, S = 4)

AmBm

x_t

x3t+h ∼ Cat(ex3t + h u3t(x_t))stochastic

AmBB

x_t+h

Now that we are equipped with adequate definitions of probability paths and velocities for the discrete case, together with a way to check that one is associated to the other, we can state the approach of discrete flow matching, which is in fact the same as for continuous flow matching:

1. Build an explicit probability path $(p_t)$ going from $p_0$ to $p_1$ (Eq. (?)) using a conditioning variable, conditional probability paths and conditional velocities.
2. Build a velocity $u_t$ generating $(p_t)$ (Eq. (?))
3. Provide an explicit loss to learn $u_t$ with a neural network (Eq. (?))

Building the probability path

Similarly to continuous flow matching, there are infinitely many possible choices for a probability path $(p_t)$ going from source $p_0$ to target $p_1$. To build a specific one, we borrow the core idea from continuous flow matching: we first choose a choose a conditioning variable $\mathbf{z} = \mathbf{x_1}$ and a conditional probability path $p_t(\cdot | \mathbf{x}_1)$ for which the exact conditional velocity $u_t(\cdot | \mathbf{x}_1)$ that generates this path is known and easy to compute. Then, we show that the resulting global velocity (obtained by marginalizing the conditional velocities) is guaranteed to generate a probability path that transports $p_0$ to $p_1$. All that’s left is to learn the global velocity as a simple conditional expectation from data, which neural networks know how to do!

In order to have factorized velocities (simpler to learn), we chose a conditional path $p_t(\cdot | \mathbf{x}_1)$ to be factorized:

$$\begin{align} & p_t(\mathbf{x} | \mathbf{x_1}) = \prod_{i=1}^S \big( \underbrace{ t \mathbb{I}(x^i = x_1^i) + (1 - t) \mathbb{I}(x^i = m) } _{:= q_t(x^j | x_1^j)} \big) \htmlData{equation-label=eq:mask-path-1}{\htmlId{eq:mask-path-1}{}} \\ \end{align}$$

This choice is similar to the linear interpolation path used in continuous flow matching $p_t(x | x_0, x_1) = t x_0 + (1 - t) x_1$, unless we cannot interpolate between two discrete tokens. Instead, we use linear interpolation in probability space: at time $t$, the $i$-th token $x^i_t$ equals the target token $x_1^i$ with probability $t$, or the mask token $m$ with probability $1 - t$. Therefore, $p_0(\cdot | \mathbf{x}_1)$ is a Dirac at the all-masked sequence, while $p_1(\cdot | \mathbf{x}_1)$ is a Dirac at $\mathbf{x}_1$. The interest of this conditional probability path is that we can derive the associated velocity in closed-form.

Note: compared to Eq. (?), one may be surprised that $u^i$ in Eq. (?) depends only on $x_i$ and not on $\mathbf{x}$. As the proof will show, this is a byproduct of the conditional probability paths factorizing over each token of the conditioning sequence $\mathbf{x}_1$.

Proof of proposition 6

We will proceed in two steps:

1. For the masked source, we show that the token-wise velocity $u_t(x^i | x_1^i)$ in Eq. (?) generates the token-wise conditional path $q_t(x^i | x_1^i)$ appearing in Eq. (?), in the sense that it satisfies the token-wise Kolmogorov equation

$$\begin{equation} \htmlData{equation-label=eq:token-kolmogorov}{\htmlId{eq:token-kolmogorov}{}} \frac{\mathrm{d}}{\mathrm{d}t} q_t(y^{i} | x_1^{i}) = \sum_{v \in \mathcal{V} } u_{t}(v | x_1^{i})_{y^{i}} \, q_t(v | x_1^{i}) \end{equation}$$

for all $y^i \in \mathcal{V}$ and all $i \in \{1, \dots, S\}$.

2. Once Eq. (?) holds, we show that the velocity

$$\begin{equation} u_t(\mathbf{x} | \mathbf{x}_1)_{\mathbf{y}} := \sum_{i=1}^S \mathbb{I}(\mathbf{y}^{\neq i} = \mathbf{x}^{\neq i}) u_t(x^i | x_1^i)_{y^i} \end{equation}$$

generates the conditional path $p_t(\mathbf{x} | \mathbf{x}_1)$, as per Eq. (?).

We start with the first statement. Since $q_t(y^i | x_1^i) := t \mathbb{I}(y^i = x_1^i) + (1 - t) \mathbb{I}(y^i = m)$, we have

$$\begin{equation} \htmlData{equation-label=eq:mask-path-derivative}{\htmlId{eq:mask-path-derivative}{}} \frac{\mathrm{d}}{\mathrm{d}t} q_t(y^{i} | x_1^{i}) = \mathbb{I}(y^{i} = x_1^{i}) - \mathbb{I}(y^{i} = m). \end{equation}$$

Equivalently, the velocity in Eq. (?) can be written as $u_{t}(v | x_1^{i})_{y^{i}} = \frac{1}{1 - t} \left( \mathbb{I}(y^{i} = x_1^{i}) - \mathbb{I}(y^{i} = m) \right) \mathbb{I}(v = m)$. Therefore, the right-hand side of Eq. (?) is

$$\begin{align} &\sum_{v \in \mathcal{V} } u_t(v | x_1^i)_{y^i} \, q_t(v | x_1^{i}) \htmlData{equation-label=eq:mask-velocity-rhs-1}{\htmlId{eq:mask-velocity-rhs-1}{}} \\ &= \sum_{v \in \mathcal{V} } \tfrac{1}{1 - t} \left( \mathbb{I}(y^{i} = x_1^{i}) - \mathbb{I}(y^{i} = m) \right) \mathbb{I}(v = m) \, q_t(v | x_1^{i}) \htmlData{equation-label=eq:mask-velocity-rhs-2}{\htmlId{eq:mask-velocity-rhs-2}{}} \\ &= \tfrac{1}{1 - t} \left( \mathbb{I}(y^{i} = x_1^{i}) - \mathbb{I}(y^{i} = m) \right) q_t(m | x_1^{i}) \htmlData{equation-label=eq:mask-velocity-rhs-3}{\htmlId{eq:mask-velocity-rhs-3}{}} \\ &= \mathbb{I}(y^{i} = x_1^{i}) - \mathbb{I}(y^{i} = m), \htmlData{equation-label=eq:mask-velocity-rhs-4}{\htmlId{eq:mask-velocity-rhs-4}{}} \end{align}$$

since $q_t(m | x_1^{i}) = 1-t$. This matches Eq. (?), so $(u_t(\cdot | x_1^i), q_t(\cdot | x_1^i))$ satisfies the token-wise Kolmogorov equation.

We now show that, using Eq. (?), the pair $(p_t(\cdot | \mathbf{x}_1), u_t(\cdot | \mathbf{x}_1))$ satisfies the Kolmogorov equation Eq. (?). Since the conditional path is factorized, we have:

$$\begin{equation} \htmlData{equation-label=eq:conditional-factorized-path}{\htmlId{eq:conditional-factorized-path}{}} p_t(\mathbf{y} | \mathbf{x}_1) = \prod_{i=1}^S q_t(y^i | x_1^i). \end{equation}$$

Differentiating this product gives:

$$\begin{align} \frac{\mathrm{d}}{\mathrm{d}t}p_t(\mathbf{y} | \mathbf{x}_1) &= \frac{\mathrm{d}}{\mathrm{d}t}\prod_{i=1}^S q_t(y^i | x_1^i) \htmlData{equation-label=eq:global-kolmogorov-proof-1}{\htmlId{eq:global-kolmogorov-proof-1}{}} \\ &= \sum_{i=1}^S \left(\frac{\mathrm{d}}{\mathrm{d}t}q_t(y^i | x_1^i)\right) \prod_{j\neq i} q_t(y^j | x_1^j). \htmlData{equation-label=eq:global-kolmogorov-proof-2}{\htmlId{eq:global-kolmogorov-proof-2}{}} \end{align}$$

Using the token-wise Kolmogorov equation Eq. (?), this becomes:

$$\begin{align} \frac{\mathrm{d}}{\mathrm{d}t}p_t(\mathbf{y} | \mathbf{x}_1) &= \sum_{i=1}^S \left[ \sum_{v\in\mathcal{V}} u_t(v | x_1^i)_{y^i}\,q_t(v | x_1^i) \right] \prod_{j\neq i} q_t(y^j | x_1^j). \htmlData{equation-label=eq:global-kolmogorov-proof-3}{\htmlId{eq:global-kolmogorov-proof-3}{}} \end{align}$$

We now compute the right-hand side of the global Kolmogorov equation using the factorized velocity Eq. (?):

$$\begin{align} &\sum_{\mathbf{x}\in\mathcal{S}} u_t(\mathbf{x} | \mathbf{x}_1)_{\mathbf{y}}\,p_t(\mathbf{x} | \mathbf{x}_1) \htmlData{equation-label=eq:global-kolmogorov-proof-4}{\htmlId{eq:global-kolmogorov-proof-4}{}} \\ &= \sum_{\mathbf{x}\in\mathcal{S}} \left[ \sum_{i=1}^S \mathbb{I}(\mathbf{y}^{\neq i}=\mathbf{x}^{\neq i})\,u_t(x^i | x_1^i)_{y^i} \right] \prod_{j=1}^S q_t(x^j | x_1^j) \htmlData{equation-label=eq:global-kolmogorov-proof-5}{\htmlId{eq:global-kolmogorov-proof-5}{}} \\ &= \sum_{i=1}^S \sum_{\mathbf{x}\in\mathcal{S}} \mathbb{I}(\mathbf{y}^{\neq i}=\mathbf{x}^{\neq i})\,u_t(x^i | x_1^i)_{y^i} \prod_{j=1}^S q_t(x^j | x_1^j). \htmlData{equation-label=eq:global-kolmogorov-proof-6}{\htmlId{eq:global-kolmogorov-proof-6}{}} \end{align}$$

For a fixed coordinate $i$, the indicator $\mathbb{I}(\mathbf{y}^{\neq i}=\mathbf{x}^{\neq i})$ forces $x^j=y^j$ for every $j\neq i$. Therefore, the sum over $\mathbf{x}$ reduces to a sum over the single free token $x^i=v$:

$$\begin{align} &\sum_{\mathbf{x}\in\mathcal{S}} \mathbb{I}(\mathbf{y}^{\neq i}=\mathbf{x}^{\neq i})\,u_t(x^i | x_1^i)_{y^i} \prod_{j=1}^S q_t(x^j | x_1^j) \htmlData{equation-label=eq:global-kolmogorov-proof-7}{\htmlId{eq:global-kolmogorov-proof-7}{}} \\ &= \sum_{v\in\mathcal{V}} u_t(v | x_1^i)_{y^i}\,q_t(v | x_1^i) \prod_{j\neq i} q_t(y^j | x_1^j). \htmlData{equation-label=eq:global-kolmogorov-proof-8}{\htmlId{eq:global-kolmogorov-proof-8}{}} \end{align}$$

Plugging this back into Eq. (?) gives exactly Eq. (?). Hence:

$$\begin{equation} \htmlData{equation-label=eq:global-kolmogorov-proof-9}{\htmlId{eq:global-kolmogorov-proof-9}{}} \frac{\mathrm{d}}{\mathrm{d}t}p_t(\mathbf{y} | \mathbf{x}_1) = \sum_{\mathbf{x}\in\mathcal{S}} u_t(\mathbf{x} | \mathbf{x}_1)_{\mathbf{y}}\,p_t(\mathbf{x} | \mathbf{x}_1). \end{equation}$$

The pair $(p_t(\cdot | \mathbf{x}_1), u_t(\cdot | \mathbf{x}_1))$ satisfies the Kolmogorov equation.

Values of conditional velocities for and uniform source and interpolation between an arbitrary $p_0$ and $p_1$

Source-Target Conditioning: Here, we condition on both the start and end points, so $z^i = (x_0^i, x_1^i)$. We choose $q_t$ as:

$$\begin{equation} \htmlData{equation-label=eq:pair-path}{\htmlId{eq:pair-path}{}} q_t(x^i | x_0^i, x_1^i) = t \mathbb{I}(x^i = x_1^i) + (1 - t) \mathbb{I}(x^i = x_0^i) \end{equation}$$

It means at time $t$, the $i$-th token $x_t^i$ equals the target token $x_1^i$ with probability $t$, or the source token $x_0^i$ with probability $1 - t$. Therefore, $p_0(\cdot | \mathbf{x}_0, \mathbf{x}_1)$ is a Dirac at $\mathbf{x}_0$, while $p_1(\cdot | \mathbf{x}_0, \mathbf{x}_1)$ is a Dirac at $\mathbf{x}_1$.

The conditional velocity field driving this transition is:

$$\begin{equation} \htmlData{equation-label=eq:pair-velocity}{\htmlId{eq:pair-velocity}{}} u_t(x^i | x_0^i, x_1^i)_v = \begin{cases} \tfrac{1}{1 - t} & \text{if } v = x_1^i \text{ and } x^i = x_0^i, \\ -\tfrac{1}{1 - t} & \text{if } v = x^i = x_0^i, \\ 0 & \text{otherwise}. \end{cases} \end{equation}$$

Here, mass moves directly from the source token $x_0^i$ toward the target token $x_1^i$. All rows except the source row are identically zero: once a token is already at state $x_1^i$, no transition occurs. Concretely, if $\mathcal{V} = \{A, B, C\}$, $x_0^i = A$, and $x_1^i = B$:

$$u_t(x^i | x_0^i = A,\, x_1^i = B)_v = \begin{array}{c|ccc} \begin{matrix} x^i \\[-2pt] \diagdown \\[-2pt] v \end{matrix} & A & B & C \\ \hline A & -\tfrac{1}{1-t} & \tfrac{1}{1-t} & 0 \\[0.3cm] B & 0 & 0 & 0 \\[0.3cm] C & 0 & 0 & 0 \end{array}$$

Uniform Source: In this case, we choose $q_t$ as:

$$\begin{equation} \htmlData{equation-label=eq:uniform-path}{\htmlId{eq:uniform-path}{}} q_t(x^i | x_1^i) = t \mathbb{I}(x^i = x_1^i) + (1 - t) \frac{1}{|\mathcal{V}|} \end{equation}$$

It means that at time $t$, the $i$-th token $x_t^i$ equals the target token $x_1^i$ with probability $t + (1 - t)\frac{1}{|\mathcal{V}|}$, while each non-target token in the vocabulary $\mathcal{V}$ has probability $(1 - t)\frac{1}{|\mathcal{V}|}$. Therefore, $p_0(\cdot | \mathbf{x}_1)$ is the uniform distribution over sequences, while $p_1(\cdot | \mathbf{x}_1)$ is a Dirac at $\mathbf{x}_1$.

The conditional velocity field driving this transition is:

$$\begin{equation} \htmlData{equation-label=eq:uniform-velocity}{\htmlId{eq:uniform-velocity}{}} u_t(x^i | x_1^i)_v = \begin{cases} \tfrac{1}{1-t} & \text{if } v = x_1^i \text{ and } x^i \neq x_1^i \\[0.3cm] -\tfrac{1}{1-t} & \text{if } v = x^i \neq x_1^i \\[0.3cm] 0 & \text{otherwise} \end{cases} \end{equation}$$

Unlike the mask source, every non-target token can transition toward $x_1^i$, since the uniform source places initial mass everywhere. Concretely, if $\mathcal{V} = \{A, B, C\}$ and $x_1^i = B$:

$$u_t(x^i | x_1^i = B)_v = \begin{array}{c|ccc} \begin{matrix} x^i \\[-2pt] \diagdown \\[-2pt] v \end{matrix} & A & B & C \\ \hline A & -\tfrac{1}{1-t} & \tfrac{1}{1-t} & 0 \\[0.1cm] B & 0 & 0 & 0 \\[0.1cm] C & 0 & \tfrac{1}{1-t} & -\tfrac{1}{1-t} \end{array}$$

Building the target velocity

Now that we have defined our conditional probability paths and velocities, we can use them to build global, unconditional ones. The overall distribution $p_t$ is obtained by marginalizing the conditional paths over target samples $\mathbf{x}_1$ drawn from $p_1$:

$$\begin{equation} \htmlData{equation-label=eq:marginal-path}{\htmlId{eq:marginal-path}{}} p_t(\mathbf{x}) := \sum_{\mathbf{x}_1} p_t(\mathbf{x} | \mathbf{x}_1) p_1(\mathbf{x}_1) \end{equation}$$

As the proof below shows, the path defined by marginalization indeed goes from $p_0$ to $p_1$.

Proof that path $(p_t)$ does go from $p_0$ to $p_1$

We verify the two boundary conditions separately:

At $t=0$: Recall the conditional path the mask source:

$$\begin{equation} q_0(x^i | x_1^i) = \mathbb{I}(x^i = m). \end{equation}$$

The value of $x_1^i$ plays no role, so $p_0(\mathbf{x} | \mathbf{x}_1)$ does not depend on $\mathbf{x}_1$:

$$\begin{equation} p_0(\mathbf{x} | \mathbf{x}_1) = \prod_{j=1}^S \mathbb{I}(x^j = m) = \delta_{m\ldots m}(\mathbf{x}) = p_0(\mathbf{x}). \end{equation}$$

Marginalizing over $\mathbf{x}_1$ then gives:

$$\begin{equation} p_{t=0}(\mathbf{x}) = \sum_{\mathbf{x}_1} p_0(\mathbf{x} | \mathbf{x}_1)\,p_1(\mathbf{x}_1) = p_0(\mathbf{x})\sum_{\mathbf{x}_1} p_1(\mathbf{x}_1) = p_0(\mathbf{x}). \end{equation}$$

At $t=1$: We have:

$$\begin{align} q_1(x^i | x_1^i) &= \mathbb{I}(x^i = x_1^i) \htmlData{equation-label=eq:endpoint-path-1}{\htmlId{eq:endpoint-path-1}{}} \\ p_1(\mathbf{x} | \mathbf{x}_1) &= \prod_{j=1}^S \mathbb{I}(x^j = x_1^j) = \mathbb{I}(\mathbf{x} = \mathbf{x}_1) \htmlData{equation-label=eq:endpoint-path-2}{\htmlId{eq:endpoint-path-2}{}} \end{align}$$

We conclude:

$$\begin{align} p_{t=1}(\mathbf{x}) &= \sum_{\mathbf{x}_1} p_1(\mathbf{x} | \mathbf{x}_1)\,p_1(\mathbf{x}_1) \\ &= \sum_{\mathbf{x}_1} \mathbb{I}(\mathbf{x} = \mathbf{x}_1)\, p_1(\mathbf{x}_1) \\ &= p_1(\mathbf{x}). \end{align}$$

We now arrive to the heart of discrete flow matching: building a velocity that transports $p_0$ to $p_1$.

Proof

We want to show

$$\begin{equation} \htmlData{equation-label=eq:marginal-velocity}{\htmlId{eq:marginal-velocity}{}} \mathbb{P}(X_{t+h}=\mathbf{y}| X_t=\mathbf{x}) = \mathbb{I}(\mathbf{x}=\mathbf{y}) + h\,u_t(\mathbf{x})_{\mathbf{y}} + o(h). \end{equation}$$

Introducing $X_1$ and using the law of total probability,

$$\begin{align} \mathbb{P}(X_{t+h}=\mathbf{y}| X_t=\mathbf{x}) \htmlData{equation-label=eq:marginal-velocity-proof-1}{\htmlId{eq:marginal-velocity-proof-1}{}} &= \sum_{\mathbf{x}_1} \mathbb{P}(X_{t+h}=\mathbf{y}, X_1=\mathbf{x}_1| X_t=\mathbf{x}) \htmlData{equation-label=eq:marginal-velocity-proof-2}{\htmlId{eq:marginal-velocity-proof-2}{}} \\ &\hspace{-8em}{}= \sum_{\mathbf{x}_1} \mathbb{P}(X_{t+h}=\mathbf{y}| X_t=\mathbf{x}, X_1=\mathbf{x}_1)\,\mathbb{P}(X_1=\mathbf{x}_1| X_t=\mathbf{x}) \htmlData{equation-label=eq:marginal-velocity-proof-3}{\htmlId{eq:marginal-velocity-proof-3}{}} \\ &\hspace{-8em}{}= \sum_{\mathbf{x}_1}\Big(\mathbb{I}(\mathbf{x}=\mathbf{y}) + h\,u_t(\mathbf{x}|\mathbf{x}_1)_{\mathbf{y}} + o(h)\Big)\,\mathbb{P}(X_1=\mathbf{x}_1| X_t=\mathbf{x}) \htmlData{equation-label=eq:marginal-velocity-proof-4}{\htmlId{eq:marginal-velocity-proof-4}{}} \\ &\hspace{-8em}{}= \mathbb{I}(\mathbf{x}=\mathbf{y}) + h\sum_{\mathbf{x}_1}u_t(\mathbf{x}|\mathbf{x}_1)_{\mathbf{y}}\,\mathbb{P}(X_1=\mathbf{x}_1| X_t=\mathbf{x}) + o(h) \htmlData{equation-label=eq:marginal-velocity-proof-5}{\htmlId{eq:marginal-velocity-proof-5}{}} \\ &\hspace{-8em}{}= \mathbb{I}(\mathbf{x}=\mathbf{y}) + h\,u_t(\mathbf{x})_{\mathbf{y}} + o(h). \htmlData{equation-label=eq:marginal-velocity-proof-6}{\htmlId{eq:marginal-velocity-proof-6}{}} \end{align}$$

where we have used in Eq. (?) the fact that $u_t(\cdot | \mathbf{x}_1)$ generates $p_t(\cdot | \mathbf{x}_1)$.

Alternative proof: apply the Kolmogorov equation to each conditional process and average over $\mathbf{x}_1$:

$$\begin{align} \frac{\mathrm{d}}{\mathrm{d}t}p_t(\mathbf{y}) \htmlData{equation-label=eq:marginal-kolmogorov-1}{\htmlId{eq:marginal-kolmogorov-1}{}} &= \sum_{\mathbf{x}_1} p_1(\mathbf{x}_1)\,\frac{\mathrm{d}}{\mathrm{d}t}p_t(\mathbf{y}|\mathbf{x}_1) \htmlData{equation-label=eq:marginal-kolmogorov-2}{\htmlId{eq:marginal-kolmogorov-2}{}} \\ &= \sum_{\mathbf{x}_1}p_1(\mathbf{x}_1)\sum_{\mathbf{x}}u_t(\mathbf{x}|\mathbf{x}_1)_{\mathbf{y}}\,p_t(\mathbf{x}|\mathbf{x}_1) \htmlData{equation-label=eq:marginal-kolmogorov-3}{\htmlId{eq:marginal-kolmogorov-3}{}} \\ &= \sum_{\mathbf{x}}\left(\sum_{\mathbf{x}_1}u_t(\mathbf{x}|\mathbf{x}_1)_{\mathbf{y}}\,\mathbb{P}(X_1=\mathbf{x}_1| X_t=\mathbf{x})\right)p_t(\mathbf{x}) \htmlData{equation-label=eq:marginal-kolmogorov-4}{\htmlId{eq:marginal-kolmogorov-4}{}} \\ &= \sum_{\mathbf{x}} u_t(\mathbf{x})_{\mathbf{y}}\,p_t(\mathbf{x}). \htmlData{equation-label=eq:marginal-kolmogorov-5}{\htmlId{eq:marginal-kolmogorov-5}{}} \end{align}$$

So $u_t$ satisfies the Kolmogorov equation for $p_t$ and therefore generates $p_t$.

This theorem follows the exact same logic as in continuous FM, and has the same practical consequence: since the target field is identified as a conditional expectation, it can be learned through a regression problem.

In continuous spaces, this is typically introduced via the $L^2$ projection view, as the conditional expectation solves

$$\begin{equation} \htmlData{equation-label=eq:l2-proj}{\htmlId{eq:l2-proj}{}} u_t(\mathbf{x}) = \underset{u_t \text{ measurable}}{\operatorname{argmin}}\ \mathbb{E}\!\left[\left\|u_t(X_t | X_1) - u(X_t)\right\|^2 \right] . \end{equation}$$

In discrete spaces, since we compare nonneg­ative rate vectors that behave like local probabilities, a Kullback-Leibler style objective is more natural. More generally, this conditional expectation is the minimizer for any Bregman divergence, and KL is the Bregman choice adapted to our setting:

$$\begin{equation} \htmlData{equation-label=eq:kl-proj}{\htmlId{eq:kl-proj}{}} u_t(\mathbf{x}) ={} \underset{u}{\operatorname{argmin}}\ \mathbb{E}\!\left[ \mathrm{KL} \left(u(X_t | X_1), u(X_t)\right) \right] \end{equation}$$

The exact conditional expectation of Eq. (?) is approximated by a neural network $u_{\theta,t}$, trained by minimizing over $\theta$:

Learning with denoiser parameterization

With straight interpolation paths $x_t := (1-t) x_0 + tx_1$, learning the ideal velocity $\mathbb{E}[X_1 - X_0 | X_t=x]$ is mathematically equivalent to learning the ideal denoiser $\mathbb{E}[X_1 | X_t = x]$, since the two are linked via $\mathbb{E}[X_1 - X_0 | X_t=x] = \frac{1}{1 - t} ( \mathbb{E}[X_1 | X_t = x] - x)$. The denoiser can be understood as predicting the most plausible clean image $x_1$ given a noisy image $x_t$ at time $t$. It turns out the same type of denoiser-velocity equivalence holds in the discrete case:

$$\begin{equation} \htmlData{equation-label=eq:velocity-posterior}{\htmlId{eq:velocity-posterior}{}} u_t^i(\mathbf{x})_v = \frac{1}{1 - t} \left( \mathbb{P}(X_1^i = v | X_t = \mathbf{x}) - \mathbb{I}(v = x^i) \right) . \end{equation}$$

In this equation, the function $\mathbf{x} \mapsto \mathbb{P}(X_1^i = v | X_t = \mathbf{x})$ can be interpreted as a denoiser: it gives the probability of the $i$-th token in the clean sequence $\mathbf{x}_1$ being $v$, given the full noisy sequence $\mathbf{x}_t$.

Proof of velocity-denoiser equivalence

We write the proof for the masked source and start from Eq. (?). In a compact form, it can be rewritten as:

$$\begin{equation} \htmlData{equation-label=eq:mask-velocity-token}{\htmlId{eq:mask-velocity-token}{}} u_t^i(x^i | x_1^i)_v = \frac{1}{1-t}\,\mathbb{I}(x^i = m)\left(\mathbb{I}(v = x_1^i) - \mathbb{I}(v = m)\right). \end{equation}$$

Case 1: $x^i \neq m$ (token already unmasked).

The indicator $\mathbb{I}(x^i = m)=0$, so $u_t^i(x^i | x_1^i)_v = 0$ for all $x_1^i, v$ and hence $u_t^i(x)_v = 0$.

It remains to check that Eq. (?) also gives zero. By Bayes’ rule,

$$\begin{equation} \htmlData{equation-label=eq:posterior-bayes}{\htmlId{eq:posterior-bayes}{}} \mathbb{P}(X_1^i = v | X_t = \mathbf{x}) = \frac{\mathbb{P}(X_t = \mathbf{x} | X_1^i = v)\,\mathbb{P}(X_1^i = v)} {\mathbb{P}(X_t = \mathbf{x})}. \end{equation}$$

Under the conditional probability path Eq. (?),

$$\begin{equation} \htmlData{equation-label=eq:mask-token-path}{\htmlId{eq:mask-token-path}{}} q_t(x^i | x_1^i) = t\,\mathbb{I}(x^i = x_1^i) + \underbrace{(1-t)\,\mathbb{I}(x^i = m)}_{=0} = t\,\mathbb{I}(x^i = x_1^i) \end{equation}$$

since $x^i \neq m$. Therefore, if $v \neq x^i$, then

$$\begin{align} \mathbb{P}(X_t = \mathbf{x} |X_1^i = v) \htmlData{equation-label=eq:posterior-zero-1}{\htmlId{eq:posterior-zero-1}{}} &= \htmlData{equation-label=eq:posterior-zero-2}{\htmlId{eq:posterior-zero-2}{}} q_t(x^i | v)\,\mathbb{P}(X_t^{\neq i} = \mathbf{x}^{\neq i} |X_1^i = v) \htmlData{equation-label=eq:posterior-zero-3}{\htmlId{eq:posterior-zero-3}{}} \\ &= 0. \htmlData{equation-label=eq:posterior-zero-4}{\htmlId{eq:posterior-zero-4}{}} \end{align}$$

and hence $\mathbb{P}(X_1^i = v |X_t = \mathbf{x})=0$. Since the conditional probabilities sum to one over $v$ , we conclude that

$$\begin{equation} \htmlData{equation-label=eq:posterior-copy}{\htmlId{eq:posterior-copy}{}} \mathbb{P}(X_1^i = v |X_t = \mathbf{x})=\mathbb{I}(v = x^i). \end{equation}$$

Substituting Eq. (?) into Eq. (?) gives

$$\begin{equation} \htmlData{equation-label=eq:velocity-zero}{\htmlId{eq:velocity-zero}{}} \frac{1}{1-t} \left( \mathbb{P}(X_1^i = v |X_t = \mathbf{x}) - \mathbb{I}(v = x^i) \right) = 0. \end{equation}$$

Thus both expressions coincide in the case $x^i \neq m$.

Case 2: $x^i = m$ (token is masked). Taking the expectation over $x_1^i$:

$$\begin{align} u_t^i(\mathbf{x})_v \htmlData{equation-label=eq:velocity-posterior-proof-1}{\htmlId{eq:velocity-posterior-proof-1}{}} &= \mathbb{E}_{x_1^i}\!\left[\frac{1}{1-t}\left(\mathbb{I}(v = x_1^i) - \mathbb{I}(v = m)\right)\right] \htmlData{equation-label=eq:velocity-posterior-proof-2}{\htmlId{eq:velocity-posterior-proof-2}{}} \\ &= \frac{1}{1-t}\left(\mathbb{P}(X_1^i = v | X_t = \mathbf{x}) - \mathbb{I}(v = m)\right) \htmlData{equation-label=eq:velocity-posterior-proof-3}{\htmlId{eq:velocity-posterior-proof-3}{}} \\ &= \frac{1}{1-t}\left(\mathbb{P}(X_1^i = v | X_t = \mathbf{x}) - \mathbb{I}(v = x^i)\right), \htmlData{equation-label=eq:velocity-posterior-proof-4}{\htmlId{eq:velocity-posterior-proof-4}{}} \end{align}$$

where the last equality uses $x^i = m$. Both cases match the claimed formula.

Because the term $\mathbb{I}(v = x^i)$ is a known constant given the current state, learning the velocity field $u$ is mathematically exactly equivalent to learning the denoiser probability $\mathbb{P}(X_1^i = v | X_t = \mathbf{x})$. Consequently, we can train a denoiser $P_{\theta, t}$ directly. Using the KL divergence between the true clean token distribution and the learned one, the loss simplifies to the standard cross-entropy [2, 3]:

Derivation

By the same argument as in Eq. (?), the true denoiser is the KL projection of the conditional clean-token distribution: $P_t(X_1^i | X_t = \mathbf{x})$ minimizes in $P$

$$\begin{equation} \htmlData{equation-label=eq:denoiser-kl-proj}{\htmlId{eq:denoiser-kl-proj}{}} \mathbb{E}_{\substack{\mathbf{x}_1 \sim p_1 \\ \mathbf{x} \sim p_t(\cdot | \mathbf{x}_1)}} \!\left[ \mathrm{KL}\left(P_t(X_1^i | X_t = \mathbf{x},\, \mathbf{x}_1),\, P(X_1^i | X_t = \mathbf{x})\right) \right] \end{equation}$$

Thus, we want to minimize:

$$\begin{equation} \htmlData{equation-label=eq:posterior-kl-cond}{\htmlId{eq:posterior-kl-cond}{}} L(\theta) \equiv \mathbb{E}_{\substack{t \sim \mathcal{U}(0,1) \\ \mathbf{x}_1 \sim p_1 \\ \mathbf{x} \sim p_t( \cdot | \mathbf{x}_1)}} \left[ \sum_{i=1}^S \mathrm{KL}\left(P_t(X_1^i | X_t = \mathbf{x},\, \mathbf{x}_1),\, P_{\theta,t}(X_1^i | X_t = \mathbf{x})\right) \right] \end{equation}$$

Given $\mathbf{x}_1$, the true conditional posterior is a one-hot: $P_t(X_1^i | X_t = \mathbf{x}, \mathbf{x}_1) = \mathbf{e}_{x_1^i}$. The KL between a one-hot and any distribution reduces to a single log term:

$$\begin{equation} \htmlData{equation-label=eq:onehot-kl}{\htmlId{eq:onehot-kl}{}} \mathrm{KL}\left(\mathbf{e}_{x_1^i},\, P_{\theta,t}(X_1^i | \mathbf{x})\right) = -\log P_{\theta,t}(X_1^i = x_1^i | \mathbf{x}) \end{equation}$$

which gives directly:

$$\begin{equation} \htmlData{equation-label=eq:cdm-loss-2}{\htmlId{eq:cdm-loss-2}{}} L_{\text{CDM}}(\theta) = -\mathbb{E}_{\substack{t \sim \mathcal{U}(0,1) \\ \mathbf{x}_1 \sim p_1 \\ \mathbf{x} \sim p_t( \cdot | \mathbf{x}_1)}} \left[ \sum_{i=1}^S \log P_{\theta,t}(X_1^i = x_1^i | X_t = \mathbf{x}) \right] \end{equation}$$

Discrete Flow Matching Sampling

S=4, 𝒱={A, B, m}, mask source, 10 Euler steps (dt=0.100) · sample #1

Velocity u_tDenoiser P(x₁|x_t)

Step 1/10t = 0.00

m

pos 1

m

pos 2

m

pos 3

m

pos 4

velocity u_t(x_t)

pos 1

pos 2

pos 3

pos 4

e_xit + h uit(x_t)

pos 1

pos 2

pos 3

pos 4

← Prev▶ PlayNext →Reset🎲 New Sample

steps

seed

slowfast

A

B

m (mask)

To wrap up, we end this section with a table summarizing the important steps in discrete flow matching, put in parallel with their equivalent for continuous flow matching.

Table 7: Continuous and discrete flow matching side by side.

Continuous flow matching	Discrete flow matching
1. Define a conditioning variable $\mathbf{z}$ with known distribution $\pi(\mathbf{z})$
$\mathbf{z}=(\mathbf{x}_0, \mathbf{x}_1)$ with $\pi(\mathbf{z}) = p_0 \times p_1$	$\mathbf{z}=\mathbf{x}_1$ with $\pi(\mathbf{z}) = p_1$ (our choice for simplicity; $\mathbf{z} = (\mathbf{x}_0, \mathbf{x}_1)$ is also possible)
2. Define a simple conditional distribution $p_t(\mathbf{x} | \mathbf{z})$
$p_t(\mathbf{x} | \mathbf{z}) = \delta_{(1 - t)\mathbf{x}_0 + t\mathbf{x}_1}(\mathbf{x})$	$p_t(\mathbf{x} | \mathbf{z}) = \prod_{j=1}^S q_t(x^j | x_1^j)$ where $q_t = \text{Cat}((1 - t) \mathbf{e}_{m} + t \mathbf{e}_{x_1^j} )$
3. Compute an associated conditional velocity field $u_t(\mathbf{x} | \mathbf{z})$
$u_t(\mathbf{x} | \mathbf{z}) = \mathbf{x}_1 - \mathbf{x}_0$	$u_t(\mathbf{x} | \mathbf{z})_{\mathbf{y}} = \sum_{i=1}^S \mathbb{I}(\mathbf{y}^{\neq i} = \mathbf{x}^{\neq i})\, u_t(x^i | x_1^i)_{y^i}$ with $u_t(x^i | x_1^i) = \frac{1}{1-t}\mathbb{I}(x^i = m)\left(\mathbf{e}_{x_1^i} - \mathbf{e}_{m}\right)$
4. Training objective
Use $L^2$: $\min_\theta \mathbb{E}_{\substack{t \sim \mathcal{U}_{[0,1]} \\ \mathbf{z} \sim \pi \\ \mathbf{x}_t \sim p_t(\cdot | \mathbf{z})}}\!\left[\Vert u_{\theta, t}(\mathbf{x}_t) - (\mathbf{x}_1 - \mathbf{x}_0) \Vert^2\right]$	Use KL: $\min_\theta \mathbb{E}_{\substack{t \sim \mathcal{U}_{[0,1]} \\ \mathbf{x}_1 \sim p_1 \\ \mathbf{x}_t \sim p_t(\cdot | \mathbf{x}_1)}}\!\left[-\sum_{i=1}^S \log P_{\theta,t}(X_1^i = x_1^i | \mathbf{x}_t)\right]$
5. Sampling Sample $\mathbf{x}_0 \sim p_0$, integration scheme for $k=0,\ldots,N-1$
Numerical Integration (Euler ODE): $\mathbf{x}_{k+1} = \mathbf{x}_k + \frac{1}{N} u_{\theta,k/N}(\mathbf{x}_k)$ If learned denoiser, use: $u_{\theta,t}(\mathbf{x}) = \frac{1}{1-t}\!\left(\mathbb{E}_\theta[X_1 | X_t = \mathbf{x}] - \mathbf{x}\right)$	Continuous Time Markov Chain: $x_{k+1}^i \sim \text{Cat}\!\left(\mathbf{e}_{x_k^i} + \frac{1}{N} u_{\theta,k/N}^i(\mathbf{x}_k)\right)$ If learned denoiser, use: $u_{\theta,t}^i(\mathbf{x}) = \frac{1}{1-t}\!\left(P_{\theta,t}(X_1^i | \mathbf{x}) - \mathbf{e}_{x^i}\right)$
6. Master Equation
Continuity Equation: $\partial_t p_t + \nabla (u_t p_t) = 0$	Kolmogorov Equation: $\frac{\mathrm{d}}{\mathrm{d}t} p_t(\mathbf{y}) = \sum_{\mathbf{x} \in \mathcal{S}} u_t(\mathbf{x})_{\mathbf{y}} \, p_t(\mathbf{x})$

To go further: planners in masked flow-matching sampling

Notes before reading this section:

1. We note the denoiser $D_{\theta,t}(\mathbf{x}_t) := P_{\theta,t}(X_1 | X_t = \mathbf{x}_t)$.
2. Because the noise level $t$ in masked flow-matching can be easily inferred from the number of masked tokens in $\mathbf{x}_t$, the denoiser does not need to know the current time in its predictions. Therefore, for any time $t$ and any given sequence $\mathbf{x}$, we write $D_{\theta,t}(\mathbf{x}) = D_{\theta}(\mathbf{x})$.

Sampling with a planner

With step size $h$, vanilla masked flow-matching sampling Eq. (?) uses $1/h$ time steps. Across a sequence of length $S$, this corresponds to $S/h$ draws from probability vectors over the vocabulary $|\mathcal{V}|$. For $h = 0.01$ and $S = 200$, generating one sample therefore involves $20,000$ such draws, which is computationally expensive. Moreover, some steps $t \to t+h$ can keep the current sequence $\mathbf{x}_t$ unchanged, implying computational waste (it can be easily observed in figure 6).

To make generation more efficient, the standard approach used for large masked flow-matching models is planner-guided sampling. Instead of relying on infinitesimal time steps, we fix the total generation steps $N$ and directly use the denoiser to reveal a fixed number of tokens at each step. For $S = 200$ and $N = 50$, we have $NS = 50 \times 200 = 10,000$ draws to do, which is half the cost of the previous example.

In other words, the planner indicates, for each position $i$ in the noisy sequence $\mathbf{x}_k$, the probability of denoising $x_k^i$ into $z^i$. If $x_k^i \neq m$, this probability is zero. For example, if $\mathbf{x}_k=(m,A,m)$ and $\mathbf{z}=(B,A,A)$, then a possible value of $G_\phi(\mathbf{z}, \mathbf{x}_k)$ is $(0.7,0,0.3)$, which means revealing position $1$ into $z^1 = B$ with probability $70\%$ and position $3$ into $z^3 = A$ with probability $30\%$.

However, we rarely decide to unmask tokens by picking a random position at each step. Instead, we can use heuristic or learned planners that rank positions according to how reliable or informative they seem to reveal.

A simple and widely used strategy in this direction is the soft greedy planner. Instead of unmasking tokens at random, it selects the positions where the denoiser is most confident in its predictions. We then relax this strict selection into a soft probability distribution using a softmax function with a temperature parameter $\tau$:

$$\begin{equation} \htmlData{equation-label=eq:planner}{\htmlId{eq:planner}{}} G_\tau(\mathbf{z}, \mathbf{x}_k)_t \propto \exp\left(\frac{1}{\tau} \log D_\theta^j(\mathbf{x}_k)_{z^j}\right) \end{equation}$$

As $\tau \to \infty$, this reduces to uniform random sampling. As $\tau \to 0$, it converges to a strict, hard selection of only the most confident token in the sampled sequence $\mathbf{z}$.

Other planning strategies

Entropy planner. Another heuristic is to score each position by the negative entropy of the denoiser distribution. Low entropy means the denoiser is concentrated on a few tokens, so the position is treated as more reliable:

$$\begin{equation} s^i(\mathbf{x}_k) = \sum_{v \in \mathcal{V}} D_\theta^i(\mathbf{x}_k)_v \log D_\theta^i(\mathbf{x}_k)_v. \end{equation}$$

The planner then applies a softmax over the still-masked positions:

$$\begin{equation} \htmlData{equation-label=eq:entropy-planner}{\htmlId{eq:entropy-planner}{}} G_\tau^{\mathrm{ent}}(\mathbf{z}, \mathbf{x}_k)_i \propto \exp\left(\frac{s^i(\mathbf{x}_k)}{\tau}\right) \end{equation}$$

Unlike soft greedy, this planner does not use the sampled token $z^i$ directly; it only uses the shape of the denoiser distribution at position $i$.

Learned planner. With a pretrained denoiser, one can also train a separate planner offline. This can be done with a cross-entropy objective of the form:

$$\begin{equation} \mathcal{L}_{G}(\phi) = \mathbb{E}\!\left[ \sum_{i=1}^{S} \mathbb{I}(x_k^i=m) \frac{1}{1-t} \mathrm{CE}\!\left(G_\phi^i(\mathbf{x}_k,\mathbf{z}),\, \mathbb{I}[z^i=x_1^i]\right) \right]. \end{equation}$$

The learned planner can then be interpreted as a consistency scorer: positions whose sampled token looks reliable under the current partially denoised context receive more planner mass.

Example: one sampling step with a soft greedy planner

To illustrate one sampling step, take the vocabulary $\mathcal{V}=\{A,B,m\}$ and sequences of length $S=3$. Suppose that after one planner step ($k=1$), the current sequence is

$$\begin{equation} \htmlData{equation-label=eq:planner-example-current}{\htmlId{eq:planner-example-current}{}} \mathbf{x}_1 = (A,m,m), \end{equation}$$

so positions $2$ and $3$ are still masked. The denoiser proposes a full candidate sequence by sampling token-wise:

$$\begin{equation} \htmlData{equation-label=eq:planner-example-candidate}{\htmlId{eq:planner-example-candidate}{}} \mathbf{z} \sim D_\theta(\mathbf{x}_1), \qquad \text{for example } \mathbf{z}=(A,A,B). \end{equation}$$

For this sampled candidate, the planner only looks at the denoiser confidence assigned to the proposed token at each still-masked position, and assigns zero mass to positions that are already unmasked.

$$\begin{align} \widetilde{G}(\mathbf{z}, \mathbf{x}_1) \htmlData{equation-label=eq:planner-example-confidence-1}{\htmlId{eq:planner-example-confidence-1}{}} &:= \begin{pmatrix} 0 \\ D_\theta^2(\mathbf{x}_1)_{z^2 = A} \\ D_\theta^3(\mathbf{x}_1)_{z^3 = B} \end{pmatrix} \htmlData{equation-label=eq:planner-example-confidence-2}{\htmlId{eq:planner-example-confidence-2}{}} \\ &= \begin{pmatrix} 0 \\ 0.6 \\ 0.7 \end{pmatrix}. \htmlData{equation-label=eq:planner-example-confidence-3}{\htmlId{eq:planner-example-confidence-3}{}} \end{align}$$

Setting the temperature to $\tau = 1/3$, the planner scales these scores by a power of $1/\tau = 3$ and normalizes them:

$$\begin{align} G_{1/3}(\mathbf{z}, \mathbf{x}_1) \htmlData{equation-label=eq:planner-example-prob-1}{\htmlId{eq:planner-example-prob-1}{}} &= \frac{1}{0.6^3 + 0.7^3} \begin{pmatrix} 0 \\ 0.6^3 \\ 0.7^3 \end{pmatrix} \htmlData{equation-label=eq:planner-example-prob-2}{\htmlId{eq:planner-example-prob-2}{}} \\ &\approx \begin{pmatrix} 0 \\ 0.386 \\ 0.614 \end{pmatrix}. \htmlData{equation-label=eq:planner-example-prob-3}{\htmlId{eq:planner-example-prob-3}{}} \end{align}$$

This means that the planner is more likely to reveal position $3$, because the sampled token $z^3=B$ received higher denoiser confidence than $z^2=A$. If we take the limit cases for the temperature, we get:

$$\begin{align} G_{0}(\mathbf{z}, \mathbf{x}_1) &= \begin{pmatrix} 0 \\ 0 \\ 1 \end{pmatrix}, \qquad G_{\infty}(\mathbf{z}, \mathbf{x}_1) &= \begin{pmatrix} 0 \\ \frac{1}{2} \\ \frac{1}{2} \end{pmatrix}. \end{align}$$

If the planner samples position $3$, then the next sequence is obtained by copying only that proposed token:

$$\begin{equation} \htmlData{equation-label=eq:planner-example-next}{\htmlId{eq:planner-example-next}{}} \mathbf{x}_2 = (A,m,B). \end{equation}$$

The figure below illustrates the sampling with a planner.

Discrete Flow Matching Sampling

S=4, 𝒱={A, B, m}, mask source, 4 planner steps (dt=0.250) · sample #1

Vanilla (Euler)Greedy Planner

Step 1/4t = 0.25

m

pos 1

m

pos 2

m

pos 3

m

pos 4

P(x₁|x_t)

pos 1

pos 2

pos 3

pos 4

Planner G_φdemasking 1 position (S/N = 4/4)

A

z₁

A

z₂

A

z₃

A

z₄

G_φ(τ=1) selection prob.

← Prev▶ PlayNext →Reset🎲 New Sample

steps(max S)

seed

τ

slowfast

A

B

m (mask)

Exposure bias

Introducing a planner-aware sampling process creates a mismatch between training and inference. During training, the standard objective masks tokens uniformly at random across all positions and orderings. During inference, the planner reveals tokens along a specific, confidence-guided trajectory.

This difference matters because uniform masking can put training weight on mask patterns that the planner would almost never visit in practice. For example, in a structured sequence where the last token is only predictable after the first two have been revealed, uniform masking may still train the model to predict that last token from $(m,m,m)$ or from $(m,m,B)$, while a confidence-based planner would first reveal easier positions and avoid asking for that token with so little context. These hard denoising tasks can inject extra error into the learned denoiser, and those errors can then be amplified during sampling when early wrong predictions are irreversible.

A fix is to reweight the denoising loss Eq. (?) toward the positions the planner would select. This gives the Planner-Aware Path Learning (PAPL) objective [4]:

Here, $w^i = G_\phi(\mathbf{z}, \mathbf{x}_t)_i$ is the normalized planner weight for position $i$, and $\alpha$ controls the strength of this guidance. The stop-gradient $\operatorname{sg}[\cdot]$ detaches these weights: for example, soft greedy weights depend on $\theta$ through $\operatorname{softmax}_i\left(\frac{1}{\tau}\log P_{\theta,t}(X_1^i = z^i | X_t = \mathbf{x}_t)\right)$, but gradients are not propagated through this dependence.

References

Notes