Variational Rectified Flow Matching (V-RFM)

• Variational Rectified Flow Matching (V-RFM) is a generative framework that extends classic flow matching by using latent variables to model multi-modal velocity fields.
• It integrates variational inference with ODE-based generation to capture distinct and intersecting sample trajectories in high-dimensional spaces.
• The approach overcomes mode-averaging limitations of deterministic flows, improving performance on tasks like high-fidelity image synthesis and language-queried audio separation.

Variational Rectified Flow Matching (V-RFM) is a generative modeling framework that extends classic rectified flow matching (RFM) by explicitly modeling multi-modal velocity vector fields via latent variables. V-RFM addresses the intrinsic multi-modality in the ground-truth velocities associated with optimal transport between distributions, preventing mode-averaging and enabling accurate sample trajectories across high-dimensional and multi-modal target distributions. It achieves this by integrating a variational inference approach with the rectified flow matching paradigm, and has been successfully deployed for tasks such as high-fidelity image generation and language-queried audio source separation (Guo et al., 13 Feb 2025, Yuan et al., 2024).

1. Foundations of Rectified Flow Matching and Mode Collapse

Classic rectified flow matching learns a deterministic velocity field $v_\theta(x_t,t)$ to deform samples $x_0$ from a source distribution $p_0$ into samples $x_1$ from a target distribution $p_1$ along a linear interpolation path: $$\frac{d x_t}{dt} = v_\theta(x_t, t),\quad x_0\sim p_0,\; x_1 \sim p_1,\; x_t = (1-t)x_0+t x_1.$$ During training, the "ground-truth" velocity at any $(x_t, t)$ is $v_{\text{gt}} = x_1 - x_0$; however, for each $(x_t, t)$, multiple $(x_0, x_1)$ pairs exist, making the ground-truth velocity field inherently multi-modal. Under the standard mean squared error (MSE) loss

$x_0$0

the optimizer is forced to predict the mean direction at each point, yielding mode-averaged, non-intersecting flows and restricting the expressivity of the generative process (Guo et al., 13 Feb 2025).

2. V-RFM: Variational Modeling of Velocity Fields

V-RFM augments the deterministic flow field with a variational latent variable $x_0$1 and models velocity as a conditional Gaussian: $x_0$2 with $x_0$3 sampled from a standard normal prior $x_0$4: $x_0$5 so that the velocity distribution at every $x_0$6 is a mixture of Gaussians, capable of capturing multiple, intersecting flow directions. This variational approach is realized in training by introducing an approximate posterior $x_0$7, typically parameterized as a Gaussian with learned mean and diagonal covariance.

The evidence lower bound (ELBO) objective for V-RFM is: $x_0$8 This objective enables the network to learn a multi-modal flow field, while the generative process at inference time is realized by sampling $x_0$9 and integrating the ODE

$p_0$0

Each $p_0$1 yields a distinct sample trajectory, naturally covering the multi-modal structure inherent in the ground-truth flows (Guo et al., 13 Feb 2025).

3. End-to-End Architectures and Loss Formulations

V-RFM can be integrated with a variational autoencoder (VAE) structure to operate in compressed latent spaces, supporting both unconditional and conditional tasks. For instance, in sound separation (Yuan et al., 2024), the architecture includes:

• A text encoder (e.g., FLAN-T5) to encode natural language queries into embeddings $p_0$2.
• A VAE encoder $p_0$3 mapping spectrograms $p_0$4 to latent vectors $p_0$5, with prior $p_0$6 and decoder $p_0$7.
• A UNet-based RFM module $p_0$8 predicting flow velocities in latent space, cross-attending to textual and mixture latents.
• A pre-trained vocoder (e.g., BigVGAN) for waveform reconstruction.

Training follows the joint loss: $p_0$9 where

$x_1$0

$x_1$1

This enables efficient learning and conditional generation in the VAE latent manifold (Yuan et al., 2024).

4. Inference and ODE Integration Procedures

At inference, the generative flow is solved via ODE integration. For unconditional or image generation tasks (Guo et al., 13 Feb 2025), one samples $x_1$2, $x_1$3 and integrates: $x_1$4 using either fixed-step Euler, adaptive solvers (e.g., Dormand–Prince), or Euler–Maruyama for SDEs. For conditional (e.g., text-to-audio) generation (Yuan et al., 2024), the steps are:

• Encode the text query into $x_1$5
• Encode the input mixture into $x_1$6
• Sample $x_1$7
• Solve $x_1$8 from $x_1$9 to $p_1$0 using $p_1$1 ODE steps
• Decode the final latent to audio via VAE decoder and vocoder

Parameter choices such as number of ODE steps $p_1$2, $p_1$3, and latent sizes are dataset and application-dependent.

5. Empirical Results and Evaluation

V-RFM demonstrates superior performance over baseline flow matching models across synthetic toy data, images, and audio tasks:

• On synthetic 1D/2D datasets, V-RFM achieves higher true log-likelihood and Parzen-window log-likelihood for all numbers of function evaluations (NFE) (Guo et al., 13 Feb 2025).
• For MNIST and CIFAR-10, V-RFM yields improved FID scores compared to classic OT-FM and instantaneous consistency flow matching (I-CFM), especially at medium/high NFE. For example, on CIFAR-10 with NFE=5, V-RFM achieves FID=25.84, compared to 36.19 for OT-FM (see table below).
• On ImageNet 256×256, class-conditional V-RFM improves FID-50K from 13.1 to 10.6 at 800k iterations (V-SiT-XL backbone); classifier-free guidance further reduces FID to 3.22 (Guo et al., 13 Feb 2025).
• In language-queried audio separation (FlowSep), V-RFM trained with 1.68k hours of audio surpasses diffusion-based baselines in both subjective and objective metrics, exhibiting higher separation quality and faster inference (Yuan et al., 2024).

NFE	OT-FM	I-CFM	V-RFM (bottleneck, KL=2e-3)
5	36.188	35.489	25.841
100	4.640	4.461	4.540
1000	3.822	3.643	3.596
adaptive	3.655	3.659	3.520

A key qualitative advantage is the recovery of naturally intersecting flow trajectories, as opposed to the artificial trajectory bending and mode collapse of deterministic flow matching.

6. Implementation and Design Considerations

Architectural choices for V-RFM are dataset and modality-specific:

• For vision, $p_1$4 is parameterized as a UNet or SiT-XL backbone, with $p_1$5 as corresponding ResNet or Transformer-based encoder. Latent sizes and KL weights vary, e.g., $p_1$6, latent dim 2–768, batch sizes 128–256, Adam or AdamW optimizers at typical learning rates $p_1$7 to $p_1$8 (Guo et al., 13 Feb 2025).
• For audio, VAE backbones from AudioLDM, FLAN-T5 text encoders, and BigVGAN vocoder are typical, with RFM modeled by 4-scale UNets and cross-attention for conditioning on text and source mixture (Yuan et al., 2024).

Critical hyperparameters include the KL regularizer (preventing over- or under-regularization of $p_1$9), the latent dimensionality (to capture sufficient multi-modality), and the ODE/SDE solver parameters (balancing quality vs. evaluation speed).

7. Advantages, Limitations, and Extensions

V-RFM offers distinct advantages:

• Faithfully captures multi-modal velocity fields to generate intersecting transport paths.
• Incur only modest additional training cost (encoder for $$\frac{d x_t}{dt} = v_\theta(x_t, t),\quad x_0\sim p_0,\; x_1 \sim p_1,\; x_t = (1-t)x_0+t x_1.$$0 and a sampling step).
• Is compatible with any ODE or SDE solver, and naturally affords controllable generation via latent variable $$\frac{d x_t}{dt} = v_\theta(x_t, t),\quad x_0\sim p_0,\; x_1 \sim p_1,\; x_t = (1-t)x_0+t x_1.$$1.

Limitations include:

• Slightly more complex ELBO training dynamics.
• Sensitivity to the choice of KL weight, latent dimension, and encoder/decoder capacity; $$\frac{d x_t}{dt} = v_\theta(x_t, t),\quad x_0\sim p_0,\; x_1 \sim p_1,\; x_t = (1-t)x_0+t x_1.$$2 can underfit or overfit if not tuned carefully.

Potential extensions are:

• Trainable or hierarchical priors on $$\frac{d x_t}{dt} = v_\theta(x_t, t),\quad x_0\sim p_0,\; x_1 \sim p_1,\; x_t = (1-t)x_0+t x_1.$$3 (e.g., $$\frac{d x_t}{dt} = v_\theta(x_t, t),\quad x_0\sim p_0,\; x_1 \sim p_1,\; x_t = (1-t)x_0+t x_1.$$4).
• Augmentation with consistency models to further lower NFE.
• Flow distillation for speedup.
• Generalization to other multi-modal conditional generation settings.

V-RFM represents a principled and flexible approach to generative modeling that restores the multi-modal structure intrinsic to task objectives such as optimal transport, conditional generation, and source separation, yielding high-fidelity, controllable, and computationally efficient sample generation (Guo et al., 13 Feb 2025, Yuan et al., 2024).