Quantum Filtering and Stabilization of Dissipative Quantum Systems via Augmented Neural Ordinary Differential Equations

Technical Summary: Quantum Filtering and Stabilization of Dissipative Quantum Systems via Augmented Neural Ordinary Differential Equations

Problem Statement
The estimation and control of open quantum systems present a fundamental challenge when the system Hamiltonian or noise model is not fully known. In realistic experimental settings, quantum trajectories are only partially accessible through continuous weak measurements, while the underlying dynamics are influenced by hidden, time-dependent, and environment-dependent dissipative processes (non-Markovian effects). Traditional control strategies, such as optimal control or Lyapunov-based methods, rely heavily on complete knowledge of the system's Hamiltonian and dissipative dynamics, rendering them less suitable for scenarios with incomplete information. Furthermore, standard machine learning approaches often lack physical interpretability and fail to explicitly recover unknown parameters, which are crucial for designing feedback control laws. The core difficulty lies in reconstructing the full quantum state (Bloch vector) and inferring hidden dissipative rates simultaneously from incomplete measurement data without an explicit physical model.

Methodology: The AQNODE Framework
The authors introduce the Augmented Quantum Neural Ordinary Differential Equation (AQNODE) framework, a data-driven approach designed to learn quantum trajectories and dissipation parameters directly from partial continuous measurement data.

1. Augmented State Representation: Unlike standard Neural ODEs that learn only observable trajectories, AQNODE defines an augmented state vector $Y_{aug}(t) = [x(t), y(t), z(t), \Delta(t), \gamma(t)]^T$. This includes the observable Bloch vector components ($x, y, z$) and the hidden, time-dependent dissipative parameters: the diffusion rate $\Delta(t)$ and the dissipation rate $\gamma(t)$. This augmentation transforms the problem from pure trajectory forecasting into a joint system identification and filtering problem.
2. Architecture: The framework consists of three main components:
   • Encoder: A neural network that maps partial measurement records ($dY(t)$), control inputs ($u(t)$), and initial conditions into a latent representation $h(t)$.
   • Neural ODE: The latent state evolves continuously in time according to a learned vector field $\dot{h}(t) = \text{MLP}_\theta([h(t), dY(t), u(t), t])$. This allows the model to capture the continuous-time dynamics and temporal smoothness inherent in physical systems.
   • Decoder: A neural network that maps the evolved latent state back to the augmented observable space, outputting estimates for both the qubit state and the hidden dissipative parameters.
3. Training and Loss: The model is trained in a supervised manner using simulated data generated from a two-level time-convolutionless (TCL) master equation with an Ohmic Lorentz-Drude reservoir. The total loss function combines the mean squared error (MSE) between predicted and ground-truth state trajectories ($L_s$) and parameter trajectories ($L_p$), weighted by hyperparameters $\kappa$ and $\beta$. The adjoint sensitivity method is used for memory-efficient gradient computation through the ODE solver.
4. Feedback Control Integration: The inferred state and parameters are utilized for real-time feedback control. The authors implement two strategies:
   • Proportional-Derivative (PD) Control: Computes control fields based on the error between the estimated state and the target state.
   • Time-Varying Linear-Quadratic Regulator (LQR): Dynamically adjusts feedback gains by solving a differential Riccati equation to minimize a cost function involving state deviation and control energy.

Key Contributions

• Simultaneous Filtering and Identification: AQNODE enables the simultaneous reconstruction of the qubit state and the inference of hidden, time-dependent dissipative rates from partial weak measurement data, without requiring an explicit physical model during inference.
• Physical Consistency and Interpretability: By embedding the system in a latent space evolved via neural ODEs, the framework ensures temporal smoothness and physical consistency, explicitly returning interpretable parameters ($\Delta(t), \gamma(t)$) rather than just a black-box trajectory fit.
• Data-Driven Control: The framework integrates learned quantum filtering with classical control strategies (PD and LQR), demonstrating that accurate state reconstruction from partial data is sufficient to steer dissipative quantum systems toward target states in real time.
• Robustness and Generalization: The approach is validated in both within-distribution (WD) and out-of-distribution (OOD) regimes, including scenarios with significant variations in environmental parameters ($\alpha, r$) and system parameters ($M, \omega_0$), as well as perturbations in initial conditions.

Results
Extensive numerical simulations on a single driven qubit interacting with a non-Markovian environment demonstrate the efficacy of AQNODE:

• State and Parameter Estimation: In the WD regime, the model achieves low Mean Squared Error (MSE) for both state components ($< 10^{-4}$) and parameters. In the OOD regime, where parameters deviate significantly from training distributions, the model maintains reasonable accuracy (MSE $< 10^{-3}$ for states), demonstrating strong generalization capabilities.
• Control Performance: When used for feedback control to steer the qubit to a target state $|0\rangle$, AQNODE-based controllers (both PD and LQR) successfully stabilize the system.
  • LQR Control: Demonstrated superior performance over PD control in terms of convergence speed, final-state deviation, and fidelity ($\geq 0.93$) even under OOD conditions. However, this improved control performance comes at the cost of higher energy expenditure; Table III indicates that LQR requires significantly larger control energy (e.g., 5624.1 vs. 2310.9 for PD in the WD regime) compared to PD control.
  • Purity and Coherence: The model accurately predicts the evolution of quantum purity and coherence, capturing the transient oscillations and steady-state behaviors induced by control fields and environmental dissipation.
• Robustness to Perturbations: The framework exhibits robustness against perturbations in the initial quantum state, with the model correcting its predictions over time to align with the ground truth, mimicking the behavior of a quantum filter.

Significance and Claims
The paper claims that AQNODE establishes a scalable, differentiable, and potentially experimentally compatible framework for the real-time modeling and control of dissipative quantum systems. By bridging machine learning with quantum mechanics, the work offers a solution to the challenge of controlling quantum systems without full knowledge of their Hamiltonians or noise models. The authors posit that this data-driven paradigm unites conventional quantum control techniques with machine learning, providing a reliable tool for quantum state estimation and stabilization. They suggest that while the current work focuses on a two-level system, the framework is extendable to higher-dimensional and complex quantum systems, potentially aiding in quantum error correction and NISQ-era simulations. The significance lies in its ability to function as a "learned quantum observer," enabling adaptive protocols in environments where exact parameter distributions are unknown.