Model-Aware Schedules Improve Generation via Fiberwise Optimal Transport
[Uczenie maszynowe] Autorzy: Luyi Jia, Boyan Zhang, Yilun Liu i 1 więcej
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 estima…
- Kategoria arxiv
- cs.LG
- Kategoria
- Uczenie maszynowe