Saved in:
Bibliographic Details
Main Authors: Liu, Pengxi, Li, Zeyu Michael, Cheng, Xiang
Format: Preprint
Published: 2026
Subjects:
Online Access:https://arxiv.org/abs/2602.19512
Tags: Add Tag
No Tags, Be the first to tag this record!
_version_ 1866912919934992384
author Liu, Pengxi
Li, Zeyu Michael
Cheng, Xiang
author_facet Liu, Pengxi
Li, Zeyu Michael
Cheng, Xiang
contents 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.
format Preprint
id arxiv_https___arxiv_org_abs_2602_19512
institution arXiv
publishDate 2026
record_format arxiv
spellingShingle Variational Trajectory Optimization of Anisotropic Diffusion Schedules
Liu, Pengxi
Li, Zeyu Michael
Cheng, Xiang
Machine Learning
Computer Vision and Pattern Recognition
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.
title Variational Trajectory Optimization of Anisotropic Diffusion Schedules
topic Machine Learning
Computer Vision and Pattern Recognition
url https://arxiv.org/abs/2602.19512