Hutchinson’s Estimator is Bad at Kronecker-Trace-Estimation

We first remark there is a great deal of literature on trace estimation without the Kronecker constraint. The design of Hutchinson’s Estimator is attributed to [Girard, 1987, Hutchinson, 1989], and some of the core papers that analyze the estimator are [Avron and Toledo, 2011, Roosta-Khorasani and Ascher, 2015, Cortinovis and Kressner, 2021].

Although relatively a new topic, trace estimation in the Kronecker-matrix-vector oracle model has recently been explored in the literature. From the physics literature, the works of [Feldman et al., 2022] and [Huang et al., 2020] analyze trace estimators when $\mathbf{A}=\otimes_{i=1}^{k}\mathbf{A}_{i}$, and shows some empirical results for their performance. From a theoretical perspective, [Vershynin, 2020] provides bounds that can be used to analyze the Kronecker-Hutchinson Estimator. However, the paper provides very coarse bounds for this problem, showing that if $k=O(d)$ then we have $\boldsymbol{\mathrm{x}}^{\intercal}\mathbf{A}\boldsymbol{\mathrm{x}}\in\operatorname{tr}(\mathbf{A})\pm O(kd^{k-1}\mathinner{\lVert\mathbf{A}\rVert}_{2})$, where $\mathinner{\lVert\mathbf{A}\rVert}_{2}$ is the operator norm of $\mathbf{A}$. To obtain a relative error $(1\pm\varepsilon)$, this suggests an overall sample complexity of $\ell=O(\frac{kd^{k-1}}{\varepsilon^{2}}\frac{\mathinner{\lVert\mathbf{A}\rVert}_{2}}{\operatorname{tr}(\mathbf{A})})=O(\frac{D}{\varepsilon^{2}}\frac{\mathinner{\lVert\mathbf{A}\rVert}_{2}}{\operatorname{tr}(\mathbf{A})})$, which is generally more than $D$ queries.

The research of [Bujanovic and Kressner, 2021] directly analyzes the Kronecker-Hutchinson Estimator when $k=2$, showing first that $\ell=O(\frac{1}{\varepsilon^{2}})$ samples suffice to ensure that $H_{\ell}(\mathbf{A})\geq(1-\varepsilon)\operatorname{tr}(\mathbf{A})$, and second that $\ell=O(\frac{d}{\varepsilon^{2}}\frac{\mathinner{\lVert A\rVert}_{2}}{\operatorname{tr}(\mathbf{A})})$ samples suffice to ensure that $H_{\ell}(\mathbf{A})\leq(1+\varepsilon)\operatorname{tr}(\mathbf{A})$. Lastly, Remark 2 of the arXiv version of [Ahle et al., 2020] is a result on using a Kronecker-Johnson-Lindenstrauss transform, and implies that $O(\frac{3^{k}}{\varepsilon^{2}})$ samples suffice for the Kronecker-Hutchinson Estimator with Rademacher vectors to achieve relative error $(1\pm\varepsilon)$. While this worst-case bound is tight, the paper gives no way to understand when this bound is lose, and for what matrices fewer samples are needed. Further, all the aforementioned theoretical papers rely on heavy-duty analysis, for instance involving Gaussian and Rademacher Chaos theory. In contrast, our contributions give tighter results using simpler techniques.

We also note that [Ahle et al., 2020] does also include a Johnson-Lindenstrauss style guarantee for working with Kronecker-shaped data using a sketch that has $O(\frac{k}{\varepsilon^{2}})$ rows. However, this sketch does not actually operate in the Kronecker-matrix-vector oracle model.

Our core technical contribution is an exact expression of the variance of Hutchinson’s Estimator when our query vectors are formed as a Kronecker product of Gaussian vectors. This also provides an upper bound on the variance when the estimator instead uses a Kronecker product of Rademacher vectors. For PSD matrices, our results imply a worst-case upper bound of $\ell=O(\frac{3^{k}}{\varepsilon^{2}})$ on the number of samples needed to estimate the trace of $\mathbf{A}$ to relative error $(1\pm\varepsilon)$. We match this upper bound with a lower bound showing that the Kronecker-Hutchinson Estimator with iid zero-mean unit-variance vectors has to use $\ell=\Omega(\frac{3^{k}}{\varepsilon^{2}})$ samples to achieve standard deviation $\varepsilon\operatorname{tr}(\mathbf{A})$. Lastly, we show that if we are allowed to use complex Kronecker-matrix-vector products instead of real matrix-vector products (i.e. if $\boldsymbol{\mathrm{x}}=\boldsymbol{\mathrm{x}}_{1}\otimes\cdots\otimes\boldsymbol{\mathrm{x}}_{k}$ where $\boldsymbol{\mathrm{x}}_{i}\in{\mathbb{C}}^{d}$), then this rate improves to $\Theta(\frac{2^{k}}{\varepsilon^{2}})$.

To more concretely state our core technical contribution, we recall the partial trace operator from quantum physics. In quantum physics, a Hermitian density matrix $\mathbf{A}\in{\mathbb{C}}^{d^{k}\times d^{k}}$ is roughly speaking used to represent the joint probability distribution for a set of $k$ particles. Suppose these particles are partitioned into two sets, so that $[k]={\mathcal{S}}\cup\bar{}{\mathcal{S}}$. Further suppose that Alice can observe the particles in ${\mathcal{S}}$, and that Bob can observe the particles in $\bar{}{\mathcal{S}}$, but they cannot observe each other’s particles. Then, the partial trace of $\mathbf{A}$ with respect to ${\mathcal{S}}$, denoted $\operatorname{tr}_{{\mathcal{S}}}(\mathbf{A})$, describes the marginal distribution of the particles that Bob observes. This is often called “tracing out the subsystem ${\mathcal{S}}$”.

We include a general formal definition of the partial trace in Section 2, but a rough intuition is easily gained by examining the $k=2$ case. Here, we can break down $\mathbf{A}={\mathbb{R}}^{d^{2}\times d^{2}}$ as a block matrix full of $\mathbf{A}_{ij}\in{\mathbb{R}}^{d\times d}$:

$$\mathbf{A}=\begin{bmatrix}\mathbf{A}_{11}&\mathbf{A}_{12}&\ldots&\mathbf{A}_{1d}\\ \mathbf{A}_{21}&\mathbf{A}_{22}&\ldots&\mathbf{A}_{2d}\\ \vdots&\vdots&\ddots&\vdots\\ \mathbf{A}_{d1}&\mathbf{A}_{d2}&\ldots&\mathbf{A}_{dd}\\ \end{bmatrix}.$$

Then, the partial trace of $\mathbf{A}$ with respect to $\{1\}$ is the sum of the diagonal blocks, i.e. $\operatorname{tr}_{\{1\}}(\mathbf{A})=\sum_{i=1}^{d}\mathbf{A}_{ii}$, and the partial trace of $\mathbf{A}$ with respect to $\{2\}$ is the matrix that contains the trace of each block matrix, i.e.

$$\operatorname{tr}_{\{2\}}(\mathbf{A})=\begin{bmatrix}\operatorname{tr}(\mathbf{A}_{11})&\operatorname{tr}(\mathbf{A}_{12})&\ldots&\operatorname{tr}(\mathbf{A}_{1d})\\ \operatorname{tr}(\mathbf{A}_{21})&\operatorname{tr}(\mathbf{A}_{22})&\ldots&\operatorname{tr}(\mathbf{A}_{2d})\\ \vdots&\vdots&\ddots&\vdots\\ \operatorname{tr}(\mathbf{A}_{d1})&\operatorname{tr}(\mathbf{A}_{d2})&\ldots&\operatorname{tr}(\mathbf{A}_{dd})\\ \end{bmatrix}.$$

We will also need to recall the partial transpose operator from quantum physics. This operator is often used in quantum physics, but lacks an intuitive physical meaning analogous to what the partial trace enjoys. Like the partial trace, the partial transpose can also be taken over any subset of particles ${\mathcal{V}}\subseteq[k]$. We denote the partial transpose of $\mathbf{A}$ with respect to ${\mathcal{V}}$ as $\mathbf{A}^{\intercal_{{\mathcal{V}}}}$, and call this “transposing the subsystem ${\mathcal{V}}$”. We have a formal definition of the partial transpose in Section 2. But, when $k=2$, the two partial transposes of $\mathbf{A}$ are:

$$\mathbf{A}^{\intercal_{\{1\}}}=\begin{bmatrix}\mathbf{A}_{11}&\mathbf{A}_{21}&\cdots&\mathbf{A}_{d1}\\ \mathbf{A}_{12}&\mathbf{A}_{22}&\cdots&\mathbf{A}_{d2}\\ \vdots&\vdots&\ddots&\vdots\\ \mathbf{A}_{1d}&\mathbf{A}_{2d}&\cdots&\mathbf{A}_{dd}\end{bmatrix}\hskip 14.22636pt\text{and}\hskip 14.22636pt\mathbf{A}^{\intercal_{\{2\}}}=\begin{bmatrix}\mathbf{A}_{11}^{\intercal}&\mathbf{A}_{12}^{\intercal}&\cdots&\mathbf{A}_{1d}^{\intercal}\\ \mathbf{A}_{21}^{\intercal}&\mathbf{A}_{22}^{\intercal}&\cdots&\mathbf{A}_{2d}^{\intercal}\\ \vdots&\vdots&\ddots&\vdots\\ \mathbf{A}_{d1}^{\intercal}&\mathbf{A}_{d2}^{\intercal}&\cdots&\mathbf{A}_{dd}^{\intercal}\end{bmatrix}.$$

Having setup both the partial trace and the partial transpose, we can now state our core theorem:

The matrix $\bar{}\mathbf{A}$ seems odd at a glance, so we first justify why it is natural. For the classical $k=1$ setting with Gaussian vectors, we can write that $\operatorname*{Var}[\boldsymbol{\mathrm{x}}^{\intercal}\mathbf{A}\boldsymbol{\mathrm{x}}]=2\mathinner{\lVert\mathbf{A}\rVert}_{F}^{2}$ for all symmetric matrices $\mathbf{A}$ [Avron and Toledo, 2011, Meyer, 2021]. For asymmetric matrices, this equality does not hold. Instead, we let $\hat{}\mathbf{A}\;{\vcentcolon=}\;\frac{\mathbf{A}+\mathbf{A}^{\intercal}}{2}$, and note that $\boldsymbol{\mathrm{x}}^{\intercal}\mathbf{A}\boldsymbol{\mathrm{x}}=\boldsymbol{\mathrm{x}}^{\intercal}\hat{}\mathbf{A}\boldsymbol{\mathrm{x}}$ for all $\boldsymbol{\mathrm{x}}\in{\mathbb{R}}^{D}$. Since $\hat{}\mathbf{A}$ is symmetric, we then find that $\operatorname*{Var}[\boldsymbol{\mathrm{x}}^{\intercal}\mathbf{A}\boldsymbol{\mathrm{x}}]=2\mathinner{\lVert\hat{}\mathbf{A}\rVert}_{F}^{2}$. Theorem 1 generalizes this idea to all possible ways to partially transpose across all possible subsystems ${\mathcal{V}}\subseteq[k]$.

Next, note that $\mathinner{\lVert\hat{}\mathbf{A}\rVert}_{F}\leq\mathinner{\lVert\mathbf{A}\rVert}_{F}$, and that symmetric matrices have $\mathbf{A}=\hat{}\mathbf{A}$. So, we know that $\operatorname*{Var}[\boldsymbol{\mathrm{x}}^{\intercal}\mathbf{A}\boldsymbol{\mathrm{x}}]=2\mathinner{\lVert\hat{}\mathbf{A}\rVert}_{F}^{2}\leq 2\mathinner{\lVert\mathbf{A}\rVert}_{F}^{2}$ holds for all matrices $\mathbf{A}$, with equality for symmetric $\mathbf{A}$. This bound can be easier to work with, since it does not require worrying about $\hat{}\mathbf{A}$. For large $k$, we similarly bound $\mathinner{\lVert\operatorname{tr}_{\mathcal{S}}(\bar{}\mathbf{A})\rVert}_{F}\leq\mathinner{\lVert\operatorname{tr}_{\mathcal{S}}(\mathbf{A})\rVert}_{F}$, providing a more approachable version of Theorem 1:

Note that many reasonable matrices will have $\mathbf{A}=\bar{}\mathbf{A}$. For instance, recall that the Kronecker-matrix-vector oracle model may be used when $\mathbf{A}$ is the sum of a small number of matrices, where each summand is a Kronecker product of $k$ matrices. If each summand is the Kronecker product of symmetric matrices, then we have $\mathbf{A}=\bar{}\mathbf{A}$.

We demonstrate the utility of Theorem 2 by using it to show the following guarantee on estimating the trace of a random rank-one matrix:

This rate of $(2+\frac{1}{d})^{k}$ is obviously expensive, but our sample complexity is lower than what is achievable with any other technique that we are aware of. Nevertheless, even without $\bar{}\mathbf{A}$, this sum of partial traces may sometimes still be too unwieldy to work with. For such cases, we provide a worst-case bound for PSD matrices:

In particular, by Chebyshev’s Inequality, Theorem 4 implies that $\ell=O(\frac{3^{k}}{\varepsilon^{2}})$ samples suffice to estimate $\operatorname{tr}(\mathbf{A})$ to relative error $(1\pm\varepsilon)$. For the case of [Bujanovic and Kressner, 2021], when $k=2$, we get the following:

This success probability can be boosted to $1-\delta$ for an arbitrary $\delta$ by picking up a $\log(\frac{1}{\delta})$ dependence in the sample complexity by returning the median of $O(\log(\frac{1}{\delta}))$ iid trials of $H_{\ell}(\mathbf{A})$. We further show that the expensive rate of $O(\frac{3^{k}}{\varepsilon^{2}})$ is essentially tight in the worst case for the Kronecker-Hutchinson Estimator:

The Kronecker product of two matrices $\mathbf{A}\in{\mathbb{R}}^{n\times m}$ and $\mathbf{B}\in{\mathbb{R}}^{p\times q}$ is $\mathbf{A}\otimes\mathbf{B}\in{\mathbb{R}}^{np\times mq}$, where

$$\mathbf{A}\otimes\mathbf{B}\;{\vcentcolon=}\;\begin{bmatrix}[\mathbf{A}]_{11}\mathbf{B}&[\mathbf{A}]_{12}\mathbf{B}&\ldots&[\mathbf{A}]_{1m}\mathbf{B}\\ [\mathbf{A}]_{21}\mathbf{B}&[\mathbf{A}]_{22}\mathbf{B}&\ldots&[\mathbf{A}]_{2m}\mathbf{B}\\ \vdots&\vdots&\ddots&\vdots\\ [\mathbf{A}]_{n1}\mathbf{B}&[\mathbf{A}]_{n2}\mathbf{B}&\ldots&[\mathbf{A}]_{nm}\mathbf{B}\end{bmatrix}$$

We also use the shorthands $\otimes_{i=1}^{k}\boldsymbol{\mathrm{x}}_{i}\;{\vcentcolon=}\;\boldsymbol{\mathrm{x}}_{1}\otimes\boldsymbol{\mathrm{x}}_{2}\otimes\cdots\otimes\boldsymbol{\mathrm{x}}_{k}$, as well as $\mathbf{A}^{\otimes k}\;{\vcentcolon=}\;\otimes_{i=1}^{k}\mathbf{A}$. The Kronecker product has three key properties that we repeatedly use:

At one point, we will also require a short lemma about the squared Frobenius norm and the Kronecker product. This lemma is roughly equivalent to noting that $\mathinner{\lVert\mathbf{A}\rVert}_{F}^{2}=\sum_{i,j=1}^{d}\mathinner{\lVert\mathbf{A}_{ij}\rVert}_{F}^{2}$ in the block matrix decomposition of $\mathbf{A}$ when $k=2$ in Section 1.2.

We take a moment to formalize the partial trace operator and characterize some useful properties of it. We start by defining the partial trace with respect to a single subsystem:

Equivalently, the partial trace operator can be defined as the unique linear operator which guarantees that for all $\mathbf{A}\in{\mathbb{R}}^{d^{i-1}\times d^{i-1}}$, $\mathbf{B}\in{\mathbb{R}}^{d\times d}$, and $\mathbf{C}\in{\mathbb{R}}^{d^{k-i}\times d^{k-i}}$, we have that $\operatorname{tr}_{\{i\}}(\mathbf{A}\otimes\mathbf{B}\otimes\mathbf{C})=\operatorname{tr}(\mathbf{B})\cdot(\mathbf{A}\otimes\mathbf{C})$.

We will often consider the partial trace of $\mathbf{A}$ with respect to a set of subsystems ${\mathcal{S}}\subseteq[k]$. We can think of this as either composing the operation in Equation 1 $|{{\mathcal{S}}}|$ many times, or as extending the summation in Equation 1 to sum over all $d^{|{{\mathcal{S}}}|}$ basis vectors that span the product of subspaces defined in ${\mathcal{S}}$. To see this equivalence, consider the partial trace of $\mathbf{A}$ with respect to ${\mathcal{S}}=\{1,2\}$:

	$\displaystyle\operatorname{tr}_{\{1,2\}}(\mathbf{A})$	$\displaystyle=\operatorname{tr}_{\{1\}}(\operatorname{tr}_{\{2\}}(\mathbf{A}))$
		$\displaystyle=\sum_{j=1}^{d}(\boldsymbol{\mathrm{e}}_{j}\otimes\mathbf{I}^{\otimes k-1})^{\intercal}\bigg{(}\sum_{\hat{j}=1}^{d}(\mathbf{I}\otimes\boldsymbol{\mathrm{e}}_{\hat{j}}\otimes\mathbf{I}^{\otimes k-2})^{\intercal}\mathbf{A}(\mathbf{I}\otimes\boldsymbol{\mathrm{e}}_{\hat{j}}\otimes\mathbf{I}^{\otimes k-2})\bigg{)}(\boldsymbol{\mathrm{e}}_{j}\otimes\mathbf{I}^{\otimes k-1})$
		$\displaystyle=\sum_{j=1}^{d}\sum_{\hat{j}=1}^{d}(\boldsymbol{\mathrm{e}}_{j}\otimes\mathbf{I}\otimes\mathbf{I}^{\otimes k-2})^{\intercal}(\mathbf{I}\otimes\boldsymbol{\mathrm{e}}_{\hat{j}}\otimes\mathbf{I}^{\otimes k-2})^{\intercal}\mathbf{A}(\mathbf{I}\otimes\boldsymbol{\mathrm{e}}_{\hat{j}}\otimes\mathbf{I}^{\otimes k-2})(\boldsymbol{\mathrm{e}}_{j}\otimes\mathbf{I}\otimes\mathbf{I}^{\otimes k-2})$
		$\displaystyle=\sum_{j=1}^{d}\sum_{\hat{j}=1}^{d}(\boldsymbol{\mathrm{e}}_{j}\otimes\boldsymbol{\mathrm{e}}_{\hat{j}}\otimes\mathbf{I}^{\otimes k-2})^{\intercal}\mathbf{A}(\boldsymbol{\mathrm{e}}_{j}\otimes\boldsymbol{\mathrm{e}}_{\hat{j}}\otimes\mathbf{I}^{\otimes k-2}).$		(Lemma 10)

Next, note that $\boldsymbol{\mathrm{e}}_{j}\otimes\boldsymbol{\mathrm{e}}_{\hat{j}}\in{\mathbb{R}}^{d^{2}}$ is a standard basis vector in ${\mathbb{R}}^{d^{2}}$ uniquely identified by the pair $(j,\hat{j})$. That is, we can replace this double sum above with a single sum over standard basis vectors in ${\mathbb{R}}^{d^{2}}$:

$\displaystyle\operatorname{tr}_{\{1,2\}}(\mathbf{A})$

$\displaystyle=\sum_{j=1}^{d^{2}}(\boldsymbol{\mathrm{e}}_{j}\otimes\mathbf{I}^{\otimes k-2})^{\intercal}\mathbf{A}(\boldsymbol{\mathrm{e}}_{j}\otimes\mathbf{I}^{\otimes k-2}).$

($\boldsymbol{\mathrm{e}}_{j}\in{\mathbb{R}}^{d^{2}}$ now)

Which now looks remarkably similar to Definition 14, but with a larger definition of “subsystem” that the standard basis vector is acting on. Formally, we define the partial trace of $\mathbf{A}$ with respect to a set ${\mathcal{S}}$ as following:

In principle, the order of subscripts does not really matter. Tracing out the first subsystem and then tracing out the second subsystem produces the exact same matrix as tracing out the second subsystem and then tracing out the first subsystem. However, notationally, a non-ascending order becomes messy since removing the second subsystem changes the indices. To see this messiness, note that the above statement would be formalized as $\operatorname{tr}_{\{1\}}(\operatorname{tr}_{\{2\}}(\mathbf{A}))=\operatorname{tr}_{\{1\}}(\operatorname{tr}_{\{1\}}(\mathbf{A}))$. To avoid this issue, we always expand the partial traces in ${\mathcal{S}}$ in sorted order as in Definition 15.

Moving on, a vital property of the partial trace is that it is trace-preserving:

Our analysis also heavily involves post-measurement reduced density matrices, which we now define, and whose long name originates from quantum physics. As mentioned in Section 1.2, if $\mathbf{A}$ is a matrix that describes the joint probability distribution of many particles (a density matrix), then $\operatorname{tr}_{\mathcal{S}}(\mathbf{A})$ describes the joint probability distribution of the particles that Bob can view. The Post-measurement Reduced Density Matrix (PMRDM) of $\mathbf{A}$ with respect to vector $\boldsymbol{\mathrm{x}}$ describes the joint distribution of the particles that Bob can view given that Alice measured her particles in the direction of the vector $\boldsymbol{\mathrm{x}}$.¹¹1For the purpose of understanding the paper, there is no particular need to understand this intuition. Those interested in learning more are recommended to read through Chapter 6 of [Aaronson, 2018]. For brevity, we refer to this matrix as the PMRDM.

Notice that the notation for the PMRDM ${\mathcal{M}}_{\{i\}}(\mathbf{A})$ does not explicitly show the vector $\boldsymbol{\mathrm{x}}_{i}$ that define the measurement. This is done for visual clarity, and is safe because the meaning of the vector $\boldsymbol{\mathrm{x}}_{i}$ will always be clear from context.²²2Namely, we always take $\boldsymbol{\mathrm{x}}_{i}$ to be the $i^{th}$ term in the expansion of $\boldsymbol{\mathrm{x}}=\otimes_{i=1}^{k}\boldsymbol{\mathrm{x}}_{i}$, where $\boldsymbol{\mathrm{x}}_{i}$ has either iid Gaussian or Rademacher entries. Further, while we could define the PMRDM with respect to an arbitrary set of subsystems ${\mathcal{S}}$, as done in Definition 15, our proofs only need to examine the PMRDM when looking at the first $i$ subsystems, so for simplicity we only define the PMRDM for that case.

We need one important property from the PMRDM, which relates it to the partial trace: