Dynamic Trace Estimation

1 Introduction . Implicit or “ matrix-free ” trace estimation is a ubiquitous computational primitive in linear algebra , which has become increasingly important in machine learning and data science . Given access to an oracle for computing matrix-vector products Ax1 , . . . , Axm between an n× n matrix A and chosen vectors x1 , . . . , xm , the goal is to compute an approximation to A ’ s trace , tr ( A ) = ∑n i=1Aii . This problem arises when A ’ s diagonal entries can not be accessed explicitly , usually because forming A is computationally prohibitive . As an example , consider A which is the Hessian matrix of a loss function involving a neural network . While forming the Hessian is infeasible when the network is large , backpropagation can be used to efficiently compute Hessian-vector products [ 32 ] . In other applications , A is a matrix function of another matrix B . For example , if B is a graph adjacency matrix , tr ( B3 ) equals six times the number of triangle in the graph [ 2 ] . ComputingA = B3 explicitly to evaluate the trace would require O ( n3 ) time , while the matrix-vector multiplication Ax = B · ( B · ( Bx ) ) only requires O ( n2 ) time . Similarly , in log-determinant approximation , useful in e.g . Bayesian log likelihood computation or determinantal point process ( DPP ) methods , we want to approximate the trace of A = log ( B ) [ 5 , 17 , 36 ] . Again , A takes O ( n3 ) time to form explicitly , but Ax = log ( B ) x can be computed in roughly O ( n2 ) time using iterative methods like the Lanczos algorithm [ 21 ] . Dynamic versions of the log-determinant estimation problem have been studied due to applications in greedy methods for DPP inference [ 15 ] . In data science and machine learning , other applications of implicit trace estimation include matrix norm and spectral sum estimation [ 16 , 41 , 31 ] , as well as methods for eigenvalue counting [ 9 ] and spectral density estimation [ 45 , 27 ] . Spectral density estimation methods typically use implicit trace estimation to estimate moments of a matrix ’ s eigenvalue distribution – i.e. , tr ( A ) , tr ( A2 ) , tr ( A3 ) , etc . – which can then be used to compute an approximation to that entire distribution . In deep learning , spectral density estimation is used to quickly analyze the spectra of weight matrices [ 33 , 28 ] or to probe information about the Hessian matrix during optimization [ 13 , 49 ] . Trace estimation has also been used for neural networks weight quantization [ 10 , 34 ] and to understand training dynamics [ 44 ] . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . 1.1 Static Trace Estimation . The mostly widely used implicit trace estimation algorithm is Hutchinson ’ s estimator [ 14 , 22 ] . Letting g1 , . . . , g ` ∈ Rn be random vectors with i.i.d . mean 0 and variance 1 entries ( e.g. , standard Gaussian or ±1 Rademachers ) , Hutchinson ’ s approximates tr ( A ) via the average h ` ( A ) = 1 ` ∑1 i=1 g T i ( Agi ) . This estimator requires ` matrix-vector multiplications to compute . Its variance can be shown to be O ( ‖A‖2F / ` ) and with high probability , when ` = O ( 1/ 2 ) , we have the error guarantee [ 3 , 30 ] : |h ` ( A ) − tr ( A ) | < ‖A‖F . ( 1 ) While improvements on Hutchinson ’ s estimator have been studied for restricted classes of matrices ( positive semidefinite , sparse , nearly low-rank , etc . ) [ 30 , 40 , 38 , 36 , 19 ] , the method is the best known for general matrices – no techniques achieve guarantee ( 1 ) with o ( 1/ 2 ) matrix-vector products . 1.2 Dynamic Trace Estimation . We explore a natural and widely applicable dynamic version of the implicit trace estimation problem : given access to a matrix-vector multiplication oracle for a dynamically changing matrix A , maintain an approximation to A ’ s trace . This problem arises in applications involving optimization in machine learning where we need to estimate the trace of a constantly changing Hessian matrix H ( or some function of H ) during model training . In other applications , A is dynamic because it is repeatedly modified by some algorithmic process . E.g. , in the transit planning method of [ 43 ] , edges are added to a network to optimally increase the “ Estrada index ” [ 11 ] . Evaluating this connectivity measure requires computing tr ( exp ( B ) ) , where B is the dynamically changing network adjacency matrix and exp ( B ) is a matrix exponential . A naive solution to the dynamic problem is to simply apply Hutchinson ’ s estimator to every snapshot of A as it changes over time . To achieve a guarantee like ( 1 ) for m time steps , we require O ( m/ 2 ) matrix-vector multiplies . The goal of this paper is to improve on this bound when the changes to A are bounded . Formally , we abstract the problem as follows : Problem 1 ( Dynamic trace estimation ) . Let A1 , ... , Am be n× n matrices satisfying : 1 . ‖Ai‖F ≤ 1 , for all i ∈ [ 1 , m ] . 2 . ‖Ai+1 −Ai‖F ≤ α , for all i ∈ [ 1 , m− 1 ] . Given implicit matrix-vector multiplication access to eachAi in sequence , the goal is to compute trace approximations t1 , . . . , tm for tr ( A1 ) , .... , tr ( Am ) such that , for each i ∈ 1 , . . . , m , P [ |ti − tr ( Ai ) | ≥ ] ≤ δ . ( 2 ) Above A1 , . . . , Am represent different snapshots of a dynamic matrix at m time steps . We require ‖Ai‖F ≤ 1 only to simplify the form of our error bounds – no explicit rescaling is necessary for matrices with larger norm . If we assume ‖Ai‖F ≤ U for some ( unknown ) upper bound U , the guarantee of ( 2 ) would simply change to involve a U terms instead of . The second condition bounds how much the matrices change over time . Again for simplicity , we assume a fixed upper bound α on the difference at each time step , but the algorithms presented in this paper will be adaptive to changing gaps between Ai and Ai+1 , and will perform better when these gaps are small on average . By triangle inequality , α ≤ 2 , but in applications we typically have α 1 , meaning that the changes in the dynamic matrix are small relative to its Frobenius norm . If this is not the case , there is no hope to improve on the naive method of applying Hutchinson ’ s estimator repeatedly to each Ai . Note on Matrix Functions . In many applications Ai = f ( Bi ) for a dynamically changing matrix B . While we may have ‖Ai+1 − Ai‖F = ‖f ( Bi+1 ) − f ( Bi ) ‖F ‖Bi+1 − Bi‖F for functions like the matrix exponential , this is not an immediate issue . To improve on Hutchinson ’ s estimator , the important requirement is simply that ‖Ai+1 − Ai‖F is small in comparison to ‖Ai+1‖F . As discussed in Section 5 , this is typically the case for application involving matrix functions . We will measure the complexity of any algorithm for solving Problem 1 in the matrix-vector multiplication oracle model of computation , meaning that we consider the cost of matrix-vector products ( which are the only way A1 , . . . , Am can be accessed ) to be significantly larger than other computational costs . We thus seek solely to minimize the number of such products used [ 39 ] . The matrix-vector oracle model has seen growing interest in recent years as it generalizes both the matrix sketching and Krylov subspace models in linear algebra , naturally captures the true computational cost of algorithms in these classes , and is amenable to proving strong lower-bounds [ 37 , 6 ] . 1.3 Main Result . Our main result is an algorithm for solving Problem 1 more efficiently than Hutchinson ’ s estimator : Theorem 1.1 . For any , δ , α ∈ ( 0 , 1 ) , the DeltaShift algorithm ( Algorithm 1 ) solves Problem 1 with O ( m · α log ( 1/δ ) 2 + log ( 1/δ ) 2 ) total matrix-vector multiplications involving A1 , . . . , Am . For large m , the first term dominates the complexity in Theorem 1.1 . For comparison , a tight analysis of Hutchinson ’ s estimator [ 29 ] establishes that the naive approach requires O ( m · log ( 1/δ ) / 2 ) , which is worse than Theorem 1.1 by a factor of α . A natural setting is when α = O ( ) , in which case Algorithm 1 requires O ( log ( 1/δ ) / ) matrix-multiplications on average over m time steps , in comparison to O ( log ( 1/δ ) / 2 ) for Hutchinson ’ s estimator , a quadratic improvement in . To prove Theorem 1.1 , we introduce a dynamic variance reduction scheme . By linearity of trace , tr ( Ai+1 ) = tr ( Ai ) + tr ( ∆i ) , where ∆i = Ai+1 −Ai . Instead of directly estimating tr ( Ai+1 ) , we combine previous estimate for tr ( Ai ) with an estimate for tr ( ∆i ) , computed via Hutchinson ’ s estimator . Each sample for Hutchinson ’ s estimator applied to ∆i requires just two matrix-vector multiplies : one with Ai and one with Ai+1 . At the same time , when ∆i has small Frobenius norm ( bounded by α ) , we can estimate its trace more accurately than tr ( Ai+1 ) 1 . While intuitive , this approach requires care to make work . In a naive implementation , error in estimating tr ( ∆1 ) , tr ( ∆2 ) , . . . , compounds over time , eliminating any computational savings . To avoid this issue , we introduce a novel damping strategy that actually estimates tr ( Ai+1 − ( 1− γ ) Ai ) ) for a positive damping factor γ . We compliment our main result with a nearly matching conditional lower bound : in Section 4 we argue that our DeltaShift method can not be improved in the dynamic setting unless Hutchinson ’ s estimator can be improved in the static setting . We also present an improvement to DeltaShift under more stringent , but commonly present , bounds on A1 , . . . , Am and each Ai+1 −Ai than Problem 1 . 1.4 Related Work . Prior work on implicit trace estimation and applications in machine learning is discussed in the beginning of this section . While there are no other methods that improve on Hutchinson ’ s estimator in the dynamic setting , the idea of variance reduction has found applications in other work on implicit trace estimation [ 1 , 12 , 26 , 30 ] . In these results , the trace of a matrix A is estimated by decomposing A = B+ ∆ where B has an easily computed trace ( e.g. , because it is low-rank ) and ‖∆‖F ‖A‖F , so tr ( ∆ ) is more easily approximated with Hutchinson ’ s estimator than tr ( A ) directly .

Summary: The authors proposed a method for estimating the trace for a sequence of matrices that only changes for a small amount at each step. The method originates from the well-known stochastic trace estimator, but bring it to a new practical setting (e.g. Hessian at learning iterations, evolving network, etc). By an interesting treatment of the iterative process (i.e. damping), the resulting algorithm could cut the computational cost by a factor of the size of the tolerable error. In addition, the authors included the recent development/improvement on the original Hutchinson estimator, to get better versions of their algorithm under more assumptions.