---
title: Normalized Flow Matching (NFM)
url: https://www.emergentmind.com/topics/normalized-flow-matching-nfm
type: topic
---

# Normalized Flow Matching (NFM)

Normalized Flow Matching (NFM) is a method for training neural ordinary differential equation (neural-ODE) based generative models through a coupling-distillation paradigm. It reframes the traditional flow matching (FM) learning process by leveraging the bijective mapping of a pretrained normalizing flow (NF) model as a structured coupling for endpoint pairs. This approach enables faster training, reduced sampling curvature, and improved sample quality compared to independent or optimal-transport (OT) couplings, with inference speed significantly surpassing that of the AR-NF teacher model [2603.09014].

## 1. Flow Matching and the Role of the Coupling Measure

Flow Matching (FM) generative modeling learns a time-dependent velocity field $v_\theta(x, t)$ to transform points from a source distribution (typically Gaussian noise, $p_{\rm noise}(z)$) to a target distribution (data, $p_{\rm data}(x)$). The core mechanism is a regression loss defined over data/noise pairings and interpolated positions:
\[
x_t = (1-t)x + t z, \quad t \in [0,1].
\]
A coupling measure $q(x, z)$ with $q(x) = p_{\rm data}(x)$ and $q(z) = p_{\rm noise}(z)$ determines how $x$ and $z$ are paired. The FM objective is given by:
\[
\mathcal{L}_{\nabla}(\theta)
= \mathbb{E}_{q(x, z, t)} \left[ \| v_\theta(x_t, t) - (z - x) \|_2^2 \right],
\]
where $q(x, z, t) = q(x, z)\, \delta(x_t - (1-t)x - t z) \, p(t)$ and $p(t)$ is typically uniform or logit-normal. The statistical structure of $q(x,z)$ has first-order impacts on speed of convergence, stability, and variance of the regression target. Standard FM uses an independent coupling, while alternative adaptive couplings such as OT have shown improvements.

## 2. Distilling Couplings from Normalizing Flow Teachers

A normalizing flow (NF), $f_{\rm NF}$, is an invertible map ($\mathbb{R}^n \to \mathbb{R}^n$) learned via maximum likelihood. The bijectivity allows construction of a quasi-deterministic coupling:
1. Slightly noised data, $x' = x + \eta\epsilon'$, with $\epsilon' \sim \mathcal{N}(0, I)$, are mapped via the NF,
2. The forward map yields $\tilde{z} = f_{\rm NF}(x', c)$, with normalization by the empirical pre-image standard deviation $\sigma_f$:
\[
z_{\epsilon'} = \frac{1}{\sigma_f} f_{\rm NF}(x + \eta\epsilon', c).
\]
3. A coupling density is defined as:
\[
q_{\rm NF}(x, z) = p_{\rm data}(x)\int \delta(z - z_{\epsilon'}) \mathcal{N}(\epsilon')\, d\epsilon'.
\]
This coupling is substituted into the FM loss, dissociating from the typical $p(x)p(z)$ or numerically expensive OT approaches. The distilled NFM objective becomes:
\[
\mathcal{L}_{\rm NFM}(\theta)
= \mathbb{E}_{x, \epsilon', t}\left[
\| v_\theta(x_t, c, t) - (z_{\epsilon'} - x) \|_2^2
\right],
\quad x_t = (1-t)x + t z_{\epsilon'}.
\]
This effectively enforces FM with a highly structured and nearly deterministic endpoint matching.

## 3. Key Equations, Training, and Inference Procedures

The essential equations underpinning FM and NFM are as follows:
- Generic path-wise FM loss:
\[
\mathcal{L}_{\nabla}(\theta)
= \mathbb{E}_{q(x, z), p(t)} \left\| v_\theta((1-t)x + t z, t) - (z - x) \right\|_2^2
\]
- NF-induced target:
\[
z_{\epsilon'}
= \frac{1}{\sigma_f} f_{\rm NF}(x + \eta \epsilon', c)
\]
- NFM distilled loss:
\[
\mathcal{L}_{\rm NFM}
= \mathbb{E}_{x, \epsilon', t} \left[ \left\| g(x_t, c, t) - (z_{\epsilon'} - x) \right\|_2^2 \right]
\]

**Training pseudocode**:
```python
# Inputs: frozen NF teacher f_NF, output-norm σ_f, noise-level η, time-sampler p(t)
# Initialize θ randomly
while not converged:
    x, c = sample_data_minibatch()
    ε′ = sample_noise()
    z = (1/σ_f) * f_NF(x + η * ε′, c)
    t = sample_time()
    x_t = (1-t) * x + t * z
    v_target = z - x
    v_pred = g_θ(x_t, c, t)
    loss = mean_squared_error(v_pred, v_target)
    θ = update_params(θ, loss)
```

**Sampling pseudocode** (Euler or Heun solver):
```python
# Given class c, S steps
z_S = sample_standard_normal()
for k in S,..,1:
    t = k/S      # or t^2 schedule
    Δt = 1/S
    x_{k-1} = x_k - Δt * g_θ(x_k, c, t)   # Euler
# Return x_0
```

## 4. Model Architectures and Hyperparameters

The teacher AR-NF model (“TarFlow”) uses Transformer-based autoregressive coupling layers (“meta-blocks”):
- Configuration example: TF-6×2+2×26-1024/2 (6 meta-blocks depth 2, 2 meta-blocks depth 26, 1024 dims, patch size 2).
- Parameter count: ≈0.82M.

The student FM network $g$:
- Shares patch embedding and Transformer block configuration with the teacher (SiT-XL/4 style), lacks invertibility constraint.
- Embeds time and class information via learned embeddings.
- Parameter count: ≈0.8M.

**Training settings**:
- Datasets: ImageNet-64, ImageNet-256 (with 256×256 images embedded via pretrained VAE).
- Teacher noise $\eta$: 0.05 (64×64), 0.10 (256×256).
- Teacher trained on 512M samples (~420 epochs), student and FM baselines on 256M (210 epochs).
- Optimizer: Adam ($\beta_1=0.9$, $\beta_2=0.999$), learning rate ≈$4\times10^{-4}$.
- Batch size: 1024.
- Time sampler: logit-normalDistr($a=-0.2$, $b=1$).
- Label dropout: $p=0.1$.
- Sampler: Euler for NFE$\leq$5, Heun otherwise.
- Guidance: classifier-free, tuned by golden-section search on FID and fixed.

## 5. Comparative Evaluation and Ablation Studies

Experimental comparisons on ImageNet-64 (class-conditional FM):

| Number of function evals (NFE) | Independent FM FID | OT (Semi-Discrete FM) | NFM (ours) |
|:-------------------------------|:------------------|:---------------------|:----------:|
| 31 (Heun)                      | 2.57              | 2.68                 | **1.78**   |
| 15 (Heun)                      | 4.80              | 3.15                 | **2.15**   |
| 7  (Heun)                      | 13.01             | 6.41                 | **3.23**   |
| 5  (Euler)                     | 17.56             | 9.29                 | **4.01**   |

On ImageNet-256 (NFE=31/15/7), NFM consistently outperforms both baseline and OT-FM across FID, with up to $5\times$ improvement in low-NFE regimes, and surpasses the AR-NF teacher FID (teacher: 1.98 on 64×64, student: 1.78).

Curvature metrics (lower $\kappa$ is straighter):

| Solver / NFE | FM $\kappa$ | OT-FM $\kappa$ | NFM $\kappa$ |
|:-------------|:-----------:|:--------------:|:------------:|
| Heun(31)     | 0.0864      | 0.0767         | **0.0435**   |
| Euler(128)   | 0.0386      | 0.0289         | **0.0181**   |

At 31 NFE, NFM achieves a $32\times$ reduction in latency over the AR-NF teacher while maintaining or improving sample quality.

**Ablations**:
- Teacher NLL performance linearly correlates with student FID.
- Optimal $\eta$ aligns between best teacher FID and best student FID; large $\eta$ ($>$0.2) harms both.
- NF latent $z$-space exhibits high non-isometry, but couplings remain effective.
- Qualitative samples show NFM behaves closer to distribution matching than strict pair-wise distillation.

## 6. Implications and Practical Significance

Normalized Flow Matching demonstrates that leveraging the invertible nature of a pretrained NF model enables a direct and structured coupling for FM learning. This provides substantial improvements:
- Faster training convergence (reduced number of iterations and loss).
- Improved sample quality (lower FID, especially with few function evaluations).
- Sampling speedup of $30$–$100\times$ over the AR-NF teacher.
- Student models can surpass the FID of their teachers, suggesting the distilled coupling leads to superior learned flows [2603.09014].

A plausible implication is that NFM’s approach to coupling reuse generalizes broadly across flow-based models, opening avenues for more efficient distillation paradigms in generative modeling. Open-source code, architectures, and training configurations are provided at github.com/apple/ml-nfm for reproducibility.

Source: https://www.emergentmind.com/topics/normalized-flow-matching-nfm