Fast surrogate modelling of EIT in atomic quantum systems using LSTM neural networks

TL;DR

The paper tackles the computational bottleneck of simulating density-matrix dynamics in Rydberg-EIT RF sensing by training a compact LSTM-based surrogate to reproduce steady-state spectra with high fidelity. The surrogate, fed by the two physical inputs $\Omega_p$ and $\Omega_{\mathrm{RF}}$, generates a 300-point transmission spectrum in milliseconds, achieving about a 5000× speed-up over full Optical Bloch Equation solvers. It attains near-unity agreement in both resonant and off-resonant regimes ($R^2$ near 1, RMSE on the order of $10^{-3}$–$10^{-2}$) and demonstrates robust generalization to unseen parameter sets. This enables real-time signal processing and optimization for deployable quantum sensors on standard CPUs, with potential extensions to experimental data and broader quantum technologies.

Abstract

Simulations of optical quantum systems are essential for the development of quantum technologies. However, these simulations are often computationally intensive, especially when repeated evaluations are required for data fitting, parameter estimation, or real-time feedback. To address this challenge, we develop a Long Short-Term Memory neural network capable of replicating the output of these simulations with high accuracy and significantly reduced computational cost. Once trained, the surrogate model produces spectra in milliseconds, providing a speed-up of 5000x relative to traditional numerical solvers. We focus on applying this technique to Doppler-broadened Electromagnetically Induced Transparency in a ladder-type scheme for Rydberg-based sensing, achieving near-unity agreement with the physics solver for resonant and off-resonant regimes. We demonstrate the effectiveness of the LSTM model on this representative optical quantum system, establishing it as a surrogate tool capable of supporting real-time signal processing and feedback-based optimisation.

Figures (5)

• Figure 1: (a) Basic experimental arrangement for a Rydberg sensor. Two counter-propagating and counter-aligned (overlapping) lasers meet at a vapour cell, reflected by the corresponding dichroic mirrors (DM). The cell is simultaneously illuminated by an external RF field, altering the EIT condition. The probe laser light is collected by a photodiode, thus providing an optical read-out of the RF strength according to the magnitude of the AT splitting. (b) Atomic levels of a typical alkali-metal Rydberg RF sensor. Ladder-type configuration with probe ($\Omega_p$), control ($\Omega_c$), and RF ($\Omega_{\mathrm{RF}}$) couplings. Key detunings, decoherence rates, and experimental parameters (atomic number density, cell length, temperature) are included in the simulation.
• Figure 2: Representative spectra from the physics-based solver using $\ce{^{87}Rb}$ atomic parameterssteck_rubidium_2003sibalic_arc_2017. Control Rabi frequency: 6 MHz. Temperature: 288 K. Cell length: 10 cm. (a) Variation with probe Rabi frequency, illustrating saturation broadening effects. (b) Variation with RF Rabi frequency, showing Autler–Townes splitting.
• Figure 3: Architecture of the surrogate LSTM model. Inputs are probe and RF Rabi frequencies. The architecture includes three dense layers (32, 64, 128 nodes), an LSTM block (300 outputs), and a final dense layer (300 linear outputs). Total model size is $\sim$7 MB.
• Figure 4: Results of surrogate model tests with $\ce{^{87}Rb}$ atomic parameterssteck_rubidium_2003sibalic_arc_2017. The fixed parameters used were: $\Omega_c = 6$ MHz and $T=288$ K with a cell length of 10 cm. (a) Testing on four of the parameter sets not included in the training or validation data. Solid lines: physics solver; dashed lines: surrogate predictions. The surrogate successfully captures amplitude variations, peak widths, separations, and offsets, while preserving smooth spectral shapes. (b) Heatmap of $R^2$ performance across the parameter grid. The surrogate model achieves $R^2>0.999$ across most of the domain, with a slight performance drop at the corners due to reduced sampling density. The model produces physically consistent outputs across the entire space.
• Figure 5: Off-resonant regime comparison using 87Rb atomic parameterssteck_rubidium_2003sibalic_arc_2017. For the training and validation sets, the fixed parameters used were: $\Omega_c = 2$ MHz, $\Omega_p=1.5$ MHz and $T=283$ K with a cell length of 10 cm. (a) Example spectra comparing solver outputs (solid) and surrogate predictions (dashed). (b) Heatmap of $R^2$ performance across RF Rabi frequency and detuning. The surrogate generalizes well, maintaining high accuracy despite the increased spectral variance introduced by detuning.