Variational Trajectory Optimization of Anisotropic Diffusion Schedules

TL;DR

This paper proposes a variational framework for anisotropic diffusion schedules in generative models, jointly training score networks and noise schedules to improve image generation quality across multiple datasets.

Contribution

It introduces a novel variational approach for anisotropic noise schedules and an efficient reverse-ODE solver, enhancing diffusion model performance.

Findings

Consistent improvement over baseline EDM across datasets

Effective joint training of score network and noise schedule

Enhanced inference with anisotropic reverse-ODE solver

Abstract

We introduce a variational framework for diffusion models with anisotropic noise schedules parameterized by a matrix-valued path $M_t(\theta)$ that allocates noise across subspaces. Central to our framework is a trajectory-level objective that jointly trains the score network and learns $M_t(\theta)$, which encompasses general parameterization classes of matrix-valued noise schedules. We further derive an estimator for the derivative with respect to $\theta$ of the score that enables efficient optimization of the $M_t(\theta)$ 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.