TI-DeepONet: Learnable Time Integration for Stable Long-Term Extrapolation

TL;DR

This paper tackles the challenge of accurate long-term extrapolation for neural operators modeling time-dependent PDEs. It introduces TI-DeepONet, which learns the instantaneous time derivative and evolves the state via numerical time integration, and its adaptive variant TI(L)-DeepONet that learns state-conditioned RK4 weights. Across Burgers', KdV, KS, and 2D inviscid Burgers' equations, the TI-based methods substantially reduce extrapolation errors and maintain stability beyond the training horizon, with TI(L)-DeepONet often delivering the best performance due to adaptive time-stepping. The framework bridges neural operator learning with classical numerical analysis, enabling continuous-time predictions and efficient inference with higher-order integrators, at the cost of increased training/inference compute which the authors address through analysis and potential optimizations.

Abstract

Accurate temporal extrapolation remains a fundamental challenge for neural operators modeling dynamical systems, where predictions must extend far beyond the training horizon. Conventional DeepONet approaches rely on two limited paradigms: fixed-horizon rollouts, which predict full spatiotemporal solutions while ignoring temporal causality, and autoregressive schemes, which accumulate errors through sequential prediction. We introduce TI-DeepONet, a framework that integrates neural operators with adaptive numerical time-stepping to preserve the Markovian structure of dynamical systems while mitigating long-term error growth. Our method shifts the learning objective from direct state prediction to approximating instantaneous time-derivative fields, which are then integrated using standard numerical solvers. This naturally enables continuous-time prediction and allows the use of higher-order integrators at inference than those used in training, improving both efficiency and accuracy. We further propose TI(L)-DeepONet, which incorporates learnable coefficients for intermediate slopes in multi-stage integration, adapting to solution-specific dynamics and enhancing fidelity. Across four canonical PDEs featuring chaotic, dissipative, dispersive, and high-dimensional behavior, TI(L)-DeepONet slightly outperforms TI-DeepONet, and both achieve major reductions in relative L2 extrapolation error: about 81% compared to autoregressive methods and 70% compared to fixed-horizon approaches. Notably, both models maintain stable predictions over temporal domains nearly twice the training interval. This work establishes a physics-aware operator learning framework that bridges neural approximation with numerical analysis principles, addressing a key gap in long-term forecasting of complex physical systems.

Figures (9)

• Figure 1: A schematic showing how autoregressive predictions in time are performed.
• Figure 2: A schematic of the proposed TI-DeepONet architecture.
• Figure 3: Temporal evolution of the relative $L_2$ error for different frameworks in all applications. Left: Comparison of baseline models with TI-DeepONet AB, which uses RK4 during training and AB2/AM3 during inference. Right: Comparison among TI-based variants. TI-DeepONet RK4 employs RK4 integration in both training and inference, while TI(L)-DeepONet adapts RK4 coefficients via a learnable scheme conditioned on the input state.
• Figure 4: 1D Burgers' equation: Performance of all frameworks in the training ($t\in[0, 0.5]$) and extrapolation ($t\in[0.5, 1]$) regimes for a representative sample.
• Figure 5: 1D KdV equation: Performance of all the frameworks in the training ($t\in[0, 2.5]$) and the extrapolation regime ($t\in[2.5, 5]$) for a representative sample.