Stochastic RF Sampler

• Stochastic RF Sampler is a sampling algorithm for rectified flow generative models that injects stochasticity to overcome the rigidity of deterministic straight-line sampling.
• The method introduces data-dependent stochastic perturbations to boost sample diversity and better adapt to low-dimensional data manifolds.
• It offers theoretical guarantees with O(k/ε) convergence and demonstrates empirical improvements in image synthesis and mode coverage, bridging RF and diffusion models.

A Stochastic RF Sampler is a sampling algorithm for rectified flow (RF) generative models that introduces stochasticity into the otherwise deterministic rectified flow sampling process. Its primary objective is to address the rigidity of standard straight-line RF sampling, thereby enabling greater generative diversity and enhancing adaptation to low-dimensional data manifolds. This approach bridges behaviors found in both rectified flow and denoising diffusion models (DDPM) by leveraging connections to the theory of stochastic localization. Stochastic RF sampling has become essential for robust, high-fidelity sample generation, especially when dealing with complex or intrinsically low-dimensional data distributions (Roy et al., 21 Jan 2026).

1. Mathematical Foundations of Rectified Flow and Stochastic RF

Rectified flow models transport samples from a noise distribution (often $\mathcal N(0, I_d)$) to a target distribution $p_1$ using a continuous-time ODE: $$\frac{dZ_t}{dt} = v_t(Z_t), \quad Z_0 \sim \mathcal N(0, I_d), ~ t \in [0,1]$$ where $v_t(x) = \mathbb E[X_1 - X_0 | X_t = x]$ and $X_t = t X_1 + (1-t) X_0$ defines the straight-line interpolation between $X_0$ (noise) and $X_1$ (data). This deterministic ODE ensures that the law of $Z_t$ matches the linearly interpolated law $p_t$ at each $t$ (Liu et al., 2022, Esser et al., 2024).

In standard RF, this process is deterministic, collapsing trajectories onto a straight path between each noise–data pair. To address the limitations in diversity and adaptation, a stochastic variant—termed here the "Stochastic RF Sampler"—is introduced by augmenting the ODE with a data-dependent stochastic term, resulting in a stochastic differential equation (SDE): $p_1$0 where $p_1$1, $p_1$2 is a standard Wiener process, and $p_1$3 is the score of the marginal at time $p_1$4 (Roy et al., 21 Jan 2026). This construction closely parallels the time-reversed SDE of DDPM and is theoretically justified by stochastic localization frameworks.

2. Motivation: Limitations of Deterministic RF and the Need for Stochastic Extension

Deterministic RF sampling exhibits the following limitations:

• Lack of sample diversity: Trajectories are strictly straight lines, resulting in low generative variability and potential mode collapse, particularly near the noise boundary (Ma et al., 10 Jun 2025).
• Poor multi-scale noise modeling: RF does not naturally encode the progressive noise schedule found in diffusion, where stochasticity is higher at early steps and anneals near the data manifold.
• Sampling complexity in high-dimension: Standard Euler discretization of the ODE converges at a rate independent of the intrinsic data dimensionality, limiting efficiency when the data lies on a low-dimensional support.

Stochastic RF sampling injects controlled stochastic perturbations into the velocity field or directly into state transitions, yielding improved diversity, better adaptation to the “intrinsic dimension” $p_1$5 of the target support, and convergence guarantees that scale as $p_1$6 for total variation error $p_1$7 (Roy et al., 21 Jan 2026). This adaptation is especially significant for realistic datasets where $p_1$8.

3. Algorithms and Discrete Implementations

3.1 Deterministic RF Sampler

Given $p_1$9 time steps $$\frac{dZ_t}{dt} = v_t(Z_t), \quad Z_0 \sim \mathcal N(0, I_d), ~ t \in [0,1]$$0, and step size $$\frac{dZ_t}{dt} = v_t(Z_t), \quad Z_0 \sim \mathcal N(0, I_d), ~ t \in [0,1]$$1: $v_t(x) = \mathbb E[X_1 - X_0 | X_t = x]$5 A nonuniform (U-shaped) time grid, with denser steps near $$\frac{dZ_t}{dt} = v_t(Z_t), \quad Z_0 \sim \mathcal N(0, I_d), ~ t \in [0,1]$$2 and $$\frac{dZ_t}{dt} = v_t(Z_t), \quad Z_0 \sim \mathcal N(0, I_d), ~ t \in [0,1]$$3, ensures improved adaptation to low-dimensional structures (Roy et al., 21 Jan 2026).

3.2 Stochastic RF Sampler

Let $$\frac{dZ_t}{dt} = v_t(Z_t), \quad Z_0 \sim \mathcal N(0, I_d), ~ t \in [0,1]$$4 approximate $$\frac{dZ_t}{dt} = v_t(Z_t), \quad Z_0 \sim \mathcal N(0, I_d), ~ t \in [0,1]$$5, $$\frac{dZ_t}{dt} = v_t(Z_t), \quad Z_0 \sim \mathcal N(0, I_d), ~ t \in [0,1]$$6 approximate the score, and $$\frac{dZ_t}{dt} = v_t(Z_t), \quad Z_0 \sim \mathcal N(0, I_d), ~ t \in [0,1]$$7: $v_t(x) = \mathbb E[X_1 - X_0 | X_t = x]$6 This stochastic step matches the reverse SDE of DDPM under an appropriate time change and provides robustness to drift estimation errors (Roy et al., 21 Jan 2026).

3.3 Momentum/Stochastic Velocity Field Sampling

An alternative approach is Discretized-RF (“momentum flow matching”), which partitions the time interval into $$\frac{dZ_t}{dt} = v_t(Z_t), \quad Z_0 \sim \mathcal N(0, I_d), ~ t \in [0,1]$$8 anchor points and injects noise into velocities rather than the state: $$\frac{dZ_t}{dt} = v_t(Z_t), \quad Z_0 \sim \mathcal N(0, I_d), ~ t \in [0,1]$$9 This yields hybrid trajectories that interpolate between strict straight lines and stochastic diffusive paths (Ma et al., 10 Jun 2025).

4. Theoretical Guarantees and Complexity

The central theoretical result is that, for data supported on a $v_t(x) = \mathbb E[X_1 - X_0 | X_t = x]$0-dimensional manifold (with certain covering assumptions), the stochastic RF sampler achieves an iteration complexity of $v_t(x) = \mathbb E[X_1 - X_0 | X_t = x]$1 up to logarithmic factors for total variation convergence (Roy et al., 21 Jan 2026). This is in contrast to standard high-dimensional samplers. The introduction of stochasticity mitigates the strict smoothness and drift-estimation constraints of deterministic RF, as noise self-corrects accumulated errors.

Moreover, the stochastic RF dynamics can be rigorously connected to stochastic localization and DDPM. This connection allows sampling in RF to inherit both the numerical stability and convergence guarantees typical of modern diffusion models—even under modest drift and score approximation error (Roy et al., 21 Jan 2026).

5. Empirical Outcomes and Practical Considerations

Empirical studies have validated that:

• On synthetic low-rank Gaussians, both deterministic and stochastic RF with U-shaped time grids outperform uniform grids dramatically in terms of total variation error.
• Stochastic RF produces higher-fidelity and less hallucinated images in text-to-image synthesis with pretrained RF models (e.g., Flux), particularly in low-step regimes and for challenging prompts (Roy et al., 21 Jan 2026).
• Momentum-based Discretized-RF (with $v_t(x) = \mathbb E[X_1 - X_0 | X_t = x]$2, $v_t(x) = \mathbb E[X_1 - X_0 | X_t = x]$3) improves FID and recall on CelebA-HQ and ImageNet compared to deterministic RF, evidencing better mode coverage and finer detail (Ma et al., 10 Jun 2025).

Recommended hyperparameters include a small $v_t(x) = \mathbb E[X_1 - X_0 | X_t = x]$4 for boundary grid density and moderate stochasticity near the noise boundary, with the nonuniform grid or anchor strategy. Excessive stochasticity can reduce quality, so balancing efficiency and diversity is dataset-dependent.

6. Connections to Diffusion Models and Stochastic Localization

There is a precise equivalence, under time-reparameterization, between the stochastic RF SDE and the reverse-time SDE of DDPM. This equivalence is formalized via the stochastic localization process, which justifies both the algorithmic construction and the convergence/rate results (Roy et al., 21 Jan 2026). Consequently, the stochastic RF sampler unifies rectified flow and score-based diffusion perspectives, offering a principled ODE–SDE interpolation for generative modeling (Ma et al., 10 Jun 2025).

7. Summary and Outlook

The Stochastic RF Sampler is a probabilistically principled, theoretically justified sampling algorithm for rectified flow models, ensuring low-discrepancy, high-diversity samples even for intrinsically low-dimensional targets. It addresses key limitations of deterministic straight-line RF sampling by introducing adaptive stochastic perturbations—ultimately combining the best aspects of flow matching and diffusion generative mechanisms. Its empirical success and flexibility promote its adoption for high-quality text-to-image, image-to-image, and domain transfer tasks on modern large-scale models (Roy et al., 21 Jan 2026, Ma et al., 10 Jun 2025).