Optimal Stochastic Trace Estimation in Generative Modeling

TL;DR

This paper introduces Hutch++, an advanced stochastic trace estimator that reduces variance and improves scalability in generative modeling, especially for high-dimensional, low-rank data, leading to higher quality outputs.

Contribution

The paper presents Hutch++, a novel trace estimation method optimized for generative models, with theoretical guarantees and practical schemes for variance reduction and scalability.

Findings

Hutch++ significantly reduces variance in trace estimation.

The method improves generative model quality across applications.

Practical schemes balance computational cost and accuracy.

Abstract

Hutchinson estimators are widely employed in training divergence-based likelihoods for diffusion models to ensure optimal transport (OT) properties. However, this estimator often suffers from high variance and scalability concerns. To address these challenges, we investigate Hutch++, an optimal stochastic trace estimator for generative models, designed to minimize training variance while maintaining transport optimality. Hutch++ is particularly effective for handling ill-conditioned matrices with large condition numbers, which commonly arise when high-dimensional data exhibits a low-dimensional structure. To mitigate the need for frequent and costly QR decompositions, we propose practical schemes that balance frequency and accuracy, backed by theoretical guarantees. Our analysis demonstrates that Hutch++ leads to generations of higher quality. Furthermore, this method exhibits effective variance reduction in various applications, including simulations, conditional time series forecasts, and image generation.