Optimal Sketching for Trace Estimation

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 . 1 Introduction . 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 can not 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 . 1.1 Our Contributions . 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/δ ) ) ) .

Summary: 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.