Second-Order Adjoint Method for Quantum Optimal Control

TL;DR

The paper addresses efficient quantum optimal control of molecular systems governed by the time-dependent Schrödinger equation by deriving and implementing a second-order adjoint method that provides exact gradients and Hessians. The authors build on a discretize-then-optimize framework, derive forward state and backward adjoint (and second-order adjoint) systems, and implement them on GPUs using JAX to achieve near-linear scaling with the number of time steps. Empirical results on small molecules (e.g., H$_2$, HeH$^+$) in STO-3G and 6-31G bases show that Hessian-enabled Newton optimization converges faster (fewer iterations and wall time) than first-order methods while delivering comparable solution quality and near-zero final-state errors. The work demonstrates the practicality of exact Hessians for quantum control with arbitrary control parameterizations and provides open-source code for broader adoption.

Abstract

We derive and implement a second-order adjoint method to compute exact gradients and Hessians for a prototypical quantum optimal control problem, that of solving for the minimal energy applied electric field that drives a molecule from a given initial state to a desired target state. For small to moderately sized systems, we demonstrate a vectorized GPU implementation of a second-order adjoint method that computes both Hessians and gradients with wall times only marginally more than those required to compute gradients via commonly used first-order adjoint methods. Pairing our second-order adjoint method with a trust region optimizer (a type of Newton method), we show that it outperforms a first-order method, requiring significantly fewer iterations and wall time to find optimal controls for four molecular systems. Our derivation of the second-order adjoint method allows for arbitrary parameterizations of the controls.

Figures (2)

• Figure 1: We give a log-log plot of wall clock time results for Algorithms \ref{['alg:firstorder']} and \ref{['alg:secondorder']}. Here $N$ represents problem size, while $J$ is the number of time steps required to go from $t=0$ to $t=T$. For each fixed choice of parameters $(N,J)$, we ran each algorithm $1000$ times; we plot the mean (circular dot) and error bar (plus/minus twice the standard deviation) of these runs. The results show that both algorithms' scaling as a function of $J$ is close to linear---see Section \ref{['sect:scaling']} in the main text for details. The upshot: for the quantum optimal control problem we studied, Hessians and gradients can be computed at similar cost to gradients alone.
• Figure 2: From the pool of $1000$maximal model (\ref{['eqn:maximal']}) solutions whose statistics were reported in Tables \ref{['tab:rawocr']}-\ref{['tab:ratios']}, the plotted solutions achieved the lowest cost. For each molecular system (choice of molecule plus basis set), we plot both the optimal field strength $f(t)$ and the magnitudes $|a_j(t)|$ of the controlled trajectory. Here $J = 200$ and $\Delta t = 0.1$. Gray components of the controlled trajectory have initial and target values of $0$; colored components' initial values differ from their targets.