Skip to content
AI Atlas
PaperActive

Model-Aware Schedules Improve Generation via Fiberwise Optimal Transport

arxiv.org/abs/2609.11842

Updated 30 min ago · first seen 11 Sept 2026

paper_01M294FPDW2M7ZPVC0N5EJAV1R

Published
11 Sept 2026
T1 · 31 min ago
arXiv
2609.11842
T1 · 31 min ago
Category
cs.LG
T1 · 31 min ago

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

Source:arXiv (Atom API + RSS)T1observed 31 min agohigh

Arxiv announce type
new

Source:arXiv (Atom API + RSS)T1observed 31 min agohigh

arXiv id
2609.11842

Source:arXiv (Atom API + RSS)T1observed 31 min agohigh

Categories
cs.LG, cs.AI

Source:arXiv (Atom API + RSS)T1observed 31 min agohigh

PDF

Source:arXiv (Atom API + RSS)T1observed 31 min agohigh

Primary category
cs.LG

Source:arXiv (Atom API + RSS)T1observed 31 min agohigh

Published
11 Sept 2026

Source:arXiv (Atom API + RSS)T1observed 31 min agohigh

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

Attributed facts

9

Source tiers

T19

Freshest observation

31 min ago

Conflicts

None