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