Variational Trajectory Optimization of Anisotropic Diffusion Schedules

2026-02-22

AI Summary

This work addresses the limitations of conventional diffusion models that employ isotropic noise schedules, which hinder flexible noise control across different subspaces and thereby constrain both generation efficiency and quality. To overcome this, the authors propose a variational framework featuring a learnable matrix-valued path \( M_t(\theta) \) that enables anisotropic noise scheduling, jointly optimizing the score network and noise strategy at the trajectory level. Key contributions include the first introduction of a learnable matrix-valued noise schedule, derivation of a gradient estimator for the score function with respect to \( \theta \), and the design of an anisotropic second-order Heun backward ODE solver. Experiments demonstrate consistent and significant improvements over the EDM baseline across all tested NFE settings on CIFAR-10, AFHQv2, FFHQ, and ImageNet-64.

Technology Category

• Computer Vision: Diffusion Models for Vision
• Machine Learning: Deep Generative Models & Autoencoders
• Search and Optimization: Learning to Search

Application Category

• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking

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.

Problem

• anisotropic diffusion
• noise schedule
• trajectory optimization
• diffusion models
• score-based generation

Innovation

• anisotropic diffusion
• variational trajectory optimization
• matrix-valued noise schedule
• score-based generative models
• reverse ODE solver