Variational Trajectory Optimization of Anisotropic Diffusion Schedules

Pengxi Liu; Zeyu Michael Li; Xiang Cheng

cs.LGFeb 23, 2026

Variational Trajectory Optimization of Anisotropic Diffusion Schedules

Authors: Pengxi Liu, Zeyu Michael Li, Xiang Cheng

Organizations: Department of Electrical and Computer Engineering, Duke University

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.

Figures & tables

Figure 1: Illustration of isotropic vs. anisotropic denoising. Top: a standard isotropic sampler denoises all directions uniformly. Bottom: an anisotropic sampler with two DCT subspaces, $V_{1}$ (low frequency) and $V_{2}$ (high frequency) (Section 6 ). Columns show intermediate reconstructions as $t$ decreases. The plot (right) displays learned subspace schedules $g_{1}(t)$ and $g_{2}(t)$ ; the former is denoised more aggressively, so low-frequency structure from $V_{1}$ emerges earlier, while high-frequency details from $V_{2}$ emerge later. Illustration only: in practice anisotropic and isotropic reconstruct different images, and the gap between $g_{1}$ and $g_{2}$ is typically smaller (see Fig. 2 – 5 ).

CIFAR-10
Method	nfe 9	nfe 11	nfe 13	nfe 15	nfe 17	nfe 35	nfe 59	nfe 79	Best FID
EDM	35.52	14.37	6.694	4.231	3.027	1.829	1.868	1.890	1.829
$g$	4.241	3.087	2.650	2.431	2.328	1.815	1.820	1.835	1.815
$\{g_{y}\}_{y=1}^{C}$	3.691	3.090	2.735	2.582	2.480	1.831	1.829	1.841	1.829
$\{(g^{\mathrm{DCT}}_{1,y},g^{\mathrm{DCT}}_{2,y})\}_{y=1}^{C}$	3.618	3.056	2.681	2.535	2.431	1.829	1.810	1.818	1.810
$(g^{\mathrm{PCA}}_{1},g^{\mathrm{PCA}}_{2})$	3.839	2.900	2.538	2.332	2.213	1.851	1.803	1.803	1.803

Table 1: FID

\downarrow

vs. NFE across datasets (50k samples). See section 7 for description of each method. For each method (both EDM and ours), we perform 3 independent random generations of 50k images and report the minimum FID across the three runs. Bold = per-NFE best. Blue = Best FID lower than EDM.

Appendix figures & tables7 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 2: CIFAR-10 learned schedule analysis. (a) PCA-based geometric mean $\sqrt{g^{\mathrm{PCA}}_{1}(t)g^{\mathrm{PCA}}_{2}(t)}$ with a log-linear reference (gray dashed; log $y$ -axis). (b) PCA anisotropy ratio $g^{\mathrm{PCA}}_{1}(t)/g^{\mathrm{PCA}}_{2}(t)$ . (c) Class-conditional isotropic schedules $g_{y}(t)/\bar{g}(t)$ , where $\bar{g}(t)=\big(\prod_{y=1}^{C}g_{y}(t)\big)^{1/C}$ with $C=\#\text{class}=10$ . (d) Class-conditional anisotropic schedules over DCT subspaces: $g^{\mathrm{DCT}}_{k,y}(t)/\bar{g}_{k}(t)$ for $k\in\{1,2\}$ (solid/dashed), where $i$ indexes classes and $\bar{g}_{k}(t)$ is the geometric mean across classes.

Figure 3: AFHQv2 learned schedule analysis. Ratio between the two learned DCT-based schedules $g^{\mathrm{DCT}}_{1}(t)$ and $g^{\mathrm{DCT}}_{2}(t)$ over time.

Figure 4: FFHQ learned schedule analysis. Ratio between the two learned DCT-based schedules $g^{\mathrm{DCT}}_{1}(t)$ and $g^{\mathrm{DCT}}_{2}(t)$ over time.

Figure 5: ImageNet-64 learned schedule analysis. (a) Class-conditional isotropic schedules shown as $g_{y}(t)/\bar{g}(t)$ , where $\bar{g}(t)$ denotes the geometric mean across classes. (b) Class-conditional anisotropic schedules summarized by the geometric mean $\sqrt{g^{\mathrm{DCT}}_{1,y}(t)g^{\mathrm{DCT}}_{2,y}(t)}$ and normalized by its class-wise geometric mean. (c) Class-conditional anisotropy ratios $g^{\mathrm{DCT}}_{1,y}/g^{\mathrm{DCT}}_{2,y}$ . (d) PCA-based (class-conditional basis) anisotropic schedule ratio $g^{\mathrm{PCA}}_{1,y}(t)/g^{\mathrm{PCA}}_{2,y}(t)$ . (e) PCA-based (class-conditional basis) class-conditional anisotropic schedules shown as $\sqrt{g^{\mathrm{PCA}}_{1,y}(t)g^{\mathrm{PCA}}_{2,y}(t)}$ normalized by the geometric mean across classes. (f) PCA-based (class-conditional basis) class-conditional anisotropy ratios $g^{\mathrm{PCA}}_{1,y}(t)/g^{\mathrm{PCA}}_{2,y}(t)$ . All ratios are plotted on a logarithmic $y$ -axis, and dashed horizontal lines indicate the reference value $1$ .

Figure 6: Class-specific PCA subspaces. Leading principal components for representative CIFAR-10 classes (cat, dog, airplane), reshaped as spatial patterns. Semantically related classes (e.g., cat and dog) exhibit qualitatively similar dominant structures, whereas visually distinct classes (e.g., airplane) display different principal directions.

Figure 7: Subspace distance embedding via classical MDS. Pairwise distances between class-specific PCA subspaces are computed using the Frobenius norm between orthogonal projectors. Classical multidimensional scaling (MDS) embeds the resulting distance matrix into two dimensions. Animal classes cluster together and are separated from vehicle classes, indicating that the learned subspaces capture semantically structured geometry.

Figure 8: The first nine 2D-DCT bases ordered by increasing frequency.

Diffusion Models

Diffusion models learn to reverse gradual data corruption, enabling flexible generation, editing, forecasting, and inference.

See all 1,169 papers

Noise Schedule Design for Diffusion Models: An Optimal Control Perspective
May 21, 2026Seo Taek Kong, Weina Wang, R. SrikantDiffusion Sampling Diffusion-Based Image Generation
ECE & CSL, University of Illinois Urbana-Champaign · Computer Science Department, Carnegie Mellon University · ECE, CSL & NCSA, University of Illinois Urbana-Champaign
Principled Design of Diffusion-based Optimizers for Inverse Problems
May 12, 2026Julio Oscanoa, Irmak Sivgin, Cagan Alkan +4Inference-Time Optimization Diffusion-Based Inverse Problems
Department of Bioengineering · Deparment of Electrical Engineering · Department of Radiology, Stanford University, USA.
Learning When to Denoise: Optimizing Asynchronous Schedules for Latent Diffusion
Jun 18, 2026Bingshuo Qian, Xiang ChengFlow Matching Latent Diffusion Models
Department of Electrical and Computer Engineering Duke University

May 21, 2026J/K move · Enter open · S save

Noise Schedule Design for Diffusion Models: An Optimal Control Perspective

Seo Taek Kong, Weina Wang, R. Srikant

ECE & CSL, University of Illinois Urbana-Champaign · Computer Science Department, Carnegie Mellon University · ECE, CSL & NCSA, University of Illinois Urbana-Champaign

Diffusion Sampling Diffusion-Based Image Generation Diffusion Timestep Scheduling Optimal Control Diffusion Models

Open paper AskPDF arXiv

We develop a principled framework for analyzing and designing noise schedules in diffusion models. We show that one can recast this design problem as an optimal control problem, whose state is the Fisher information of the diffusion process, which evolves according to an ODE, and whose control input is the noise schedule. The objective of the optimal control problem is a functional involving the Fisher information, which is shown to be an upper bound on the Kullback-Leibler sampling error. By solving this optimal control problem, we obtain sufficient conditions on noise schedules under which

\tilde{\mathcal{O}}(d/n)

sampling error is achievable, where

d

is the data dimension and

n

is the number of discretization steps. While existing theoretical works also prove that

\tilde{\mathcal{O}}(d/n)

sampling error bounds are achievable, these results hold for specific noise schedules that differ from those used in practice. Under additional assumptions on the data distribution, we derive a family of noise schedules that recovers the functional forms of schedules widely used in practice. We show that these schedules achieve the

\tilde{\mathcal{O}}(d/n)

error bound when their parameters are chosen according to our analysis. Our theory offers a principled approach to noise schedule design, yielding explicit functional forms and guidance for selecting their parameters.

Abstract

Diffusion Models

Explore similar work

Noise Schedule Design for Diffusion Models: An Optimal Control Perspective

Principled Design of Diffusion-based Optimizers for Inverse Problems

Learning When to Denoise: Optimizing Asynchronous Schedules for Latent Diffusion