Variational Trajectory Optimization of Anisotropic Diffusion Schedules¶

会议: NeurIPS 2026
arXiv: 2602.19512
论文: https://arxiv.org/abs/2602.19512
代码: https://github.com/lizeyu090312/anisotropic-diffusion-paper
领域: 其他
关键词: 扩散模型, 各向异性调度, 变分优化, 轨迹得分匹配, 逆向ODE求解器

一句话总结¶

本文提出了各向异性扩散调度的变分轨迹优化框架，通过引入矩阵值正定噪声协方差路径与空间方向导数梯度估计器，联合学习非对称噪声分配与得分网络，并在二阶 Heun 采样下在多项图像生成基准上取得优于 EDM 的表现。

研究背景与动机¶

现代扩散生成模型（如 DDPM、Score SDE 与 EDM）绝大多数基于各向同性（Isotropic）的高斯加噪假设，即在任意给定的扩散时刻 $t$，注入的前向噪声协方差矩阵均为标量与单位阵的乘积 $\sigma(t)^2 I$。然而，真实数据（尤其是自然图像）的几何结构具有高度各向异性：自然图像的能量显著富集于低频空间分量，高频细节更易受到高斯扰动淹没；同时在多尺度与隐空间生成任务中，粗粒度轮廓与细粒度纹理也天然适宜以不同的时间步调进行去噪。如果强制所有空间与频域方向在同一时刻承受相同的加噪与去噪强度，将严重制约扩散模型对复杂数据流形几何的适配效率。

各向异性扩散模型试图通过矩阵值噪声协方差路径 $M_t(\theta)$ 替代标量噪声方差，从而在不同的正交子空间（如 DCT 频段或类条件 PCA 基）与扩散时间之间灵活调配去噪精力。然而，设计合理的各向异性调度空间异常巨大（半正定矩阵锥上的时序路径），更严峻的理论矛盾在于：各向同性扩散中改变标量时间调度本质上仅是沿同一维边际分布流形重参数化时间轴；而对于矩阵值路径 $M_t(\theta)$，不同的参数化调度将诱导完全不同的边际分布族 $p_t(x; \theta) = p_0 * \mathcal{N}(0, M_t(\theta))$，进而彻底改变反向去噪的目标得分场。这导致手工设计的各向异性调度极度脆弱，亟需一种原则性的数据驱动优化准则。此外，在双层优化中对调度参数求导需要计算最优得分网络关于调度的隐式梯度，而常规网络无法直接提供关于调度参数的偏导数。

本文的切入角度是建立一个轨迹级的变分得分匹配框架，不仅定义了衡量理想去噪动力学与实际网络速度场失配的全局目标，而且利用热传导方程的几何性质，将调度参数的导数转化为得分网络输入的高阶空间方向导数。核心 idea：将各向同性时间标量推广为半正定单调矩阵值噪声轨迹 $M_t(\theta)$，构建轨迹级变分得分匹配损失以联合训练得分网络与调度参数，并基于高阶空间方向导数导出无偏梯度估计器与矩阵化二阶 Heun 求解器。

方法详解¶

整体框架¶

本文方法的整体流程涵盖加噪扩散前向建模、流场参数化、轨迹级双层优化训练以及逆向矩阵 ODE 采样四个阶段。给定输入数据 $x_0 \sim p_0$ 与参数化矩阵值噪声路径 $M_t(\theta)$，系统在正交子空间上对数据施加非对称扰动；得分网络通过尺度稳定的流场形式预测去噪向量；优化阶段通过空间高阶导数估计器更新调度参数 $\theta$ 并同步优化网络参数 $\phi$；推理采样阶段则使用基于协方差平方根增量的矩阵二阶 Heun 求解器还原数据。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入数据与参数化调度<br/>x0 与 Mt(theta)"] --> B["变分轨迹级得分匹配目标<br/>度量逆向轨迹上的速度场失配"]
    B --> C["尺度稳定流参数化<br/>标准化向量场以平衡不同噪声尺度下的梯度方差"]
    C --> D["基于方向导数的调度梯度估计器<br/>利用高阶空间导数求解双层优化外层梯度"]
    D --> E["矩阵轨迹二阶 Heun 求解器<br/>基于协方差平方根增量的高精度采样"]
    E --> F["输出生成样本 x0"]

关键设计¶

1. 变分轨迹级得分匹配目标：度量逆向轨迹上的速度场失配

传统标量扩散调度仅调节一维时间重参数化，而矩阵值路径 $M_t(\theta)$ 的变动直接改变了边际分布族 $p_t(x; \theta) = p_0 * \mathcal{N}(0, M_t(\theta))$。为了能够原则性地联合评估得分网络拟合与调度路径优劣，受 Girsanov 路径空间测度变换控制启发，本文在方差保持（VP）坐标系下构建了理想逆向速度场 $v$ 与网络诱导速度场 $\bar{v}$ 的路径平方失配积分。通过单样本代理速度替换不可计算的真实得分，最终推导得出紧凑的轨迹级加权得分匹配目标： $$L(\theta, \phi) := \mathbb{E}_{t \sim \text{Unif}[0, T],\, x_0 \sim p_0,\, \epsilon \sim \mathcal{N}(0, I)} \left[ \left\| W_t(\theta) \left( M_t(\theta)^{1/2} \text{net}\left(x_0 + M_t(\theta)^{1/2}\epsilon,\, t,\, \phi\right) + \epsilon \right) \right\|_2^2 \right]$$ 其中权重矩阵定义为 $W_t(\theta) := (I + M_t(\theta))^{-1/2} (\partial_t M_t(\theta)) M_t(\theta)^{-1/2}$。在任意固定的半正定单调调度 $M_t(\theta)$ 下，当模型容量和数据无限时，该损失的全局极小点严格恢复真实数据得分 $\nabla_x \log p_t(x; \theta)$。这一设计成功将各向异性矩阵调度的评估转化为显式的二次优化目标。

2. 尺度稳定流参数化：标准化向量场以平衡不同噪声尺度下的梯度方差

在扩散模型中，得分向量的模长随噪声水平急剧变化，其理论范数正比于 $\|M_t(\theta)^{-1/2}\|_2$。在小噪声区域（$t \to 0$），这一数值的剧烈爆炸会导致损失函数梯度出现严重的方差膨胀，阻碍调度参数 $\theta$ 的平稳收敛。针对这一痛点，本文引入尺度标准化的向量场 $\text{flow}(x, t, \phi) := M_t(\theta)^{1/2} \text{net}(x, t, \phi)$。流场 $\text{flow}$ 在整个扩散时序演化过程中模长保持近乎时间常数。在此参数化下，轨迹损失等价重写为： $$L(\theta, \phi) = \mathbb{E}_{x_0, \epsilon, t} \left[ \left\| W_t(\theta) \left( \text{flow}(x_t, t, \phi) + \epsilon \right) \right\|_2^2 \right]$$ 该设计使网络拟合目标直接对齐标准正态扰动 $-\epsilon$，消除了不同噪声尺度下梯度的病态量纲不一致，显著提升了联合优化数值稳定性。

3. 基于方向导数的调度梯度估计器：利用高阶空间导数求解双层优化外层梯度

优化调度参数 $\theta$ 本质上是一个双层优化问题：$\min_\theta H(\theta) := L(\theta, \phi^*(\theta))$。对 $H(\theta)$ 求导不仅包含损失对 $M_t(\theta)$ 的显式偏导，更包含隐式依赖项 $\mathbb{E}[\langle \nabla_v \ell,\, \partial_\theta \text{flow}^*(x_t, t, \phi^*(\theta)) \rangle]$。常规双层优化需回溯展开优化轨迹或计算海森逆，开销巨大且网络并不直接暴露对 $\theta$ 的导数。本文发现高斯卷积边际分布满足广义热传导方程 $\partial_\theta p_t(x; \theta) = \frac{1}{2} \text{div}(p_t \partial_\theta M_t \nabla \log p_t)$，由此推导出关键恒等式，将 $\partial_\theta \text{flow}^*$ 精确转化为对空间输入 $x$ 的高阶方向导数： $$\partial_\theta \text{flow}(x, t, \phi^*) = \frac{1}{2}\sum_{i=1}^d \partial_r \partial_s \text{flow}\Big(x + r e_i + s \partial_\theta M_t e_i\Big)\Big|_{r=s=0} + \partial_s \text{flow}\Big(x + s M_t^{-1/2} \partial_\theta M_t \text{flow}\Big)\Big|_{s=0} + \frac{1}{2} M_t^{-1}(\partial_\theta M_t) \text{flow}$$ 该估计器完全独立于调度参数维度 $\dim(\theta)$，仅需 3 次空间反向传播（向量-雅可比积 VJP）即可完成无偏估计，使得矩阵调度的端到端一阶梯度更新变得轻量可行。

4. 矩阵轨迹二阶 Heun 求解器：基于协方差平方根增量的高精度采样

在推理阶段，若直接采用一阶欧拉法求解逆向矩阵 ODE，在低步数（如 NFE $\le 15$）下截断误差极为显著。针对可交换正交子空间协方差族（如 2D-DCT 或 PCA 投影），逆向 ODE 满足 $d\bar{x}_t = -\partial_t(M_t^{1/2}) \text{flow}(\bar{x}_t, t, \phi) dt$，其自然时间步长即为协方差平方根增量 $\Delta U_k := M_{t_k}^{1/2} - M_{t_{k-1}}^{1/2}$。本文在相邻步 $[t_{k-1}, t_k]$ 间引入预测评估点 $\hat{t}_k$，构造仿射流场积分，推导得出了闭式矩阵二阶校正步： $$\bar{x}_{t_{k-1}} = \bar{x}_{t_k} + \Delta U_k u_k + \frac{1}{2} (\Delta U_k)^2 (\hat{U}_k - U_k)^{-1} (\hat{u}_k - u_k)$$ 当预测点选取为端点 $\hat{t}_k = t_{k-1}$ 时，式子自然化简为矩阵版 Heun 梯形步 $\bar{x}_{t_{k-1}} = \bar{x}_{t_k} + \Delta U_k \frac{u_k + \hat{u}_k}{2}$。该求解器在保证每步 2 次网络评估的同时，可在相邻采样步间复用端点向量（$K$ 步仅需 $2K-1$ 次 NFE），配合子空间对角化计算，在极低计算开销下实现了高精度的矩阵轨迹数值积分。

损失函数 / 训练策略¶

训练采用两阶段微调范式：网络权重直接加载预训练好的 EDM 权重进行热启动，调度参数 $\theta$ 则初始化自标量等向调度模型。使用 Adam 优化器（$\beta_1 = 0.9, \beta_2 = 0.999$）联合更新网络参数 $\phi$ 与调度参数 $\theta$，两者采用解耦的学习率（调度学习率随网络学习率等比例退火）。训练过程中设置线性预热与后半程学习率衰减，并在反向传播后显式执行 NaN/Inf 梯度清洗。对于类条件调度参数，仅在当前 batch 中出现的类别上更新梯度，并通过分布式 DDP All-Reduce 严格同步所有卡上的类特异性调度。全局微调统一控制在 1.2M 图像处理量以内，评测阶段统一使用网络的 EMA 权重。

实验关键数据¶

主实验¶

在 CIFAR-10、AFHQv2、FFHQ 和 ImageNet-64 四个经典图像生成基准上，评估了 50k 样本生成下的 FID 指标（3 次独立随机运行的最低值），对比基线为官方 EDM 模型。

数据集	评估指标	EDM 基线	本文最佳变体	配置说明	提升幅度 (FID 降幅)
CIFAR-10 (32x32)	Best FID ↓	1.829 (NFE 35)	1.803 (NFE 59)	PCA 各向异性调度 $(g_1^{\text{PCA}}, g_2^{\text{PCA}})$	-0.026
AFHQv2 (64x64)	Best FID ↓	2.042 (NFE 79)	2.010 (NFE 119)	DCT 各向异性调度 $(g_1^{\text{DCT}}, g_2^{\text{DCT}})$	-0.032
FFHQ (64x64)	Best FID ↓	2.374 (NFE 119)	2.242 (NFE 79)	学习所得等向调度 $g$	-0.132
ImageNet-64 (64x64)	Best FID ↓	2.276 (NFE 119)	2.238 (NFE 119)	类条件 DCT 各向异性 $\{(g_{1,y}^{\text{DCT}}, g_{2,y}^{\text{DCT}})\}_{y=1}^C$	-0.038

消融实验¶

消融实验重点探究了不同调度参数化形式（从各向同性 EDM、学习所得标量 $g$、到类条件各向异性）在极低网络函数评估次数（NFE 9、11、13、15）及充分采样阶段的表现（数值摘自 CIFAR-10 与 ImageNet-64）。

数据集	调度配置变体	NFE 9 FID ↓	NFE 11 FID ↓	NFE 15 FID ↓	Best FID ↓	说明
CIFAR-10	EDM (标准等向基线)	35.52	14.37	4.231	1.829	极低步数下离散化误差严重崩溃
CIFAR-10	学得等向调度 $g$	4.241	3.087	2.431	1.815	标量优化极大缓解低 NFE 性能
CIFAR-10	类条件等向 $\{g_y\}$	3.691	3.090	2.582	1.829	类别特异性时间重参数化
CIFAR-10	类条件 DCT $\{(g_{1,y}^{\text{DCT}}, g_{2,y}^{\text{DCT}})\}$	3.618	3.056	2.535	1.810	固定频段下类别特异性加噪
CIFAR-10	PCA 各向异性 $(g_1^{\text{PCA}}, g_2^{\text{PCA}})$	3.839	2.900	2.332	1.803	跨 NFE 综合表现最佳
ImageNet-64	EDM (标准等向基线)	35.32	15.19	5.458	2.276	复杂多类数据集上基线崩溃严重
ImageNet-64	类条件等向 $\{g_y\}$	8.416	5.907	4.079	2.296	低 NFE 下显著收敛加速
ImageNet-64	PCA 各向异性 $(g_1^{\text{PCA}}, g_2^{\text{PCA}})$	29.47	15.55	5.666	2.240	共享调度在复杂流形下低步数受限
ImageNet-64	类条件 PCA $\{(g_{1,y}^{\text{PCA}}, g_{2,y}^{\text{PCA}})\}$	8.366	5.823	3.992	2.239	类特异性基与调度联合优化
ImageNet-64	类条件 DCT $\{(g_{1,y}^{\text{DCT}}, g_{2,y}^{\text{DCT}})\}$	8.340	5.808	3.983	2.238	低步数与极限采样双双达到最优

关键发现¶

极低 NFE 采样收益极大：在计算预算受限的极限场景（NFE 9），EDM 在 CIFAR-10 和 ImageNet-64 上的 FID 分别退化为 35.52 和 35.32，而优化调度后的模型分别达 3.618 和 8.340，生成图像质量实现了数量级跨越的质感提升。
高复杂度数据依赖类条件各向异性：在简单的 CIFAR-10 上，基于全局共享调度的 PCA 各向异性即可达到最佳的 1.803 FID；而在类别极其多样、视觉结构高度离散的 ImageNet-64 上，必须依靠类特异性调度 $\{(g_{1,y}, g_{2,y})\}$ 才能取得最优采样质量（2.238）。
频域退化先慢后快规律：学得的 DCT 各向异性比例曲线显示，低频成分在前向过程中被赋予更平缓的加噪曲线而在反向过程中更早去噪重构，高频细节则集中在较晚时刻还原，从数据驱动角度验证了粗粒度至细粒度分层生成的几何规律。

亮点与洞察¶

空间高阶导数避开双层优化死结：巧妙利用高斯热传导方程的偏微分方程结构，将对模型参数的不可计算隐式导数转化为网络输入的高阶空间方向导数。仅需 3 次 VJP 反向传播，便解决了高维矩阵时间调度的无偏梯度计算难题。
协方差增量自然导出矩阵 Heun 求解器：将 EDM 中标量噪声阶梯严谨推广为正定矩阵平方根的矩阵差分步长 $\Delta U_k$。利用子空间正交投影可交换性，避免了采样过程中的密集矩阵求逆开销。
轻量微调兼顾通用迁移性：无需从头预训练高昂的扩散骨干，直接在标准 EDM 权重上以仅 1.2M 图像微调便能使网络得分场与新的非对称几何路径充分对齐，即插即用属性突出。

局限与展望¶

正交子空间分解的结构约束：为保证矩阵平方根与逆的高效计算，当前理论实现强烈依赖于子空间投影矩阵可交换（即共享正交基）的假设，尚未探索任意非对称、动态流形投影的泛化形式。
调度自由度受限于手工指定基底：目前子空间基底主要依赖先验设计的 2D-DCT 块或预先统计的类条件 PCA，子空间投影算子本身尚未与网络联合端到端优化。
模态扩展潜力待验证：当前验证主要集中于中低分辨率图像（$\le 64\times 64$），未来在视频时序相关先验、3D 几何生成以及隐空间扩散模型（LDM）中的应用有待深入挖掘。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 首次提出连续矩阵值扩散调度的变分轨迹优化框架，并推导了基于空间高阶导数的无偏梯度估计器。
实验充分度: ⭐⭐⭐⭐☆ 覆盖 4 个标准数据集与多种 NFE 设置，消融设计系统严密；若能在高分辨率 LDM 上进一步验证会更加完备。
写作质量: ⭐⭐⭐⭐⭐ 数学推导严谨精炼，从连续流场物理测度直通离散求解器闭式解，逻辑环环相扣。
价值: ⭐⭐⭐⭐☆ 为扩散模型噪声路径设计开辟了突破标量假设的通用变分范式，对极低步数快速生成具有重要参考价值。