Variational Trajectory Optimization of Anisotropic Diffusion Schedules
This work addresses a domain-specific problem in generative modeling for researchers and practitioners, offering incremental improvements in diffusion model performance.
The paper tackles the problem of optimizing anisotropic noise schedules in diffusion models by introducing a variational framework that jointly trains the score network and learns a matrix-valued schedule, resulting in consistent improvements over baseline models across multiple datasets like CIFAR-10 and ImageNet-64.
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.