Train once and generalize: Zero-shot quantum state preparation with RL

Authors: Krishna Agaram, Siddhant Midha, Adrian Müller, Vikas Garg
Primary area: applications to physical sciences (physics, chemistry, biology, etc.)
Keywords: Quantum State Preparation, Deep Reinforcement Learning, Zero-shot Inference, Off-the-shelf Algorithms, Generalization

Abstract:
Quantum state preparation forms an essential cornerstone of quantum information science and quantum algorithms. Designing efficient and scalable methods for approximate state preparation on near-term quantum devices remains a significant challenge, with worst-case hardness results compounding this difficulty. In this work, we propose a deep reinforcement learning framework for quantum state preparation, capable of immediate inference of arbitrary stabilizer states at a fixed system size post a training phase. Our approach scales substantially beyond previous works by leveraging a novel reward function. In our experiments on stabilizer states up to nine qubits, our trained agent successfully prepares nearly all previously unseen states, despite being trained on less than $10^{-3}$\% of the state space -- demonstrating significant generalization to novel states. Benchmarking shows our model produces stabilizer circuits with size $60$\% that of existing algorithms, setting a new state of the art in circuit efficiency. Furthermore, we show that this performance advantage is consistent across states with varying entanglement content. We also analyze the rate of increase of entanglement entropy across the prepared circuit, obtaining insight into the quantum entanglement dynamics generated by our trained agent. Finally, we prove our agent generalizes to (almost) the entire space of stabilizer states.

Decision: reject
Reviewer scores: 5, 3, 6, 6, 3, 3

Reviews and discussion:
<thread>
creation date:
1730669817647

thread_type:
official_review

<thread_object>
creator: reviewer_crcx

<content>
summary:
This work utilizes a reinforcement learning PPO agent to solve the quantum state preparation problem. The key idea involves a potential based reward that is based on goalposts - earlier fidelity penalties are weighted less to encourage exploration and the preparation of partially suboptimal but intermediately necessary states. Training is conducted on stabilizer circuits, which can be completely classically simulated and only involve the H, CNOT and phase gate variants, the placement of which are used as actions. To reduce the set of possible end-target states the RL agent trains to reverse the state preparation, starting from the target state towards the |0> state. After training a one-shot model is evaluated on a study of random target states and a set of states relevant for quantum error correction is shown to yield circuits preparing valid states with fewer gates than analytical models. An analysis on the entangling entropy is given empirically and a upper bound is approximated via Monte-Carlo estimations.

strengths:
The ideas of this paper are presented understandable and mostly concise, the format is clean and the structure is logical. The background is explained very extensively. While I think the application of RL for stabilizer states is has fundamental issues, see below, I find the concept of reverse construction for a fixed end-state an interesting idea. I also think the extra effort for the study of entangling entropy and the attempt at a generalization bound, albeit quite loosely estimated, deserve praise and should be more common in empirical QRL papers.

weaknesses:
Nonetheless, I think this paper has three major issues:

- While fundamental explanations are generally good, the first 5 1/2 pages are simply explanation of fundamental QC / RL basics (excluding Figure 2). The first original section starts at the bottom of page 6, which is not efficient use of space. Much of Chapter 2 could be formulated much more concise or moved to an appendix. In general the information and result density is rather light for a A* publication with 10 pages. On the other hand the actual RL algorithm realization, and the training details should be made much more clear. From what I understand the observation is the whole stabilizer tableau (flattened probably?) and the actions are discrete choices over the set of all applicable gate-actions - or gate-target pairs for two-qubit gates. I am not clear if the addition of g_i in L343 is included somehow, or how often the actions are chosen. Once for each qubit for every layer? Some of the training details like hyperparameters are in the appendix, but are critical information for a self-contained paper and should be moved to the main-paper. (E.g. the training episodes, the max-T curriculum for increasing n, the fidelity relaxation and the fact that the otherwise-case of the reward was not even used as described in the paper.)
- The goalpost potential-reward as part of the contribution feels quite over-engineered and I’d argue this work is lacking a critical evaluation in the RL context. Similarly, while the reverse generation thought is interesting, without a comparison to the forward-generation training (zero to target) I find it hard to judge any significance. A study of the generation with a simpler fidelity based reward as is found in the literature (even the ones cited here, Kölle et al., Gabor et al., Wu et al.) would have helped, as would a test with a simpler, fixed step-based regularizer instead of the goalpost. A direct comparison of ‘backward’ vs. ‘forward’ generation training would also help support the claims made here.
- Finally, I think the RL application for stabilizer circuits in particular has fundamental issues. While they are indeed used in QEC, the point and the reason for that is, that they are completely and very efficiently classically simulated. (The abstract of Aaronson & Gottesman mentions thousands of qubits.) In other words, any efficiency gain w.r.t. circuit size, generalization ability or else, could also be classically brute-forced. Without application to circuits with unsimulatable quantum-effects (including rotation gates, e.g.) efficient one-shot generalization is an interesting insight, but lacking overall impact in terms of QC relevancy.
---
Minor Issues:
- L256 Typo: parameters: [a] the target …
- Inconsistent use of ‘Fig.’ and ‘Figure’ references.

questions:
- This max { } function notation is confusing to me, L323 (e.g.) uses max{a : b}, L324 uses max{a | b}? How is this supposed to be read and understood?
- Why was n=9 picked as the benchmark and not something higher, since stabilizer circuits can be so efficiently computed? L517 mentions that the he 9-qubit agent was not trained to convergence, why is that?
- How were the observation and action choices implemented in detail?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730707782167

thread_type:
official_review

<thread_object>
creator: reviewer_gvx9

<content>
summary:
The manuscript presents a deep reinforcement learning approach for quantum state preparation (QSP) . The authors design a novel reward function with guarantees that significantly scales beyond previous works, allowing for generalization to unseen states. They demonstrate their method on stabilizer states up to nine qubits, achieving state-of-the-art performance by preparing target states with varying degrees of entanglement.

strengths:
- This manuscript introduces a novel novel reward function with provable guarantees, which adds a level of theoretical robustness to the approach.
- The proposed method's success on stabilizer states, which are crucial for quantum error correction and other quantum information processes, suggests its applicability.
- The paper demonstrates that their method generates circuits that are up to 60% shorter than those produced by existing algorithms in some cases, which is an improvement in efficiency.

weaknesses:
- I believe the term "zero-shot" in the title of this paper is misleading.Zero-shot learning (ZSL) is a problem setup in deep learning where, at test time, a learner observes samples from classes which were **not** observed during training, and needs to predict the class that they belong to. However, in the context of this paper, training is conducted on a set $\mathcal S$ of $n$-qubit states. From my perspective, this setup aligns more with traditional supervised learning rather than what is commonly referred to as "zero-shot" learning.
- The authors claim that "The key idea that allows us to obtain zero-shot inference is the following: a sequence of actions starting in state |ψ⟩ leading to the state |0⟩ yields the inverse of a circuit preparing |ψ⟩ starting from |0⟩". However, I believe this concept is not original to this paper. For example, a similar idea was proposed in [1]. Yet, the authors have failed to cite it.
- I doubt the authors overstated their results. It's important to distinguish between the preparation of arbitrary quantum states and the preparation of arbitrary stabilizer states. The latter is a more specialized and arguably less challenging problem compared to the general case of arbitrary state preparation. Stabilizer states, due to their specific properties and structure, may be more amenable to efficient preparation methods, which might not directly translate to the broader and more complex task of preparing any arbitrary quantum state.

[1] Huang, Hsin-Yuan, et al. "Learning shallow quantum circuits." *Proceedings of the 56th Annual ACM Symposium on Theory of Computing*. 2024.

questions:
Please see the "Weaknesses" above.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730712985713

thread_type:
official_review

<thread_object>
creator: reviewer_hnhq

<content>
summary:
This manuscript studied the problem of quantum state preparation (QSP). While theoretical lower bounds for the required operations have been established, efficiently approximating quantum states on near-term quantum computers remains an open question. The authors proposed a reinforcement learning (RL) algorithm for QSP, focusing on the preparation of stabilizer states— a set of states of significant practical interest—using Clifford gates. The authors introduced a unified target state by replacing $|\psi\rangle$ with $|0\rangle$ and approximating backwards, referred as "Zero-shot" in the manuscript. Additionally, the authors proposed a moving-goalpost reward (MGR) function that aligns the maximum cumulative reward with the highest final fidelity. On the experimental side, the authors conducted experiments with up to 9 qubits, with and without connectivity restrictions, achieving better performance in terms of gate count compared to referenced algorithms. Theoretically, they proved a loose lower bound on the probability of generalization success.

strengths:
1. The problem studied is highly relevant to the quantum computing community. The manuscript presented a RL algorithm to efficiently prepare stabilizer state, which is valuable for future research.

2.  The numerical results is promising compared to conventional algiorthm, even with connectivity constraints, showing the potential utility on near-term quantum computers.

weaknesses:
1. The nanuscript lack discussing about relevant literatures. The "Zero-shot" trick is not new; amany existing works have adopted it in reinforcement learning for quantum circuit synthesis and compiling, with QSP being only a subset [1,2]. 

[1] Zhang, Yuan-Hang, et al. "Topological quantum compiling with reinforcement learning." Physical Review Letters 125.17 (2020): 170501.
[2] Qiuhao, Chen, et al. "Efficient and practical quantum compiler towards multi-qubit systems with deep reinforcement learning." Quantum Science and Technology (2024).

2. The claim in the abstract that "To our knowledge, this is the first work to prepare arbitrary stabilizer states on more than two qubits without re-training" is questionable. A recent work [3] also proposed a reinforcement learning algorithm for Clifford circuit synthesis.

[3] Kremer, David, et al. "Practical and efficient quantum circuit synthesis and transpiling with Reinforcement Learning." arXiv preprint arXiv:2405.13196 (2024).

3. The authors established a loose lower bound for the probability of generalization success (as indicated in the last equation of the main text). In Table 2, the probability decreases with qubit size, raising concerns about the generalizability of the proposed algorithm. Additionally, in Line 517, the authors stated, "We did not train the 9-qubit agent to convergence...," which raises further concerns about the trainability of the proposed RL algorithm.

questions:
1. The term "Zero-shot" is ambiguous. Typically, "zero-shot" refers to training without class labels. In the manuscript, the authors seem to be referring to the unified target trick, which should be clarified.

2. In the introduction, the authors claim that their RL method is applicable to "arbitrary states." However, the states studied are specifically stabilizer states. The authors should either correct this claim or provide results for additional categories of states.

3. How are the experimental results related to the theoretical lower bound of generalization? The authors should consider displaying the success probability achieved during their experiments.

4, What are the definitions of $q$ and $N$ in Lines 517-518? Additionally, the last equation should be numbered for clarity.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730727094516

thread_type:
official_review

<thread_object>
creator: reviewer_3scv

<content>
summary:
This study presents a zero-shot reinforcement learning (RL) algorithm that incorporates a novel reward function specifically designed for quantum state preparation (QSP). The authors concentrate on stabilizer state preparation tasks involving up to 9 qubits. Notably, the proposed algorithm demonstrates strong generalization capabilities on unseen data, utilizing a training dataset that is 0.001% smaller than the total state space. Furthermore, the quantum circuits generated by this algorithm achieve a circuit size reduction of over 60% compared to existing methods.

strengths:
1. The authors introduce a novel reward function that is inspired by the intrinsic nature of the QSP problem, integrating it into the RL framework to enhance performance in terms of circuit size reduction compared to existing algorithms.
2. The proposed algorithm demonstrates zero-shot generalization, significantly decreasing computational complexity compared to traditional RL algorithms, which require retraining for each new target state.
3. The authors provide a theoretical analysis of the generalization capabilities of the proposed algorithm.

weaknesses:
1. As someone who is not well-versed in RL algorithms, I found it somewhat challenging to follow the overall algorithms framework. The authors present the general RL algorithms separately before introducing their specific formulation for QSP tasks along with the novel reward function. It would enhance clarity if the authors summarized the algorithm in a table format.
2. The authors primarily use the total number of gates in the circuit, including both single-qubit and two-qubit gates, as the main evaluation metric for circuit size. However, two-qubit gates are more challenging to implement in practical devices and are more susceptible to noise. It would be beneficial to separately compare the counts of single-qubit and two-qubit gates in the circuits generated by the proposed algorithm against those produced by existing algorithms.
3. Typos: Line 256: "a the target".

questions:
The questions are included in Weaknesses.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730814822613

thread_type:
official_review

<thread_object>
creator: reviewer_5w9q

<content>
summary:
In this paper, the authors employ reinforcement learning (RL) to the task of quantum state preparation (QSP) on physically relevant states. Additionally, their method enables zero-shot preparation for any target state when the system size is fixed. 
Numerical results for systems with up to $9$ qubits show that their method generates circuits up to $60$% shorter than the baseline methods evaluated in the study.

strengths:
The paper does a reasonable job introducing the quantum state preparation task and, from the presented difficulties of the domain, manages to derive its adaptations to standard RL employed to tackle this domain. The adaptations presented are sound and straight-forward. The results give a nice guideline for future research in that direction.

weaknesses:
- The presentation of the paper is somewhat convoluted and challenging to follow; clearer descriptions would be helpful.
  - The contributions could be better highlighted, such as the introduction of the moving goalpost reward (MGR) for the QSP task.
  - The methodology and results, presented through plots, are introduced quite late in the paper. Providing clearer scenario descriptions of the experiments would improve readability.
  - A clearer contrast with related work would be helpful.
  - Figure legends could be enlarged for improved readability.
  - To aid understanding of the MGR function, a figure illustrating its workings would be valuable.

- The paper appears to misinterpret the QSP definition by assuming access to the target state $|\psi\rangle$. Typically, in QSP, one has access only to a succinct (classical) description of the target state, not the circuit that implements it. However, in this approach, the authors start each episode with $|\psi\rangle$ as the initial state and then find a circuit to prepare $|\mathbf{0}\rangle^{\otimes n}$ for zero-shot state preparation. This approach may not be physically feasible in a lab; if the circuit for $|\psi\rangle$ is known, the rationale for using RL (or any method) to find a circuit that prepares $|\mathbf{0}\rangle^{\otimes n}$ and then inverts it is unclear. It might be argued that the given $|\psi\rangle$ circuit is too deep or not NISQ-friendly, necessitating a more compact architecture, though this would still be costly in a real lab setting.

- Insufficient comparison with relevant state-of-the-art methods for the quantum state preparation task:
  - RL methods: [1-4]
  - ML methods: [5,6]
  - SAT-based methods: [7,8]

- The algorithm's performance under common noise models, such as state preparation and measurement (SPAM), bit flip, phase flip, depolarizing noise, and thermal relaxation is not evaluated. These noise sources are standard in quantum computing.

- Several claims in the paper lack theoretical or numerical support. Ablation studies providing numerical evidence would help support these claims:
  - Line 306-308: "However, the cumulative reward obtained ..... terminating the episode."
  - Line 803-806: "Finally, for the linear-connectivity agents, .... faster with this term."

- The paper lacks open-source code and a reproducibility statement.


[1] Fosel, Thomas, et al. "Quantum circuit optimization with deep reinforcement learning", arXiv:2103.07585 (2021)\
[2] Patel, Yash J., et al. "Curriculum reinforcement learning for quantum architecture search under hardware errors",  ICLR (2024)\
[3] Zen, Remmy, et al. "Quantum Circuit Discovery for Fault-Tolerant Logical State Preparation with Reinforcement Learning", arXiv:2402.17761 (2024)\
[4] Kremer, David, et al. "Practical and efficient quantum circuit synthesis and transpiling with Reinforcement
Learning", arXiv:2405.13196 (2024)\
[5] Wang, Hanrui, et al. "Quantumnas: Noise-adaptive search for robust quantum circuits." 2022 IEEE International Symposium on High-Performance Computer Architecture (HPCA)\
[6] Wu, Wenjie, et al. "Quantumdarts: differentiable quantum architecture search for variational quantum algorithms", ICML (2023)\
[7] Peham, Tom, et al. "Depth-Optimal Synthesis of Clifford Circuits with SAT Solvers", QCE (2023)\
[8] Schneider, Sarah, et al. "A SAT Encoding for Optimal Clifford Circuit Synthesis", ASPDAC (2023)

questions:
- see above.

- HSH gate in the gateset:
  - What is the rationale behind including HSH gate in the gateset $G$ (action space)? Because in principle the agent should learn the symmetry on $S$ w.r.t. to $X$ and $Z$ operations by itself. I believe this inductive bias is introduced during the learning process while the whole point of using RL is not to have any kind of bias.

  - Also HSH gate has three single qubit gates. In Table 1, do you count it as one gate for all the methods? Do you use the same gateset for all the methods for a fair comparison?

  - Can you provide an ablation study with and without HSH gate to see how the number of gates differ between these two settings for your method?

- Why is there no circuit size numbers for $9$-qubit DRL (local gates) in Table 1?

- In Fig. 3(d), $n=9$, can you please explain the behaviour of the circuit size for depth $t=1, 2$ signficantly larger for your method compared to other two baselines?

- How does the algorithm perform for the QSP task when the target states which are complex and physically relevant but are not classically simulable?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730973054511

thread_type:
official_review

<thread_object>
creator: reviewer_pcnz

<content>
summary:
The paper proposes a zero-shot quantum state preparation approach. RL has been used.
A moving goal post reward function has been designed. Training is performed on less than 0.0001% of the
state space and then generalized.

strengths:
1. Upto 60% shorter circuit depths are obtained.
2. Generalized to more than 2 qubits without re-training.

weaknesses:
1. Proof of Prop. 2 of convergence is not completely convincing. Conditions 
under which it is valid need rigorous statement.
2. No guarantee that the potential function will be strictly convex. No consideration
is given to this.
3. It is not clear why retraining is not required. How general is this is not stated.

questions:
1. Why retraining is not required needs further explanation and study?
2. Why was the 9 qubit agent not trained till convergence?
</content>
</thread_object>

</thread>
<thread>
creation date:
1732676154379

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General Response by Authors (1/2)

comment:
We are grateful to the reviewers for their time and insightful comments; and to the area, senior area, and program chairs for their service.
We are glad to note that the reviewers appreciated several aspects of this work including **efficiency** (PcnZ, 3SCv, hnhQ, GvX9), the novel **reward** function (3SCv, GvX9), **theoretical analysis** (3SCv, GvX9, crcX) and **relevance** to the quantum community (hnhQ, GvX9).

We have tried our best to address all the questions and concerns raised by the reviewers including the minor ones. We summarize below how we have addressed the key concerns here in addition to the point-by-point reviews to the reviewers.

1. **Stabilizer and non-stabilizer states**. This comment is in regards to some referees' concerns about the general applicability of the framework and the reason for implementation on Clifford circuits. We apologize for the possible lack of clear explanation in the manuscript, and clarify this point in detail here as well as in the paper. We would like to stress the fact that the classical simulability of stabilizer states does not imply triviality of state preparation. Moreover, there is a plethora of new physics in quantum dynamics that has seen light in the recent years owing to the classical simulability of Clifford circuits, discoveries which would not have been possible without this numerical convenience (for instance see [A-D]). Yet, many of those aspects generalize to non-Clifford dynamics and circuits. The takeaway is that by leveraging the classical simulation possibility of a class of quantum states, which form a 3-design and capture many properties of universal quantum dynamics, we are able to gain enhanced insight into the possibilities of leveraging reinforcement learning methods for quantum state preparation. Indeed, the story is not complete and this method needs to be studied in the full generality of the Hilbert space. Thus, looking forward, we provide important future directions to achieve this generality while concluding the paper.

[A] Li, Yaodong, Xiao Chen, and Matthew PA Fisher. "Quantum Zeno effect and the many-body entanglement transition." Physical Review B 98.20 (2018): 205136.

[B] Li, Yaodong, Xiao Chen, and Matthew PA Fisher. "Measurement-driven entanglement transition in hybrid quantum circuits." Physical Review B 100.13 (2019): 134306.

[C] "Measurement-induced entanglement and teleportation on a noisy quantum processor." Nature 622, no. 7983 (2023): 481-486.

[D] Bao, Yimu, Maxwell Block, and Ehud Altman. "Finite-time teleportation phase transition in random quantum circuits." Physical Review Letters 132.3 (2024): 030401.

2. **The QSP setting**. In response to comments regarding the information the agent has access to, we would like to clarify that we are indeed working in the typical QSP setting, but with relaxed conditions for this work where we are allowed the full (classical) stabilizer tableau of the state to the algorithm. This is already relevant when one wants to implement a circuit conditioned on connectivity and gate-sets to prepare a state one knows about in advance in the lab. The algorithm provides the circuit that the experimentalist would implement. Further, we would like to point out that the inversion procedure is only for training the algorithm; after the fact, we only need access to the target state to find a circuit to generate it, reaffirming the QSP paradigm. Moreover, while one would ideally like the query model to be, say, local Pauli measurements on the target state, we have a relaxed query model wherein the full tableau is available as a query. This work indeed establishes the fact that training very general agents using reinforcement learning is possible. An immediate future work that we are pursuing is changing the query model and hence bringing it closer the setting of preparing states given succinct classical descriptions.

3. **Zero-shot terminology**. We agree that usage of "zero-shot" without some additional context could be confusing. Indeed, multiple zero-shot settings have received attention in the community such as predicting images from a class without having been presented any examples of that class during training, or more recently, in the context of so called 'zero-shot prompting' in LLMs.  Here, we use the term "zero-shot" to convey the successful preparation of a state without ever having access to it. We have clarified our usage of this definition of zero-shot in the introduction of the revised manuscript. To quote, "*An agent that does not need re-training to prepare unseen
states will be called zero-shot in this work, to emphasize successful generalization to states not seen during training.*"
</content>
</thread_object>

</thread>


Metareview:
metareview:
The authors introduce a deep reinforcement learning (RL) framework for quantum state preparation (QSP), specifically targeting stabilizer states. By employing a novel moving-goalpost reward (MGR) function and a reverse-generation approach, the method demonstrates zero-shot generalization, enabling the preparation of unseen stabilizer states without retraining. Simulations conducted on systems with up to nine qubits highlight the effectiveness of the proposed method.

While the integration of reinforcement learning with stabilizer state preparation is noteworthy, the submission is hindered by several critical issues:
1. The study mostly focuses on stabilizer states, limiting its broader applicability. 
2. The theoretical guarantees for generalization success are considered loose and lack rigorous validation.
3. The absence of evaluations under realistic quantum noise models, such as depolarizing or SPAM noise, undermines the practical relevance.
4. The presentation was critiqued for devoting excessive space to background material while omitting critical details of the RL implementation and methodology.