AdaRecurr: State-Adaptive Recurrent Computation for Neural PDE Simulations

Neural simulators can substantially reduce the cost of PDE simulation, but autoregressive rollouts typically apply the same amount of inference computation to every physical transition even though different states may benefit differently from additional computation. Some states require only limited recurrent refinement, whereas others continue to improve with deeper refinement. This raises a natural question: can computation be allocated dynamically according to the current physical state as the trajectory evolves? We introduce AdaRecurr, a state-adaptive recurrent computation framework that turns recurrent depth from a global test-time parameter shared across an entire rollout into an online computation budget selected separately for each physical transition. AdaRecurr builds on a pretrained recurrent neural PDE simulator that can perform a variable number of recurrent refinement iterations for each prediction. Keeping this simulator frozen, AdaRecurr observes the current physical field together with the latent refinement already executed and sequentially decides whether to STOP or CONTINUE. Continuing reuses the current latent state and performs further recurrent refinement, while stopping accepts the current prediction and advances the physical rollout. The controller is trained on self-generated autoregressive trajectories to balance prediction error against additional recurrent computation. Across three neural PDE benchmarks, AdaRecurr reduces trajectory-level relative $L_2$ error by 3.6–15.3% over fixed recurrent-depth inference at comparable mean recurrent depth. On viscous Burgers, states receiving more computation exhibit a 36.2% higher mean aggregate estimated shock strength than states receiving less computation. In a controlled compute-demand heterogeneity experiment, we further observe larger gains from adaptive allocation when recurrent-computation requirements vary more strongly across physical states. These results show that test-time recurrent computation can be allocated dynamically along evolving physical trajectories rather than using a single fixed recurrent depth for the entire rollout.