Variational Trajectory Optimization of Anisotropic Diffusion Schedules

A variational framework learns matrix-valued anisotropic diffusion schedules and improves EDM across CIFAR-10, AFHQv2, FFHQ, and ImageNet-64.

Pengxi Liu, Zeyu M Li, Xiang Cheng

Published 2026Atlanta Poster Session 1 · Wed, Dec 9, 10:00 AM–1:00 PM local time · Hall C1arXiv ↗OpenReview ↗

72%

OverallHighly rated

OverallHighly ratedVote to see the scoreThe exact score shows once you've voted, so every vote is your own call. The first half of each home page shelf shows its scores.

Readers

–

Only vote on papers you've read. Sign in with GitHub to vote.

AI panel8/20reviewers recommend it

lenient 2/5

medium 5/10

strict 1/5

AI panel?Vote to see what the 20 AI reviewers said

Abstract

We introduce a variational framework for diffusion models with anisotropic noise schedules parameterized by a matrix-valued path $M_t(θ)$ that allocates noise across subspaces. Central to our framework is a trajectory-level objective that jointly trains the score network and learns $M_t(θ)$, which encompasses general parameterization classes of matrix-valued noise schedules. We further derive an estimator for the derivative with respect to $θ$ of the score that enables efficient optimization of the $M_t(θ)$ schedule. For inference, we develop an efficiently-implementable reverse-ODE solver that is an anisotropic generalization of the second-order Heun discretization algorithm. Across CIFAR-10, AFHQv2, FFHQ, and ImageNet-64, our method consistently improves upon the baseline EDM model in all NFE regimes. Code is available at https://github.com/lizeyu090312/anisotropic-diffusion-paper.