Bootstrapping Classical Shadows for Neural Quantum State Tomography

A promising way to efficiently parameterize a pure quantum state $\ket{\psi}$ of multiple qubits is with a neural network parameterized by a weight vector $\lambda$ as

where $s\in\mathcal{B}_{n}=\{0,1\}^{n}$ is a binary vector in the computational basis. Similarly, the larger state space of density matrices can be parameterized with a neural network. Among competing approaches, the one based on the idea of purification [10],

is particularly appealing, as the resulting density matrix is physical. Here, $\bar{s}$ are auxiliary input variables representing the extended dimensions of the Hilbert space of the purified state. Therefore, their number $n_{\bar{s}}$ constrains the entropy of $\rho_{\lambda}$, with $n_{s}$ being a theoretically proven sufficient upper bound for capturing any density matrix over the physical input variables $s$. We treat $n_{\bar{s}}$ as a hyperparameter on the same footing as other parameters affecting the neural network’s architecture and hence its expressive power.

When the neural network parameters $\lambda$ in the ansatz Eqs. 1 and 2 are embedded within a generative model based on an autoregressive architecture, the resulting model can be very expressive and amenable to an efficient extraction of information about the underlying wave function as long as the quantity of interest can be cast as an unbiased estimator over samples from the model. For example, the purity of $\rho_{\lambda}$ can be estimated by evaluating the Swap operator between the physical and auxiliary qubits [11, 12, 13].

The classical shadow formalism provides an efficient method for predicting various properties using relatively few measurements. A classical shadow is created by repeatedly performing a simple procedure: at every iteration $i$ a random unitary transformation $\rho\mapsto U_{i}\rho U_{i}^{\dagger}$ is applied, after which all the qubits in the computational basis are measured. We can efficiently make a record of the pairs $(U_{i},s_{i})$ in classical memory. Let $\rho_{i}=U^{\dagger}\outerproduct{s_{i}}{s_{i}}U$ be the associated density matrix. The mapping $\mathcal{M}$:$~{}\rho\mapsto\mathbb{E}[\rho_{i}]$ is invertible if the measurement protocol is tomographically complete. The set $\mathcal{S}_{\rho}=\{\mathcal{M}^{-1}(\rho_{i})\}_{i=1}^{N}$ is called a classical shadow of $\rho$ of size $N$ and can be used to construct an unbiased estimator $\hat{\rho}=\mathbb{E}[\mathcal{M}^{-1}(\rho_{i})]$ of $\rho$.

The classical shadow $\mathcal{S}_{\rho}$ suffices to predict $M$ arbitrary linear and nonlinear functions $O_{i}$ of the state $\rho$ up to an additive error $\epsilon$ if it is of size $\Omega(\log(M)~{}{\rm max}_{i}~{}||O_{i}||^{2}_{\rm shadow}/\epsilon^{2})$. The shadow norm depends on the ensemble from which the unitaries $U_{i}$ are drawn for creating the classical shadow [1] and has efficient bounds, particularly if the ensemble satisfies a unitary 3-design property. For example, this is the case if the unitaries are drawn from the Clifford set as in Ref. [1] or using analog evolutions of different lengths of time as in Ref. [14].

Two particularly interesting ensembles are the Pauli and Clifford groups. These ensembles are computationally favourable because of their efficient classical simulations, according to the Gottesman–Knill theorem. The inverse twirling maps for the $n$-qubit Clifford and Pauli groups take the forms $\mathcal{M}^{-1}_{n}\left(X\right)=(2^{n}+1)X-\mathbb{I}$ and $\otimes_{k=1}^{n}\mathcal{M}_{1}^{-1}$, respectively. The shadow norm scales well for the Clifford ensemble, while the Pauli ensemble is more experimentally friendly. For the Pauli ensemble, the shadow norm scales exponentially poorly with respect to the locality of an observable, rendering the procedure inefficient in estimating high-weight observables. A middle-ground solution suggested in the literature [4, 5, 6] is to use shallow subsets of the Clifford group.

In this paper, we show that an alternative approach to improving the predictive power of Pauli ensembles (i.e., the shallowest possible shadows) is using a neural network (or other parameterized ansatze). Whereas the raw classical shadow is powerful at predicting many observables efficiently, it is not always sample efficient due to the unphysical nature of the state it provides. This is because, while shadow tomography yields a trace-one quantum state, not all eigenvalues are guaranteed to be positive. However, neural quantum states allow for natively imposing physicality constraints on a reconstructed state.

Another advantage of the unitary design property of the Clifford ensemble is that it allows for the efficient prediction of nonlinear functions of a quantum state. This is particularly useful in estimating entropy-based quantities such as the Rényi entropies, since they involve two copies of the state, and thus quadratic functions of it. We show that neural networks also improve the predictive power of such nonlinear quantities when combined with the purification technique explained in Section II.1.

Prior to studying the training of neural quantum states with our new methods, we explore the benefits of using shadow weights and stabilizer-based importance sampling. To this end, we focus on the key training metrics: the informational content and the statistical properties of the gradient estimator.

To facilitate our analysis, we consider an empirical proxy for the data distribution $p^{\rm{data}}_{\rm{emp}}(\phi)\propto p^{\rm{data}}(\phi)\mathbb{I}_{\phi\in\widetilde{\mathcal{D}}}$, where $\mathbb{I}_{\phi\in\widetilde{\mathcal{D}}}$ is the indicator function that excludes unobserved stabilizers. This distribution represents the most-accurate reconstruction of the probability measure $p^{\rm{data}}$ given a dataset $\mathcal{D}$, assuming no prior knowledge about $p^{\rm{data}}$. To assess the informational content in the shadow weights, we consider the KL divergence ${\rm{KL}}\left(p^{\rm{data}}_{\rm{emp}}||p^{\rm{target}}\right)$ between this typically unknown distribution and $p^{\rm{target}}$, which denotes either of the two empirically accessible distributions used for training: $p^{\rm{emp}}$ based on Eq. 6 or $p^{\rm{sh}}$ specified by Eq. 7. Specifically, we study the dependence of ${\rm{KL}}\left(p^{\rm{data}}_{\rm{emp}}||p^{\rm{target}}\right)$ on the size of the dataset relative to the size of the system by scaling the number of qubits given a fixed dataset, as shown in Fig. 1(a), and scaling the number of measurements in the dataset for a fixed number of qubits, as depicted in Fig. 1(b).

The advantage of shadow weights in terms of estimating $p^{\rm{data}}_{\rm{emp}}$ is evident for smaller systems, as shown in Fig. 1 (a), though it gradually diminishes as the systems grow larger. In contrast, Fig. 1(b) indicates that this benefit can be reclaimed by increasing the number of measurements performed. A crucial factor in interpreting the plots in Fig. 1 is considering the difference in the configuration space dimensions of $p^{\rm{data}}_{\rm{emp}}$ associated with Pauli and Clifford measurements. Notably, Pauli measurements can reconstruct $4^{n}$ distinct stabilizers, whereas Clifford measurements can reconstruct $2^{0.5+\mathcal{O}(1)n^{2}}$ distinct stabilizers [15]. For $n=6$ qubits, the dataset sizes considered in Fig. 1(b) must be compared to the dimension of the stabilizer space. For the Pauli distribution, the dataset ranges from 5% to 500% of the total set size of distinct stabilizers. However, for the Clifford distribution, the maximum number of measurements is less than $0.001\%$ of the corresponding stabilizer set size. This significant difference explains the slight scaling advantage of the $p^{\rm{emp}}$ Pauli measurements over the $p^{\rm{emp}}$ Clifford measurements, as the former have a higher likelihood of obtaining repeated stabilizers, while the latter degenerate to a uniform distribution. For shadow weights, we observe a scaling advantage over empirical distributions in both regimes. However, unlike in the case of the empirical distributions, the Clifford distribution shows a slight improvement over the Pauli distribution. Such an advantage of Clifford shadow weights over Pauli ones, despite being in a regime with less data, is consistent with the theoretical expectation of an exponential benefit in predicting high-weight, low-rank observables (i.e., other stabilizers) [1]. Therefore, given a sufficiently large measurement dataset, training with shadow weights-based cross-entropy is expected to produce neural quantum states that better capture the underlying quantum state.

We now evaluate the advantages of the proposed stabilizer-based method of Section II.3.2, over the neural state-based method of Eq. 10, for Monte Carlo estimation of the overlap, which is used to compute the training gradient. For benchmarking purposes, we analyze the errors in the gradient estimates relative to the true gradient, obtained through an exhaustive evaluation of the inner product of Eq. 9. Specifically, we focus on the error in the gradient direction and distinguish between errors calculated over a mini-batch and a full batch of measurements. Figure 2 illustrates the evolution of these errors as the neural quantum state is trained using the true gradient. Figure 2(a) shows that both estimators struggle most with capturing the correct gradient direction during the early stages of training, where the neural quantum state is likely the most different from the target state. However, after several epochs, as the neural state becomes more similar to the target, the stabilizer-based estimation significantly improves, while the neural state estimation remains relatively high and inconsistent. Throughout the entire training process, the stabilizer-based estimator consistently provides estimates with much lower variance, which is essential for training stability. Additionally, our proposed estimator is consistently more accurate than the neural-based estimator. Being an unbiased estimator, the stochastic error due to mini-batching is expected to be cancelled out when averaged over the mini-batches and thus can be lowered using smaller learning rates, albeit at the cost of slower training dynamics. Figure 2(b) shows whether this expected behaviour holds for the given estimators. There is a clear difference between them: the stabilizer-based estimator functions as an unbiased estimator, reducing gradient errors by averaging over mini-batches. In contrast, the neural state-based estimator exhibits a systematic bias throughout the entire training duration, resulting in an extreme accumulation of errors from the true gradient. To eliminate this finite-sample bias, a significant increase in the samples generated by the neural state would be necessary. Therefore, given a fixed computational resource budget, our benchmarking results demonstrate that the stabilizer-based gradient estimation has a significant advantage over the neural state-based gradient estimation, highlighting its potential for optimization.

In this section, we focus on the reconstruction of pure GHZ states from Pauli and Clifford measurements as a benchmark for the training enhancements we have proposed. In particular, we use the shadow-based cross-entropy ($\mathcal{L}_{\text{SCE}}$), the empirical cross-entropy ($\mathcal{L}_{\text{ECE}}$), and the infidelity ($\mathcal{L}_{\text{inf}}$) as the loss functions. In addition, we employ the estimators Eqs. (10) and (II.3.2) in evaluating the overlap between the neural network and stabilizer states. The superscripts “NS” (for neural state-based samples) and “SS” (for stabilizer-based samples) appear with the corresponding loss function $\mathcal{L}$ to distinguish between the two estimators.

We illustrate the trainability and the generalization power of these training methods in Fig. 3 via the optimization progress curves of loss functions evaluated on the training set (upper row) and infidelity with respect to the true states (lower row), over 50 epochs. For Pauli measurements, plotted in Fig. 3(a) and (b), the superiority of cross-entropy-based training over the infidelity-based approach is apparent. The latter struggles to achieve significant progress in approximating the underlying quantum state. Moreover, the benefits of the cross-entropy-based loss function become even more pronounced when considering computational costs. Both loss functions rely on computing the overlap between the neural quantum state and stabilizers (see Eq. 9). For the infidelity loss function, each Pauli measurement leads to a proliferation of stabilizers that grows super-exponentially with the number of qubits (see Eq. 4). In contrast, this prohibitive scaling is entirely absent for both cross-entropy-based loss functions, as each Pauli measurement contributes a single stabilizer. In a previous work [2], we attempted to address the prohibitive scaling of the infidelity-based approach through Monte Carlo sampling of stabilizers from a classical shadow state. However, this strategy led to unstable training. Additionally, we note that stabilizer-based training improves in-training performance across all loss functions (the results for neural-state-based cross-entropy loss functions have been omitted for visual clarity). Moreover, these improvements during training translate to enhanced generalization only for cross-entropy loss functions. In fact, for the cross-entropy loss function, stabilizer-based training proves to be crucial.

We now discuss the Clifford measurements, beginning with the benchmarking performed on the six-qubit GHZ state shown in Fig. 3(c) and (d). Our previous observations regarding the advantages of stabilizer-based sampling generalize to this scenario and, in fact, are applicable to the infidelity loss function as well. Indeed, the loss function $\mathcal{L}_{\text{inf}}^{\text{SS}}$ leads to an improved generalization with the final infidelity of roughly an order of magnitude compared to the loss function $\mathcal{L}_{\text{inf}}^{\text{NS}}$, as shown in Fig. 3(d). Interestingly, this improvement in generalization is achieved despite their corresponding loss functions yielding similar values at the end of training. Further examination of Fig. 3(c) and (d) reveals that our loss function $\mathcal{L}_{\text{ECE}}^{\text{SS}}$ greatly outperforms the previously proposed shadow-based loss function $\mathcal{L}_{\text{inf}}^{\text{NS}}$ [2] in both the convergence rate and generalization. In addition, a comparison between the loss functions $\mathcal{L}_{\text{SCE}}^{\text{SS}}$ and $\mathcal{L}_{\text{ECE}}^{\text{SS}}$ reveals that our cross-entropy estimator leads to a superior convergence rate and lower infidelity, illustrating the benefits of shadow weights extracted from the classical shadow state (see Eq. (7)). Overall, both loss functions $\mathcal{L}_{\text{SCE}}^{\text{SS}}$ and $\mathcal{L}_{\text{inf}}^{\text{SS}}$ demonstrate the best convergence rates. The loss functions $\mathcal{L}_{\text{SCE}}^{\text{SS}}$ and $\mathcal{L}_{\text{ECE}}^{\text{SS}}$ exhibit the greatest generalization, approaching an infidelity of $0.01$, while our estimator shows the greatest stability in reaching states in proximity to the true states. An even more drastic difference in performance is observed for the eight-qubit GHZ state, shown in Fig. 3(e) and (f). In this case, the loss functions $\mathcal{L}_{\text{inf}}^{\text{NS}}$, $\mathcal{L}_{\text{inf}}^{\text{SS}}$, and $\mathcal{L}_{\text{ECE}}^{\text{SS}}$ struggle to escape local minima. While all of the loss functions result in infidelities that remain stuck at $0.99$, only the loss function $\mathcal{L}_{\text{SCE}}^{\text{SS}}$ stands out, achieving an infidelity as low as $0.05$. To summarize, our results indicate that our proposed method provides the best training objective for neural networks trained to approximate an underlying quantum state from limited measurements.

In this section, we describe our application of the methods introduced earlier, namely, the shadow-based cross-entropy loss function with stabilizer-based sampling, to mixed quantum states generated by noisy circuits. Specifically, we consider a circuit that prepares a six-qubit GHZ state beginning from the all-zeroes state $\ket{0}^{\otimes 6}$. The circuit is composed of a Hadamard gate applied to the first qubit, followed by a series of CNOT gates applied to the consecutive pairs of qubits. For the noise model, we assume two-qubit depolarization noise on the CNOT gates. The depolarizing noise channel $\rho\mapsto(1-\frac{16}{15}p)\rho+\frac{16}{15}p\mathbb{I}/4$ is parameterized via a noise strength parameter $p$. Motivated by the experimental relevance of our technique, we consider Pauli measurement datasets, as those are less noisy than Clifford measurements. As the ansatz for the density matrix, we employ the purification ansatz Eq. (2) with a transformer neural network used to parameterize $\psi_{\lambda}(s,\bar{s})$. This approach allows us to impose physical constraints on the ansatz.