StFT: Spatio-temporal Fourier Transformer for Long-term Dynamics Prediction
Authors: Da Long
Shandian Zhe
Samuel Williams
Leonid Oliker
Zhe Bai
Venue: Accepted by TMLR ()

Abstract: Simulating the long-term dynamics of multi-scale and multi-physics systems poses a significant challenge in understanding complex phenomena across science and engineering. The complexity arises from the intricate interactions between scales and the interplay of diverse physical processes, which manifest in PDEs through coupled, nonlinear terms that govern the evolution of multiple physical fields across scales. Neural operators have shown potential in short-term prediction of such complex spatio-temporal dynamics; however, 
achieving stable high-fidelity predictions and providing robust uncertainty quantification over extended time horizons remains an open and unsolved area of research. These limitations often lead to stability degradation with rapid error accumulation, particularly in long-term forecasting of systems characterized by multi-scale behaviors involving dynamics of different orders. To address these challenges, we propose an autoregressive Spatio-temporal Fourier Transformer (StFT), in which each transformer block is designed to learn the system dynamics at a distinct scale through a dual-path architecture that integrates frequency-domain and spatio-temporal representations. By leveraging a structured hierarchy of StFT blocks, the resulting model explicitly captures the underlying dynamics across both macro- and micro- spatial scales. Furthermore, a generative residual correction mechanism is introduced to learn a probabilistic refinement temporally while simultaneously quantifying prediction uncertainties, enhancing both the accuracy and reliability of long-term probabilistic forecasting. Evaluations conducted on three benchmark datasets (plasma, fluid, and atmospheric dynamics) demonstrate the advantages of our approach over state-of-the-art ML methods.

Published in Transactions on Machine Learning Research (01/2026)
StFT: Spatio-temporal Fourier Transformer for Long-term Dynamics Prediction
Da Long
do. do.long@utah. edu
School of Computing, University of Utah
Shandian Zhe
u6014464 @utthh edu
School of Computing, University of Utah
Samuel Williams
swwilliam@@bblggr
Applied Mathematics and Computational Research Division, Lourence Berkeley National Laboratory
Leonid Oliker
loliker gov
Applied Mathematics and Compulational Research Division, Laurence Berkeley National Laboratory
Zhe Bai
ahe zhebai@lbl.gg gov
Applied Mathematics and Computational Research Division, Laurence Berkeley National Laboratory
Reviewed on OpenReview: https: // openreview. net/ forum? id= a9Cb0ri20N
Abstract
Simulating the long-tern bong-temm dynamics of multi-scale and multi-physics systems poses signifi- cant challenge in understanding complex phenomena across science and engineering. The complexity arises from the intricate interactions between scales and the interplay of diverse physical processes, which manifest PDEs through coupled, nonlincar terms that govern the evolution of multiple physical fields across scales. Neural operators have shown potential in short-term prediction of such complex spatio-temporal dynamics; however, achieving stable high-fidelity predictions and providing robust uncertainty quantification over extended time horizons remains open and unsolved area of research. These limitations often lead to stability degradation with rapid error accumulation, particularly in long-term forecasting of systems characterized by multi-scale behaviors involving dynamics of different orders. To address these challenges, we propose an autoregressive Spatio-temporal Fourier Transformer (StFT), in which each transformer block is designed to learn the system dynamics at distinct scale through dual-path architecture that integrates frequency-domain and spatio-temporal representations. By leveraging a structured hicrarchy of StFT blocks, the resulting model explicitly captures the underlying dynamics across both macro- and micro- spatial scales. Furthermore, generative residual correction mechanism is introduced to learn probabilistic refinement temporally while simultancously quantifying prediction uncertainties, enhancing both the accuracy and reliability of long-term probabilistic forecasting. Evaluations con- ducted on three benchmark datasets (plasma, fluid, and atmospheric dynamics) demonstrate the advantages of our approach over state-of-the-art ML methods.
"Corresponding author (shehai&(lil.go))

---

Published in Transactions on Machine Learning Research (01/2026)
1 Introduction
Predicting the long-term spatio-temporal dynamics of systems governed by partial differential equations (PDEs) is a cornerstone of scientific and engineering research, with broad applications in fields such as earth system modeling, plasma science, fluid dynamics, and beyond. Traditional approaches rely heavily on numerical solvers, which discretize the domain and iteratively solve PDEs using methods including finite clement, finite volume and spectral methods (Tadmor 2012). While effective in many scenarios, these techniques face limitations when applied to multiphysics systems characterized by complex dynamics and multiscale behaviors. They require substantial computational resources and exhibit poor scalability with increasing problem size, rendering them impractical for high-dimensional, largo-scale, long-term physics systems due to excessive computational costs and memory demands.
Recent advances in deep learning have revolutionized the field of PDE modeling by introducing data-driven methodologies that significantly accelerate computations for science while maintaining high accuracy. Inspired by the universal approximation theorem (Chen & Chen 1995). neural operators that learn the mapping between two function spaces have demonstrated great success in simulating various PDE systems across multiple scientific disciplines without retraining for new conditions (Li et al. 2020 Lu et al. 2021 Building the success of transformers in natural language processing and computer vision [Vaswani] 2017 Dosovitskiy neural operators process multiple imput functions while enabling arbitrary querying of output function locations, offering enhanced flexibility in handling complex functional mappings (Hao
2020), transformer-base
et al. problems, including weather forecasting, turbulent fluid dynamics, and boiling phenomena (Pathak et al.,
2023 Li et. al.) 2022). A series of neural operalors have been developed to address complex scientific
2022
Li ct al.] 2023a Bi ct al. 2023 Hassan et al. 2023: Lin et al. 2021).
Despite the success of these methods, accurate and long-term predictions of complex physical systems remain challenging, primarily due to the requirements for numerical stability, high-fidelity modeling, and reliable uncertainty quantification over extended horizons. The inherent multi-scale nature and multi-physics complexity of such systems necessitate methodologies that can efficiently represent and integrate dynamics across disparate spatial and temporal scales while simultaneously capturing the complex interactions between distinct physical processes, such as the influence of micro-scale turbulence on macro-scale flow in fluids and combustion (Poters 2009 Natrajan et al. 2007). For largo-scale atmospheric pressure systems, high-pressure ridges and low- pressure troughs play crucial role in shaping local weather patterns; inaccurate representation of those structures can cause significant errors in forecasting rainfall, wind speed, and temperature (Wang let al. 2006 Barlow et al.) 2019). In magnetically confined plasmas, multiphysics arises from the coupling of physical processes that govern plasma behavior, including electromagnetic fields, turbulence transport, thermodynamics, and particle interactions that are potentially coupled with kinetic models in high-fidelity simulations. The magnetohydrodynamii (MHD) instabilitics caused by current or pressure gradients can limit burning plasma performance, and threaten fusion device integrity (von Goeler et al., 1974 Graves el al., 2012 Seo et al., (2024). Furthermore, integrating uncertainty quantification (UQ) into modeling frameworks is essential for assessing the confidence and reliability of predictions in such complex systems (Cheung et al., 2011 Scher & Messori, 2018 Kruger et al. 2024). Although neural operators present advantages over traditional approaches, they still encounter challenges associated with the demands for scientific. fidelity and stability, especially when the underlying physics exhibit rapid changes or high-frequency components. These issues are further intensified in high-resolution simulations of multi-scale scenarios. Recent efforts to address these limitations include NeC,Net, which encodes a high-order numerical scheme with boundary condition encoding into neural networks (Wang et al., (0244)), and Dyffusion, which trains a forccasting network and an interpolation network that allows lor continuous time sampling and multi-step prediction for long-range forecasting [Rühling Cachay et al.) 2023). However, most existing neural operators lack built-in mechanisms for uncertainty quantification, which is particularly critical for reliable modeling of long-term dynamics, where even small errors can propagate across scales and result in significant inaccuracies.
Existing approaches for predicting spatio-tempont temporal dynamics can be broadly be classified into two primary categorics. The first category compriscs models that directly forecast future states at fixed time horizons using a sequence of past observations (Wang et al.] [224bb] Kontolati et al. 2024). The second category includes models that utilize an autoregressive manner, which addresses challenges of scaling and fitting complexities
2

---

Published in Transactions on Machine Learning Research (01/2026)
continnous-time emulator (Pathak et al. 2022 Rühling Cachay et al. 2023 Lippe et al.
2023 McCabe
et al. 2023 Generally, prediction errors incurred in the short term can accumulate, leading to instability and reduced accuracy in long-term forecasts. To mitigate these issues, previous work has proposed techniques such as the pushforward trick, invariance preservation, and iterative refinement (MeCabe et al., 2023 Lippe
jet al. 2023 Brandstetter et al. 2022). Nevertheless, the development of multi-scale modeling frameworks for long-term dynamic prediction remains crucial for capturing the interactions across scales and enhancing prediction accuracy. Concurrently, incorporating uncertainty quantification is critical for identifying the spatial and temporal regions where predictive confidence deteriorates, especially in complex systems where localized uncertainties can influence global dynamics over time.
In this work, we introduce an auloregressive Spatio-temporal Fourier Transformer (StFT), for long- range forecasting of multi-scale and multi-physic systems. At each level of the spatial hierarchy, one StFT block models the physical dynamics associated with a distinct spatial scale or receptive field, as inferred from the spatiotemporal data. Each StFT block adopts a dual-path architecture: (1) the frequency path captures large-scale dynamics by operating in the Fourier domain, focusing low-frequency components that, are critical for modeling global behavior. (2) the spatio-temporal path operates in the full physical space incorporating all spatio-temporal features to capture fine-scale features. Through a hicrarchical composition of StFT blocks across multiple scales, augmented by a generative residual correction block, the resulting model learns the intricate interactions both within the same scales and across different. scales. Moreover, it produces uncertainty estimates at each spatial and temporal point, enabling assessment of prediction confidence throughout the forecast. The cascading StFT blocks enable our model to predict high-resolution dynamics across a. spectrum of varying scales in correlated physical processes. By integrating SiFT within auto-regressive framework, our method achieves superior accuracy in long-term prodictions compared to existing stale-oftthe-arl autoregressive baselines. Our contributions are summarized follows:
We propose Spalio-temporal Fourier transformer (SLFT), a novel ML model that learns underlying dynamics across spatial scales for multi-physics systems via a dual-path (frequency and spatio- temporal path) architecture, which effectively captures both the global, large-scale structures and local, fine-scale features.
We propose SUFT-F, which incorporates a probabilistic residual correction mechanism to refine the forecasting of StFT temporally and provide pointwise uncertainty quantification.
We propose an overlapping tokenizer and a detokenizer that share regions between adjacent patches, improving spatial smoothing and reducing discontinuity artifacts.
We demonstrate the effectiveness of StFT in an autoregressive framework on a diverse set of applications inchuding the plasma, fluid, and atmospheric dynamics. Evaluating performance across variables, StFT outperforms the best baselines. Its probabilistic variant, StFT-F further improves average forecasting accuracy by 5% and produces uncertainty estimates that are empirically calibrated, as demonstrated through confidence-basse evaluation.
2 Related Work
Neural Operators. Neural operator architectures and their variants have been proposed, including Fourier
neural operators (Li et al.] 2020 Gupta et al. 2021 Tran et al.] 2021 Cao al. 2023
Li et al.
2023b
Rahman a..,
2022), DeepONet (Lu al. 2021: Wang et al.) 2021 Jin et al.
2022;
Wang et al.
2022
Kontolati et al. 2024 Prasthofer et al., 2022), transformer based operators (Hao et al.
2023
Cao
2021
Li et al. 2022),
and image-to-image operators (Gupta & BBandstetter 2022 Long al.
2024 U-Net,
fundamentally hicrarchical structure model, which has inspired several neural operators (Rahman et al.
2022: Liu et al.
2022b; Gupla & Brandsteller 2022), allows solutions to multi-scale PDEs by hierarchically
aggregating feature representations of progressively coarser spatial resolutions. Recent work in computer vision (Liu et al.
2021 Fan et al.] 2021 Zhang et al. 2022) have introduced methods for extracting multi-scale features through hierarchical architectures. However, these hicrarchical models do not explicitly model or
forecast the multi-scale structures of physical processes, which limit the transparency and interpretability of
3

---

Published in Transactions on Machine Learning Research (01/2026)
their representations across scales. In contrast, our method begins with coarse approximation that captures large-scale, low-frequency phenomena, and incrementally refines the representation over layers to resolve finer details. These structured decompositions enable error diagnosis, enhance interpretability of model performance different scales, and allow for targeted improvements with an explicit refinement mechanism.
Generative Models. Generative models, especially diffusion models have demonstrated success various domains, including vision, audio, robotics (Ho al. 2020 Song et al., 2020 Tian al. 2024 Kong et al, 2020 Wolleb et al.] 2022). and relevance to spatio-temporal dynamics prediction (Ho et al.. 2022 Voleti
2023 As alternative approach in generative modeling, flow matching has been introduced to support efficient sampling and has since been applicd to video generation (Lipman et al. 2022
et al. 2022 Singer et al. 2022 Rühling Cachay et al.
Liu et al. 2022a, Albergo & Vanden-Eijuden 2022 Polyak et al. 2024; Esser et al. 2024]. While video generation typically explores a range of creative and diverse possibilities from text or image prompts, forccasting spatial-temporal dynamics driven by PDEs necessitates more than mere statistical resemblance requires each prodiction is firmly grounded the underlying physics. To achieve accurate forecasting while capturing the inherent stochasticity of physical processes, our work incorporates flow matching block following the proceeding StFT blocks. This enables the model to align its prodiction distribution with the underlying physical dynamics and generate calibrated uncertainty estimates via confidence-base based metrics.
3 Method
Formulation. We consider an autoregressive formulation for long-term multi-scale spatiotemporal physical processes. We define vector ū, representing the historical snapshots of the multi-physics variables at timestamps from t- k + 1 to t, in total of & snapshots of ün [u, Ue-1,..., uc4+1] specifically. We formulate the probabilistic one-step forward neural operator StFT-F as
u++1 Faz (ūc) Te41,Te41 Po, Fa, (ū.(),),
(1)
where Foa denote the StFT operator, a deterministic forecasting parameterized by 9a, and Pg, is the generative flow matching block parameterized by 0g for refining the forecasting of StFT while quantifying uncertainty estimates. Fo, represents the deterministic evolution of the system that encapsulates the dynamics from multi-scale spatial refinement in StFT. The residual refinement. Po, (rl (rl Fo, (ū.)) captures the probabilistic nature of the residual from the generative model. It represents prediction distribution conditioned on the current state ūt and the deterministic prediction Fo,( modeling the uncertainty or variations that missed by the deterministic component. The residual rz+1 calculates deviations from the deterministic prediction, and its distribution allows the model to account for noise or inherent stochasticity in the physical processes. Therefore, by sampling residual rt, our model learns stochastic trajectories from data. Besides providing prediction uncertainties, these stochastic trajectories can help study the long-term behavior, stabilities, and bifurcations in stochastic systems [Lucor et al. 2003; Krämer et al. 2022
Figure 1 presents the overview of the SUFTFF, the overlapping tokenizer, and the design of the SUFT block. Algorithm and 2 detail the design of our model. The following subsections introduce StFT with the overlapping tokenizer creniizerrdetoooee /detoknnzerr and the residual correction mechanism based on flow matching.
3.1 StFT Block
Overlapping tokenizer. Each StFT block first tokenizes the discretized functions ü, of shape T' W H C, where T is the temporal dimension, W, H are the spatial dimensions, and C denotes the mumber of physical variables. We apply tokenization along the spatial dimensions for each variable. To enhance spatial continuity while minimizing visual artifacts, we propose to use an overlapping tokenizer (OLT) and detokenizer (OLDT) that allows adjacent patches to share boundaries through overlapping regions. For instance, as shown in Figure [ a 3 x 3 input generates four 2 2 patches with a 1 X 1 overlap, where the overlapping areas (indicated in gray) are shared between patches. During detokenization, overlapping regions are reconstructed by averaging the corresponding values from neighboring patches. This strategy effectively mitigates discontinuity issues at the patch boundaries, which is particularly important for accurately representing smooth and
4

---

Published in Transactions on Machine Learning Research (01/2026)
Variables
Temporal Stacking
Spatial Refinement
L
W
it
StFT Block
1t+
1t
H
Overlapping Tokenizer
Coarse-to-Fine Composition
Sampled Noise
Variable Mixer
FFT & Filter
Spatiotemporal Embedder
Freq path
patt
S- path
Freq Embedder
Flow Matching
F 0
Transformer Layer
Transformer
Layer
Ped
Residual Refinement
Linear
Element-wise Add
W+++
StFT-F
Overlapping Detokenizer
Overlapping Tokenizer
üt historical snapshots
2
2
u *++1 I-th layer in SIFT block
1+1 #+11
4
6
F o
deterministic prediction
9
8 8
Pos
probabilistic refinement
SIFT
Figure 1: Left: Overview of the proposed StFT and StFT-F model. The model predicts ue+1 using the past k snapshots ūe employing L spatial refinements from coarse to fine scales through the proposed Spatiotemporal Fourier Transformer (SLFT) blocks. Bottom left: illustration of the overlapping tokenizer, where the patch size is 2 x 2, and the overlapping number is 1 x 1. Right: Illustration of the proposed StFT block. First, the past snapshots ūn and coarser prediction from the previous layer are lemporally stacked. The stacked discretizations are passed to the overlapping tokenizer to generate tokens for each variable. Next, tokens corresponding to different variables in the same spatial are mixed through a variable mixer. Two paths of transformation, frequency path and spatiotemporal path, process the frequency embeddings and the spatiotemporal embeddings respectively. The finer prediction for timestamp obtained after passing through the overlapping detokenizer.
continuous target functions. Furthermore, incorporating shared boundaries into patch embeddings enriches feature representation and extraction from fine-scale structures.
Variable mixer. Each StFT block is designed to handle specific scale; therefore, by employing a specific patch size, we partition the input at a corresponding level of granularity. To ensure that the first block captures the coarsest features or the largest. scale, we set the patch size plo x Ph, to a large value, allowing it to model broad spatial structures effectively. As a result, O pr x
C) patches are fed into the variable PT x
mixer, where patches corresponding to different physical processes but sharing the same spatial domain are mixed into single token. Following this step, two transformation paths are performed: one is in the spatio-temporal domain, which operates self-attention on spatio-temporal embeddings; the other is in the frequency domain, which operates on frequency embeddings.
Frequency embeddings. The tokens are first processed by a 2D/3D Fourier transform, where only low- frequency components are retained. These low-frequency components are then passed through a frequency embedder to obtain frequency embeddings ft. Subsequently, these frequency embeddings are fed to the
5

---

Published in Transactions on Machine Learning Research (01/2026)
standard transformer layers for mixing information and nonlinear transformation in the frequency domain. Finally, an inverse 2D/3D Fourier transform and linear projection are applied to map the frequency embeddings back to the spatio-temporal domain.
Spatio-temporal embeddings. The same set of tokens first pass through a apatio-tempore embedder, after which the spatio-temporal embeddings & are processed by multiple standard transformer layers for mixing correlations and nonlinear transformation. Finally a separate linear projection is applied to get predictions for each patch of U
Next, an overlapping detokenizer yields the first block prediction "/+1- Each token represents a significant. portion of the historical snapshots, encapsulating macroscopic structural features. This coarse-level partitioning reduces the complexity of modeling fine-grained details. By maintaining a lower granularity, the model prioritizes structural coherence over extraneous details, enabling to focus on capturing and predicting global relationships between regions more effoctively.
3.2 A Hierarchy of StFT Blocks
In the subsequent StFT blocks, we shift our focus to smaller scales with fine details. Consequently, we concatenate the prediction ul41 with the input üt and consider this combination as the input for the next StFT block. We further subdivide each patch from the previous StFT block into smaller patches.
The smaller patches allow for less information to be aggregated within single patch, thereby minimizing the risk of losing local variations and enhancing the richness and informativeness of the fine-scale representation.
Coerse to fine petches
By leveraging the finer gramularity of the patches to focus smaller regions, it allows the model to better localize features and capture their details. In addition, conditioning on the coarser prediction allows the models to iteratively refine its estimates, beginning with a broad global summary of the corresponding regions. Through the repeated subdivision of the patches, the model progressively refines its predictions across multiple scales. As shown in Figure []] the StFT model predicts u and applies total of L multi-scale spatial refinements.
A
patc
Subdivide
Subdivide
Figure 2: A patch is subdivided into smaller patches for hierar- chical learning.
3.3 Residual
Refinement and Uncertainty Estimation Based on Flow Matching
Finally, the model refines its deterministic predictions through a rectified flow block, which belongs to the family of flow matching models (Liu 2022 Liu et al. 2022a). Flow matching is formulated as an ordinary differential equation in time [0, 1J, (x) uz (w) (x)),, where the learnable velocity field v, directs the transformation of each sample X, from source distribution Po, typically Gaussian distribution, toward the target distribution Xı Pı with p1 representing the data distribution. If we prescribe the velocity field e, such that it guides every sample along a straight-line trajectory from Xo to X1, it is referred to as a rectified flow. In this case, X, represents the linear interpolation across the entire timespan between X, and X1, which can be expressed as X, = TX, + (1 - No- We employ a parameterized Me, to approximate Pr, leading to the following learning objective: [(B,) Er,x,.X,
(X1 Xx) 2 In our model, the rectified flow block takes the deterministic prediction u4+1 from the composition of L SLFT blocks and the observations ú, conditioning inputs. Its objective is to generate the distribution of residuals Tr+1 y E, 2/411 where is the ground truth for the solution at t + 1. Our training loss then becomes:
Ex.-NNo.,, (ūg, Fo, (o,,, (r441 -x)))
(2)
where X, the linear interpolation between the source sample Xo and the target Te+1-

---

Published in Transactions on Machine Learning Research (01/2026)
Algorithm 1: SIFT
Algorithm StFT block: Freg. and S-t. paths
Inputs:
history in = (u..... (4-k+1)
Frequency Path: Tt1 Mar--aarileMMMierr ((2c)o)
Initialize: blocks e [1, Ll, patch sizes phi truncation modes magaw overlaps ohi,o ub41 88
J. < T--FFFFilter(E,, Ma,
8 fiz 6- (ree-PerBBbeeddr((..))
None, M+41 as 0 8 var idx
f TranssormerBlock (fr)
5 ( ma, ma)
for 1 = I..... L do
+ Limear ((04-111)))
Fe TemporalStackingg(. 1)
7 Spatiotemporal Path:
{ac) OLT(2e, Ph.a, Oh,w);
F2,2 + Variaaabeeeeeeieerreaaa (12))o)
(2++1hu; Invoke Freg. & S-t. paths
8 c < S+Embedder((r.))
1/41 OLDT((2e+))e Plase oh.ss)
9 6 t TransformorBlock
M141 ue41 + 4/41
{24122 +- Linear
10 Return: Me+1
11 Merge: (z(-1)o + (Fem1a - (24112)
4 Experiments
4.1 Datasets
In this section, we consider three spalio-temporal multi-physics systems arising from time-dependent PDEs of a variety of complexities, including a high-dimensional plasma dynamics system based on reconstructed equilibrium of DIII-D experimental discharges (Bai et al. 2025). a 2D incompressible Navier-Stokes equation in velocity-pressiie pressure form within a square domain driven by an external force, and a viscous shallow-water equation modeling the dynamics of large-scale atmospheric flows on a spherical domain. The problem setup and data generation are detailed in Appendix B
4.2 Experimental Setup
Long-term multi-physics prediction up to a horizon of 244 timesteps. Our goal is to simulate long-time trajectories given a few initial observations. This task is particularly challenging due to the multiple correlated variables present in the Navier-Stokes and plasma magnetolydrodynamiis (MHD), with the test trajectories consisting of snapshots that vary from 71 to 244. We employ an autoregressive framework for all the methods: during training, each model utilizes five historical snapshots to predict the next one in a forward pass. At test time, given the initial five snapshots of trajectory, all models autoregressively generate the entire trajectories by iteratively predicting future states based their own previous outputs.
Baseline setup. We evaluate these datasets using the following well-known and state-of- thre--off-th--tttte methods for comparison: autorogressive AFNO, autoregressive Fourier Neural Operator (AR-FNO), autoregressive U-Net (AR-UNet),, and autoregressive vision transformer (AR-Vi) (iT). More specifically, for AFNO, we used the latest. variant from the DPOT work (Hao et al., 2024), which is enhanced with a temporal aggregation layer and improved expressivity through the removal of enforced sparsity. For FNO, we use the authors' open-source implementation. For U-Net, employ the implementation of the modified U-Net as evaluated in the recent BubbleML. work (Hassan et al. 2023), where the modified U-Net was initially used in PDEBench work [Takamoto et al. 2022). and demonstrated superior performance over their baselines.
Validation and hyperparameter tuning. We divide the trajectories of each dataset into training, test, and validation sets. For each method, we identify the tunable hyperparameters, specify a range for each hyperparameter, and conduct a grid scarch. For our methods, we implement both StFT, the deterministic component of our model, and StFT-F, our model with generative residual refinement block. We use the AdamW optimizer with learning rate of 4 to train those models on an A100 GPU. We ensure that all models are fairly and thoroughly trained by imposing an identical time budget across all models. More specifically, we impose 24 GPU hours on the plasma MHD and Navier-Stokes datasets, and 48 GPU hours on the shallow-water dataset. A comprehensive list of all the hyperparameters along with their respective ranges is provided in the Appendix
Evaluation metrics. First, we evaluate the forecasting performance by calculating the mean Le relative error. For StFT-F, in order to obtain the mean prediction, we generate 50 stochastic predictions at each
7

---

Published in Transactions on Machine Learning Research (01/2026)
x10
AR-UNet
4.00
Mean L, Relative Error
10
AR-VIT
AR-FNO
Mean L, Relative Error
1.00
2.00
AR-SIFT-F
10-1
AR-UNet
AR-VIT
0.10
AR-FNO
AR-DPOT
10-2
AR-StFT-F
threshold:
0.1
0.0
0.2
0.4
0.6
0.8 1.0
0.2
0.4
0.6
0.8
1.0
Timestamps
Timestamps
Figure 3: Results of autoregressive prediction in L2 relative enror (log scale) across the timespan: (left) perturbed parallel vector potential JA₁ plasma MHDD and (right) magnitude of velocity in shallow-water equation. The shaded region indicates the uncertainty distribution of in the relative error of StFT--F For a given error threshold, StFT-F maintains accuracy over least twice the time horizon compared the baselines.
Table I: Quantitative results of model and baselines in the autoregressive framework: relative L2 three spatiotemporal prediction systems, AR-StF'T refers to the deterministic model's results. AR-StFT-F denotes the probabilistic model with residual refinement of StFT. All models have been subjected to hyperparameter tuning to ensure fair and optimal performance comparisons.
Dataset
Variable(s)
AR-SIFTT
AR-S(FT-F AR-DPOT
AR-FNO
AR-VIT AR-UNet,
2.80e-2
2.240-2
1.0c-1
2.28e-1
1.73e-1
L.0cc0
d.A
2.45e-2
2.30c-2
8.36e-2
2.30e-1
3.24e-1
8.13e-1
6 Bi
3.05e-2
2.66e-2
8.98e-2
2.33e-1
1.95e-1
7.79e-1
Plasma MHD
8ne
2.8e-22
2.45e-2
8.61e-2
2.33e-1
2.08e-1
1.01e0
ön:
3.28e-2
2.93e-2
8.76e-2
2.33e-1
2.18e-1
1.04e0
duc
3.99e-2
3.73e-2
9.59e-2
3.18e-1
2.99e-1
6.96e-1
u
3.38e-2
3.30e-2
4.67e-2
4.46e-2
5.09e-2
6.16e-2
Navier-Stokes
3.60e-2
3.17e-2
4.52e-2
4.57e-2
4.60e-2
6.15e-2
P
5.16e-2
4.44e-2
6.18e-2
5.90e-2
7.03e-2
7.84-22
Shallow-Water
V
6.25e-2
6.53e-2
7.97e-2
9.53e-2
1.33e-1
2.02e-1
autoregressive step, and then inject their mean into the next step. Additionally, we assess the uncertainty quantification capability of StFT-F by sampling 100 trajectories for each test case. For each trajectory, at each autoregressive step, we generate a single prediction, which is subsequently fed into the next autoregressive step to iteratively forecast the full sequence.
4.3 Main Results
Table [ shows the forecasting performance of all the models on the three applications. We present several test trajectories and visualize their uncertainties across the forecasting time horizon estimated by StFT-F, as illustrated in Figure 5. StFT performs significantly better than all other baselines across all physical processes in the three applications. In the plasma MHD dataset, the test trajectory has a total of six coupled physics variables and 244 snapshots. On average, SUFT achieves a reduction in error by factor of three compared to the best baseline, AR-DPOT((AFNO) Although AR-FNO maintains high resolution in its long-term prediction, it fails to capture the correct dynamics of mode evolution, leading to out-of-phase predictions shown in Figure 4
We examine the error growth by plotting the L2 relative errors over time, as illustrated in Figure 3 for several representative variables. In the bottom figure for the shallow-water datasot, AR-FNO first appears slightly better than all other methods during the short term from timestamp 0 0.2, and StFT-F shows superior performance soon after. For the plasma dataset, StFT-F demonstrates dominance starting from i 0.2 with
8

---

Published in Transactions on Machine Learning Research (01/2026)
Ground
AR SIFT-F
AR
AR DPOT
AR FNO
AR
AR
Mean Relative Error
T-0
1.9
AR-S:FT AR-SUFTFF
T-59
0.5
0.0
0.2
0.4
0.6
0.8
1.0
Timestamps
Coarse scale
Meso scale
Fine scale
T-119
0.0
100
Contributing Factor (96)
T-179
-0.5
20
T-238
-10
Plasma MHD
Navier-Stokes
Shallow-water
Figure 4: Left top: Comparing StFT and STFT-F in Le relative error across timestamps for two plasma. MHD variables with uncertainty bands of StFTT-F. Left bottom: The contribution of each StFT block. From bottom to top, the patch size decreases. For this plasma data, the first two levels are sufficient to capture the multi-scale structures, whereas for Navier-Stokes and shallow-water equations, the finest scale contributes Right: Spatiotemporal evolution of perturbed electrostatic potential öΦ predicted by antoregressive models. StFT and StFT-F remain accurate past T 59, whereas baseline methods exhibit out-of-phase predictions.
significantly.
a stable performance, while the errors of all baseline methods begin to propagate from that point onward, resulting in a rapid decline compared to StFT-F, Notably, StFT-F exhibits long-term term stability relative to the other methods. As shown in Figure both of our methods accurately capture the dominant mode phase, and StFT-F generates predictions that are more closely aligned with the ground truth compared StFT. We also compare the error For the shallow-wate water and Navier-Stokes datasets, average, StFT reduces the errors by average of 27% and 25%, respectively. In Navicr-Stoke and plasma datasets, StFTT- not only surpasses StFT but also offers the additional capability of uncertainty quantification, ion, achieving error reductions of 10% in both cases. In the shallow-water dataset, we observe a slight increase in error with StFT-F. These results indicate that our method achieves superior long-term stability and accuracy among all other methods. More visualizations and results are included in Appendix
time comparing StFT and StFT-F. StFT-F begins to prevail from 0.6
4.4 Uncertainty Quantification
Figure 5 presents the distribution of the empirical standard viation along with the mean prediction. As observed, regions with large errors correspond to those exhibiting significant uncertainties predicted by StFT-F. Additionally, it is evident that uncertainties increase with time. This aligns with our expectation, errors accumulate during the autoregressive forecasting process. Besides the empirical evaluation relying on the standard deviation plots, we further measure the ro- bustness of StFT'ss uncertainty quantification using confidence intervals as shown in Table 2 Specifically, we use the predicted uncortainty to compute empirical coverage by measuring the proportion of ground truth values that fall within certain confidence interval around the predicted mean. This provides more rigorous evaluation of StFT's ability to capture uncertainties. Details about. CI coverage and results regarding each
standard de-
Table 2: Average confidence interval (90% and 95%) coverage.
Dataset
Cl: 90%
Cl: 95%
Plasma MHD
89.4%
92.5%
Shallow- Wattrr
89.5%
98.0%
9

---

Published in Transactions on Machine Learning Research (01/2026)
(b)
(c)
Figure 5: Evaluation of forecasting for three applications: ground truth, StFT-F prediction, residual, and uncertainty time shown for initial (left) and final (right) states. Variables include: (a) dΦ in plasma MHID (b) in Navier-Stokes, and (c) V in shallow-water equations.
For the 90% confidence interval, average coverage is very close to ideal value 90%, demonstrating that StFT-F well-calibratod around the 90% confidence interval.
variable are included in Appendix E
4.5 Ablation Study
Multi-scale and frequency path. To evaluate the effectiveness of the hierarchical structure and the frequency path in StFT, we conduct an ablation study employing mono-scale versus multi-scale models, and employing models with and without the frequency path. Experimental results are shown in Table 3a
Table 3: Ablation study of StFT. (a) Evaluates model structure and frequency path. (b) Shows performance impact of using more hierarchical scales.
(a) Effect mono-/multi-scale structures and frequency path
Model
Plasma MHD
Shallow- Water
Mono-scale
0.0805 0.105
2.3729 0.101 / 0.0975
Multi-scale
0.0404
0.0956
Convergence study on number of scales. We evaluate the performance of StFT on the number of scales for the plasma MHD dataset. As shown in Table 3b it is observed that increasing the number of scales leads to a reduction in prodiction error a certain point, with the two-layer con- figuration achieving the lowest error. The inclusion of additional scales results in slight error increase, indicating conver- gence in model performance. Additionally, conduct an ablation study on the overlapping tokenizer. More details are included in Appendix C
Multi-scalc + F
0.0307
0.0625
(b) L2 error vs. number of scales (plasma MHD).
# Scales L2 Rel. Error
1
0.0805
2
0.0307
0.0385 3
0.0391
4.6 Additional Results
Computational and Space Complexity. We compare StFT with other baselines regarding the inference FLOPS per sample, training time and inference time per batch, and peak memory usage during training. We also compare StFT and StFT without the overlapping tokenizer, and observed this enhancement comes with negligible impact compute cost, in Appendix D
Contribution of Each Scale. To assess the contribution of each StFT block for a specific scale in fitting Iylla the training data, quantify the weight of each block by: W, Tyllz where y; represents the prediction from the i-th StFT block, and y denotes the ground truth. We normalize the contributions, and present the
10

---

Published in Transactions on Machine Learning Research (01/2026)
A greater contributing factor from the fine-scale layer in StFT is observed in the Navier-Stokes and shallow-water equations, attributed to the sharper changes and smaller scale structures inherent in the dynamics of higher-order nonlinearities.
contributions in Figure 1
5 Conclusion
In this paper, we propose a spatio-temporal Fourier transformer (StFTT) for multi-scale and multi-physiee long-term dynamics forecasting. Specifically, cach StFT block tailored to address a particular spatial scale, and through a hierarchical composition of multiple StFT blocks spanning different scales, StFT learns the interplay between multiple scales and interactions between multiple physical processes, resulting in stable and accurate long-term dynamics forecasting in an autoregressive manner. Furthermore, we propose and demonstrate the use of generative residual correction mechanism, which enables meaningful quantification of uncertainties in the predictive model. Despite demonstrating superior forecasting ability in SciMI., the model is basod on regular grids, which constraints its applicability to irregular grids, and future work will focus on extending the framework to handle irregular geometries. The polential in improving its performance includes model parallelism across the multi-scale StFT blocks and extending StFT-F to end-to-end training.
Acknowledgements
ZB, DL, SW, and LO acknowledge support from the U.S. Department of Energy, Office of Science, Sci- DAC/ DAC/Advanced Scientific Computing Research under Award Number DE-AC02-05CH1122 with DL receiving additional support from the Margolis Foundation and MURT AFOSR under grant FA9550-20-1-0335 S7 was supported by MURI AFOSR grant FA9550-20-1-035, Margolis Foundation, and NSF Career Award IIS-2046295. This research used resources of the National Energy Research Scientific Computing Center (NERSC), U.S. Department of Energy Office of Science User Facility located at Lawrence Berkeley National Laboratory, operated under Contract Number DE-ACO2-05CHI122 1231, well as resources from the NCSA Delta GPU cluster from NCF Access program under award number IIS-2046295.
Broader Impact Statement
Turbulence remains one of the great unsolved problems in physics. Yet turbulence, whether driven by gravity, heating, magnetic fields, manifests in physical phenomena spanning multi-scale fluid flow to plasma fusion reactors to planetary atmospheres to convective layers in stars to the flow of interstellar gas in stellar nurseries that span trillions of kilometers. Where mathematics failed, data-driven machine learning models provide a pathway to understanding turbulence and gaining insights into not only the origins and ultimate fate of stars and planctary ecosystems, but also opens pathways to ncarly infinite sources of clean energy. Our StFT model presented in this paper is a slepping stone towards obviating the exponentially increasing computational costs of simulating geophysical systems (an inherently chaotic dynamical system requiring vast ensembles of high-resolution small time step simulations) as well realization of digital twins to aid in the design and operational control of tokamak- based fusion power plants.
References
Michael Albergo and Eric Vanden-Fijnden. Building normalizing flows with stochastic interpolants. arXiv preprint ar Xrx::2200.. 15571, 2022.
Zhe Bai, Xishuo Wei, William Tang, Leonid Oliker, Zhihong Lin, and Samuel Williams. Transfer learning nonlinear plasma dynamic transitions in low dimensional embeddings via deep neural networks. Machine Learning: Science and Technology, 6(2):025015, 2025.
Mathew Barlow, William Gutowski, John R Gyakum, Richard W Katz, Young-Kwon. Lim, Russ S Schumacher, Michael F Wehner, Lauric Agel, Michael Bosilovich, Allison Collow, et al. North american extreme precipitation events and related large-scale metcorological patterns: review of statistical methods, dynamics, modeling, and trends. Climate Dynamics, 53:6835 (18855-6775, 2019.
11

---

Published in Transactions on Machine Learning Research (01/2026)
Kaifeng Bi, Lingxi Xie, Hengheng Zhang, Xin Chen, Xiaotao Gu, and Qi Tian. Accurate medium-range global weather forecasting with 3d neural networks. Nature, 619(7970):533 -5385 2023.
Johannes Brandstetter, Daniel Worrall, and Max Welling, Message passing neural pde solvers. arXiv preprint arXiv:2202.0339 2022.
Keaton J Burns, Geoffrey M Vasil, Jeffrey S Oishi, Daniel Lecoanet, and Benjamin P Brown. Dedalus: A flexible framework for numerical simulations with spectral methods. Physical Review Research, 2(2):023068, 2020.
Qianying Cao, Somdatta Goswami, and George Em Karniadakis, Lno: Laplace neural operator for solving differential equations. arXiv preprint Xiv:220903 10528, 2023.
Shuhao Cao. Choose a transformer: Fourier or galerkin. Advances in neural information processing systems, 14:24924-24940, 2021.
Tianping Chen and Hong Chen. Universal approximation to nonlinear operators by neural networks with arbitrary activation functions and its application to dynamical systems. IEEE transactions on newral melworks, 6(4):911 ((4):911-9114, 997,, 1995.
Sai Hung Cheung, Todd A Oliver, Ernesto E Prudencio, Serge Prudhomme, and Robert Moser. Bayesian uncertainty analysis with applications to turbulence modeling. Reliability Engineering & System Safety, (9):1137 1149, 2011.
Alexey Dosovitskiy. An image is worth 16x16 words: Transformers for image recognition at scale. arXiv preprint Xrxi:22010. 11999, 2020.
Patrick Esser, Sumith Kulal, Andreas Blattmann, Rahim Entezari, Jonas Müller, Harry Saini, Yam Levi, Dominik Lorenz, Axel Sauer, Frederic Boesel, et al. Scaling rectified flow transformers for high-resolution image synthosis, 2024. URL https://arzio.. org/abs haddb//40..0/// 22000.0000006 1, 2024.
Haoqi Fan, Bo Xiong, Karttikeya Mangalam, Yanghao Li, Zhicheng Yan, Jitendra Malik, and Christoph Feichtenhofer. Multiscale vision transformers. In Proceedings of the IEEE/CVF international conference on compuler vision, pp. 6824 6835, 2021.
JP Graves, IT Chapman, S Coda, M Lennholm, M Albergante, and M Jucker. Control of magnetohydrodyy namic stability by phase space engineering of energetic ions in tokamak plasmas. Nature commumications, 3(1):624, 2012.
Gaurav Gupta, Xiongye Xiao, and Paul Bogdan. Multiwavvlett-asedd operator learning for differential equations. Advances in neural information processing systems, 34:24048 24062, 2021.
Jayesh K Gupta and Johannes Brandstetter. Towards malt--ppttioeemponlsssall generalized pde modeling. Xiv preprint arXiv:2209. 15616, 2022.
Zhongkai Hao, Zhengyi Wang, Hang Su, Chengyang Ying, Yinpeng Dong, Songming Liu, Ze Cheng, Jian Song, and Jun Zhu. Gnot: A general neural operator transformer for operator learning. In International Conference on Machine Learning, pp. 12556 12569. PMLR, 2023.
Zhongkai Hao, Chang Su, Songming Liu, Julius Berner, Chongyang Ying, Hang Su, Anima Anandkumar, Jian Song, and Jun Zhu. Dpot: Auto-regressive denoising operator transformer for large-scale pde pre-training. arXiu preprint ar Xiv:24 02..090.00333 03542, 2024.
Sheikh Md Shakeel Hassan, Arthur Feeney, Akash Dhruv, Jihoon Kim, Youngjon Suh, Jaiyoung Ryu, Yoonjin Won, and Aparna Chandramowlishwaran. Bubbleml: a multi-physics datasct and benchmarks for machine learning. XrXiv preprinl arXiv:2307.1466 14623, 2023.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advonces in neural information processing systems, 33:6840 08600,66611118 2020.
12

---

Published in Transactions on Machine Learning Research (01/2026)
Jonathan Ho, Tim Salimans, Alexey Gritsenko, William Chan, Mohammad Norouzi, and David J Fleet. Video diffusion models. Advonces in Neural Information Processing Systems, 35:8633 8646, 2022.
Pengzhan Jin, Shuai Meng, and Lu Lu. Mionet: Learning multiple-input operators via tensor product. SIAM Journal Scientific Computing, 14(6)-3340-A33544 2022.
Zhifeng Kong, Wei Ping, Jiaji Huang, Kexin Zhao, and Bryan Catanzaro. Diffwave: versatile diffusion model for audio synthesis. arXiv preprint orXiv:2009.0979 09761, 2020.
Katiana Kontolati, Somdatta Goswami, George Em Karniadakis, and Michael D Shields. Learning nonlinear operators in latent spaces for real-time predictions of complex dynamics in physical systems. Nature Communications, 15(1):1101, 2024.
Nicholas Krämer, Jonathan Schmidt, and Philipp Hennig. Probabilistic numerical method of lines for time-dependent partial differential equations. In International Conference on Artificial Intelligence and Statistics, pp. 625-639. PMLR, 2022.
SE Kruger, Jarrod Leddy, EC Howell, Sandeep Madireddy, C Akcay, T Bechtel Amara, J McClenaghan, LL Lao, David Orozco, SP Smith, et al. Thinking bayesian for plasma physicists. Physics of Plasmas, 31 (5), 2024.
Zhijie Li, Wenhui Peng, Zelong Yuan, and Jianchun Wang. Long-term predictions of turbulence by implicit u-net enhanced fourier neural operator. Physics of Fluids, 35(7), 2023a.
Zijie Li, Kazem Meidani, and Amir Barati Farimani. Transformer for partial differential equations' operator learning. aixiv preprint arXim:2005.1367 2022.
Zongyi Li, Nikola Kovachki, Kamyar Azizzadenesheli, Burigede Liu, Kaushik Bhattacharya, Andrew Stuart, and Anima Anandkumar. Fourier neural operator for parametric partial differential equations. arXiv preprint XrXi::2010. 08895, 2020.
Zongyi Li, Daniel Zhengyu Huang, Burigede Liu, and Anima Anandkumar. Fourier neural operator with learned deformations for pdes on general geometries. Journal of Machine Learning Research, 24(388):1-26, 2023b.
Chensen Lin, Zhen Li, Lu Lu, Shengze Cai, Martin Maxey, and George Em Karniadakis. Operator learning for predicting multiscale bubble growth dynamics. The Journal of Chemicol Physics, 154(10), 2021.
Yaron Lipman, Ricky TQ Chen, Heli Hamu, Maximilian Nickel, and Matt Le. Flow matching for gencrative modeling. arXiv preprint arXiv:2210.027: 2022.
Phillip Lippe, Bas Veeling, Paris Perdikaris, Richard Turner, and Johannes Brandstetter. Pde-refiner: Achieving accurate long rollouts with neural pdc solvers. Advances in Neural Information Processing Systems, 36:67398-6744 2023.
Qiang Liu. Rectified flow: A marginal preserving approach to optimal transport. arXiv preprinl arXiv:2209.14550 2022.
Xingchao Liu, Chengyue Gong, and Qiang Liu. Flow straight and fast: Learning to generate and transfer data with rectified flow. arXiv preprint arXiv:22000050 030003, 2022a.
Xinliang Liu, Bo Xu, and Lei Zhang, Ht-net: Hierarchical transformer based operator learning model for multiscale pdcs. 2022b.
Ze Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, and Baining Guo. Swin transformer: Hierarchical vision transformer using shifted windows. In Proceedings of the IEEE/CVF international conference on computer vision, pp. 10012-10022, 2021.
Da Long, Zhitong Xu, Guang Yang, Akil Narayan, and Shandian Zhe. Arbitrarll--conditioned diffusion for multi-physics emulation. arXio preprint axX:::2410113794 2024.
13

---

Published in Transactions on Machine Learning Research (01/2026)
Lu Lu, Pengzhan Jin, Guofei Pang, Zhongqiang Zhang, and George Em Karniadakis. Learning nonlinear operators via deeponet based on the universal approximation theorem of operators. Nalure machine intelligence, 3(3):218 229, 2021.
D Lucor, Xiu, C-H Su, and GE Karniadakis. Predictability and uncertainty in cfd. International Journal for Numerical Methods in Fluids, 43(5))483 505, 2003.
Michael McCabe, Peter Harrington, Shashank Subramanian, and Jed Brown. Towards stability of autoregres- sive neural operators. arXiv preprint arXiv:2306. 10619, 2023.
Vinay K Natrajan, Eiichiro Yamaguchi, and Kenneth T Christensen, Statistical and structural similarities between micro-and macroscale wall turbulence. Microftuidics and Nanoftuidics, 3:89-100, 2007.
Jaideep Pathak, Shashank Subramanian, Peter Harrington, Sanjeev Raja, Ashesh Chattopadhyar, Morteza Mardani, Thorsten Kurth, David Hall, Zongyi Li, Kamyar Azizzadenesheli, et al. Fourcastnet: A global data-driven high-resolution weather model using adaptive fourier neural operators. arXiv preprint ar X20: 11214, 2022.
William Peebles and Saining Xie. Scalable diffussio models with transformers. In Proceedings of the IEEE/CVF International Conference Computer Vision, pp. 4195-4205, 4205, 2023.
Norbert Peters. Multiscale combustion and turbulence. Proceedings of the Combustion Institude, 32(1):1-255 2009.
Adam Polyak, Amit Zohar, Andrew Brown, Andros Tjandra, Animesh Sinha, Ann Lee, Apoorv Vyas, Bowen Shi, Chih-Yao Ma, Ching-Yao Chuang, al. Movie gen: A cast of media foundation models. arXic preprin! arXiv:2410.13790 2024.
Michael Prasthofer, Tim De Ryck, and Siddhartha Mishra. Variable-input deep operator networks. arXiv preprint Xrv: 11404, 2022.
Md Ashiqur Rahman, Zachary E Ross, and Kamyar Azizzadenesheli, U-no: U-shaped neural operators. arXiv preprinl Xrx::22204.111 111227 2022.
Salva Rühling Cachay, Bo Zhao, Hailey Joren, and Rose Yu. Dyffusion: A dynamics dynamics-informed diffusion model for spatiotemporal forecasting. Advances in neural information processing systems, 36:45259 45287, 2023.
Sebastian Scher and Gabriele Messori. Predicting weather forecast uncertainty with machine learning. Quarterly Journal of the Royal Meleorological Sociely, 144(717):2830 2841, 2018.
Jaemin Seo, yaugKyeun Kim, Azarakhsh Jalalvand, Rory Conlin, Andrew Rothstein, Joseph Abbate, Keith Erickson, Josiah Wai, Ricardo Shousha, and Egemen Kolemen. Avoiding fusion plasma tearing instability with deep reinforcement learning. Nature, 626(8000):746 751, 2024.
Uriel Singer, Adam Polyak, Thomas Hayes, Xi Yin, Jic An, Songyang Zhang, Qiyuan Hu, Harry Yang, Oron Ashual, Oran Gafni, et al. Make-a-video: Texl-to-video generation without texl-video data. arXiv preprin! arXiv:2209.14792, 2022.
Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. arXiv preprint ar Wt:::2011.133000 2020.
Eitan Tadmor. review of numerical methods for nonlinear partial differential equations. Bulletin of the American Mathematical Society, 49(4): 494405507-554 2012.
Makoto Takamoto, Timothy Praditia, Raphael Leiteritz, Daniel MacKinlay, Francesco Alesiani, Dirk Pflüieer and Mathias Niepert. Pdebench: An extensive benchmark for scientific machine learning. Advonces in Neural Information Processing Systems, 35:1596-108 1611, 2022.
14

---

Published in Transactions on Machine Learning Research (01/2026)
Sibo Tian, Minghui Zheng, and Xiao Liang. Transfusion: A practical and effective transformer-bras based diffusion model for 3d human motion prediction. TEEE Robolies and Aulomalion Letters, 2024.
Alasdair Tran, Alexander Mathews, Lexing Xie, and Cheng Soon Ong. Factorized fourier neural operators. arXiv preprint arXiv:2111.1980) 2021.
A Vaswani. Attention is all you need. Advances in Neural Information Processing Systems, 2017.
Vikram Voleti, Alexia Selcocur-Martinees and Chris Pal. Mevd-maskod conditional vidco diffusion for prediction, generation, and interpolation. Advances in neural information processing systems, 35:23371- 23385, 2022.
Stodiek von Goeler, W Stodiek, and N Sauthoff. Studies of internal disruptions and m= oscillations in tokamak discharges with soft-x-ray tecniques. Physical Review Letters, 33(20):1201, 1974.
Bin Wang, Brian Hoskins, and Bin Wang. Large-scale almospheric dynamics. The Asian Monsoon, pp- 357-415, 2006.
Qi Wang, Pu Ren, Hao Zhou, Xin- Yang Liu, Zhiwen Deng, Yi Zhang, Ruizhi Chengze, Hongsheng Liu, Zidong Wang, Jian-Xun Wang, Ji-Rong Wen, Hao Sun, and Yang Liu. 2nec"2net: Pde-preserved coarse correction network for efficient prediction of spatiotemporal dynamics. In A. Globerson, L. Mackey, D. Belgrave, A. Fan, U. Paquet, J. Tomezak, and C. Zhang (eds.), Advances in Neural Information Processing Systems, volume 37, pp. 68897-6888 (8897-68925. Curran Associates, Inc., 2024a.
Sifan Wang, Hanwen Wang, and Paris Perdikaris. Learning the solution operator of parametric partial differential equations with physies-informet deeponets. Science advances, =ea)800588 2021.
Sifan Wang, Hanwen Wang, and Paris Perdikaris. Improved architectures and training algorithms for deep operator networks. Journal of Scientific Computing, 92(2):35, 2022.
Sifan Wang, Jacob H Seidman, Shyam Sankaran, Hanwen Wang, George J Pappas, and Paris Perdikaris. Bridg- ing operator learning and conditioned neural fields: A unifying perspective. arXiv preprint 2rX 2024b.
Julia Wolleb, Florentin Bieder, Robin Sandkühler, and Philippe C Cattin. Diffusion models for medical anomaly detection. In International Conference on Medical image computing and computer-assisted intervention, pp. 35 45. Springer, 2022.
Zizhao Zhang, Han Zhang, Long Zhao, Ting Chen, Sercan Ö Arik, and Tomas Pfister. Nested hierarchical transformer: Towards accurate, data-efficient and interpretable visual understanding. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 36, pp. 3417 3425, 2022.
A Appendix
B Problem Setup and Datasets
B.1 Plasma magnetohydrodynamiic (MHD) equations
We consider maguetolyydoddyammic (MHD) equations that characterize the plasma instabilities in fusion tokamaks. The coupled multi-physics system includes the continuity equation solving charge density őn, Poisson's equation solving do, the Ampere's law to solving dug, the Faraday's is law with the assumption Ej- to solving jAp; and the perpendicular force balance equation to solving 8 The first continuity equation
15

---

Published in Transactions on Machine Learning Research (01/2026)
for gyrocenter charge density is expressed as,
dôn
VBu
+
Bo
noV+
+ +B_
0t
Bo
Bo
Bo
Bo
(6P) 8P1) V
VoP
+
BB
Bo
(3)
iP
by
V
bo
bo
ba
VoBII
Po B
eBo
e
bu
VôB
Po = 0,
e
where n is the density, B is the magnetic field, uj is the parallel flow velocity, and P is the pressure. The perturbed quantities are denoted by with the equilibrium states including temperature, density, magnetic field and the flux surface from the reconstruction of DIII-D experiments. Here, ōn on. + q.0n:// qn stands for the difference ion and electron density, and őnn = òule + tgoupp// denotes the difference of ion and electron flow. We have v. = by x V (6P) + 8P_) (nom, where my is the electron mass, and Ωg = eBo/me is the electron cyclotron frequency. The perturbed electron parallel flow dup can be solved from Ampere's 's law,
duj
vi
μgeno
where Mo is the permeability of vacuum. dAp is the perturbed vector potential. In the single fluid model, Ej = 0 is assuned. Then öA can be solved from
0A
bo
VΦ,
(5)
Øt
and the electrostatic potential o can be solved from gyrokinetic Poisson's equation (the quasi-neutrality condition)
eôn
(6)
where is the speed of light, VA is the Alfvén velocity, and eD is the dielectric constant of vacuum. The parallel magnetic perturbation 8B is given by the perpendicular force balance,
8B B,
8P1 3e OPo deb
(7)
Po
2 txhu Po
The perturbed pressure in the fluid limit can be calculated by
OP. 8rh - 2 Bo
8B Po:
6P1
Ueo
(8)
OPo 8rh
8B Po-
8P Oibo
Bo
In these equations, vo and deh is the equilibrium and perturbed magnetic flux, and the evolution of deb is solved from
08b
de
99)
01
Ooo
where ag is from the Clebsch representation of B field, and Bg Vito x Vag. We run lincar gyrokinctic simulation with a 100 x 250 x 24 mesh in radial, poloidal and parallel directions. The time step is set to At 0.0058o/C. = 1.483 x 10 ss. We keep both n = 0, 1 modes, generate a trajectory 128, 000 time steps, and collect the data every 100 snapshots. We focus on emulating the dynamics of clectrostatic potential 86, parallel vector potential { electron number density ōn,, ion number density ōn,, and electron velocity du. in their trajectories.
16

---

Published in Transactions on Machine Learning Research (01/2026)
Ground truhh
AR DROT
Ground tnuth
AR SIFTF
AR SIFT
AR DROT
AR FND
AR VIT
AR UNet
T-119
T-119
T-238
T-238
Figure 6: Temporal evolution normalized perturbed parallel vector potential dAi and perturbed electron density on. contours predicted by different models: StFT-F, StFT, FNO, ViT and U-Net. Significant phase differences between the predictions of the models appear after 59, where StFT and StFT-F perform stable across the forecasting time horizon.
B.2 2D incompressible Navier-Stokes equations
We consider the 2D incompressible Navier-Stokes (NS) equation on a rectangular domain (z.y) E [0, 1]3,
Ou 0p
du
Ou
+
V2u + f(x,y),
Öt
dr
Oy
Re
Dv 0p
De
dv
1
V2 f(x,y),
(10)
oy
dr
Oy
Re
Du dr O,
dx
Oy
where u and u represent the velocity components in the x and y directions, and p represents the pressure. fxr y) is the source term, and we set it to e (00(((000.5)) (x-0.5) The Reynolds number is set to 1000. We run a finite difference solver to compute the solutions on a 50 50 spatial grid, with the temporal domain discretized into a total of 101 timestamps over [0, We generated total of 100 trajectories by sampling the four boundary conditions uniformly from (..1.
B.3 Spherical shallow-wate water equations
We consider the viscous shallow-water equations modeling the dynamics of large-scale atmospheric flows:
DV
Dt fk
-V Dh AV V +4 x ZE
(11)
(0
DI
where is the velocity vector tangential to the spherical surface, the unit vector normal to the surface, is the thickness of the fluid layer, f 2 sin the Coriolis parameter (E being the Earth's angular
velocity), g the gravitational acceleration, and v is the diffusion coefficient. The equations defined over
spherical domain Ω = (A, Φ), with longitude À and latitude ø
17

---

Published in Transactions on Machine Learning Research (01/2026)
AR UNet
AR VIT
SUFTFF
SIFT
DPOT
FNO
UNeb
ME T-1
SAMI
T-24
T-24
T-48
-00 T-48
-5000
T-72
T-72
T-96
.22 T-96
Figure 7: incompressible Navier-Stokes equation: pointwise error the predicted evolution velocity component and pressure contours across different models: StFT-F, StFT, FNO, VIT and U-Net. For long-term predictions, StFT and StFT-F demonstrate lower residuals compared to other models.
As an initial condition, zonal flow typical of a mid-latitude tropospheric jet is defined for the velocity component u as a function of latitude d:
0. Mmaa exp [0-2oll0--0
Φ S 00:
u(
Φ1
Φ z Φ1
where Umax the maximum zonal velocity, do and Φ1 represent the southern and northern boundaries of the jet in radians, and expl-4/(01 oo) normalizes Umax the midpoint of the jet. To induce barotropic instability, localized Gaussian perturbation is added to the height field, expressed as:
(
A
27
U2 si 27
)
M(A, 6.1 = 0) = c.co(()) exp
exp
B
where - < X < and parameters h, dh,, and 8 control the shape and location of the perturbation. The parameters a and 3 are sampled from uniform distributions ~ U[0.1, and 8 ~~ U[0.03, 0.2]. We ran the solver from Dedalus Burns et al. (2020) on a 256 x 256 spherical grid, and the temporal dimension is discretized into 72 timestamps. We have a total of 200 trajectories by sampling a and B.
C Ablation Study
C.1 The Hierarchical Structure and The Frequency Path
To evaluate the effectiveness of the hicrarchical structure and the frequency path in StFT, we conduct an ablation study the analysis of the joint influence and trade-off between patch size and the number of retained Fourier modes. First, only keep layer of StFT while changing the number of Fourier modes k. Second, we keep the hicrarchical structure, and tuning & in each hicrarchical layer.
Table shows the L2 relative errors averaged over all the variables. Note that I is the coarsest level, and la is the finest level. F stands for the frequency path StFT blocks. These results demonstrate the effectiveness of the hicrarchical composition StFT blocks and the frequency component in StFT block. With the hicrarchical composition, for Plasma MHD, the error drops 0.0307 from 0.0805, and for shallow-ward water equations (SWE), the error drops to 0.0625 from 0.0975 with k 8 retained modes in the Fourier space. We observe that the
18

---

Published in Transactions on Machine Learning Research (01/2026)
AR StFT-F
AR SiFT
AR DPOT
AR
AR VIT
AR UNetb
FNO
T=1
0.1
T-17
T-34
0.0
T=50
-0.1
T-67
Figure 8: Spherical shallow-water equations: pointwise error the temporal evolution velocity field predicted by all the autoregressive models: SUFT-F, SUFT, FNO, ViT and U-Net. The prediction error exhibits temporal growth trend, with our model StFT and StFT-F consistently demonstrate lower residuals over the forecasting time horizon.
(b)
10
AN
Figure 9: Additional evaluation of forecasting: ground truth, StFT-F prediction, residual, and uncertainty over time
shown for initial (left.) and final (right) states. Variables include: (a) one in plasma MHD, (b) p in Navier-Stokes.
19

---

Published in Transactions on Machine Learning Research (01/2026)
Model
(Multi-) Scale Patch Size
Fourier Modes in F
Datasct
==
== 8
k 16
MHD
4: (128,128)
0.0926
0.0805
0.0542
MHD
12: (64,64)
0.1206
0.1050
0.1131
MHD
41 + l2: ((128, (1,,,,11),,,,,,,
0.0544
0.0307
0.0578
SWE
1: 1228,
0.5110
2.5729
47.178
SWE
12: (64,64)
0.1123
0.1010
0.14744
SWE
1:: (32,32)
0.0984
0.0975
0.5455
SWE
4+4 12 + 1::
0.0766
0.0625
0.0829
Table 4: Ablation study results, We run models with combinations of StFT blocks and the retained Fourier modes in the frequency path. l is the coarsest level, and ly is the finest level. stands for the mumber of retained Fourier modes in the frequency path of StFT blocks.
fine-level layer setting in the SWE achieves the best performance among the single layer results, and the multi-layer settings further decrease the prediction error. Higher Fourier modes contribute substantially to the coarsc-layer representations in the MHD system, reflecting perturbations that are distributed over a broader spatial domain. In contrast, for the SWE problem, high frequency components are more spatially localized, predominantly confined to small regions. The interaction between refinement. level and Fourier modes indicate that the model performance generally improves with the hierarchical composition, while the retained Fourier modes must be tuned match the dominant frequency content present within each StFT block.
C.2 Convergence on the Number Scales
We evaluate the performance of StFT on the number scales for the plasma MHD dataset. Specifically, we run StFT with one scale (patch size of 128), three scales (patch sizes of 128, 64, and 48), and four scales (patch sizes of 128, 64, 48, and 32). As shown in Table 7 it is observed that increasing the number of scales leads a reduction in prediction error a certain point, with the two-layer configuration achieving the lowest error (0.0307). Beyond this, the inclusion of additional scales results in slight increase in error, indicating convergence in model performance, and further increasing the number of scales may not provide performance gain.
Table 5: Effect of the number of scales L2 relative error of the overall prediction the plasma MHD.
Number of Scales
L2 Relative Error
0.0805
0.0307
0.0385
0.0391
C.3 The Overlapping Tokenizer
We conduct an ablation study the overlapping tokenizer design, evaluating both predictive performance and computational cost, as reported in Tables 5 and 6 We evaluate the computational cost including the inference FLOPS per sample, Lraining/inference time per batch and peak memory usage. We used a fixed batch size of 20 for all models. The results are summarized in the following two tables. StFT-O stands for StFT with the overlapping tokenizer and StFT-NO for the opposite. All experiments were performed on an NVIDIA A100 GPU. Incorporating the overlapping tokenizer leads to substantial improvements in accuracy, in plasma MHD and shallow-water datasets, with errors reductions of 68% and 10% respectively.
20

---

Published in Transactions on Machine Learning Research (01/2026)
Notably, this enhancement comes with negligible impact on computational complexity, inference time and peak memory usage.
Table 6: Effect of the overlapping tokenizer on the prediction error of plasma MHD and shallow-water datasets.
Method
Plasma MHD
Shallow- Water
With overlapping tokenizer
0.0307
0.0625 0.0986
Without overlapping tokenizer
0.0700
Table 7: Computational complexity comparison StFT with and without overlapping tokenizer.
Dataset
Method
GFLOPs
Training Iter (s)
Inference / Iter (s)
Peak Memory (GB)
MHD MHD
SUFT-O 0.704
0.181
0.0455
9.49
StFT-NO
0.704
0.157
0.0454
9.49
SWE SWE
StFT-O StFT-NO
4.30
0.189
0.0319
12.2
4.30
0.189
0.0319
12.2
D
Computational and Space Complexity
We compare StFT with other baselines regarding the inference FLOPS per sample, training time per batch, inference time per batch, and peak memory usage (during training). In plasma MHD and shallow-water datasets, compared to FNO, StFT 20x smaller, and 8x smaller in computation (FLOPS) respectively. While achieving the highest prediction accuracy, StFT has the same order of magnitude of FLOPs as UNet and ViT in plasma MHD. the shallow-water equation (SWE), however UNet much more computationally expensive 11x more than StFT while its relative error is 3x larger than StFT.. StFT is also efficient in both training and inference time. For plasma MHD, StFT achieves 30% reduction in training time and 50% reduction in inference time compared to FNO. For SWE, SIFT is 40% faster in training and 60% faster than FNO in inference. Although StFTT incorporates dual paths operating in the frequency domain and the spatio-temporal domain, its peak memory usage remains comparable to that of FNO in the plasma MHD, and is reduced by 56% compared to FNO in SWE, and 26% less than UNet in SWE. Figure 10 compares the inference time and FLOPPs versus mean L, relative error for all models.
Table 8: Comparison of methods plasma MHD and shallow-water datasets regarding the computational and space complexity.
Dataset
Method
GFLOPs
Training / Iter (s)
Inference Iter (s)
Peak Memory (GB)
MHD
StFT
0.704
0.181
0.0455
9.49
MHD
DPOT
1.10
0.0676
0.0120
3.20
MHD
FNO
13.8
0.262
0.0900
10.1
MHD
UNET
0.462
0.0263
0.00230
1.38
MHD
VIT
0.338
0.0171
0.00237
1.65
SWE
SLET
4.30
0.189
0.0319
12.2
SWE
DPOT
1.61
0.0427
0.0133
8.43
SWE
FNO
34.6
0.325
0.0886
30.2
SWE
UNET
48.5
0.107
0.0257
16.5
SWE
VIT
0.90
0.0130
0.00347
7.45
21

---

Published in Transactions on Machine Learning Research (01/2026)
1
eR-UNet
AR-ViT
AR-FNO
AR-DPOT
AR-StFT
100
100
Mean L, Relative Error
Mean L, Relative Error
10-1
10
10-2
10-2
10-2
10 1
100
10
Inference Time (s)
GFLOPs
100
100
Mean L2 Relative Error
Mean L2 Relative Error
10
10-1
100
101
10-2
10-1
GFLOPs
Inference Time (s)
Figure 10: Comparison models in terms mean L2 relative error versus inference time (left) and FLOPs (right) for plasma MHD (top) and shallow-wate (bottom).
22

---

Published in Transactions on Machine Learning Research (01/2026)
E Signficance of UQ on StFT-F
For the Plasma MHD dataset, we measure the average coverage for the first 20 autoregressively predicted snapshots (each row represents one physical variable), as shown in Table In the last row, we report the average coverage across all physical variables. For the 90% confidence interval, average coverage is very close to the ideal value of 90%, demonstrating that StFT-F is well-calibrated around the 90% confidence interval. For the 95% confidence interval, the average coverage under-confident by 2.5%, suggesting the intervals may be slightly narrow. For the shallow-water dataset, we provide the results in Table 10 where we measure the average coverage for the first 10 predicted snapshots. For the 90% confidence interval, average coverage is also very close to the ideal value of 90%, which shows that StFT- is also well-calibrated for the shallow-water dataset.
Table 9: Confidence interval coverage on the plasma MHD dataset.
Variable
Cl: 90%
Cl: 95%
0.906
0.937
5.A
0.977
0.988
8B
0.900
0.930
ōne
0.933
0.962
8n;
0.867
0.910
due
0.781
0.823
Average Coverage
89.4%
92.5%
Table 10: Confidence interval coverage on the shallow-water dataset.
Variable CI: 90%
CI: 95%
V
0.895
0.980
Table 11: Training/validde validation/oeess data splits and the training budget for all models measured an A100 GPU.
Dataset
Total
Split (Train / Val Test)
Training Budget
Plasma MHD
traj. (1,221 snapshots)
927 50
244
24h
Navier-Stokes
100 (raj-(101 snaps/traj)
80 10 10
24h
Shallow-water
200 traj. (72 snaps/traj] traj)
160 / 20 20
48h
F
Experimental Details
Traiiin/valliaaionnn data sets. For the plasma MHD data, we split the trajectory of 1221 snapshots into a training set (the first 927 snapshots), a validation set (the middle 50 snapshots), and a test. set. (the last 244 snapshots). For the Navior-Stokes datasct, we have a total of 100 trajectorics (101 snapshots for each trajectory), and split them into 80 trajectories for training, 10 for validation, and 10 for testing, For the shallow-water dataset, have total of 200 trajectories (72 snapshots for each trajectory), and split them into 160 trajectories for training, 20 for validation, and 20 for testing. To ensure a fair comparison, we impose fixed training budget across all models. Specifically, a 48-hour limit mesaured on one A100 GPU was set. for the shallow-water equation dataset, while a 24-hour limit is applied to both the plasma MHD and Navier-Stokes equation datasets. Table II summarizes the datasets and the training budget for all models.
Generative residual correction block. We follow a two-step training protocol in training StFTTFF first, train St.T thoroughly with the training budget, and then we train the generative residual correction block
23

---

Published in Transactions on Machine Learning Research (01/2026)
for another 200 epochs. We employ rectified flow to learn distributions of the residuals given the prediction of StFT and the history snapshots. We implement a similar structure to the Diffusion Transformer (DiT) as the backbone model (Peebles & Xie] 2023). In each DřT block, we apply adaptive layer normalization before a self-attention layer and an MLP layer. We use adaLN-Zero for time conditioning. For the history snapshots uy and the prediction (ūe) of SIFT, these conditions are incorporated as extra input tokens.
Hyperparameters. For StFT on the plasma dynamics, 3D FFT is used to encode the spatio-tempore inputs in the frequency path. We use the patch size of 128 for the the first StFT block and and 64 for the second StFT block. The overlapping size is set to 1. The hidden dimension is set to 128. The depth for cach StFT block is set to 6. We keep the lowest 8 frequencies for cach spatial dimension. For the rectified flow block, the depth is set to 8, and the hidden dimension is set to 128. For the Navier-Stokes equation, StFT uses a patch size of 25 for the coarse block (the first block), 13 for the middle block, and 8 for the last block. The overlapping size is set to O, and the frequency path is not used. For each block in the hierarchical structure, the depth is set to 8, and we use hidden dimension of 512. In the rectified flow block, we use a depth of 4 and set. the hidden dimension to 128. For the shallow-water equation, three levels of SLFT blocks are employed, and their patch sizes are set to 128, 64, and 32 respectively. For each block, the depth is set to 6, and the hidden dimension is set to 512. We use the 2D FFT to encode the spatio-temporal inputs, and the lowest 8 frequencies are kept for each spatial dimension. The overlapping size is set to 1. For the rectified flow model, we use a depth of 8 and a hidden dimension of 128.
Table 12: Hyperparameter search range for each method.
Method
Hyperparameter Search
Hidden dimension: [256, 512] Patch size: [8, 16, 36]
DPOT
Depth: 6
Heads: 4
Modes: [16, 20, 24]
AR-FNO
Layers: [4.
Hidden dimension: 256
Hidden dimension: 2256,112
AR-VIT
Patch size: [16, 32, 64]
AR-UNet
Bottleneck hidden dimension: [64, 128, 256, 512]
Baselines. For all the baselines, we run all models with the same training time budget, as detailed in Table 11 For DPOT, we vary the hidden dimension from [256, 512), and the patch size from [8, 16,32], We set the depth 6, and the number of heads 4, which is the default setting from the author's implementation. For AR- the number of modes are selected through search in [16, 20, 24], the number of layers are searched in [4, 5], and the hidden dimension is set to 256. For AR-VIT, we vary the hidden dimension from [256, 512], and the patch size from [16, 116,32.64 64). For AR-UNet, the hidden dimension of bottlencck emboddings are searched (64, 128, 256, 512). Table 12 summarizes the hyperparameter choices for each method.
G More Visualization Results
Figure illustrates the ground truth and predicted temporal evolution of normalized perturbed parallel vector potential 84 and perturbed electron density on, in plasam MHD using AR-StFT, AR-StFT-F, AR-DPOT, AR-FNO, AR-VIT and AR-UNet methods. StFT and StFT-F perform stable across the forecasting time horizon. Figure X and g show the pointwise error of all the models compared to the ground truth data in the Navier-Stokes and shallow-water equations, where StFT and StFT-F demonstrate lower residuals compared to other baseline models.
In addition, we compare STFT with StFT-F for the autoregressive prediction across different timestamps for velocity u in Navier-Stokes equation, shown in Figure [1] A slightly larger crror is observed between = 0.4 and 0.6, which can be attributed to the training objective of StFT. As StFT is optimized using a pointwise loss function, it is encouraged to produce predictions that closely match the most likely outcome.
24

---

Published in Transactions on Machine Learning Research (01/2026)
x10-2
Mean L2 Relative Error
AR-StFT
AR-StFT-F
0.2
0.4
0.6
0.8
1.0
Timestamps
Figure 11: Results of comparing StFT and StFT-F for the autoregressive prediction in Le relative error across different timestamps for velocity u in Navier-Stokes equation. The shaded region represents the standard deviation distribution of the relative error of StFT-F. However, the uncertainty is negligible that it is not visually disceruible. SUFT-F demonstrates better performance in the latter stages the forecasting time horizon compared to SUFT.
In contrast, StFT-FF is designed learn the full distribution of the target, potentially introducing higher variance in its predictions. This distributional modeling, while beneficial for uncertainty quantification, may result in marginally increased errors as it captures characteristics beyond the mean behavior.
25


Reviews:
[{"review_id": "7GpvUG9sNP", "reviewer": "Reviewer_js98", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper introduces StFT, a spatio-temporal Fourier Transformer for long-term forecasting of multi-scale, multi-physics PDE systems. Each StFT block models dynamics at a distinct spatial scale using a dual-path design—one path operates in the Fourier domain and the other in the spatio-temporal domain. The blocks are composed coarse-to-fine to explicitly encode cross-scale interactions. Evaluations on plasma MHD, 2D incompressible Navier–Stokes, and spherical shallow-water show improved long-horizon accuracy.\n\nStrengths\n\n1. Clear, modular architecture that separates global (frequency) and local (spatio-temporal) processing; hierarchy aligns well with PDE multi-scale structure.\n2. Practical overlapping tokenizer to improve spatial continuity with negligible overhead (as shown in ablations).\n3. Uncertainty quantification via rectified flow yields useful, empirically calibrated intervals on multiple datasets.\n\nWeaknesses\n\n1. Plasma experiments appear to rely on a single long trajectory, limiting cross-condition generalization evidence (Table 11).\n2. The proposed method currently targets regular grids; authors note this as a limitation.\n3. Some tasks show slightly higher point error at certain horizons due to distributional modeling.", "claims_and_evidence_yn": "Yes", "claims_explanation": "The core architectural claims are well-substantiated: the paper explains the dual-path block, hierarchical composition, and overlapping tokenizer clearly, with algorithms and diagrams. Empirically, the autoregressive setup is consistent across methods; budgets and baselines are described, and long rollouts are used.", "audience_yn": "Yes", "audience_explanation": "TMLR’s readership includes researchers in scientific ML, operator learning, forecasting, and uncertainty quantification. The paper advances multi-scale neural operators with an interpretable hierarchy and introduces a pragmatic UQ mechanism for long-horizon rollouts, both of broad interest to this community.", "requested_changes": "1. The plasma evaluation is based on a single trajectory split into train/val/test. Please add experiments across multiple discharges/conditions if possible.\n2. Expand ablations to cover: (i) frequency truncation bandwidth, (ii) patch sizes/overlap hyperparameters, and (iii) number of scales beyond the brief convergence table.\n3. Consider adding (a) a non-autoregressive baseline (e.g., continuous-time emulator trained for multi-step), and (b) a physics-informed operator (if feasible). Current baselines are strong but largely data-driven.", "broader_impact_concerns": "No concerns.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "365yDZrDBv", "reviewer": "Reviewer_Cb2H", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper presents StFT (Spatio-temporal Fourier Transformer), a hierarchical transformer framework for long-term forecasting of multi-scale physical dynamics governed by PDEs. Each StFT block models global dependencies through a Fourier-domain path and local dynamics through a spatio-temporal path, capturing interactions across spatial and temporal scales. The method also introduces an overlapping tokenizer to ensure smooth patch transitions and a probabilistic extension, StFT-F, that employs flow matching for residual correction and calibrated uncertainty quantification.\n\nExtensive experiments on diverse PDE systems (magnetohydrodynamics, Navier–Stokes, and shallow-water equations) demonstrate substantial improvements in long-horizon accuracy and stability over strong Fourier and transformer baselines.", "claims_and_evidence_yn": "Yes", "claims_explanation": "The claims are well supported by clear and comprehensive evidence. Experiments on three PDE-based benchmarks show consistent gains in long-term accuracy, stability, and uncertainty calibration over strong baselines. Ablation studies validate the role of key components, and results are clearly presented. Minor details on computational cost could be expanded, but overall, the evidence convincingly supports the paper’s claims.", "audience_yn": "Yes", "audience_explanation": "The paper would interest TMLR readers working on scientific machine learning, spatio-temporal modeling, and physics-informed deep learning, as it advances long-term dynamics prediction and uncertainty quantification.", "requested_changes": "While the paper demonstrates strong predictive performance, it provides limited quantitative discussion of computational cost. Expanding this section would make the contribution more transparent and practically relevant. The authors should report runtime, memory usage, and parameter counts compared to key baselines (e.g., FNO, AFNO, and transformer variants) under identical settings. A scaling analysis showing how training and inference time grow with grid resolution, sequence length, and the number of StFT layers would also help readers assess the method’s efficiency for large-scale or real-time simulations. Including these comparisons, either in the main text or supplementary material, would strengthen the empirical rigor of the paper.", "broader_impact_concerns": "NA", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "wFsjTkfIf8", "reviewer": "Reviewer_yNL3", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper introduces StFT (Spatio-temporal Fourier Transformer) and its probabilistic extension StFT-F for forecasting the evolution of spatial fields and uncertainty quantification. The model block combines a frequency path to capture low-frequency structure and a spatio-temporal path for local details. Predictions are produced from the past k snapshots of the variables and refined hierarchically from coarse to fine scales. Additionally, a rectified flow-matching block is trained on the residuals between the deterministic prediction and the true next state, enabling uncertainty quantification. Experiments on three spatio-temporal multi-physics systems compare StFT-F with autoregressive baselines such as AR-FNO, AR-UNet, and AR-ViT, AR-DPOT.\n\nStrengths\n-  The multiscale hierarchy can help reduce compounding error.\n-  The model leverages spectral and spatial representations. \n-  The rectified residual block provides additional refinement and uncertainty estimation.\n-  Numerical results show improved performance compared to benchmarks.\n-  The paper is easy to follow overall.\n\nWeeknesses\n-  Baselines do not include recent models (e.g., PDE-Refiner or Dyffusion) despite being mentioned in the introduction. \n-  Multiple overlapping components (spectral, temporal, diffusion) make it difficult to isolate the contribution of each.\n-  The interplay between spatial resolution in the hierarchical refinement and the number of retained Fourier components is not clear.\n-  Uncertainty quantification is performed by sampling from the residual network, which can yield a different mean prediction than propagation of the deterministic mean.\n- Results on uncertainty calibration are limited, no comparison with recent models with UQ.", "claims_and_evidence_yn": "No", "claims_explanation": "- Limited baseline coverage.  The baselines do not include recent models with uncertainty quantification (e.g., PDE-Refiner, DeepONet-Grid-UQ), even though such approaches are discussed in the introduction. As a result, the comparison focuses on standard deterministic autoregressive models and may not fully reflect the current state of the field.\n- Unclear contribution of individual components.  The model combines spectral, spatial–temporal, and diffusion modules, but their respective roles and interactions are not clear. \n- Interplay between spatial resolution and Fourier truncation.  The relationship between the spatial resolution used in the hierarchical refinement and the number of retained Fourier components is not analyzed. Since both affect the model’s effective bandwidth, their joint influence on stability and accuracy is unclear.\n- Uncertainty estimation based on a separate rollout regime.  Uncertainty quantification is performed by sampling from the residual network, which can yield a different mean trajectory from the mean propagation.  \n- Limited analysis of uncertainty calibration and temporal dependence.  It is only quantified in terms of coverage %, e.g. no ECE is given.  Also it is not clear (qualitatively) how it propagates in time.", "audience_yn": "No", "audience_explanation": "While the empirical results are promising, the study omits comparisons with established methods. The work requires substantial revision and broader evaluation.", "requested_changes": "- Include proper comparison with methods (e.g., PDE-Refiner, Dyffusion, see above).\n- Explain the relation between patch size, refinement level, and number of retained Fourier modes. \n- Specify that the flow-matching residual refines one-step predictions; UQ sampling may shift the mean.\n- Add ablations isolating spatial, spectral, and residual components.\n- Provide quantitative UQ and calibration analysis over time and comparison with other methods.", "broader_impact_concerns": "no concerns", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Decision: Accept (poster)