利用扩散教师实现期望的方差缩减

论文信息

标题: Variance Reduction for Expectations with Diffusion Teachers

作者: Jesse Bettencourt, Xindi Wu, Matan Atzmon, et al.

发布日期: 2026-05-20

PDF 链接: 下载 PDF

背景与动机

扩散模型已成为图像、视频、3D 内容生成的核心引擎。在许多下游应用中，预训练的扩散模型被当作“冻结的教师”，为文本到 3D 优化、单步蒸馏、数据归因等任务提供梯度信号。这些梯度本质上是关于噪声层级和随机噪声的蒙特卡洛期望。每个样本都需要昂贵的上游计算（渲染、仿真、编码），使得估计量的方差直接主导了总计算成本。在实验室规模下，这往往意味着六到七位数的预算。

已有方差降低工作主要集中于教师训练阶段的损失重加权和噪声调度设计，而下游使用者通常直接继承教师的噪声分布，采用临时的平均或引入偏差，缺乏系统化的方差分析和计算分配视角。论文提出的 CARV (Compute-Aware Variance Reduction) 框架正是为了回答三个核心问题：哪些成分主导方差？如何在不偏置目标函数的前提下降低方差？在固定预算下，如何在廉价操作（加噪、去噪）和昂贵操作（渲染、编码）之间进行权衡？

CARV 框架与方法细节

CARV 的核心思想是将扩散教师梯度视为一个分层蒙特卡洛估计问题，并通过三种互补的无偏技术显著提升单位计算的方差效率。

计算重用与分层估计

扩散梯度估计涉及两类随机性：昂贵的上游操作（如 3D 渲染、生成器前传）和廉价的扩散噪声（时间步 $t$ 和高斯噪声 $\boldsymbol{\epsilon}$ ）。朴素方法中对每个梯度样本都重新执行完整流水线，成本高昂。CARV 引入一种 分层蒙特卡洛估计器：首先缓存一定数量的昂贵上游输出，然后对每个缓存结果独立地重采样多次扩散噪声，再平均得到梯度贡献。例如，在 SDS（Score Distillation Sampling）任务中，对 $R$ 个渲染视角各生成一次潜码 $\mathbf{z}^{(r)}$ ，然后对每个视角进行 $K$ 次独立的加噪与去噪，产生 $R \times K$ 个有效样本。最终的梯度估计量为：

\hat{\nabla}_{\boldsymbol{\theta}}^{\mathrm{reuse}} = \frac{1}{R}\sum_{r=1}^{R}\Big(\frac{1}{K}\sum_{k=1}^{K} \mathbf{f}(\mathbf{x}^{(r)}, t^{(r,k)}, \boldsymbol{\epsilon}^{(r,k)})\Big)\frac{\partial \mathbf{x}^{(r)}}{\partial \boldsymbol{\theta}}

当渲染和编码成本远高于去噪成本时，这一策略几乎在不增加总计算量的前提下显著增加了有效样本数，从而大幅降低方差。论文将该思路推广到单步蒸馏（DMD）等场景，只要存在“昂贵状态”与“廉价噪声”的可分离性，即可发挥作用。

时间步重要性采样

扩散教师的梯度对时间步 $t$ 的依赖高度非均匀，直接使用均匀采样会浪费大量样本在贡献微小的噪声层级上。CARV 利用重要性采样，从提议分布 $q(t)$ 采样并乘以似然比 $p(t)/q(t)$ 。最优提议 $q^\star(t) \propto p(t)\sqrt{\mathbb{E}[\|\mathbf{f}(t,\boldsymbol{\xi})\|_2^2\mid t]}$ 虽理论上不可行，但作者发现 SDS 的权重函数 $w_{\mathrm{SDS}}(t)$ （包含 $\alpha_t$ 和调度因子）与实际参数梯度的范数高度相关，因此直接使用 $q(t) \propto p(t) w_{\mathrm{SDS}}(t)$ 作为廉价而高效的代理。该提议无需额外回传，仅通过修正权重即可保持无偏，在实验中能带来约 $1.2\times$ 的方差效率提升。

分层采样与逆 CDF 构造

分层采样将时间步的定义域划分为若干层，每层内强制采样，以避免样本聚集并确保对整个噪声范围的均衡覆盖。CARV 进一步将分层与重要性采样结合，形成 分层‑重要性采样：在提议分布的量子化空间（ $[0,1]$ 均匀分箱）上分层，通过逆 CDF 映射回非均匀的 $t$ 值。对每个渲染视角 $\mathbf{x}^{(r)}$ ，设定 $B$ 个等概率层，抽取

t_b^{(r)} = \mathrm{CDF}_q^{-1}\!\left(\frac{b-1+\xi_b^{(r)}}{B}\right), \quad \xi_b^{(r)}\sim\mathcal{U}(0,1)

配合相应的 $\tilde{w}(t_b^{(r)})$ 和独立噪声 $\boldsymbol{\epsilon}_b^{(r)}$ ，形成无偏的低方差估计。该方法几乎不增加计算开销，却能抑制由于重要性采样带来的样本波动，并在实验中额外提供约 $1.1-1.2\times$ 的相对效率提升。

创新点与贡献

CARV 的贡献可归纳为：

分层蒙特卡洛估计器：首次系统地在扩散教师梯度中区分昂贵上游与廉价噪声，并通过缓存‑重噪机制实现计算复用，这是有效计算乘数（ECM）达到 $2-3\times$ 的主要来源。
显式权重驱动的提议分布：直接使用任务自带的权重函数（如 SDS 的 $w_{\mathrm{SDS}}(t)$ ）作为重要性采样的代理，无需额外建模，在文本到 3D 和低引导强度下效果尤为显著。
分层‑逆 CDF 采样：将传统分层采样与重要性采样无缝结合，在提议分位数空间上分层，确保了统计效率与实现简洁性的统一。
计算感知的方差度量框架：引入有效计算乘数（ECM）和相对效率（RE），统一评估不同采样策略在异质计算成本下的真实收益，为从业者提供了直观的比较基准。

实验结果分析

论文在三个典型任务上进行了系统性验证：

文本到 3D 优化（SDS）：在 threestudio 框架内，均匀基线、单独的重要性加权、单独的分层、以及二者的组合四种策略被全面评测。仅计算重用一项即可使单位方差的等效计算提升约 $2.6\times$ ，而结合 IW 和分层后可达 $3.3\times$ 。CLIP 分数在等量迭代下显著提升，达到收敛质量所需时间几乎减半。可视化结果也证实了低方差梯度带来的几何与纹理改进。
单步蒸馏（DMD）：对 DiT‑XL/2 教师进行蒸馏时，重采样将参数梯度方差降低 $3.4\!-\!16\times$ ，分层进一步降低 $1\!-\!2\times$ 。然而，FID 并未随方差降低而改善，表明此处蒙特卡洛噪声已不再是瓶颈，辅助损失和生成器输入多样性限制了最终性能。这一“负结果”反而划定了 CARV 技术的适用范围：当方差不再是主导因素时，单纯降低梯度波动无济于事。
视频数据归因：在 MOTIVE 框架下，分层采样在有限的噪声‑时间预算内使影响排序与真实梯度的相关系数大幅提升，等效计算乘数在合理预算（16‑64 个时间步）下超过 $2\times$ ，有力证实了分层在不增加编码成本的前提下对排名稳定性的贡献。

实践应用建议

明确任务瓶颈：如果渲染或生成器前传是计算瓶颈，优先启用计算重用（ $K>1$ ）；如果仅有廉价去噪和加噪，分层和重要性采样也能以零成本提升效率。
选择适合的提议分布：对于 SDS 类任务，直接使用 $w_{\mathrm{SDS}}(t)$ ；若权重非单调（如 DMD 的归一化因子），可暂不引入重要性采样，仅依靠分层和重用。
实施分层时用逆 CDF：配合自动微分库的 rsample 接口，可直接在分位数空间分层，减少代码侵入。
利用方差测量框架：在开发新管道时，先用 Welford 算法在线监测梯度方差，快速评估不同 $(R,K)$ 组合的实际 ECM，避免凭经验试错。

局限与未来发展

CARV 技术仅在蒙特卡洛方差是主要瓶颈时才有效，当系统被其他因素（如生成器多样性、辅助损失、优化动力学）限制时，收益会消失。此外，分层和重要性采样的收益会随着 $K$ 增大而饱和，因为跨渲染的方差最终成为主导。未来工作可探索自适应调度 $K$ 与 $R$ ，以及将相似思想应用于基于流匹配或一致性模型的教师，甚至联合优化权重函数与采样方案，以进一步压榨计算效率。

总结

CARV 为依赖扩散教师梯度的流水线提供了一套简洁、无偏且极其实用的方差降低工具集。通过分层蒙特卡洛重用、基于任务权重的重要性采样和分层‑逆 CDF 构造，它无需改变优化目标即可实现 $2-3\times$ 的有效计算加速。在文本到 3D 和数据归因任务中，它带来了质量与效率的明确提升；同时，它在单步蒸馏中的“有限作为”也界定了技术适用的边界。这种系统化的计算感知视角，不仅为当下扩散应用提供了即时指导，也为未来更复杂的生成管道中的方差管理奠定了方法论基础。