Improving Rectified Flow with Boundary Conditions (2506.15864v1)

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.

Boundary-Enforced Rectified Flow Models: Enhancing Generative Model Accuracy in High-Dimensional Spaces

The paper entitled "Improving Rectified Flow with Boundary Conditions" addresses a prominent challenge in flow-based generative models, specifically related to enforcing boundary conditions within Rectified Flow models. It proposes the Boundary-enforced Rectified Flow Model (Boundary RF Model) as a solution to improve generative performance by ensuring boundary conditions are satisfied without extensive modifications to existing model architectures. This approach holds significant implications for the reliability and accuracy of generative modeling, especially in high-dimensional applications.

Overview of Rectified Flow Models

Rectified Flow models have emerged as a notable alternative to diffusion models, focusing on solving an Ordinary Differential Equation (ODE) to transform a noise distribution into a target data distribution. These models rely on learning a velocity field that governs the trajectory from noise to data, allowing efficient numerical integration for generating samples. However, the process of directly modeling this velocity with unconstrained neural networks has shown to often violate theoretical boundary conditions, thereby deviating from the expected ODE behavior, particularly near the terminal points of stochastic sampling paths.

Proposed Boundary-Enforced Rectified Flow Models

The authors identify that the existing implementations of Rectified Flow models struggle with maintaining boundary conditions, especially at the terminal time ($t=1$). The Boundary RF Model rectifies this by implementing constraints that enforce the correct behavior at the boundaries: the velocity field becomes the identity map when applied to clean images as $t$ approaches 1. This enforcement is achieved through parameterization variants that introduce minimal changes to existing architectures while making theoretical boundaries hold in practice.

Two primary variants are presented: the Mask-based Boundary RF Model, which incorporates constraints via informed parameterization; and the Subtraction-based Boundary RF Model, which simplifies enforcement by relying on a single parameterization to achieve stable end behavior.

Empirical Validation and Results

Through comprehensive experiments, the Boundary RF Model demonstrates consistent superiority in image generation tasks over the vanilla RF model. Specifically, improvements are quantified using metrics such as the Fréchet Inception Distance (FID), where the Boundary RF Model achieves an 8.01% reduction in FID on ImageNet using ODE sampling compared to the baseline. Additionally, systematic ablation studies confirm the importance of boundary condition enforcement, showing that both Mask-based and Subtraction-based models stabilize the model and control approximation errors in stochastic environments.

Implications and Future Directions

The work markedly indicates that by enforcing boundary conditions, one can achieve improved generation quality and stability in sampling procedures. This not only advances the technical capability of Rectified Flow models but also provides a robust architecture for future model developments.

Looking ahead, potential directions include scaling these enhanced models to larger, higher-resolution tasks and extending their application to other domains, such as text-to-image generation or video synthesis. Additionally, exploring adaptive strategies for boundary enforcement could yield further performance improvements and task-specific refinements.

In summary, this research reinforces the notion that small architectural adjustments, grounded in theoretical rigor, can have a transformative impact on the efficacy of generative models. The Boundary RF Model sets a precedent for the thoughtful incorporation of theoretical constraints into model design, enhancing the robustness and quality of output in complex, high-dimensional tasks.