VRFM: Variational Rectified Flow Matching

• VRFM is a generative modeling framework that introduces latent-variable parameterizations to capture multi-modal velocity fields for improved sampling.
• It employs a variational encoder and an ODE-based integration to transform a simple base distribution into a target, yielding more accurate and controllable flows.
• Empirical results in image synthesis and audio source separation highlight VRFM’s enhanced sample quality and controllability compared to deterministic approaches.

Variational Rectified Flow Matching (VRFM) is a generative modeling framework that enhances classic rectified flow matching by introducing latent-variable parameterizations of velocity vector fields in order to represent and sample from multi-modal flow directions. In contrast to deterministic flow matching, VRFM enables more expressive, controllable, and accurate transformations between a simple base distribution and a target distribution by explicitly modeling locally ambiguous or multi-modal flow fields. VRFM has demonstrated empirical benefits across synthetic, image, and audio modalities, notably including large-scale image synthesis and language-queried audio source separation tasks (Guo et al., 13 Feb 2025, Yuan et al., 2024).

1. Foundations: Rectified Flow Matching

Classic rectified flow matching (RFM) transforms samples from a simple source distribution $p_0(x_0)$ (typically isotropic Gaussian) to a target data distribution $p_1(x_1)$ by integrating a learned velocity vector field $v_\theta(x, t)$ over the interval $t \in [0, 1]$, with the dynamics specified by the ordinary differential equation (ODE):

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

The training procedure relies on randomly coupling pairs $(x_0, x_1)$ drawn independently from $p_0$ and $p_1$, generating linear interpolants:

$x_t = (1-t)x_0 + t x_1$

and using the difference $v = x_1 - x_0$ as the "ground-truth" velocity at position $p_1(x_1)$0. The standard mean-square error loss is:

$p_1(x_1)$1

This results in a learned deterministic velocity field $p_1(x_1)$2, which, at any position $p_1(x_1)$3, regresses toward the average direction of all possible paired flows.

2. Limitations of Deterministic Flow Matching

In multi-modal transformation tasks, the set of “ground-truth” velocity vectors associated with a specific $p_1(x_1)$4 may be multi-directional due to random pairings of $p_1(x_1)$5 and $p_1(x_1)$6. The deterministic (single-vector) field produced by minimizing the $p_1(x_1)$7 loss cannot represent this ambiguity; instead, it averages over possible flow directions, introducing artifacts such as curved or misaligned flows, as evidenced by the baseline's performance on synthetic 1D/2D tasks and CIFAR-10 (Guo et al., 13 Feb 2025). These averaging effects result in suboptimal sample quality, reduced likelihoods, and less controllable transformations, particularly in the presence of inherent data multi-modality.

3. The VRFM Framework: Modeling Multi-modal Velocity Fields

Variational Rectified Flow Matching addresses multi-modality by introducing a latent variable $p_1(x_1)$8 and modeling the conditional velocity distribution as:

$p_1(x_1)$9

This forms a Gaussian mixture over possible velocity directions at each location. Since $v_\theta(x, t)$0 is unobserved during training, VRFM uses a variational encoder $v_\theta(x, t)$1, which produces a posterior Gaussian:

$v_\theta(x, t)$2

The training objective maximizes the marginal log-likelihood of the observed velocities using a standard evidence lower bound (ELBO):

$v_\theta(x, t)$3

This approach yields a latent-dependent, multi-modal velocity field, with multi-modality recoverable at inference by sampling different $v_\theta(x, t)$4 vectors.

4. Training and Inference Procedures

Training

The typical VRFM training loop involves the following steps:

• Sample $v_\theta(x, t)$5 and $v_\theta(x, t)$6.
• Sample $v_\theta(x, t)$7 and compute $v_\theta(x, t)$8.
• Encode $v_\theta(x, t)$9 via reparameterization.
• Minimize the VRFM loss $t \in [0, 1]$0 via stochastic gradient descent.

Model parameterizations vary by domain; for images, UNet or Transformer backbones are employed, and for audio separation tasks (FlowSep), a U-Net with cross-attention on text embeddings is used (Yuan et al., 2024, Guo et al., 13 Feb 2025).

Inference

At inference,

• Sample $t \in [0, 1]$1 (e.g., Gaussian noise).
• Draw $t \in [0, 1]$2.
• Integrate the ODE $t \in [0, 1]$3 from $t \in [0, 1]$4 to $t \in [0, 1]$5 using methods such as Euler or Dopri5.
• The final point $t \in [0, 1]$6 represents a sample from the target distribution.

In domains such as conditional audio separation, the flow operates in a pre-trained variational autoencoder (VAE) latent space and is conditioned on additional inputs, such as a mixture encoding and text query embedding (Yuan et al., 2024).

5. Empirical Results Across Modalities

Empirical studies on synthetic 1D/2D data, MNIST, CIFAR-10, and ImageNet demonstrate the advantages of VRFM over deterministic RFM:

Domain	Metric	RFM Baseline	VRFM
1D, 2D	Velocity var.	Collapses to mean	Matches ground-truth multi-modality
MNIST	FID	Higher, no z control	Lower, smooth/controllable by z
CIFAR-10	FID@NFE=2	166.7	117.7 (adaptive-norm, KL=5e-3)
ImageNet	FID@50K	17.2→14.6 (w/o cfg)	17.2→14.6; with cfg: 5.40→4.91

VRFM achieves straighter, more intersecting flows, higher log-likelihoods, and improved sample quality (LL, FID, Inception Score) across all tested image scales (Guo et al., 13 Feb 2025).

In audio source separation (FlowSep), RFM in VAE latent space underpins generative models that outperform discriminative approaches in separation quality and efficiency, with no reported VRFM extension in FlowSep to explicitly model multi-modal velocities (Yuan et al., 2024). This suggests current instantiations for separation may still assume a unimodal velocity field in latent space.

6. Architecture, Conditioning, and Hyperparameters

Model architectures for VRFM are domain-specific. For images, UNets, conv–ResNets, and large-scale Transformers (e.g., SiT-XL) are used, with the posterior encoder implemented as 3–5 block MLP or Transformer variants. For language-queried source separation in audio, FlowSep employs a U-Net with cross-attention for text query integration, based on frozen FLAN-T5 features and mixture latents (Yuan et al., 2024).

Hyperparameter selection includes:

• Learning rates: 1e-3 (synthetic, MNIST), 2e-4 (CIFAR-10), 1e-4 (ImageNet)
• KL weights: dataset-dependent, e.g., 1e-3 for MNIST, 2e-3 for ImageNet
• Batch sizes: 128–256 (images); 8 (FlowSep)
• Steps: up to 800K for ImageNet; up to 1M for FlowSep (Yuan et al., 2024, Guo et al., 13 Feb 2025)
• Inference: Number of function evaluations (NFE) is tunable, e.g., 10 (FlowSep), and as low as 2 sufficient for some image settings.

7. Advantages, Limitations, and Future Directions

Advantages:

• Explicitly models multi-modal velocity fields, recovering correct local ambiguity and enabling diverse sampling.
• Produces straighter and more intersectional flows, resulting in easier ODE integration and improved convergence.
• Provides controllable and disentangled sampling via the latent code $t \in [0, 1]$7, demonstrated by style and content control in MNIST and CIFAR-10.

Limitations:

• Increases computational overhead due to the auxiliary posterior encoder and the KL divergence term in the ELBO.
• Sensitivity to hyperparameters such as latent dimensionality and KL weight, which require per-task tuning.
• It remains an open problem to optimally combine VRFM with consistency-based accelerations or to extend it to stochastic (SDE-based) flows (Guo et al., 13 Feb 2025).

Future Directions include: integrating VRFM with consistency or distillation frameworks to further reduce inference cost, learning adaptive priors $t \in [0, 1]$8, extending to hierarchical or stochastic settings, and broadening applications to novel domains such as graphs and 3D data (Guo et al., 13 Feb 2025).

---

VRFM provides a principled, variational generalization of rectified flow matching, equipped to address the multi-modality inherent in complex generative modeling tasks. Its capacity for learning and sampling diverse flow directions has been empirically validated and is a subject of continuing research advancement in both image and audio generative modeling (Guo et al., 13 Feb 2025, Yuan et al., 2024).