Machine learning meets $\mathfrak{su}(n)$ Lie algebra: Enhancing quantum dynamics learning with exact trace conservation

Abstract

Machine learning (ML) has emerged as a promising tool for simulating quantum dissipative dynamics. However, existing methods often struggle to enforce key physical constraints, such as trace conservation, when modeling reduced density matrices (RDMs). While Physics-Informed Neural Networks (PINN) aim to address these challenges, they frequently fail to achieve full physical consistency. In this work, we introduce a novel approach that leverages the $\mathfrak{su}(n)$ Lie algebra to represent RDMs as a combination of an identity matrix and $n^2 - 1$ Hermitian, traceless, and orthogonal basis operators,where $n$ is the system's dimension. By learning only the coefficients associated with this basis, our framework inherently ensures exact trace conservation, as the traceless nature of the basis restricts the trace contribution solely to the identity matrix. This eliminates the need for explicit trace-preserving penalty terms in the loss function, simplifying optimization and improving learning efficiency. We validate our approach on two benchmark quantum systems: the spin-boson model and the Fenna-Matthews-Olson complex. By comparing the performance of four neural network (NN) architectures -- Purely Data-driven Physics-Uninformed Neural Networks (PUNN), $\mathfrak{su}(n)$ Lie algebra-based PUNN ($\mathfrak{su}(n)$-PUNN), traditional PINN, and $\mathfrak{su}(n)$ Lie algebra-based PINN ($\mathfrak{su}(n)$-PINN) -- we highlight the limitations of conventional methods and demonstrate the superior accuracy, robustness, and efficiency of our approach in learning quantum dissipative dynamics.