Kraus Constrained Sequence Learning For Quantum Trajectories from Continuous Measurement

TL;DR

This work proposes a Kraus-structured output layer that converts the hidden representation of a generic sequence backbone into a completely positive trace preserving (CPTP) quantum operation, yielding physically valid state updates by construction, across diverse backbones.

Abstract

Real-time reconstruction of conditional quantum states from continuous measurement records is a fundamental requirement for quantum feedback control, yet standard stochastic master equation (SME) solvers require exact model specification, known system parameters, and are sensitive to parameter mismatch. While neural sequence models can fit these stochastic dynamics, the unconstrained predictors can violate physicality such as positivity or trace constraints, leading to unstable rollouts and unphysical estimates. We propose a Kraus-structured output layer that converts the hidden representation of a generic sequence backbone into a completely positive trace preserving (CPTP) quantum operation, yielding physically valid state updates by construction. We instantiate this layer across diverse backbones, RNN, GRU, LSTM, TCN, ESN and Mamba; including Neural ODE as a comparative baseline, on stochastic trajectories characterized by parameter drift. Our evaluation reveals distinct trade-offs between gating mechanisms, linear recurrence, and global attention. Across all models, Kraus-LSTM achieves the strongest results, improving state estimation quality by 7% over its unconstrained counterpart while guaranteeing physically valid predictions in non-stationary regimes.

Figures (3)

• Figure 1: Comparative performance of Baseline (unconstrained) vs. Kraus-constrained models. Error bars represent one standard deviation. The Kraus Head provides a consistent fidelity lift for recurrent backbones, effectively acting as a physics-informed regularizer
• Figure 2: Visualization of the Switching Dataset. (Top) The normalized homodyne measurement record $y_t$ provided as input to the models. (Bottom) Ground-truth Bloch vector components showing the Phase 1 ($\sigma_x$-rotation) to Phase 2 ($\sigma_y$-rotation) transition at the vertical dashed line. The stochastic "jitter" in the Bloch components illustrates the measurement back-action (kicks) modeled by the SME.
• Figure 3: Layer-wise gradient norms for Kraus-RNN (left) and Kraus-GRU (right) computed from trained checkpoints. The Kraus-RNN exhibits severe gradient imbalance, with input_proj.weight nearly vanishing ($\sim\!10^{-2.5}$) while recurrent weights in later layers receive substantially larger gradients. In contrast, Kraus-GRU maintains comparatively uniform gradient flow across all layers.