Variance Reduction for Expectations with Diffusion Teachers

Jesse Bettencourt, Xindi Wu, Matan Atzmon, James Lucas, Jonathan Lorraine · ICML 2026 SPIGM Workshop

Abstract

Pretrained diffusion models serve as frozen teachers feeding downstream pipelines such as 3D generation, single-step distillation, and data attribution. The teacher gradients these pipelines consume are Monte Carlo expectations over noise levels and Gaussian noise samples; their estimator variance dominates compute cost because each draw requires expensive upstream work (rendering, simulation, encoding). We introduce CARV, a compute-aware variance-accounting framework that motivates a hierarchical Monte Carlo estimator: amortize the expensive upstream computation over cheap diffusion-noise resamples, sharpened by timestep importance sampling and a stratified inverse-CDF construction. In 3D generation and attribution experiments, CARV delivers 2 to 3 times effective compute multipliers (most from amortized reuse, with about 25% additional from importance sampling and stratification) without changing the objective. In single-step distillation, the same techniques cut gradient variance by an order of magnitude but do not improve downstream FID, indicating that MC variance is no longer the bottleneck in that setting.

Links