Research map / Image generation / Fast diffusion sampling
Bellman Optimal Stepsize Straightening of Flow-Matching Models
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes BOSS, which uses a dynamic-programming search over Euler timestamps to find Bellman-optimal stepsizes for a pretrained rectified-flow model, then straightens the velocity network on those stepsizes with about 12,000 finetuning iterations.
Problem
Rectified-flow/flow-matching models need many function evaluations for sampling, and the standard reflow distillation needs very costly retraining (hundreds of thousands of iterations) to reach good few-step quality.
Approach
Select the Euler timestamp schedule for a fixed NFE budget by minimizing an estimated local truncation error over anchor timestamps via a dynamic program (Bellman recursion), then finetune the velocity network so that the K-step Euler path matches the fine-grained uniform trajectory on that schedule.
Key results
- Bellman Euler reduces FID vs uniform Euler at 4 NFE on CelebA-HQ from 158.95 to 92.03 and on LSUN-Bedroom from 84.35 to 61.60.
- BOSS-6 achieves FID 18.67 on CelebA-HQ vs 43.57 for Uniform-Reflow and 34.56 for Distill-6-Reflow.
- On CIFAR-10, BOSS-6 (4.80) is slightly worse than Uniform-Reflow (4.33) and Distill-6-Reflow (4.35).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Replace uniform Euler stepsizes with Bellman optimal stepsizes | pretrained Rectified Flow · CIFAR-10, 4/6/8 NFE | FID: 51.95 -> 47.57; 25.69 -> 23.35; 16.82 -> 15.74 | Helped |
| Replace uniform Euler stepsizes with Bellman optimal stepsizes | pretrained Rectified Flow · CelebA-HQ, 4/6/8 NFE | FID: 158.95 -> 92.03; 127.01 -> 72.54; 109.42 -> 49.80 | Helped |
| Replace uniform Euler stepsizes with Bellman optimal stepsizes | pretrained Rectified Flow · LSUN-Church, 4/6/8 NFE | FID: 106.94 -> 80.91; 53.85 -> 45.09; 34.74 -> 33.22 | Helped |
| Replace uniform Euler stepsizes with Bellman optimal stepsizes | pretrained Rectified Flow · AFHQ-Cat, 4/6/8 NFE | FID: 68.95 -> 45.54; 61.50 -> 36.15; 56.96 -> 33.94 | Helped |
| Replace uniform Euler stepsizes with Bellman optimal stepsizes | pretrained Rectified Flow · LSUN-Bedroom, 4/6/8 NFE | FID: 84.35 -> 61.60; 39.19 -> 35.35; 32.15 -> 25.80 | Helped |
| Apply Bellman stepsize straightening instead of the strongest baseline in the table | 6-step flow-matching retraining · CIFAR-10 | FID: 4.33 -> 4.80 | Hurt |
| Apply Bellman stepsize straightening instead of the strongest baseline in the table | 6-step flow-matching retraining · CelebA-HQ | FID: 34.56 -> 18.67 | Helped |
| Apply Bellman stepsize straightening instead of the strongest baseline in the table | 6-step flow-matching retraining · LSUN-Church | FID: 34.52 -> 17.43 | Helped |
| Apply Bellman stepsize straightening instead of the strongest baseline in the table | 6-step flow-matching retraining · AFHQ-Cat | FID: 46.24 -> 26.10 | Helped |
Surprising findings
- On CIFAR-10, BOSS-6 does not beat Uniform-Reflow or Distill-6-Reflow, performing slightly worse.
- Bellman stepsizes transferred from LSUN-Church still beat uniform Euler on CelebA-HQ and LSUN-Bedroom.
Limitations
- Finding the Bellman stepsizes still requires an additional training/computation cost.
- On CIFAR-10 the method is not better than standard reflow baselines.
- Uses an approximation of local truncation error computed on the Euler trajectory with a single batch (N=100).
Experimental setup
Datasets: CIFAR-10 CelebA-HQ LSUN-Church LSUN-Bedroom AFHQ-Cat
Benchmarks: CIFAR-10 CelebA-HQ LSUN-Church LSUN-Bedroom AFHQ-Cat
Metrics: FID
Baselines: Uniform Euler Uniform Heun RK45 Uniform-Reflow Distill-K-Reflow
Builds on: Rectified Flow reflow DDPM++ Clean-FID LoRA
Closest papers by meaning
- Align Your Flow: Scaling Continuous-Time Flow Map Distillation (NeurIPS 2025 · similarity 0.89)
- Self-Corrected Flow Distillation for Consistent One-Step and Few-Step Image Generation (AAAI 2025 · similarity 0.89)
- Optimal Flow Matching: Learning Straight Trajectories in Just One Step (NeurIPS 2024 · similarity 0.88)
- FlowTurbo: Towards Real-time Flow-Based Image Generation with Velocity Refiner (NeurIPS 2024 · similarity 0.88)
- Blockwise Flow Matching: Improving Flow Matching Models For Efficient High-Quality Generation (NeurIPS 2025 · similarity 0.88)
- Faster Inference of Flow-Based Generative Models via Improved Data-Noise Coupling (ICLR 2025 · similarity 0.88)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.