Updates for Hutch++

At SOSA 2021, I presented the paper Hutch++: Optimal Stochastic Trace Estimation (with Cameron Musco, Christopher Musco, and David P. Woodruff). This webpage contains some short updates that may be useful to people.

Table of Contents:

Hutch++ for Undergrads

I believe that the analysis of both Hutchinson's Estimator and Hutch++ is simple and approachable. This section contains the analysis an advanced undergraduate should hopefully be able to understand.

Part 1: The Fundamental Shape of Hutchinson's Estimator

We can easily verify two properties of this classical estimator for symmetric $\boldsymbol{ A}$:

This variance is in term of the Frobenius norm of $\boldsymbol{ A}$. So, we expect that Hutchinson's Estimator satisfies $H_\ell(\boldsymbol{ A}) \in \text{tr}(\boldsymbol{ A}) \pm \frac{\sqrt2}{\sqrt\ell} \| \boldsymbol{ A}\|_{ F}$ with high probability.^[1] Typically in computer science we are given an error tolerance $\varepsilon$, and want to know how many matrix-vector products we need to compute. (i.e. how small can we make $\ell$ while still being accurate?) In order to build a guarantee here, we assume $\boldsymbol{ A}$ is positive semi-definite (PSD), which in turn implies that $\| \boldsymbol{ A}\|_{ F} \leq \text{tr}(\boldsymbol{ A})$. With this in mind, we can prove the standard result for Hutchinson's Estimator:

So, if we want to ensure $H_\ell(\boldsymbol{ A})\in(1\pm\varepsilon)\text{tr}(\boldsymbol{ A})$ with good probability, we need the standard deviation to have $\frac{\sqrt 2}{\sqrt \ell}\text{tr}(\boldsymbol{ A}) \leq \varepsilon \text{tr}(\boldsymbol{ A})$, so we need $\ell=O(\frac{1}{\varepsilon^2})$ samples.

Part 2: Hutchinson's versus the Top Few Eigenvalues

This was a section in the first draft of the paper, but was unfortunately cut. This is a slow, methodical intuition for the geometry Hutch++ takes advantage of.

Hutchinson's Estimator is powerfull, and gives a nice rate of convergence. However, the proof of Lemma 2 has a suspicious step. If we write out three lines of analysis:

Looking at these inequalities,

• The analysis of line 1 is always tight, since we know the variance of $H_\ell(\boldsymbol{ A})$ exactly.

• The analysis of line 3 is always tight, since we just set $\ell$ to the smallest value that gets error $\varepsilon$.

• The analysis of line 2 is not always tight.

The second line is tight only if $\| \boldsymbol{ A}\|_F \approx \text{tr}(\boldsymbol{ A})$. That is, the analysis above only tells us that Hutchinson's Estimator needs $O(\frac1\varepsilon)$ samples if $\boldsymbol{ A}$ is the kind of matrix that has $\| \boldsymbol{ A}\|_F \approx \text{tr}(\boldsymbol{ A})$. So, what kind of matrix is that?

We now finish the discussion with a comparison of two ideas:

• So, Hutchinson's Estimator only needs to use many samples (i.e. $O(\frac1\varepsilon)$ samples) if $\boldsymbol{ A}$ has very special structure: it has a small number of large eigenvalues.

• If a matrix has a small number of large eigenvalues, the trace must be well approximated by the sum of those eigenvalues.

With this, we conclude the fundamental intuition behind Hutch++:

This leads us to pick the following rough design of an algorithm:

1. Find a good low-rank approximation $\tilde\boldsymbol{ A}_k$

2. Notice that $\text{tr}(\boldsymbol{ A}) = \text{tr}(\tilde \boldsymbol{ A}_k) + \text{tr}(\boldsymbol{ A} - \tilde\boldsymbol{ A}_k)$

3. Compute $\text{tr}(\tilde \boldsymbol{ A}_k)$ exactly

4. Approximate $\text{tr}(\boldsymbol{ A}-\tilde\boldsymbol{ A}_k)$ with Hutchinson's Estimator

5. Return $\text{Hutch++}(\boldsymbol{ A}) = \text{tr}(\tilde\boldsymbol{ A}_k) + H_\ell(\boldsymbol{ A}-\tilde\boldsymbol{ A}_k)$

In the next section, we state a formal version of Claim 1 (see Lemma 3), show how to compute such a matrix $\tilde\boldsymbol{ A}_k$ (see $\boldsymbol{ Q}\boldsymbol{ Q}^\intercal\boldsymbol{ A}$ in Theorem 1), and bound the complexity of the Hutch++ estimator (see Theorem 2).

Part 3: The Variance of Hutch++

Special thanks to Joel A. Tropp for working most of this out, and encouraging us to look into sharpening the constants in the analysis of Hutch++.

By expressing the variance of Hutch++ with all of its constants laid bare, and by using a very simple analysis, this analysis will hopefully allow practitioners to easily the exact parameters in Hutch++ code.

Before bounding the variance of Hutch++, we include the proof of one lemma and import another theorem:

We also import the following theorem, which controls the expected error from our low-rank approximation, and which we do not prove:

Now, we proceed to the analysis of the variance of Hutch++:

Part 4: Practitioner Advice

The variance in Theorem 2 still has two parameters in it, which leave a bit of ambiguity for practitioners. So, we quickly analyze the choice of $k$ and $\ell$ that minimizes our bound on the variance.

Formally, suppose we are allowed to compute exactly $m$ matrix-vector products with $\boldsymbol{ A}$. Then,

• Computing $\boldsymbol{ Q}$ uses $k+p=2k+1$ products

• Computing $\text{tr}(\boldsymbol{ Q}^\intercal\boldsymbol{ A}\boldsymbol{ Q})$ uses $k+p=2k+1$ products

• Computing $H_\ell( (\boldsymbol{ I}-\boldsymbol{ Q}\boldsymbol{ Q}^\intercal)\boldsymbol{ A})$ uses $\ell$ products

So, we have $m=2(2k+1)+\ell=4k+2+\ell$. We can then verify that $k=\frac{m-2}{8}$ and $\ell=\frac{m}{2}-1$ minimizes the variance. Notably, this is equivalent to setting $\ell=4k$, which looks more intuitive. This produces a final variance of $\text{Var}[\text{Hutch++}(\boldsymbol{ A})]\leq\frac{16}{(m-2)^2}\text{tr}^2(\boldsymbol{ A})$.

Nyström-Hutch++

One additional approach to implementing Hutch++ for PSD matrices considers the Nyström method, where we use the approximation

Note that we can compute $\tilde\boldsymbol{ A}_k=\boldsymbol{ Y}(\boldsymbol{ Q}^\intercal\boldsymbol{ Y})^{-1}\boldsymbol{ Y}^\intercal$ where $\boldsymbol{ Y}=\boldsymbol{ A}\boldsymbol{ Q}$ contains all the matrix-vector products with $\boldsymbol{ A}$. We can compute $\text{tr}(\tilde\boldsymbol{ A}_k)$ efficiently:

The matrix inversion is only for a $(k+p) \times (k+p)$ sized matrix, which is easy. We can then define the NYS-Hutch++ estimator as

We can write out this Hutchinson's Estimator step as

where $\boldsymbol{ G}$ has random sign-bit entries and $\boldsymbol{ R}\;{\vcentcolon=}\;\boldsymbol{ Y}^\intercal\boldsymbol{ G}$. Formally, we can write this out as the following algorithm:

Or, in matlab:

function trace_est=simple_nystrom_hutchplusplus(A, num_queries)
    S = 2*randi(2,size(A,1),ceil(num_queries/4))-3;
    G = 2*randi(2,size(A,1),floor(num_queries/2))-3;
    [Q,~] = qr(A*S,0);
    Y = A*Q;
    R = Y' * G;
    QYinv = pinv(Q' * Y);
    trace_est = trace(QYinv * (Y'*Y)) + 1/size(G,2)*[trace(G'*A*G) - trace(R'*QYinv*R)];
end

Epilogue

Code

• We provide code on Github, written in MATLAB.

• Akshay Agrawal also has an implementation on Github using pytorch, though we haven't verified this code ourselves.

• Email me if you want to add your implementation to the list!

Quiz Questions

If you have read Hutch++, then quiz yourself by trying to answer some intuitive high-level questions:

1. Why is the constant-factor approximation on our low-rank approximation $\|A-\tilde{A}_k\|_F \leq 2 \|A-A_k\|_F$ sufficient? Why don't we need a relative error approximation like $\|A-\tilde{A}_k\|_F \leq (1+\varepsilon) \|A-A_k\|_F$?

Footnotes

^[1] We analyze variance in this page for convenience, but all of these confidence intervals hold with high probability (i.e. with logarithmic dependence on the failure probability), and this is widely analyzed in the formal publications.

References

• Gilbert, Strauss, Tropp, and Vershynin. One sketch for all: fast algorithms for compressed sensing. STOC 2007.

• Halko, Martinsson, and Tropp. Finding structure with randomness: Probabilistic algorithms for constructing approximate matrix decompositions. SIAM review 53 no. 2 2011.

• Avron and Toledo. Randomized algorithms for estimating the trace of an implicit symmetric positive semi-definite matrix. Journal of the ACM (JACM) 58, no. 2 2011.