LiMemFlow: Memory-Enhanced Flow Matching for Lithium-Ion Transport in Solid-State Electrolytes

Understanding and predicting long-horizon lithium-ion transport in solid-state electrolytes is important for screening high-conductivity materials. This is especially important near room temperature (300 K), where practical battery operation is most relevant. Although machine-learning interatomic potentials reduce force-evaluation cost, they still rely on sequential small-step integration, motivating large-step generative propagators for long-horizon trajectory prediction. Existing large-step propagators based on flow matching and diffusion models typically formulate trajectory generation as a position-only Markov transition, discarding velocity-induced temporal dependence and collective slow processes that strongly affect diffusion statistics at 300 K. We introduce LiMemFlow, a memory-enhanced flow matching model for lithium-ion trajectory prediction. LiMemFlow combines a PET-based geometric backbone with two complementary memory modules: a point-wise history memory that encodes recent atomic displacement sequences and a pair-wise history memory that stores historical relative motion between neighboring atoms. These memory signals are injected into the equivariant propagator through residual cross-attention and updated during autoregressive rollout. We evaluate LiMemFlow on two trajectory datasets of lithium-containing crystalline structures: our newly computed 300 K dataset and 600–1200 K lithium-transport benchmarks. At 300 K, LiMemFlow reduces $$\begin{gathered} log[ \\ mathrmMSD\_ \\ mathrmLi] \end{gathered}$$ MAE by 20.0% and $$\begin{gathered} log[ \\ mathrmMSD\_ \\ mathrmLi \\ text-curve] \end{gathered}$$ MAE by 17.0% compared with the strongest LiFlow baseline. Consistent with the expected role of memory in slow, path-dependent transport, explicit memory conditioning yields its clearest gains at 300–600 K, including the practically relevant near-room-temperature regime, while the gains diminish at higher temperatures and LiMemFlow remains competitive with the strongest LiFlow baseline at 1000–1200 K.