Block Flow: Learning Straight Flow on Data Blocks

TL;DR

This paper introduces block matching, a novel method that partitions data using label information to learn straighter flows in flow-matching models, improving sampling efficiency and reducing errors.

Contribution

The paper proposes block matching to partition data with labels, controlling flow curvature and enhancing generative performance in flow-matching models.

Findings

Achieves competitive generation performance with fewer errors.

Effectively balances diversity and numerical stability.

Provides a flexible regularization strategy for flow curvature control.

Abstract

Flow-matching models provide a powerful framework for various applications, offering efficient sampling and flexible probability path modeling. These models are characterized by flows with low curvature in learned generative trajectories, which results in reduced truncation error at each sampling step. To further reduce curvature, we propose block matching. This novel approach leverages label information to partition the data distribution into blocks and match them with a prior distribution parameterized using the same label information, thereby learning straighter flows. We demonstrate that the variance of the prior distribution can control the curvature upper bound of forward trajectories in flow-matching models. By designing flexible regularization strategies to adjust this variance, we achieve optimal generation performance, effectively balancing the trade-off between maintaining diversity in generated samples and minimizing numerical solver errors. Our results demonstrate competitive performance with models of the same parameter scale.Code is available at \url{https://github.com/wpp13749/block_flow}.