PaperlayerView exact revision on arXiv

Source-linked AI summary

Model-Aware Schedules Improve Generation via Fiberwise Optimal Transport

Luyi Jia, Boyan Zhang, Yilun Liu, Steffen Rulands

arXiv:2609.11842v1cs.LGcs.AI

TL;DR

Model-agnostic kinetic schedules ignore prediction error, so the paper introduces fiberwise optimal-transport risk to construct model-aware traversal schedules. Combining this risk with kinetic action gives closed-form allocation that improves DDPM and flow-matching generation, while normalized profiles support reusable analytic templates.

  • Problem

    Existing kinetic schedule baselines are model-agnostic and ignore prediction error, motivating a model-dependent correction based on signal/noise decomposition transport.

  • Method

    The method averages base-preserving optimal-transport costs between true and predictor-induced decompositions within affine fibers, then combines the resulting risk with coefficient-path kinetics for closed-form time allocation.

  • Results

    Across DDPM and flow-matching settings, model-aware schedules consistently improve strong baselines, including a 38.6% relative FID reduction for flow matching on CIFAR-10 at 16 function evaluations.

  • Takeaways & Limitations

    Aligned normalized risk profiles and allocation deformations motivate a frozen analytic schedule template that retains model-aware improvement without further risk estimation or model-specific fitting.

  • Takeaways & Limitations

    The shared cross-system profile shape remains theoretically unexplained, and end-to-end evaluation of pretrained larger conditional models requires retraining under the resulting schedules.

Abstract

from arXiv · show

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.

1 INTRODUCTION

The paper introduces a model-aware schedule that augments model-agnostic coefficient-path kinetics with fiberwise prediction risk. This yields closed-form time allocation, supports general prediction targets, and reveals aligned normalized risk and allocation profiles across settings.

  • 1 INTRODUCTION: Fiberwise prediction risk measures optimal-transport discrepancy between true and predictor-induced signal/noise decompositions at fixed states.The construction preserves the marginal distribution while comparing compatible decompositions within affine fibers.
  • 1 INTRODUCTION: Closed-form schedule allocation combines fiberwise prediction risk with coefficient-path kinetic action on a fixed coefficient curve.The risk term reduces time in higher-risk regions, while kinetics prevent arbitrarily rapid traversal.
  • 1 INTRODUCTION: Normalized fiberwise-risk profiles and training-time allocation deformations align across independently trained models and settings in kinetic reference coordinates.These findings motivate a frozen analytic allocation template for reusable schedule design.
  • 1 INTRODUCTION: The risk formulation extends to general linear prediction targets and can be estimated directly from corresponding prediction errors.Under the normalized symmetric product metric, it reduces to a weighted noise-prediction MSE.
  • 1 INTRODUCTION: The schedules consistently improve DDPM and flow-matching baselines across prediction targets, configurations, datasets, and architectures.The reported flow-matching improvement reaches a 38.6% relative FID reduction at 16 function evaluations.

2 RELATED WORK

Related work studies schedule, path, solver-grid, and kinetic-energy design from complementary perspectives. The paper instead estimates fiberwise risk once from a baseline checkpoint and uses it to construct a fixed schedule.

  • 2 RELATED WORK: Prior methods optimize inference grids, trajectory regularity, conditional entropy, or velocity fields, whereas this work targets schedule construction.BOSS combines dynamic programming for grid selection with velocity-field fine-tuning.
  • 2 RELATED WORK: Existing schedule methods include log-SNR importance sampling, online rate equalization, prescribed sine traversals, and spatial-Lipschitz traversal optimization.The paper distinguishes these approaches from its fixed, baseline-checkpoint-based schedule estimation.
  • 2 RELATED WORK: Optimal-transport analyses have optimized coefficient-path or marginal kinetic energy, including Cond-OT baselines and jointly optimized paths and traversals.The cited flow-matching evaluations use independent endpoints rather than minibatch OT coupling.

3 BACKGROUND

The paper represents diffusion and flow-matching processes as affine paths whose coefficient curves can be traversed at different speeds. Coefficient-path kinetics provide model-agnostic baselines and a reference coordinate for schedule comparisons.

  • 3 BACKGROUND: Affine probability paths combine data and Gaussian noise as x_t = m_tx_0 + s_tϵ, while schedules specify the coefficient path over time.For VP/DDPM, m_t = √ᾱ_t and s_t = √(1 − ᾱ_t).
  • 3 BACKGROUND: The coefficient curve is geometric and independent of traversal speed, so schedule design includes choosing how quickly to move along that curve.Different endpoint couplings can preserve the affine representation while changing the prediction problem.
  • 3 BACKGROUND: Figure 1 summarizes how within-fiber decomposition discrepancies and coefficient-path kinetics jointly determine model-aware schedule allocation.Its marker spacing is schematic rather than a quantitative axis encoding.
  • 3 BACKGROUND: Coefficient-path kinetics use the Euclidean action of coefficient motion and connect to expected pathwise kinetic action under stated endpoint assumptions.With independent endpoints and matched second moments, the two actions coincide after normalization.
  • 3 BACKGROUND: Minimizing kinetic action yields constant-speed traversal of the shortest admissible coefficient curve, defining the kinetic reference coordinate τ.The resulting baselines are Cond-OT and, under VP constraints, the exact cosine parameterization.

4 FIBERWISE PREDICTION RISK VIA OPTIMAL TRANSPORT

Fiberwise prediction risk compares true and predictor-induced decompositions that share the same observed state. Base-preserving optimal transport averages this within-fiber discrepancy and yields an error-estimable risk profile applicable across prediction targets.

  • 4.1 DECOMPOSITION FIBERS AND FIBERWISE OPTIMAL TRANSPORT: Affine fibers F_x collect all signal/noise decompositions whose components sum to a fixed observed state x.The decomposition bundle groups these fibers and projects each decomposition back to its base state.
  • 4.1 DECOMPOSITION FIBERS AND FIBERWISE OPTIMAL TRANSPORT: At fixed (t, x), true and predictor-induced decompositions differ only vertically within the same fiber.The horizontal component instead describes motion of the observed state along the probability path.
  • 4.1 DECOMPOSITION FIBERS AND FIBERWISE OPTIMAL TRANSPORT: Fiberwise risk averages squared Wasserstein transport costs from the conditional law of true decompositions to the predictor-induced point mass within each fiber.This is the conditional mean squared vertical discrepancy averaged over x drawn from p_t.
  • 4.2 RISK FOR PREDICTION TARGETS: Under the normalized symmetric product metric, the fiberwise risk becomes a weighted noise-prediction MSE that is directly estimable from prediction errors.The base-preserving construction leaves the marginal p_t unchanged.
  • 4.2 RISK FOR PREDICTION TARGETS: For a general linear target y_t = a_tx_0 + b_tϵ, changing the prediction target changes only the coordinate expression of the same fiberwise risk.The formulation covers the DDPM and flow-matching targets considered in the paper.
  • 4.2 RISK FOR PREDICTION TARGETS: The risk is expressed as R(τ) in each model-agnostic kinetic reference coordinate for use in schedule-time allocation.This coordinate supports comparisons across models and settings.

5 MODEL-AWARE SCHEDULE DESIGN

The model-aware formulation minimizes coefficient-path kinetic action under a schedule-time-integrated fiberwise prediction-risk budget. On a fixed coefficient curve, it converts an estimated risk profile into a closed-form traversal allocation while preserving practical DDPM and flow-matching reference coordinates.

  • 5.1 RISK-CONSTRAINED KINETIC FORMULATION: The formulation combines coefficient-path kinetic action with a model-dependent fiberwise prediction-risk constraint, rather than optimizing kinetic action alone.Under stated conditions, the objective can also be viewed through a tractable Wasserstein generation-error sensitivity bound.
  • 5.2 FIXED-CURVE MODEL-AWARE ALLOCATION: The one-shot construction estimates risk from a baseline checkpoint, holds it fixed during retraining, and typically changes allocations only marginally when risk is re-estimated.This supports using a frozen risk profile for schedule construction.
  • 5.2 FIXED-CURVE MODEL-AWARE ALLOCATION: On a fixed coefficient curve, the KKT condition yields a closed-form allocation density that maps an estimated risk profile into schedule time.The risk profile and coefficient-space speed are held fixed while traversal is optimized, with normalization enforced by a multiplier.
  • 5.2 FIXED-CURVE MODEL-AWARE ALLOCATION: Higher risk receives less schedule time when λ > 0 in the kinetic reference coordinate, while the kinetic term prevents arbitrarily rapid traversal.The resulting deformation depends jointly on the risk profile, its absolute scale, and λ relative to the kinetic term.
  • 5.2 FIXED-CURVE MODEL-AWARE ALLOCATION: For DDPMs, the method deforms traversal of an offset-normalized cosine approximation to the VP kinetic baseline, while flow matching uses Cond-OT time as its exact reference coordinate.Reparameterization changes both the state–time map and the linear velocity prediction target.

6 EXPERIMENTS

Across DDPM and flow-matching experiments, model-aware schedules improve FID across datasets, samplers, targets, architectures, and training configurations. Normalized risk profiles and allocation deformations also align across evaluated models, supporting a reusable analytic schedule template.

  • Main results: 16.4% FID reduction is reached by DDPM on CIFAR-10 at 16 NFE with DPM++3M, while all five reported configurations improve on both CIFAR-10 and ImageNet-64.ImageNet-64 reductions at 16 NFE range from 0.91 to 1.75 FID, with DPM++3M also improving at 32 and 64 NFE.
  • Main results: 38.6% mean relative FID reduction is achieved for CIFAR-10 flow matching, while all four CIFAR-10 configurations and all seven ImageNet-64 configurations improve at the reported checkpoints.CIFAR-10 reductions span 6.9%–38.6%; ImageNet-64 midpoint and Heun3 improve by 24.7%–29.8%.
  • Target and architecture transfer: Model-aware schedules transfer across prediction targets and architectures, improving all six DDPM v-prediction configurations by 7.4%–14.7% and all five U-ViT-S/2 configurations.The U-ViT-S/2 result includes a 12.2% reduction at 16 NFE without architecture-specific retuning.
  • Coupling and training-time sampling: Across endpoint couplings and training-time sampling distributions, model-aware schedules reduce FID by 31.3%–39.4% while preserving the baseline ordering of these design choices.OT coupling remains better than independent coupling, and logit-normal, uniform, and U-shaped sampling retain their ordering.
  • Robustness: Risk estimates from epoch 100 preserve the shared normalized risk shape, changing final FID by only 0.02 for DDPM and 0.06 for flow matching.The reported gains persist across tested training stages and nonzero λ values.
  • Shared risk-profile shapes and allocation deformations: The frozen analytic template recovers most model-aware gains across datasets and model families, including 109.7% of the DDPM gain on ImageNet-64.Recovery is 98.1%, 98.0%, and 75.8% for the three CIFAR-10 settings, and 76.2% for ImageNet-64 OT-CFM.

7 DISCUSSION AND LIMITATIONS

The paper reports that combining coefficient-path kinetics with fiberwise prediction risk improves generation quality across evaluated models and settings. Shared normalized risk profiles motivate a frozen analytic schedule, while theoretical explanation and broader end-to-end evaluation remain open.

  • Model-aware schedules improve generation quality across models, training settings, and sampling settings.
  • Shared risk shapes and allocation deformations suggest empirical universality and motivate a usable analytic schedule without further risk estimation or model-specific fitting.
  • The shared profile shape and quantitative cross-system agreement remain theoretically unexplained.
  • End-to-end evaluation for larger conditional latent diffusion and 2-RF models requires retraining them under the resulting schedules.

AI USE STATEMENT

The authors disclose that ChatGPT assisted with editing, coding, figure generation and refinement, and literature research, but not with research ideation or core scientific decisions.

  • ChatGPT assisted with text editing, coding, Figure 1 generation and refinement, and literature research, with all outputs reviewed and verified by the authors.
  • ChatGPT was not involved in research ideation, theoretical results, methodology, experimental design, or citation selection and verification.

A.1 RELATING PATHWISE ACTION TO COEFFICIENT-PATH KINETICS UNDER GENERAL ENDPOINT JOINT LAWS

For fixed endpoint joint laws, coefficient-path kinetics relates to pathwise kinetic action through bounds or exact proportionality factors. The standard Cond-OT path minimizes coefficient-path kinetics, while the base–fiber split isolates decomposition discrepancies from base motion.

  • For a fixed nondegenerate endpoint joint law, coefficient-path kinetics minimizes an upper bound on pathwise action, but not necessarily pathwise action itself.
  • Under independent, matched-second-moment endpoints, the VP constraint preserves the common second-moment scale.
  • For any fixed endpoint joint law, endpoint moments and coupling affect only the proportionality factor relating pathwise action to coefficient-path kinetics along a fixed curve.
  • The standard Cond-OT coefficient path minimizes both coefficient-path kinetics and pathwise action when only coefficient endpoints are prescribed.
  • The base–fiber tangent split makes same-base decomposition discrepancies purely vertical, while expected horizontal kinetic energy is one quarter of expected pathwise kinetic energy.

B.2 FULL FIBERWISE OPTIMAL-TRANSPORT DERIVATION

The appendix derives fiberwise prediction risk as base-preserving optimal transport between decomposition laws sharing the same state marginal. It then derives the fixed-curve allocation rule through constrained optimization and KKT conditions.

  • Fiberwise prediction risk compares true and predictor-induced decomposition laws within each affine fiber over a shared base state.
  • The total-space optimal-transport formulation is equivalent because infinite cross-base cost restricts finite-cost couplings to the same fiber.
  • The fixed-curve schedule problem combines kinetic allocation with schedule-time-integrated fiberwise risk under normalization and risk constraints.
  • KKT conditions show that an inactive risk constraint recovers the kinetic baseline, whereas an active constraint selects a point on the kinetic–risk tradeoff.
  • The discrete stationarity condition yields the coefficient-space allocation rule used by the one-shot construction.

C.3 EFFECTIVE GEOMETRY AND RISK-RECALIBRATION PROPERTIES

The optimal allocation traverses a fixed coefficient curve at constant speed under a risk-modified geometry. Its deformation depends on the risk profile's shape, not its absolute affine calibration.

  • Effective geometry: The model-aware optimum is constant-speed traversal under a risk-modified coefficient-space line element.Cumulative risk-modified length determines the schedule time assigned to each reference location.
  • Risk recalibration: Positive affine recalibration of the risk profile leaves the optimal allocation unchanged.Absolute risk scale and offset do not by themselves determine the schedule deformation.
  • Risk recalibration: Identical allocation densities imply that two risk profiles agree up to an additive offset absorbed by normalization.

D.1 MODEL-AWARE SCHEDULE CONSTRUCTION AND TRAINING PIPELINE

The one-shot construction estimates fiberwise risk from a baseline checkpoint, converts interval risks and coefficient-space lengths into a closed-form reparameterization, and retrains from scratch with that schedule fixed. Across prediction targets, architectures, samplers, and coupling settings, the resulting model-aware schedules improve the reported FID comparisons.

  • Construction pipeline: The one-shot pipeline estimates risk from a fixed baseline checkpoint, constructs a monotone reparameterization, and trains a new model from scratch with it fixed.The baseline checkpoint is used for schedule construction rather than continued training.
  • Construction pipeline: The interval allocation combines coefficient-space arc length with estimated fiberwise risk before cumulative allocations define t = Φ(τ).The implementation evaluates representative interval risks and solves for the allocation normalization parameter.
  • DDPM implementation: The DDPM implementation uses the practical VP reference coordinate and preserves the baseline timestep-sampling distribution while changing diffusion coefficients.The native grid has T = 1000 steps, with standard βj capping and matched zero-terminal-SNR processing for v-prediction.
  • Flow matching implementation: The flow-matching implementation fixes the Cond-OT curve and changes its traversal while holding endpoint coupling and training-time sampling distribution fixed.Reparameterization changes both the state–time map and the corresponding velocity target.
  • Robustness results: The model-aware schedule improves every reported v-prediction configuration, with relative FID reductions from 7.4% to 14.7%.The gains persist across second- and third-order DPM-Solver++ multistep solvers and therefore under a different prediction target.
  • Robustness results: The schedule also improves all five architecture-control settings and all four alternative flow-matching settings, with relative FID reductions reaching 39.4%.The architecture control includes a U-ViT-S/2 model, while the flow-matching comparisons vary endpoint coupling and training-time sampling distributions.

E.6 TRAINING-STAGE AND TRADEOFF-WEIGHT SENSITIVITY

Model-aware gains persist across training stages, tradeoff weights, checkpoints, integrators, datasets, and pretrained diagnostics, while normalized risk and allocation shapes remain broadly aligned. The analytic template recovers most DDPM gains but leaves additional flow-matching gains from setting-specific risk structure.

  • Sensitivity: Model-aware schedules improve across tested training stages and nonzero tradeoff weights, including an approximately 3-FID flow-matching improvement across epochs 400–900.At epoch 400, all tested nonzero λ values improve over the standard Cond-OT parameterization.
  • Checkpoint sensitivity: Epoch-100 risk estimates retain essentially the full later-checkpoint improvement, changing final FID by only 0.02 for DDPM and 0.06 for flow matching.The risk profile remains sufficiently stable across substantially different training stages for the resulting schedule to remain effective.
  • Finite-step integration: Finite-step effects are integrator- and setting-dependent: model-aware schedules help higher-order integrators on CIFAR-10, worsen Euler there, and improve Euler on ImageNet-64.The ImageNet-64 result indicates that the CIFAR-10 Euler degradation is not a systematic first-order incompatibility.
  • Cross-system agreement: Normalized risk profiles and allocation deformations remain closely aligned across prediction targets, checkpoints, architectures, endpoint couplings, sampling distributions, and both evaluated datasets.Pretrained DiT-XL/2 and InstaFlow 2-RF diagnostics also retain close normalized-risk agreement without retraining or modified-schedule FID evaluation.
  • Asymmetry: Risk-profile asymmetry shifts normalized risk centers toward the data side while allocation centers move toward the noise side, consistent with reciprocal-square-root risk weighting.The asymmetry is retained by model-aware allocations and is absent from the symmetric analytic template.
  • Analytic template: The analytic template recovers 98.1% and 98.0% of CIFAR-10 DDPM gains, 75.8% for matched OT-CFM, 109.7% for ImageNet-64 DDPM, and 76.2% for ImageNet-64 OT-CFM.The results indicate that shared deformation explains most DDPM improvement while flow matching retains additional setting-specific gains.
Loading 2609.11842v1…