Optimal Stochastic Trace Estimation in Generative Modeling (AISTATS 2025)

By Xinyang Liu*¹, Hengrong Du*² , Wei Deng³, Ruqi Zhang^1
1Purdue University, ²UC Irvine, ²Morgan Stanley
*Equal contribution

Why Hutch++ in Generative Modeling?

• Lower variance and error in trace estimation compared to Hutchinson trace estimator.
• In generative modeling, 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.
• Our practical acceleration technique that balance frequency and accuracy, backed by theoretical guarantees.
• Our analysis demonstrates that Hutch++ leads to generations of higher quality.

Code Structure

We provide two separate projects of GruM for three types of graph generation tasks:

• Simulations with Neural ODE model
• Image data modeling with advanced SB-based diffusion models

Each projects consists of the following:

ffjord-plus-plus : Code for density estimation of toy distributions with FFJORD++

sb-fbsde-plus-plus : Code for image data modeling with SB-FBSDE++

We provide the details in README.md for each projects.

Acknowledgements

This repository was heavily built off of FFJORD, SB-FBSDE and VSDM.

Citation

@article{liu2025optimal,
  title={Optimal Stochastic Trace Estimation in Generative Modeling},
  author={Liu, Xinyang and Du, Hengrong and Deng, Wei and Zhang, Ruqi},
  journal={arXiv preprint arXiv:2502.18808},
  year={2025}
}