Variational Rectified Flow Matching

TL;DR

This paper introduces Variational Rectified Flow Matching, a novel framework that models multi-modal velocity vector-fields to improve flow-based generative modeling, demonstrating superior results on various datasets.

Contribution

It extends classic rectified flow matching by learning multi-modal velocity vector-fields through a variational approach, capturing complex flow directions.

Findings

Effective on synthetic data, MNIST, CIFAR-10, and ImageNet.

Learns multi-modal flow directions, outperforming traditional methods.

Produces more accurate and diverse sample generation.

Abstract

We study Variational Rectified Flow Matching, a framework that enhances classic rectified flow matching by modeling multi-modal velocity vector-fields. At inference time, classic rectified flow matching 'moves' samples from a source distribution to the target distribution by solving an ordinary differential equation via integration along a velocity vector-field. At training time, the velocity vector-field is learnt by linearly interpolating between coupled samples one drawn from the source and one drawn from the target distribution randomly. This leads to ''ground-truth'' velocity vector-fields that point in different directions at the same location, i.e., the velocity vector-fields are multi-modal/ambiguous. However, since training uses a standard mean-squared-error loss, the learnt velocity vector-field averages ''ground-truth'' directions and isn't multi-modal. In contrast, variational rectified flow matching learns and samples from multi-modal flow directions. We show on synthetic data, MNIST, CIFAR-10, and ImageNet that variational rectified flow matching leads to compelling results.