Physics-informed Temporal Alignment for Auto-regressive PDE Foundation Models
Authors: Congcong Zhu
Xiaoyan Xu
Jiayue Han
Jingrun Chen
Venue: ICML 2025 poster (2025)
Keywords: PDE foundation model
Physics-informed constrain
Inverse problem
Self-supervised learning
Primary area: deep_learning->foundation_models
TL;DR: A self-supervised learning framework for solving the error accumulation issue in PDE foundation models

Abstract: Auto-regressive partial differential equation (PDE) foundation models have shown great potential in handling time-dependent data. However, these models suffer from error accumulation caused by the shortcut problem deeply rooted in auto-regressive prediction. The challenge becomes particularly evident for out-of-distribution data, as the pretraining performance may approach random model initialization for downstream tasks with long-term dynamics. To deal with this problem, we propose physics-informed temporal alignment (PITA), a self-supervised learning framework inspired by inverse problem solving. Specifically, PITA aligns the physical dynamics discovered at different time steps on each given PDE trajectory by integrating physics-informed constraints into the self-supervision signal. The alignment is derived from observation data without relying on known physics priors, indicating strong generalization ability to out-of-distribution data. Extensive experiments show that PITA significantly enhances the accuracy and robustness of existing foundation models on diverse time-dependent PDE data. The code is available at \url{https://github.com/SCAILab-USTC/PITA}.

Physics-informed Temporal Alignment for Auto-regressive PDE Foundation Models
Congcong Zhu 122
Xiaoyan Xu
Jiayue Han Jingrun Chen 523
Abstract
Library
0,0700 409
Auto-regressive partial differential equation (PDE) foundation models have shown great po- tential in handling time-dependent data. How- ever, these models suffer from error accumula- tion caused by the shortcut problem deeply rooted in auto-regressive prediction. The challenge be- comes particularly evident for out-o--distribution data, the pretraining performance may ap- proach random model initialization for down- stream tasks with long-term dynamics. To deal with this problem, we propose physics-informed temporal alignment (PITA), self-supervised learning framework inspired by inverse problem solving. Specifically, PITA aligns the physical dynamics discovered at different time steps on each given PDE trajectory by integrating physics- informed constraints into the self-supervision sig- nal. The alignment is derived from observa- tion data without relying on known physics pri- ors, indicating strong generalization ability to -ut-of-distributio data. Extensive experiments show that PITA significantly enhances the accu- racy and robustness of existing foundation mod- els on diverse dieeneepndentt PDE data. The code is available at https: :/ //////g/ //ggttttb..ooo com/ SCAILab-OSTC/E
Idaanified Dynamies
Foundation Model
Cooppeess
Downt-
Idianiifaal
Merge
sumple
Compressed
Dynamics
11
(a) Auto-regresshve Prediction
(b) Physics informed Temporal Alignment
Figure 1: Insight of the proposed framework. Existing (a) Auto-regressive prediction may cause error accumulation for long-term PDE data. (b) Physics-informed temporal alignment is integrated to handle this problem.
& Kutz, 2024), PDE foundation models (Shen et al., 2024; Song et al., 2024; Gupta & Brandstetter, 2022) have gained increasing prominence in scientific machine learning. These foundational models leverage neural operators as surrogate models perform PDE trajectory prediction, which runs significantly faster than the traditional numerical solvers. Among the various PDE foundation models, auto-regressive prediction, i.e., predicting the next behavior based on past dynamics data, is one of the most promising strategies for training or pretraining, which aims to endow the models with generalization ability to downstream tasks (Geneva Zabaras, 2020). However, it is observed that auto-regressive prediction may introduce deeply-rooted shorteut problem, where the model chooses simple and mendacious solu- tion to approximately satisfy the optimization objective by duplicating previous dynamics. More details are discussed in Appendix B. Since the duplicating operation is easy for neural networks to learn, the model may overlook capturing long-term dynamics, resulting in error accumulation. There- fore, the long-term accurate prediction of auto-regressive PDE foundation models remains a significant challenge when facing downstream tasks with long trajectory data.
1. Introduction
With the ongoing advancement in computational capabilities and data-driven methodologies (Cicirello, 2024; Brunton
"Equal contribution 'School of Artificial Intelligence and Data Science, University of Science and Technology of China Suzhou Institute for Advanced Research, Unive rsity Science and Tech- nology of China 'Suzhou Big Data & AI Research and Engineering Center DDepartment of Radiation Oncology. University of Kansas Medical Center School of Mathematical Sciences. University of Science and Technology China. Correspondence Jingrun Chen <jingrunchen ustc.edu.en.>.
As we know, all trajectories of PDE data are simulated based on the underlying physical dynamics. Thus, the unique un- derlying physical equation can govern any segment in the corresponding PDE trajectory. This property naturally facil-
Proceedings of the ndnd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Physics-informed Temporal Alignment for Auto-regressive PDE Foundation Models
itates embedding physical constraints into the PDE trajec- tory prediction. Early research such as (Raissi et al., 2019; Cuomo et al., 2022) utilizes known physics information to construct optimization objectives, enabling precise solutions for specific PDE trajectories. However, this approach may overfit the training data with known physics laws, prevent- ing the models from serving as foundation for downstream tasks. To this end, some foundation models (Song et al., 2024; Cao et al., 2024; Zhou et al., 2024; Sun et al., 2024) introduce PDE symbols as a condition prior to encoding into embedded representations by neural networks, thus guid- ing the prediction results. Although these methods flexibly encode various known PDE symbols to impose physical constraints on the prediction results, they merely perform conditional mapping based on known equation priors rather than capturing the intrinsic dynamics that represent PDE tra- jectory data. They may encounter significant performance degradation for downstream tasks without known physical priors. Therefore, exploring universal physical constraints is critically important for PDE foundation models.
way, PITA seamlessly unifies data-driven forecasting with physics-based modeling, offering a versatile framework for reliable long-term dynamics prediction.
2. Related Work
2.1. PDE Foundation Models
PDE foundation models have shown immense potential in many fields, such as fluid dynamics (Liu et al., 2024b; Luo et al., 2023), geophysics (Liu & Ma, 2024), solid mechan- ies (Yizheng et al., 2024), and chemical reactor modelling (Wang & Wu, 2024). Recently, PDE foundation models have been proposed to enable operator learning across dif- ferent PDE families, incorporating pretraining and finetun- ing to enhance generalization. For instance, PIMRL (Wan et al., 2025) introduces a two-stages multi-scale learning framework that leverages multi-scale data for spatiotempo- ral dynamics prediction. PROSE (Sun et al., 2024) applies an encoder-decoder framework to integrate numerical data and equation embeddings. while PROSE-FD (Liu et al., 2024b) extends this approach to develop a foundation model for fluid dynamics. Moreover, MPP (McCabe et al., 2023) adopts an auto-regressive approach for pretraining and fine- tuning on time-dependent PDE datasets. To further investi- gate the generalization of auto-regressive prediction, DPOT (Hao et al., 2024) introduces a denoising pretraining strategy to improve transferability for downstream tasks. Despite the aforementioned advancements, there are many other works (Yang et al., 2023; Subramanian et al., 2024; Song et al., 2024; Ye et al., 2024b;a; Chen et al., 2024; Cao et al., 2024; Shen et al., 2024; Zhao et al., 2023) contributing to the development of PDE foundation models across various domains.
Inspired by solving inverse problem (Rudy et al., 2017: Stephany & Earls, 2022), we thought over how to derive unknown dynamics systems from observation data and trans- form them into physics-informed constraints during auto- regressive prediction. To achieve this, we propose a self- supervised learning framework named physics-informet temporal alignment (PITA), which integrates governing laws discovery into ruto-regres prediction, as shown in Fig- ure 1. In detail, PITA first discovers the governing PDE equation given the initial values with an unknown dynamics system by time series measurements. This process identi- fies the key derivative terms and parameters that form the structure and explicit expression of the PDE. Then, the auto- regressive prediction results are merged into the initial value sequence and used to rediscover the governing equation. Since the initial value sequence and the merged sequence are represented by the same underlying physical law, any discrepancy between the discovered governing equations corresponding to them indicates that auto-regressive predic- tion results may contain dynamics deviation. Based on this insight, we utilize the discrepancies of physical laws discov- ered from observation data and ruto-regressive predictions to supervise the ruto-regressive PDE models. Moreover, the discovered governing equation can also serve as physies- informed regularization, added to the optimization objective to supervise the auto-regressive predictions to follow the physical laws described. Finally, leveraging uncertainty- based weighting and alternating direction optimization, the physical supervisions derived from observation data are in- tegrated into data-driven auto-regressive prediction. This ensures that the predictions follow both the observed tem- poral dynamics and the underlying physics laws, making PITA applicable to diverse time-dependent PDE data. In this
2.2. Shortcut Problem and Error Accumulation
Shortcuts commonly emerge as decision rules allowing mod- els to excel on standard benchmarks while struggling to generalize under more complex testing conditions. This dis- crepancy between intended and learned solutions (Geirhos et al., 2020) leads to severe error accumulation, particu- larly in long temporal sequences. Several frameworks have been proposed to address shorteut learning. For example, the LTGR framework (Du et al., 2021) prevents overconfi- dent predictions shortcut samples, while COMI (Zhao et al., 2024) reduces the model's reliance on shortcuts by integrating standard empirical risk minimization to enhance its ability to extract underlying information. Additionally, ietuuiionbbsed and auto-regressive generative classifiers have been proposed to address shortcut issues by modeling both causal and spurious features (Li et al., 2024). Other works (Luo al., 2021; Brown al., 2023; Du et al., 2021; Dagaev et al., 2023; Robinson et al., 2021; Hermann et al., 2023; Scimeca et al., 2023; Chuah et al., 2022) have also
2

---

Physics-informed Temporal Alignment for Auto-regressive PDE Foundation Models
explored various strategies detect and rectify shorteut learning.
tiotemporal domain, constructing tailored function libraries, and employing sparse regression to identify dominant dy- namics. Finally, the loss functions comprising data loss, physics loss, and consistency loss are introduced in Sec. 3.4. Together, these losses guide the model to balance among learning from observed data, enforcing physical constraints, and ensuring temporal coherence.
2.3. Data-driven Inverse Problems
Modeling complex dynamical systems has traditionally relied on PDEs. Recently, data-driven techniques have emerged to uncover the underlying PDEs efficiently (Long et al., 2018; Brunton & Kutz, 2024). Sparse regression methods have leveraged finite differences to approximate algo- rithms such as the Douglas-Rachtoon algorithm (Schaeffer, 2017) and Sparse Threshold Ridge Regression (STRidge) (Rudy al., 2017) identify PDE coefficients.
3.2. Auto-regressive Prediction
partial derivatives and employed
In the auto-regressive prediction, a neural operator Ge. pa- rameterized by weights 0, takes Tin frames as input and predicts the next frame based on the previous frames:
de+1 Go (tcce)
Physics-informed neural network (PINN) (Raissi et al., 2019) has been integrated into PDE discovery to address these limitations. For instance, DeepMoD (Both et al., 2021) combines PINNs and SINDy by introducing regularization terms into the PINN loss function. PINN-SR (Chen et al., 2021) adopts a similar loss function structure and introduces an alternating direction optimization training strategy.
where uiz represents ( When t= Tin the input sequence denotes the ground truth frames, ensuring that the model starts with accurate information. 0 denotes the parameters of the neural operator. To predict the (t + 2)- th frame, the neural operator uses Tin frames along with the predicted frame út-1 as input. By iteratively performing this step, the model rolls out time window of length Tar. which can be customized. In our experiments, the roll-in window length is set to T.n 10, and the roll-out window length is set to Tar = 1 or Tar = 10. However, directly applying this strategy may lead to the shortcut problem, causing the accumulation of errors propagated across the time window. Our framework employs physics-informed temporal alignment strategy to address this issue.
3. Methodology
3.1. Overview
The general form of PDEs is considered as follows:
Ou(x,,) t) F(u)(x, t),
Ot u(x 0) =,(x), [2u] t) = O, x E 02,
x x Ω
3.3. Governing Equations Discovery
Following the uuto-regressive predictions, the process of PDE discovery is employed to identify the underlying phys- ical laws from the data while ensuring temporal alignment. First, the data is downsampled in both the spatial and tem- poral domains to enhance computational efficiency. Next, a nonlinear library is constructed, incorporating candidate terms that may represent the underlying physical dynamics. Finally. sparse regression is applied to derive the funda- mental form of the PDE. Note that downsampling is not employed for the initial input sequence, so as ensure the accuracy of the PDE discovered from ground truth data.
where xc E Ω Rd denotes the spatial variable, [0. x Ω + gdd is the solution of time-dependent PDE. Fx,) = F(t, x, u. uzu, U.. ) is a differential operator with spatial derivative terms. The initial condition is given by uo(z): : .2 + Rd,, and the boundary condition is defined by the operator B.
We propose PITA, physics-informed temporal alignment strategy for PDE foundation models. PITA integrates auto-regressive prediction and PDE discovery into a self- supervised learning framework, as shown in Figure 2. This physics-informed alignment ensures that the predictions conform to both the observed temporal dynamics and the underlying PDEs, thereby enhancing predictive accuracy and maintaining physical fidelity.
Downsampling Data: Before obtaining the PDE coeffi- cients, the data generated through the auto-regressive pre- diction is downsampled to reduce computational time. For each ue we assume that the function u, is discretized over n points, represented as X,, {01, 1, ,,, Jn a, R. In the spatial domain, the grid is randomly sampled to re- duce the number of points to one-quarter of its original size. The sampled grid is denoted as XT, = {zx, Li, D... }. where w = [n/4] and {i,, 1 S f1,2,... n) rep- resents the set of sampled indices. The selection of w will be discussed in Sec. 4.4. The derivatives are computed
The auto-regressive prediction, detailed in Sec. 3.2, enables iterative future state predictions by progressively leveraging prior outputs. To mitigate error accumulation inherent in the ruto-regressive approach, PDE discovery is incorporated as discussed in Sec. 3.3. This integration encompasses several key steps, including compressing data across the full spa-
3

---

Physics-informed Temporal Alignment for Auto-regressive PDE Foundation Models
Nonlinear Library of Data and Derivatives
8,0 - (0)
Sparse Regression
0./10)
Consistency Supervision
Back-
[dentified Dynamics
ADO
Propagation
Anto-regressive PDE Foundation Model
Compress
Identified Compressed
Dynamics
Physics Supervision
- (
Uncertainty-Rueed Weights
Compressed Spanse Regression
sample
Merge
Duta
May,
Supervision
Myy
Compressed Library
Figure 2: Work flow of PITA. The proposed framework integrates auto-regressive prediction and PDE discovery with self-supervised learning: (1) The pretrained PDE model takes the initial temporal states {ur)z=1 input and predicts future states (ú, Taar as output in an auto-regressive manner; (2) Data-driven PDE discovery is then performed on the compressed imput sequence to infer the governing equations. Temporal alignment is achieved by matching the discovered physical laws from predictions with those obtained from the ground truth sequence; (3) The loss function consists of three parts, i.e., data loss Datas physics loss Cphy, and consistency loss Coon with uncertainty-based strategy employed adjust the weights dynamically.
using a small number of spatially localized points near each measurement position via polynomial interpolation (Rudy et al., 2017). In the temporal domain, the most recent Tc frames are retained, enabling the capture critical temporal information that more relevant than the entire sequence. As shown in Figure 2, the diagonal matrix D signifies the and Dij for + j. where non-zero diagonal elements indi- cate retaining data. For simplicity, the matrix D is omitted in the following contexts.
Sparse Regression: After constructing the libraries for the PDE, the equation can be expressed as
a = (
(3)
where (50) is the library of candidate terms, and RSxC represents the sparse coefficient matrix. Assuming that only few key terms dominate the underlying dynamics, the discovery problem is formulated as sparse regression to identify the sparse coefficient matrix A:
downsampling process, with Dij E {0, 1) for
= arg min (0) a, 112 + || (4) A
Building Libraries of Candidate Terms: Before con- structing libraries, a preparatory step involves rearranging all the compressed data ÚA) E RexmxC into matrix U(O) (=xm)xCC which represents C physical variables collected over spatial locations and time points.
This optimization problem is addressed using sparse re- gression (Rudy et al., 2017: Ma et al., 2023; Zhang, 2008). For example, if C = / where U = [U1,- ... U A =
A.,... Ay1. where U; E ((xm))x A E R5x1 Equa-
tion (3) can be rewritten as
Next, we construct a library (0) E (xxx)) consisting of 5 predefined candidate linear and nonlinear terms, along with partial derivatives for the PDE (Rudy et al., 2017). For instance, (00) may include time derivatives, spatial derivatives, and n-th degree polynomial terms (Rudy al., 2017), assembled in matrix represented by
(0, O,U,(0)) 20) [A
Ar). (5)
Thus the coefficient matrix A is solved iteratively through each X, as detailed in Algorithm 1. This approach in- tegrates ridge regression with hard threshold to enforce sparsity, striking a balance between model complexity and accuracy.
= , ((0))_U()) 11. (2)
3.4. Integrating Data and Physics Supervisions
where finite differences are used to compute the derivatives. Each column of (0)) contains all the values of partic- ular candidate function across all (n x m²) space-time grid points where the data are collected.
The total loss function integrates three parts: the data loss Datas the physics loss L and the consistency loss LCon-

---

Physics-informed Temporal Alignment for Auto-regressive PDE Foundation Models
Algorithm 1 Sparse regression for Equation (4) and (5). Input: Time derivative vector (,U, candidate func- tion library matrix (()), threshold tolerance B, maxi- mum iteration number K. Initialize: X; t (0)a,U,()), 1, I.k for k < K do
the different learning tasks of the model. Determining the penalty coefficients for each term is challenging, model performance across tasks strongly depends on the relative weighting of each loss (Sener & Koltun, 2018; Chen et al., 2018). Instead of relying on traditional hyperparameter tuning, we employ an uncertainty- bbasdd multi-task learning strategy (Kendall et al., 2018; Liebel & Körner, 2018) to adjust these weights dynamically.
Determine two groups of indices of coefficients in A
P = {o E P: A < 3), Q = {u E Q [ > 83.
Impose sparsity on small values:
Let ó {ō1, 52, d5} be the parameters for the weights of the three losses in our task. The weighted total loss function becomes
Update A.:
arg min $ -,, 2 Update k= 1
Data (0) + PPyy (0)
269
282
(9)
end for
Output: The best solution X; = X U.A
+ 26
Lcon (0) + log 81 81828..
These parameters are updated simultancously with the model parameters during training. Minimizing this objective with respect to 61, 82, and 63 allows the model adaptively learn the relative weights of the three losses based on the data, which enables robustness across varying datasets.
Data Loss: The data loss is defined
DData (0)
1
Ti
0 T., (0)ll2
Tar
1..TT..l l2
where ú.,TT,,(0 (0) represents the predicted solution and Ui+T; is the ground truth data. This term quantifies the discrepancy between predictions and observations. Previous works have employed similar approaches such as (Raissi al., 2019; Chen et al., 2021: Both et al., 2021).
3.5. Alternating Direction Optimization
The total loss function described in Equation (9) exhibits an implicit and complex form, which complicates the direct resolution of the optimization problem due to the lu regular- ization, rendering it NP-hard (Chen et al., 2021). Although relaxing the Eo term through the less stringent & regulariza- tion enhances well ool-possessess and allows for optimization in continuous space, this may lead to false-positive identifica- tion, hindering the accurate realization of the sparsity of the PDE coefficients (Berg & Nyström, 2019; Both et al., 2021). Inspired by (Chen et al., 2021), we employ an alternating direction optimization ( (ADO) algorithm to decompose the overall optimization problem into a series of manageable subproblems, enabling the sequential optimization of 0 and A over several alternating iterations (denoted as k). For instance, in the (k 1)-th alternating iteration, the sparse coefficient matrix A in Equation (11) is updated to A using sparse regression, based on the parameters 02. of the neural operator Go from the previous iteration:
Physies Loss: Inspired by (Rudy et al., 2017; Brunton et al., 2016; Chen et al., 2021; Raissi et al., 2019), the physics loss is expressed as
Tor (0) (.) (0)A.11 +a |A (7)
where the residual term ensures that the product of the com- puted coefficients A; and the function library (0) ap- proximates the time-dependent term (O). The to norm enforces sparsity in the coefficients.
Consistency Loss: The consistency loss defined as
Tar
CCon(0) (0)
Σ - 1,00)
(8)
A = arg min Φ 34- O,U 0 1gt |4 (10)
A
where A denotes the true coefficients from ground-truth data and (0) are from the i-th time window, both obtained via sparse regression as detailed in Algorithm 1. This loss enforces consistency between the discovered physics laws represented by (0) and real physics laws represented by ^ To the best of our knowledge, it is the first work to use time alignment as a constraint for PDE foundation models.
Next, the parameters 0 in the current iteration then updated to 0x+1 based the solved AR+1,
1
0 0.+1, := arg min
Cata (0, (0,4811)
20
0.5
+
Cpyy (0, (0,4.+1)
(1I)
28
Unnertainty-Basee Weighted Loss Function: The total loss function combines these three supervisions to balance
+ Lcon (0, (044) + log 82 8263].
28
5

---

Physics-informed Temporal Alignment for Auto-regressive PDE Foundation Models
Table I: Comparison of PITA and existing auto-regressive strategies with Tor = 1. The underlined results represent the predictions of PITA, while the bold results indicate the optimal outcomes for the same model. '-' means that the result is unavailable. Notably, five datasets exhibit long trajectory characteristics, with detailed descriptions provided in Appendix C.
Finetune Sirategy
END
FDLBeach
PDEBennb-CCNS/ CNS-(M.
FDEArena
CFDBench
Model
le:5
le-1
1e-3
DR
SWE
1,0.1
1,0.01
M1
0.1.0.1
0.1.0.11
M0.1
NS-Force
NS
Preduct whvet injeclory
DFOFH
Auto regress
0J05200
0.00355
0.00350
0.01326
0,00206
0,011200
0,019550
0,61535
0.01440
0.01550
0001500
0.10266
0.09100
0,00391
7M
PITA
0,06629
0.00499
0,00218
0.02119
0,00202
0.01793
0.01565
0.01679
0.01080
0.01250
0.01165
0.11988
0.06465
0,00360
DPOT-S
Auto-regiess
0.03200
0.00641
0.0030T
0.01239
6,00210
0,01250
0.01670
0.01480
0.0152
001260
0.01390
0,07598
0.08670
0.003882
RIM
PITA
0.02240
0.00232
0.00157
0.01346
640150
0.01057
0/01217
0.01157
001750
0.01140
0.01455
0.06769
0.84190
MPP-TI
Aute-regress
0,03510
0.00645
0.0554]
0.0161
7M
0.03364
0.64370
0.04225
MPP-S
Auto-regress
0,02944
0,00281
0.04257
0.02012
30M
PITA
0.9902
0/00274
0.64051
0.01906
DPOT-M
Auto regress
0,02290
0.00385
0,00297
0,11200
0.00998
0,01460
0,01230
0.06100
0.00947
0.01280
0.05571
0.02940
0.00373
122M
PITA
0.03199
0.00185
0.00193
GA0054
0-01100
0:016633
0.01067
0.01010
0.00958
0,05829
0,02191
0.002299
FNO-M
Auto-Pogress
010836
0.005 77
0:00202
000000
0.22559
0.03345
0,15557
011200
0.04449
0106255
0.15196
6,15590
0:01382
170M
PITA
0J97668
0.00552
0.00152
0.10947
0.43483
0
0.10234
0.05666
0.07955
0.15944
0.14066
U.10756
DPOT-L
Auto-regress
01021 30
0,00400
0,01310
0.01195
0.01600
0.00900
0.01253
0005995
0.03000
500M
PITA
0.01059
0.00198
000100
0,01001
0.00995
0.0999X
0.02004
0.01871
0.01978
0,04965
0.020148
0.0NE302
Product long Erojectosy
DROT- Auto-regress
7M
PITA
0003600
0.017188
0000500
0.00327
0001000
0,01090
081100
101000
6,01000
0.30034
0.31012
DROT-S
Aute-regress
0.02370
0.00137
0200235
0.26800
30M
PITA
0.00854
0.00223
F,00000
0,00137
0.22620
MPP-Ti
regress
683291
7M
PITA
GA2174
MPP-S
Auto- regress
0201631
3UM
PITA
0,00950
DPOT-M
Auto-rogress
0.01.2
0.00335
021030
6000227
0.17200
122M
PITA
0.00669
0.00133
0.16390
FNO-M M
TCETeSS
0.01761
0.00125
101100
0.00020
0,00924
0.13156
170M
0,01710
610664
0.38731
DPOTL
Auto- regress
0.0104
0.00323
0.00729
6,00170
0.000000
500M
PITA
0,00708
0.00199
0.00670
0,00121
0.16488
This alternation between the suboptimal solutions will con- verge towards a high-quality optimization result that satisfies global convergence.
Training and Evaluation: All experiments are carried out on single A800 GPU with 80 GB of memory. We apply the commonly used scale-independent normalized root mean squared error (nRMSE) (Takamoto et al., 2023; Hao et al., 2023) to measure the quality of the prediction, which is defined as follows,
4. Experiments
Datasets: To enable comparisons with baselines employing the auto-regressive strategy, select 12 datasets from four different sources: 3 datasets from FNO (Li et al., 2020), 6 datasets from PDEBench (Takamoto et al., 2022), 2 datasets from PDEArena (Gupta& & Brandstetter, 2022), and 1 dataset from CFDBench (Luo et al., 2023). For generalization tasks, the Burgers' equation from (Boussif et al., 2022) is included as an additional dataset. Additional details regarding the datasets used can be found in the Appendix C.
Tiest
|m (0)l2
nRMSE-
(12)
Treat
||u. 11:112
where Trest indicates the length of testing data in the tem- poral domain.
4.1. State-of-the-Art Results
The in-distribution performance of PITA is evaluated and presented in Table 1. The testing datasets comprise PDE data that share the same boundary conditions and parame- ters as the training data but differ in their initial conditions. Overall, PITA demonstrates -fate-of-the-art (SOTA) perfor- mance in both short trajectory predictions (10 steps) and long trajectory predictions (longer than 10 steps). A detailed analysis of the results is provided below.
We compare PITA with ruto-regressive base- lines, primarily focusing on PDE foundation models such as DPOT (Hao et al., 2024) and MPP (McCabe et al., 2023). For DPOT, we select models of sizes Ti, S, M. and L, which have been pretrained on 12 datasets. In the case of MPP, we choose the Ti and S models, pretrained on 10 datasets. Both models are designed to work predominantly with 2D datasets. Additionally, we assess our method on single- family models (Shen et al., 2024), where pretraining and finetuning are conducted on a single dataset instead of mul- tiple datasets. FNO-M (Li et al., 2020) serves as the repre- sentative baseline. The sizes and parameters of the various models are detailed in Appendix D.I. These ensure com- prehensive and fair evaluation across diverse PDE settings.
Baselines:
When performing short tra- jectory prediction, PITA achieves SOTA performance across most datasets and various model sizes, particularly with medium and large models. For small-sized models (Ti, S), PITA achieves the best results in 75% of prediction tasks, with an average improvement of 11.21% % in testing outcomes. Notably, the maximum improvement observed was an im-
Short Trajectory Prediction:
6

---

Physics-informed Temporal Alignment for Auto-regressive PDE Foundation Models
pressive 63.81% with the DPOT-S model the FNO-NS- le-4 dataset. For medium and large-sized models (M, L.), PITA excels in 77.78% of the prediction tasks, achieving an average reduction in nRMSE of 13.53%, with the best improvement recorded at 53.36% using the DPOT-L model on the FNO-NS-Ie-3 dataset. These results indicate that although PITA is not specifically tailored for PDE data with relatively short trajectories, it still exhibits excellent perfor- mance compared to finetuning approaches that rely solely on the auto-regressive approach. This highlights PITA's ro- bustness and effectiveness in handling variety of prediction tasks, contributing to its overall superiority in performance.
0.015
DPOT with PITA
0.014
DPOT
0.013
nRMSE
0.012
0.011
0.010
0.009
0.008
0.007
7
30
100
500
Model Size (M)
Long Trajectory Prediction: By integrating temporal alignment with physics-informed constraints, PITA demon- strates superior performance over traditional auto-regressive training methods, effectively addressing the shorteut prob- lem commonly encountered in long trajectory prediction.
Figure 3: Results of scalability for different model sizes.
datasets, although errors display a gradual increase with temporal progression, PITA substantially reduces error ac- cumulation compared to conventional auto-regressive bench- marks. These findings provide strong empirical evidence that PITA'S multi-scale consistency mechanism effectively mitigates enror propagation across extended temporal se- quences in long-term forecasting tasks.
Overall, PITA reduces the total accumulated error signif- icantly compared to the auto-regressive approach, achiev- ing an improvement of 30.22% over the long trajectory, which translates to an average decrease of 2.00% error for each time step. The shortcut problem is particularly pronounced in the Ie-- dataset, where the average percentage of error accumulation per prediction step reaches 34.31%. In contrast, the shortcut problem is less significant in the PDEBench-SWE and PDEBench-DR datasets that are longest trajectory datsets (91 steps). This discrepancy can be attributed a main factor that the dynamics exhibits smaller variations in long term trajectory, making it easier for the model to capture the temporal dynamics (Takamoto et al., 2022). Despite this, PITA still demonstrates meaning- ful performance improvements over all the baselines. We further analyze the relevance between the performance gains of PITA and foundation model size. The experiments are conducted on PDEBench-DR with the longest trajectory, as shown in Figure 3. We can see that in long trajectory predic- tions, PITA demonstrates scalability. It even outperforms the baseline with 100M parameters when using only 30M parameters, showcasing parameter efficiency in modeling long-term dynamics. Additional scaling experiments on long trajectory datasets are presented in Appendix E.3.
4.3. Generalizing to Downstream tasks
To investigate the generalization performance of PITA on downstream tasks, we consider two experimental settings: (1) PDEs of the same type as those in the pretraining datasets, but with different coefficients, (2) PDEs of a dif- ferent type not present in the pretraining datasets.
In the first setting, we assess PITA'S generalization capac- ity using the compressible Navier-Stokes equations from PDEBench (Takamoto et al., 2022) with different shear viscosity and bulk viscosity Ç. The finetuning results, detailed in the fourth column of Table 2, demonstrate that PITA outperforms all other settings, achieving an average reduction in error by 5.93%. This exceptional performance highlights PITA's capability to proficiently adapt to the in- tricate challenges posed by compressible fluid dynamics.
In the second setting, we use the viscous Burgers' equation (Boussif et al., 2022) for evaluation. The pretrained model is finetuned on this dataset for 500 epochs employing both auto-regressive training and PITA. The corresponding re- sults, presented the third column of Table 2, demonstrate that PITA consistently outperforms across all model sizes, achieving an average error reduction of 37.8% compared to the original auto-regressive approach. The incorporation of PDE discovery enables PITA to effectively learn the under- lying physical laws of the dataset without prior knowledge, indicating enhanced adaptability to various types of PDEs. These findings underscore PITA's robust adaptability and practical value for broad range of physical modeling tasks.
4.2. Solution to Error Accumulations
Following the error visualization methodology introduced in (Christlieb et al., 2016), we compute and plot the rolling- step mean squared error (MSE) for each long-term fore- casting dataset, as illustrated in Figure 10. This rolling window highlights temporal trends in prediction quality and makes error accumulation more apparent. Our anal- ysis demonstrates that prediction enrors on datasets such as PDEBench-SWE. FNO-NS- and FNO-NS- 1c-4 ex- hibit no significant error accumulation over time. For other
7

---

Physics-informed Temporal Alignment for Auto-regressive PDE Foundation Models
FDEBesch-SWE Stepa)
PDE Steps)
NS-Foore Steps)
MSF MS Pf Time Step
1.M02
PITA
Auto regressire
MSE Pfr Time Step.
PITA
Asto-ingrosävee
MSE ofPer Time Step
0.11
Auto regresiive
PITA
(M02
0.0076
0.09
1,018
0.66
60066
0.0018
0,0035
0.012
0.02
Time Step (1)
Time Step (1)
Time Step (1)
PNO-NS-1e3 Dataset (20 Steps)
FNO-NS-Io-4 Datasel (20 Steps)
MSE Pfr Time Step
5.000
Aun-nggrossve
MSE of Per Time Step
0.12
PTTA
regressive
000
0.10
0.000
0.66
2.000
4.000
0.0
2.5
7.5
16.0
12.5 15.0
175
0.0
255
5.0
7.5
10.0
12.5
15.0
17.
Time Stop (1)
Time Step (1)
Figure 4: Rolling step MSE long trajectory datasets.
Table 2: Generalization on two downstream tasks of PITA compared with baselines. More comprehensive generaliza- tion experiments are presented in Appendix E.4.
components. Notably, the physics loss Lpky plays piv- otal role in guiding the model to adhere to the underlying physical laws, resulting in 44.97% accuracy promotion. The consistency loss LCon leads to a 44.14. 10% increase in accu- racy. It maintains temporal alignment and is essential for the stability of the learned solution, ensuring smooth transitions in long-term predictions.
Pretrained Model
Finetune Strategy
Burgers
PDE-Bench CNS-(M, n)
1,1 x 10
Auto-regress
0.00904
0.18681
DPOT-Ti
PITA 0.00894
0.16498
Auto-regress
0.00533
0.13372
DPOT-S
PITA 0.00458
0,12941
Task 2: Effects of Spatial Downsampling. To boost com- putational efficiency, PITA adopts downsampling strategy to randomly select grids in the spatial domain. The sampled grid is represented as X = fæ.,, T2z. T }. where w = [n/[] denotes the number of sampled points, and I represents the downsampling factor. The value of is se- lected from (1, 4). Experimental results indicate that larger downsampling factors (i.e., fewer sampled grids) of- ten lead to better performance, which can be attributed to the balance between data sparsity and generalization (Liu et al., 2024a). Fewer sampled grids encourage the model to focus on key spatial features, effectively mitigating overfitting to fine- ime-ggrinedd noise present in denser grids.
Auto-regress
0.00378
0.14607
DPOT-M
PITA 0.00362
0.14453
4.4. Ablation Studies
We performed five ablation studies to assess the impact of different design choices PITA, training small-scale models on the FNO-NS-J 1E-3 dataset for Tasks 1-4 and on the Burgers equation for Task 5.
Task 1: Effectiveness of Loss Components. To highlight the indispensability of each supervision, we systematically evaluate the pretrained model under various configurations. Specifically, Data refers to using only data loss, following the auto-regressive strategy, Data + Pky denotes the com- bination of data loss and physics loss, while D + LCon represents the combination of data loss and consistency loss. Cotta represents the inclusion of all three compo- nents, aligning with the PITA training strategy. Except for the Data setting, all configurations adopt the multi-task learning strategy described Sec. 3.4. As shown in Ta- ble 3, PITA achieves the best results across all settings, underscoring the effectiveness of integrating multiple loss
Task 3: Effects of Temporal Downsampling. To enhance computational efficiency, PITA employs downsampling strategy that selects frames within the temporal domain. The sampled solution is expressed as {4) where T represents the last frame and Te denotes the number of frames retained during downsampling. The values of To are selected from {10, 7, 5, 3). Notably, the best results for both prediction tasks oceur at Tc 7. Conversely, increasing or decreasing To from this optimal point results in elevated nRMSE values. While excessive downsampling can lead to
8

---

Physics-informed Temporal Alignment for Auto-regressive PDE Foundation Models
significant loss of critical information, retaining too many frames may detract from the model's s efficiency. Given that the results for Tc = 3 and Tc 7 are relatively close, we opted to use Tc = 3 for the majority of our results presented in Table to optimize computational costs.
Table 3: Results of the ablation studies. For each task, only the specified settings are different with the other hyperpa- rameters and training configurations remaining the same.
Task No.
Settings
nRMSE
nRMSE
(10 Step)
(Full Length)
Task 4: Effects of Multi-Task Learning Strategy. The in- fluence of weighting different components of the total loss is evaluated. The data loss weight is set 1, the physics loss and consistency loss are weighted by on and 022, respectively. We compare manually chosen weights (a1, .2) with the automatically adjusted strategy. The tested manual settings for (01, (22) are (0.5, 0.5), (0.3, 0.7), and (0.7, 0.3). As pre- sented in Table 3, PITA achieves an average improvement of 9.06% over all manual weight settings, suggesting the ef- fectiveness of its adaptive strategy in dynamically balancing competing objectives. This ensures each loss component contributes optimally to the learning process. Furthermore, PITA exhibits greater robustness to weight variations com- pared to existing methods (Raissi et al., 2019), highlighting its superior ability to generalize across diverse PDE data.
Data Lrhy + Data
0.00301
0.00437
0.00166
0.00240
TASK| 1
Lcon + Data
0,00165
0.00249
L
0.00157
0.00223
1= 1 1 1=
0.00169
0.00238
TASK 2
0.00168
0.00237
i 4
0.00157
0.00223
Tc 10
0.00161
0.00229
Tc = 7
0.00155
0.00221
TASK 3
Tc = 5
0.00170
0.00240
Tc =
0.00157
0.00223
(0.5, 0.5)
0.00172
0.00243
(0.3, 0.7)
0.00175
0.00246
TASK 4
(0.7, 0.3)
0,00171
0.00242
Automatic
0.00157
0.00223
None
0.00325
0.01091
Task 5: Effects of Key Library Terms. We assessed how key library terms affects PITA's performance in inferring physical laws by conducting experiments with distinet con- figurations on the Burgers equation dataset, which was ex- cluded from pre- training. The Burgers equation is composed of a time-dependent term 0,u,& a diffusion term VV u and a convective term udu, which are indispensable for capturing the underlying dynamics of the system (Gao & Zou, 2017). Table 3 compares three configurations: "None" removes all first- and second-order derivatives (no convective or diffu- sion terms), "One-Order" removes only second-order deriva- tives (diffusion eliminated, convection retained), and "Full" uses the complete library. It is evident that the full library achieves state-of-the-arr predictive accuracy. In the "One- Order" configuration where only the convective term is retained, the performance experiences a slight degradation. Further removal of both the convective and diffusion terms leads to a modest compromise in the model's ability to cap- ture the underlying dynamics. Nevertheless, the model still achieves a 4 41.42% improvement over the baseline, demon- strating the robustness of the approach even with limited incorporation of physical knowledge. These results concur with Task 1 Is ablation of physics loss, where PITA continues to improve over the baseline via data and consistency losses even when physics loss is ineffective.
One-Order
0.00287
0.00862
TASK 5
Complete 0.00152
0.00776
Auto-regress 0.00431
0.01857
size, indicating that PITA remains suitable for larger models. For example, PITA may take approximately 42.32% and 22.8% longer to process for DPOT-Ti and DPOT-L, respec- tively. Thus, an appropriate downsampling strategy in PDE discovery is advisable to balance effectiveness and computa- tional consumption. Second, the performance of PITA may also be limited by the method used for PDE discovery. The completeness of the library of terms and the level of noise in the observational data may affect the model's performance.
5. Conclusion
In this paper, we present PITA, a novel approach for PDE foundation models. PITA is designed to address the lim- itations of traditional auto-regressive methods by mitigat- ing shortcut issues. Instead of relying prior knowledge embedded, PITA discovers underlying physical laws in a data-driven manner and incorporates these laws into opti- mization objectives, compelling the model to accurately learn the governing physics. Extensive experiments are con- ducted to rigorously validate PITA on foundation models pretrained across diverse datasets scaling up to 500 million parameters, well as individual neural operators tailored to specific PDEs. PITA consistently achieves trate-of-the- art performance across multiple datasets, demonstrating its effectiveness across models of different sizes. It also ex- cels in various downstream tasks, highlighting its strong generalization capabilities.
4.5. Limitations
Despite the advantages mentioned above, the proposed PITA still has some limitations. First, PITA may require more computational cost, which results from an intricate gradient propagation process due to the additional two losses. Impor- tantly, this extra computation time does not scale with model

---

Physics-informed Temporal Alignment for Auto-regressive PDE Foundation Models
Impact Statement
in Neural Information Processing Systems, 35:31972- 31985, 2022.
From a scientific perspective, PITA significantly advances doysics-informet machine learning by providing a versatile framework that effectively bridges data-driven methodolo- gies with fundamental physical principles. Its capability to discover governing equations directly from data serves as a powerful tool for analyzing complex phenomena, partic- ularly in scenarios where explicit models are incomplete or unavailable. However, it is essential to acknowledge that predictions made by neural networks for physical sys- tems inherently involve approximation errors and often lack interpretability.
Brown, A., Tomasev, N., Freyberg, J., Liu, Y., Karthike- salingam, A., and Schrouff, J. Detecting shortcut learning for fair medical ai using shortcut testing. Nature Commu- nications, 14(1):4314, 2023.
Brunton, S. L. and Kutz, J. N. Promising directions of machine learning for partial differential equations. Nature Computational Science, 4(7):483-494, 2024.
Brunton, S. L.,, Proctor, J. L., and Kutz, J. N. Discovering governing equations from data by sparse identification of nonlinear dynamical systems. Proceedings of the Na- tional Academy of Sciences, 1 113(15):3932-33
PITA shows significant potential for numerous future ap- plications in numerical computing. First, as a development to the traditional auto-regressive approach, PITA can be seamlessly applied to various PDE foundation models, sup- porting both pretraining and finetuning processes. Second, for domain-specific applications like fluid dynamics, library terms be tailored according to governing equations such as Navier-Stokes, Thus enables more accurate solutions and better aligmment with physical laws. Third, by incorporat- ing PDE discovery, PITA opens up opportunities to utilize other inverse problem-solving techniques. such as initial state recovery and parameter inference, improve forward problem-solving.
2016.
Budišić, M., Mohr, R., and Mezić, I. Applied koopman- ism. Chaos: An Interdisciplinary Journal of Nonlinear Science, 22(4), 2012.
Cao, D., Ye, W., Zhang. Y., and Liu, Y. Timedit: General- purpose diffusion transformers for time series foundation model. arXiv preprint arXiv:2409.02322 2024.
Chen, T., Zhou, H., Li, Y., Wang, H., Gao, C., Shi, R., Zhang, S., and Li, J. Omniarch: Building the foun- dation model for scientific computing. arXiv preprint arXiv:2402. /6014, 2024.
Acknowledgements
Chen, Z., Badrinarayanan, V., Lee, C.-Y., and Rabinovich, A. Gradnorm: Gradient normalization for adaptive loss balancing in deep multitask networks. In International Conference on Machine Learning, pp. 794-803 PMLR, 2018.
This work was supported in part by the National Natu- ral Science Foundation of China (NSFC) Major Research Plan on Interpretable and General Purpose Next-generation Artificial Intelligence under Grant No. 92370205, NSFC Grant 12425113, the Natural Science Foundation of Jiangsu Province under Grant BK20240462, the China Postdoc- toral Science Foundation under Grants 2023M733427 and 2023TQ0349, and the Jiangsu Funding Program for Ex- cellent Postdoctoral Talent. We also acknowledge support from the Key Laboratory of the Ministry of Education for Mathematical Foundations and Applications of Digital Tech- nology, USTC, and the valuable discussions and assistance provided by Zijun Ding.
Chen, Z., Liu, Y., and Sun, H. Physics-informed learning of governing equations from scarce data. Nature Communi- cations, 12(1))6136, 2021.
Christlieb, A., Guo, W., and Jiang, Y. A weno-based method of lines transpose approach for vlasov simulations. Jour- nal of Computational Physics, 327:337-367, 2016.
Chuah, W., Tennakoon, R., Hoseinnezhad, R., Bab- Hadiashar, A., and Suter, D. Itsa: An information- theoretic approach automatic shortcut avoidance and domain generalization in stereo matching networks. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 13022-13032, 2022.
References
Berg, J. and Nyström, K. Data-driven discovery of PDEs in complex datasets. Journal of Computational Physics, 384:239-252, 2019.
Cicirello, A. Physics-enhanced machine learning: posi- tion paper for dynamical systems investigations. arXiv preprint arXiv:2405.055: 45590097..7... 2024.
Both, G.-J., Choudhury, S., Sens, P., and Kusters, R. Deep- mod: Deep learning for model discovery in noisy data. Journal of Computational Physics, 428: 109985, 2021.
Cuomo, S., Di Cola, V. S., Giampaolo, F., Rozza, G., Raissi, M., and Piccialli, F. Scientific machine learning through physics-infore -inforeed neural networks: Where we are and
Boussif, O.,, Bengio, Y., Benabbou, L., and Assouline, D.
Magnet: Mesh agnostic neural PDE solver. Advances
10

---

Physics-informed Temporal Alignment for Auto-regressive PDE Foundation Models
what's next. Journal of Scientific Computing, 92(3):88, 2022.
Kendall, A., Gal, Y., and Cipolla, R. Multi-task learning using uncertainty weigh losses for scene geometry and semantics. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pp. 7482- 7491, 2018
Dagaev, N., Roads, B. D., Luo, X., Barry, D. N., Patil, R., and Love, B. C. too-good-to-be-tee prior to reduce shortcut reliance. Pattern Recognition Letters, 166:164- 171, 2023.
Knowles, I. and Renka, R. J. Methods for numerical dif- ferentiation of noisy data. Electron. J. Differ. Equ, 21: 235-246 2014.
Donoho, D. L. Compressed sensing. IEEE Transactions Information Theory, 52(4): 2006.
Li, A. C.,, Kumar, A., and Pathak, D. Generative classifiers avoid shortcut solutions. In ICML 2024 Workshop on Structured Probabilistic Inference Generative Model- ing. 2024.
Du, M., Manjunatha, V., Jain, R., Deshpande, R., Dernon- court, F,, Gu, J., Sun, T., and Hu, X. Towards interpreting and mitigating shorteut learning behavior of nlu models. arXiv preprint arXiv:2103.06922. 2021.
Li, Z., Kovachki, N., Azizzadenesheli, K., Liu, B., Bhat- tacharya, K., Stuart, A., and Anandkumar, A. Fourier neural operator for parametric partial differential equa- tions. arXiv preprint arXiv:2010.08895, 2020.
Gao, Q. and Zou, M. An analytical solution for two and three dimensional nonlinear burgers' equation. Applied Mathematical Modelling, 45:255-2700 2017.
Liebel, L. and Körner, M. Auxiliary tasks in multi-task learning. arXiv preprint arXiv: 1805.0633,, 2018.
Geirhos, R., Jacobsen, J.-H., Michaelis, C.,, Zemel, R., Bren- del, W., Bethge, M., and Wichmann, F. A. Shortcut learn- ing in deep neural networks. Nature Machine Intelligence, 2(11):665-673, 2020.
Liu, J., Ram, P., Yao, Y., Liu, G., Liu, Y., SHARMA, P., Liu,
S., al. Model sparsity can simplify machine unlearning. Advances Neural Information Processing Systems, 36, 2024a.
Geneva, N. and Zabaras, N. Modeling the dynamics of PDE systems with thysics-constrained deep auto- regressive networks. Journal of Computational Physics, 403:109056, 2020.
Liu, Q. and Ma, Foundation models for geophysics: Re- views and perspectives. arXiv preprint arXiv:2006.03163, 2024.
Golub, and Pereyra, V. The differentiation pseudo- inverses and nonlinear least squares problems whose vari- ables separate. SIAM Journal on Numerical Analysis, 10 (2):413 -0:2413-33,,. 1973.
Liu, Y., Sun, J., He, X., Pinney, G., Zhang, Z., and Schaeffer, H. Prose-fd: A multimodal PDE foundation model for learning multiple operators for forecasting fluid dynamics. arXiv preprint urXiv:2409.098/1 2024b.
Gupta, J. K. and Brandstetter, spatiotemporal-scale generalized PDE modeling, arXiv preprint arXiv:2209. 1tt.:222299111166. 156/6, 2022.
Towards multi-
Long, Z., Lu, Y., Ma, X., and Dong, B. PDE-net: Learning PDEs from data. In International Conference on Machine Learning, pp. 3208-3222 -216. PMLR, 2018.
Hao, Z., Yao, J., Su, C., Su, H., Wang, 7.,, Lu, F., Xia, 7 Zhang, Y., Liu, S., Lu, L., et al. Pinnacle: A comprehen- sive benchmark of physics-informed neural networks for solving pdes. arXiv preprint arXiv:2306.088 08827, 2023.
Luo, X., Wei, L., Wen, L., Yang, J., Xie, L., Xu, Z., and
Tian, Q. Rectifying the shortcut learning of background for few-shot learning. Advances in Neural Information Processing Systems, 34: 14:13373-13085 2021.
Luo, Y., Chen, Y., and Zhang. Z. Cfdbench: A comprehen- sive benchmark for machine learning methods in fluid dynamics. arXiv preprini urXiv:2310.05963. 2023.
Hao, Z., Su, C., Liu, S., Berner, J., Ying, C., Su, H., Anand- kumar, A., Song, J., and Zhu, J. Dpot: regresseeee denoising operator transformer for large-scale pde pre- training, arXiv preprint ar arXiv:2403.03542 2024.
Ma, H., Lu, X., and Zhang, L. Extracting parametric dy- namics from time-series data. Nonlinear Dynamics, 111 (16):15177-15199, 2023.
Hermann, K. L., Mobahi, H., Fel, T., and Mozer, M. C. On the foundations of shortcut learning. arXiv preprint arXiv:2310. 16..228 2023.
McCabe, M., Blancard, B. R.-S., Parker, L. H., Ohana, R., Cranmer, M., Bietti, A., Eickenberg, M., Golkar, S., Krawezik, G., Lanusse, F., et al. Multiple physics pre- training for physical surrogate models. arXiv preprint arXiv:2310.0290 2023.
Huang, K. and Aviyente S. Sparse representation for signal classification. Advances in Neural Information Process- ing Systems, 19, 2006.

---

Physics-informed Temporal Alignment for Auto-regressive PDE Foundation Models
Noack, B. R., Afanasiev, K., Morzyúski, M., Tadmor, G., and Thiele, F. A hierarchy of now-dimensional models for the transient and post-transient cylinder wake. Journal of Fluid Mechanics, 497:335-333 2003.
Takamoto, M., Praditia, T.,, Leiteritz, R., MacKinlay, D., Alesiani, F,, Pflüger, D., and Niepert, M. PDEbench: An extensive benchmark for scientific machine learning. Advances in Neural Information Processing Systems, 35: 1596-1611, 2022.
Raissi, M., Perdikaris, P., and Karniadakis, G. E. Physies- informed neural networks: A deep learning framework for solving forward and inverse problems involving nonlinear partial differential equations. Journal of Computational Physics, 378:686-707, 2019.
Takamoto, M., Alesiani, F., and Niepert, M. Learning neural PDE solvers with parameter-guided channel attention. International Conference on Machine Learning, pp.
33448-33444 PMLR, 2023.
Robinson, J., Sun, L., Yu, K., Batmanghelich,, K., Jegelka, S., and Sra, S. Can contrastive learning avoid shorteut solutions? Advances in Neural Information Processing Systems, 34:4974-4986, 2021.
Tibshirani, R. The lasso method for variable selection the cox model. Statistics in Medicine, 16(4):385-395, 1997.
Wan, H., Wang, Q. Mi, Y., and Sun, H. Pimrl: Physics- informed multi-scale recurrent learning for spatiotempo- ral prediction. arXiv preprint arXiv:2503.10253, 2025.
Rudy, S. H., Brunton, S. L., Proctor, J. L., and Kutz, N. Data-driven discovery of partial differential equations. Science Advances, 3(4):e ((4):e6602611 602614, 2017.
Wang. Z. and Wu, Z. Foundation model for chemical process modeling: Meta-learning with physics-informed adaptation. arXiv preprint arXiv:2405.. 1752, 2024.
Schaeffer, H. Learning partial differential equations via data discovery and sparse optimization. Proceedings of the Royal Society A: Mathematical, Physical and Engineer- ing Sciences, (7321197):20100444 2017.
Wen, J., Fang. X., Cui. J., Fei, L., Yan, K., Chen, Y.. and Xu, Y. Robust sparse linear discriminant analysis. IEEE Transactions on Circuits and Systems for Video Technol- ogy, 29(2):390-4033 2018
Scimeca, L., Rubinstein, A., Teney, D., Oh, S. J., Nicoli- cioiu, A. M., and Bengio, Y. Mitigating shorteut learn- ing with diffusion counterfactuals and diverse ensembles. arXiv preprint arXiv:2311. 16176, 2023.
Yang. L.., Liu, S., Meng, T., and Osher, S. J. In-context operator learning with data prompts for differential equa- tion problems. Proceedings of the National Academy of 2023.
Sener, O. and Koltun, V. Multi-task learning as multi- objective optimization. Advances in Neural Information Processing Systems, 31, 2018
Sciences, 220(3)):2310144111
Ye, Z., Huang, X., Chen, L., Liu, H., Wang, Z., and Dong, B. PDEformer: Towards a foundation model for one- dimensional partial differential equations. arXiv preprint arXiv:2402 /2652, 2024a.
Shen, J., Marwah, T., and Talwalkar, A. Ups: Efficiently building foundation models for PDE solving via cross- modal adaptation. In ICML 2024 Al for Science Work- shop, 2024.
Ye, Z., Huang, X., Chen, L., Liu, Z., Wu, B., Liu, H., Wang, Z., and Dong, PDEformer-l: foundation model for one-dimensional partial differential equations. arXiv preprint arXiv:2407.06666 2024b.
Song. Z., Yuan, J., and Yang, H. Fmint: Bridging human designed and data pretrained models for differential equa- tion foundation model. arXiv preprint arXiv:2404.1460 14688, 2024.
Yizheng, W., Xiaoying, Z.. Rabezuk, T., and Yinghua, Ai for PDEs in solid mechanics: A review. Advances in Mechanics, 54(3):1-57, 2024.
Stephany, R. and Earls, C. PDE-read: Human-readable partial differential equation discovery using deep learning. Neural Networks, 154:360-388 3822, 2022.
Subramanian, S., Harrington, P., Keutzer, K., Bhimji, W., Morozov, D., Mahoney, M. W., and Gholami, A. To- wards foundation models for scientific machine learning: Characterizing scaling and transfer behavior. Advances in Neural Information Processing Systems, 36, 2024.
Zhang, T. Adaptive forward-backwart greedy algorithm for sparse learning with linear models. Advances in Neural Information Processing Systems, 21, 2008
Zhao, L., Liu, Q., Yue, L., Chen, W., Chen, L., Sun, R., and Song, C. Comi: Correct and mitigate shortcut learning behavior in deep neural networks, In Proceedings of the 47th International ACM SIGIR Conference on Research and Development in Information Retrieval, pp. 218 2282 2024.
Sun, J., Liu, Y., Zhang, Z., and Schaeffer, H. Towards a foundation model for partial differential equation: Multi- operator learning and extrapolation. arXiv preprint arXiv:2404. 12355, 2024.
12

---

Physics-informed Temporal Alignment for Auto-regressive PDE Foundation Models
Zhao, Z., Ding, X., and Prakash, B. A. Pinnsformer: A transformer-based framework for physics-informed neu- ral networks, arXiv preprint arXiv:2307.11833 2023.
Zhou, H., Ma, Y., Wu, H., Wang. H., and Long, M. Uni- solver: PDE-conditional transformers are universal PDE solvers, arXiv preprint arXiv:2405.17522 2024.
13

---

Physics-informed Temporal Alignment for Auto-regressive PDE Foundation Models
A. Table of Notations
The primary notations used this paper are listed in Table 4.
Table 4: Mathematical Notations.
Notations for Sec. 3.1
Symbol
Definition
æ
Spatial variable in the domain f C R"
Ω
Spatial domain of PDE.
0n
Boundary of spatial domain Q
u(x,
Time-dependent function defined over the spatial domain Q for the time interval [0, T). Discretized grid of n points in 22, represented as {æ 22...., Rn), where each æ; 124. Total number of spatiotemporal PDEs considered. Dimensionality of the solution space of u.
X,,
N
du
dx
Dimensionality of the spatial domain for the k-th PDE. Initial condition for u at t O.
uo(a)
Flu] (x, t)
Differential operator acting on u and its spatial derivatives.
B (ar, t)
Boundary condition operator applied boundary points a: E at2.
u(x,
Solution of PDE time f, defined in domain 52 with values in Rdd.
Notations for Sec. 3.2
Symbol
Definition
0
Parameters of neural operator. In subsequent sections refers the parameters of foundation models.
Go
Neural operator parameterized by weights 0.
Tin
Roll-in window length, representing the number of input frames used ground truth. Roll-out window length, representing the number of frames predicted iteratively by Go.
Tar
ucT
Set of frames (u,);= used input. Predicted frame time step T.
ür
ur
Ground truth frame at time step T.
Notations for Sec. 3.3
Symbol
Definition
Lp norm.
X
Sampled grid for spatial domain.
w
Number of downsampled grid points.
C
Number of physical variables.
U
Matrix representing the compressed data of C' channels across m time points and n spatial locations. Library matrix containing candidate terms and partial derivatives for the PDE.
S
Total number of candidate terms in the library. Matrix of the discovered PDE coefficients.
A
Φi
Pscudo-inverse of the library matrix Φ
Symbols used in the formulas but not listed in the tables are clarified in the main text.
B. Illustration of the Shortcut Problem
This section investigates the origin of the shortcut problem and its consequential error accumulation auto-regressive prediction frameworks. When predicting the T-th frame, the prediction error at the (7-1))th time step propagates to subsequent predictions. This propagation triggers cumulative error amplification over successive time steps, which is critical limitation of auto-regressive mechanisms. As illustrated in Figure 5, within prediction window of length Tar. errors originating from the initial frame amplify iteratively, thereby degrading the accuracy of later predictions. Such accumulation manifests as severe deviations or unphysical artefacts in dynamical system simulations, particularly when predicted states violate the fundamental physical laws represented by the governing PDEs. For instance, in fluid dynamics or climate modeling applications, these artefacts may manifest as non- non-physical vorticity patterns or temperature distributions that diverge from observed system behavior. The compounding errors not only distort short-term predictions but also erode the validity of extrapolated trajectories over extended timescales. These artefacts ultimately compromise the reliability of
14

---

Physics-informed Temporal Alignment for Auto-regressive PDE Foundation Models
long-term scientific forecasts, raising significant challenges for applications requiring high-fidelity numerical simulations Furthermore, the error accumulation phenomenon arises from shortcut learning in auto-regressive approaches, where models
Input time window
Rollout time window
u (u², u
1
u u
u
-1'11
u T+1
u +Tar
u
u
T-1
u
2
u w (u ..,,u-11 concat(u", u)
at+1
TT
' T
Eroor
Ervr
Ertr
T+1
7 Tar
T+1 = (.. u²)
Accumulative Error
Figure 5: Illustration of Error Propagation in Auto-regressive Prediction.
prioritize copying values from the (T-1)-th frame to predict the T-th frame, rather than learning the underlying physical dynamics governed by PDEs. To quantify this behavior, we analyze two long-trajectory datasets using auto-regressive prediction, where the nRMSE notably higher compared to other methods, as shown in Figure In this figure, the pink 1) measures the structural similarity index (SSIM) between
curve labeled Temporal SSIM (predicted vs predicted consecutive predicted frames, út-1 and ú, The blue curve labeled Frame SSIM (predicted vs ground truth quantifies the SSIM between the predicted frame ú; and the ground truth frame
At initial time steps, the blue curve dominates the pink curve (Figure 6), indicating that û, aligns more closely with the ground truth u, than with the prior prediction ú; This suggests the model initially captures physical dynamics encoded the PDEs rather than relying on superficial shorteut features. However, at later stages, the pink curve surpasses the blue curve, revealing that i, increasingly resembles ú instead of u,, thereby exacerbating shortcut-driven error propagation. This divergence reflects the model's failure learn true physical features instead, it recursively replicates prior predictions, introducing incremental errors at each step. Over extended trajectories, the repeated copying of erroneous frame values leads to severe етror accumulation, particularly in systems requiring long-term stability, such as turbulence modeling or climate projections. Addressing this limitation is critical for improving the robustness of auto-regressive methods in scientific computing.
C. Details of Datasets
C.I. FNO-NS->
This benchmark dataset considers the 2D Navier-Stokes equation for viscous, incompressible fluid in vorticity form the unit torus. The equation is written
O. Vw vAw f.
(13)
V u 0,
where w is vorticity field, u is the velocity field, f is the external force. The only varying component in this dataset is the viscosity coefficient v, which takes values from the set {1 x 10-5,10 1 x 10-4, 1 x 10-3), corresponding to the datasets FNO-NS-Ie-5, FNO-NS-Ie-4, and FNO-NS- 1e-3. For PNO-NS-Ie-S the total length of the testing data is 20 steps, while for both FNO-NS- 1e-4 and ENO-NS-Le-3, the length is 30 steps. The task involves predicting future vorticity steps w t) given the initial 10 steps, where (ar, 1) [0, 1]2 x [0, T).
15

---

Physics-informed Temporal Alignment for Auto-regressive PDE Foundation Models
1.000
0.95
Temponal SSIM (pred vs (t-))
0.999
Prame SSIM (pred GT
0.90
Average SSIM
0.998
Average SSIM
0.85
0.997
0.75
0.995
0.70
0.65
0.994
Temporal (pred pred (I-1)
0.60
0.993
Frams SSLM (peed GT
0.55
20
40
60
80
10
20
30
40
Time Step (t)
Time Step (t)
Figure 6: Intrinsic connection between error accumulation and shorteut learning. On the left, present the shorteut issue the PDEBench-DD -DR dataset with trajectory length being 91. On the right, we present the shorteut issue on the NS2D-Foree dataset with trajectory length being 46.
C.2. PDF:Bench-CNS-MM..
Different from FNO-v, this benchmark dataset considers the 2D Navier-Stokes equation for viscous, compressible fluid. The equation is written as
0.p+ + (pv) = 0,
P p0.v + v. Vv) = =-Vp n+v (5+ vv
(14)
pu
( c+ c
po?
0g
+ +
2 v o =,
2
where P represents the mass density, v is the velocity, p denotes the gas pressure, and p/(T 1) is the internal energy, where r = 5/3. The viscous stress tensor is denoted by o', and n,. S are the shear and bulk viscosity, respectively. The Mach number M is defined by M = /ul/vTp/p TTp/p.
This benchmark examines the varying components of Mach number M and shear viscosity 7) with values consisting of (1, 0.1), (1, 0.01), (0.1, 0.1), and (0.1, 0.01). The total length the dataset 21 steps. The prediction length for the short trajectory is set to 11, whereas for the other datasets, the prediction length for the short trajectory 10. Given the initial 10 steps, the objective predict the velocity v( t). pressure p(æ, t), and density p(æ, t) fields, where (x,() t) [0, 1 1].
C.3. PDEBench-SWE
The saallow-water equations, derived from the general Navier-Stokes equations, present suitable framework for modelling free-surface flow problems. The equation is written
0,h + 0,hu + 0,hhv = 0.
0,hu + 0g
2
2 ghh
+0,whh hg,,
(15)
0, + 0, h+ ggh
+ +.uuhh -ghhh,,
where u and u represent the velocities in the horizontal and vertical directions, describes the water depth, and indicates a spatially varying bathymetry. The total length of the testing dataset comprises 101 steps. Given the initial 10 steps, the objective is to predict the water depth ((x, t) within the domain [-2.5, 2.5 2 x [0, 1].
16

---

Physics-informed Temporal Alignment for Auto-regressive PDE Foundation Models
C.4. PDEBench-DR
This benchmark dataset considers diffusion-reaction type PDE that combines diffusion process and rapid evolution from source term. The equation is expressed as
0,u Du + Duy Rus 0,0 = D yg Ra
(16)
where Du and D, are the diffusion coefficient for the activator and inhibitor, respectively, R., Ru (u, v) and R. R. (u, v) are the activator and inhibitor reaction function, respectively. The total length of the testing dataset comprises 101 steps. Given the initial 10 steps, the task is to predict the density fields u, where the domain is [-1,1]3 [0.5].
C.5. PDEArena
This benchmark dataset focuses on the incompressible Navier-Stokes equations in velocity function and vorticity stream formulation,
dv v Vu Vo Vp+ f,
aL
(17)
V v
where a represents the velocity flow fields. f denotes the external force, and p is the internal pressure. This benchmark consists of two subsets: the PDEArena-NS, which includes a fixed external force, and the PDEArena-NS-Fortt which incorporates a variable external force. For the PDEArena-NS dataset, the length the testing data 14 steps, while for the PDEArena-NS Force dataset, the length is 56 steps. The objective for both datasets predict the velocity v(x t), pressure p(x, 1), and density p(ær 1) fields using the initial 10 steps, where (ac, L) 32] [0, 24].
C.6. CFDBench-NS
This benchmark dataset focuses on the incompressible fluid dynamics on domains with irregular geometries
0
0t ppr) + V (pux) = -Vp - V p Vu (V
(18)
V (pu) 0,
where p is the density and μ is the dynamic viscosity, u is the velocity field, and is the pressure. The total length of the testing dataset consists of 20 steps. The task is to predict the velocity field u(æ, t) given the initial 10 time steps.
C.7. Magnet-
Viscous Viscous Burgers
To evaluate whether PITA can effectively handle downstream tasks, we consider the Burgers' equation, which models the dynamics of a field u incorporating nonlinear advection and diffusion. The equation is written as:
0.u +uu = 3Au.
5
(19)
(00, x,
Σ
A₂ sin (2mlf /64 + x cos (2m1" + +z))
where u represents the field, is the diffusion coefficient as Aj [-0.5, 0.5), 17, {1,2,3) 3}. and $7.0% Φb [0, 2 ). The total length of the testing dataset consists of 50 steps. The task predict the field (ax t) over the domain (æx, E [0.25, 63.7 [0, 9.8), given the initial 10 steps.
(0, 0.2). and the coefficients of initial conditions are sampled
D. Training Details
D.1. Model Configuration
Below, we provide the details of the various model configurations and scales used in the paper.
DPOT: In Table 5, we provide the training configurations of DPOT (Hao et al., 2024).
17

---

Physics-informed Temporal Alignment for Auto-regressive PDE Foundation Models
Table 5: Configurations of DPOT of different sizes.
Size
Attention Dim
MLP Dim
Layers
Heads
Params
Tiny
512
512
7M
Small
1024
1024
6
8
30M
Medium
1024
4096
12
122M
Large
1536
6144
24
16
509M
FNO: In Table 6, we provide the training configurations of FNO (Li et al., 2020), where Model and Mode2 represent the number of Fourier modes used in i and y spatial dimension respectively.
Table 6: Configurations of FNO.
Size
Model
Mode2
Width
Depth
Params
Medium
8
512
170M
MPP: In Table 7, we provide the training configurations of MPP (McCabe et al., 2023).
Table 7: Configurations of MPP.
Size
Embed Dim
MLP Dim
Heads
Blocks
Patch Size
Params
Tiny
192
768
3
12
[16,16]
7.6M
Small
384
1536
6
12
[16,16]
29M
D.2. Hyperparameters
The following training hyperparameters are used across all experiments, except where explicitly stated otherwise.
Table 8: Training Hyperparameter Settings Across Models and Strategies.
Model
DPOT
MPP
FNO
Strategy
Auto-regressive
PITA
Auto-regressive
PITA
Auto-regressive
PITA
Batch Size
20
20
24
24
20
20
Gradient Clipping
10000
10000
1
1
Dropout
0
0
0.1
0.1
0
Initial Learning Rate
le-3
1e-3
le-3
le-3
1e-3
le-3
Optimizer
Adam
Adam
AdamW
AdamW
Adam
Adam
Learning Rate Schedule
Cycle
Cycle
Cycle
Cycle
Cycle
Cycle
Weight Decay
le-6
1e-6
5e-2
5e-2
le-6
1e-6
Warmup Epoch
50
50
5
50
50
optimizer momentum
(0.9,099
(0.9,0.9)
(0.9,0.999)
(0.9, 0.999)
(0.9,0.9)
(0.9,0.9)
E. Supplementary Experimental Results
E.1. Data Efficiency Analysis of Downstream Tasks
To establish that PITA can be effectively applied to downstream tasks with limited and costly data acquisition, we conducted a series of experiments with varying sample sizes, analyzing data efficiency as shown in Table 9, Our experiments utilize the DPOT-S model. In general, both PITA and the auto-regressive model exhibit improved performance with an increased number of samples, whether training from scratch or finetuning from pretrained models. Notably, when finetuning from pretrained models, PITA achieves comparable prediction accuracy the auto-regressive model, utilizing only 500
18

---

Physics-informed Temporal Alignment for Auto-regressive PDE Foundation Models
samples-half the number required by the latter, which requires 1000 samples reach a similar level of accuracy. Moreover, when training from scratch, PITA demonstrates state-of-the-art performance, achieving an average reduction in nRMSE by 1 11.48. This significant improvement highlights PITA's efficiency and suitability for scenarios where data is scarce and expensive to obtain.
Table 9: Comparisons of Data Efficiency on Downstream Tasks.
Finetuning from Pretrained Model
Training from Scratch
Samples Size
Auto-regress
PITA
Auto-regress
PITA
10
0.1092
0.10273
0.14395
0.13873
100
0,01077
0.00998
0.01538
0.01405
500
0.00549
0,00514
0.00640
0.00586
800
0.00543
0.00451
0.00636
0.00531
1000
0.00518
0.00504
0.00629
0.00502
E.2. Performance Comparison of PITA and Auto-regressive Methods on Long Trajectory Datasets
To demonstrate that PITA effectively helps the shortcut problem commonly encountered in long trajectory datasets and achieves superior performance, we present the prediction results for both auto-regressive and PITA methods in Figure 7. The nRMSE for each method across the datasets is calculated as the average of the prediction results from all models, as detailed Table The results indicate that PITA consistently outperforms other methods across all long trajectory datasets, achieving an average nRMSE reduction of 29.15%. Except for the PDEBench-DR and PDEArena-NS-Fooe datasets, PITA also significantly outperforms the ruto-regressive training approach, achieving an average nRMSE reduction of 15.93%. While the performance of PITA the two aforementioned datasets is relatively close that of auto-regressive, the results remain competitive.
Auto-regressive prediction 10 steps PTTA prediction of 10 steps
Auto Auto-regred regressive prediction full length
PITA predistion of full longth
nRMSE
10
PDEBench-SWE
PDEBench DR
INO-NS-Ie-3 Datasets with Long Trajectory
INO-NS-Io-4
PDEArma NS-Force
Figure 7: Results of PITA and Rege-Regeessive Methods on Long Trajectory Datasets.
E.3. Scaling Experiments
In this section, we analyze the scaling results from Table 1, which are illustrated in Figure 8. It is observed that as the model size increases, the average nRMSE across the five datasets with long trajectories decreases, generally following a scaling law. However, on two larger-scale datasets (FNO-NS-Ie-3 and FNO-NS-Ie-A PITA performs less stably on the larger-sized baselines. The main reason may be that the DPOT baseline uses cyclic learning rate decay incorporating very high learning rates during training. This may lead gradient explosion or overly aggressive weight updates, affecting the convergence of models with more parameters. Furthermore, DPOT performs finetuning for all model sizes with only 500 epochs and sets the gradient clipping value to 10,000. This prevents PITA from converging stably on larger datasets with larger model sizes. We found that this issue can be mitigated after adjusting the training settings. For example, when we provide sufficient training, PITA shows more stable convergence results, as shown in Table 10 of Appendix E.4. To demonstrate PITA's plug-and-pe aod--and-ply performance more fairly, we still follow DPOT's training settings in the experiments, without making any adjustments to PITA, except for special declarations.
19

---

Physics-informed Temporal Alignment for Auto-regressive PDE Foundation Models
0.0024
wit
0.006
DPCT with PITA DPOT
DPCT
0.0022
0.0020
nRMSE
MRES
0.004
0.0018
0.0016
0.003
0.0014
0.002
0.0012
30
100
500
30
100
500
Model Size (M)
Model Size (M)
DPCT with PITA
0.035
DPOT with PITA
DPCT
DPOT
0.28
0.030
0.26
nRMSE
nRMSE
0.025
0.24
0.020
0.22
0.20
0.18
0.16
30
100
30 Model Size (M)
500
Model Size (M)
Figure& 8: Results of scaling experiments for different dataset sizes. In the first row, we present results for PDEBench-SWV (left) and FNO-NS-1e-3 (right). In the second row, we show results for NS-Force (left) and FNO-NS-Ie-4 (right).
E.4. Comparative Analysis of Training from Scratch and Finetuning Pretrained Models
To compare the performance of pretrained models with those trained from scratch, we selected four datasets, shown in Table 10. Testing on the PDEBench-DR and FNO-NS-Ie-5 datasets represents in-distribution testing, while the PDEArena- SWE and Burgers serve as the onttt---dittiittion test. It is noteworthy that although the shallow-water equation was included in the training dataset for the DPOT model, the only physics variable observed during training was the water depth. In the out-of-distribution context, in addition to predicting the surface height h, the model will also predict 88 steps of zonal velocity u, meridional velocity v. pressure field p. and wind vorticity field g.
Firstly, for in-distribution test, the results from finetuning are slightly better than those achieved by training from scratch using the auto-regressive and PITA methods. This improvement is attributed to the former's 's training over 2500 epochs, compared to only 1,500 epochs for the latter. With a reduction of 19.684% in nRMSE, the training cost for finetuning is 1.5 times that of training from scratch using PITA or the auto-regressive method. This suggests that large scale pretraining on PDE datasets may not be effective and does not scale well with an increased number of training epochs. Secondly, in out-of-distribution test, the results from finetuning using pretrained models are generally worse than those from training from scratch. This indicates that during the pretraining process, the model primarily Iearns the physical laws from the training datasets, lacking the capacity to generalize to unseen tasks where the underlying physical principles are unknown. Moreover, the physics laws governing different families of PDEs share little similarity, even among PDEs within the same family but with varying parameters. As result, a foundation model may struggle on distt--iiittiittto tests. Lastly, when training from scratch, the results are generally better than those obtained using the auto-regressive method. This observation highlights PITA's potential application in PDE foundation models.
E.5. Comparison of PITA and Auto-regressive Methods (Tar = 10)
In this section, we present the training results for Tar = 10. Each dataset's full trajectory is predicted, with the roll-out
prediction length fixed at T., = 10. Consequently, Table 11 contains only the results for long trajectory predictions. is observed that when the roll-out window length is set to 10, the results are generally poorer compared to those from Tar = 1. This decline in performance can be attributed to the shortcut problem that occurs during training.
20

---

Physics-informed Temporal Alignment for Auto-regressive PDE Foundation Models
Table 10: Experimental Results of Pretrained Models Versus Models Trained from Scratch Using Auto-regressive and PITA Methods. Finetune in the table indicates that the model is finetuned on a pretrained model of a specific size. Train from Scratch (Auto-regress) refers to models trained using the auto-regressive approach without any pretrained models. Train from Scratch (PITA) denotes models trained using the PITA approach without any pretrained models.
In Distribution
Out of Distribution
PDEBench-DR
FNO-NS-Te:
PDEArema
Burgers
Model
Strategy
Epoch
500
1000
1500
500
1000
1500
500
1000
1500
500
1000
1500
Finetune (ugo-r)gress)
0.12281
0.04249
0,00522
0.11144
0.06398
0,03595
18,46497
5.24994
1.01142
0,03712
0.02543
0.00455
DPOT-T
Train from Scratch (Aulio regness)
0.29705
0.05467
0.00651
0.18219
0.07222
0.04072
5,98731
3.20792
0.99900
0.06691
0.01526
0.00199
Train from Scratch (PITA)
0.42291
0.03921
0.00577
0.15501
0,07312
0,04009
2,07097
3.89625
0.99710
0.06945
0.01318
0.00374
Finetune (Aulo regress)
0.14761
0.02113
0.004566
0.11242
0304399
0.03054
44,70683
2.69744
039112
0,04679
0.01627
010004
DPOTS
Train from Scratch regress)
0.20054
0.05483
0,00617
0.11495
0.04570
0,03771
8.0087
2.668499
0.99414
0.06251
0.02703
0.00383
Train from Scratch (PITA)
0.80244
0.05853
0.00807
0.11741
0.06509
0.03701
16.16088
1.87741
0.99051
0,04121
0.01943
0.00307
Table 11: Comparisons of PITA and Auto-regressive Strategies Across Various Models and Datasets (Tar = 10).
FNO.NS lo-4
FDB1Bench
Pittbench CNS-ity C)
PDEArea
CFDbBench
lo.5
lo.3
DR
SWF
1.011
1.001
B.1.D.1
0.1.0.01
MB.1
NS-Foree
Small Medels
DPOFER
0:0779
033665
0.0342
0.0063
01515
010343
0,0434
0,0492
012133
0J0352
0.344
0.0762
QUMEZT
0,0897
0,0692
0,0058
0.0286
0,0032
0,8505
0.0492
0.0531
0/0622
0,0019
B
Auto regress
0.0879
0.0106
0,002
0.0312
0.0036
0.0678
0.0280
0.0479
0.0439
0.0357
0.0358
0.2185
0.0570
0.0028
30M
PITA
0,0790
0.025
0,113
0.0016
0.8548
0,0011
0.0425
0.0496
00311
0.0358
0,2335
0.0621
0
DPOTEM
Ante-repres
D.0755
B3402
0.0311
0.8554
0,0032
0.512
0.216
0.036
0,0517
00216
0J666
0.2256
0.0540
0.099
122M
PITA
0,0630
0,0285
0.0019
0.0306
0.0020
0.8468
0.0175
0.0322
0,0532
0.008
0.0370
0.1878
0.0500
0.0022
FNO-M
Auto-regress
0.0023
0.0450
0.0233
0.3054
0.3446
0.0041
170M
0.1300
0.0148
0
0.0030
0.8478
0.0414
0.0446
0.0631
00310
0.0466
0.2485
0.1997
0.0049
Large Models
DPOT-L
0J0031
0.0544
0,0018
UXGUN
0U0344
008267
0.0436
0.1741
0,0562
0U0036
0.0036
0.0019
0.0312
0.0417
0.1675
0,0509
0.0039
E.6. Comparison of Sparse Regression and Other Numerical Algorithms
The choose of sparse regression framework for PDE discovery is primarily based on its theoretical and experimental advantages. Firstly, the la regularization in sparse regression in our PDE discovery step inherently suppresses noise by enforcing sparsity in the candidate coefficient space. As noted in compressed sensing theory (Donoho, 2006), sparse regularization, particularly tu regularization, prunes amall-magnitude terms caused by observational noise, effectively driving them to zero while preserving dominant dynamical terms (Huang & Aviyente, 2006; Wen et al., 2018). Additionally, prior work (Chen al., 2021; Rudy et al., 2017) has provided experimental evidence that sparse regression can successfully identify govemning equations with high accuracy even in the presence of noisy data. The fo-norm regularization adopted in this paper systematically eliminates noise-dominated low-magnitude terms through a hard thresholding mechanism.
Secondly, compared to the sparse regression algorithm. existing regression methods exhibit significant limitations in complex correlation and noise scenarios. While the classical LASSO (Tibshirani, 1997) method achieves sparsity through l regularization, its convex relaxation property tends to randomly select collinear features when high correlations exist between columns of the data matrix (as demonstrated in the typical PDE identification scenario in reference (Knowles & Renka, 2014)), leading unstable identification of true physical patterns and resulting in non-unique solutions and pseudo-sparsity phenomena. The sequentially thresholded least squares (STLS) method (Budišić et al., 2012) improves sparse identification through a recursive thresholding mechanism, but its iterative process based on ordinary least squares lacks regularization mechanisms. In the presence of ill-conditioned matrices or highly correlated features (as shown in the experiments of reference (Budišić et al., 2012)), the computation of unregularized inverse matrices amplifies noise sensitivity and leads to cumulative parameter estimation bias with increasing iterations. Regarding interpretability, although traditional ridge regression improves matrix condition numbers and mitigates collinearity effects through & regularization (as theoretically analyzed in reference (Noack et al., 2003)), its inherent non-sparse solution property fundamentally conflicts with the interpretability requirements of physical models.
Thirdly, in our experiments, when employing other non-sparse methods such as pseudo-invers or least squares regression (Golub& & Pereyra, 1973), the resulting coefficient matrices exhibit significant non non-physical oscillations due to the lack of sparsity constraints in the solution space. This leads to an increase in the second-order norm of the residual terms by 2 or 3 orders of magnitude. Such ill-conditioned solutions make it difficult to effectively normalize LPyy and can also trigger gradient explosion phenomena during back propagation.
21

---

Physics-informed Temporal Alignment for Auto-regressive PDE Foundation Models
Table 12: Comparisons of Numerical Algorithms Adopted for PITA in the DPOT-Ti Model.
PDEBench-DR
FNO-NS-Ie-5
Numerical Algorithm
Test Loss (91 Steps)
Test Loss (IO Steps)
Pseudoinverse
0.02016
0.08971
Least Squares Method
0.06903
0.09470
Sparse Regression 0,01090
0.05629
E.7. Comparison of Computational Cost Between PITA and Auto-regressive Methods
PITA differs from conventional auto-segressive methods by incorporating a PDE discovery procedure, which introduces additional training-time computation and CPU memory usage. We measured per-batch training time, the extra time incurred by the PDE discovery process ("STRidge time" in Table 13), and CPU memory consumption. Given that GPU memory usage remains invariant relative the baseline, it is therefore excluded from our measurements. The results summarized in Table 13 and Figure 9.
As shown in Figure 9 (right), the overhead due to PDE discovery (detailed in Section 3.3) is constant and independent of model size. During inference, PITA matches the baseline's efficiency-beeause discovery and alignment occur only during training-and and consistently outperforms it on long-term prediction benchmarks (see Figure 7).
Figure 9 (left) shows that PITA requires no additional GPU memory but incurs a fixed, moderate increase in CPU memory for candidate library construction and sparse regression. This decoupling of computational overhead from model dimensionality enables PITA scale to billion-prametter architectures without prohibitive resource demands. For example, 500 M- parameter, L-level model incurs only 0.517s extra per batch 20% increase) while achieving a 31.61% % improvement in long-trajectory prediction accuracy.
Table 13: Computation Cost Comparison Between PITA and Auto-regressive.
Size
Model
Training Time (s)
STRidge Time (s)
CPU Memory Consumption (MB)
PITA
0.4124
0.0959
1429.43
Ti
Auto-regressive
0.2941
1060.53
PITA
0.5783
0.0847
1533.21
S
Auto-regressive
0.3464
1095
PITA
1.2012
0.1124
2207.5
M
Auto-regressive 0.8924
1937.48
PITA
3.0381
0.1098
2581.84
L
Auto-regressive
2.5207
2203.63
E.8. Influence of the Percentage of Remained Library Terms
To evaluate the impact of the incomplete library, we conducted experiments with subsampled candidate term sets, as shown in Table 14. When randomly retaining 50% of the library terms, the test loss on PDEBench-SSV SWE increases from baseline of 0.00137 (achieved with full-term libraries) to 0.00213. Notably, this performance degradation diminishes when preserving 80% of library terms, where the test loss stabilizes 0.00135 statistically comparable the full-library configuration. Crucially, this finding indicates that the library architecture requires no specialized customization for individual PDE systems.
Table 14: Influence of the Percentage of Remained Library Terms Test Loss.
Percentage the Remained Library Terms
Test Loss (10 Steps)
Test Loss (Full Length)
50%
0,00209
0.00213
80%
0,00145
0.00135
100%
0.00150
0.00137
22

---

Physics-informed Temporal Alignment for Auto-regressive PDE Foundation Models
CPU Mcmory Consumption (MB)
2600
DPOT with PITA
Sparse Regression Processing Time Size
DPOT
2400
0.1124
0.1098
2200
0,0959
Time (seconds)
2000
0,0824
0,08
1800
0,05
1600
1400
0.04
1200
0.02
1000
30
500
0.00
Model Size (M)
Model Size
Figure 9: Results of Computational Costs Between PITA and Auto-regressive Method. In the left, we present the CPU memory consumption of each approach. In the right, we show the additional computational time for the sparse regression process of PITA.
E.9. Influence Accuracy of the Discovered Physics Accuracy on PITA's Performance
In this discussion, we designed specific experiment to evaluate PITA's robustness against dynamic failure, where we deliberately injected noise into the gradient computation PITA's physics loss to simulate inaccurate or corrupted dynamics. As illustrated in Table 15, under such direct attack the physics constraint, PITA still exhibits strong robustness.
Table 15: Influence of Discovered Physics Accuracy on PITA Under Noise Attack in Terms of Physics Loss.
PDEBench-SWE
Method
nRMSE (10 Steps)
nRMSE (91 Steps)
PITA (Salt Pepper Noise) PITA (Without Noise)
0.00224
0.00256
0.00202
0,00199
Auto-regressive
0.00208
0.00241
We further validate robustness by perturbing the PDE spatial grid (see Table 16), showing that even with an intentionally distorted grid, the method maintains performance parity with unperturbed settings. Even when the PDE discovery fails, the multi-task learning framework mitigates its negative impact, ensuring that the model degrades gracefully to the baseline performance.
Table 16: Influence of Discovered Physies Accuracy on PITA in the DPOT-Ti Model.
PDEBench-SWWi
PDEBench-DR
Grid Type
nRMSE (10 Steps)
nRMSE (91 Steps)
nRMSE (10 Steps)
nRMSE (91 Steps)
Wrong Grid
0.0024
0.00231 0.00199
0.02115
0.01684
Correct Grid
0.00202
0.02119
0.0109
E.10. Performance on Modeling Chaotic Dynamical Systems
To rigorously evaluate PITA's capacity for modeling chaotic dynamical systems under extended temporal extrapolation, we conducted controlled experiments on synthetic 2D Kolmogorov turbulence flows, which is a canonical benchmark for chaotic PDE systems exhibiting multiscale energy cascades and nonlinear dissipation processes. Our turbulence simulation dataset contains 300 spatiotemporal trajectories with 64 x 64 grid resolution and 150 temporal steps, intentionally designed to challenge long-term forecasting fidelity. Both methodologies leverage the same pretrained DPOT-Ti model. As shown in Table 17, our results demonstrate that PITA significantly outperforms the auto-regressive baseline in capturing the intricate patterns of chaotic fluid evolution. Specifically, PITA exhibits superior alignment with the underlying physical dynamics, showcasing its enhanced ability to model complex, long-term behaviors in turbulent systems.
23

---

Physics-informed Temporal Alignment for Auto-regressive PDE Foundation Models
Table 17: Test Loss on Kolmogorov Turbulence Flow.
Test Loss (10 Steps)
Test Loss (150 Steps)
Auto-regressive
0.2581
0.4459
PITA
0.2247
0.3152
E.11. Robustness Analysis of PITA under Noise and Data Masking
One of the most significant challenges facing PDE foundation models is their ability to cope with noisy or incomplete data in real-world scenarios. To assess PITA's robustness, we firstly conduct a series of experiments in which Gaussian noise is injected into the dataset at three different levels: 0.05, 0.005 and 0.0005. As summarized in Table 18, PITA consistently outperforms all baseline methods under every noise condition, achieving an average accuracy improvement of 51.79 %. Notably, even at the highest noise level of 0.05, PITA maintains stable performance, demonstrating much smaller degradation compared to its peers. These results not only confirm PITA's strong resilience data corruption but also highlight its practical potential for deployment in applications where high-quality measurements are difficult to obtain.
Table 18: Test Loss with Noisy Data at Different Noise Levels on the DPOT-Ti Model with the FNO-NS-Ie-3 Dataset.
Auto-regressive
PITA
Noise Scale
nRMSE (10 Steps)
nRMSE (20 Steps)
nRMSE (10 Steps)
aRMSE (20 Steps)
0.05
0.32331
0.45795
0.2017
0.3214
0.005
0.01947
0.0631
0.00828
0.01241
0.0005
0.00402
0.00608
0.00237
0.00333
Secondly, we evaluate PITA'S capability to reconstruct solutions from incomplete observations. Specifically, 25% of the PDE dataset is randomly masked, and the model is tasked with predicting the full field, including within the oceluded regions. As reported in Table 19, PITA significantly outperforms the auto-regressive baseline under these conditions. While all methods experience some performance degradation when confronted with missing data, PITA's relative drop in accuracy is markedly lower. We attribute this resilience to PITA's integrated PDE- PDE-discovery module, which infers the governing dynamics even from downsampled, partially missing inputs, enabling it to generalize beyond the observed measurements. These results demonstrate PITA's enhanced ability to recover accurate solutions in data-sparse settings, highlighting its potential for deployment in real-world scenarios where measurement gaps are ubiquitous.
Table 19: Test Loss with Incomplete Data on the DPOT-Ti Model with the PDEBench- Dataset.
Auto-regressive
PITA
Mask Ratio
nRMSE (I0 Steps) nRMSE (91 Steps)
nRMSE (10 Steps)
nRMSE (91 Steps)
0.25
0.02176
0.1264
0.0112
0.0801
E. Error Bars of PITA on Long Trajectory Datasets
To ascertain that the improvements of PITA demonstrated in Table 1 stem specifically from its PDE discovery framework and alternating direction optimization algorithm, we evaluate the performance of PITA against auto-regressive methods on datasets featuring extended temporal trajectories. Such datasets inherently pose challenges for long-term prediction and downstream tasks due to their temporal complexity. By assessing multiple model scales, further investigate how variance patterns and reliability metrics correlate with model complexity. As illustrated in Figure 10, the results confirm that PITA's enhancements are statistically significant and intrinsically tied to its architectural design, rather than stochastic fluctuations.
F. Visualizations of Trajectory
In this section, we visualize datasets predictions presented in Table
24

---

Physics-informed Temporal Alignment for Auto-regressive PDE Foundation Models
Swaddad Deviction on FDE SWB
Standard Deviction on PDEBeoch
Standard Destatio on FNO 16:5
6:00
6.055
Mon Vilue
6.05
Mau Valle
0.004
Men
600
0.044
0.010
0.03
6.022
0.01
6.000
0.000
BPSTSS
DPUTE Modd
PNG-
UNTIL
DROTS
IROEM Model
INO-
DNPIELL
PNO-M
2NTH
Model
Sandaad Devistivn on FNO-NS-IE4
Standend Devistion on POEAAm-NN--ooor
07000
0.4
0.025
Maan Wiee
Moan Nale
0.000
0.055
0000
0.1
0.00
0.000
0001-11
DIYTES
PN0-M
DIOTE
6:0
DI01-
DINTEM Model
IND-M
DIVEL
Model
Figure 10: Mean Value and Standard Deviation on Datasets with Long Trajectory.
F.1. PDEBench-DR
In this section, we visualize two quantities from the diffusion-reaction equation: the density fields u and v.
25

---

Physics-informed Temporal Alignment for Auto-regressive PDE Foundation Models
t+1
t+11
c+16
1+:1
t+26
to:31
1436
t+:41
1146
t+51
L456
L+61
L+66
[+7]
[+76
[+81
L+86
[+9]
F.2. FNO-NS-V
Here we visualize the vorticity w of the incompressible Navier-Stokes equation, with different viscous coefficient v
1 1x 10 1 10 10
== 1 10-3
>= 10
26

---

Physics-informed Temporal Alignment for Auto-regressive PDE Foundation Models
I+2
6+3
1+d
t+5
t+7
t+8
[+9
t+10
IMIA heddi dion
Ground Iruzhs
b+11
t+12
t+13
t+14
t+15
t+16
t+17
t+18
t+19
t+20
PIIA Prodic tion
Gooand Trutt
I+2
t+3
++4
145
6+6
t+7
t+8
t+9
t+10
PITA Medictio
Greand Iouhh
t+11
[+12
L+13
t+14
t+15
t+16
t+17
t+18
t+19
(+20
I11A Prodiciin
Groand Ituthh
27

---

Physics-informed Temporal Alignment for Auto-regressive PDE Foundation Models
10-5
[+1
t+2
t43
t44
t+5
1+6
t+7
t+8
t+9
t+10
PITA Prediction
Ground Touth
F33. PDEArena
Here we visulize the velocity v. pressure p, and density fields p from two datasets from PDEArema-NS1/ NS- NS-Force and NS.
NS-Force
[+]
[+11
t+21
[+31
[+4]
FITA Prodiction
Ground Trum
t+11
t+16
t+21
t+31
t+36
t+41
t+44
PITAA Prodi cton
Ground Trads
t+1
t+6
t+11
t+16
t+21
t+26
t+31
t+:36
t+44
t146
PITA Prodicton
Ground Touth
28

---

Physics-informed Temporal Alignment for Auto-regressive PDE Foundation Models
NS
t+1
t+22
t+3
t+4
Ground Prediction
t+1
t+22
t+3
t+4
Preciction
Ground
t+1
h+2
t+3
t+4
Preciction
Ground Truth
29

---

Physics-informed Temporal Alignment for Auto-regressive PDE Foundation Models
F.4. PDEBench-CNS
Here we visualize the velocity fields and v, pressure p and density p of the compressible naive-stokes equation with different parameters.
(), = 0.01)
t+1
t+2
t+3
t+4
1+5
[+6
t+7
1+8
1+9
t+110
t+11
++7
[+
149
1+100
[+11
1+1
1+2
1+3
1+4
1+5
146
1+7
1+8
149
1+100
I+11
30

---

Physics-informed Temporal Alignment for Auto-regressive PDE Foundation Models
(M, (0.1, 0.01)
++1
++2
1+3
1+4
145
146
1+7
1+8
149
14:10
[+11
++1
1+2
t+3
t+7
[+
1+9
1+10
[+11
H+1
1+2
1+3
1+
1+5
1+6
1+7
1+1
1+9
1+100
C-11
31

---

Physics-informed Temporal Alignment for Auto-regressive PDE Foundation Models
(M, - (1,
[+1
1+2
++3
1+4
145
[-6
14-
1+8
144
14:10
[+11
++1
1+2
t+3
[+6
t+7
1+8
1+9
t+10
[+11
H+1
++2
t+3
[+4
1+5
1:46
1+7
1+8
1+9
1+110
C-11
32

---

Physics-informed Temporal Alignment for Auto-regressive PDE Foundation Models
(M, (0.1,0.)
[+1
1+2
1+3
145
T-6
1+7
1+8
149
14110
[+11
++1
1+2
t+3
[+6
t+7
1+8
1+9
1+10
[+11
H+1
1+2
1+3
1+4
1+5
1+6
1+7
1+-
1+9
1+110
C-11
33

---

Physics-informed Temporal Alignment for Auto-regressive PDE Foundation Models
(M,n) - 10
1+1
1+2
1+3
1+4
1455
146
1+-
1+8
149
14-10
[+11
++1
1+2
t+3
[+6
t+7
1+1
1+9
1+10
h+11
++1
1+2
1+3
1+4
1+5
1-46
t+7
1+1
1+9
1+100
1+11
34

---

Physics-informed Temporal Alignment for Auto-regressive PDE Foundation Models
F.5. Magnet- Viscous Burgers
Saagnut-Viscous
Here visualize the field from viscous burgers' equation.
PrlA Prodician
Ground truth
[+11
6+12
t+13
++14
t+15
[+16
++17
t+18
t+19
t-20
PYTA Poealeeeeeeo dan
Ground Irath
t+21
L+22
t+23
1+24
L+25
L+26
L+227
t+28
L+29
1+30
PTAA Prodde on
Ground Irutts
t+:31
++32
t+-33
h+34
t+35
1+:36
++-37
t+38
t+-39
t=40
PTTA Prackation
Ground Irunds
35

---

Physics-informed Temporal Alignment for Auto-regressive PDE Foundation Models
F.6. CFDBench
Here visualize the velocity and pressure fields
RITA Predictio
Ground Truth
[+1
I+2
L+3
++
L+5
L+6
L+7
1+8
1+9
I+10
FFAA Prodictinn
Ground Truhh
F.7. PDEBench-SWE
Here visualize water depth h from the whallow-water equation.
L+-6
t+11
E+16
L+21
[+26
t+31
L+36
[+:4]
[+46
FITA Prosicton
Gound Ibver
PITA Frodiction
Grunnd mruh
36


Reviews:
[{"review_id": "NJ7EXNJjvv", "reviewer": "Reviewer_SLyt", "summary": "The paper introduces Physics-informed Temporal Alignment (PITA), a new self-supervised learning framework aimed at improving autoregressive PDE foundation models. The authors identify a \"shortcut\" issue common in autoregressive models, where the model takes easy solutions, leading to accumulated prediction errors, especially in long trajectories or out-of-distribution data. To fix this issue, PITA combines physics-informed constraints with traditional autoregressive predictions. This is done by first discovering PDE equations from observation data and then aligning the predicted trajectories with these discovered physics constraints.\n\nThe main contributions are: (1) integrating PDE discovery with autoregressive prediction to solve the \"shortcut\" problem; (2) proposing a combined loss function with data, physics, and consistency terms, balanced using uncertainty-based weighting; and (3) demonstrating strong results across several PDE benchmarks (e.g., PDEBench, PDEArena, CFDBench), achieving better accuracy than existing foundation models, especially for long-term predictions.", "questions": "", "limitations": "", "rating": 3, "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": "1.Regarding the additional cost of PDE discovery and alignment.\n\nReply 1: We appreciate the reviewer highlighting computational scaling. For fixed batch sizes, the additional time from PDE discovery and alignment remains constant, independent of dataset scale or trajectory length, since it depends solely on parameter $T_{ar}$ (Sec. 3.2). Furthermore, for single-dataset fine-tuning, pre-training incurs baseline cost only, while PITA overhead occurs exclusively during fine-tuning, making its overall training time negligible.\n\n2.Regarding additional computational cost and its limitations.\n\nReply 2: Thanks for your valuable suggestions. As shown in https://anonymous.4open.science/r/PITA_1/Table1.pdf, the time consumption of the PDE discovery and alignment steps is constant and does not scale with the model size. Moreover, PITA performs PDE discovery and alignment only during the training phase, so its inference efficiency is identical to the baseline model. Due to space limitations, please refer to the discussion in Reply 3 to Reviewer y7j5.\n\n3.Regarding noise or incomplete observational data.\n\nReply 3: We appreciate your concern about noise sensitivity. The L0 regularization (Sparse Regression) in our PDE discovery step inherently suppresses noise by enforcing sparsity in the candidate coefficient space. As noted in compressed sensing theory [5], sparse regularization, particularly L0 regularization, prunes small-magnitude terms caused by observational noise, effectively driving them to zero while preserving dominant dynamical terms [1, 2]. Additionally, prior work [3, 4] has provided experimental evidence that sparse regression can successfully identify governing equations with high accuracy even in the presence of noisy data. Regarding incomplete data, PITA employs local sampling in both time and space. To further demonstrate the robustness of PITA, we conducted experiments under two challenging settings: (1) noisy data with noise levels of 0.05, 0.005, 0.0005, and (2) incomplete data where 25% of the observations were randomly masked. As shown in the results https://anonymous.4open.science/r/PITA_1/Table6.pdf and Table7.pdf, PITA consistently outperforms baselines even under these adverse conditions. This illustrated the robustness of L0 regularization against noisy or incomplete observational data. \n\n4.Regarding the choice of the sparse regression approach.\n\nReply 4: As previously discussed in terms of noise robustness, sparse regression achieves the separation of physical law extraction from noise interference through explicit sparsity constraints. The L0-norm regularization adopted in this paper systematically eliminates noise-dominated low-magnitude terms through a hard thresholding mechanism. Whereas LASSO’s L1 regularization suffers from collinear feature bias due to convex relaxation [1], and STLS lacks inherent regularization in its recursive ordinary least squares framework, our method achieves principled sparsity control. Notably, while ridge regression’s L2 penalty improves matrix conditioning [3], it fundamentally conflicts with sparsity requirements through non-zero coefficient retention. This structured sparsity enforcement proves essential for disentangling physical laws from noise interference. Besides, in our experiments https://anonymous.4open.science/r/PITA_1/Table5.pdf, when employing other non-sparse methods such as pseudo-inverse or least squares, the resulting coefficient matrices exhibit significant non-physical oscillations due to the lack of sparsity constraints in the solution space. Such ill-conditioned solutions make it difficult to effectively normalize $L_{Phy}$ and can also trigger gradient explosion phenomena during backpropagation.\n\n5.Regarding the failure cases of PDE discovery\n\nReply 5: The discussion regarding noise and incomplete data has already been provided in Reply 3. In this discussion, we designed a specific experiment to evaluate PITA's robustness against dynamic failure, where we deliberately injected noise into the gradient computation of PITA's physics loss to simulate inaccurate or corrupted dynamics. As shown in https://anonymous.4open.science/r/PITA_1/Table9.pdf, even under such a direct attack on the physics constraint, PITA still exhibits strong robustness. We further validate robustness by perturbing the PDE spatial grid (see https://anonymous.4open.science/r/PITA_1/Table8.pdf), showing that even with an intentionally distorted grid, the method maintains performance parity with unperturbed settings. As discussed with Reply 3, even when the PDE discovery fails, the multi-task learning framework mitigates its negative impact, ensuring that the model degrades gracefully to the baseline performance.\n\n[1] Sparse representation for signal classification.\n\n[2] Robust sparse linear discriminant analysis.\n\n[3] Physics-informed learning of governing equations from scarce data.\n\n[4] Data-driven discovery of partial differential equations.\n\n[5] Compressed sensing.", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "Yes, it uses standard PDE datasets for evaluation, such as PDEBench, PDEArena, CFDBench.", "theoretical_claims": "The paper does not provide theoretical analysis or proofs; it is primarily empirical", "experimental_designs_or_analyses": "The comparisons with existing autoregressive PDE foundation models, such as DPOT and MPP, provide a reasonable basis for assessing the effectiveness of PITA. The use of ablation studies to analyze the impact of different components, including loss functions, downsampling strategies, and model sizes, strengthens the validity of the experimental results. While PITA improves accuracy, the additional cost of PDE discovery and alignment is not fully analyzed in terms of runtime or scalability to larger datasets and longer trajectories.", "relation_to_broader_scientific_literature": "The paper proposes physics-informed temporal alignment (PITA), a self-supervised learning framework to improve the accuracy of existing PDE foundation models.", "essential_references_not_discussed": "No", "other_strengths_and_weaknesses": "Strengths: (1) The idea of combining PDE discovery with autoregressive modeling is innovative, addressing the common shortcut issue clearly and effectively. (2) The paper is well-motivated, providing good insight into why existing autoregressive approaches might fail on long-term trajectories.\n\nLimitations: (1) Although the approach significantly improves performance, it could be computationally intensive due to the additional complexity introduced by PDE discovery and alignment steps, which might limit its applicability in practice. (2) The sensitivity of the method to noisy or incomplete observational data isn't fully addressed; this aspect might substantially impact its applicability in real-world scenarios.", "other_comments_or_suggestions": "N/A", "questions_for_authors": "1. How sensitive is the PDE discovery step in PITA to noise or incomplete observational data?\n2. Why did you choose the sparse regression approach specifically for discovering PDE terms? Would other methods perform similarly or better? \n3. How would PITA perform if PDE discovery fails to find the correct governing equations, especially when dealing with noisy or limited data?", "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": "", "extra_scores": {}}, {"review_id": "G9UbKNESYZ", "reviewer": "Reviewer_Xe1t", "summary": "This paper focuses on the shortcut bias that can occur in PDE auto regressive models. The proposed method suggests enhanced predictions using some physical knowledge extracted from data, using sparse regression. The results show some improvement in the rollout prediction.\n\n## update after rebuttal\nAfter the rebuttal discussion, I am in favor of acceptance of the paper.", "questions": "", "limitations": "", "rating": 4, "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": "1.Regarding the comparison of the shortcut problem.\n\nReply 1: Thanks for your insightful comments. We followed the error visualization methodology in [1], plotting rolling-step MSE for each long-term dataset, as shown in the Shortcut folder of https://anonymous.4open.science/r/PITA_1/ (see 'Shortcut' folder). The results show that on datasets such as PDEBench-SWE, FNO-NS-1e-3, and FNO-NS-1e-4, the prediction error does not exhibit significant accumulation as time progresses. While errors exhibit a slight increasing trend over time in other datasets, PITA still reduces error accumulation compared to the auto-regressive method. This observation strongly indicates that PITA effectively mitigates error propagation over extended temporal sequences.\n\n2.Regarding additional computational cost. \n\nReply 2: Due to space limitations, please refer to the discussion of this topic in Reply 3 to Reviewer y7j5. We kindly invite you to consult that section for a detailed discussion.\n \n3.Regarding illustrating the shortcut.\n\nReply 3: Thanks for your suggestion. In the revision, we will incorporate a detailed explanation of the shortcut problem, including its definition, standard approaches, and its specific implications for PDE research. \n\n4.Regarding the performance of the PITA with incomplete library terms.\n\nReply 4: To evaluate the impact of the incomplete library, we conducted experiments with subsampled candidate term sets, as shown in https://anonymous.4open.science/r/PITA_1/Table3.pdf. When randomly retaining 50% of the library terms, the test loss on PDEBench-SWE increases from a baseline of 0.00137 (achieved with full-term libraries) to 0.00213. Notably, this performance degradation diminishes when preserving 80% of library terms, where the test loss stabilizes at 0.00135 – statistically comparable to the full-library configuration.  Crucially, our findings indicate that the library architecture requires no specialized customization for individual PDE systems. \n\n5.Regarding Line 149, right column.\n\nReply 5: We sincerely apologize for the oversight. The parameter w specifically denotes the resolution of the downsampled grid. \n\n6.Regarding Lines 163, 164.\n\nReply 6: Thanks for your meticulous attention to conceptual consistency. The notation θ in lines 163–164 indeed refers to the trainable parameters of models (e.g., FNO, DPOT, MPP), which aligns precisely with the parameter definition in Line 244. \n\n7.Regarding the stability of physical losses in the training objective.\n\nReply 7: Prior studies highlight a trade-off between data loss and physics loss in multi-task optimization, caused by conflicting gradient directions, which can result in unstable training or slow convergence [3,6]. \nInspired by [2], we adopt an uncertainty-based weighting strategy to balance the three losses, as detailed in Sec. 3.4. As shown in Table 3, PITA achieves a 48.97% improvement in test accuracy compared to the auto-regressive baseline. Furthermore, compared to manually specified weights for different losses (Table 3, Task 4), the uncertainty-based weighting strategy shows significant advantages.\n\n8.Regarding removing the data loss.\n\nReply 8: PITA fundamentally relies on autoregressive, data-driven signals. Removing data loss disrupts its ability to capture intrinsic PDE dynamics. Abandoning the data-driven paradigm entirely would require reconstructing the foundation model architecture, like Spline-PINN, or stronger physical priors, potentially limiting generalizability to unknown dynamics.\n\n9.Regarding the computation of the sparse regression.\n\nReply 9: Sparse regression indeed performs on temporally and spatially downsampled trajectories, as systematically validated in Table 3 of the submission. Temporal downsampling and spatial subsampling force the model to capture essential system dynamics and scale-invariant patterns while reducing noise-induced overfitting, which aligns with findings in PDE-Find [5].\n\n10.Regarding 𝜆∗ and incomplete library functions.\n\nReply 10: Both $\\lambda^*$ and $\\lambda_i(\\theta)$ are computed via sparse regression algorithms, as detailed in Algorithm 1, with the distinction that $\\lambda^*$is derived from ground-truth data while $\\lambda_i(\\theta)$originates from predicted data.\t\nMoreover, we kindly refer the reviewer to Reply 4 for a detailed discussion on the issue of incomplete candidate libraries.\n\n[1] A WENO-based method of lines transpose approach for Vlasov simulations.\n\n[2] Multi-task learning using uncertainty to weigh losses for scene geometry and semantics.\n\n[3] Multi-task learning as multi-objective optimization.\n\n[4] Spline-pinn: Approaching pdes without data using fast, physics-informed hermite-spline cnns.\n\n[5] Data-driven discovery of partial differential equations.\n\n[6] PINN training using biobjective optimization: The trade-off between data loss and residual loss.", "claims_and_evidence": "The paper shows extensive experiment on the gain in term of performances of the proposed method. However, I think that some experiment highlighting the shortcut problem would benefit to the argumentation. While it is detailed as being the main focus of the paper, it is hard for the reader to identify whether the proposed method succeeds to solve this issue. Some visualizations are proposed in the appendix, but maybe some focused on these could be a first illustration? It is hard to identify from mse rollout if the proposed method effectively handles this problem or not.", "methods_and_evaluation_criteria": "The proposed method is evaluated on several standard and recent models and datasets of the literature.", "theoretical_claims": "NA: no theoretical claims in the paper.", "experimental_designs_or_analyses": "Experiments highlight the improvement in performances when using PITA. However, no comparison on the shortcut problem is provided, which could help in assessing the effectiveness and understanding the effect of the proposed method on the predicted trajectories. Moreover, since the method involves solving sparse regression, I think providing training and inference time comparison is an important point.", "relation_to_broader_scientific_literature": "Maybe that providing an illustration of what the shortcut problem is and more importantly, what does it implies for PDE would give insights to the reader (using standard methods of the literature).", "essential_references_not_discussed": "NA up to my knowledge", "other_strengths_and_weaknesses": "- I appreciate the limitation section, providing insights on the limitations of the methods. While reading the paper, the 2 main limitations mentioned were the one I thought about. However, I think this would be interesting to illustrate these limitations with some ablation in the appendices. \n- No computational time comparison is provided, while it is a great problem for PDE solving (it is discussed in the limitations section). Moreover, adding regression suggest an increase of the computational time. This should be detailed in the appendix. \n- As discussed in the limitation section, It would complement the argumentation to evaluate the performances of the PITA method with uncomplete library terms, to evaluate the sensibility of the method to the prior assumption made as well as wrt the library size. \n- The paper is well written, easy to follow and shows extensive experiments, well detailed.", "other_comments_or_suggestions": "- Line 149 right column, w not introduced before\n- Lines 163, 164, theta not introduced, are they the parameters of the models being used for the prediction task? Ie DPOT, MPP…? (are they the same parameters as line 244 right column?)", "questions_for_authors": "- 1/ It is well known in the literature that adding physical losses in the training objective can lead to unstable or long optimization. How does behaves PITA with respect to this? Does it complicate training? \n- 2/ How behaves PITA when removing the data loss? Are L_con and L_Phy sufficient? \n- 3/ From what I understood, the sparse regression is computed on downsampled trajectories? \n- 4/ How are built $\\lambda^*$? What happens if the true PDE coefficients are taken from an incomplete library functions?", "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": "", "extra_scores": {}}, {"review_id": "sPtas69yz2", "reviewer": "Reviewer_y7j5", "summary": "The paper *\"Physics-informed Temporal Alignment for Auto-regressive PDE Foundation Models\"* introduces **Physics-Informed Temporal Alignment (PITA)**, a self-supervised learning framework designed to address the error accumulation problem in auto-regressive PDE foundation models. Instead of relying on predefined physics priors, PITA **discovers governing equations from observation data** and integrates them into the model’s training process through a combination of **data loss, physics loss, and consistency loss**. By aligning the physical dynamics at different time steps, PITA ensures that predictions remain consistent with underlying physical laws, reducing shortcut learning and improving generalization to out-of-distribution data. Extensive experiments across **12 datasets** show that PITA significantly enhances both short and long-term prediction accuracy, outperforming existing auto-regressive methods while maintaining scalability.", "questions": "", "limitations": "", "rating": 3, "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": "1.Regarding Long-Horizon Stability\n\nReply 1：To rigorously evaluate PITA’s capacity for modeling chaotic dynamical systems under extended temporal extrapolation, we conducted controlled experiments on synthetic 2D Kolmogorov turbulence flows, which is a canonical benchmark for chaotic PDE systems exhibiting multiscale energy cascades and nonlinear dissipation processes. Our turbulence simulation dataset contains 300 spatiotemporal trajectories with 64×64 grid resolution and 150 temporal steps, intentionally designed to challenge long-term forecasting fidelity. Both methodologies leverage the same pretrained DPOT-Ti model. As shown in https://anonymous.4open.science/r/PITA_1/Table4.pdf, our results demonstrate that PITA significantly outperforms the auto-regressive baseline in capturing the intricate patterns of chaotic fluid evolution. Specifically, PITA exhibits superior alignment with the underlying physical dynamics, showcasing its enhanced ability to model complex, long-term behaviors in turbulent systems. \n\n2.Regarding Essential References.\n\nReply2: Thanks for your valuable suggestion. We apologize for having cited an outdated version of OmniArch as:\nBuilding flexible machine learning models for scientific computing at scale.We will ensure that the recommended references are properly discussed in the main text of the revised manuscript. \n\n3.Regarding the computational cost. \n\nReply 3: We appreciate the reviewer’s concern regarding computational cost. The additional overhead from PDE discovery and temporal alignment primarily arises from three steps detailed in Section 3.3: building candidate libraries, sparse regression, and alternating direction optimization. As shown in https://anonymous.4open.science/r/PITA_1/Train.pdf, the time overhead of these steps remains constant and does not scale with model size.\nRegarding inference efficiency, since PDE discovery and alignment occur only during training, PITA's inference efficiency matches that of the baseline, while consistently outperforming it on long-term predictions (Figure 4, Appendix D.2).\nIn terms of memory usage, PITA incurs no additional GPU memory overhead compared to auto-regressive methods. It does require moderate extra CPU memory for building candidate libraries and sparse regression (https://anonymous.4open.science/r/PITA_1/CPU.pdf). Importantly, this CPU memory overhead is fixed.\nCollectively, these features render PITA fundamentally scalable: its computational overhead is decoupled from model dimensionality, effectively overcoming traditional memory limitations in physics-informed methods. This architecture-agnostic scalability supports seamless adaptation to billion-parameter models with bounded resource consumption. For instance, a 500M-parameter L-level model incurs only 0.517 additional seconds per batch (~20% increase), yet achieves a notable 31.61% performance improvement on long-trajectory prediction (https://anonymous.4open.science/r/PITA_1/Table1.pdf).\n\n4.Regarding the performance of PITA tested on underlying PDEs that exhibit (1) complex, high-dimensional, or non-sparse structures; and (2) high nonlinearity or structures that do not conform to a sparse representation.\n\nReply 4: The experimental validation of the PITA framework is conducted on a series of high-dimensional PDE datasets spanning diverse physical systems, including fluid dynamics, reaction-diffusion processes, and nonlinear wave phenomena, which are summarized as follows:\na.FNO-NS-ν and PDEArena Datasets\nDerived from the 2D incompressible Navier-Stokes equations, these datasets involve strongly coupled velocity-pressure fields and nonlinear advection terms, posing challenges in modeling multiscale interactions. The time-varying external force $f(x,t)$ in PDE-Arena-NS-Force further introduces non-stationary dynamics, requiring adaptive modeling capabilities. PITA achieves a 21.82% improvement in prediction accuracy compared to autoregressive baselines.\nb.PDEBench-DR\nThe reaction terms $R_u(u,v)$ and $R_v(u,v)$ induce chaotic dynamics and Turing patterns, reflecting the complexity of high-order nonlinear systems. PITA demonstrates a 15.48% performance gain in resolving multistable patterns and chaotic trajectories.\nc.Viscous Burgers Equation\nCharacterized by shock formation and nonlinear advection-dominated regimes, this dataset validates robustness in handling discontinuous solutions. PITA achieves a 19.41% error reduction in shock-capturing tasks compared to traditional solvers.\nd.As discussed in Reply 1 regarding the Kolmogorov Turbulence Flow, it is a canonical benchmark for chaotic PDE systems exhibiting multiscale energy cascades and nonlinear dissipation processes. As shown in https://anonymous.4open.science/r/PITA_1/Table4.pdf, our results demonstrate that PITA significantly outperforms the auto-regressive baseline in capturing the intricate patterns of chaotic fluid evolution.", "claims_and_evidence": "Yes. I do not see significant problems in this paper.", "methods_and_evaluation_criteria": "Overall, the proposed methodology and evaluation criteria are well-justified, relevant, and appropriate for assessing the challenges associated with auto-regressive PDE modeling.\n\nHowever,  the governing equation discovery process introduces an additional computational step after each auto-regressive prediction. While this ensures physics-informed corrections, it may significantly increase training time, particularly for large-scale PDE simulations. The paper does not quantify the computational cost trade-off between accuracy improvements and increased training time. There is a need for benchmarks on training efficiency, memory consumption, and scalability across different model sizes.", "theoretical_claims": "No theoretical claims.", "experimental_designs_or_analyses": "Generally, the experiments are comprehensive. \n\nHowever, I have additional two questions:\n\n1. Long-Horizon Stability – There is no explicit analysis of how PITA performs in very long rollouts (e.g., 100+ steps), particularly for chaotic PDEs, where error accumulation could still occur.\n\n2. Computational Overhead – The training time, inference speed, and memory usage of PITA are not reported, making it unclear how computationally expensive the PDE discovery process is compared to baseline models.", "relation_to_broader_scientific_literature": "[1] This approach aligns with the broader scientific literature on Physics-Informed Neural Networks (PINNs), which embed physical laws directly into neural network training to solve both forward and inverse problems involving PDEs. PINNs leverage governing equations to guide the learning process, ensuring that neural network solutions adhere to known physical principles. ​\n\n[2] The development of Neural Operators represents a significant advancement in this domain. Neural Operators are designed to learn mappings between infinite-dimensional function spaces, enabling the modeling of complex physical systems with improved efficiency and accuracy. They have been applied to various scientific and engineering disciplines, including turbulent flow modeling and computational mechanics, demonstrating their capability to handle complex PDE solutions across varying input conditions and geometries.​\n\n[3] The Theory of Functional Connections (TFC) has been applied in neural networks to enhance the performance of PINNs by effectively eliminating constraints from the optimization process. This integration significantly improves computational efficiency and accuracy, enabling the resolution of complex problems with greater ease.", "essential_references_not_discussed": "OmniArch: Building the Foundation Model for Scientific Computing\n\nAuthors: Tianyu Chen, Haoyi Zhou, Ying Li, et al.​\nVenue: arXiv preprint arXiv:2402.16014, 2024.​\nRelevance: OmniArch proposes a unified architecture to tackle multi-scale and multi-physics scientific computing challenges by integrating physical alignment into the modeling process. This approach aligns with PITA's goal of incorporating physics-based constraints to improve model predictions. ​\n\n\nPINNsFormer: A Transformer-Based Framework for Physics-Informed Neural Networks\n\nAuthors: Zhiyuan Zhao, Xueying Ding, B. Aditya Prakash​\nVenue: arXiv preprint arXiv:2307.11833, 2023.​\nRelevance: PINNsFormer introduces a framework that leverages transformer architectures to capture temporal dependencies in physics-informed neural networks. This methodology is pertinent to PITA's focus on aligning model predictions with temporal dynamics governed by PDEs. ​\n\n\nPIMRL: Physics-Informed Multi-Scale Recurrent Learning for Spatiotemporal Prediction\n\nAuthors: Wan Han, Qi Wang, Hao Sun​\nVenue: International Conference on Learning Representations (ICLR), 2025.​\nRelevance: PIMRL introduces a multi-scale learning framework that effectively leverages multi-scale data for spatiotemporal dynamics prediction, addressing challenges similar to those targeted by PITA. ​", "other_strengths_and_weaknesses": "Please see above reviews.", "other_comments_or_suggestions": "Please see above reviews.", "questions_for_authors": "Could you provide detailed metrics on the computational cost associated with the governing equation discovery process within PITA, including training time, memory usage, and scalability across different model sizes?\n\nHow does PITA perform when the underlying PDEs exhibit complex, high-dimensional, or non-sparse structures? Have you tested the framework on systems where the governing equations are highly nonlinear or do not conform to a sparse representation?", "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": "", "extra_scores": {}}, {"review_id": "CQwVP2Z9IL", "reviewer": "Reviewer_qJJU", "summary": "The authors studies autoregressive model which are used to make predictions about time series governed by a PDE system. To solve the shortcut problem for autoregressive models, the authors propose to use the training sequence to learn the governing equation, then train the autoregressive model so that they fit the training data and also the physical law learned.", "questions": "", "limitations": "", "rating": 3, "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": "1.Regarding the ablation studies on more datasets.\n\nReply 1: We conducted additional experiments on three datasets that span different PDE domains and exhibit varied physical characteristics, with an average step length of 84. The results (https://anonymous.4open.science/r/PITA_1/Table2.pdf) consistently corroborate our main findings, further reinforcing the reliability of our conclusions.\n\n2.Regarding the error bounds.\n\nReply 2: Thanks for your insightful question. PITA is fundamentally a data-driven PDE-solving approach built upon auto-regressive baselines. To our knowledge, theoretical analysis of error bounds in data-driven PDE methods remains an open and challenging problem, which is currently beyond the scope of our work. Notably, baseline methods such as DPOT, FNO, and MPP also do not provide discussions on error bounds. Nonetheless, we fully agree with the reviewer’s point and will carefully consider error analysis in future.\n\n3.Regarding the PDE learning component and the robustness of PITA.\nReply 3: Even if equation discovery fails, our multi-task learning automatically down-weights the physics loss, causing the model to degrade into a dual-constraint mode governed by data loss and consistency loss. Here, the consistency loss specifically aligns macroscopic data patterns derived from the discovered dynamics. As shown in Table 3 of our submission, even without physics loss, PITA still surpasses the baseline by 43.02%. For longer trajectories, the consistency loss is also critical, as demonstrated in https://anonymous.4open.science/r/PITA_1/Table2.pdf.\n\nAdditionally, we conducted a stress test by directly injecting noise into the physics-loss gradients, simulating severely incorrect physics guidance. Results in https://anonymous.4open.science/r/PITA_1/Table6.pdf and Table7.pdf show the robustness, and Table9.pdf indicates that PITA regrades as baseline performance under such attacks. \n\nFinally, while both PITA and PINN-based methods leverage physics constraints, PINNs require explicitly formulated governing equations (e.g., PINN-SR[1]). In contrast, PITA discovers the underlying PDE directly from data without physics priors, making PDE discovery not merely complementary but essential for learning unknown dynamics.\n\n4.Regarding the evaluation on more complex dataset.\n\nReply 4: We conducted experiments on synthetic 2D Kolmogorov turbulence flows—a more complex dataset with 150 temporal steps. Using DPOT-Ti model, our results (https://anonymous.4open.science/r/PITA_1/Table4.pdf) demonstrate that PITA substantially outperforms the baseline. Due to space constraints, please refer to the discussion in Reply 4 to Reviewer y7j5.\n\n5.Regarding the additional computational cost.\n\nReply 5: Thanks for your valuable suggestions. As shown in https://anonymous.4open.science/r/PITA_1/Train.pdf , the time consumption of the PDE discovery and alignment steps is constant and does not scale with the model size. Moreover, PITA performs PDE discovery and alignment only during the training phase, so its inference efficiency is identical to the baseline. Due to space limitations, please refer to the discussion in Reply 3 to Reviewer y7j5.\n\n6.Regarding the shortcut problem and long-term trajectory accuracy.\n\nReply 6: As discussed in Reply 4, PITA exhibits significant advantages in long-trajectory data. Following the accumulative error visualization in [1], we plotted the rolling-step MSE. Results (https://anonymous.4open.science/r/PITA_1/, see 'Shortcut' folder) show that the error remains stable and comparable to earlier time steps. Due to space constraints, please refer to the discussion in Reply 1 to Reviewer Xe1t.\n\n7.Regarding other types of time series observations.\n\nReply 7: We agree that the PITA framework can indeed be extended and adapted to other types of time series data. Your feedback has been incredibly insightful, and we are committed to incorporating this direction into future work.\n\n8.Regarding the design of the PDE discovery term.\n\nReply 8: We adopted the PDE-FIND library, which is widely utilized by PDE discovery methods such as SINDy, as a general and standard library in the field without requiring specialized design. As shown in https://anonymous.4open.science/r/PITA_1/Table3.pdf, experiments demonstrate that PITA remains effective even when more than half of the candidate terms in the library are removed randomly. This indicates PITA does not depend on carefully tailored terms.\n\n9.Regarding the proof of loss term being able to reduce the shortcut problem.\n\nReply 9: As discussed in Reply 3, 4, and 6, results provide strong evidence of PITA's promising potential for shortcut learning.\nWe attribute this to two key components: (1) the physics loss, enforcing fine-grained alignment of physical dynamics, and (2) the consistency loss, aligning macroscopic data patterns over temporal windows, jointly mitigating error accumulation.\n\n[1] Physics-informed learning of governing equations from scarce data.", "claims_and_evidence": "All the results should include error bounds -- otherwise it is difficult to see whether the improvement is due to random chance or due to the actual methods.\n\nThe ablation studies seems to be okay, since it directly investigates the performance of each introduced components, however could also be tried for other tested dataset as well to make it more convincing.\n\nI am also curious how important the PDE learning component is, and how robust PITA is to the component. Specifically, if the PDE regression component is not sufficient to model the true equation, how much would the learned dynamics suffer? Furthermore, it may be that not learning the PDE in full but rather providing a better PDE regularisation may also be enough (this is shown in Table 3 Task 1, where Lphy + Ldata does perform worse than PITA but not that much worse, and could potentially be improved to bypass directly learning the PDE in the first place too).", "methods_and_evaluation_criteria": "The benchmarks seem to be reasonable for physics-informed problems, however I'm not sure if more complex datasets (e.g. more high-dimensional or noisier real-life data) should also be used beyond those that are already typically studied in simpler PINNs and deep operator models. I do think that the fluid dynamics data might be trending towards these scenarios though and may be okay if these are also the norms used in existing benchmarks.\n\nI notice that the authors have mentioned that the method requires more computation than other benchmarks. While the authors have given some numbers to this in the limitations section, they could be reported further as well for a clearer picture of the additional computational requirements for the methods.", "theoretical_claims": "There are none in the paper.", "experimental_designs_or_analyses": "In terms of the shortcut problem that is mentioned by the authors -- while the results of PITA seems to be positive, it is unclear whether this is due to its success in solving the shortcut problem, or whether the loss term is able to simply better capture the physics bias in general. I'm not sure if the fact that long term trajectory accuracy is sufficient for this (there may be other reasons long-term trajectories perform poorer in other models).", "relation_to_broader_scientific_literature": "This work seems to be relevant to learning physical systems with data and making predictions based on them. They therefore seem to be quite relevant to scientific applications where those types of PDEs are studied. Due to the equation discovery portion they could also be used for other types of time series observations as well, however this does not seem to be focused on in the current paper.", "essential_references_not_discussed": "None as far as I am aware of.", "other_strengths_and_weaknesses": "No other points beyond what was addressed above.", "other_comments_or_suggestions": "None.", "questions_for_authors": "Most of the questions have been raised from above but will mention some of them here again:\n\n1. Regarding the importance of the PDE solver in this case -- how reliant is the method on a adequately good design of the PDE discovery term? \n\n2. Are the authors able to provide stronger proof that the loss term is able to reduce the shortcut problem, rather than just merely reducing the prediction error and attributing this to being due to the shortcut problem being solved?", "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": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: The paper introduces a self-supervised framework to mitigate error accumulation in auto-regressive PDE surrogates. The core contribution lies in integrating losses derived from governing equations with the standard data loss used for training surrogates. These governing equations are discovered from both the data and the model predictions using symbolic regression. Two loss terms are introduced: a physical loss corresponding to the inferred equation, and a so-called consistency loss that measures the similarity between the equation inferred from historical data and the one inferred from the model’s predictions. This framework is compatible with existing models, and experiments are conducted on various foundation model architectures. Results on 12 datasets demonstrate the positive impact of this approach on reducing error accumulation.

Reviewers requested additional experiments and clarifications regarding the added complexity of the method relative to the backbone models it builds upon, as well as its robustness to observational noise and to errors or incompleteness in the discovered equations. The authors provided detailed responses and supplementary experiments addressing these concerns. The reviewers’ feedback on the responses was generally positive: two out of four increased their scores, and all expressed an inclination toward acceptance.
I recommend that the authors carefully incorporate the reviewers’ suggestions and revise the manuscript, particularly to improve clarity in some imprecise sections. I recommend acceptance.
Decision: Accept (poster)