Variance reduction of diffusion model's gradients with Taylor approximation-based control variate

TL;DR

This paper introduces a Taylor approximation-based control variate to reduce gradient variance in diffusion models, improving training stability and efficiency for high-dimensional data generation.

Contribution

It proposes a novel control variate derived from Taylor expansion, with theoretical proof of equivalence and empirical validation on both low and high-dimensional problems.

Findings

Significant variance reduction in gradients.

Improved training stability in diffusion models.

Effective on both low and high-dimensional data.

Abstract

Score-based models, trained with denoising score matching, are remarkably effective in generating high dimensional data. However, the high variance of their training objective hinders optimisation. We attempt to reduce it with a control variate, derived via a $k$-th order Taylor expansion on the training objective and its gradient. We prove an equivalence between the two and demonstrate empirically the effectiveness of our approach on a low dimensional problem setting; and study its effect on larger problems.