A Modern Analysis of Hutchinson’s Trace Estimator.

Abstract

The paper establishes the new state-of-art in the accuracy analysis of Hutchinson’s trace estimator. Leveraging tools that have not been previously used in this context, particularly hypercontractive inequalities and concentration properties of sub-gamma distributions, we offer an elegant and modular analysis, as well as numerically superior bounds. Besides these improvements, this work aims to better popularize the aforementioned techniques within the CS community.

I Introduction

I-A Background

Estimating the matrix trace is a problem of fundamental interest [1, 2, 3] and arises in many problems such as approximating spectral properties of matrices [4, 5, 6], solving partial differential equations [7, 8, 9, 10, 11], error evaluation in machine-learning [12], and combinatorial counting [13] to mention a few. For readers particularly interested in data science or optimization, it is of critical interest as the hessian trace stores valuable information about curvature. To give a meaningful example, consider fitting a linear classifier on the MNIST dataset [14] of hand-written digits, widely used as a benchmark and a toy problem in machine learning. Since images are in resolution $28\times 28$ and grouped in $10$ classes, the number of parameters is $m=28\cdot 28\cdot 10=7840$, and the size of the hessian matrix is $m^2\approx 6\cdot {10}^6$. The diagonal average, proportional to the trace, equals the average eigenvalue (valuable information), whereas individual rows or columns are nearly zero (up to random noise). This is illustrated in Footnote 2.

Figure 1: Hessian matrix of the linear classifier trained on MNIST data. The diagonal stores valuable information as opposed to rows and columns: entries in rows and colums are nearly zero, whereas the diagonal average is $\approx 0.25$. The Tensorflow code is available as a Python Notebook at GitHub.

As seen from the above example, already toy examples lead to huge matrices; for larger problems storing and inspecting such matrices is impossible. Fortunately, it is possible to estimate the trace (equal to the diagonal sum, or equivalently to the sum of eingenvalues) without knowing the full matrix. The requirement is that one can compute efficiently matrix-vector products; in the case of the hessian such products can be computed by by auto-differentiation [15] in all popular machine-learning frameworks such as Tensorflow [16], PyTorch [17] and JAX [18]. Under the assumption that an $m\times m$ matrix $A$ is symmetric positive semi-definite (which applies to all hessians), the popular estimator due to Hutchinson [1] is

${\mathrm{tr}}_H(A)\triangleq z^{T}Az,z{\sim }^{iid}{\{-1,1\}}^m,$ (1)

where the components of the vector $z$ are Rademacher random variables, equal to $+1$ or $-1$ with equal probability. The above estimator uses only one sampled $v$ and thus is noisy (of high-variance), so it is usually boosted by averaging $n$ independent trials (for suitably large $n$). Formally:

${\mathrm{tr}}_{H^{(n)}}(A)\triangleq {\displaystyle \frac{1}{n}}{\displaystyle \sum _{i=1}^n}{z}_i^{T}A{z}_i,z_i{\sim }^{iid}{\{-1,1\}}^m.$ (2)

The estimators are unbiased, that is correct on average:

$𝔼{\mathrm{tr}}_H(A)=𝔼{\mathrm{tr}}_{H^{(n)}}(A)=\mathrm{tr}(A),$ (3)

and this is fairly easy to prove. The focus of this work is on a more challenging problem of understanding their concentration properties. Here we ask the following question:

What sample size $n$ guarantees the desired relative error $ϵ$ at the confidence level of $1-\delta$?

Formally, for the error ${\mathrm{err}}_{H^{(n)}}(A)\triangleq \frac{{\mathrm{tr}}_{H^{(n)}}(A)}{\mathrm{tr}(A)}-1$, and fixed $ϵ,\delta$, we want possibly small $n=n(ϵ,\delta )$ such that

$|{\mathrm{err}}_{H^{(n)}}(A)|⩽ϵ\text{with prob. at least}1-\delta .$ (4)

I-B Related Work

The estimator is already more than 30 years old [1], and although alternatives exist (such as methods based on Lanczos quadrature [19]), it is provably best in terms of variance [20] and arguably wins in simplicity, being preferred by developers of statistical/learning software [21]. Quite surprisingly, it had been used for a while without a rigorous assessment of accuracy, until the work of Avron and Toledo [3], who established the finite sample size guarantees. Their result was later improved by Roosta-Khorasani and Ascher [20], who essentially got rid of the lossy proof step involving a crude union bound. Their approach is based on the Chernoff-like direct calculations, and offers the bound $n=O({ϵ}^{-2}\log (1/\delta ))$.

II Contribution

II-A Summary

In this work we offer a modular, modern and more accurate analysis of Hutchinson’s estimator. The improvements upon prior works can be summarized as follows:

• Modularity and Novelty of Techniques. In the first step, we explicitly link the estimator accuracy to the dispersion of the Rademacher Quadratic Chaos with a unitary matrix kernel. We then obtain a bound on such chaoses by means of Hypercontractive Inequality, an important result in Boolean Fourier Analysis [22], originally due to Aline Bonami [23]. As the final step we express this bound as the sub-gamma property (see the works of Boucheron [24]), which makes it convenient (and accurate) to conclude the final concentration result. Moreover, we state our results for the multiple-sample and one-shot estimators; this is of interest, since the base building block may be boosted differently than by averaging (see for example the median algorithm [25]).

  Thus, in contrast to ad hoc calculations in prior works, we are able to use established tools from the field of high-dimensional probability and Fourier analysis. Our transparent approach opens a way for further refinements.

• Superior Accuracy. With the dedicated tools we obtain bounds that are numerically better up to an order of magnitude. They are stated as elementary formulas, convenient to use in practical applications of trace estimation.

II-B Preliminaries

Below we explain the notation and definitions used when formulating our results presented in the next section.

The $d$-th norm of a r.v. $X$ is ${\Vert X\Vert }_d\triangleq {(𝔼{|X|}^d)}^{1/d}$; it is a valid norm (e.g. it obeys the triangle inequality) when $d⩾1$ [25]. A useful tool for studying concentration is the moment generating function, defined as ${\mathrm{MGF}}_X(t)=𝔼\exp (tX)$ (a function of real parameter $t$). In modern high-dimensional probability one classifies distributions according to the behaviour of $\mathrm{MGF}$; in particular we call $X$ sub-gamma with variance factor $v$ and scale $c$ when $\log {\mathrm{MGF}}_{|X|}(t)⩽\frac{v{t}^2}{2(1-ct)}$, and denote by $X\in \mathrm{\Gamma }(v,c)$ [24]; the same formula holds for the centered gamma distribution, hence the name.

II-C Results

In what follows we assume that $A$ is non-zero symmetric positive semi-definite matrix. Then, necessarily, $\mathrm{tr}(A)>0$.

II-C1 One-Shot Estimator

Below we present the following bound on the accuracy of the estimator in (1):

Theorem 1 (Hutchison’s Estimator 1-Sample Bound).

For any integer $d⩾2$, the relative error of Hutchison’s Estimator (1) obeys the following bound

${\Vert {\mathrm{err}}_H(A)\Vert }_d⩽d-1.$ (5)

In particular, this implies the sub-gamma behaviour

${\mathrm{err}}_H(A)\in \mathrm{\Gamma }(1,{\displaystyle \frac{8}{3}}),$ (6)

which in turns gives the following probability tail bound

$\mathrm{Pr}[|{\mathrm{err}}_H(A)|⩾ϵ]⩽{\mathrm{e}}^{-\frac{{ϵ}^2}{2\left(1-\frac{8}{3}ϵ\right)}}.$ (7)

The proof goes along the following lines: a) we express the error in form of Rademacher chaos b) we use hypercontractive inequalities to bound its moment; this establishes the part (5) c) we plug the moment bound into Taylor’s expansion of the moment generating function and then estimate the series, arriving at the sub-gamma condition (6). By the tail properties of sub-gamma distributions, we conclude the tail bound (7).

II-C2 Multiple-Sample Estimator

Next, we move to the multiple-sample estimator in (2). We establish the following

Theorem 2 (Hutchinson’s Estimator $n$-Sample Bound).

The relative error of $n$-sample Hutchinson’s Estimator has the following sub-gamma behaviour

${\mathrm{err}}_H(A)\in \mathrm{\Gamma }({\displaystyle \frac{1}{n}},{\displaystyle \frac{8}{3}}).$ (8)

In particular, this implies the following bound on the error tail probability, valid for any :

$\mathrm{Pr}[|{\mathrm{err}}_H^{(n)}(A)|⩾ϵ]⩽{\mathrm{e}}^{-\frac{n{ϵ}^2}{2\left(1-\frac{8}{3}ϵ\right)}}.$ (9)

This result is obtained from Theorem 1, by using extra facts on the concentration of sums of sub-gamma random variables: essentially the variance factor decreases form $1$ to $\frac{1}{n}$ for sums of length $n$ (just as the variance of $n$ iid terms).

Corollary 1 (Sample Size for Hutchinson’s Estimator).

The relative error is absolutely bounded by $ϵ$ with probability $1-\delta$, provided that the sample size $n$ is bigger or equal to

$n(ϵ,\delta )={\displaystyle \frac{2\left(1-\frac{8}{3}ϵ\right)\log \left(\frac{1}{\delta }\right)}{{ϵ}^2}}.$ (10)

This holds for any and .

II-D Numerical Comparison with Related Work

Below we demonstrate numerical improvements with respect to the prior works  [3, 20]. The bounds, with exact constants, are summarized in Table I below.

Author	Sample Size $n=n(ϵ,\delta )$
this work	$\frac{2\left(1-\frac{8}{3}ϵ\right)\log \left(\frac{1}{\delta }\right)}{{ϵ}^2}$
[20]	$\frac{6\log \left(\frac{2}{\delta }\right)}{{ϵ}^2}$
[3]	$\frac{6\log \left(\frac{2}{\delta }\right)+6\mathrm{r}\mathrm{a}\mathrm{n}\mathrm{k}(A)}{{ϵ}^2}$

TABLE I: Bounds on the accuracy of Hutchinson’s Estimator.

Furthermore, in Figure 2, we present a detailed evaluation of how the sample size $n$ depends on the error $ϵ$, when the confidence parameter $\delta$ is fixed. The setup for this experiment is the discussed hessian of a toy classifier on MNIST.

Figure 2: Sample size for Hutchinson’s Estimator. We assume $\delta ={10}^{-3}$ which results in confidence of $0.999$. The matrix size is $7840$, which corresponds to the Hessian of the linear MNIST classification model (input data are black-white images in the resolution of $28\times 28$, grouped in $10$ classes).

III Auxiliary Results

III-1 Matrix Theory

We will need the following fact on the decomposition of symmetric matrices [26]. Recall that a matrix $B$ is orthonormal when $B^T=B^{-1}$ (${\cdot }^T$ denotes transposition); in particular, the columns and rows of $B$ are of unit length.

Lemma 1 (Symmetric Matrix Factorization).

Any symmetric real matrix $A$ can be written as $A=B^T\mathrm{\Lambda }B$ where $B$ is an orthonormal real matrix and $\mathrm{\Lambda }$ is diagonal, consisting of (necessarily real) eigenvalues of $A$.

III-2 Quadratic Chaos

We will need the following special case of the celebrated Hypercontractivity Inequality [23, 22]:

Lemma 2 (($2,d)$-Hypercontractivity).

Let $F$ be a polynomial of degree $2$ in Rademacher variables $Z_i$ (e.g. $F={\sum }_{i\ne j}{a}_{i,j}{Z}_i{Z}_j$ for some weights $a_{i,j}$). Then

${\Vert F\Vert }_d⩽(d-1){\Vert F\Vert }_2.$ (11)

III-3 Sub-Gamma Random Variables

The tail probability of sub-gamma distributions can be estimated from the MGF bound by the Crammer-Chernoff method, as follows:

Lemma 3 (Sub-Gamma Tails).

Let $X\in \mathrm{\Gamma }(v,c)$, then

$\forall ϵ>0:\mathrm{Pr}[|X|⩾ϵ]⩽{\mathrm{e}}^{-\frac{{ϵ}^2}{v+cϵ}}.$ (12)

Moreover, for sub-gamma distributions the following composition property holds for sums of independent terms:

Lemma 4 (Sub-Gamma Sums).

Let random variables $X_i\in \mathrm{\Gamma }(v_i,c_i)$ be independent, then

${\displaystyle \sum _i}{X}_i\in \mathrm{\Gamma }(v^{\prime },c^{\prime }),v^{\prime }\triangleq {\displaystyle \sum _i}{v}_i,c^{\prime }\triangleq \underset{i}{\max }{c}_i.$ (13)

The proofs of these lemmas can be found in [24].

IV Proofs

IV-A Proof of Theorem 1

Define the following random variable

$Y=z^{T}Az-𝔼[z^{T}Az]=z^{T}Az-\mathrm{tr}(A).$ (14)

then our task is to bound the concentration of $Y$.

IV-A1 Reduction to Off-Diagonal Quadratic Chaos

Since $A$ is symmetric, by Lemma 1 we have $A=B^T\mathrm{\Lambda }B$ where $B$ is orthonormal and $\mathrm{\Lambda }$ is diagonal made of the eigenvalues of $A$. Since, in addition, the matrix $A$ is positive semi-definite, its eigenvalues are non-negative. Thus, we can write

$z^{T}Az={(Bz)}^T\mathrm{\Lambda }(Bz)={\Vert {\mathrm{\Lambda }}^{1/2}Bz\Vert }_2^2.$ (15)

This can be decomposed as

$z^{T}Az={\displaystyle \sum _i}{Y}_i,Y_i$ $\triangleq {({\mathrm{\Lambda }}^{1/2}Bz)}_i^2.$ (16)

To further simplify, denote by ${\lambda }_1,\mathrm{\dots },{\lambda }_m$ all the eigenvalues of $A$; these are precisely the diagonal entries of $\mathrm{\Lambda }$. We note that $Y_i={\lambda }_i{({\sum }_j{B}_{i,j}{z}_j)}^2$, and $𝔼Y_i={\lambda }_i{\sum }_j{B}_{i,j}^2$ (because $𝔼z_j^2=1$). Therefore we obtain

$Y_i-𝔼Y_i={\lambda }_i{\displaystyle \sum _{j\ne j^{\prime }}}{B}_{i,j}{B}_{i,j^{\prime }}{z}_j{z}_{j^{\prime }}.$ (17)

Thus, every $Y_i$ is a quadratic form in $z_j$ with zero-diagonal.

IV-A2 Bounding Quadratic Chaos

For every fixed $i$ we apply Lemma 2 to $Y_i-𝔼Y_i$, which is a polynomial of degree $2$ in Rademacher random variables $z_i$ (note that the off-diagonal property guarantees the degree is exactly two).

Since the random variables $z_j{z}_{j^{\prime }}$ indexed by tuples $(j,j^{\prime })$ for $j\ne j^{\prime }$ are uncorrelated, we have ${\Vert Y_i-𝔼Y_i\Vert }_2^2={\lambda }_i^2{\sum }_{j\ne j^{\prime }}{B}_{i,j}^2{B}_{i,j^{\prime }}^2$; but $B$ is orthonormal, so ${\sum }_{i,j}{B}_{i,j}^2=1$ and ${\Vert Y_i-𝔼Y_i\Vert }_2^2⩽{\lambda }_i^2$. Thus, we obtain

${\Vert Y_i-𝔼Y_i\Vert }_d⩽{\lambda }_i^2(d-1).$ (18)

By the triangle inequality

${\Vert {\displaystyle \sum _i}(Y-𝔼Y_i)\Vert }_d⩽(d-1){\displaystyle \sum _i}{\lambda }_i.$ (19)

Finally by the definition of $Y_i$ and the standard identity $\mathrm{tr}(A)={\sum }_i{\lambda }_i$ from linear algebra [26]

${\Vert z^{T}Az-\mathrm{tr}(A)\Vert }_d⩽(d-1)\mathrm{tr}(A).$ (20)

Dividing the both sides of this inequality by $\mathrm{tr}(A)$, we complete the proof of the first part of Theorem 1.

IV-A3 Bounding Moment Generating Function

Let $E=z^{T}Az/\mathrm{tr}(A)-1$ be the relative error of the estimator. Then by the previous result, the Taylor expansion ${\mathrm{e}}^{tx}={\sum }_{d⩾0}{(tx)}^d/d!$, and the fact that $𝔼E=0$ we have

$MGF(E)⩽1+{\displaystyle \sum _{d⩾2}}{\displaystyle \frac{t^d{(d-1)}^d}{d!}}.$ (21)

Let $a_d={(d-1)}^d/d!$. Then we have

${\displaystyle \frac{a_{d+1}}{a_d}}={\displaystyle \frac{d^{d+1}}{{(d-1)}^d}}\cdot {\displaystyle \frac{1}{d+1}}.$ (22)

This ratio is decreasing when $d⩾2$, and thus

${\displaystyle \frac{a_{d+1}}{a_d}}⩽{\displaystyle \frac{8}{3}}.$ (23)

Therefore, it follows that

$\mathrm{MGF}(E)⩽1+{\displaystyle \frac{t^2}{2(1-\frac{8}{3}t)}},$ (24)

which, by the standard inequality $\log (1+u)⩽u$, implies

$\log \mathrm{MGF}(E)⩽{\displaystyle \frac{t^2}{2(1-\frac{8}{3}t)}},$ (25)

so that, by definition, $E\in \mathrm{\Gamma }(1,8/3)$. This completes the proof of the second part of Theorem 1.

IV-A4 Concluding Concentration Properties

The third part of Theorem 1 follows directly by Lemma 3.

IV-B Proof of Theorem 2

We first apply the result on sum of independent sub-gamma distributions from Lemma 4, which proves the first part of the theorem. The second part follows by the result on sub-gamma tails, stated in Lemma 3.

IV-C Proof of Corollary 1

The corollary follows by equating the right-hand side of the tail bound from Theorem 2, and equating with $\delta$. Solving with respect to $n$ gives the claimed formula on the sample size.

V Conclusion

This work establishes a new state-of-art bound on the classical trace estimator due to Hutchinson. The core idea is the clever usage of bounds on Rademacher Chaos and the theory of sub-gamma distributions. The results are numerically superior and of immediate interest to any problems involving trace estimation. Besides, the author hopes to contribute to better popularization of the discussed techniques, novel in the problem context, within the computer science community.