Memory Needs Control: Adaptive Memory Gate for Solving PDEs

Neural operators have emerged as a powerful data-driven approach for solving time-dependent partial differential equations (PDEs). Among recent advances, memory-augmented neural operators explicitly incorporate past states and have achieved remarkable performance under low-resolution observation settings. However, existing approaches use a fixed memory weight, which sets how much of the past is mixed into the current state. We show that the benefit of memory depends on the resolution and the physical parameters. Empirically, on Kuramoto–Sivashinsky (KS), no single memory weight is optimal across resolutions and viscosities. Theoretically, the Bayes gain of history is bounded by the unresolved spectral energy. We propose AMGFNO, which controls the memory weight with an adaptive memory gate (AMG). The AMG is the product of a content gate and a frequency-aware gate. The gate is learned during training: It opens where much of the spectrum is unresolved and nearly closes where almost all of the spectrum is resolved. On 20 settings spanning 1D, 2D, and 3D PDEs (KS, Burgers', Navier–Stokes, and compressible CFD), AMGFNO is best or tied in 17. On KS at $\nu =0.10$, $f=32$, AMGFNO reduces nRMSE by 30% (0.036 to 0.025) over the best fixed-memory baseline. The gate is not specific to S4. Across S4, S5, LRU, LSTM, and a causal Transformer memory, the gate lowers the error of the model without the gate in 38 of 45 settings.