Variational Trajectory Optimization of Anisotropic Diffusion Schedules

📄 arXiv: 2602.19512v1 📥 PDF

作者: Pengxi Liu, Zeyu Michael Li, Xiang Cheng

分类: cs.LG, cs.CV

发布日期: 2026-02-23

🔗 代码/项目: GITHUB

💡 一句话要点

提出变分框架优化各向异性扩散模型，提升图像生成质量与效率。

🎯 匹配领域: 支柱一：机器人控制 (Robot Control)

关键词: 扩散模型 各向异性噪声 变分推断 图像生成 逆ODE求解器 矩阵值路径 score网络

📋 核心要点

  1. 现有扩散模型噪声schedule通常是各向同性的，限制了模型在不同特征空间上的表达能力和优化效率。
  2. 论文提出变分框架，通过学习矩阵值噪声schedule，实现各向异性噪声控制，从而更精细地调整扩散过程。
  3. 实验表明，该方法在多个图像数据集上，显著提升了生成质量和采样效率，超越了现有基线模型。

📝 摘要（中文）

本文提出了一个变分框架，用于优化具有各向异性噪声schedule的扩散模型。该schedule由一个矩阵值路径$M_t(θ)$参数化，该路径在子空间中分配噪声。该框架的核心是一个轨迹级别的目标函数，它联合训练score网络并学习$M_t(θ)$，涵盖了矩阵值噪声schedule的通用参数化类别。此外，我们还推导出了score函数关于$θ$的导数的估计器，从而能够有效地优化$M_t(θ)$ schedule。在推理方面，我们开发了一种可高效实现的逆ODE求解器，它是二阶Heun离散化算法的各向异性推广。在CIFAR-10、AFHQv2、FFHQ和ImageNet-64数据集上，我们的方法在所有NFE（Number of Function Evaluations）机制中始终优于基线EDM模型。

🔬 方法详解

问题定义：现有的扩散模型通常采用各向同性的噪声 schedule，即在所有特征维度上添加相同强度的噪声。这种方式忽略了不同特征维度之间的差异，限制了模型的表达能力和优化效率。例如，某些特征可能对噪声更敏感，而另一些特征则需要更强的噪声才能有效探索。

核心思路：论文的核心思路是通过学习一个矩阵值的噪声 schedule $M_t(θ)$，实现各向异性噪声控制。$M_t(θ)$ 可以根据特征的重要性或敏感度，在不同的子空间中分配不同强度的噪声。通过优化 $M_t(θ)$，可以更有效地引导扩散过程，从而提高生成质量和采样效率。

技术框架：整体框架包含两个主要部分：score网络的训练和噪声schedule $M_t(θ)$的优化。首先，使用轨迹级别的目标函数联合训练 score 网络和学习 $M_t(θ)$。然后，推导出一个 score 函数关于 $θ$ 的导数估计器，用于高效地优化 $M_t(θ)$。在推理阶段，使用一个各向异性的逆 ODE 求解器，该求解器是二阶 Heun 离散化算法的推广。

关键创新：最重要的创新点在于提出了一个变分框架，可以学习各向异性的噪声 schedule。与传统的各向同性噪声 schedule 相比，该方法能够更精细地控制扩散过程，从而提高生成质量和采样效率。此外，论文还提出了一个高效的导数估计器，用于优化噪声 schedule。

关键设计：论文使用矩阵值路径 $M_t(θ)$ 来参数化噪声 schedule。该路径可以是任意的矩阵值函数，例如线性函数、多项式函数或神经网络。为了优化 $M_t(θ)$，论文推导出了 score 函数关于 $θ$ 的导数估计器。该估计器基于变分推断的思想，可以有效地计算导数。在推理阶段，论文使用一个各向异性的逆 ODE 求解器，该求解器是二阶 Heun 离散化算法的推广。该求解器可以有效地处理各向异性的噪声 schedule。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，该方法在 CIFAR-10、AFHQv2、FFHQ 和 ImageNet-64 等数据集上，在所有 NFE 机制中始终优于基线 EDM 模型。例如，在 CIFAR-10 数据集上，该方法在相同的 NFE 下，FID 分数显著优于 EDM 模型，表明生成质量得到了显著提升。

🎯 应用场景

该研究成果可广泛应用于图像生成、图像编辑、图像修复等领域。通过学习各向异性的噪声 schedule，可以生成更高质量、更逼真的图像。此外，该方法还可以用于加速扩散模型的采样过程，提高生成效率。未来，该方法有望应用于更多领域，例如视频生成、3D 模型生成等。

📄 摘要（原文）

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.