RayFlow: Instance-Aware Diffusion Acceleration via Adaptive Flow Trajectories

1. Introduction

• Traditional diffusion needs many sampling steps for high quality, while shared-Gaussian convergence can overlap paths and make outcomes unstable.The paper identifies varying expectations across timesteps as another obstacle to step reduction without quality degradation.Source: §1 · p. 1
• Rectified flow [23] and rectified diffusion [44] use straight and first-order curved sample-noise paths, respectively, but such matching methods have stated drawbacks.The paper identifies path inconsistency, reduced diversity, and a lack of theoretical derivation for sampling stability.Source: §1 · p. 1
• RayFlow sends each sample along a unique path toward its target mean, reducing path overlap and sampling randomness while targeting stable reconstruction.The paper says its design maximizes path probability between the starting point, target mean, and origin.Source: §1 · p. 2
• Time Sampler uses Stochastic Stein Discrepancies and neural networks to select crucial training timesteps and reduce computational redundancy.It approximates an optimal timestep distribution intended to minimize variance of the training-loss estimator.Source: §1 · p. 2
• The introduction also presents practical training and sampling algorithms, including a one-step variant, and a theoretical analysis of path probabilities and sampling stability.The stated contributions include optimizing parameters for sampling stability.Source: §1 · p. 2
• Experiments report that RayFlow improves image quality, generation efficiency, and controllability compared with existing acceleration algorithms.The introduction provides no quantitative values for these comparisons.Source: §1 · p. 2

2.1. Denoising Diffusion Probabilistic Models [10]

• DDPMs learn a data distribution through T discrete noising steps followed by iterative denoising to reconstruct the original sample.The forward process reaches a standard Gaussian after T steps.Source: §2.1 · p. 2

2.2. Simulation-Free Flow Matching

• Flow matching learns an ODE mapping from a noise distribution to a data distribution by regressing a marginal velocity field.The paper states that diffusion and flow-matching training differ only by a weighting factor.Source: §2.2 · p. 3

2.3. Flow Trajectories

• The trajectory examples include variance-preserving diffusion and rectified flow, whose forward paths use square-root interpolation and linear interpolation, respectively.The cited passages specify the corresponding forward processes and marginal velocity fields.Source: §2.3 · p. 3

2.4. Importance Sampling

• Importance sampling draws from an alternative distribution when standard Monte Carlo is inefficient, with the aim of reducing estimator variance.The estimator is described as unbiased, and its sampling distribution must cover regions where p(x)ξ(x) is nonzero.Source: §2.4 · p. 3

3.1. RayFlow

• RayFlow transforms data through a Markov chain toward a target distribution N(ϵµ, σ^2I), where ϵµ is a pretrained mean vector.The method describes forward and backward processes and a probability-flow path.Source: §3.1 · pp. 3–4
• Maximizing the path probability—the probability of a forward-and-backward transformation—provides optimal target-distribution parameters and can reduce sampling instability.The theorem gives a way to find parameters that maximize this probability.Source: §3.1 · p. 4

3.2. Timestep Sampling

• Uniform timestep sampling is described as inefficient because of high variance and redundant computation; Time Sampler targets a distribution that minimizes the variance of the training-loss estimate.The optimal sampling distribution is presented in Proposition 3.Source: §3.2 · p. 4
• Time Sampler combines SSD with a neural network to approximate a data-dependent optimal timestep distribution despite limited samples.The network parameterizes particle locations, which are iteratively updated by minimizing the KL divergence to the target distribution.Source: §3.2 · pp. 4–5
• The particle-update terms are described as driving particles toward the target distribution while preventing them from getting too close.The update aims to minimize distance until convergence.Source: §3.2 · p. 5

• Algorithm 2 is titled RayFlow Distillation Sampling.Source: §17 · p. 5
• The procedure takes sampling steps K and prompt c as inputs.Source: §17 · p. 5
• The procedure outputs a sampling result ˆx0.Source: §17 · p. 5

• Algorithm 3 is titled RayFlow One-step Sampling.Source: §5 · p. 5
• The procedure takes prompt c as its input.Source: §5 · p. 5
• The procedure outputs a sampling result ˆx0.Source: §5 · p. 5

3.3. Algorithms

• Training constructs synthetic noise-sample pairs, then trains a denoiser to remove noise at each timestep using learned time weights.The Time Sampler predicts important denoising timesteps.Source: §3.3 · p. 6
• Standard sampling starts from pure noise ˆxK and repeatedly predicts noise to refine the output over K denoising steps.Source: §3.3 · p. 6
• The one-step variant generates the final output directly with a single denoising step at timestep T.The algorithms section identifies this as the efficiency-oriented alternative.Source: §1 · p. 5, §3.3 · p. 6

4.1. Implementation Details

• Experiments train on curated LAION/COYO subsets and evaluate four datasets using SD15, SDXL, and PixArt with aesthetic, CLIP, Image Reward, and Pick Score metrics.Source: §4.1 · p. 6

4.2. Main Results

• Across datasets and metrics, adding Time Sampler outperforms its counterpart, and its learned distribution adaptively focuses on different diffusion timesteps.Source: §4.2 · p. 6, Figure 4 · p. 6
• On SDXL/COCO-5k, σ=0.3 performs better than 0.1 or 0.5; the authors suggest moderate variance balances diversity and sampling stability.Source: §4.2 · pp. 6–7, Figure 3 · p. 6
• Across SD15, SDXL, and PixArt, RayFlow variants generally lead the reported comparisons, with SDXL-PeRFlow stronger at one step and SDXL-Ray better in multi-step settings.PixArt-Ray has the best FID and aesthetic scores; SD15-Ray performs best across 1–8 sampling steps.Source: §4.2 · p. 7, Table 1 · p. 7
• Against ReFlow, PeRFlow, RDiff, and InstaFlow, RayFlow leads most SD15 metrics and achieves state-of-the-art on SDXL in two-step comparisons; LoRA scores better in Image Reward than UNet.Source: §4.2 · p. 7, Table 2 · p. 7
• In diverse text-to-image scenes, one-NFE RayFlow images are comparable or superior to multi-step models, with fine-detail and visual-coherence strengths.Source: §4.2 · pp. 7–8, Figure 5 · p. 8
• Combining ϵµ and Tϑ achieves the best results across SD15, SDXL, and PixArt; ϵµ alone yields limited gains, while Tϑ alone remains suboptimal.Source: §4.2 · p. 8, Table 3 · p. 8

5. Conclusion

• RayFlow guides samples along unique paths toward precomputed target means, while Time Sampler selects crucial training timesteps; experiments report gains in quality, speed, and computational efficiency.The conclusion also identifies extending RayFlow to other modalities and improving target-mean determination as future work.Source: §5 · p. 8

A.1. Flow Trajectory

• The flow-trajectory appendix describes a forward process that adds noise to image data and a backward process that denoises it.Its derivation uses a target distribution with mean ϵµ and variance σ²I.Source: §A.1 · pp. 13–14

A.2. Probability Path

• The probability-path appendix defines the path probability as the forward probability of reaching the target multiplied by the backward probability of returning to the starting sample.For the intermediate conditional distribution, it states that the Gaussian mean and variance are the parameters to calculate.Source: §A.2 · pp. 16–17

A.3. Optimal Probability Path

• For a sample and target distribution, the appendix identifies optimal parameters that maximize the path probability.The optimal path is stated to be the optimal probability path.Source: §A.3 · pp. 18, 20

A.4. Optimal Denoiser

• The denoising-loss optimization is convex, with a unique solution identified by setting the gradient with respect to ϵθ to zero.Source: §A.4 · pp. 20–21

B.1. Stein Identity

• The Stein identity states that the expected Stein operator expression is zero for a bounded, smooth function, with the score defined as the gradient of the log density.Source: §B.1 · p. 21

B.2.2 Stein Discrepancy Measure

• The Stein discrepancy is positive when the distributions differ, but is generally difficult to compute directly.Source: §B.2.2 · p. 22

B.2.3 RKHS Framework

• The appendix introduces RKHS theory to make the discrepancy measure computationally tractable.Source: §B.2.3 · p. 22

B.2.4 Kernelized Stein Discrepancy (KSD)

• The KSD formulation provides a computationally tractable kernel-based measure of discrepancy between distributions.Source: §B.2.4 · p. 22

B.3. Optimal sampling Distribution

• The appendix states an optimal sampling distribution for Eq.(5) that minimizes variance.The result is described as applying to any probability distribution p.Source: §B.3 · pp. 22–23
• The proof formulates the variance objective with a normalized sampling distribution and derives a condition by differentiating with respect to q(t).Source: §B.3 · pp. 22–23
• The proof uses the Cauchy-Schwarz inequality to establish that q*(t|x0, ϵµ) minimizes variance for the stated problem.Source: §B.3 · p. 23