Gradient-based Quantum Control and Engineering with Adjoint Sensitivity

Controlling quantum systems lies at the heart of many quantum technologies, from quantum sensing to quantum information processing. Recently, there has been a growing interest in optimization quantum control using automatic differentiation and neural networks to achieve high fidelity quantum operations. Typically, these algorithms rely on storing computational graphs of mathematical operations to leverage backpropagation, which incurs large memory overhead due to said storage.

In this study, we present a new approach based on adjoint sensitivity, commonly employed in inverse-design for photonics and differential-equation-based machine learning to compute gradients over differential equation solvers, for quantum controls. We show that owing to the low memory usage and an ability to optimize over either continuous functions or pulse sequences, adjoint sensitivity outperforms many approaches such as GRAPE in searching for optimal control sequences. We demonstrate several examples in generating quantum states and gates of multi-qubit systems, as well as Hamiltonian engineering to cancel out qubit frequency fluctuations. Our optimization is performed using JAX and Diffrax, two open-source Python libraries that combine efficient adjoint sensitivity and high-performance computing.