We apply a machine-learning-enhanced greybox framework to a quantum optimal control protocol for open quantum systems. Combining a whitebox physical model with a neural-network blackbox trained on synthetic data, the method captures non-Markovian noise efects and achieves gate fidelities above 90% under Random Telegraph and Ornstein-Uhlenbeck noise. Critical issues of the approach are discussed. This work showcases an attention-based machine-learning-enhanced greybox framework for quantum optimal control, designed to improve the manipulation of open quantum systems subject to complex noise [1, 2]. Quantum control [3] is essential for quantum technologies such as computation, communication, and sensing, yet achieving robust control remains challenging when the system interacts with non-Markovian environments that are hard or even impossible to characterize. The proposed greybox model combines: • a whitebox component that captures the analytically tractable portion of the system dynamics using known physical principles; • a blackbox component implemented via neural networks, trained to learn the unmodelled efects of the environment from data.
We consider a single qubit subject to classical dephasing noise along the -axis and driven by external control fields. In the interaction picture, the dynamics is described by the time-dependent Hamiltonian () = ctrl() + () , where is the coupling strength between the qubit and the noise, and () is a classical stochastic process modeling dephasing noise [5]. Specifically, in this work we consider () to be either a Random Telegraph Noise (RTN) process or an Ornstein-Uhlenbeck (OU) process.
The control Hamiltonian ctrl() implements a drive along the and -axes
ctrl() = () + () , each control field (), with ∈ {, }, consisting of 5 Gaussian-shaped pulses.
Two diferent stochastic processes were considered and compared, namely an RTN process and an OU process. They are characterised by their power spectrum (), which is the Fourier transform of the two-point correlation function ⟨ () ()⟩, and in both cases has a Lorentzian shape [4, 6] () ≈
4 4 2 + 2 (1) where is the switching rate for the RTN process and 1/ is the correlation time of the OU process. The two stochastic processes difer since the latter is Gaussian, while the former is not, and it is known that this has an impact on dynamic protocols of protection against noise, such as spin-echo [4, 7].
The proposed greybox model integrates analytical knowledge of the quantum system with a transformerbased neural network. This hybrid architecture includes two components: • A whitebox part, which enforces the known unitary dynamics of the driven qubit and the associated measurement process; • A blackbox neural network, trained to model the influence of the environment on the system’s evolution.
Inputs and Outputs The model takes as input the amplitudes of five Gaussian control pulses applied along each of the and axes, for a total of ten real parameters. The pulse widths and positions are ifxed. The output consists of six gate fidelities, each associated with a diferent target from a universal set of single-qubit gates.
Model Architecture The blackbox core is a lightweight transformer encoder. It processes the input pulse parameters and predicts a set of noise-related parameters that are fed into whitebox layers implementing: • Hamiltonian construction and time evolution based on discretized control fields; • Expectation value calculation over a tomographically complete set of initial states; • Process matrix reconstruction and fidelity estimation.
In addition, the model includes specialized output heads that refine the predicted expectation values before computing the final fidelities.
Training Strategy Only the blackbox layers contain trainable parameters. The network is trained using the Adam optimizer to minimize the mean squared error across the six predicted fidelities. Training is supervised and based on synthetic data generated by simulating noisy quantum dynamics. Whitebox constraints ensure physically consistent predictions throughout.
A schematic overview of the architecture is shown in Fig. 1.
Separate models were trained across varying values of the coupling strength . This provides insight on the efectiveness of the graybox approach as a function of the Markovianity of the quantum map describing time evolution under the efect of noise, which in this case is parametrized by the ratio / [8, 4].
RTN Case. The model showed low training and test MSE across all gates, with prediction errors increasing with but remaining in the 10− 2–10− 3 range, indicating robust generalisation. As an emulator in the optimal control pipeline, it enabled the design of control pulses achieving fidelities above 99% for the lowest and above 90% for the highest, with minor gate-dependent variations. OU Case. The model exhibited similar performance, with low and stable MSE values across all , confirming robustness to diferent noise types. Optimal control results mirrored those of the RTN case, with fidelities exceeding 99% at low and remaining above 90% even at stronger coupling. While fidelity declines under higher noise, the model continues to support efective pulse design; future improvements may benefit from larger datasets or more advanced strategies.
Our result validates the graybox approach, showing that the optimization framework we have chosen is very efective in suppressing efects of low-frequency noise ( / > 1), but less efective for noise yielding Markovian maps (/ < 1). Apparently, Gaussianity does not have an impact in this case, but we expect that the picture may change when considering 1/ noise [9, 4, 6, 10] resulting from a set of stochastic processes with diferent .
A natural development of this work is applying the method to two-qubit gates, addressing the efect of both time- and space-correlated noise [11, 12, 13]. Two major issues to be investigated are the scalability of the approach to larger quantum architectures and the ability to reproduce asymptotic results known from the theory of dynamical decoupling [14, 4, 12].
During the preparation of this work, the authors used ChatGPT in order to: perform grammar and spelling checks, paraphrase and reword. After using this tool, the authors reviewed and edited the content as needed and take full responsibility for the content of the publication.