arXiv subjects

Jiaji Zhang

Publications and source records attributed to Jiaji Zhang.

At least 19 recordsLinked to original sources

arXiv2026

Reinforcement Learning Assisted Quantum Simulation of Many-Body Excited States and Real-Time Dynamics

The computation of electronic excited states and real-time quantum dynamics of many-fermion systems is among the most promising applications of near-term quantum computing. In this work, we generalize the reinforcement learning contracted quantum eigensolver (RL-CQE), previously developed for ground-state problems, to electronic excited states and real-time quantum dynamics, in which a deep Q-network agent adaptively selects the two-body operators at each iteration, yielding more compact ans\"{a}tze and improved robustness with respect to critical hyperparameters. A key feature of the algorithm is a scalable state representation based on the ACSE residuals, whose dimension grows with the one-particle basis but remains independent of the number of targeted excited states. We also verify the equivalence of sign-free qubit operators in the excited-state setting, extending a result previously established for ground-state problems. Our RL-CQE for time evolution derives from a constant-scaling ansatz that represents the wave function with a fixed number of unitary transformations independent of simulation time $t$, enabled by the shared unitary structure of the purified ensemble treatment of excited states. Benchmarks on chemical systems demonstrate chemical accuracy with minimal operator counts across a range of bond lengths.

quant-ph↗

arXiv2026

Shiva-DiT: Residual-Based Differentiable Top-$k$ Selection for Efficient Diffusion Transformers

Diffusion Transformers (DiTs) are costly at high resolution because self-attention scales quadratically with token sequence length. Existing pruning methods do not jointly provide end-to-end learnability, low training overhead, and deterministic token counts for predictable token-dependent computation. We propose Shiva-DiT, based on Residual-Based Differentiable Top-k Selection. Its forward pass executes hard top-k selection, while a residual-aware straight-through estimator propagates gradients to both token scores and the budget k without evaluating a second backbone path. A Context-Aware Router and Adaptive Ratio Policy learn layer- and timestep-dependent retention schedules under a target average budget. Experiments on SD3-Medium, Flux.1-dev, and PixArt-{\Sigma} show consistent reductions in FLOPs and measured latency. On SD3-Medium, Shiva-DiT provides four fidelity-latency operating points and reaches a 1.54x wall-clock speedup with competitive fidelity.

cs.LG↗

arXiv2025

Toward Robust and Efficient ML-Based GPU Caching for Modern Inference

In modern GPU inference, cache efficiency remains a major bottleneck, and heuristic policies such as \textsc{LRU} can perform far worse than the offline optimum. Existing learning-based caching systems improve hit rates mainly through predictor design, but often follow learned predictions blindly, making performance unreliable when predictions are inaccurate. In contrast, emerging learning-augmented caching algorithms~\cite{pmlr-v80-lykouris18a,mitzenmacher2022algorithms} provide performance guarantees by carefully integrating predictions into caching policies, achieving both \emph{consistency} (near-optimality under perfect predictions) and \emph{robustness} (bounded worst-case performance under prediction errors). However, deployment remains challenging. A practical algorithm should satisfy strict time and space efficiency constraints, which some theoretical work overlooks, while also incurring low deployment overhead. We propose learning-augmented LRU, a deployment-oriented learning-augmented caching algorithm that guarantees \emph{1-consistency} and \emph{$O(k)$-robustness}, incurs low time and space overhead, and maintains strong compatibility. We further build a GPU cache, called \textsc{LCR}, on top of learning-augmented LRU to benefit from its theoretical guarantees and translate them into practical performance. In experiments, \textsc{LCR} reduces P99 time-to-first-token (TTFT) by up to 28.3\% on LLM workloads and increases throughput by up to 24.2\% on deep learning recommendation (DLRM) workloads. Even with poor predictions, performance degrades gracefully and remains close to \textsc{LRU}, demonstrating robustness with practical value.

cs.LG↗

arXiv2025

Neural Network Solution of Non-Markovian Quantum State Diffusion and Operator Construction of Quantum Stochastic Process

Non-Markovian quantum state diffusion provides a wavefunction-based framework for modeling open quantum systems. In this work, we introduce a novel machine learning approach based on an operator construction algorithm. This algorithm employs a neural network as a universal generator to reconstruct the stochastic time evolution operator from an ensemble of quantum trajectories. Unlike conventional machine learning methods that merely approximate time-dependent wavefunctions or expectation values, our operator-based approach yields broader applications and enhanced interpretability of the stochastic process. We benchmark the algorithm on the spin-boson model across diverse spectral densities, demonstrating its accuracy. Furthermore, we showcase the operator's utility in calculating absorption spectra and reconstructing reduced density matrices at extended timescales. These results establish a new paradigm for the application of machine learning in quantum dynamics.

quant-ph↗

arXiv2025

Robustifying Learning-Augmented Caching Efficiently without Compromising 1-Consistency

The online caching problem aims to minimize cache misses when serving a sequence of requests under a limited cache size. While naive learning-augmented caching algorithms achieve ideal $1$-consistency, they lack robustness guarantees. Existing robustification methods either sacrifice $1$-consistency or introduce excessive computational overhead. In this paper, we introduce Guard, a lightweight robustification framework that enhances the robustness of a broad class of learning-augmented caching algorithms to $2H_{k-1} + 2$, while preserving their $1$-consistency. Guard achieves the current best-known trade-off between consistency and robustness, with only O(1) additional per-request overhead, thereby maintaining the original time complexity of the base algorithm. Extensive experiments across multiple real-world datasets and prediction models validate the effectiveness of Guard in practice.

cs.DS↗

arXiv2025

SegQuant: A Semantics-Aware and Generalizable Quantization Framework for Diffusion Models

Diffusion models have demonstrated exceptional generative capabilities but are computationally intensive, posing significant challenges for deployment in resource-constrained or latency-sensitive environments. Quantization offers an effective means to reduce model size and computational cost, with post-training quantization (PTQ) being particularly appealing due to its compatibility with pre-trained models without requiring retraining or training data. However, existing PTQ methods for diffusion models often rely on architecture-specific heuristics that limit their generalizability and hinder integration with industrial deployment pipelines. To address these limitations, we propose SegQuant, a unified quantization framework that adaptively combines complementary techniques to enhance cross-model versatility. SegQuant consists of a segment-aware, graph-based quantization strategy (SegLinear) that captures structural semantics and spatial heterogeneity, along with a dual-scale quantization scheme (DualScale) that preserves polarity-asymmetric activations, which is crucial for maintaining visual fidelity in generated outputs. SegQuant is broadly applicable beyond Transformer-based diffusion models, achieving strong performance while ensuring seamless compatibility with mainstream deployment tools.

cs.CV↗

arXiv2025

Twin-Space Representation of Classical Mapping Model in the Constraint Phase Space Representation: Numerically Exact Approach to Open Quantum Systems

The constraint coordinate-momentum \textit{phase space} (CPS) has recently been developed to study nonadiabatic dynamics in gas-phase and condensed-phase molecular systems. Although the CPS formulation is exact for describing the discrete (electronic/ vibrational/spin) state degrees of freedom (DOFs), when system-bath models in condense phase are studied, previous works often employ the discretization of environmental bath DOFs, which breaks the time irreversibility and may make it difficult to obtain numerically converged results in the long-time limit. In this paper, we develop an exact trajectory-based phase space approach by adopting the twin-space (TS) formulation of quantum statistical mechanics, in which the density operator of the reduced system is transformed to the wavefunction of an expanded system with twice the DOFs. The classical mapping model (CMM) is then used to map the Hamiltonian of the expanded system to its equivalent classical counterpart on CPS. To demonstrate the applicability of the TS-CMM approach, we compare simulated population dynamics and nonlinear spectra for a few benchmark condensed phase system-bath models with those obtained from the hierarchical equations of motion method, which shows that our approach yields accurate dynamics of open quantum systems.

physics.chem-ph↗

arXiv2024

Simulating Quantum Many-Body States with Neural-Network Exponential Ansatz

Preparing quantum many-body states on classical or quantum devices is a very challenging task that requires accounting for exponentially large Hilbert spaces. Although this complexity can be managed with exponential ans\"atze (such as in the coupled-cluster method), these approaches are often tailored to specific systems, which limits their universality. Recent work has shown that the contracted Schr\"odinger equation enables the construction of universal, formally exact exponential ans\"atze for quantum many-body physics. However, while the ansatz is capable of resolving arbitrary quantum systems, it still requires a full calculation of its parameters whenever the underlying Hamiltonian changes, even slightly. Here, inspired by recent progress in operator learning, we develop a surrogate neural network solver that generates the exponential ansatz parameters using the Hamiltonian parameters as inputs, eliminating the need for repetitive computations. We illustrate the effectiveness of this approach by training neural networks of several quantum many-body systems, including the Fermi-Hubbard model.

quant-ph↗

arXiv2024

Neural Quantum Propagators for Driven-Dissipative Quantum Dynamics

Describing the dynamics of strong-laser driven open quantum systems is a very challenging task that requires the solution of highly involved equations of motion. While machine learning techniques are being applied with some success to simulate the time evolution of individual quantum states, their use to approximate time-dependent operators (that can evolve various states) remains largely unexplored. In this work, we develop driven neural quantum propagators (NQP), a universal neural network framework that solves driven-dissipative quantum dynamics by approximating propagators rather than wavefunctions or density matrices. NQP can handle arbitrary initial quantum states, adapt to various external fields, and simulate long-time dynamics, even when trained on far shorter time windows. Furthermore, by appropriately configuring the external fields, our trained NQP can be transferred to systems governed by different Hamiltonians. We demonstrate the effectiveness of our approach by studying the spin-boson and the three-state transition Gamma models.

quant-ph↗

arXiv2024

Tracking the Electron Density Changes in Excited States -- A Computational Study on Pyrazine

The development of X-ray free-electron lasers (XFELs) has enabled ultrafast X-ray diffraction (XRD) experiments, which are capable of resolving electronic/vibrational transitions and structural changes in molecules, or capturing molecular movies. While time-resolved XRD has received increasing attention, the extraction of information content from signals is challenging and requires theoretical support. In this work, we combined X-ray scattering theory and trajectory surface hopping approach to resolve dynamical changes in the electronic structure of photo-excited molecules by studying time evolution of electron density changes between electronic excited states and ground state. Using pyrazine molecule as an example, we show that key features of reaction pathways can be identified, enabling the capture of structural changes associated with electronic transitions for a photo-excited molecule.

physics.chem-ph↗

arXiv2024

Non-markovian neural quantum propagator and its application to the simulation of ultrafast nonlinear spectra

The accurate solution of dissipative quantum dynamics plays an important role on the simulation of open quantum systems. Here we propose a machine-learning-based universal solver for the hierarchical equations of motion, one of the most widely used approaches which takes into account non-markovian effects and nonperturbative system-environment interactions in a numerically exact manner. We develop a neural quantum propagator model by utilizing the neural network architecture, which avoids time-consuming iterations and can be used to evolve any initial quantum state for arbitrarily long times. To demonstrate the efficacy of our model, we apply it to the simulation of population dynamics and linear and two-dimensional spectra of the Fenna-Matthews-Olson complex.

physics.chem-ph↗

arXiv2024

Any-step Dynamics Model Improves Future Predictions for Online and Offline Reinforcement Learning

Model-based methods in reinforcement learning offer a promising approach to enhance data efficiency by facilitating policy exploration within a dynamics model. However, accurately predicting sequential steps in the dynamics model remains a challenge due to the bootstrapping prediction, which attributes the next state to the prediction of the current state. This leads to accumulated errors during model roll-out. In this paper, we propose the Any-step Dynamics Model (ADM) to mitigate the compounding error by reducing bootstrapping prediction to direct prediction. ADM allows for the use of variable-length plans as inputs for predicting future states without frequent bootstrapping. We design two algorithms, ADMPO-ON and ADMPO-OFF, which apply ADM in online and offline model-based frameworks, respectively. In the online setting, ADMPO-ON demonstrates improved sample efficiency compared to previous state-of-the-art methods. In the offline setting, ADMPO-OFF not only demonstrates superior performance compared to recent state-of-the-art offline approaches but also offers better quantification of model uncertainty using only a single ADM.

cs.LG↗

arXiv2024

Knowledgeable Agents by Offline Reinforcement Learning from Large Language Model Rollouts

Reinforcement learning (RL) trains agents to accomplish complex tasks through environmental interaction data, but its capacity is also limited by the scope of the available data. To obtain a knowledgeable agent, a promising approach is to leverage the knowledge from large language models (LLMs). Despite previous studies combining LLMs with RL, seamless integration of the two components remains challenging due to their semantic gap. This paper introduces a novel method, Knowledgeable Agents from Language Model Rollouts (KALM), which extracts knowledge from LLMs in the form of imaginary rollouts that can be easily learned by the agent through offline reinforcement learning methods. The primary challenge of KALM lies in LLM grounding, as LLMs are inherently limited to textual data, whereas environmental data often comprise numerical vectors unseen to LLMs. To address this, KALM fine-tunes the LLM to perform various tasks based on environmental data, including bidirectional translation between natural language descriptions of skills and their corresponding rollout data. This grounding process enhances the LLM's comprehension of environmental dynamics, enabling it to generate diverse and meaningful imaginary rollouts that reflect novel skills. Initial empirical evaluations on the CLEVR-Robot environment demonstrate that KALM enables agents to complete complex rephrasings of task goals and extend their capabilities to novel tasks requiring unprecedented optimal behaviors. KALM achieves a success rate of 46% in executing tasks with unseen goals, substantially surpassing the 26% success rate achieved by baseline methods. Furthermore, KALM effectively enables the LLM to comprehend environmental dynamics, resulting in the generation of meaningful imaginary rollouts that reflect novel skills and demonstrate the seamless integration of large language models and reinforcement learning.

cs.LG↗

arXiv2024

Artificial-intelligence-based surrogate solution of dissipative quantum dynamics: physics-informed reconstruction of the universal propagator

The accurate (or even approximate) solution of the equations that govern the dynamics of dissipative quantum systems remains a challenging task for quantum science. While several algorithms have been designed to solve those equations with different degrees of flexibility, they rely mainly on highly expensive iterative schemes. Most recently, deep neural networks have been used for quantum dynamics but current architectures are highly dependent on the physics of the particular system and usually limited to population dynamics. Here we introduce an artificial-intelligence-based surrogate model that solves dissipative quantum dynamics by parameterizing quantum propagators as Fourier neural operators, which we train using both dataset and physics-informed loss functions. Compared with conventional algorithms, our quantum neural propagator avoids time-consuming iterations and provides a universal super-operator that can be used to evolve any initial quantum state for arbitrarily long times. To illustrate the wide applicability of the approach, we employ our quantum neural propagator to compute population dynamics and time-correlation functions of the Fenna-Matthews-Olson complex.

quant-ph↗

arXiv2023

Episodic Return Decomposition by Difference of Implicitly Assigned Sub-Trajectory Reward

Real-world decision-making problems are usually accompanied by delayed rewards, which affects the sample efficiency of Reinforcement Learning, especially in the extremely delayed case where the only feedback is the episodic reward obtained at the end of an episode. Episodic return decomposition is a promising way to deal with the episodic-reward setting. Several corresponding algorithms have shown remarkable effectiveness of the learned step-wise proxy rewards from return decomposition. However, these existing methods lack either attribution or representation capacity, leading to inefficient decomposition in the case of long-term episodes. In this paper, we propose a novel episodic return decomposition method called Diaster (Difference of implicitly assigned sub-trajectory reward). Diaster decomposes any episodic reward into credits of two divided sub-trajectories at any cut point, and the step-wise proxy rewards come from differences in expectation. We theoretically and empirically verify that the decomposed proxy reward function can guide the policy to be nearly optimal. Experimental results show that our method outperforms previous state-of-the-art methods in terms of both sample efficiency and performance.

cs.LG↗

arXiv2023

Coherent two-dimensional THz magnetic resonance spectroscopies for molecular magnets: Analysis of Dzyaloshinskii-Moriya interaction

To investigate the novel quantum dynamic behaviors of magnetic materials that arise from complex spin-spin interactions, it is necessary to probe the magnetic response at a speed greater than the spin-relaxation and dephasing processes. Recently developed two-dimensional (2D) terahertz magnetic resonance (THz-MR) spectroscopy techniques use the magnetic components of laser pulses, and this allows investigation of the details of the ultrafast dynamics of spin systems. For such investigations, quantum treatment -- not only of the spin system itself but also of the environment surrounding the spin system -- is important. In our method, based on the theory of multidimensional optical spectroscopy, we formulate nonlinear THz-MR spectra using an approach based on the numerically rigorous hierarchical equations of motion. We conduct numerical calculations of both linear (1D) and 2D THz-MR spectra for a linear chiral spin chain. The pitch and direction of chirality (clockwise or anticlockwise) are determined by the strength and sign of the Dzyaloshinskii-Moriya interaction (DMI). We show that not only the strength but also the sign of the DMI can be evaluated through the use of 2D THz-MR spectroscopic measurements, while 1D measurements allow us to determine only the strength.

cond-mat.mes-hall↗

arXiv2022

Model-based Reinforcement Learning with Multi-step Plan Value Estimation

A promising way to improve the sample efficiency of reinforcement learning is model-based methods, in which many explorations and evaluations can happen in the learned models to save real-world samples. However, when the learned model has a non-negligible model error, sequential steps in the model are hard to be accurately evaluated, limiting the model's utilization. This paper proposes to alleviate this issue by introducing multi-step plans to replace multi-step actions for model-based RL. We employ the multi-step plan value estimation, which evaluates the expected discounted return after executing a sequence of action plans at a given state, and updates the policy by directly computing the multi-step policy gradient via plan value estimation. The new model-based reinforcement learning algorithm MPPVE (Model-based Planning Policy Learning with Multi-step Plan Value Estimation) shows a better utilization of the learned model and achieves a better sample efficiency than state-of-the-art model-based RL approaches.

cs.LG↗

arXiv2021

Probing photoinduced proton coupled electron transfer process by means of two-dimensional resonant electronic-vibrational spectroscopy

We develop a detailed theoretical model of photo-induced proton-coupled electron transfer (PPCET) processes, which are at the basis of solar energy harvesting in biological systems and photovoltaic materials. Our model enables to analyze the dynamics and the efficiency of a PPCET reaction under the influence of a thermal environment by disentangling the contribution of the fundamental electron transfer (ET) and proton transfer (PT) steps. In order to study quantum dynamics of the PPCET process under an interaction with non-Markovian environment we employ the hierarchical equations of motion (HEOM). We calculate transient absorption spectroscopy (TAS) and a newly defined two-dimensional resonant electronic-vibrational spectroscopy (2DREVS) signals in order to study the nonequilibrium reaction dynamics. Our results show that different transition pathways can be separated by TAS and 2DREVS.

physics.chem-ph↗