Optimal Sketching for Trace Estimation
Year: 2022

Abstract:
Matrix trace estimation is ubiquitous in machine learning applications and has traditionally relied on Hutchinson’s method, which requires O(log(1/δ)/ ) matrixvector product queries to achieve a (1± )-multiplicative approximation to tr(A) with failure probability δ on positive-semidefinite input matrices A. Recently, the Hutch++ algorithm was proposed, which reduces the number of matrix-vector queries from O(1/ ) to the optimal O(1/ ), and the algorithm succeeds with constant probability. However, in the high probability setting, the non-adaptive Hutch++ algorithm suffers an extraO( √ log(1/δ)) multiplicative factor in its query complexity. Non-adaptive methods are important, as they correspond to sketching algorithms, which are mergeable, highly parallelizable, and provide low-memory streaming algorithms as well as low-communication distributed protocols. In this work, we close the gap between non-adaptive and adaptive algorithms, showing that even non-adaptive algorithms can achieve O( √ log(1/δ)/ + log(1/δ)) matrixvector products. In addition, we prove matching lower bounds demonstrating that, up to a log log(1/δ) factor, no further improvement in the dependence on δ or is possible by any non-adaptive algorithm. Finally, our experiments demonstrate the superior performance of our sketch over the adaptive Hutch++ algorithm, which is less parallelizable, as well as over the non-adaptive Hutchinson’s method.

text: Matrix trace estimation is ubiquitous in machine learning applications and has traditionally relied on Hutchinson’s method, which requires O(log(1/δ)/ 2) matrixvector product queries to achieve a (1± )-multiplicative approximation to tr(A) with failure probability δ on positive-semidefinite input matrices A. Recently, the Hutch++ algorithm was proposed, which reduces the number of matrix-vector queries from O(1/ 2) to the optimal O(1/ ), and the algorithm succeeds with constant probability. However, in the high probability setting, the non-adaptive Hutch++ algorithm suffers an extraO( √ log(1/δ)) multiplicative factor in its query complexity. Non-adaptive methods are important, as they correspond to sketching algorithms, which are mergeable, highly parallelizable, and provide low-memory streaming algorithms as well as low-communication distributed protocols. In this work, we close the gap between non-adaptive and adaptive algorithms, showing that even non-adaptive algorithms can achieve O( √ log(1/δ)/ + log(1/δ)) matrixvector products. In addition, we prove matching lower bounds demonstrating that, up to a log log(1/δ) factor, no further improvement in the dependence on δ or is possible by any non-adaptive algorithm. Finally, our experiments demonstrate the superior performance of our sketch over the adaptive Hutch++ algorithm, which is less parallelizable, as well as over the non-adaptive Hutchinson’s method.

heading: 1 Introduction
text: The problem of implicit matrix trace estimation arises naturally in a wide range of applications [1]. For example, during the training of Gaussian Process, a popular non-parametric kernel-based method, the calculation of the marginal log-likelihood contains a heavy-computation term, i.e., the log determinant of the covariance matrix, log(det(K)), where K ∈ Rn×n, and n is the number of data points. The canonical way of computing log(det(K)) is via Cholesky decomposition on K, whose time complexity is O(n3). Since log(det(K)) = ∑n i=1 log(λi), where λi’s are the eigenvalues of K, one can compute tr(log(K)) instead. Trace estimation combined with polynomial approximation (e.g., the Chebyshev polynomial or Stochastic Lanczos Quadrature) to log [2], or trace estimation combined with maximum entropy estimation [3] provide fast ways of estimating tr(log(K)) for large-scale data. Other popular applications of implicit trace estimation include counting triangles and computing the Estrada Index in graphs [4, 5], approximating the generalized rank of a matrix [6], and studying non-convex loss landscapes from the Hessian matrix of large neural networks (NNs) [7, 8].
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
To define the problem, we consider the matrix-vector product model as formalized in [9, 10], where there is a real symmetric input matrix A ∈ Rn×n that cannot be explicitly presented but one has oracle access to A via matrix-vector queries, i.e., one can obtain Aq for any desired query vector q ∈ Rn. For example, due to a tremendous amount of trainable parameters of large NNs, it is often prohibitive to compute or store the entire Hessian matrix H with respect to some loss function from the parameters [7], which is often used to study the non-convex loss landscape. However, with Pearlmutter’s trick [11] one can compute Hq for any chosen vector q. The goal is to efficiently estimate the trace of A, denoted by tr(A), up to error, i.e., to compute a quantity within (1± )tr(A). For efficiency, such algorithms are randomized and succeed with probability at least 1 − δ. The minimum number of queries q required to solve the problem is referred to as the query complexity.
Computing matrix-vector products Aq through oracle access, however, can be costly. For example, computing Hessian-vector products Hq on large NNs takes approximately twice the time of backpropagation. When estimating the eigendensity of H , one computes tr(f(H)) for some density function f , and needs repeated access to the matrix-vector product oracle. As a result, even with Pearlmutter’s trick and distributed computation on modern GPUs, it takes 20 hours to compute the eigendensity of a single Hessian H with respect to the cross-entropy loss on the CIFAR-10 dataset [12], from a set of fixed weights for ResNet-18 [13] which has approximately 11 million parameters [7]. Thus, it is important to understand the fundamental limits of implicit trace estimation as the query complexity in terms of the desired approximation error and the failure probability δ.
Hutchinson’s method [14], a simple yet elegant randomized algorithm, is the ubiquitous work force for implicit trace estimation. Letting Q = [q1, . . . ,qq] ∈ Rn×q be q vectors with i.i.d. Gaussian or Rademacher (i.e., ±1 with equal probability) random variables, Hutchinson’s method returns an estimate of tr(A) as 1q ∑q i=1 q T i Aqi = 1 q tr(Q
TAQ). Although Hutchinson’s method dates back to 1990, it is surprisingly not well-understood on positive semi-definite (PSD) matrices. It was originally shown that for PSD matrices A with the qi being Gaussian random variables, in order to obtain a multiplicative (1 ± ) approximation to tr(A) with probability at least 1 − δ, O(log(1/δ)/ 2) matrix-vector queries suffice [15].
A recent work [16] proposes a variance-reduced version of Hutchinson’s method that shows only O(1/ ) matrix-vector queries are needed to achieve a (1± )-approximation to any PSD matrix with constant success probability, in contrast to theO(1/ 2) matrix-vector queries needed for Hutchinson’s original method. The key observation is that the variance of the estimated trace in Hutchinson’s method is largest when there is a large gap between the top few eigenvalues and the remaining ones. Thus, by splitting the number of matrix-vector queries between approximating the top O(1/ ) eigenvalues, i.e., by computing a rank-O(1/ ) approximation to A, and performing trace estimation on the remaining part of the spectrum, one needs only O(1/ ) queries in total to achieve a (1± ) approximation to tr(A). Furthermore, [16] shows Ω(1/ ) queries are in fact necessary for any trace estimation algorithm, up to a logarithmic factor, for algorithms succeeding with constant success probability. While [16] mainly focuses on the improvement on in the query complexity with constant failure probability, we focus on the dependence on the failure probability δ.
Algorithm 1 Hutch++: Stochastic trace estimation with adaptive matrix-vector queries 1: Input: Matrix-vector multiplication oracle for
PSD matrix A ∈ Rn×n. Number m of queries. 2: Output: Approximation to tr(A). 3: Sample S ∈ Rn× m 3 and G ∈ Rn× m 3 with i.i.d.
N (0, 1) entries. 4: Compute an orthonormal basis Q ∈ Rn× m 3 for
the span of AS via QR decomposition. 5: return t = tr(QTAQ) + 3
m tr(GT (I −
QQT )A(I −QQT )G).
Algorithm 2 NA-Hutch++: Stochastic trace estimation with non-adaptive matrix-vector queries 1: Input: Matrix-vector multiplication oracle for
PSD matrix A ∈ Rn×n. Number m of queries. 2: Output: Approximation to tr(A). 3: Fix constants c1, c2, c3 such that c1 < c2 and c1 +
c2 + c3 = 1. 4: Sample S ∈ Rn×c1m, R ∈ Rn×c2m, and G ∈
Rn×c3m, with i.i.d. N (0, 1) entries. 5: Z = AR, W = AS 6: return t = tr((STZ)†(W TZ)) +
1 c3m (tr(GTAG)− tr(GTZ(STZ)†W TG)).
Achieving a low failure probability δ is important in applications where failures are highly undesirable, and the low failure probability regime is well-studied in related areas such as compressed sensing [17], data stream algorithms [18, 19], distribution testing [20], and so on. While one can always reduce the failure probability from a constant to δ by performingO(log(1/δ)) independent repetitions
and taking the median, this multiplicative overhead of O(log(1/δ)) can cause a huge slowdown in practice, e.g., in the examples above involving large Hessians.
Two algorithms were proposed in [16]: Hutch++ (Algorithm 1), which requires adaptively chosen matrix-vector queries and NA-Hutch++ (Algorithm 2) which only requires non-adaptively chosen queries. We call the matrix-vector queries adaptively chosen if subsequent queries are dependent on previous queries q and observations Aq, whereas the algorithm is non-adaptive if all queries can be chosen at once without any prior information about A. Note that Hutchinson’s method uses only non-adaptive queries. [16] shows that Hutch++ can use O( √ log(1/δ)/ + log(1/δ)) adaptive matrix-vector queries to achieve (1± ) approximation with probability at least 1− δ, while NA-Hutch++ can use O(log(1/δ)/ ) non-adaptive queries. Thus, in many parameter regimes the non-adaptive algorithm suffers an extra √ log(1/δ) multiplicative factor over the adaptive algorithm.
It is important to understand the query complexity of non-adaptive algorithms for trace estimation because the advantages of non-adaptivity are plentiful: algorithms that require only non-adaptive queries can be easily parallelized across multiple machines while algorithms with adaptive queries are inherently sequential. Furthermore, non-adaptive algorithms correspond to sketching algorithms which are the basis for many streaming algorithms with low memory [21] or distributed protocols with low-communication overhead (for an example application to low rank approximation, see [22]). We note that there are numerous works on estimating matrix norms in a data stream [23, 24, 25, 26], most of which use trace estimation as a subroutine.

heading: 1.1 Our Contributions
text: Improving the Non-adaptive Query Complexity. We give an improved analysis of the query complexity of the non-adaptive trace estimation algorithm NA-Hutch++ (Algorithm 2), based on a new low-rank approximation algorithm and analysis in the high probability regime, instead of applying an off-the-shelf low-rank approximation algorithm as in [16]. Instead of O(log(1/δ)/ ) queries as shown in [16], we show that O( √ log(1/δ)/ + log(1/δ)) non-adaptive queries suffice to achieve a multiplicative (1 ± ) approximation of the trace with probability at least 1 − δ, which matches the query complexity of the adaptive trace estimation algorithm Hutch++. Since our algorithm is non-adaptive, it can be used in subroutines in streaming and distributed settings for estimating the trace, with lower memory than was previously possible for the same failure probability. Theorem 1.1 (Restatement of Theorem 3.1). Let A be any PSD matrix. If NA-Hutch++ is im-
plemented with m = O (√ log(1/δ) + log(1/δ) ) matrix-vector multiplication queries, then with
probability 1− δ, the output t of NA-Hutch++ satisfies (1− )tr(A) ≤ t ≤ (1 + )tr(A).
The improved dependence on δ is perhaps surprising in the non-adaptive setting, as simply repeating a constant-probability algorithm would give an O(log(1/δ)/ ) dependence. Our non-adaptive algorithm is as good as the best known adaptive algorithm, and much better than previous nonadaptive algorithms [16, 14]. The key difference between our analysis and the analysis in [16] is in the number of non-adaptive matrix-vector queries we need to obtain an O(1)-approximate rank-k approximation to A in Frobenius norm.
Specifically, to reduce the total number of matrix-vector queries, our queries are split between (1) computing Ã, a rank-k approximation to the matrix A, and (2) performing trace estimation on A−Ã. Let Ak = minrank-kA ‖A−Ak‖F be the best rank-k approximation to A in Frobenius norm. For our algorithm to work, we require ‖A− Ã‖ ≤ O(1)‖A−Ak‖F with probability 1− δ. Previous results from [27] show the number of non-adaptive queries required to compute Ã is O(k log(1/δ)), where each query is an i.i.d. Gaussian or Rademacher vector. We prove O(k + log(1/δ)) non-adaptive Gaussian query vectors suffice to compute Ã. Low rank approximation requires both a so-called subspace embedding and an approximate matrix product guarantee (see, e.g., [28], for a survey on sketching for low rank approximation), and we show both hold with the desired probability, with some case analysis, for Gaussian queries. A technical overview can be found in Section 3.
The improvement on the number of non-adaptive queries to achieve O(1)-approximate rank-k approximation has many other implications, which can be of an independent interest. For example, since low-rank approximation algorithms are extensively used in streaming algorithms suitable for low-memory settings, this new result directly improves the space complexity of the state-of-the-art
streaming algorithm for Principle Component Analysis (PCA) [29] from O(d · (k log(1/δ))) to O(d · (k + log(1/δ))) for constant approximation error , where d is the dimension of the input. Lower Bound. Previously, no lower bounds were known on the query complexity in terms of δ in a high probability setting. In this work, we give a novel matching lower bound for non-adaptive (i.e., sketching) algorithms for trace estimation, with novel techniques based on a new family of hard input distributions, showing that our improved O( √ log(1/δ)/ + log(1/δ)) upper bound is optimal, up to a log log(1/δ) factor, for any ∈ (0, 1). The methods previously used to prove an Ω(1/ ) lower bound with constant success probability (up to logarithmic factors) in [16] do not apply in the high probability setting. Indeed, [16] gives two lower bound methods based on a reduction from two types of problems: (1) a communication complexity problem, and (2) a distribution testing problem between clean and negatively spiked random covariance matrices. Technique (1) does not apply since there is not a multi-round lower bound for the Gap-Hamming communication problem used in [16] that depends on δ. One might think that since we are proving a non-adaptive lower bound, we could use a non-adaptive lower bound for Gap-Hamming (which exists, see [18]), but this is wrong because even the non-adaptive lower bound in [16] uses a 2-round lower bound for Gap-Hamming, and there is no such lower bound known in terms of δ. Technique (2) also does not apply, as it involves a 1/ × 1/ matrix, which can be recovered exactly with 1/ queries; further, increasing the matrix dimensions would break the lower bound as their two cases would no longer need to be distinguished. Thus, such a hard input distribution fails to show the additive Ω(log(1/δ)) term in the lower bound.
Our starting point for a hard instance is a family of Wigner matrices (see Definition 2.1) shifted by an identity matrix so that they are PSD. However, due to strong concentration properties of these matrices, they can only be used to provide a lower bound of Ω( √ log(1/δ)/ ) when < 1/ √ log(1/δ). Indeed, setting δ to be a constant in this case recovers the Ω(1/ ) lower bound shown in [16] but via a completely different technique. For larger , we consider a new distribution testing problem between clean Wigner matrices and the same distribution with a large rank-1 noisy PSD matrix, and then argue with probability roughly δ, all non-adaptive queries have unusually tiny correlation with this rank-1 matrix, thus making it indistinguishable between the two distributions. This gives the desired additive Ω(log(1/δ)) lower bound, up to a log log(1/δ) factor.
Theorem 1.2 (Restatement of Theorem 4.1). SupposeA is a non-adaptive query-based algorithm that returns a (1± )-multiplicative estimate to tr(A) for any PSD matrix A with probability at least 1−δ. Then, the number of matrix-vector queries must be at least m = Ω (√ log(1/δ) + log(1/δ)
log(log(1/δ))
) .

heading: 1.2 Related Work
text: A summary of prior work on the query complexity of trace estimation of PSD matrices is given in Table 1. For the upper bounds, prior to the work of [30], the analysis of implicit trace estimation mainly focused on the variance of estimation with different types of query vectors. [30] gave the first upper bound on the query complexity. The work of [15] improved the bounds in [30]. On the lower bound side, although [15] gives a necessary condition on the query complexity for Gaussian query vectors, this condition does not directly translate to a bound on the minimum number of query vectors. The work of [16] gives the first lower bound on the query complexity in terms of but only works for constant failure probability.

heading: 2 Problem Setting
text: Notation. A matrix A ∈ Rn×n is symmetric positive semi-definite (PSD) if it is real, symmetric and has non-negative eigenvalues. Hence, x>Ax ≥ 0 for all x ∈ Rn. Let tr(A) = ∑n i=1 Aii
denote the trace of A. Let ‖A‖F = ( ∑n i=1 ∑n j=1 A 2 ij)
1/2 denote the Frobenius norm and ‖A‖op = sup‖v‖2=1 ‖Av‖2 denote the operator norm of A. Let N (µ, σ
2) denote the Gaussian distribution with mean µ and variance σ2. Our analysis extensively relies on the following facts:
Definition 2.1 (Gaussian and Wigner Random Matrices). We let G ∼ N (n) denote an n×n random Gaussian matrix with i.i.d. N (0, 1) entries. We let W ∼ W(n) = G + GT denote an n× n Wigner matrix, where G ∼ N (n). Fact 2.1 (Rotational Invariance of a standard Gaussian). Let R ∈ Rn×n be an orthornormal matrix. Let g ∈ Rn be a random vector with i.i.d. N (0, 1) entries. Then Rg has the same distribution as g. Fact 2.2 (Upper and Lower Gaussian Tail Bounds). Letting Z ∼ N (0, 1) be a univariate Gaussian random variable, for any t > 0, Pr[|Z| ≥ t] = Θ(t−1 exp(− t 2
2 )).
3 An Improved Analysis of NA-Hutch++
Suppose we are trying to compute a sketch so as to estimate the trace of a matrix A up to a (1 ± )-factor with success probability at least 1 − δ. Note that we focus on the case where we make matrix-vector queries non-adaptively. For any algorithm that accomplishes this with small constant failure probability, one can simply repeat this procedure O(log(1/δ)) times to amplify the success probability to 1− δ. Since these queries are non-adaptive and must be presented before any observations are made, it seems intuitive that the number of non-adaptive queries of NA-Hutch++ (Algorithm 2) should be O(log(1/δ)/ ) as shown in [16]. In this section, we give a proof sketch as to why this can be reduced to O( √ log(1/δ)/ + log(1/δ)) as stated in Theorem 3.1. All proof details are provided in the supplementary material.
Theorem 3.1. Let A be a PSD matrix. If NA-Hutch++ is implemented with m = O( √
log(1/δ)/ + log(1/δ)) matrix-vector multiplication queries, then with probability 1−δ, the output of NA-Hutch++, denoted by t, satisfies (1− )tr(A) ≤ t ≤ (1 + )tr(A).
NA-Hutch++ splits its matrix-vector queries between computing an O(1)-approximate rank-k approximation Ã and performing Hutchinson’s estimate on the residual matrix A− Ã containing the small eigenvalues. The trade-off between the rank k and the number l of queries spent on estimating the small eigenvalues is summarized in Theorem 3.2.
Theorem 3.2 (Theorem 4 of [16]). Let A ∈ Rn×n be PSD, δ ∈ (0, 12 ), l ∈ N, k ∈ N. Let Ã and ∆ be any matrices with tr(A) = tr(Ã) + tr(∆) and ‖∆‖F ≤ O(1)‖A − Ak‖F where Ak = arg minrank k Ak ‖A−Ak‖F . Let Hl(M) denote Hutchinson’s trace estimator with l queries on matrix M . For fixed constants c, C, if l ≥ c log( 1δ ), then with probability 1 − δ, for Z =
tr(Ã) +Hl(∆), we have |Z − tr(A)| ≤ C √ log(1/δ) kl · tr(A).
The total number of matrix-vector queries directly depends on the number of non-adaptive queries required to compute an O(1)-approximate rank-k approximation Ã. Consider S ∈ Rn×c1m,R ∈ Rn×c2m for some constants c1, c2 > 0 as defined in Algorithm 2, and set our low rank approximation of A to be Ã = AR(STAR)†(AS)T . The standard analysis [16] applies a result from streaming low-rank approximation in [27], which requires m = O(k log(1/δ)) to get ‖A− Ã‖F ≤ O(1)‖A− Ak‖F with probability 1− δ. [16] then sets k = O(1/ ) and l = O(log(1/δ)/ ) in Theorem 3.2 to get a (1 ± ) approximation to tr(A). However, the right-hand side of Theorem 3.2 suggests the optimal split between k and l should be k = l. The reason [16] cannot achieve such an optimal split is due to a large number m of queries to compute the O(1)-approximate rank k-approximation. We give an improved analysis of this result, which may be of independent interest.
To get O(1) low rank approximation error, we need the non-adaptive query matrices S, R to satisfy two properties: the subspace embedding property (see Lemma 3.3), and an approximate matrix product for orthogonal subspaces (see Lemma 3.4). While it is known that m = O(k + log(1/δ))
suffices to achieve the first property, we show that m = O(k + log(1/δ)) suffices to achieve the second property when S,R are matrices with i.i.d. Gaussian random variables, stated in Lemma 3.4. Lemma 3.3 (Subspace Embedding (Theorem 6 of [28])). Given δ ∈ (0, 12 ) and ∈ (0, 1). Let S ∈ Rr×n be a random matrix with i.i.d. Gaussian random variables N (0, 1r ). Then for any fixed d-dimensional subspace A ∈ Rn×d, and for r = O((d + log( 1δ ))/
2), the following holds with probability 1− δ simultaneously for all x ∈ Rd, ‖SAx‖2 = (1± )‖Ax‖2 Lemma 3.4 (Approximate Matrix Product for Orthogonal Subspaces). Given δ ∈ (0, 12 ), let U ∈ Rn×k,W ∈ Rn×p be two matrices with orthonormal columns such that UTW = 0, p ≥ max(k, log(1/δ)), rank(U) = k and rank(W ) = p. Let S ∈ Rr×n be a random matrix with i.i.d. Gaussian random variables N (0, 1r ). For r = O(k + log( 1 δ )), the following holds with probability 1− δ, ‖UTSTSW ‖F ≤ O(1)‖W ‖F .
Note that we will apply the above two lemmas with constant . The proof intuition is as follows: consider a sketch matrix S of size r with i.i.d. N (0, 1r ) random variables as in Lemma 3.4. The range of U ∈ Rn×k corresponds to an orthonormal basis of a rank-k low rank approximation to A, and the range of W ∈ Rn×p is the orthogonal complement. Note that both SU and SW are random matrices consisting of i.i.d. N (0, 1r ) random variables and thus the task is to bound the size, in Frobenius norm, of the product of two random Gaussian matrices with high probability. Intuitively, the size of the matrix product is proportional to the rank k and inversely proportional to our sketch size r. The overall failure probability δ, however, is inversely proportional to k, since as k grows, the matrix product involves summing over more squared Gaussian random variables, i.e., χ2 random variables, and thus becomes even more concentrated. We show that for k ≥ log(1/δ), a sketch size of O(k) suffices since the failure probability for each χ2 random variable is small enough to pay a union bound over k terms. On the other hand, when k < log(1/δ), we show that r = O(log(1/δ)) suffices for the union bound. Combining the two cases gives r = O(k + log(1/δ)).
Having shown the above, we next show that the low rank approximation error, i.e., ‖A− Ã‖F , is upper bounded by: 1) the inflation in eigenvalues by applying a sketch matrix S as in Lemma 3.3; and 2) the approximate product of the range of a low rank approximation to A and its orthogonal complement, as in Lemma 3.4. Together these show that m = O(k + log(1/δ)) suffices for Ã to be an O(1)-approximate rank-k approximation to A with probability 1 − δ, as stated in Theorem 3.5. Note that in both Lemma 3.3 and Lemma 3.4, the entries of the random matrix are scaled Gaussian random variables N (0, 1r ). However, when one sets the low rank approximation as Ã = AR(STAR)†(AS)T , the scale cancels and one can choose standard Gaussians in the sketching matrix for convenience as in Theorem 3.5. Theorem 3.5. Let A ∈ Rn×n be an arbitrary PSD matrix. Let Ak = arg minrank-kAk ‖A−Ak‖F be the optimal rank-k approximation to A in Frobenius norm. If S ∈ Rn×m and R ∈ Rn×cm are random matrices with i.i.d. N (0, 1) entries for some fixed constant c > 0 withm = O(k+log(1/δ)), then with probability 1 − δ, the matrix Ã = (AR)(STAR)†(AS)T satisfies ‖A − Ã‖F ≤ O(1)‖A−Ak‖F .
This improved result enables us to choose k = l = O( √
log(1/δ)/ ) in Theorem 3.2, and combined with Theorem 3.5, this shows that only O( √ log(1/δ)/ + log(1/δ)) matrix-vector queries are needed to output a number in (1± )tr(A) with probability 1− δ, as we conclude in Theorem 3.1.

heading: 4 Lower Bounds
text: In this section, we show that our upper bound on the query complexity of non-adaptive trace estimation is tight, up to a factor of O(log log(1/δ)).
Theorem 4.1 (Lower Bound for Non-Adaptive Queries). Let ∈ (0, 1). Any algorithm that accesses a real PSD matrix A through matrix-vector multiplication queries Aq1,Aq2, . . . ,Aqm, where q1, . . . ,qm are real-valued, non-adaptively chosen vectors, requires
m = Ω
(√ log(1/δ)
+ log(1/δ)
log log(1/δ)
) queries to output an estimate t such that with probability at least
1− δ, (1− )tr(A) ≤ t ≤ (1 + )tr(A).
Our lower bound hinges on two separate cases: we first show an Ω( √
log(1/δ)/ ) lower bound in Section 4.1 whenever = O(1/ √ log(1/δ)). Second, we show an Ω( log(1/δ)log log(1/δ) ) lower bound in
Section 4.2 that applies to any ∈ (0, 1). Observe that for < 1/ √
log(1/δ), the first lower bound holds; for ≥ 1/ √ log(1/δ), our second lower bound dominates. Therefore, combining both lower
bounds implies that for every and δ, the query complexity of O( √
log(1/δ)/ + log(1/δ)) for non-adaptive trace estimation is tight, up to a log log(1/δ) factor.
We now give a proof sketch of the two lower bounds. All details are in the supplementary material. Our lower bounds crucially make use of rotational invariance of the Gaussian distribution (see Fact 2.1) to argue that the first q queries are, w.l.o.g., the standard basis vectors e1, ..., eq. Note that our queries can be assumed to be orthonormal. Both lower bounds use the family of n× n Wigner matrices (see Definition 2.1) with shifted mean, i.e., W + C · I for some C > 0 depending on ‖W ‖op, as part of the hard input distribution. The mean shift ensures that our ultimate instance is PSD with high probability.

heading: 4.1 Case 1: Lower Bound for Small
text: The first lower bound is based on the observation that due to rotational invariance, the not-yet-queried part of W is distributed almost identically to W , up to some mean shift, conditioned on the queried known part, no matter how the queries are chosen. The sum of diagonal entries of the not-yet-queried part is Gaussian, and this still has too much deviation to determine the overall trace of the input up to a (1± ) factor when n = √ log(1/δ)/ and < 1/ √ log(1/δ).
Theorem 4.2 (Lower Bound for Small ). For any PSD matrix A and all = O(1/ √
log(1/δ)), any algorithm that succeeds with probability at least 1 − δ in outputting an estimate t such that (1− )tr(A) ≤ t ≤ (1 + )tr(A), requires m = Ω( √ log(1/δ)/ ) matrix-vector queries.

heading: 4.2 Case 2: Lower Bound for Every
text: The second lower bound presented in Theorem 4.3 is shown via reduction to a distribution testing problem between two distributions presented in Problem 4.4. Theorem 4.3 (Lower Bound on Non-adaptive Queries for PSD Matrices). Let ∈ (0, 1). Any algorithm that accesses a real, PSD matrix A through matrix-vector queries Aq1,Aq2, . . . ,Aqm, where q1, . . . ,qm are real-valued non-adaptively chosen vectors, requires m = Ω( log(1/δ) log log(1/δ) ) to output an estimate t such that with probability at least 1− δ, (1− )tr(A) ≤ t ≤ (1 + )tr(A). In the distribution testing problem, we consider Wigner matrices W ∼ W(log(1/δ)) shifted by Θ( √
log(1/δ))I . The problem requires an algorithm for distinguishing between a sample Q from this Wigner distribution and a sample P from this distribution shifted by a random rank-1 PSD matrix. The rank-1 matrix is the outer product of a random vector with itself and is chosen to provide a constant factor gap between the trace of P and Q. Problem 4.4 (Hard PSD Matrix Distribution Test). Given δ ∈ (0, 12 ), set n = log(1/δ). Choose g ∈ Rn to be an independent random vector with i.i.d. N (0, 1) entries. Consider two distributions:
• Distribution P on matrices { C log3/2( 1δ ) · 1 ‖g‖22 ggT + W + 2 √ log( 1δ )I }
, for some fixed constant C > 1.
• Distribution Q on matrices { W + 2 √ log( 1δ )I } .
where W ∼ W(n) as in Definition 2.1. Let A be a random matrix drawn from either P or Q with equal probability. Consider any algorithm which, for a fixed query matrix Q ∈ Rn×q , observes AQ, and guesses if A ∼ P or A ∼ Q with success probability at least 1− δ.
We then show in Lemma 4.5 that any algorithm which succeeds with probability 1− δ in distinguishing P from Q requires Ω( log(1/δ)log log(1/δ) ) non-adaptive matrix-vector queries.
Due to rotational invariance and since queries are non-adaptive, the first q queries are the first q standard unit vectors. By Fact 2.2, with probability at least 1log(1/δ) , however, a single coordinate of g
has absolute value at most 1log(1/δ) . By independence, with probability at least ( 1 log(1/δ) ) q , all of the first q coordinates of g are simultaneously small, and thus give the algorithm almost no information to distinguish P from Q; this probability is δ if q = O( log(1/δ)log log(1/δ) ). Lemma 4.5 (Hardness of Problem 4.4). For a non-adaptive query matrix Q ∈ Rn×q as in Problem 4.4, given δ ∈ (0, 12 ), for n = log(1/δ), if q = o( log(1/δ) log log(1/δ) ), no algorithm can solve Problem 4.4 with success probability 1− δ.

heading: 5 Experiments
text: 1Part I: Comparison of Failure Probability and Running Time We give sequential and parallel implementations of the non-adaptive trace estimation algorithm NA-Hutch++ (Algorithm 2), the adaptive algorithm Hutch++ (Algorithm 1) and Hutchinson’s method [14]. We specifically explore the benefits of the non-adaptive algorithm in a parallel setting, where all algorithms have parallel access to a matrix-vector oracle. All the code is included in the supplementary material and will be publicly released.
Metrics. We say an estimate failed if on input matrix A, the estimate t returned by an algorithm falls into either case: t < (1− )tr(A) or t > (1 + )tr(A). We measure the performance of each algorithm by: 1) the number of failed estimates across 100 random trials, 2) the total wall-clock time to perform 100 trials with sequential execution, and 3) the total wall-clock time to perform 100 trials with parallel execution.
Datasets and Applications. We consider different applications of trace estimation from synthetic to real-world datasets. In many applications, trace estimation is used to estimate not only tr(A), but also tr(f(A)) for some function f : R→ R. Letting A = V ΣV T be the eigendecomposition of A, we have f(A) := V f(Σ)V T , where f(Σ) denotes applying f to each of the eigenvalues. Due to the expensive computation of eigendecompositions of large matrices, the matrix-vector multiplication f(A)v is often estimated by polynomials implicitly computed via an oracle algorithm for a random vector v. The Lanczos algorithm is a very popular choice due to its superior performance (e.g. [31, 2, 7]). We compare the performance of our trace estimation algorithms on the following applications and datasets, and use the Lanczos algorithm as the matrix-vector oracle on a random vector v in some particular cases.
• Fast Decay Spectrum. We first consider a synthetic dataset of size 5000 with a fast decaying spectrum, following [16], which is a diagonal matrix A with i-th diagonal entry Aii = 1/i
2. Matrices with fast decaying spectrum will cause high variance in the estimated trace of Huthinson, but low variance for Hutch++ and NA-Hutch++. The matrix-vector oracle is simply Av.
• Graph Estrada Index. Given a binary adjacency matrix A ∈ {0, 1}n×n of a graph, the Graph Estrada Index is defined as tr(exp(A)), which measures the strength of connectivity within the graph. Following [16], we use roget’s Thesaurus semantic graph2 with 1022 nodes, which was originally studied in [5], and use the Lanczos algorithm with 40 steps to approximate exp(A)v as the matrix-vector oracle.
• Graph Triangle Counting. Given a binary adjacency matrix A ∈ {0, 1}n×n of a graph, the number of triangles in the graph is 1/6 · tr(A3). This is an important graph summary with numerous applications in graph-mining and social network analysis (e.g. [32, 33]). We use arxiv_cm, the Condense Matter collaboration network dataset from arXiv 3. This is a common benchmark graph with 23, 133 nodes and 173, 361 triangles. The matrix-vector oracle is A3v. Note that A3 in this case is not necessarily a PSD matrix.
• Log-likelihood Estimation for Gaussian Process. When performing maximum likelihood estimation (MLE) to optimize the hyperparameters of a kernel matrix A for Gaussian Processes, one needs to compute the gradient of the log-determininant of A, which involves estimating tr(A−1) [2]. Following [2], we use the precipitation4 dataset, which consists
1Our code is available at: https://github.com/11hifish/OptSketchTraceEst 2http://vlado.fmf.uni-lj.si/pub/networks/data/ 3https://snap.stanford.edu/data/ca-CondMat.html 4https://catalog.data.gov/dataset/u-s-hourly-precipitation-data
of the measured amount of precipitation during a day collected from 5,500 weather stations in the US in 2010. We sample 1,000 data points, and construct a covariance matrix A using the RBF kernel with length scale 1. We use the Lanczos algorithm with 40 steps as in [2] to approximate A−1v as the matrix-vector oracle.
Implementation. We use random vectors with i.i.d. N (0, 1) entries as the query vectors for all algorithms. NA-Hutch++ requires additional hyperparameters to specify how the queries are split between random matrices S,R,G (see Algorithm 2). We set c1 = c3 = 14 and c2 = 1 2 as [16] suggests. For each setting, we conduct 10 random runs and report the mean number of failed estimates across 100 trials and the mean total wall-clock time (in seconds) conducting 100 trials with one standard deviation. For all of our experiments, we fix the error parameter = 0.01 and measure the performance of each algorithm with {10, 30, 50, . . . , 130, 150} queries on synthetic, roget and precipitation, and with {100, 200, . . . , 700, 800} queries on arxiv_cm which has a significantly larger size. The parallel versions are implemented using Python multiprocessing5 package. Due to the large size of arxiv_cm, we use sparse_dot_mkl6, a Python wrapper for Intel Math Kernel Library (MKL) which supports fast sparse matrix-vector multiplications, to implement the matrix-vector oracle for this dataset. During the experiments, we launch a pool of 40 worker processes in our parallel execution. All experiments are conducted on machines with 40 CPU cores.
Results and Discussion. The results of Hutch++, NA-Hutch++ and Hutchinson over the 4 datasets are presented in Figure 1. The performance of all algorithms is consistent across different datasets with different matrix-vector oracles, and even on a non-PSD instance from arxiv_cm. Given the same number of queries, Hutch++ and NA-Hutch++ both give significantly fewer failed estimates than Hutchinson, particularly on PSD instances. It is not surprising to see that Hutchinson fails to
5https://docs.python.org/3/library/multiprocessing.html 6https://github.com/flatironinstitute/sparse_dot
achieve a (1± )-approximation to the trace most of the time due to the high variance in its estimation, given a small number of queries and a high accuracy requirement ( = 0.01).
For computational costs, the difference in running time of all algorithms is insignificant in our sequential execution. In our parallel execution, however, Hutch++ becomes significantly slower than the other two, NA-Hutch++ and Hutchinson, which have very little difference in their parallel running time. Hutch++ suffers from slow running time due to its adaptively chosen queries, despite the fact that Hutch++ consistently gives the least number of failed estimates.
It is not hard to see that NA-Hutch++ gives the best trade-off between a high success probability in estimating an accurate trace with only a few number of queries, and a fast parallel running time due to the use of non-adaptive queries, which makes NA-Hutch++ more practical on large, real-world datasets. We remark that although the Lanczos algorithm is adaptive itself, even with a sequential matrix-vector oracle, our non-adaptive trace estimation can still exploit much more parallelism than adaptive methods, as shown by our experiments.
Part II: Comparison of Performance on Log Determinant Estimation We give an additional experiment to compare the performance of Hutch++, NA-Hutch++ and Hutchinson on estimating log(det(K)) = tr(log(K)), for some covariance matrix K. Estimating log(det(K)) is required when computing the marginal log-likelihood in large-scale Gaussian Process models. Recently, [3] proposed a maximum entropy estimation based method for log determinant estimation, which uses Hutchinson’s trace estimation as a subroutine to estimate up to the k-th moments of the eigenvalues, given a fixed k. The i-th moment of the eigenvalues is E[λi] = 1n tr(K
i), where K is an n × n PSD matrix, and λ is the vector of eigenvalues. [3] shows that their proposed approach outperforms traditional Chebyshev/Lanczos polynomials for computing log(det(K)) in terms of absolute value of the relative error, i.e., abs (estimated log determinant - true log determinant)/abs(true log determinant).
We compare the estimated log determinant of a covariance matrix with different trace estimation subroutines for estimating the moments of the eigenvalues. We use 2 PSD matrices from the UFL Sparse Matrix Collection7: bcsstk20 (size 485 × 485) and bcsstm08 (size 1074 × 1074), with varying max moments {10, 15, . . . , 30} and 30 matrix-vector queries. We repeated each run 100 times and reported the mean estimated log determinant with each trace estimation subroutine. While an improved estimate of the eigenvalue moments does not necessarily lead to an improved estimate of the log determinant, it is not hard to show that an accurate moment estimation does lead to improved log determinant estimation in extreme cases where the eigenspectrum of K contains a few very large eigenvalues. Such a case will cause Hutchinson’s method to have very large variance, while our method reduces the variance by first removing the large eigenvalues. The eigenspectrums of both input matrices and the results are presented in Figure 2.

heading: 6 Conclusion
text: We determine an optimal Θ( √
log(1/δ)/ + log(1/δ)) bound on the number of queries to achieve (1± ) approximation of the trace with probability 1− δ for non-adaptive trace estimation algorithms, up to a log log(1/δ) factor. This involves both designing a new algorithm, as well as proving a new lower bound. We conduct experiments on synthetic and real-world datasets and confirm that our non-adaptive algorithm has a higher success probability compared to Hutchinson’s method for the same sketch size, and has a significantly faster parallel running time compared to adaptive algorithms.
7https://sparse.tamu.edu/
Acknowledgments and Disclosure of Funding
We would like to thank the anonymous reviewers for their feedback. We are also grateful to Raphael Meyer for many detailed comments on the lower bound proofs. D. Woodruff was supported by NSF CCF-1815840, Office of Naval Research grant N00014-18-1-2562, and a Simons Investigator Award.

Review summaries:
This paper contains two main contributions. First, they provided an improved analysis on the NA-Hutch++ algorithm to bridge the gap between adaptive and non-adaptive algorithm. Secondly, they also provided a nearly matching lower-bound on the complexity of the problem, closing the gap.

This paper gives improved bounds for estimating the trace of a matrix using matrix-vector products. It turns what was previous a multiplicative factor of log(1 / failure probability) into an additive one, and shows that this bound is tight up to a loglog factor. These methods are evaluated on a wide range of matrices, several of which are implicitly represented (via the Lanczos method), and a better success probability was observed.

This paper presents an improved analysis for Hutch++, a stochastic trace estimation method proposed recently. In particular, the paper presents a new analysis for the non-adaptive variant of the Hutch++ algorithm and show that even the non-adaptive algorithm achieves optimal dependency on the failure probability \delta (i.e., O(\sqrt{\log(1/\delta)}))for the number of samples/matvecs. The paper also presents an improved lower bound for the number of matrix-vector queries required in terms of \delta and error tolerance \epsilon. Some numerical results are also presented to illustrate the performance of the methods.

The paper studies the problem of estimating the trace of an input PSD matrix with the perspective of improving the query complexity of matrix-vector products required to compute an eps-multiplicative approximation of the trace. There have been two types of randomized algorithms solving this problem: adaptive and non-adaptive.  The former uses samples that depend on previous ones and therefore, by definition, is sequential. The latter doesn't, and is therefore amenable to parallelization, clearly a useful/desirable feature to have.   However, the prior lowest query complexity results for adaptive algorithms were lower than those for non-adaptive algorithms, by a factor of sqrt(log(1/delta)), where delta is the error probability. This paper focuses on non-adaptive algorithms for trace estimation of PSD matrices. Within the framework of this problem, it has two key contributions.   1. The first is an improved analysis of a non-adaptive algorithm proposed in a previous paper. The analysis is an improvement in the sense that it is able to show that this non-adaptive algorithm in fact needs just as many queries as the best adaptive algorithm, thus shaving off a factor of sqrt(log(1/delta)) from the prior analysis.   2. The second contribution is a proof that the new result obtained is in fact optimal, thereby completely resolving the question of non-adaptive query complexity of matrix-vector products for trace estimation of PSD matrices.   In the process of 1., the paper also provides the following technical contribution: a proof that to estimate the k-rank approximation of a PSD matrix with probability 1-delta, it suffices to use O(k + log(1/delta)) queries, down from the previous best query complexity of O(k log (1/delta)). The paper explains how, to achieve k-rank approximation, one needs matrices that satisfy two properties: (1) a subspace embedding property, and (2) a property about orthogonal matrix products. While property (1) was satisfiable by the improved query complexity before, it wasn't known that (2) was. This paper shows, by carefully arguing about tail bounds of chi-squared distributed variables, that (2) is also in fact true for O(k + log(1/delta)) queries.