[Paper Review] XTrace: Making the most of every sample in stochastic trace estimation

Ethan N. Epperly, Joel A. Tropp+1 more|arXiv (Cornell University)|Jan 19, 2023

Random Matrices and Applications31 references4 citations

TL;DR

This paper introduces XTrace and XNysTrace, novel stochastic trace estimators that leverage the exchangeability principle and variance reduction to achieve significantly lower estimation error than existing methods like Girard–Hutchinson and Hutch++. By constructing symmetric, exchangeable estimators from matrix-vector products, the algorithms reduce variance by orders of magnitude for a fixed number of matvecs, with theoretical guarantees tied to the matrix spectrum and empirical validation across diverse test matrices.

ABSTRACT

The implicit trace estimation problem asks for an approximation of the trace of a square matrix, accessed via matrix-vector products (matvecs). This paper designs new randomized algorithms, XTrace and XNysTrace, for the trace estimation problem by exploiting both variance reduction and the exchangeability principle. For a fixed budget of matvecs, numerical experiments show that the new methods can achieve errors that are orders of magnitude smaller than existing algorithms, such as the Girard-Hutchinson estimator or the Hutch++ estimator. A theoretical analysis confirms the benefits by offering a precise description of the performance of these algorithms as a function of the spectrum of the input matrix. The paper also develops an exchangeable estimator, XDiag, for approximating the diagonal of a square matrix using matvecs.

Motivation & Objective

• Address the fundamental limitation of existing stochastic trace estimators, which fail to satisfy the exchangeability principle and thus do not minimize variance.
• Develop new randomized algorithms that exploit exchangeability and variance reduction to improve trace estimation accuracy under a fixed matvec budget.
• Provide a theoretical analysis linking estimator performance to the spectral properties of the input matrix.
• Introduce XDiag, an exchangeable estimator for approximating the diagonal of a matrix using matrix-vector products.
• Demonstrate through numerical experiments that the new methods achieve error reductions of multiple orders of magnitude compared to state-of-the-art approaches.

Proposed method

• Propose XTrace, a symmetric, exchangeable estimator that computes the trace as a function of multiple matrix-vector products with i.i.d. random vectors, ensuring minimum variance under exchangeability.
• Integrate low-rank approximation via truncated SVD to reduce variance, using the residual matrix to refine the trace estimate.
• Design XNysTrace as a Nyström-based variant that uses structured random projections to improve convergence for matrices with rapidly decaying spectra.
• Apply the exchangeability principle to derive estimators that are symmetric functions of the random test vectors, ensuring optimal statistical properties.
• Introduce a posteriori error estimator that reflects actual estimation error by leveraging sample correlations across independent runs.
• Use the theoretical framework of U-statistics and exchangeable random variables to justify the optimality of the proposed estimators.

Experimental results

Research questions

• RQ1 Can the exchangeability principle be systematically applied to design better stochastic trace estimators with lower variance?
• RQ2 How does enforcing symmetry in the estimator function improve performance compared to non-exchangeable methods like Girard–Hutchinson?
• RQ3 To what extent can variance reduction via low-rank approximation and structured sampling improve trace estimation accuracy?
• RQ4 How do the new estimators perform relative to Hutch++ and other state-of-the-art methods across different matrix spectral structures?
• RQ5 Can an exchangeable diagonal estimator (XDiag) be constructed that matches or exceeds the accuracy of existing diagonal approximation techniques?

Key findings

• For a fixed number of matrix-vector products, XTrace and XNysTrace reduce estimation error by multiple orders of magnitude compared to the Girard–Hutchinson and Hutch++ estimators.
• Theoretical analysis shows that the variance of XTrace scales favorably with the spectrum of the matrix, particularly when singular values decay slowly.
• Empirical results demonstrate that XTrace correlations are bounded and remain low (e.g., ≤0.06) even for matrices with rapidly decaying eigenvalues, indicating strong statistical stability.
• The posterior error estimator $ ext{err}^2$ accurately reflects the true mean squared error, with average estimates correct within a factor of 1.2 in practice.
• XNysTrace achieves superior performance for matrices with fast singular value decay due to its structured sampling approach, outperforming XTrace and Hutch++ in such regimes.
• XDiag provides a new, exchangeable method for diagonal approximation that maintains low variance and is competitive with existing techniques.