Research map / Image generation / Fast diffusion sampling

Variational Trajectory Optimization of Anisotropic Diffusion Schedules

NeurIPS 2026 · Image generation › Fast diffusion sampling · approach: Diffusion theory and design space · 0 citations (OpenAlex)

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Learns a matrix-valued (anisotropic) noise schedule M_t jointly with the score network via a trajectory-level loss, with a plug-in estimator for schedule gradients and an anisotropic Heun solver, reporting FID gains over EDM on CIFAR-10, AFHQv2, FFHQ and ImageNet-64.

Problem

Standard diffusion models use isotropic (scalar) noise schedules, which allocate denoising effort uniformly across directions; learning a matrix-valued noise trajectory is hard because changing the schedule changes the whole marginal family and the optimal score.

Approach

A trajectory-level weighted score-matching loss with matrix weight W_t(theta) jointly trains the score network and schedule parameters; the schedule gradient is estimated from x-directional derivatives of the network (three backward passes), with a flow parameterization M_t^{1/2} times the score for variance reduction, and an anisotropic generalization of Heun's method for sampling.

Key results

• CIFAR-10 best FID improves from 1.829 (EDM) to 1.803 with learned PCA two-subspace schedule.
• FFHQ best FID improves from 2.374 (EDM) to 2.242 with a learned isotropic schedule; AFHQv2 from 2.042 to 2.010 with DCT anisotropic schedule.
• ImageNet-64 best FID improves from 2.276 to 2.238 with class-conditional DCT anisotropic schedules.

Surprising findings

• On FFHQ, the learned isotropic schedule achieves a better best FID than the anisotropic DCT variant (2.242 vs 2.313).

Limitations

• Gains are small (often under 0.1 FID) and FID is reported as minimum over three seeds.
• Anisotropic schedules are restricted to commuting, orthogonal-projector families (e.g., DCT or PCA subspaces) in experiments.
• Improvements are reported only at the best NFE for some datasets and the EDM baseline is at its recommended settings.

Experimental setup

Datasets: CIFAR-10 AFHQv2 FFHQ ImageNet-64

Benchmarks: CIFAR-10 AFHQv2 FFHQ ImageNet-64

Metrics: FID

Baselines: EDM learned isotropic schedule class-conditional isotropic schedules

Builds on: EDM VDM MuLAN

Closest papers by meaning

• Variational Control for Guidance in Diffusion Models (ICML 2025 · similarity 0.88)
• Score-Optimal Diffusion Schedules (NeurIPS 2024 · similarity 0.87)
• On the Trajectory Regularity of ODE-based Diffusion Sampling (ICML 2024 · similarity 0.87)
• Test-Time Scaling of Diffusion Models via Noise Trajectory Search (NeurIPS 2025 · similarity 0.87)
• Align Your Steps: Optimizing Sampling Schedules in Diffusion Models (ICML 2024 · similarity 0.87)
• The Geometry of Noise: Why Diffusion Models Don't Need Noise Conditioning (NeurIPS 2026 · similarity 0.87)