Model-Aware Schedules Improve Generation via Fiberwise Optimal Transport

Abstract

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.

Authors 4

Luyi Jia, Boyan Zhang, Yilun Liu, Steffen Rulands

Specification

Official page

arxiv.org/abs/2609.11842

Source:arXiv (Atom API + RSS)·T1·observed 1 h ago·high

Arxiv announce type

new

Source:arXiv (Atom API + RSS)·T1·observed 1 h ago·high

arXiv id

2609.11842

Source:arXiv (Atom API + RSS)·T1·observed 1 h ago·high

Categories

cs.LG, cs.AI

Source:arXiv (Atom API + RSS)·T1·observed 1 h ago·high

PDF

arxiv.org/pdf/2609.11842

Source:arXiv (Atom API + RSS)·T1·observed 1 h ago·high

Primary category

cs.LG

Source:arXiv (Atom API + RSS)·T1·observed 1 h ago·high

Published

11 Sept 2026

Source:arXiv (Atom API + RSS)·T1·observed 1 h ago·high

Each value shows its source, tier and observation time. Conflicting claims are kept side by side and flagged — never averaged. How AI Atlas records facts →

Provenance