Machine Learning-aided Optimal Control of a noisy qubit

Abstract: We apply a graybox machine-learning framework to model and control a qubit undergoing Markovian and non-Markovian dynamics from environmental noise. The approach combines physics-informed equations with a lightweight transformer neural network based on the self-attention mechanism. The model is trained on simulated data and learns an effective operator that predicts observables accurately, even in the presence of memory effects. We benchmark both non-Gaussian random-telegraph noise and Gaussian Ornstein-Uhlenbeck noise and achieve low prediction errors even in challenging noise coupling regimes. Using the model as a dynamics emulator, we perform gradient-based optimal control to identify pulse sequences implementing a universal set of single-qubit gates, achieving fidelities above 99% for the lowest considered value of the coupling and remaining above 90% for the highest.