A Hybridizable Neural Time Integrator for Stable Autoregressive Forecasting
This work addresses the challenge of stable autoregressive forecasting for chaotic dynamical systems, which is critical for scientific foundation models in domains like fusion energy.
The paper introduces a hybridizable neural time integrator that embeds an autoregressive transformer within a shooting-based mixed finite element scheme, achieving provable stability for long-horizon forecasting of chaotic systems. The method reduces model parameters by 65× and achieves a 9,000× speedup over particle-in-cell simulation for a fusion component surrogate.
For autoregressive modeling of chaotic dynamical systems over long time horizons, the stability of both training and inference is a major challenge in building scientific foundation models. We present a hybrid technique in which an autoregressive transformer is embedded within a novel shooting-based mixed finite element scheme, exposing topological structure that enables provable stability. For forward problems, we prove preservation of discrete energies, while for training we prove uniform bounds on gradients, provably avoiding the exploding gradient problem. Combined with a vision transformer, this yields latent tokens admitting structure-preserving dynamics. We outperform modern foundation models with a $65\times$ reduction in model parameters and long-horizon forecasting of chaotic systems. A "mini-foundation" model of a fusion component shows that 12 simulations suffice to train a real-time surrogate, achieving a $9{,}000\times$ speedup over particle-in-cell simulation.