Variational Rectified Flow Matching

• The paper introduces VRFM, which explicitly models multi-modal velocity fields using latent variables to improve generative sample quality over classic methods.
• VRFM employs a variational ELBO combining reconstruction loss and KL regularization to optimize neural ODE-based transport between distributions.
• Empirical results on datasets like MNIST, CIFAR-10, and ImageNet demonstrate VRFM’s superior FID, log-likelihood, and sampling efficiency.

Variational Rectified Flow Matching (VRFM) is a generative modeling framework that extends rectified flow matching by explicitly accounting for the inherent multi-modality of velocity vector-fields encountered when transporting samples between probability distributions. Unlike classic rectified flow matching, which collapses diverse optimal transport directions into uni-modal field estimates through mean-squared error regression, VRFM introduces a latent variable to parameterize and sample from a mixture of plausible flows at each point along a distribution-matching trajectory. This approach more faithfully captures the structure of complex, multi-modal data distributions and provides improvements in both sample quality and sampling efficiency across a broad range of synthetic and real-world domains.

1. Theoretical Foundations

Rectified flow matching provides a continuous interpolation mechanism between a source $p_0(x_0)$ and target data distribution $p_1(x_1)$. For any coupled pair $(x_0, x_1)$, points along the linear interpolation

$x_t = (1-t)x_0 + t x_1, \qquad t \in [0, 1]$

are moved through the vector field

$v(x_t, t) = x_1 - x_0.$

A neural velocity field $v_\theta(x, t)$ parameterizes this transport, defining the generative process through the ODE

$\frac{dx}{dt} = v_\theta(x, t), \qquad x(0) = x_0,$

with the evolving density $p_t(x)$ governed by the transport PDE

$$\frac{\partial \log p_t(x)}{\partial t} = -\mathrm{div}\,v_\theta(x,t).$$

In classic flow matching, when two different starting/ending pairs $(x_0, x_1)$ and $p_1(x_1)$0 are linearly interpolated to identical $p_1(x_1)$1, the set of valid velocities at that location is multi-modal. However, mean-squared error loss enforces a regression to the mean, erasing this structure and biasing the learned field.

VRFM augments this structure by associating a latent variable $p_1(x_1)$2 with each instance. Given $p_1(x_1)$3, the velocity field is modeled conditionally as $p_1(x_1)$4, producing a Gaussian mixture marginal over $p_1(x_1)$5.

2. Variational Formulation and Training Objective

The generative process aims to maximize the marginal likelihood of the "ground-truth" velocity $p_1(x_1)$6 for a given $p_1(x_1)$7: $p_1(x_1)$8 Introducing an encoder $p_1(x_1)$9, training proceeds by maximizing the evidence lower bound (ELBO): $(x_0, x_1)$0 Given the Gaussian form, the ELBO reduces (up to a constant) to a reconstruction (MSE) term plus a KL regularizer: $(x_0, x_1)$1 where $(x_0, x_1)$2 modulates the regularization trade-off.

3. Model Architecture and Sampling Procedure

Both the velocity network $(x_0, x_1)$3 and posterior $(x_0, x_1)$4 are implemented using neural architectures suitable for the data domain:

• Velocity network ($(x_0, x_1)$5):
  • Input: $(x_0, x_1)$6. Time $(x_0, x_1)$7 is encoded (e.g., sinusoidal + projection), and $(x_0, x_1)$8 is processed via a small MLP.
  • Backbone: MLP for 1D/2D, convolutional-ResNet for MNIST, UNet/Transformer for higher-dimensional domains (e.g., CIFAR-10, ImageNet).
  • Output: mean velocity vector in the data domain.
• Posterior network ($(x_0, x_1)$9):
  • Input: any combination of $x_t = (1-t)x_0 + t x_1, \qquad t \in [0, 1]$0.
  • Architecture: analogous backbone ending in mean ($x_t = (1-t)x_0 + t x_1, \qquad t \in [0, 1]$1) and log-variance ($x_t = (1-t)x_0 + t x_1, \qquad t \in [0, 1]$2) heads for reparameterization.

The sampling process for generation is as follows:

1. Sample $x_t = (1-t)x_0 + t x_1, \qquad t \in [0, 1]$3 and $x_t = (1-t)x_0 + t x_1, \qquad t \in [0, 1]$4.
2. Numerically integrate $x_t = (1-t)x_0 + t x_1, \qquad t \in [0, 1]$5 from $x_t = (1-t)x_0 + t x_1, \qquad t \in [0, 1]$6 to $x_t = (1-t)x_0 + t x_1, \qquad t \in [0, 1]$7.
3. Output $x_t = (1-t)x_0 + t x_1, \qquad t \in [0, 1]$8, which is a sample from the model's approximation to $x_t = (1-t)x_0 + t x_1, \qquad t \in [0, 1]$9.

4. Empirical Results and Benchmarks

VRFM yields superior empirical performance on synthetic and real datasets:

Task	Metric	Classic FM	VRFM
1D Gaussian→bimodal	Log-likelihood (PW)	Lower	Higher
2D circle transport	Likelihood/qualitative	Lower	Higher
MNIST (28×28)	FID vs. NFE	Higher	Lower
CIFAR-10 (32×32)	FID (NFE=5)	≈35.5	≈28.9
CIFAR-10 (adaptive)	FID	3.66	3.55
ImageNet (256²)	FID-50K@400K	17.2	14.6
ImageNet (256²)	FID-50K@800K	13.1	10.6

On MNIST, VRFM exhibits smooth 2D manifolds in latent space, allowing manipulation of digit style via $v(x_t, t) = x_1 - x_0.$0 and content diversity via $v(x_t, t) = x_1 - x_0.$1. On CIFAR-10 and ImageNet, VRFM improves FID at both fixed and adaptive NFE, and supports controllable style-content synthesis by conditioning on $v(x_t, t) = x_1 - x_0.$2. With classifier-free guidance on ImageNet, V-SiT-XL improves FID further (e.g., 3.22 vs. 3.43 at 800K steps).

5. Algorithmic Implementation

Training Algorithm

$v_\theta(x, t)$9

Inference Algorithm

$\frac{dx}{dt} = v_\theta(x, t), \qquad x(0) = x_0,$0

6. Hyperparameter Regimes and Ablations

Key ablation findings include:

• KL-weight ($v(x_t, t) = x_1 - x_0.$3): Typical values are $v(x_t, t) = x_1 - x_0.$4 to $v(x_t, t) = x_1 - x_0.$5 for CIFAR, $v(x_t, t) = x_1 - x_0.$6 for MNIST.
• Posterior conditioning: Best performance when conditioning on both $v(x_t, t) = x_1 - x_0.$7 and $v(x_t, t) = x_1 - x_0.$8, or $v(x_t, t) = x_1 - x_0.$9 + $v_\theta(x, t)$0.
• Fusion mechanisms (CIFAR-10): Adaptive Norm (adding $v_\theta(x, t)$1 to time embedding) and Bottleneck Sum (injecting $v_\theta(x, t)$2 at lowest U-Net resolution) both effective.
• Latent dimension: 1D/2D for small images; 768 for CIFAR-10; 1152 for ImageNet.
• Batch size: 256–512 for images; 1,000 for synthetic.
• Optimization: AdamW, learning rate $v_\theta(x, t)$3–$v_\theta(x, t)$4, weight decay $v_\theta(x, t)$5.
• Training steps: 20K (synthetic), 100K (MNIST), 600K (CIFAR-10), 800K (ImageNet).
• Encoder size: Up to 6.7% size retains most of the performance.

7. Significance and Extensions

VRFM addresses the multi-modality of transport vector-fields in neural ODE-based generative modeling. By leveraging a variational ELBO on Gaussian mixture velocity predictions, VRFM provides:

• Higher sample quality (lower FID and higher log-likelihood) than classic flow matching, especially at low NFE regimes, streamlining sampling for practical deployment.
• A simple latent-based style-content control mechanism: fixing $v_\theta(x, t)$6 controls style, $v_\theta(x, t)$7 modulates content.
• Applicability to high-dimensional generative tasks, including complex image synthesis on benchmarks such as CIFAR-10 and ImageNet.

The use of VRFM as a building block in multi-stage and multimodal generative systems, such as audio synthesis pipelines for text-to-room impulse response generation (Vosoughi et al., 25 Oct 2025), demonstrates its utility as an ODE-based generative operator in latent space. Potential future directions include structured latent variables for more expressive control, and integration with non-linear or curved interpolation trajectories to better model the geometry of intricate data manifolds.

Summary Table: Conceptual Comparison

Feature	Classic Rectified Flow Matching	Variational Rectified Flow Matching
Velocity Modeling	Uni-modal (mean-squared error)	Multi-modal (latent-indexed)
Loss	MSE	Variational ELBO
Sampling	ODE – one velocity per location	ODE with latent-sampled velocities
Style Control	Not inherent	Direct via latent $v_\theta(x, t)$8
Empirical Results	FID/log-likelihood: baseline	FID/log-likelihood: improved
Guidance	Not explicit	Compatible with classifier-free

The introduction of explicit latent variables and a variational training regimen in VRFM fundamentally enhances the ability of flow-matching models to replicate complex, multi-modal data distributions and supports new advances in efficient conditional generative modeling (Guo et al., 13 Feb 2025).