---
title: Flow Marching Transformer for PDEs
url: https://www.emergentmind.com/topics/flow-marching-transformer-fmt
type: topic
---

# Flow Marching Transformer for PDEs

Searching arXiv for the specified paper and closely related work to ground the article.
arxiv_search: query="2509.18611 OR \"Flow marching for a generative PDE foundation model\""

Flow Marching Transformer (FMT) is the core neural backbone introduced in “Flow marching for a generative PDE foundation model” for generative modeling of PDE-governed spatiotemporal dynamics [2509.18611]. Within the broader Flow Marching framework, FMT implements the conditional flow-marching velocity field $g_\theta(x,t,h)$ that transports a noisy latent state toward its clean successor, thereby coupling neural operator learning with flow matching. The model operates on latent embeddings produced by a Physics-Pretrained Variational Autoencoder (P2VAE), combines a diffusion-forcing history filter with a multi-scale latent temporal pyramid, and is trained on a corpus of approximately $2.5\text{M}$ trajectories spanning 12 distinct PDE families. Its stated purpose is to reduce long-term rollout drift while enabling uncertainty-aware ensemble generation, in contrast to deterministic Transformer-based PDE foundation models that lack generative flexibility [2509.18611].

## 1. Functional role in the Flow Marching framework

FMT occupies the predictive core of the Flow Marching pipeline. The P2VAE first maps each physical frame $x_s \in \mathbb{R}^{H \times W \times C}$ to a latent grid $y_s \in \mathbb{R}^{H' \times W' \times C'}$. For a block of four consecutive times $\{s,s+1,s+2,s+3\}$, noisy latent states are constructed, passed through a small GRU or RNN with parameters $\phi$, and summarized into history codes $h_i$ that serve as the “diffusion-forcing” state. FMT then processes the multi-scale latent tokens together with temporal conditioning and the history summary to predict a preconditioned residual of the form
$$(1-t)\,g_\theta(x_{s,t}^k,t,h_{s-1}) \approx x_{s+1} - x_{s,t}^k.$$

The model is therefore not a direct next-step regressor in the deterministic neural-operator sense. Instead, it learns a unified velocity field over jointly sampled physical time and noise level. In the formulation given for FMT, the transformer stack fuses noisy latent observations from several recent time points, while the RNN state provides a learned summary of prior dynamics. A plausible implication is that the model treats short-term state evolution and uncertainty transport as parts of a single conditional generative process rather than as separate forecasting and noise-injection stages [2509.18611].

## 2. Flow-marching objective and diffusion-forcing formulation

For each physical step $s \rightarrow s+1$, Flow Marching samples
$$t \sim \mathrm{Uniform}(0,1), \qquad k \sim \mathrm{Uniform}(0,1), \qquad z \sim \mathcal{N}(0,I),$$
and constructs the noisy state
$$x_{s,t}^k = x_s + t(x_{s+1}-x_s) - (1-t)(1-k)(x_s-z).$$
After P2VAE encoding, the analogous expression is used in latent space for $y_{s,t}^k$.

The instantaneous ground-truth velocity is defined as
$$u_t^k = \frac{d}{dt}x_t^k = \frac{x_{s+1}-x_t^k}{1-t},$$
and the paper states the decomposition
$$u_t^k = (x_{s+1} - kx_s) + (1-t)(1-k)^2 \nabla_x \log q_t^k(x_t^k).$$
To avoid feeding $k$ into the network, the training rule adopts the simpler frame-interpolation target
$$u_t := \frac{x_{s+1}-x_{s,t}^k}{1-t}.$$

Preconditioning near $t \rightarrow 1$ yields the flow-marching objective
$$
L_{\mathrm{FM}} = \frac{1}{2}\,\mathbb{E}_{(s,t,k,z)}
\left[
\left\|
(1-t)\,g_\theta(x_{s,t}^k,t,h_{s-1}) - (x_{s+1}-x_{s,t}^k)
\right\|^2
\right].
$$
The diffusion-forcing filter updates the latent history according to
$$h_s = p_\phi(h_s \mid h_{s-1}, x_{s,t_s}^{k_s}, t_s),$$
and the conditional objective becomes
$$
L_{\mathrm{CFM}} = \frac{1}{2}\,\mathbb{E}_{s,t_s,k_s,z,h_{s-1}}
\left[
\left\|
(1-t_s)\,g_\theta(x_{s,t_s}^{k_s},t_s,h_{s-1}) - (x_{s+1}-x_{s,t_s}^{k_s})
\right\|^2
\right].
$$

A recurrent misconception would be to view $k$ as an explicit conditioning variable required by the backbone. In the formulation used for FMT, the model is trained with noisy states generated from jointly sampled $(t,k,z)$, but the prediction target is chosen so that the network need not receive $k$ directly [2509.18611].

## 3. Transformer architecture and conditioning pathway

FMT is described as a SiT-style transformer with AdaLN-Zero conditioning. Its inputs are the flattened noisy latents drawn from four pyramid scales corresponding to four consecutive time points. Positional and time embeddings for each $t_i$ are added, and at each transformer layer the history code is injected through AdaLN-Zero. The implementation uses multi-head self-attention with FlashAttention v2, SwiGLU feed-forwards, and RMSNorm throughout.

The high-level execution sequence is fixed. First, the P2VAE encoder $E_\omega$ maps physical states to latent grids. Second, noisy latent versions of adjacent states are formed using the sampled noise levels and Gaussian perturbations. Third, a small GRU or RNN consumes each noisy latent together with its sampled time to produce history codes. Fourth, the stacked latent-token sequence is processed by the transformer. Finally, the output token sequence is reassembled into a grid of the same shape as the most recent latent and linearly projected to the velocity field $g_\theta$ [2509.18611].

The model therefore combines two distinct forms of temporal conditioning: explicit tokenized latent context from recent frames and an implicit recurrent state summarizing the past. The paper states that all architectural modules follow standard implementations, and on that basis FMT is said to scale from 6 M to 138 M parameters with linear compute scaling [2509.18611].

## 4. Latent temporal pyramids and the P2VAE interface

The latent temporal pyramid is central to FMT’s efficiency. The rationale given is that physics is often near-Markovian: recent frames matter more, while older frames can be processed more coarsely. For four consecutive latent frames $y_0,y_1,y_2,y_3$, the construction is:

| Frame level | Spatial treatment | Role |
|---|---|---|
| Level 0 | Downsample $y_0$ by factor 8 | Oldest coarse context |
| Level 1 | Downsample $y_1$ by factor 4 | Intermediate context |
| Level 2 | Downsample $y_2$ by factor 2 | Recent context |
| Level 3 | Use $y_3$ at full latent resolution | Current detailed state |

These token sets are stacked along the sequence dimension, augmented with time embeddings, and passed into the transformer. Self-attention then fuses information across all pyramid scales, allowing the model to route fine-scale information from the most recent state while retaining coarse summaries of earlier states. A plausible implication is that FMT reallocates attention budget toward the temporally local information most relevant for near-Markovian dynamics.

The P2VAE provides the latent space in which FMT operates. The encoder-decoder pair is
$$E_\omega(x)=y, \qquad D_\omega(y)=\hat{x},$$
with a latent downscaling of $128 \times 128 \rightarrow 16 \times 16$ and channel compression $c3 \rightarrow c16$. The VAE is pretrained using the ELBO-style loss
$$
L_{\mathrm{VAE}} = \frac{1}{2}\,\mathbb{E}_x\left[\|x-D_\omega(E_\omega(x))\|^2\right]
+ \beta\,D_{\mathrm{KL}}(q_\omega(y|x)\|p(y)),
$$
with $\beta = 10^{-3}$ during P2VAE pretraining. Once these weights are frozen, every physical state $x_s$ is replaced by its latent $y_s = E_\omega(x_s)$, and the paper states that FMT then sees only sequences of latent states, reducing memory and compute by approximately $12\times$ [2509.18611].

## 5. Training regime and sampling procedures

Training is organized in two stages. First, the P2VAE is pretrained for 100k steps to minimize $L_{\mathrm{VAE}}$. Second, the P2VAE is frozen and FMT together with the diffusion-forcing RNN is trained for 100k steps to minimize $L_{\mathrm{CFM}}$. For joint fine-tuning on a new system, the total loss is
$$L_{\mathrm{VAE}}(\omega) + L_{\mathrm{CFM}}(\theta,\phi),$$
with a stop-gradient placed on $y$ to preserve the VAE [2509.18611].

For deterministic long-term rollouts, sampling follows a probability-flow ODE. Given a clean history $h_{s-1}$ from $y_{0:s-1}$, the current step sets $k_i = 1$ so that no initial-condition noise is injected. The procedure initializes $t \leftarrow 0$ and $x \leftarrow x_s$, then iterates for $N=100$ with $dt=0.01$:
$$x \leftarrow x + dt \cdot g_\theta(x,t,h_{s-1}), \qquad t \leftarrow t + dt.$$
The resulting terminal state approximates $x_{s+1}$. The paper states that this procedure exactly integrates the continuity equation and is numerically stable.

For aleatoric uncertainty, the model uses the reverse-time SDE adapted from PF-ODE:
$$
dx = \left[g_\theta + \frac{1}{2}\eta^2(1-t)\left(x-x_s-tg_\theta\right)\right]dt + \eta(1-t)\,d\bar{w}_t,
$$
discretized by Euler–Maruyama. Here $\eta$ controls the noise amplitude. A separate initial-condition uncertainty control is provided by setting $k_s<1$ only at the final step while keeping the history clean. The paper explicitly states that ensembles generated by varying $k_3$ isolate initial-condition uncertainty from aleatoric noise [2509.18611].

## 6. Computational profile, benchmarks, and interpretation

The efficiency claim for FMT is derived from its pyramid attention structure. A naive 4-frame video model with 256 tokens per frame incurs attention cost of order $O((4 \cdot 256)^2)$ per layer. FMT instead uses token counts corresponding to the four pyramid resolutions, with
$$(2^2)^2 + (4^2)^2 + (8^2)^2 + (16^2)^2,$$
in place of
$$(1024)^2 = (4 \cdot 16^2)^2.$$
The stated efficiency gain is
$$
\eta = \frac{(4 \times 16^2)^2}{2^4 + 4^4 + 8^4 + 16^4} \approx 15.
$$
Empirically, the paper reports that FMT attains up to $15\times$ throughput increase compared to a bi-directional self-attention video-diffusion baseline [2509.18611].

The reported downstream results include few-shot adaptation on unseen Kolmogorov turbulence at $\mathrm{Re}=222$. With only 5k finetune steps, the model yields $L2$ relative error $\approx 0.0836$ and $\mathrm{VRMSE} \approx 0.1053$, whereas training from scratch under the same budget gives $0.1342 / 0.2367$. For long-term rollout on PDEArena–NS and PDEBench–CNS–Low/High, FMT-B with 42 M parameters is reported to outperform VICON-88 M uniformly; on PA-NS at step 10, the average L2RE is 0.2234 versus 0.3618. The uncertainty-stratified ensembles are described as showing coherent spreads as $k$ or $\eta$ vary, without spurious artifacts [2509.18611].

These results frame FMT as a generative PDE foundation-model backbone rather than a conventional deterministic sequence predictor. A common misunderstanding would be to classify it as a full-length video diffusion model. The paper instead characterizes it as a flow-matching model in latent space, equipped with diffusion forcing and temporal pyramids, and designed specifically to support stable rollouts and uncertainty-aware ensemble generation in PDE settings.

Source: https://www.emergentmind.com/topics/flow-marching-transformer-fmt