Improving Rectified Flow with Boundary Conditions

2025-06-18

AI Summary

Unconstrained velocity field modeling in Rectified Flow often violates boundary conditions—e.g., non-zero terminal velocities—leading to ODE solution drift and amplified sampling errors. To address this, we propose the first explicit boundary-constrained formulation for Rectified Flow, enforcing physically consistent and numerically stable velocity field learning by incorporating boundary regularizers (e.g., zero-velocity constraints at endpoints) with minimal code modifications to existing frameworks. Our method jointly optimizes neural velocity fields, ODE/SDE solvers, and boundary-aware regularization. Evaluated on ImageNet, it achieves an 8.01% FID reduction under deterministic ODE sampling and an 8.98% FID reduction under stochastic SDE sampling, significantly improving generation quality and sampling robustness. The core contribution is the first boundary-enhanced Rectified Flow paradigm, rigorously grounded in dynamical systems theory while maintaining practical deployability across mainstream diffusion-based generative models.

Technology Category

• Computer Vision: Diffusion Models for Vision
• Search and Optimization: Sampling/Simulation-based Search
• Machine Learning: Learning with Manifolds

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Graph Algorithms and Modeling for the Web: Efficient manipulation of static and dynamic Web-related graphs
• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking

Abstract

Rectified Flow offers a simple and effective approach to high-quality generative modeling by learning a velocity field. However, we identify a limitation in directly modeling the velocity with an unconstrained neural network: the learned velocity often fails to satisfy certain boundary conditions, leading to inaccurate velocity field estimations that deviate from the desired ODE. This issue is particularly critical during stochastic sampling at inference, as the score function's errors are amplified near the boundary. To mitigate this, we propose a Boundary-enforced Rectified Flow Model (Boundary RF Model), in which we enforce boundary conditions with a minimal code modification. Boundary RF Model improves performance over vanilla RF model, demonstrating 8.01% improvement in FID score on ImageNet using ODE sampling and 8.98% improvement using SDE sampling.

Problem

• Rectified Flow fails to satisfy boundary conditions
• Unconstrained neural network causes inaccurate velocity field
• Boundary conditions improve generative model performance

Innovation

• Enforces boundary conditions in Rectified Flow
• Minimal code modification for boundary enforcement
• Improves FID scores significantly on ImageNet