Search papers, labs, and topics across Lattice.
The paper introduces a variational framework for diffusion models that optimizes anisotropic noise schedules, parameterized by a matrix-valued path, to allocate noise across subspaces. This framework jointly trains the score network and learns the matrix-valued noise schedule using a trajectory-level objective and a derived estimator for the score's derivative with respect to the schedule parameters. The method also develops an efficient reverse-ODE solver, generalizing the Heun discretization, and demonstrates improved performance over the EDM baseline across several image datasets in various NFE regimes.
Anisotropic noise schedules, optimized via a novel variational framework, can significantly boost diffusion model performance across various datasets and NFE regimes.
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.
