Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Block Flow: Learning Straight Flow on Data Blocks

Zibin Wang, Zhiyuan Ouyang, Xiangyun Zhang

TL;DR

Block Flow introduces label-informed block matching to learn straighter flows in flow-matching models by pairing data blocks $p(x_0|y)$ with localized Gaussian priors $p(z|y)$. It establishes a theoretical link between forward-trajectory curvature and prior variance, and delivers regularization-based training schemes (Full Alignment and Hybrid Information Alignment) to control this curvature while preserving sample diversity. The method demonstrates competitive results on CIFAR-10 and MNIST with parameter-efficient architectures, highlighting the potential for faster sampling with reduced truncation error. Overall, Block Flow provides a flexible framework to improve generation quality and efficiency by exploiting label structure through learnable Gaussian-mixture priors.

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}.

Block Flow: Learning Straight Flow on Data Blocks

TL;DR

Block Flow introduces label-informed block matching to learn straighter flows in flow-matching models by pairing data blocks with localized Gaussian priors . It establishes a theoretical link between forward-trajectory curvature and prior variance, and delivers regularization-based training schemes (Full Alignment and Hybrid Information Alignment) to control this curvature while preserving sample diversity. The method demonstrates competitive results on CIFAR-10 and MNIST with parameter-efficient architectures, highlighting the potential for faster sampling with reduced truncation error. Overall, Block Flow provides a flexible framework to improve generation quality and efficiency by exploiting label structure through learnable Gaussian-mixture priors.

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}.
Paper Structure (14 sections, 2 theorems, 26 equations, 8 figures, 4 tables, 2 algorithms)

This paper contains 14 sections, 2 theorems, 26 equations, 8 figures, 4 tables, 2 algorithms.

Table of Contents

  1. Introduction
  2. Background
  3. Method
  4. Block Matching
  5. Curvature Control for Forward Trajectories
  6. Regularization Strategy in Training
  7. Full Alignment Scheme
  8. Hybrid Information Alignment Scheme
  9. Related Works
  10. Experiment
  11. Discussion
  12. Conclusion
  13. Proof of Proposition
  14. Experiment Detail

Key Result

Proposition 1

If the probability distribution of $x_c$ is a Dirac delta function, i.e., $x_c \sim \delta(x - c),$ then:

Figures (8)

  • Figure 1: Forward and reverse trajectories of 1-rectified flow liu2022flowstraightfastlearning, fast ODE lee2023minimizing, and block flow on a toy example. Intersections in forward trajectories lead to averaging effects, causing the resulting reverse trajectories to deviate from the true data distribution. We partition the toy example into two blocks and match each block with a separate Gaussian distribution.By employing block matching, block flow significantly reduces cross-label intersections in forward trajectories, resulting in straighter reverse trajectories.
  • Figure 2: A visualization of the block flow training
  • Figure 3: Quantitative results of FABR with varying $beta$ on the CIFAR-10 dataset.
  • Figure 4: Comparison of Block Flow on CIFAR-10 and MNIST datasets
  • Figure 5: Results of block flow (FANR) with the $L_2$ norm on CIFAR-10 datase
  • ...and 3 more figures

Theorems & Definitions (4)

  • Proposition 1
  • Proposition 2
  • proof
  • proof