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

Krishna Agaram, Siddhant Midha, Adrian Müller, Vikas Garg

Keywords: Quantum State Preparation, Deep Reinforcement Learning, Zero-shot Inference, Off-the-shelf Algorithms, Generalization

TL;DR: Reinforcement learning framework for zero-shot quantum state preparation.

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.

Primary Area: applications to physical sciences (physics, chemistry, biology, etc.)

Code Of Ethics: I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics.

Submission Guidelines: I certify that this submission complies with the submission instructions as described on https://iclr.cc/Conferences/2025/AuthorGuide.

Anonymous Url: I certify that there is no URL (e.g., github page) that could be used to find authors’ identity.

No Acknowledgement Section: I certify that there is no acknowledgement section in this submission for double blind review.

Submission Number: 13925