Technical Summary: Improving Rectified Flow with Boundary Conditions

Problem Statement

Rectified Flow (RF), also known as Flow Matching, has emerged as a powerful generative modeling framework that learns a velocity field to transform a noise distribution into a target data distribution via an Ordinary Differential Equation (ODE). Despite its success in tasks like text-to-image and video generation, the authors identify a critical limitation in standard RF implementations: the learned velocity field often fails to satisfy theoretically derived boundary conditions.

Specifically, under linear interpolation between noise ($t=0$) and data ($t=1$), the optimal velocity field should satisfy:

1. Left Boundary ($t=0$): $\mathbf{v}(\mathbf{x}, 0) = \mathbb{E}[\mathbf{X}_1] - \mathbf{x}$
2. Right Boundary ($t=1$): $\mathbf{v}(\mathbf{x}, 1) = \mathbf{x}$

Standard neural network parameterizations of the velocity field do not inherently enforce these constraints. As visualized in the paper, the predicted velocity at $t=1$ often deviates from the expected identity map. This violation is particularly detrimental during stochastic sampling (e.g., using Langevin dynamics). Since the score function estimation (via Tweedie's formula) involves a division by $(1-t)$, an inaccurate velocity near $t=1$ causes the score estimate to diverge. This leads to unstable sampling, resulting in over-smoothed images or "cartoonish" artifacts, especially when using larger step sizes or noise scales.

Methodology

To address this, the authors propose the Boundary-enforced Rectified Flow Model (Boundary RF Model). This approach explicitly incorporates boundary constraints into the velocity field parameterization through two distinct variants:

1. Mask-based Boundary RF Model

This variant enforces both boundary conditions ($t=0$ and $t=1$) by structuring the velocity field as a linear combination of boundary terms and a learnable neural network output:
$$\mathbf{v}(\mathbf{x}, t) = g(t) \cdot (C - \mathbf{x}) + f(t) \cdot \mathbf{x} + h(t) \cdot m_\theta(\mathbf{x}, t)$$
Where:

• $m_\theta(\mathbf{x}, t)$ is a standard neural network (e.g., U-Net or DiT).
• $C = \mathbb{E}[\mathbf{X}_1]$.
• $f(t), g(t), h(t)$ are scalar functions designed to vanish or activate at specific boundaries (e.g., $g(1)=f(0)=h(0)=h(1)=0$ and $g(0)=f(1)=1$).
• The authors suggest a "Standard Cosine" choice: $g(t) = \cos(\frac{\pi}{2}t)$, $f(t) = \sin(\frac{\pi}{2}t)$, and $h(t) = \sin(\pi t)$.

This construction guarantees that the boundary conditions are satisfied by design, regardless of the neural network's output.

2. Subtraction-based Boundary RF Model

Motivated by the observation that enforcing the $t=1$ condition alone yields significant empirical benefits, the authors propose a simpler, single-boundary variant:
$$\mathbf{v}(\mathbf{x}, t) = \mathbf{x} + m_\theta(\mathbf{x}, t) - m_\theta(\mathbf{x}, 1)$$
This formulation inherently satisfies $\mathbf{v}(\mathbf{x}, 1) = \mathbf{x}$ because the term $m_\theta(\mathbf{x}, 1)$ cancels out at $t=1$. While this approach offers design simplicity and flexibility, it incurs a computational cost, requiring an additional forward pass to compute $m_\theta(\mathbf{x}, 1)$ for every evaluation.

Impact on Stochastic Sampling

By enforcing $\mathbf{v}(\mathbf{x}, 1) = \mathbf{x}$, the Boundary RF Model ensures that the numerator in Tweedie's formula ($t\mathbf{v}(\mathbf{x}, t) - \mathbf{x}$) approaches zero as $t \to 1$. This prevents the score function $\nabla \log \rho_t$ from diverging, thereby stabilizing stochastic samplers (SDE samplers) and allowing for robust generation even with larger noise scales or step sizes.

Key Contributions

1. Identification of a Limitation: The paper highlights that unconstrained velocity fields in Rectified Flow violate theoretical boundary conditions, leading to inaccurate ODE trajectories and unstable stochastic sampling.
2. Novel Parameterization: The introduction of the Boundary RF Model, which offers two practical variants (Mask-based and Subtraction-based) to enforce these conditions with minimal code modification.
3. Stabilization of Score Functions: Demonstrating that boundary enforcement directly stabilizes the score function near $t=1$, mitigating the divergence issues that plague stochastic sampling in vanilla RF.
4. Empirical Validation: Comprehensive experiments showing consistent improvements in generation quality and sampling stability across different datasets and model scales.

Experimental Results

The authors evaluated their models on CIFAR-10 and ImageNet (256$\times$256 and 512$\times$512) using Diffusion Transformer (DiT) architectures.

• Quantitative Improvements:

  • On ImageNet 256$\times$256, the Subtraction-based Boundary RF Model achieved an FID of 6.32 (ODE sampling) and 6.38 (SDE sampling), compared to the vanilla RF baseline of 6.87 and 6.79, respectively.
  • The Mask-based model achieved an FID of 6.63 (ODE) and 6.18 (SDE).
  • Improvements were also observed in Inception Score (IS), Precision, and Recall across both datasets.
  • The Subtraction-based variant consistently achieved the highest IS scores.

• Sampling Robustness:

  • In experiments using the Overshooting Sampler with varying noise scales ($c$), Boundary RF models demonstrated significantly more stable performance and less degradation in sample quality compared to vanilla RF as stochasticity increased.
  • Qualitative results showed that Boundary RF models produce sharper images with fewer artifacts and better object fidelity, particularly in stochastic sampling settings where vanilla RF tends to over-smooth.

• Scalability:

  • The benefits of the Boundary RF Model were consistent when scaling to larger models (DiT-L, DiT-XL) and higher resolutions (512$\times$512).

Significance and Claims

The paper claims that Boundary RF Model offers a simple, readily implementable solution to a fundamental limitation in Rectified Flow. By explicitly enforcing boundary conditions, the method:

1. Improves Generation Quality: Consistently outperforms vanilla RF in FID and other metrics.
2. Enables Stable Stochastic Sampling: Resolves the divergence of the score function near the terminal time, allowing for more robust sampling strategies.
3. Is Easy to Adapt: Can be integrated into existing RF models (like those using U-Net or DiT) with minimal architectural changes.

The authors position this work as a necessary refinement for practical deployment of Rectified Flow, particularly in scenarios requiring high-quality stochastic sampling. They suggest that while their focus is on RF, the core concept of enforcing boundary conditions could potentially be extended to diffusion models with minor adaptations. Future work is proposed to explore adaptive boundary enforcement and application to text-to-video generation.