# Optimal Stochastic Trace Estimation in Generative Modeling

arXiv: https://arxiv.org/abs/2502.18808 · HTML: https://mlsanity.com/paper/2502.18808/

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Placed in Diffusion foundations and sampling, one of 14 directions in Diffusion models.
- Builds on 46 of its 76 references that are in the map.
- First posted on arXiv 26 Feb 2025.

## Placed in

- Diffusion models / [Diffusion foundations and sampling](https://mlsanity.com/diffusion/diffusion-foundations-and-sampling/)

## Builds on (most cited first, 30 of 46)

- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Score-Based Generative Modeling through Stochastic Differential Equations](https://mlsanity.com/paper/2011.13456/) (2020)
- [Communication-Efficient Learning of Deep Networks from Decentralized Data](https://mlsanity.com/paper/1602.05629/) (2016)
- [Hierarchical Text-Conditional Image Generation with CLIP Latents](https://mlsanity.com/paper/2204.06125/) (2022)
- [Flow Matching for Generative Modeling](https://mlsanity.com/paper/2210.02747/) (2022)
- [Generative Modeling by Estimating Gradients of the Data Distribution](https://mlsanity.com/paper/1907.05600/) (2019)
- [Auto-Encoding Variational Bayes](https://mlsanity.com/paper/1312.6114/) (2013)
- [Neural Ordinary Differential Equations](https://mlsanity.com/paper/1806.07366/) (2018)
- [Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow](https://mlsanity.com/paper/2209.03003/) (2022)
- [Density estimation using Real NVP](https://mlsanity.com/paper/1605.08803/) (2016)
- [Imagen Video: High Definition Video Generation with Diffusion Models](https://mlsanity.com/paper/2210.02303/) (2022)
- [DiffWave: A Versatile Diffusion Model for Audio Synthesis](https://mlsanity.com/paper/2009.09761/) (2020)
- [On the Convergence of FedAvg on Non-IID Data](https://mlsanity.com/paper/1907.02189/) (2019)
- [NICE: Non-linear Independent Components Estimation](https://mlsanity.com/paper/1410.8516/) (2014)
- [Building Normalizing Flows with Stochastic Interpolants](https://mlsanity.com/paper/2209.15571/) (2022)
- [Stochastic Interpolants: A Unifying Framework for Flows and Diffusions](https://mlsanity.com/paper/2303.08797/) (2023)
- [Improving and generalizing flow-based generative models with minibatch optimal transport](https://mlsanity.com/paper/2302.00482/) (2023)
- [Diffusion Schr"odinger Bridge with Applications to Score-Based Generative Modeling](https://mlsanity.com/paper/2106.01357/) (2021)
- [FFJORD: Free-form Continuous Dynamics for Scalable Reversible Generative Models](https://mlsanity.com/paper/1810.01367/) (2018)
- [CSDI: Conditional Score-based Diffusion Models for Probabilistic Time Series Imputation](https://mlsanity.com/paper/2107.03502/) (2021)
- [Recurrent Neural Networks for Multivariate Time Series with Missing Values](https://mlsanity.com/paper/1606.01865/) (2016)
- [Sliced Score Matching: A Scalable Approach to Density and Score Estimation](https://mlsanity.com/paper/1905.07088/) (2019)
- [Autoregressive Denoising Diffusion Models for Multivariate Probabilistic Time Series Forecasting](https://mlsanity.com/paper/2101.12072/) (2021)
- [BRITS: Bidirectional Recurrent Imputation for Time Series](https://mlsanity.com/paper/1805.10572/) (2018)
- [Rectified Flow: A Marginal Preserving Approach to Optimal Transport](https://mlsanity.com/paper/2209.14577/) (2022)
- [Multisample Flow Matching: Straightening Flows with Minibatch Couplings](https://mlsanity.com/paper/2304.14772/) (2023)
- [Convergence of denoising diffusion models under the manifold hypothesis](https://mlsanity.com/paper/2208.05314/) (2022)
- [Likelihood Training of Schr"odinger Bridge using Forward-Backward SDEs Theory](https://mlsanity.com/paper/2110.11291/) (2021)
- [GP-VAE: Deep Probabilistic Time Series Imputation](https://mlsanity.com/paper/1907.04155/) (2019)

## Built on by (most cited first, 2 of 2)

- [Few-Step Boltzmann Generators via Scalable Likelihood Flow Maps](https://mlsanity.com/paper/2606.29110/) (2026)
- [StAD: Stein Amortized Divergence for Fast Likelihoods with Diffusion and Flow](https://mlsanity.com/paper/2605.16486/) (2026)

## Abstract (arXiv)

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.