Machine learning of measurement schemes for efficient quantum observable estimation

Technical Summary: Machine Learning of Measurement Schemes for Efficient Quantum Observable Estimation

Problem Statement
Estimating the expectation value of observables is a fundamental subroutine in quantum computing, particularly for near-term algorithms like the Variational Quantum Eigensolver (VQE). However, the destructive nature of quantum measurements necessitates preparing quantum states multiple times (shots), creating a performance bottleneck. While observables are typically decomposed into sums of Pauli strings ($O = \sum_j a_j P_j$), naive estimation requires a number of measurements scaling quadratically with the $L_1$ norm of the coefficients ($\|\vec{a}\|_1$). Existing methods to reduce this cost generally fall into two categories:

1. Grouping methods (e.g., LDF, OGM): These output a probability distribution over a small set of measurement groups. They often rely on hand-crafted heuristics and struggle with molecular Hamiltonians containing exponentially many terms.
2. Classical Shadow methods (e.g., LBCS): These use distributions over all Pauli measurements. While scalable, they often employ simple structures that fail to achieve optimal shot efficiency for complex molecular systems.

The core challenge is how to enhance the representation power of measurement schemes to achieve better shot efficiency for complicated observables without relying on rigid, hand-crafted heuristics.

Methodology
The authors propose a framework to learn measurement schemes directly from the observable structure using machine learning, specifically stochastic gradient descent (SGD) and a two time-scale update rule (TTUR).

1. Composite Measurement Scheme (CMS):
   The authors define a "measurement scheme" as a generator that outputs a set of measurements. They introduce the Composite Measurement Scheme, which combines multiple simple measurement schemes (sub-schemes) via a trainable mixing probability distribution $\vec{r}$.

   • C-LBCS (Composite-Locally Biased Classical Shadow): As a concrete realization, they propose C-LBCS. This method learns a mixture of Locally Biased Classical Shadow (LBCS) schemes. Each sub-scheme $S_k$ is an LBCS with adjustable per-qubit Pauli distributions $\{\beta^k_i\}$. The composite scheme samples a sub-scheme according to $\vec{r}$ and then generates a Pauli string based on that sub-scheme's distribution.

2. Cost Function (Average One-Shot Variance):
   To optimize the scheme without knowledge of the target quantum state, the authors define a cost function $V$ based on the average one-shot variance over the entire Hilbert space (using the Haar measure).

   • The cost function is derived as $V = \sum_j a_j^2 \frac{1}{h_j}$, where $h_j$ is the average probability that the $j$-th Pauli term is "covered" (measured) by the scheme.
   • For C-LBCS, $h_j$ is decomposed as $\sum_k r_k h^k_j$, where $h^k_j$ is the coverage probability of the $k$-th sub-scheme. This structure allows the cost function to be optimized with respect to both the mixing weights $\vec{r}$ and the sub-scheme parameters $\{\theta^k_i\}$.

3. Optimization Strategies:
   The optimization landscape is non-convex overall but exhibits a bi-convex structure with respect to the mixing weights $\vec{r}$ and the coverage probabilities $\{h^k_j\}$. To navigate this, the authors employ:

   • Gradient Rescaling: To prevent sub-schemes with low mixing weights ($r_k \approx 0$) from being "frozen out" during training, the gradient for sub-scheme parameters is rescaled by $1/r_k$.
   • Two Time-Scale Update Rule (TTUR): Inspired by GAN training, $\vec{r}$ and $\{\theta^k_i\}$ are updated with different learning rates (specifically, the learning rate for $\vec{r}$ is set 10 times smaller than for the sub-scheme parameters) to stabilize convergence.
   • Stochastic Gradient Descent (SGD): To handle observables with thousands of terms (common in electronic structure problems), the training uses mini-batches of Pauli terms, enabling scalability and GPU acceleration.

Key Results
The authors numerically demonstrate the efficacy of C-LBCS on molecular systems ranging from small molecules to CO$_2$ (30 qubits, $\sim$11,000 terms).

• Performance: C-LBCS consistently outperforms state-of-the-art methods, including Derandomized Classical Shadow (Derand), Overlapped Group Measurement (OGM), and ShadowGrouping (SG), across various molecular Hamiltonians (LiH, H$_6$, H$_2$O, NH$_3$, N$_2$, C$_2$H$_2$, C$_2$H$_4$, CO$_2$) and encoding transformations (Jordan-Wigner and Bravyi-Kitaev).
• Scalability: The method scales effectively with the number of sub-schemes. Increasing the number of sub-schemes allows for a systematic reduction in variance, outperforming SG even when the number of sub-schemes is comparable to the number of groups in OGM.
• Robustness: The training is robust to small batch sizes (down to 1/128 of the total terms), indicating suitability for large-scale observables. The ablation studies confirm that both gradient rescaling and TTUR are critical for achieving the best results.
• Efficiency: Unlike greedy sequential methods (like SG), C-LBCS training is fully parallelizable on GPUs. The classical overhead for generating measurements post-training is low ($O(n_S + n_q)$), and the one-time training cost is amortized over many measurement rounds in typical workflows like VQE.

Significance and Claims
The paper claims to open a "reliable and scalable path toward efficient observable estimation on large quantum systems."

• Shift from Heuristics to Learning: The primary contribution is shifting the paradigm from hand-crafted, bottom-up heuristics to a top-down, end-to-end learning framework that directly optimizes measurement schemes for specific observables.
• Universality: The authors prove that C-LBCS is universal; with $3^{n_q}$ sub-schemes, it can simulate any distribution of Pauli measurements acting non-trivially on every qubit. However, they note that in practice, far fewer sub-schemes are needed to achieve state-of-the-art efficiency.
• Practicality: The method addresses the bottleneck of measurement efficiency in near-term quantum algorithms without requiring prior knowledge of the quantum state. The authors emphasize that while the cost function is defined in the limit of infinite measurements, the resulting schemes provide superior performance even in finite-shot regimes for the tested systems.

The authors remain modest regarding limitations, noting that the cost function assumes an infinite measurement limit; if the total measurement budget is very small relative to the number of terms, the trained scheme may not be strictly optimal for that specific finite budget. Additionally, the approach relies on the assumption that the target state is unknown, which is standard for VQE but requires the cost function to be state-independent.