# Model-Aware Schedules Improve Generation via Fiberwise Optimal Transport

arXiv: https://arxiv.org/abs/2609.11842 · HTML: https://mlsanity.com/paper/2609.11842/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in Flow matching, time series and science, one of 14 directions in Diffusion models.
- Builds on 35 of its 43 references that are in the map.
- First posted on arXiv 10 Sep 2026.

## Placed in

- Diffusion models / [Flow matching, time series and science](https://mlsanity.com/diffusion/flow-matching-time-series-and-science/)

## Builds on (most cited first, 30 of 35)

- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [U-Net: Convolutional Networks for Biomedical Image Segmentation](https://mlsanity.com/paper/1505.04597/) (2015)
- [High-Resolution Image Synthesis with Latent Diffusion Models](https://mlsanity.com/paper/2112.10752/) (2021)
- [Microsoft COCO: Common Objects in Context](https://mlsanity.com/paper/1405.0312/) (2014)
- [Denoising Diffusion Implicit Models](https://mlsanity.com/paper/2010.02502/) (2020)
- [Diffusion Models Beat GANs on Image Synthesis](https://mlsanity.com/paper/2105.05233/) (2021)
- [Scalable Diffusion Models with Transformers](https://mlsanity.com/paper/2212.09748/) (2022)
- [Classifier-Free Diffusion Guidance](https://mlsanity.com/paper/2207.12598/) (2022)
- [Flow Matching for Generative Modeling](https://mlsanity.com/paper/2210.02747/) (2022)
- [Scaling Rectified Flow Transformers for High-Resolution Image Synthesis](https://mlsanity.com/paper/2403.03206/) (2024)
- [Improved Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2102.09672/) (2021)
- [Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow](https://mlsanity.com/paper/2209.03003/) (2022)
- [DPM-Solver++: Fast Solver for Guided Sampling of Diffusion Probabilistic Models](https://mlsanity.com/paper/2211.01095/) (2022)
- [Improving and generalizing flow-based generative models with minibatch optimal transport](https://mlsanity.com/paper/2302.00482/) (2023)
- [All are Worth Words: A ViT Backbone for Diffusion Models](https://mlsanity.com/paper/2209.12152/) (2022)
- [InstaFlow: One Step is Enough for High-Quality Diffusion-Based Text-to-Image Generation](https://mlsanity.com/paper/2309.06380/) (2023)
- [A Downsampled Variant of ImageNet as an Alternative to the CIFAR datasets](https://mlsanity.com/paper/1707.08819/) (2017)
- [Common Diffusion Noise Schedules and Sample Steps are Flawed](https://mlsanity.com/paper/2305.08891/) (2023)
- [Multisample Flow Matching: Straightening Flows with Minibatch Couplings](https://mlsanity.com/paper/2304.14772/) (2023)
- [Improving the Training of Rectified Flows](https://mlsanity.com/paper/2405.20320/) (2024)
- [Align Your Steps: Optimizing Sampling Schedules in Diffusion Models](https://mlsanity.com/paper/2404.14507/) (2024)
- [On the Trajectory Regularity of ODE-based Diffusion Sampling](https://mlsanity.com/paper/2405.11326/) (2024)
- [On Kinetic Optimal Probability Paths for Generative Models](https://mlsanity.com/paper/2306.06626/) (2023)
- [Improved Noise Schedule for Diffusion Training](https://mlsanity.com/paper/2407.03297/) (2024)
- [Conditional Wasserstein Distances with Applications in Bayesian OT Flow Matching](https://mlsanity.com/paper/2403.18705/) (2024)
- [LayoutFlow: Flow Matching for Layout Generation](https://mlsanity.com/paper/2403.18187/) (2024)
- [Lipschitz-Guided Design of Interpolation Schedules in Generative Models](https://mlsanity.com/paper/2509.01629/) (2025)
- [Entropic Time Schedulers for Generative Diffusion Models](https://mlsanity.com/paper/2504.13612/) (2025)
- [Optimal Scheduling of Dynamic Transport](https://mlsanity.com/paper/2504.14425/) (2025)
- [Bellman Optimal Stepsize Straightening of Flow-Matching Models](https://mlsanity.com/paper/2312.16414/) (2023)

## Abstract (arXiv)

Diffusion and flow-matching schedules control the signal and noise coefficients that mix data and noise along affine probability paths. Minimizing a kinetic action defined on coefficient paths, motivated by optimal transport, helps explain strong baselines but remains model-agnostic and ignores prediction error. Here we introduce a model-aware schedule construction based on fiberwise optimal transport. At a fixed time and state on the probability path, compatible signal/noise decompositions form an affine fiber. We define a fiberwise prediction risk by averaging optimal-transport costs between the true and predictor-induced decompositions within these fibers. On a fixed coefficient curve, combining this risk with coefficient-path kinetic action yields a closed-form optimal time allocation. This construction extends to general linear prediction targets, and the risk profile can be estimated from an early baseline checkpoint. We evaluate DDPMs and flow matching across prediction targets, training configurations, risk-estimation checkpoints, datasets, and architectures. Our model-aware schedules consistently outperform strong baselines, including a 38.6% relative FID reduction for flow matching on CIFAR-10 at 16 function evaluations. Each model-agnostic kinetic baseline determines its own kinetic reference coordinate. In these coordinates, fiberwise-risk profiles from independently trained models in different settings align closely after normalization to unit area. The resulting schedule deformations used in training also align, suggesting empirical universality across the evaluated models and settings. Pretrained-checkpoint diagnostics extend this normalized-risk agreement to larger conditional latent diffusion and 2-RF models. A frozen analytic allocation template retains most of the model-aware improvement without further risk estimation or model-specific fitting.