{"id":1399,"date":"2023-01-26T19:29:34","date_gmt":"2023-01-26T19:29:34","guid":{"rendered":"https:\/\/www.ethanepperly.com\/?p=1399"},"modified":"2024-01-28T02:56:22","modified_gmt":"2024-01-28T02:56:22","slug":"stochastic-trace-estimation","status":"publish","type":"post","link":"https:\/\/www.ethanepperly.com\/index.php\/2023\/01\/26\/stochastic-trace-estimation\/","title":{"rendered":"Stochastic Trace Estimation"},"content":{"rendered":"\n

I am delighted to share that me, Joel A. Tropp<\/a>, and Robert J. Webber<\/a>‘s paper XTrace: Making the Most of Every Sample in Stochastic Trace Estimation<\/a><\/em> has recently been released as a preprint on arXiv. In it, we consider the implicit trace estimation problem:<\/p>\n\n\n\n

Implicit trace estimation problem:<\/strong> Given access to a square matrix \"A\" via the matrix\u2013vector product operation \"\omega \mapsto A\omega\", estimate its trace \"\tr A = \sum_{i=1}^n A_{ii}\".<\/p><\/blockquote>\n\n\n\n

Algorithms for this task have many uses such as log-determinant computations in machine learning<\/a>, partition function calculations in statistical physics<\/a>, and generalized cross validation for smoothing splines<\/a>. I described another application to counting triangles in a large network<\/a> in a previous blog post<\/a>.<\/p>\n\n\n\n

Our paper presents new trace estimators XTrace and XNysTrace which are highly efficient, producing accurate trace approximations using a small budget of matrix\u2013vector products. In addition, these algorithms are fast to run<\/a> and are supported by theoretical results which explain their excellent performance. I really hope that you will check out the paper<\/a> to learn more about these estimators!<\/p>\n\n\n\n

For the rest of this post, I’m going to talk about the most basic stochastic trace estimation algorithm, the Girard<\/a>\u2013Hutchinson<\/a> estimator. This seemingly simple algorithm exhibits a number of nuances and forms the backbone for more sophisticated trace estimates such as Hutch++<\/a>, Nystr\u00f6m++<\/a>, XTrace, and XNysTrace<\/a>. Toward the end, this blog post will be fairly mathematical, but I hope that the beginning will be fairly accessible to all.<\/p>\n\n\n\n

Girard\u2013Hutchinson Estimator: The Basics<\/h2>\n\n\n\n

The Girard<\/a>\u2013Hutchinson<\/a> estimator for the trace of a square matrix \"A\" is<\/p>\n\n\n\n

(1) <\/span>   <\/span>\"\[\hat{\tr} = \frac{1}{m} \sum_{i=1}^m \omega_i^* A \omega_i. \]\"<\/p><\/p>\n\n\n\n

Here, \"\omega_1,\ldots,\omega_m\" are random vectors, usually chosen to be statistically independent<\/a>, and \"{}^*\" denotes the conjugate transpose<\/a> of a vector or matrix. The Girard\u2013Hutchinson estimator only depends on the matrix \"A\" through the matrix\u2013vector products \"A\omega_1,\ldots,A\omega_m\".<\/p>\n\n\n\n

Unbiasedness<\/h3>\n\n\n\n

Provided the random vectors are isotropic<\/em><\/a><\/p>\n\n\n\n

(2) <\/span>   <\/span>\"\[\mathbb{E} [\omega_i\omega_i^*] = I, \]\"<\/p><\/p>\n\n\n\n

the Girard\u2013Hutchinson estimator is unbiased<\/a>:<\/p>\n\n\n\n

(3) <\/span>   <\/span>\"\[\mathbb{E} [\hat{\tr}] = \tr A.\]\"<\/p><\/p>\n\n\n\n

Let us confirm this claim in some detail. First, we use linearity of expectation<\/a> to evaluate<\/p>\n\n\n\n

(4) <\/span>   <\/span>\"\[\mathbb{E} [\hat{\tr}] = \mathbb{E} \left[ \frac{1}{m} \sum_{i=1}^m \omega_i^*A\omega_i \right] = \frac{1}{m} \sum_{i=1}^m \mathbb{E} \left[ \omega_i^* A \omega_i\right]. \]\"<\/p><\/p>\n\n\n\n

Therefore, to prove that \"\mathbb{E} [\hat{\tr}] = \tr A\", it is sufficient to prove that \"\mathbb{E} \left[\omega_i^*A\omega_i\right] = \tr A\" for each \"i\".<\/p>\n\n\n\n

When working with traces, there are two tricks that solve 90% of derivations. The first trick is that, if we view a number as a \"1\times 1\" matrix, then a number equals its trace, \"x = \tr x\". The second trick is the cyclic property<\/a>: For a \"k\times p\" matrix \"B\" and a \"p\times k\" matrix \"C\", we have \"\tr (BC) = \tr (CB)\". The cyclic property should be handled with care<\/a> when one works with a product of three or more matrices. For instance, we have<\/p>\n\n\n\n

  <\/span>   <\/span>\"\[\tr[BCD] = \tr[(BC)D] = \tr[D(BC)] = \tr[DBC].\]\"<\/p><\/p>\n\n\n\n

However,<\/p>\n\n\n\n

  <\/span>   <\/span>\"\[\tr [BCD] \ne \tr[CBD] \quad \text{in general}.\]\"<\/p><\/p>\n\n\n\n

One should think of the matrix product \"BCD\" as beads on a closed loop of string. One can move the last bead \"D\" to the front of the other two, \"\tr [BCD] = \tr[DBC]\", but not interchange two beads, \"\tr[BCD] \ne \tr[CBD]\".<\/p>\n\n\n\n

With this trick in hand, let’s return to proving that \"\mathbb{E} \left[\omega_i^*A\omega_i\right] = \tr A\" for every \"i\". Apply our two tricks:<\/p>\n\n\n\n

  <\/span>   <\/span>\"\[\mathbb{E} \left[\omega_i^*A\omega_i\right] = \mathbb{E} \tr \left[\omega_i^*A\omega_i\right] = \mathbb{E} \tr \left[A\omega_i\omega_i^*\right].\]\"<\/p><\/p>\n\n\n\n

The expectation is a linear<\/em> operation and the matrix \"A\" is non-random, so we can bring the expectation into the trace as<\/p>\n\n\n\n

  <\/span>   <\/span>\"\[\mathbb{E} \left[\omega_i^*A\omega_i\right] = \mathbb{E} \tr \left[A\omega_i\omega_i^*\right] = \tr(A \mathbb{E}[\omega_i\omega_i^*] ).\]\"<\/p><\/p>\n\n\n\n

Invoke the isotropy condition (2) and conclude:<\/p>\n\n\n\n

  <\/span>   <\/span>\"\[\mathbb{E} \left[\omega_i^*A\omega_i\right] = \tr(A \mathbb{E}[\omega_i\omega_i^*] ) = \tr(A\cdot I) = \tr A.\]\"<\/p><\/p>\n\n\n\n

Plugging this into (4) confirms the unbiasedness claim (3).<\/p>\n\n\n\n

Variance<\/h3>\n\n\n\n

Continue to assume that the \"\omega_i\"‘s are isotropic (3) and now assume that \"\omega_1,\ldots,\omega_m\" are independent. By independence<\/a>, the variance<\/a> can be written as<\/p>\n\n\n\n

  <\/span>   <\/span>\"\[\Var(\hat{\tr}) = \frac{1}{m^2} \sum_{i=1}^m \Var(\omega_i^*A\omega_i).\]\"<\/p><\/p>\n\n\n\n

Assuming that \"\omega_1,\ldots,\omega_m\" are identically distributed<\/em> \"\omega_1,\ldots,\omega_m \sim \omega\", we then get<\/p>\n\n\n\n

  <\/span>   <\/span>\"\[\Var(\hat{\tr}) = \frac{1}{m} \Var(\omega^*A\omega).\]\"<\/p><\/p>\n\n\n\n

The variance decreases like \"1/m\", which is characteristic of Monte Carlo<\/em>-type algorithms<\/a>. Since \"\hat{\tr}\" is unbiased (i.e, (3)), this means that the mean square<\/em> error decays like \"1/m\" so the average error (more precisely root-mean-square error) decays like<\/p>\n\n\n\n

  <\/span>   <\/span>\"\[\left| \hat{\tr} - \tr A \right| \lessapprox \frac{\mathrm{const}}{\sqrt{m}}.\]\"<\/p><\/p>\n\n\n\n

This type of convergence is very slow. If I want to decrease the error by a factor of \"10\", I must do \"100\times\" the work!<\/p>\n\n\n\n

Variance-reduced trace estimators like Hutch++<\/a> and our new trace estimator XTrace<\/a> improve the rate of convergence substantially. Even in the worst case, Hutch++ and XTrace reduce the variance at a rate \"1/m^2\" and (root-mean-square) error at rates \"1/m\":<\/p>\n\n\n\n

  <\/span>   <\/span>\"\[\Var(\hat{\tr}_{\text{H++ or X}}) \le \frac{\mathrm{const}}{m^2},\quad \left| \hat{\tr}_{\text{H++ or X}} - \tr A \right| \lessapprox \frac{\mathrm{const}}{m}.\]\"<\/p><\/p>\n\n\n\n

For matrices with rapidly decreasing singular values, the variance and error can decrease much faster than this.<\/p>\n\n\n\n

Variance Formulas<\/h2>\n\n\n\n

As the rate of convergence for the Girard\u2013Hutchinson estimator is so slow, it is imperative to pick a distribution on test vectors \"\omega\" that makes the variance of the single\u2013sample estimate \"\omega^*A\omega\" as low as possible. In this section, we will provide several explicit formulas for the variance of the Girard\u2013Hutchinson estimator. Derivations of these formulas will appear at the end of this post. These variance formulas help illuminate the benefits and drawbacks of different test vector distributions.<\/p>\n\n\n\n

To express the formulas, we will need some notation. For a complex number \"z = a + bi\" we use \"\Re(z) = a\" and \"\Im(z) = b\" to denote the real and imaginary parts. The variance of a random complex number \"z\" is<\/p>\n\n\n\n

  <\/span>   <\/span>\"\[\Var(z) := \mathbb{E} |z - \mathbb{E} z|^2 = \Var(\Re z) + \Var(\Im z).\]\"<\/p><\/p>\n\n\n\n

The Frobenius norm<\/a> of a matrix \"A\" is<\/p>\n\n\n\n

  <\/span>   <\/span>\"\[\left\|A\right\|_{\rm F}^2 = \sum_{i,j} |A_{ij}|^2.\]\"<\/p><\/p>\n\n\n\n

If \"A\" is real symmetric<\/a> or complex Hermitian<\/a> with (real<\/a>) eigenvalues \"\lambda_1,\ldots,\lambda_n\", we have<\/p>\n\n\n\n

(5) <\/span>   <\/span>\"\[\left\|A\right\|_{\rm F}^2 = \sum_{i=1}^n \lambda_i^2. \]\"<\/p><\/p>\n\n\n\n

\"A^\top\" denotes the ordinary transpose<\/a> of \"A\" and \"A^*\" denotes the conjugate transpose<\/a> of \"A\".<\/p>\n\n\n\n

Real-Valued Test Vectors<\/h3>\n\n\n\n

We first focus on real-valued test vectors \"\omega\". Since \"\omega\" is real, we can use the ordinary transpose \"{}^\top\" rather than the conjugate transpose \"{}^*\". Since \"\omega^\top A\omega\" is a number, it is equal to its own transpose:<\/p>\n\n\n\n

  <\/span>   <\/span>\"\[\omega^\top A \omega = (\omega^\top A \omega)^\top = \omega^\top A^\top \omega.\]\"<\/p><\/p>\n\n\n\n

Therefore,<\/p>\n\n\n\n

  <\/span>   <\/span>\"\[\omega^\top A\omega = \frac{\omega^\top A \omega + \omega^\top A^\top \omega}{2} = \omega^\top \left( \frac{A + A^\top}{2} \right)\omega.\]\"<\/p><\/p>\n\n\n\n

The Girard\u2013Hutchinson trace estimator applied to \"A\" is the same as the Girard\u2013Hutchinson estimator applied to the symmetric part<\/a> of \"A\", \"(A+A^\top)/2\".<\/p>\n\n\n\n

For the following results, assume \"A\" is symmetric, \"A = A^\top\".<\/p>\n\n\n\n

  1. Real Gaussian:<\/strong> \"\omega_1,\ldots,\omega_m\" are independent standard normal random vectors<\/a>.

      <\/span>   <\/span>\"\[\Var(\omega^\top A\omega) = 2 \left\|A\right\|_{\rm F}^2.\]\"<\/p><\/li>

  2. Uniform signs (Rademachers):<\/strong> \"\omega_1,\ldots,\omega_m\" are independent random vectors with uniform \"\pm 1\" coordinates<\/a>.

      <\/span>   <\/span>\"\[\Var(\omega^\top A \omega) = 2\sum_{i\ne j} |A_{ij}|^2.\]\"<\/p><\/li>

  3. Real sphere:<\/strong> Assume \"\omega_1,\ldots,\omega_n\" are uniformly distributed on the real sphere of radius \"\sqrt{n}\": \"\omega \sim \text{Uniform} \{x\in \mathbb{R}^n : x^\top x = n\}\".

      <\/span>   <\/span>\"\[\Var(\omega^\top A\omega) = \frac{2n}{n+2} \left( \left\|A\right\|_{\rm F}^2 - \frac{1}{n} |\tr A|^2 \right).\]\"<\/p><\/li><\/ol>\n\n\n\n

    These formulas continue to hold for nonsymmetric \"A\" by replacing \"A\" by its symmetric part \"(A+A^\top)/2\" on the right-hand sides of these variance formulas.<\/p>\n\n\n\n

    Complex-Valued Test Vectors<\/h3>\n\n\n\n

    We now move our focus to complex-valued test vectors \"\omega\". As a rule of thumb, one should typically expect that the variance for complex-valued test vectors applied to a real symmetric matrix \"A\" is about half the natural real counterpart\u2014e.g., for complex Gaussians, you get about half the variance than with real Gaussians.<\/p>\n\n\n\n

    A square complex matrix has a Cartesian decomposition<\/a><\/p>\n\n\n\n

      <\/span>   <\/span>\"\[A = A^{\rm H} + i A^{\rm SH}\]\"<\/p><\/p>\n\n\n\n

    where<\/p>\n\n\n\n

      <\/span>   <\/span>\"\[A^{\rm H} = \frac{A+A^*}{2} ,\quad A^{\rm SH} = \frac{A - A^*}{2i}\]\"<\/p><\/p>\n\n\n\n

    denote the Hermitian and skew-Hermitian parts of \"A\". Similar to how the imaginary part<\/em> of a complex number is real<\/em>, the skew-Hermitian<\/em> part of a complex matrix is Hermitian<\/em> (and \"i A^{\rm SH}\" is skew-Hermitian). Since \"A^{\rm H}\" and \"A^{\rm SH}\" are both Hermitian, we have<\/p>\n\n\n\n

      <\/span>   <\/span>\"\[\Re(\omega^* A\omega) = \omega^* A^{\rm H} \omega, \quad \Im (\omega^* A \omega) = \omega^* A^{\rm SH} \omega.\]\"<\/p><\/p>\n\n\n\n

    Consequently, the variance of \"\omega^*A \omega\" can be broken into Hermitian and skew-Hermitian parts:<\/p>\n\n\n\n

      <\/span>   <\/span>\"\[\Var(\omega^* A\omega) = \Var(\omega^* A^{\rm H}\omega) + \Var(\omega^* A^{\rm SH}\omega).\]\"<\/p><\/p>\n\n\n\n

    For this reason, we will state the variance formulas only for Hermitian \"A\", with the formula for general \"A\" following from the Cartesian decomposition.<\/p>\n\n\n\n

    For the following results, assume \"A\" is Hermitian, \"A = A^*\".<\/p>\n\n\n\n

    1. Complex Gaussian:<\/strong> \"\omega_1,\ldots,\omega_m\" are independent standard complex random vectors<\/a>, i.e., each \"\omega_i\" has iid entries distributed as \"(g_1+ig_2)/\sqrt{2}\" for \"g_1,g_2\" standard normal random variables.

        <\/span>   <\/span>\"\[\Var(\omega^* A\omega) = \left\|A\right\|_{\rm F}^2.\]\"<\/p><\/li>

    2. Uniform phases (Steinhauses):<\/strong> \"\omega_1,\ldots,\omega_m\" are independent random vectors whose entries are uniform on the complex unit circle \"\{ z \in \complex : |z| \}\".

        <\/span>   <\/span>\"\[\Var(\omega^* A \omega) = \sum_{i\ne j} |A_{ij}|^2.\]\"<\/p><\/li>

    3. Complex sphere:<\/strong> Assume \"\omega_1,\ldots,\omega_n\" are uniformly distributed on the complex sphere of radius \"\sqrt{n}\": \"\omega \sim \text{Uniform} \{x\in \complex^n : x^* x = n\}\".

        <\/span>   <\/span>\"\[\Var(\omega^* A\omega) = \frac{n}{n+1} \left( \left\|A\right\|_{\rm F}^2 - \frac{1}{n} |\tr A|^2 \right).\]\"<\/p><\/li><\/ol>\n\n\n\n

      Optimality Properties<\/h2>\n\n\n\n

      Let us finally address the question of what the best<\/em> choice of test vectors is for the Girard\u2013Hutchinson estimator. We will state two results with different restrictions on \"\omega_1,\ldots,\omega_m\".<\/p>\n\n\n\n

      Our first result, due to Hutchinson<\/a>, is valid for real symmetric matrices with real test vectors.<\/p>\n\n\n\n

      Optimality (independent test vectors with independent coordinates).<\/strong> If the test vectors \"\omega_1,\ldots,\omega_m \in \mathbb{R}^n\" are isotropic (2), independent from each other, and have independent entries, then for any fixed real symmetric matrix \"A\", the minimum variance for \"\hat{\tr}\" is obtained when \"\omega_1,\ldots,\omega_m\" are populated with random signs \"(\omega_i)_j \sim \textnormal{Uniform} \{\pm 1\}\".<\/p><\/blockquote>\n\n\n\n

      The next optimality results will have real and complex versions. To present the results for \"\mathbb{R}\"-valued and an \"\complex\"-valued test vectors on unified footing, let \"\field\" denote either \"\mathbb{R}\" or \"\complex\". We let a \"\field\"-Hermitian matrix be either a real symmetric matrix (if \"\field = \mathbb{R}\") or a complex Hermitian matrix (if \"\field = \complex\"). Let a \"\field\"-unitary matrix be either a real orthogonal matrix<\/a> (if \"\field = \mathbb{R}\") or a complex unitary matrix<\/a> (if \"\field = \complex\").<\/p>\n\n\n\n

      The condition that the vectors \"\omega_1,\ldots,\omega_m\" have independent entries<\/em> is often too restrictive in practice. It rules out, for instance, the case of uniform vectors on the sphere. If we relax this condition, we get a different optimal distribution:<\/p>\n\n\n\n

      Optimality (independent test vectors).<\/strong> Consider any set \"\mathscr{A}\" of \"\field\"-Hermitian matrices which is invariant under \"\field\"-unitary similary transformations:

        <\/span>   <\/span>\"\[\text{If $A \in \mathscr{A}$ and $U$ is $\field$-unitary, then $U^*AU \in \mathscr{A}$.}\]\"<\/p>Assume that the test vectors \"\omega_1,\ldots,\omega_m\" are independent and isotropic (2). The worst-case variance \"\sup_{A \in \mathscr{A}} \Var(\hat{\tr})\" is minimized by choosing \"\omega_1,\ldots,\omega_m\" uniformly on the \"\field\"-sphere: \"\omega_1,\ldots,\omega_m \sim \text{Uniform} \{ x \in \field^n : x^*x =n \}\".<\/p><\/blockquote>\n\n\n\n

      More simply, if you wants your stochastic trace estimator to be effective for a class<\/em> of inputs \"\mathscr{A}\" (closed under \"\field\"-unitary similarity transformations) rather than a single input matrix \"A\", then the best distribution are test vectors drawn uniformly from the sphere. Examples of classes of matrices \"\mathscr{A}\" include:<\/p>\n\n\n\n