Active Learning with Selective Time-Step Acquisition for PDEs
Authors: Yegon Kim
Hyunsu Kim
Gyeonghoon Ko
Juho Lee
Venue: ICML 2025 poster (2025)
Keywords: Active learning
Partial Differential Equation (PDE)
Primary area: general_machine_learning->online_learning_active_learning_and_bandits
TL;DR: We propose an active learning framework that reduces the data acquisition cost for PDE surrogate modeling by selectively querying time steps along a trajectory from numerical solvers.

Abstract: Accurately solving partial differential equations (PDEs) is critical to understanding complex scientific and engineering phenomena, yet traditional numerical solvers are computationally expensive. Surrogate models offer a more efficient alternative, but their development is hindered by the cost of generating sufficient training data from numerical solvers. In this paper, we present a novel framework for active learning (AL) in PDE surrogate modeling that reduces this cost. Unlike the existing AL methods for PDEs that always acquire entire PDE trajectories, our approach strategically generates only the most important time steps with the numerical solver, while employing the surrogate model to approximate the remaining steps. This dramatically reduces the cost incurred by each trajectory and thus allows the active learning algorithm to try out a more diverse set of trajectories given the same budget. To accommodate this novel framework, we develop an acquisition function that estimates the utility of a set of time steps by approximating its resulting variance reduction. We demonstrate the effectiveness of our method on several benchmark PDEs, including the Burgers' equation, Korteweg–De Vries equation, Kuramoto–Sivashinsky equation, the incompressible Navier-Stokes equation, and the compressible Navier-Stokes equation.
Experiments show that our approach improves performance by large margins over the best existing method. Our method not only reduces average error but also the 99\%, 95\%, and 50\% quantiles of error, which is rare for an AL algorithm. All in all, our approach offers a data-efficient solution to surrogate modeling for PDEs.

Active Learning with Selective Time-Step Acquisition for PDEs
Yegon Kim Hyunsu Kim
Gyeonghoon Ko
Juho Lee'
Abstract
1. Introduction
Accurately solving partial differential equations (PDEs) is critical understanding complex sci- entific and engineering phenomena, yet tradi- tional numerical solvers are computationally ex- pensive. Surrogate models offer a more effi- cient alternative, but their development is hin- dered by the cost of generating sufficient train- ing data from numerical solvers. this paper, we present a novel framework for active learn- ing (AL) in PDE surrogate modeling that reduces this cost. Unlike the existing AL. methods for PDEs that always acquire entire PDE trajecto- ries, our approach strategically generates only the most important time steps with the numeri- cal solver, while employing the surrogate model to approximate the remaining steps. This dra- matically reduces the cost incurred by each tra- jectory and thus allows the active learning algo- rithm to try out a more diverse set of trajecto- ries given the same budget. To accommodate this novel framework, we develop an acquisi- tion function that estimates the utility of a set of time steps by approximating its resulting vari- ance reduction. We demonstrate the effective- ness of our method on several benchmark PDEs, including the Burgers' equation, Korteweg- De Vries equation, Kuramoto-SSvaahinsk. equation, the incompressible Navier-Stokes equation, and the compressible Navier-Stokes equation. Exper- iments show that our approach improves perfor- the best existing method. Our method not only reduces average error but also the 99%, 95%, and 50% quantiles of error, which is rare for an AL algorithm. All in all, our approach offers a data-efficien solu- tion to surrogate modeling for PDEs.
In many scientific and engineering applications, accurately solving partial differential equations (PDEs) in the form of trajectories of states evolving over time is essential for understanding complex phenomena (Holton and Hakim, 2013; Atkins et al., 2023; Murray, 2007; Wilmott et al., 1995). The traditional approach involves running nu- merical solvers, which provide accurate solutions but are computationally costly, taking several hours, days or even weeks to run depending on the complexity of the problem (Cleaver et al., 2016; Cowan al., 2001). As result, there is significant interest in developing surrogate models (Greydanus et al., 2019; Bar-Sinai et al., 2019; Sanchez- Gonzalez et al., 2020; Z. Li et al., 2020; Brandstetter et al., 2022b; Lippe et al., 2024) that can approximate the solutions more efficiently. Surrogate models are obtained by solving regression tasks on some "ground truth" data. The ground truth data for PDEs are generated by numeri- cal solvers, which costly compared to those of standard regression problems. As result, the expense of data quisition presents a major bottleneck in the development of surrogate models for PDEs.
Active Learning (AL, Chernoff, 1959; MacKay, 1992; Set- tles, 2009) can address this challenge by adaptively ac- quiring the most informative inputs, effectively reducing the amount of ground-truth data required to obtain a high- quality surrogate model. However, there is a general lack of research in AL for regression tasks (Dongrui Wu, 2018; Holzmüller et al., 2023), let alone PDEs. Existing studies on AL for PDEs have predominantly dealt with univariate outputs such as energy (Pestourie et al., 2020; Pickering et al., 2022), or predictions at single, fixed time point (Bajracharya et al., 2024; Dongxia Wu et al., 2023). To our surprise, the only work directly addressing AL for pre- diction of trajectories is that of Musekamp et al. (2024). In this work, the surrogate model is set as an autoregres- sive model that predicts the evolved state of a PDE at time + Al given a state at an arbitrary time point 4, and trained on data acquired by existing regression-basee AL methods (Holzmüller al., 2023). Specifically, at each round of acquisition, the AL method chooses initial con- ditions from which entire trajectories are acquired. How- ever, we argue that querying all the states in a trajectory is not sample-efficient, especially for autoregressive surro-
mance by large margins
Korea Advanced Institute of Science and Technol- ogy, Dacjeon, Korea. Correspondence to: Yegon Kim <ycgonkim@ke @ kwist.ac.kr....... kc>, Juho Lee <juholee@ krr>
Proceedings of the 1nnd International Conference Machine Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Active Learning with Selective Time-Step Acquisition for PDEs
ramoto-Sivsshinsy equation, the incompressible Navier- Stokes equation and the compressible Navier-Stokes equa- tion. The compressible Navier-Stokes equation is par- ticularly challenging task due to its high nonlinearity and turbulent behavior. Additionally, we analyze the behavior of our AL method, providing insights into the factors that contribute to its effectiveness. Our code publicly avail- able at https: / / githittubb www./ooooooeeoooooooooooooooeooe stap.
catar Voa
2. Background
Figure 1: PCA of surrogate model hidden layer's activa- tion patterns on states of the incompressible Navier-Stokes equation. The left figure highlights states within 10 trajec- tories, and the right figure highlights the same number of states chosen randomly.
2.1. Preliminaries
We consider PDEs with one time dimension t E [0, and possibly multiple spatial dimensions a [2,,,,,... E X where X is the spatial domain such as the unit interval. These can be written in the form
gate models.
Acquiring entire trajectories is inefficient mainly for two reasons. First, states within a trajectory are often strongly correlated, undermining their diversity or the joint informa- tion gain (Houlsby et al., 2011; Kirsch et al., 2019). Fig. shows a side-by-side comparison of PDE states selected within a few trajectories and the same number of states se- lected randomly from many trajectories. States within tra- jectories exhibit clustering and hence lower diversity com- pared to the randomly selected ones. Further analysis and details are given in Appendix D.2. Secondly, even if the states are not strongly correlated, it can be the case that only certain time steps of a trajectory are the most informative due to the dynamics of the PDE. In both cases, noting that the main cost is running the numerical solver, it would be ideal selectively acquire only the most important time steps with the numerical solver, for fraction of the cost of acquiring the entire trajectory. However, this is usually impossible without querying all the time steps that come earlier.
dau F(t, x, u, U1 Uwwu....)
(1)
where u. : [0,T] + Rd is a solution the PDE. We are also given a specific boundary condition and fixed time interval At. If the PDE is well-posed (W. Evans, 1988), there exists, for each to E R, an evolution operator Gt, which maps an initial condition u == u(4o,' .) to the solu- tion u == u(to + At, .). For simplicity, we only consider time-independent PDEs, for which the evolution operator G, is the same for all t, say G. Iterating over G multi- ple times, we can obtain trajectory (u²); where u' := G+)) [ut] with GG) being the i-th iterate of G. In practice, G is implemented by various numerical meth- ods, adequately chosen for the given PDE, as elaborated in Appendix B.1.
L. ;=1 of length L,
We train a neural surrogate model & with input-output pairs (u, G[ul) from the numerical solver G. Active learning aims to build a high quality training dataset by adaptively selecting informative inputs to be fed into the solver G. Prior work (Musekamp et al., 2024) operates framework where initial conditions w" are selected from pool P. from which full trajectories of length L are ob- tained. For instance, Query-by-Commitee (QbC, Seung et al., 1992) queries initial conditions u that maximize the predictive uncertainty estimated from a committee of M models.
In this paper, we propose a novel framework for data ac- quisition that circumvents the constraint of having to query all time steps in a trajectory, along with an AL strategy that leverages this flexibility. Our method combines both nu- merical solver and a surrogate model to acquire data along a trajectory with reduced cost. Specifically, it selects which time steps along a trajectory to query to the solver, while using the surrogate model to approximate the remaining steps. To do so, we develop a novel acquisition function that guides our AL strategy in choosing which time steps to query to the numerical solver in each trajectory.
the the
1 M
aquc(u²)
Overall, our framework, equipped with the novel AL strat- egy. significantly improves surrogate model performance over previous methods. We validate our approach through extensive experiments on benchmarks, including the Burg- ers' equation, the Korteweg-DD De Vries equation, the Ku-
m=
i=1
where Ui is the prediction of the 7th state from the th surrogate model in the committee and ä := 1 M in is the mean prediction from the committee.
M
2

---

Active Learning with Selective Time-Step Acquisition for PDEs
Algorithm 1 Overview of STAP
Initial Conditions
Training Data
Require: Pool P of initial conditions, budget B per round, num- ber of rounds R. numerical solver G, trajectory length L, ini- tial training dataset D
Select (e.g. QbC)
Add data from Solver
Retraän
Ensure: Trained surrogate model &
1: Train C on D
Solocted
Selective T'mme-Step Acquisition
2:
for round = to I do
3: cost
FNO (Burrgate Model
4: while cost B do
5:
Choose initial condition u from P
> P 3.3
de Alaauade
6:
p. + P {u
PNO
Numerical Solver
7:
Choose 5 (81, 8L)
$$ 3.2 and 3.3
8:
á u
9:
for = to I. do
Figure 2: Illustrated overview of STAP. This illustration
G[a
if true
describes one round of AI.
10:
C(a'T1]
if = false
11:
if = true then
12:
D DU {(ū 1 ú²)}
2.2. Problem Setting
13:
cost cost +
Our goal is obtain surrogate model & that approximates the expensive numerical solver G with low error
14:
end
15:
end for
16: end while
17: Train C on D
E Bulrp(u)) err (u))) (G(0) (G'0(tt))) DE
(3)
18:
end for
where crr(:.)) is an error metric. Obtaining the surrogate model requires sampling training data from the numerical solver, which incurs nontrivial cost. AL aims to im- prove sample efficiency by sampling only the most impor- tant data. In particular, AL utilizes the current surrogate model G, or a committee of surrogate models [Gmm] M m=1 to inform its choice. After acquiring the data chosen by AL, the surrogate G is retrained with the new dataset.
dition u from the pool P, similar to the existing AL meth- ods for PDE trajectories. However, while existing methods acquire the entire trajectory starting from the chosen ini- tial condition u (Musekamp et al., 2024), our method quires partial trajectory. Specifically, we select subset of time steps to simulate from u", rather than acquiring the full trajectory. The rationale behind this approach is that, given a fixed budget, acquiring as many trajectories as pos- sible-albeit partially from different initial conditions often more beneficial than fully acquiring fewer trajecto- ries. This strategy enables more efficient exploration of the data space and improves the overall sample efficiency of the framework.
We assume that there exists pool P of initial conditions nu o At each round of AL, we train committee of M surro- gate models {Gm m==1 with the training dataset collected from G up to that round of AL. We then use this commit- tee to select a batch of inputs to be queried to the solver G, and add the pairs of queried inputs and resulting outputs to the training dataset. The cost at each round, defined as the number of inputs queried G, is limited to a certain bud- get B. We aim achieve low errors at each round, AL strategy would ideally acquire data with cost as close to or equal to the budget (S. Li et al., 2022a).
More specifically, for a given initial condition u.,, we de- fine a boolean sequence of length L, S = (s1...., st), which we refer to as the sampling pattern. For example, S could be (true, false, true). The sampling pattern specifies that data will be acquired only time steps corre- sponding to true values while skipping those marked false.
3. Selective Time-Step Acquisition for PDEs
After selecting the sampling pattern S, the next step is to acquire the PDE trajectory. While acquiring a full trajec- tory is straightforward using a numerical solver G, obtain- ing a partial trajectory corresponding
3.1. Framework of Data Acquisition
We present our method, Selective Time-Step Acquisition for PDEs (STAP), which operates under framework of data acquisition that is much more sample efficient than previous works. Algorithm provides an overview of our framework. Fig. 2 also provides an illustrated version of the overview. We start with a surrogate model G or a commit- tee of surrogate models (Gm 1M 1m=1, trained with the initial dataset D. At every round of AL, we choose an initial con-
S can be tricky. We want run the solver G only for the time steps speci- fied by S (those with true patterns), but the solver requires the skipped time steps (those with false patterns) as inter- mediate inputs. If we give up and just run the solver for all time steps for this reason, we wouldn't be saving any cost. To address this, we use a simple heuristic: for the skipped time steps, we replace the simulation with predic-
3

---

Active Learning with Selective Time-Step Acquisition for PDEs
tions from the surrogate model (we use the average surro- gate G 1 M
variance reduction between all the distinct pairs in the com- mittee:
M Im=1 Gm when we have committee). That
is, starting with u" = u., we iterate over 1 1 L:
a(u", S) M(M - 1) a.BE .M) Mazttb
Σ
R(a, S). (6)
G] if 8; true
(4)
G(
if s; = false.
We add to our dataset D only the pairs obtained with the solver G, namely (a'11,0 ú') with 8, true.
We observe that our acquisition function simplifies to QbC in Eq. when S acquires all the time steps, differing only by a constant factor of two. This occurs because, in that case, us,S,a = u.. which makes the second term in the summand of Eq. 5 vanish. Consequently, we can interpret our acquisition function as a generalization of QbC that ac- commodates for the selection of time steps.
In comparison to full trajectory acquisition, which requires L executions of the numerical solver, our strategy invokes the numerical solver ||||| := Σ [[s; = true] times and utilizes the surrogate model I. ||S|| times. Since the sur- rogate model is significantly cheaper to run than the nu- merical solver, this approach substantially reduces the cost of acquisition, enabling us to explore more initial condi- tions within the same budget. In fact, as discussed in $ 2.2, we define the acquisition cost precisely as ||5||| We repeat expanding our training dataset with new initial conditions and sampling patterns until the cost incurred in the current round reaches budget B. At the end of each round, we retrain the surrogate G with the expanded training set D.
As an additional sanity check, consider the scenario where S does not sample any time steps. In this situation, Us,S.a = ús. leading the two terms in the summand to cancel each other out, resulting in zero variance reduction. Since acquiring no data should yield zero utility, we con- firm that our acquisition function behaves as expected in this limiting case. Appendix A.2 further details the precise motivation behind the design of our acquisition function.
Previous methods listed in Musekamp et al. (2024) can be considered a special case of ours where the sampling pat- tern S is full of true entries. Our framework is therefore a strict generalization of previous works. In the remainder of this section, we describe how STAP adaptively chooses initial conditions u and sampling patterns S.
3.3. Batch Acquisition Algorithm
Since retraining on new dataset with every new data point is costly, and acquiring data in a parallel manner can be more efficient than when done in sequential manner, tive learning algorithms should accommodate batch acqui- sition, that is, the acquisition of multiple data points at each round of active learning (Kirsch et al., 2019). There are two main challenges in designing a batch acquisition AL methods. First, simply maximizing the acquisition values of individual data points doesn't maximize the utility of the batch. Numerous works report that picking instances that maximize individual acquisition values can severely under- perform compared to methods that take into account the interactions between those instances (Kirsch et al., 2019; Ash et al., 2019). The problem is chiefly attributed to the lack of diversity and representativeness (Dongrui Wu, 2018) caused by oversampling of small, high value regions (Smith et al., 2023). Secondly, the AL method needs to search over large combinatorial space of size ((IP) ()). where [P] and B are the pool size and batch size, re- spectively. The search method should therefore be scal- able. Most existing batch AI. algorithms address both is- sues (Holzmüller et al., 2023).
3.2. Acquisiiion Function
To adaptively select the sampling pattern S with the ini- tial condition u., we propose a novel acquisition func- tion a(u", S) that assesses the utility of S. Given com- mittee {Gm))M consider (G., Gi)) for some * b E (I....,M) MM. We define the utility of the sampling pattern S for the pair (Ga.G)) as the resulting variance reduc- sion in the pair's rolled-out trajectories. Specifically, let ú. and ún be the trajectories estimated by G. and GG start- ing from u. Next, let án be the trajectory rolled-out using Eq. 4 but with G and G replaced by G, and G.. re- spectively. In other words, Gi is used at time steps where the sampling pattern indicates false and G, otherwise. This trajectory is designed approximate GG that's updated it- self with data from G. at time steps for which the sampling pattern indicates true. The variance reduction is defined as
I.
With the acquisition function defined in $ 3.2, we present batch acquisition algorithm given pool P of initial con- ditions. We define the cost of batch {uu.. 5 as the total number of queries to the solver, ビØ= ||S If we were to maximize the sum of individual acquisition val- ues N a(u". Sy)) under a budget constraint E, || S B, there is a known approximate solution Salkin and De
R(a, b. 5) : Σ (1 -i]l==
ll..sall?). (5)
The sampling pattern S that maximizes R(a, b, S) is the one where the current models G. and Go disagree the most, and acquiring data from S effectively reduces this discrep- ancy. Our acquisition function is defined as the average
4

---

Active Learning with Selective Time-Step Acquisition for PDEs
Kluyver (1975) that greedity maximizes the cost-weighted acquisition function a" (u", total cost exceeds the budget B. However, as discussed above, it's questionable whether the sum of individual ac- quisition values is actually a good representative for the utility of a batch. Moreover, even if we employ the ap- proximate greedy solution, we are still searching over the product pool of the sampling pattern S and the pool P of initial conditions, whose size is on the order of (P))))). Both terms impose significant computational burden on op- timizing the cost-weighted objective a" a'nnn, SS)
and Dongxia Wu al. (2023) apply their AL methods to single-state prediction. Bajracharya et al. (2024) explores the use of active learning in tasks of predicting steady states of PDEs, which can be seen as predicting single states at + cc. Finally, Musekamp et al. (2024) experiments with active learning in predicting PDE trajectories with autore- gressive models.
SS/\SS) until the
Active selection of time points. While our work is the first to propose time step selection in active learning (AL) for PDEs, the concept of selecting time points has been ex- plored in other contexts. For example, in physics-informed neural networks (PINNs), active selection of collocation points for training has been widely studied (Arthurs and King, 2021; Gao and Wang, 2023; Mao and Meng, 2023; C. Wu al., 2023; Turinici, 2024). "Labels" for PINNs, or the residual loss, can be calculated directly at any time point using closed form equations. There are also methods in Bayesian experimental design (BED) that choose obser- vation times that maximize information gain about param- eters of interest (Singh et al., 2005; Cook et al., 2008). In those works, a trajectory is already "there", but the cost attributed the act of observing a time point. In contrast, in our setting, we cannot directly acquire time point, be- cause there is cost in the simulation of the trajectory.
We therefore propose STAP as an add-on to existing AL methods that acquire full trajectories. By using STAP as an add-on, the diversity and representativennss promoted by existing AL methods (Holzmüller al., 2023; Kirsch et al., 2023; Musekamp al., 2024) are upheld, and the size of the optimization space for STAP is reduced to O(2²). Specifically, a full-trajectory AL method A, which we call a base method, first selects an initial condition u". Musekamp et al. (2024) introduces several possibilities for such a method, including QbC (Seung et al., 1992), Largest Cluster Maximum Distance (LCMD, Holzmüller et al., 2023), Core-Set (Sener and Savarese, 2017), and stochas- tic batch active learning (SBAL, Kirsch et al., 2023). We then optimize the cost- weighte acquisition function a" c"(u,, S) over the sampling pattern S while holding u." fixed, and add the pair (u", SS to the current batch. We iter- ate this two-stage process until the cost of the batch reaches our budget B. Additionally, if the cost ever exceeds the budget after adding a pair, we truncate the sampling pat- tern so that the cost is exactly equal to the budget.
Multi-fidelity AL. Closely related to our work is multi- fidelity active learning (S. Li et al., 2022b; Dongxia Wu et al., 2023; Hernandez-Garte where outputs are acquired varying fidelities for each in- put, with associated costs inherent to each fidelity. In our context, the task of actively selecting sampling pattern for given initial condition can be seen as fidelity selection problem, where acquiring all time steps corresponds to the highest fidelity but also incurs the highest cost.
et al., 2023; S. Li et al., 2024),
We use a greedy algorithm to optimize the wost-weighted acquisition function a" (u², over S in the combinato- rial space of size O.2 In the greedy algorithm, we start by initializing S with all entries set to true. At each step of the greedy algorithm. we propose a neighboring pattern S' by applying a bit-flip mutation, where each bit of S is flipped with a probability of e. The proposal is accepted only if the acquisition value a" (u", 8') is higher than the current value a" a""", S). This process of proposal and ac- ceptance/rejection is repeated T times. We use T = 100 and e = 0.1 throughout our experiments. A more con- cise summary of the batch acquisition algorithm is given in Appendix A.3. The algorithmic complexity of batch acqui- sition is discussed in 5.6.
5. Experiments
5.1. Baseline AL Methods
To compare with our method, we experiment with AL for full trajectory sampling introduced in Musekamp al. (2024). Random sampling from the pool set is the simplest method. QbC (Seung et al., 1992) is a simple active learn- ing algorithm that selects points according to maximum disagreement among members of a committee. LCMD (Holzmüller et al., 2023) an AL algorithm that uses a feature map. We concatenate the last hidden layer activa- tions of committee members at all time steps of trajec- tory, and sketch the concatenated features to dimension of 512 using random projection. Kirsch et al. (2023) pro- poses SBAL, which randomly samples data points x with a probability distribution proportional to its temperature- scaled acquisition value p(æ) x a(x)m.. We use the ac- quisition function of QbC with temperature m L. We
4. Related Work
AL for PDEs. The works by Pestouric et al. (2020), Picker- ing et al. (2022), and Gajjar al. (2022) apply active learn- ing to problems involving PDEs, but their tasks are limited to predicting a particular quantity of interest, such the maximum value of an evolved state. S. Li et al. (2024)
5

---

Active Learning with Selective Time-Step Acquisition for PDEs
Burgers
KdV
KS
INS
CNS
-2.0
0.25
0.50
2.8
-2.5
0.25
0.00
-0.2
-3.0
0.00
-0255
-0.4
2.4
-355
0.25
0.50
-4.0
-4.5
-0.50
-0.75
-0.8
2.0
5 Round
10
10
5
10
0
5 Round
10
0
5
10
Round
Round
Round
Random
QbC
SBAL
LCMD
SBAL+STAP
Figure 3: Log RMSE of AL strategies, measured across 10 rounds of acquisition. Each round incurs constant cost of data acquisition, namely the budget B.
leave out Core-Set (Sener and Savarese, 2017) because it generally underperforms compared the above methods, according both Holzmüller et al. (2023) and Musekamp et al. (2024).
Table 1: Log RMSE of AL strategies, averaged across 10 rounds of acquisition
Burgers
KdV
KS
INS
CNS
Random
-2.881-0000
0. an
-0.258 c0
-0.422.0000
2.603emms
Q&C
-3.121a0065
0.266a0:221
-0.268
-0.385p0.04
2.844
5.2. Target PDEs
LCMD SBAL
-2.847unee7
0.256xxce
0.046 21013
-0.320uaon
2.736uni
-3.388uaes
0.030wacss
-0.275
-0461aa0
2.422a0ms
We evaluate our method on range of PDEs. The first is the Burgers' equation, which shows either diffusive behavior or shock phenomena depending on the viscos- ity parameter. Next, we test the nonlinear Korteweg- De Vries (KdV) equation, which is known for exhibiting soli- tary wave pulses with weak interactions (Zabusky and Kruskal, 1965). We then apply our method to the Ku- ramoto- Sivassi-iv (KS) equation, another nonlinear PDE in one dimension, notable for its chaotic dynamics. Lastly, we examine two forms of the Navier-Stokes equation in two spatial dimensions: the incompressible Navier-Stoker equation (INS) in vorticity form and the compressible Sovier-Stokes equation (CNS). Both present significant challenges due to their strong nonlinearity and turbulent behavior. All equations are solved with periodic boundary conditions. Our choice of PDEs and their parameters is de- signed to validate our research across various flows exhibit- ing complex behaviors, including shock waves, chaotic dynamics, and turbulence. Additional details are in Ap- pendix B.
SHALASTAP
-3.674.0
-0.088-0.cp
-0.349 10
-0.525mn
2.36300
FNO is simply trained on ground truth input-output pairs from the solver G without backpropagating through two or more time steps. All models were trained with Adam (Kingma, 2014) for 100 epochs, using a learning rate of 10-3, a batch size of 32, and a cosine annealing scheduler (Loshchilov and Hutter, 2016).
5.4. Results
We compare between the four baselines introduced in $ 5,1, and our method combined with SBAL (SBAL+STAP). The pool set has 10,000 initial conditions, and we always start with an initial dataset of 32 fully sampled trajectories. The initial conditions in the test set are sampled from the same distribution as those in the pool set. An ensemble size of M = is used, as it has been shown be sufficient for good AL performance (Pickering al., 2022; Musekamp et al., 2024). We perform 10 rounds of acquisition, and the budget of each round is set to B = 8 x L where L is the length of a full trajectory. This means that at every round of AL, full trajectory algorithms sample 8 trajectories, each of length L. All experiments were conducted on 8 NVIDIA GeForee RTX 2080 Ti GPUs, and the results are averages from 5 seed values. Appendix C provides full report of all metrics on all methods, while here we summarize the most interesting results.
5.3. Surrogate Models
We use Fourier Neural Operators (FNO, Z. Li et al., 2020) for our surrogate model G. In particular, we train it to pre- dict the differences between states in adjacent time steps, following Musekamp al. (2024). All models have four hidden layers with 64 channels in each layer. We use 32, 256, 128, 16, and 32 Fourier modes for Burgers, KdV, KS, INS and CNS equations, respectively. We also normalize the data according to the initial dataset's 's mean and standard deviation over all temporal and spatial dimensions. The
Fig. 3 shows plots of the committee's 's logarithmic RMSE across the 10 rounds of acquisition, and Table summa- rizes the results with mean logarithmic RMSEs, where the
6

---

Active Learning with Selective Time-Step Acquisition for PDES
Mathod
Methad
-3.0
Random
Random
ObC
QbC
SOAL
SBAL
SEAL
SBAL+STAP
Rardom
-4.0
-2
SBAL
-3
-3
-4.5
SOAL+STAP
2
6
10
0
2
6
10
2
6
10
Round
Round
Round
(a) 99% quantile
(b) 95% quantile
(c) 50% quantile
Figure Quantiles of log RMSE on Burgers measured across rounds of acquisition.
12
(a) Burgers
(b) KdV
(c) KS
(d) INS
(c) CNS
Figure Timesteps chosen by SBAL+STAP. Each row corresponds an acquired trajectory, where the black cells indicate the selected time steps. We show twenty trajectories acquired in the first rounds of active learning.
mean is taken over all 10 rounds. We can observe from the plots that SBAL+STAP outperforms other AL baselines in a robust manner. Most notably, it improves performance on the KS equation, where no other baseline improves signifi- cantly over Random selection.
Table Log RMSE with Bernoulli sampling averaged across rounds of acquisition.
Burgers
KdV
KS
INS
CNS
333888omm
0.030-omm
00.255mo01
2.422-non
+STAP
-36674omm
-0,088
-0,349
-0.525:o.mc
2.363-0011
Following Holzmüller al. (2023), we also report the 99%, 95%, and 50% quantiles of log RMSE. Quantiles useful for analyzing the behavior of AI. strategies. AL. methods tend to improve performance on points with extreme er- rors, or in other words the top quantiles, while not so much in the middle quantiles. This is why AL methods perform differently depending on the algorithm and the nature of the task. Fig. shows their respective plots for Burgers. All baseline methods improve performance over random sampling in the 99% and 95% quantiles, and this is true even for LCMD which doesn't improve over random sam- pling average error. the 50% quantile, none of the baseline methods reduce the error below random sampling. but STAP surprisingly outperforms random sampling. This trend holds for all equations, detailed in Appendix We therefore infer that STAP isn't simply sacrificing the surrogate model's performance in some trajectories to improve its performance in others.
+Ber(1/16)
3.223
0.053.a014
-0.365 arce
40.529.om
2.375
4Ber(1/8)
-3. 3.15220m
0.049-0014
-0.359-a
-0.525+0
2.370-0011
+Ber(1/4)
3. 3.102aaam
0.018
-0.346
40.515uomm
2.372.a013
+Ber(1/2)
-3.458au
-0.064-000
-0.324aa.c0
-0.500g0
2.392a0011
selection of the later time steps. The distributions still show clear differences between tasks, such as the number of cho- sen time steps per trajectory or the frequency of later time steps. This suggests that STAP choosing time steps in highly adaptive manner.
We then ask ourselves: what if we perform random selec- tion, for instance with a probability for every time step? We call this method Bernoulli sampling, or Ber(p), where each entry of is true with probability p. Table 2 sum- marizes the performance Ber(p) for p = 1/16, 1/8, 1/4, and 1/2. In general, Bernoulli sampling improves over the base method SBAL, but it can also severely underperform at certain values of p, such as for KdV. Still, for each PDE, there exists at least one value of which Bernoulli sam- pling provides an advantage over the base method SBAL. In other words, sparse sampling of time steps itself has an inherent advantage over full full-traeectory sampling, as first hypothesized with our analysis in Fig.
5.5. Random Bernoulli Sampling of Time Steps
We plot in Fig. 5 the distribution of time steps that STAP chooses. The plot clearly shows the general tendency of STAP to acquire the early time steps, with an occasional
7

---

Active Learning with Selective Time-Step Acquisition for PDES
Table 3: Wall-clock time of each procedure during batch selection in INS. Note that these are not the costs of data acquisition, but the computational cost of batch selection algorithms.
Table 5: Log RMSE with multistep FNO averaged across 10 rounds of acquisition.
Random
SBAL
+STAP
Burgers 8L/8
-1.670=0.098
-18893=00083
+2.058+0.028
Random
QbC
LCMD
SBAL
KdV 81/8
1.402+0.029
1. 1,404=0.024
1.364+0.043
Time (scoodds)
0.1
62.5+0.1
100.3+0.1
62.5+0.1
KS 81.8
1.255
1.232=00012
1.156+0.008
KS L/2
1.340s0004
1.335=0.011
1.288
+STAP
+STAP MF
+STAP
INS L/2
1.124
1.118=0.007
1.081 +0.012
CNS L/2
3.593+0.023
3.594=0.044
3.420+0.050
132.7=0.7
91.5+0.1
15.4+0.1
(scconds)
0-
0-
0-
Table Log RMSE of more efficient STAP variants aver- aged across 10 rounds.
5-
SBAL
+STAP
+STAP MF
+STAP 10
Burgers KdV
-3,388am.022
-3.674.0071
-3.60sa.177
-3. 3.52240002
0.030g0.029
-0.08820000
-0.065an.04
-0.1 1820024
15
15
15
KS
-0,275-0.014
-0.349-aos
-0.326.0.04
-0.3 060000
INS
-0.422-0.010
00
-0.529.0.00
-0.502+0004
CNS
2.422a0:445
2.363a0016
2.360
2.401au004
20
20 20
How does STAP manage to be on par with or outperform the best Ber(p) on every task? One might hypothesize that the frequency of chosen time steps what matters most for performance and that STAP somehow finds the optimal frequency p to sample with. We have measured the frequency with which STAP samples time steps in the first round of AL, and obtained 0.35, 0.11, 0.19, 0.22, and 0.16 for Burgers, KdV, KS, INS, and CNS, respectively. One can see that for KdV, despite Ber(p) performing better with higher p and best with = 1/2, STAP outperforms all Ber(p) with frequency 0.11, which is close to 1/8. other words, the specific time steps that are sampled also matter as much as the overall frequency. These observa- tions altogether suggest that STAP adaptively chooses not only the frequency of the time steps acquire, but also their specific locations. We report the full results in Ap- pendix D.3, along with a variant of Bernoulli sampling that enforces acquiring consecutive initial time steps.
Figure 6: From left to right, ground truth trajectory on KS, trajectory predicted by surrogate model trained with trajectories, and trajectory predicted by surrogate model trained with one trajectory.
ternative is to decrease the number of greedy optimization steps T from 100 to 10, which reduces the cost by factor of 10. We call this variant STAP 10. The wall-clock time of each baseline method and STAP is measured with a sin- gle NVIDIA GeForee RTX 2080 Ti GPU, and summarized in Table 3. The performance of SBAL with STAP and its two variants are summarized in Table 4. Note that STAP 10 incurs only fraction of computational cost over the baseline methods, while still achieving significant boost in performance over its base method.
5.7. Multi-step model
5.6. Computational Complexity of STAP
We have done experiments with multi-step FNO which re- ceives N timesteps as input and outputs N timesteps. To perform STAP, we group the total number of timesteps into non-overlapping clusters of N timesteps, and perform STAP as if each cluster is one timestep with N channels. We have performed two variants of experiments. In the first, we divide timestep in our main experiment into 8 smaller timesteps, so that a total of L timesteps turn into SL timesteps, and train 8-in-8-out models. In the second variant, we keep L timesteps the same but train 2-in-2-out models. Fig. and Table shows the log RMSE for Burg-
The time complexity of computing our acquisition function for a single instance of (w",5) ) is O L). Since we op- timize the acquisition function with 7 steps, and we can acquire at most B initial conditions, the time complexity of our batch acquisition algorithm is O(MLBT) in the worst case. We can parallelize the optimization of multi- ple Sj's i to a certain extent using graphics processing unit (GPU), which can significantly alleviate the burden of B.
We can further reduce the cost by at most a factor of M with STAP MF described Appendix B.4. Yet another al-
8

---

Active Learning with Selective Time-Step Acquisition for PDEs
Burgers
KdV
KS
INS
CNS
1.5
1.5
4.5
1.5
1.0
1.0
4.0
1.0
0.5
0.5
3.5
0.5
3.0
-3
0.0
0.0
0.0
2.5
-0.5
-0.5
5 Round
10
10
10
10
5
10
Round
Round
Round
Round
Random
SBAL
SBAL+STAP
Figure 7: Active learning with initial training dataset containing one trajectory, as opposed to 32 trajectories in the main experiment. STAP is robust to initially inaccrrate surrogate models.
MaiI 120 DUSCOIIE
Malic two LEETURCIONN
can be properly evaluated. We find that only several of the states sampled by SBAL+STAP diverge significantly from the random ground truth states. In other words, the sur- rogate model synthesizes erroneous inputs when looked at trajectory-wise, but when looked at individually, they aren't t out-of-distribution enough to harm STAP significantly.
6. Conclusion
In this paper, we presented novel framework for active learning in surrogate modeling of partial differential equa- tion (PDE) trajectories, significantly reducing the cost of data acquisition while maintaining or improving model ac- curacy. By selectively querying only a subset of time steps in a PDE trajectory, our method STAP enables the acquisi- tion of informative data fraction of the cost of acquiring entire trajectories. We introduced new acquisition func- tion that estimates the utility of set of time steps based on variance reduction, effectively guiding the selection pro- cess in an adaptive manner. Through extensive experiments on benchmark PDEs, including the Burgers equation, Ko- rteweg-De Vries equation, Auramoto-Sivashine tion, incompressible Navier-Stokes equation, and com- pressible Navier-Stokes equation, we demonstrated that our approach consistently outperforms existing AL methods, providing more cost-efficient and accurate solution for PDE surrogate modeling.
Figure 8: PCA of FNO hidden layer activation for PDE states sampled by SBAL (left) and SBAL+STAP (right) during the first round of active learning on KS.
ers, KdV, and KS of the first variant, and KS, INS and CNS of the second variant.
5.8. Out-of-distribbtion Synthetic Inputs
Inaccurate surrogate models might synthesize inputs that lie far from the ground truth distribution, harming the rep- resentativeness (Dongrui Wu, 2018). Indeed, under limited training data, the surrogate model outputs visibly erroncous trajectories. Fig. 6 are comparisons of the ground truth tra- jectories and predictions from surrogate models trained on 32 trajectories and one trajectory, respectively.
equa-
To test how much this error harms STAP, we perform exper- iments where the initial training dataset contains 1 trajec- tory, compared to 32 used in our main experiments. Fig. 7 shows the log RMSEs To our own surprise, we find that SBAL+STAP still outperforms Random and SBAL, except for INS in the early rounds.
Our results show that STAP can significantly enhance sur- rogate modeling for PDEs, particularly in scenarios where the numerical solver is computationally expensive. We fur- ther showed that the success of STAP is driven by its ability to prioritize both diverse and informative time steps. Mov- ing forward, this framework could be extended to more complex systems and integrated with other machine learn- ing techniques, providing broader applicability in scientific and engineering simulations. Future work may also explore alternative acquisition functions and applications to simu- lations outside the domain of PDEs.
Fig. 8 shows comparisons of FNO activation PCAs for PDE states sampled by SBAL and SBAL+STAP during the first round of active learning on KS Note that the PCA was fitted to ground truth states in random trajec- tories (blue points), independent from the states sampled by SBAL and SBAL+STAP, so that distirifsssriituiooneesss

---

Active Learning with Selective Time-Step Acquisition for PDEs
Acknowledgements
Bar-Sinai, Yohai et al. (2019). "Learning data-driven dis- cretizations for partial differential equations". In: Pro- ceedings of the National Academy of Sciences 116.31, pp. 15344-15349 (cit. on 1).
This work was partly supported by Institute of Informa- tion & communications Technology Planning & Evalua- tion(IITP) grant funded by the Korea goverimen((MSIT) N.o.RS-2024-005092779 Global AI Frontier Lab; No.RS- 2019-11190075, Artificial Intelligence Graduate School Program(KAIST) No.RS-2022-1122018. Development and Study of AI Technologies to Inexpensively Conform to Evolving Policy on Ethics), and the National Research Foundation of Korea(NRF) grant funded by the Korea gov- emmenn(MSSTT) (NRF-2022R 21--2022RIASA70S ASAA7833907 12). We also thank Hyungi Lee for his thoughtful discussion.
Brandstetter, Johannes, Max Welling, and Daniel E Wor- rall (2022a). "Lie point symmetry data augmentation for neural PDE solvers". In: International Conference on Machine Learning. PMLR, pp. 2241-2256 (cit. on pp. 15, 16).
Brandstetter, Johannes, Daniel Worrall, and Max Welling (2022b). "Message passing neural PDE solvers". In: arXiv preprint arXiv:220.003376 (cit. on p. 1).
Impact Statement
Chassignet, Eric P et al. (2007). "The HYCOM (hybrid coordinate ocean model) data assimilative system". In: Journal of Marine Systems 65.1-4, pp. 60-83 (cit. on p. 16).
We propose a new method that improves the cost efficiency of acquiring data for building a surrogate model PDE trajectories. Although our approach doesn't have a direct positive or negative impact in ethical societal aspects, it accelerates the process of building surrogate model for an arbitrary PDE. This could be used for good, such as medical simulations, environmental modeling, and opti- mizing engineering designs, potentially leading to advance- ments in healthcare, sustainability, and technological in- novation. However, like many technologies, this method could also be misused in domains where rapid simulations could have harmful consequences, such as the development of hazardous materials. Therefore, researchers and practi- tioners should apply these methods with consideration of their broader societal implications, aiming to ensure that the benefits of the technology are used responsibly and eth- ically.
Chernoff, Herman (1959). "Sequential Design of Experi- ments". In: The Annals of Mathematical Statistics 30.3, pp. 755-770 (cit. on p. 1).
Cleaver, Timothy A et al. (2016). "Using design of ex- periments methods for applied computational fluid dy- namics: case study". In: Quality Engineering 28.3, pp. 280-292 (cit. on p. 1).
Cook, Alex R, Gavin Gibson, and Christopher A Gilligan (2008). "Optimal observation times in experimental epi- demic processes".. In: Biometrics 64.3, pp. 860-868 (cit. p. 5).
Cowan, Timothy J, Andrew S Arena Jr. and Kajal K Gupta (2001). "Accelerating computational fluid dynam- ics based acroelastic predictions using system identifi- cation". In: Journal of Aircraft 38.1, pp. 81-87 (cit. on
References
Arthurs, Christopher and Andrew P King (2021). "Active training of physics-informed neural networks to aggre- gate and interpolate parametric solutions to the Navier- Stokes cquations". In: Journal of Computational Physics 438. p. 10366 (cit. on p. 5).
p. 1).
Evans, WD (1988). PARTIAL DIFFERENTIAL EQUA- TIONS (cit. on p. 2).
Ash, Jordan et al. (2019). "Deep batch active learning by diverse, uncertain gradient lower bounds". In: arXiv preprint /006::9900.0000 (cit. on p. 4).
Gajjar, Aarshvi, Chinmay Hegde, and Christopher P Musco (2022). "Provable active learning of neural networks for parametric PDEs". In: The Symbiosis of Deep Learning and Differential Equations Il (cit. on 5)
Atkins, Peter William, Julio De Paula, and James Keeler (2023). Atkins' physical chemistry. Oxford university press (cit. on p. 1).
Gao, Wenhan and Chunmei Wang (2023). "Active learn- ing based sampling for high-dimensional nonlinear par- tial differential equations". In: Journal of Computational Physics 475, p. 111848 (cit. on P. 5).
Bajracharya, Pradeep et al. (2024). "Feasibility study active learning of smart surrogates for scientific sim- ulations". In: arXiv preprint arXiv:2407.07674 (cit. on pp. 1, 5).
Gelbart, Michael A, Jasper Snoek, and Ryan P Adams (2014). "Bayesian optimization with unknown con-
10

---

Active Learning with Selective Time-Step Acquisition for PDES
straints". In: arXiv preprint arXiv: (cit. on p. 17).
Kirsch, Andreas, Joost Van Amersfoort, and Yarin Gal (2019). "Batchbald: Efficient and diverse batch acquisi- tion for deep bayesian active learning". In: Advances in neural information processing systems 32 (cit. on pp. 2, 4).
Goswami, Somdatta
(2022). "Deep transfer opera- tor learning for partial differential equations under con- ditional shift", In: Nature Machine Intelligence 4.12,
pp. 1155-1164 (cit. on p. 17).
Kirsch, Andreas et al. (2023). "Stochastic Batch Acqui- sition: Simple Baseline for Deep Active Learning". In: Transactions on Machine Learning Research. ISSN: 2835-8856 (cit. on p. 5).
Greydanus, Samuel, Misko Dzamba, and Jason Yosinski (2019). "Hamiltonian neural networks". In: Advances in neural information processing systems 32 (cit. on p. 1).
Kusner, Matt J, Brooks Paige, and José Miguel Hernández- Lobato (2017). "Grammar variational autoencoder". In: International conference on machine learning. PMLR, pp. 1945-1954 (cit. p. 16).
Hernandez-Garcia, Alex et al. (2023). "Multi-fidelity active learning with gflownets". In: arXiv preprint MrXiv:2306.11755 (cit. on onp.3) 5).
Hernández-Lobato, José Miguel et al. (2015). "Predictive entropy search for Bayesian optimization with unknown constraints". In: International conference on machine learning. PMLR, pp. 1699-1707 (cit. on p. 17).
Li, Shibo et al. (2022a). "Batch multi fultii-iieeiity active learn- ing with budget constraints", In: Advances in Neural In- formation Processing Systems 35, pp. 995-1007 (cit. on p. 3)
Holton, James R and Gregory Hakim (2013). introduc- lion to dynamic meteorology. Vol. 88. Academic press (cit. on p. 1).
Li, Shibo et al. (2022b). "Deep Multi-Fidelity Active Learning of High-Dimension Outputs". In: Interna- tional Conference on Artificial Intelligence and Statis- tics, AISTATS 2022, 28-30 March 2022, Virtual Event (cit. on pp. 5, 13).
Holzmüller, David et al. (2023). "A framework and bench- mark for deep batch active learning for regression". In: Journal of Machine Learning Research 24. 24,164, pp. 1-81 (cit. on pp. 1, 1.4-7,10 17).
Li, Shibo et al. (2024). "Multi-Resolution Active Learn- ing of Fourier Neural Operators". In: International Con- ference Artificial Intelligence and Statistics. PMLR, pp. 2440-24424442 2448 (cit. on 5).
Houlsby, Neil et al. (2011). "Bayesian active learning for classification and preference learning". In: arXiv preprint arXiv:1112.574. (cit. on 2)
Li, Zongyi et al. (2020). "Fourier neural operator for para- metric partial differential equations". In: arXiv preprint 200.:22000.08895 (cit. on pp. 1, 6, 16).
Jarrin, Nicolas et al. (2006). "A for generating inflow conditions for large-eddy simula- tions". In: International Journal of Heat and Fluid Flow 27.4, pp. 585-593 (cit. p. 16).
method
Lippe, Phillip et al. (2024). "Pde-refiner: Achieving accu- rate long rollouts with neural pde solvers". In: Advances in Neural Information Processing Systems 36 (cit. on p. 1).
Jumper, John al. (2021). "Highly accurate protein struc- ture prediction with AlphaFoldd" In: nature 596.7873, pp. 583-589 (cit. on 16).
Loshchilov, Ilya and Frank Hutter (2016). "Sgdr: Stochas- tic gradient descent with warm restarts". In: arXiv preprint 1608.0660 (cit. on p. 6).
Kalnay, E (2003). Atmospheric Modeling, Data Assimila- tion and Predictability. Vol. 341. Cambridge University Press (cit. on 16).
MacKay, David JC (1992). "Information- objective functions for active data selection". In: Neural compu- lation 4.4, pp. 590-604 (cit. on p. 1).
Kassam, Aly-Khan and Lloyd N Trefethen (2005). "Fourth- order time-stepping for stiff PDEs". In: SIAM Journal on Scientific Computing 26.4, pp. 1214-1233 (cit. on p. 16).
Mao, Zhiping and Xuhui Meng (2023). "Physics-inforee informed neural networks with residual/gradient-asse adaptive sampling methods for solving partial differential equa- tions with sharp solutions". In: Applied Mathematics and Mechanics 44.7, pp. 1069-1084 (cit. on p. 5).
Kingma, Diederik P (2014). "Adam: A method for stochas- tic optimization". In: arXiv preprint 141.6.1412.6980 (cit. on p. 6).

---

Active Learning with Selective Time-Step Acquisition for PDEs
Murray, James D (2007). Mathematical biology: I. An in- troduction. Vol. 17. Springer Science & Business Media (cit. on p. l).
Takamoto, Makoto al. (2022). "Pdebench: An exten- sive benchmark for scientific machine learning". In: Ad- vances in Neural Information Processing Systems 35, pp. 1596-1611 (cit. on pp. 15, 16).
Musekamp, Daniel et al. (2024). "Active Learning for Neu- ral PDE Solvers". In: arXiv preprint arX0v:2408.01336 (cit. on pp. 1-6).
Taylor, Karl E, Ronald J Stouffer, and Gerald A Meehl (2012). "An overview of CMIP5 and the experiment de- sign". In: Bulletin of the American meteorological Soci- ety 93.4, pp. 485-498 (cit. on p. 16).
Pestourie. Raphaël et al. (2020). "Active learning of deep surrogates for PDEs: application to metasurface design". In: npj Computational Materials 6.1, p. 164 (cit. on pp. 5).
Turinici, Gabriel (2024). "Optimal time sampling in physics-informed neural networks". In: arXiv preprint arXiv:2404.18780 (cit. on p. 5).
Pickering, Ethan et al. (2022). "Discovering and forecast- ing extreme events via active learning in neural opera- tors". In: Nature Computational Science 2.12, pp. 823- 833 (cit. on pp. 5, 6).
Wilmott, Paul, Sam Howison, and Jeff Dewynne (1995). The mathematics of financial derivatives: student in- troduction. Cambridge university press (cit. p. 1).
Roy, Nicholas and Andrew McCallum (2001). "Toward op- timal active learning through monte carlo estimation of error reduction". In: leml, williamstown 2.441-444 p. 4
Wu, Chenxi al. (2023). "A comprehensive study of non-adaptive and residual-based adaptive sampling for physics-informed neural networks". In: Computer Methods in Applied Mechanics and Engineering 403, p. 115671 (cit. on 5).
(cit. on p. 13).
Salkin, Harvey M and Cornelis A De Kluyver (1975). "The knapsack problem: survey". In: Naval Research Logis- tics Quarterly 22.1, pp. 127-144 (cit. on 4).
Wu, Dongrui (2018). "Pool-based sequential active learn- ing for regression". In: IEEE transactions on neural net- works and learning systems 30.5, pp. 1348-1359 (cit. on pp. 1, 4, 9).
Sanchez-Gonzalez, Alvaro et al. (2020). "Learning to sim- ulate complex physics with graph networks". In: In- ternational conference on machine learning. PMLR, pp. 8459-8468 (cit. on p. 1).
Wu, Dongxia
(2023). "Disentangled multi-fidelity deep bayesian active learning". In: International Con- ference on Machine Learning. PMLR, pp. 37624-37663 (cit. pp. 5).
Sener, Ozan and Silvio Savarese (2017). "Active learning for convolutional neural networks: A core-set approach". In: arXiv preprint :tti::/700.0044 (cit. on pp. 5, 6).
Zabusky, Norman J and Martin D Kruskal (1965). "Interac- tion of" solitons" in collisionless plasma and the recur- rence of initial states". In: Physical review letters 15.6, p. 240 (cit. on pp. 6, 15).
Settles, Burr (2009). Active learning literature survey. Tech. rep. University of Wisconsin-Madisoe Department of Computer Sciences (cit. on pp. 1, 13).
Seung, H Sebastian, Manfred Opper, and Haim Sompolin- sky (1992). "Query by committee". In: Proceedings of the fifth annual workshop on Computational learning theory, pp. 287-294 (cit. on pp. 5)
Singh, Rohit et al. (2005). "Active learning for sampling in time-series experiments with application to gene expres- sion analysis". In: Proceedings of the 22nd international conference Machine learning, pp. 832-839 (cit. p. 5).
Smith, Freddie Bickford et al. (2023). "Prediction-orieeneed bayesian active learning". In: International Confer- ence on Artificial Intelligence and Statistics. PMLR, pp. 7331-7348 (cit. on p. 4).
12

---

Active Learning with Selective Time-Step Acquisition for PDEs
A. Further explanation of Acquisition with STAP
A.1. Different types PDE active learning
u
y
u
u
(a) Qol
(b) Single-state
u
u !
u
(c) Autoregressive
Figure 10: Task settings assumed by previous works in active learning of PDEs.
There are three primary tasks in active learning for PDEs, each depending on the type of surrogate model being trained. The first task, univariate Quantity of Interest (Qol) prediction, focuses on learning model to directly predict a scalar Qol, denoted as y, from an initial condition ua. The second task, single-state prediction, involves learning a model to predict single state transition from 22 u over fixed time interval At. The third task, autoregressive trajectory prediction, aims approximate the ground truth evolution operator G using a surrogate model predict the entire time evolution of the states. Fig. 10 provides a visual comparison of the three tasks. In this paper, focus the autoregressive trajectory prediction task.
A.2. Motivation behind the Acquisition Function
Here we detail the motivation behind our acquisition function defined in $ 3.2. First, one can imagine several alternative acquisition functions.
The most straightforward alternative is to simply use the sum of the variances at time points for which b; = true. The variances are larger for the later time steps since they accumulate, and in our preliminary experiments, we found that this is catastrophic as undersampling the earlier time steps leads to the sampled trajectory being very out-of-distriuution, and hence the trained surrogate model underperforming on the test distribution.
quickly became clear us that we need some kind of measure of "how much total uncertainty will be reduced by sampling these time steps", instead of "how uncertain is model on these time steps?" This would help select sampling patterns that reduce the cottoffdditrrbution-ness introduced by G. One way to approximate this is to use mutual information, as used by S. Li et al. (2022b). In other words, we would rollout M trajectories with M surrogate models, and compute the mutual information between time steps for which by = true and all time steps. However, in preliminary experiments, we found that this method underperforms, which we hypothesize because relying simply on the covariance matrix of the committee between time steps is not good enough method for computing the posterior uncertainty.
We identified two "pathways" through which sampling a time step reduces uncertainty in the remaining time steps. First, there is the "indirect" pathway: sampling time step will reduce the model's uncertainty similar inputs, hence reducing uncertainty on the remaining time steps. This is what is approximated by mutual information. Then, there is the "direct" pathway: sampling time step i gives out the i + 1 th state, which starts chain reaction of reducing model uncertainty on all successive states. Note that these two pathways are not distinct from a strictly theoretical view, but are rather two ways of approximating uncertainty reduction.
The direct pathway motivated our acquisition function based on variance reduction. variance reduction, we calculate the posterior uncertainty by rolling out the trajectories with N surrogate models, but collapse into one surrogate model at time steps for which b; = true. This effectively computes the reduced uncertainty due to the effect of the direct pathway. With experiments, we confirmed that this acquisition function behaves just like we wanted: it is slightly biased towards sampling the earlier time steps, and it chooses an appropriate frequency of time steps to sample that leads to good performance.
To elaborate, our acquisition function is an approximation to the expected error reduction (EER), which statistically near-optimal for active learning (Settles, 2009; Roy and McCallum, 2001). The EER measures how much the model's gen-
13

---

Active Learning with Selective Time-Step Acquisition for PDEs
Algorithm Batch Acquisition Algorithm of STAP
Require:
Budget B, base active learning algorithm A, probability e, number of iterations T for greedy optimization, pool P of initial conditions, cost function cost(-) for batches.
Ensure: A batch B of initial conditions and sampling patterns.
1: B 8 Ø
2: while cost( < B do
3: Acquire an initial condition u with A.
4: Initialize S & (true,.... true).
5: for i 1 to T do
6: C (C1, (Ci...., CL) where
i.dd Ber(s).
7:
5' = S + C
8:
if af"" 2 a" (u", then
9:
S
10:
end end for
12:
if
||S]| + (B) > B then
13:
Keep only the first (B cost(B)) trues from S and flip the remaining trues. end if
14:
15:
B + BU BU ((u,,)}
16:
end while
eralization error is likely reduced after updating on hypothetically acquired data. We model our hypothetical belief about the ground truth solver as a #uniform categorical distribution" over the ensemble [G We assume that acquiring the trajectory of u" with sampling pattern S only reduces generalization error on the trajectory of u". The current general- ization error expected to be the average of || ào over b. We make a second assumption that the hypothetically acquired data S,a will update the model such that the model predicts the trajectory úb.5.a given w". This gives us the expected reduction error (ā -0,- - 11 ds l| averaged across a and b, which equal to our acquisition function.
A.3. Batch Acquisition Algorithm
Algorithm summarizes the batch selection algorithm STAP. Starting with an empty batch B, the algorithm repeatedly selects initial conditions and their sampling patterns until reaching the budget limit. It first uses the base active learning method A to choose initial condition n., Then, it optimizes which time steps sample through greedy procedure: starting with pattern S that samples all time steps (all true values), it performs T iterations of random mutations. In each iteration, it generates candidate pattern S' by randomly flipping entries in S with probability (using binary mask C where each entry is drawn from Bernoulli distribution and the XOR operation @). If this new pattern achieves better value according to the cost-weighted acquisition function a", it becomes the current pattern. To ensure the budget isn't tt exceeded, if adding the current pattern would go over budget, the algorithm truncates by keeping only enough true values exactly meet the budget. The pair of initial condition and its optimized sampling pattern (u", SS) is then added the batch B.
B. Experimental details
B.1. Details on PDEs
In this section, we describe the PDEs used in our experiments. Each of these equations plays critical role in modeling physical phenomena and showcases diverse behaviors, from diffusion and soliton dynamics to chaotic systems and fluid flow. Examples of PDE trajectories are shown Fig. 11.
Burgers' Equation The one-dimensional Burgers' equation is expressed as
uu uZ, (v/T)0znu
where u = u(a, t) represents the field and v 2 is the viscosity parameter. The Burgers' equation exhibits shock phenomena, characterized by the gradients of becoming extremely sharp or even discontinuous when v These
14

---

Active Learning with Selective Time-Step Acquisition for PDEs
50
50
300
130
150
200
200
250
250
50
100
125
(a) Burgers
(b) KdV
(c) KS
0
20
0
20
0
20
20
20
0
20
0
20
(d) INS
20
20
20
0
20
20
20
0
20
(c) CNS
Figure 11: Example trajectories of PDEs. (a), (b), (c): Horizontal and vertical axes represent the temporal and spatial
domain. (d), (e): Two-dimensional states at six time points are shown.
shocks arise due the the advection term ud,u, while the presence of the diffusion term aggu prevents the formation of discontinuities in the wave. We set 0.01 ensure the formation of sharp gradients while maintaining numerical stability. For simulations, we adopt the implementation by Takamoto et al. (2022), which employs finite difference method (FDM) to compute the two terms.
Korteweg-De Vries (KdV) Equation The second equation we study is the Korteweg De Vries (KdV) equation, given by:
0,u + ud.44 + Ōzzn O,
where u uux I) represents a wave profile evolving over space and time. This nonlinear PDE describes the evolution of shallow water waves, and its most famous characteristic is the presence maintain their shape over long distances and weak interactions with other waves (Zabusky and Kruskal, 1965). Solitons have important applications in fluid dynamics, plasma physics, and optical fiber communications. The KdV equation's 'ss nonlinearity and third-order spatial derivative (izza) allow it to capture complex wave behavior. The equation is also known for conserving key quantities like energy. We solve this equation using the pseudospectral method with Dormand-Prince solver as in Brandstetter et al. (2022a).
stable wave packets that
Kuramoto-Sivashinsly (KS) Equation
The Kuramoto Sivashinsky (KS) equation a fourth-order nonlinear PDE,
Curamoto-SSvvshinskyy
written as:
D.u + Orzit + Orrzzu ud,u 0,
15

---

Active Learning with Selective Time-Step Acquisition for PDEs
where u(x; 1) is the evolving field in space and time. The KS equation is known for its chaotic behavior and is used model phenomena such as flame front propagation, plasma instabilities, and thin film dynamics. Its chaotic nature arises from the interplay between destabilizing nonlinear terms and stabilizing higher-order diffusion terms. The equation is particularly challenging to solve due to its sensitivity to initial conditions and long- term unpredictability. To handle this complexity, we use the Exponential Time Differencing (ETD) fourth-order Kunge-Kutta method, as introduced by Kassam and Trefethen (2005). This numerical method is well-suited for stiff PDEs like the KS equation.
Incompressible Navier-Stokes (INS) Equation We consider the vorticity form of the 2D incompressible Navier-Stokes (INS) equation, which governs the motion of incompressible and viscous fluid flows. The equation is given by:
du++ + V Vu = vu=vVuu+f V.v=0.
(10)
where u(z11, I22 1) is the vorticity, v is the velocity field, v is the kinematic viscosity, and [(11, is an external forcing term. This equation describe the behavior of incompressible fluid flow. playing a central role in understanding turbulence, weather patterns, and aerodynamics. The external forcing term 322) is set
f(x) = 0.1 (sin(27(1) + x2)) + cos (2 + x2))),
(11)
which injects energy into the system, driving complex fluid dynamics. We use kinematic viscosity of = 2 x 10-1 ensuring a sufficiently turbulent regime. In our experiments, we adapt the Crank method implemented by Z. Li et al. (2020).
Compressible Navier-Stokes (CNS) Equation Finally, we consider the 2D compressible Navier-Stokes (CNS) equation, which extends the incompressible case by incorporating density variations. We adapt the form used in Takamoto et al. (2022), which is given by:
dep + V (0o) = 0,
1
p(0:0+ + U. -v) = - + +p+ + (S+ + 3 VVV v),
(12)
0.(e+ 1 500²)+ + V VI 1 _00²)0 - u. v-oll = 0,
where it has four variables (density p. velocity v- = (Va, vyy) and pressure p) and two fixed parameters (shear viscosity ? and bulk viscosity (). The terms a' and are the viscous stress tensor and the internal energy, respectively. We set 17 = < = 0.05. The numerical solution was computed using a second-order scheme in both time and space, employing a high-resolution reconstruction method for the inviscid terms and a central differencing approach for the viscous terms.
Initial conditions As per Brandstetter et al. (2022a), states are first sampled from a simple distribution and then evolved for certain time to obtain the initial conditions. The evolved initial conditions are more realistic than the sampled states, in that they are more likely to be observed under a system governed by the respective PDEs. This procedure hence approximates applications where the initial conditions of interest are realistic states either from observed data (Jumper et al., 2021; Kalnay, 2003; Chassignet et al., 2007; Taylor et al., 2012) or carefully crafted synthetic data (Jarrin et al., 2006; Kusner et al., 2017). For 1D equations, the states are sampled from truncated Fourier series with random coefficients (Brandstetter et al., 2022a), and for the 2D INS equation, states are sampled from a Gaussian random field as described in Z. Li et al. (2020). Similarly, for the CNS equation, the initial conditions for all four variables are constructed by superposing random sinusoidal waves, with density and pressure renormalized to ensure positivity. The lengths and discretizations of trajectories summarized in Table 6.
Let U stand for the uniform distribution. For Burgers and KdV, the initial condition is in the form ΣN A, sin(2 6.). The amplitudes and phases are always sampled from U U(O, 1I) and U[[ 27)). For Burgers, N = 2 and k ~ U1,, 3, 4), and for KdV, N = 10 and K; ~ U(1 2, 3). For KS and INS, the initial conditions are gaussian random fields drawn from (0, 25 25 - + 25/) and (O, (3/2 + 191)- -2.5) respectively. For CNS, we first sample E (1,2,3,43 Ag sin(2.rkr / T. + ) the amplitudes and phases sampled uniformly at random for each channel (rho, p and v). Then, we renormalize rho, to lie within (1 + Ap) and Po + Ap), respectively, where Po ~ U[0.1, 10)). Ap ~ ([0.013, 0.26)), A, ~ U[0.044. 0.8)) and To := Po/Po ~ U[0.11, 10)). The v is also computed by superposing si- nusoidal waves, but with amplitudes chosen so that the initial condition has the given initial Mach number M ~ U(001, 1)).
16

---

Active Learning with Selective Time-Step Acquisition for PDEs
Table 6: Domain lengths and discretizations for trajectory learning.
PDE
Domain Length (T
Resolution (L, No)
Burgers
(2.0, 1.0)
(13, 256)
KdV
(52.0, 128.0)
(13, 256)
KS
(13.0, 1.0)
(26, 256)
INS
(13.0, 1.0, 1.0)
(26, 32, 32)
CNS
(0.5, 1.0, 1.0)
(26, 32, 32)
Table 7: Acquired datasize in KdV
Round
2
3
7
10
SBAL
416
520
624
728
832
936
1040
1144
1248
1352
1456
SBAL+STAP
416
507
611
715
819
923
1027
1131
1235
1339
1443
B.2. Error Metries
The test always consists of 1,000 trajectories, on which several error metries are defined. The RMSE defined trajectory
1/ (xx(xx) tt (2,)
(13)
LN 1=1 j=1
Similarly, the NRMSE is defined as
Ei. Ju (x;)) (x=)))
(14)
En, 1/ (x,)
and the MAE as
L. Σ Σ 1u'xxx (x₂) a (xx))-
(15)
The metrics are averaged across all trajectories in the test set. We also report their logarithmic values averaged across all AL rounds, following Holzmüller et al. (2023). Note that we do not use committee's 's mean prediction for computing the metries, but instead compute the metrics for each model and report their average.
B.3. Simulation Instability
was observed that using STAP on the KdV equation, the simulation crashes on small subset of synthetic inputs. Analysis reveals that these synthetic inputs have unusually large norms and particularly appear in later parts of trajectories due to accumulated error. We do not attempt fix this problem explicitly due to the risk of wer-complicating our method, and simply refrain from adding these time steps to the training dataset. This means that STAP actually acquires smaller number of time steps than the budget B per round of acquisition, which could be problematic when large subset of inputs do crash. However, we find that this is not the case, and the number of such inputs is small enough that STAP can outperform other baselines. We report the comparison of datasize across rounds in Table 7, for single experiment. We can see that 13 time steps were left out the first round due to instability, and no instability occurred in the rounds after.
Since queries that crash incur cost, they should be avoided as much as possible. Previous works Bayesian optimization (Gelbart al., 2014; Hemández--oobtto al., 2015) propose methods learn these unknown constraints. Alternatively, one could simply test out large, random inputs. In fact, we find that the maximum absolute value of an input being above 10 is robust criterion for predicting that the solver will crash. Either way, could simply filter out time steps that fall outside of these constraints during runtime of the solver, and use the freed up budget on acquiring other trajectories. Another possible approach is to impose physical constraints on the surrogate model (Goswami et al., 2022) that reduces the risk of outputting abnormal synthetic inputs. For instance, the KdV equation is mmergy-conservngg and when this prior knowledge is encoded into the surrogate model, the synthetic inputs would never be abnormally large like we experienced with our naive surrogate models.
17

---

Active Learning with Selective Time-Step Acquisition for PDEs
B.4. STAP MF
We can also define simpler acquisition function in the spirit of mean-field approximation. We take the mean model -= Em Gm. and define the variance reduction R(G, b, .) between G and model G, in the same way as before. We then average the variance reduction between the mean model and all models in the committee:
M
OMF (u. ()
R)),
(16)
M b=1
which reduces the computational cost by factor of M in the best case. We call this modified version STAP P MF.
C. Full Report of Results Main Experiment
Table 8: Mean log metrics for Burgers' Equation
RMSE
NRMSE
MAE
99%
95%
50%
Random
20.8120.000
3.955 060
-4.813 067
00.28820. 145
-1.701:0.0m
-3.924.10.00
Random+STAP
-0.477=00061
-4.43320.00 078
-5.39420.00 0.061
-1.010=0.10
-2.55120.100
-4.27811..0
SBAL
-3.388=00055
-4.08120.08)
-5.332
-1.677 065
-2.375
-3.88510.00f
SBAL+STAP
-0.674=0.711
-1.306=0.066
-5.585=0.066
2.1885
-2.853=0.077
-4.017=0.070
QbC
30.21
3.8055 062
-1.483=0.00 0.001
-2.098=200rr
-3.571+0.068
QbC+STAP
3.333=0.064
-3.941 =0. 069
-5.250=a.0ms 050
-1.874+0.0s
-2.153=20.655
-3.671+0.066
LCMD
-2.847 00007
-3.555 00.022
-4.819 00.00
-1.160 041
-1.758:0.013
-3.338
LCMD+STAP
2.92550.0611
3.5466
-
-1.281
-1.874
-3.360+0.079
Table Mean log metrics for KdV Equation
RMSE
NRMSE
MAE
95%
50%
Random
0.191
-1.19320.000
-2.034+0 045
2.449 047
1.395
-1.19620 043
Random+STAP
0..044
-14425=0.0447
-2.228 us6
1.885 033
1.296 038
-1.424+0.04
SBAL
0.030+t0.029
-1.282#4.000
-2.139+a.u27
1.875+a..0g
1.267=0.0o
-1.267#0.00
SBAL+STAP
0.088
-1.378
-2.239
1.731 #0.000
1.280 043
-1.378
QbC
0.266.10.027
-1.019=0.02
-1.87910.00
1.859 037
1.251
-1.01920.
QbC+STAP
0.134
-1.130=0.00
-2.00410 035
1.721:0 031
1.120=0.08 030
-1.286:1.00 035
LCMD
0.256:10.053
-1.03320.008
-1.87910.00 033
1.868 034
1.322:0..0.
-1.100:0.00 0.0088
LCMDASTAP
0.286
-0.978=0.034
-1.821+0.039
1.799
1.128=0.034
-1.129+0.032
Table 10: Mean log metrics for KS Equation
RMSE
NRMSE
MAE
99%
95%
50%
Random
-0.558
-..683
-2.165
1.097
0.752=0.00
-0.575 004
RTAP
-0.335++..014
-1.759-0.v14
-2.248
1.060#0.01z
0.691#u.oe
-0.662+0.01
SBAL
-01
-1.700=0.014
-2.18410.0m
1.086 017
0.732=0..m 023
-0.59420 012
SBAL+STAP
-0.34920.00 003
-1.774=0.003
-2.265-10.0 003
1.04220.00 011
0.672 0122
-0.67420 008
QbC
00.268 004
-1.693=20.00
-2.178:10.00 0044
1.077 008
0.739=0.00 013
-0.58210 008
QbC+STAP
0.331
-1.756=0.014
-2.246+0.014
1.050
0.681 = o20
-0.650
LCMD
0.01642.015
-1.378=a.015
-1.829+0.015
1.20444.00 o00
0.95-1=0.011
-0.203=0.00
LCMDASTAP
0.138
-1.561.0.m16
-.0033+0.017
1.139
00.14
-0.431 20.016
We provide a full report of all results from the main experiment. Table 8, Table 9, Table 10, Table 11, and Table 12 show the full results on Burgers, KdV, KS, and NS equations, respectively. Fig. PDEs.
shows the plots RMSE quantiles on all
18

---

Active Learning with Selective Time-Step Acquisition for PDEs
Random
QBC
SBAL
-3.5
LOMD
Mothod
Method
SBAL+STAP
Random
Rardom
-2
ObC
SBAL
-4.0
SOAL
QbL:
LCMD
LCMD
-3
SBAL+ STAP
-4.5
SBAL+
2
6
10
6
10
2
10
Round
Round
Round
(a) 99% quantile, Burgers
(b) 95% quantile, Burgers
(c) 50% quantile, Burgers
3.0
Method
2.00
Random
2.5
QaC:
1.75
-0.5
LOMD
Method
Method
Random:
1.25
QhC
-1.0
QbC
SBAL
LCMD
ICMD
1.5
0.75
-1.5
2
10
Round
Round
Round
(d) 99% quantile, KdV
(c) 95% quantile, KdV
(f) 50% quantile, KdV
1.3
1.0
0.0
Random
1.2
1.1
Method
0.8
Method
SEAL
LCWD
1.0
QbC
0.6
QbC
SBAL+5114 STAP
0.9
LCMD
0.4
LCMD
-1.0
S
SBAL+STAP
0.8
2
6
8
10
2
6
10
2
8
10
Round
Round
Round
(g) 99% quantile. KS
(h) 95% quantile, KS
(i) 50% quantile, KS
Method
Methoc
Method
0.75
ObC
Random
0.4
handom
Q&C
Random
SBAL
SEAL
0.50
LCMD
LCMD
0044
LCVD
+STAP
SBAL+STAP
SBAL+5TA
0.25
0.0
0.00
-0.2
-0.8
-0.4
-1.0
2
8
10
2
8
10
2
8
10
Round
Round
Round
(j) 99% quantile, INS
(k) 95% quantile, INS
(l) 50% quantile, INS
Mothod
4.0
2.6
4.75
Random
QbC
3.8
4.50
50AL
LCMP
3.6
2.2
Methidd
Rardom
Random
4.25
SBAL+STAP
3.4
QbC
2.0
SBAL
SBAL
3.2
LEMO
1.8
LCMD
3.75
3.0
SBAL+STAP
53AL+5TA7
10
10
Round
Round
Round
(m) 99% quantile, CNS
(n) 95% quantile, CNS
(0) 50% quantile, CNS
Figure 12: Log RMSE quantiles
19

---

Active Learning with Selective Time-Step Acquisition for PDEs
Table 11: Mean log metrics for INS Equation
RMSE
NRMSE
MAE
99%
95%
50%
Random
0.42220.0100
-2.71440.010
-.6998t..010
0.515:0.002
0.13420.012
-0.520:20.000
Random=STAP
00.489+0.012
-2.780+0.012
-2.76420.012
0.4112+0.042
0.037=0..25
-0.581
SBAL
(0.46110.01
-2-021
-2.736:00.009
0.37220.039
0.043 22.424
-0.54320.004
SBAL+STAP
0..255+0.000
-2.814++0.0us
-2.797+0.000
0,268
-0.043++0.00r
-0.601+0.003
QbC
0.385
-2.675+t0.u1s
-2.663+0.016
+.282++0.024
0.047=0.0m
-0.439
OBC+STAP
0.466++0.011
-2.757
-2.746+0.010
0.205++0.021
0.036
-0.518+v.010
LCMD
0.320
-2.007
-0.587
0.532++0.19
0.202
-0.400
LCMD+STAP
00.77810.012
-2.66510.012
-2.645+0.011
0.47611.0117
0.145
-0.46210.011
Table 12: Mean log metrics for CNS Equation
RMSE
NRMSE
MAE
99%
95%
50%
Random
2.603#0.034
2.498=
0.732==0.aa1
4.61820.057
3.806=0.050
2.041++0.000
Random+STAP
2.461=0.017
-2.377=0.015
-0.863+0.013
4.458=0.057
3.659=0.02m
1.905++0.030
SBAL
2.422:10.045
-2.482:0.040
-0.850
4. 4.206 038
3.497:0.064
2.024:10:04
SBAL+STAP
2.363:20.01
-2.501
-0.932:100015
0.096
3.423
1. 1.99010..031
QbC
2.84410.021
-2.102.0.024
-0.476.10.018
4.329
3. 3.899
2,539 1.0015
QhC+STAP
2.78820.017
-2.06320.021
-0.53420.015
4.249=0.017
3.783=0.010
2.518+0.022
LCMD
2.736=0.02
-2.227=0.02
-0.58240.025
4.263=0.021
3.80320.032
2.401 #0.032
LCMD4STAP
2.73920.00 002
-2.14720.00
-0.57220.002
4.14320.017
3.72620.017
2.47820.015
D. A Additional experiments
D.1. Active learning with 20 Rounds
We have performed the main experiment for 20 rounds instead of 10, on Random, SBAL, and SBAL+STAP. The results are shown in Fig. 14. We observe that the gap between SBAL and SBAL+STAP keeps widening, except in the KdV equation.
D.2. Diversity of Sparsely Selected Time Steps
We provide simple analysis to show that time steps sampled in sparse manner are more diverse than time steps from entire trajectories. Out of 128 trajectories, we first randomly chose 10 trajectories, which contains L 10 states. Then, out of all L x 128 states, we randomly chose L x 10 states. The first choice represents full trajectory sampling, and the latter represents spare time steps sampling. We probe an FNO surrogate model trained on all the 128 trajectories its hidden layer, and observe the hidden layer activation each of the 128 states. The result is shown in Fig. 15, where black points represent states from the fully sampled trajectories and red points represent sparsely selected states. The latter states visibly more diverse, which partially explains how sampling time steps in sparse manner from trajectories can benefit surrogate model.
D.3. Random Bernoulli Sampling of Time Steps
We provide the whole list of results with Berouuli sampling described in 5.5. Also, enforce consecutive initial time steps sampling by bringing all the true entries in S to the beginning. We call this method Initial Bernoulli sampling, or Initial Ber(p). We report the results with SBAL in Table 13 and Table 14, Initial Bernoulli sampling always performs the worst, possibly because they rarely see the time steps at the end.
D.4. Time-dependent Incompressible Navier-Stokes
We have performed an experiment on a time-dependent incompressible Navier Stokes equation, simply by using the time- dependent external force in our current INS equation. The new forcing term is a sinusoidal mixture two spatial co- ordinates and the temporal coordinate. Fig. 16 shows the log RMSE of Random, SBAL, and SBAL+STAP, Since our methodology aligns closely to the time-dependent PDEs, our method significantly outperforms the others.
20

---

Active Learning with Selective Time-Step Acquisition for PDEs
Burgers 8L/8
KdV 8L/8
KS BL/8
KS L/2
INS L/2
CNS L/2
-1.0
1.8
1.4
4.00
-1.5
1.4
1.4
1.2
3.75
-2.0
12
1.2
1.0
3.50
1.0
3.25
0.8
-2.5
0.0
0.8
3.00
0
10
5 Round
10
5 Round
10
5
10
10
10
Round
Round
Round
Round
Random
SBAL
SBAL+STAP
Figure 13: Log RMSE of AL strategies with multi-step FNO, across 10 rounds of acquisition.
Burgers
KdV
KS
INS
CNS
1.0
3.00
-2
0.00
2.75
0.5
0.0
-0.25
-3
2.50
0.0
-0.50
-0.5
2.25
005
-0.75
-1.0
-1.00
-1.0
1.75
0
10 Round
20
0
10 Round
20
10 Round
20
0
10 Round
20
10 Round
20
Random
SBAL
SBAL+STAP
Figure 14: Log RMSE of AL strategies, measured across 20 rounds of acquisition.
E. Comparison of Costs with Our Numerical Solvers
Our paper does not claim direct computational speedups our benchmark PDEs; instead, it relies on the benchmark PDEs as proxies reflecting realistic, expensive simulations. Our surrogate metric, the number of numerical solver-simulated timesteps, effectively represents relative computational savings in realistic, expensive simulations.
However, we provide the analysis of computational cost on our benchmark PDEs for the sake of completeness. We compare the wall clock time of Random and SBAL+STAP, and for meaningful comparison, we cut the SBAL+STAP experiment when it reaches below the RMSE of the final Random surrogate model.
Table 15 compares the wall clock times between non-AL and AL, and Table 16 decomposes them into data acquisition, model training, and data selection. We find that AL reduces the total cost in CNS, where acquisition is relatively expensive. On other benchmarks, the training and data selection costs dominate, as expected. We want to stress yet again that the benchmarks were intentionally chosen to be inexpensive, to enable fast experimentation.
We provide the condition under which AL reduces total cost. Suppose AL improves data efficiency by E over non-AL. Define T acquire, Timann the acquisition time and training time per unit data, and Tretect as the data selection time per round. The total cost of non-AL is
(22 a Tacquire + N (1) train Turain
and for AL,
N (2)
acquire T Tcquire + train Terain + MT
where (0)
acquire are the number of acquired data, and (6) train are the total number of training examples (counting duplicates),
21

---

Active Learning with Selective Time-Step Acquisition for PDEs
10
(a) Burgers
(b) KdV
(c) KS
(d) INS
INS
1700
0
300
800
2000
(e) CNS
Figure 15: PCA of FNO hidden layer's activation pattern for both entire trajectories (black) and sparsely sampled time steps (red)
22

---

Active Learning with Selective Time-Step Acquisition for PDEs
Table 13: Bernoulli sampling
SBAL
+STAP
+Ber(1/16) +Ber(1/8)
+Ber(1/4)
+Ber(1/2)
Burgers
RMSE -0038820052
30.07420077
33.23120.100 2..15220.207
-3.10220.41
-3.45820.007
NRMSE -4.081 +0.061
-..006+0.00
-3.847 +0.171
-3.782+0.22 259
-3.734+0.447
-4.122+0.0m +0.078
MAE
-5.332+a.046
-5.585+0.066
-5.16940.104 25.085+0.272 -5.037+0.444
-5.397+0.062
Kdv
RMSE 0.030
-0.088+0.0au
0.049#a.014 0.053+v.014
0.018+v.024
-0.06440.001
NRMSE -1.282+a.u50 -1.378+0.0au
-1.254 -1.257#0.014 -1.288+0.uz
-1.370
MAE
-2.139+0.u27
20.23
-0.0824.u116
-20083+v.918
-2.120+0.uz
-2.207
KS
RMSE -0.275+0.014
-0.319
0.3655+0.0s 0.359
-0.446+0.00s
-0.32440.007
NRMSE -1.700+0.014 -1.774.0.003
-1.790+0.000 -..88440.00
-1.77140.00s
-1.749940.07
MAE
-2.18410.014
-2.26510.008
-2.28210.
-2.27610.006
-2.26210.00 009
-2.23710.000
INS
RMSE -0.422:0.010
-0.525+0.00s 005
- 1..08
-0.525
-0.515+0.007 -0.500
NRMSE -2.75110.012 -2.814
-2.818 10.008
-2.814.0.006
- 00T
-2.789
MAE
-2.736#0.009
-0.797
-2.794 t0.008
00.990
-2.781+0.005
-2.769
CNS
RMSE 2.422+0.045
2.363=#.016
2.375+0.069
2.370#0.01s 2.372+0.012
2.392+4.00
NRMSE -2.482+a.00
-2.501 +0.051
-2.540
-2.521+o.00 -2.226+0.024
-2.519
MAE -0.880+0.041 -0.93220.015
-0.92520.019 -0.921
-0.92320.012 -0.90620.010
and M the number of rounds. With initial datasize D and acquired datasize B per round,
(1) acquire
BM
(0 D + BM
(2)
BM/E
acquire
M/E
(2) train
Σ
(D+ + B rrund)
round=0
For AL to reduce the total cost, the setting would need satisfy
(1)) aequire Tacquire
(1) Irain Tiran
N-(2) T
Vxcquire + A-(2) Timin + MT
"uquiree
select
Table 17 lists these values, and whether they satisfy the condition above.
23

---

Active Learning with Selective Time-Step Acquisition for PDEs
Table 14: Initial Bernoulli sampling
Initial Ber(1/16) Initial Ber(1/8) Initial Ber(1/4) Initial Ber(1/2)
Burgers
RMSE
-3.73710.00 10.034
-3.6861100644
-3.63410.024
-0.05110.043
NRMSE
-4.36840.049
-4.32210.00
-4.274.0.027
-4.22210.056
MAE
-5.654+0.027
-5.607+0.06
-5.561=0.018
-5.483+0.031
KdV
RMSE
0.032:0.016
-0.015
-0.001
0.011
NRMSE
-1.278+0.011
-1.321+0.01
-1.303=0.018
-1.294+0.014
MAE
-2.15010.016
-2.19710.014
-2.18110.015
10.168
KS
RMSE
-0.302+0.00
-0.293+0 008
-0.287=0.005
-0.283+0.008
NRMSE
-1.72810.00s
-1.71910.
-1.71310.005
-1.708
MAE
-2.21640.008
-2.206 0.008
-2.199+0.007
-2.100
INS
RMSE
-0.28240.008
-0.294+0.010
-0.329 -0.397+0.008
NRMSE
-2.57510.000
-2.58710.010
-2.62120.005
-2.68810.008
MAE
-0.597+0.004
-2.609+0.000
-2.643=0.001
-2.703=00.00
CNS
RMSE
2.431:1.0017
2.420:10.009
2..8910.011
2.420
NRMSE
-2.49240.020
-0.493+0.011
-2.545
-2.501+0.00
MAE
-0.873+0.015
-0.883+0.00
-0.918+0.009
-0.886+0.008
Time-dependent INS
0.0
-0.2
-0.4
-0.6
10
Round
Random
SBAL
SBAL+STAP
Figure 16: Log RMSE of AL. strategies time-dependent INS.
Table 15: Total wall clock time until target RMSE (in seconds)
Equation Random
STAP
Burgers
237
441
KdV
834
1469
KS
385
2920
INS
526
3702
CNS
4053
3476
24

---

Active Learning with Selective Time-Step Acquisition for PDEs
Table 16: Total wall clock time decomposed into acpuistioon/raaiinn/ellllooooo (in seconds)
Equation
Random
SBAL+STAP
Burgers
90/147/0
27/234/180
KdV
670/164/0
455/664/350
KS
40/345/0
35/2075/810
INS
190/336/0
160/1750/1792
CNS
3570/483/0
1448/972/1056
Table 17: Variables for cost analysis
Equation
E
Tacquire
Timin
Twleet
Satisfied
Burgers
3.33
0.087
0.101
60
F
KdV
1.43
0.654
0.106
50
F
KS
1.11
0.005
0.116
90
F
INS
1.25
0.077
0.112
224
F
CNS
2.50
1.760
0.157
264
25


Reviews:
[{"review_id": "IFOy2jhCqo", "reviewer": "Reviewer_vQAV", "summary": "This paper develops an acquisition function that estimates the utility of a set of time steps and utilizes it for batch active learning in training surrogate models for PDEs. The empirical results show that the proposed method outperforms the baselines.", "questions": "", "limitations": "", "rating": 2, "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": "We appreciate your thoughtful questions and feedback.\n\n> Theoretical analysis of the acquisition function\n\nOur acquisition function is an approximation to the expected error reduction (EER), which is statistically near-optimal for active learning [1][2]. The EER measures how much the model’s generalization error is likely reduced after updating on hypothetically acquired data.  We model our hypothetical belief about the ground truth solver as a *uniform categorical distribution* over the ensemble $ \\\\{\\hat{G}\\_a\\\\}\\_{a=1}^M $. We assume that acquiring the trajectory of $u^0$ with sampling pattern $S$ only reduces generalization error on the trajectory of $u^0$. The current generalization error is expected to be the average of $  \\| \\hat u\\_a - \\hat u\\_b \\|^2  $ over $b$. We make a second assumption that the hypothetically acquired data $ \\hat{u}\\_{b,S,a} $ will update the model such that the model predicts the trajectory $ \\hat{u}\\_{b,S,a} $ given $u^0$. This gives us the expected reduction in error $ \\| \\hat u\\_a - \\hat u\\_{b} \\|^2 - \\| \\hat u\\_a - \\hat u\\_{b,S,a} \\|^2 $ averaged across $a$ and $b$, which is equal to our acquisition function. Although proving an optimality bound is outside of our expertise, we emphasize that our design of the acquisition function was guided by this exact theoretical consideration.\n\n> OOD data due to poor quality of surrogate models in early stages\n\nThank you for raising an important point. Please see Common Response 2 in our response to reviewer f7JP.\n\n> Performance metrics\n\nhttps://anonymous.4open.science/api/repo/icml_rebuttal-E9AB/file/diff.png?v=0041f607\n\nThe linked image shows the improvement in log RMSE over Random (no active learning). In the KS equation, SBAL+STAP improves over Random five times as much compared to SBAL.\n\n|Equation|Random|SBAL|SBAL+STAP|\n|-|-|-|-|\n|Burgers|0.1522|0.2277|0.2451|\n|KdV|0.0831|0.1052|0.1132|\n|KS|0.1001|0.1031|0.1094| \n|INS|0.0806|0.0869|0.0912|\n|CNS|0.0544|0.0795|0.0884|\n\nHowever, we agree that this quantity might be misleading to readers, and thus provide a table of the measure of data efficiency, defined as the average reduction in log RMSE per round. Overall, SBAL+STAP improves data efficiency by about 10 percent compared to SBAL, which is the SOTA algorithm. For a simulation that takes ten days to run, this would amount to saving a whole day. We encourage you to look at the results in Figure 4 of Musekamp et al. [3], Figure 3(b) of Li et al. [4] (the yellow and blue lines correspond to no active learning and active learning), and Figure 5 of Bajracharya et al. [5]. Our method provides arguably the largest and the most robust performance gain reported in the PDE active learning literature.\n\n> random time step selection outperforms STAP\n\nAs shown in Table 2, no single choice of $p$ for random time step selection outperforms full trajectory sampling on all PDE datasets, sometimes even degrading the performance. Unless there is a way to adaptively select $p$, random time step selection is not a viable active learning method, and was used in our paper solely for the purpose of analysis.\n\n> Experiments with more iterations\n\nhttps://anonymous.4open.science/api/repo/icml_rebuttal-E9AB/file/20_rounds.png?v=7ac767fc\n\nWe have performed the main experiment for 20 rounds instead of 10, on Random, SBAL, and SBAL+STAP. We observe that the gap between SBAL and SBAL+STAP keeps widening, except in the KdV equation.\n\n> Experiments with time-dependent PDEs\n\nhttps://anonymous.4open.science/api/repo/icml_rebuttal-E9AB/file/time_dependent_ins.png?v=f9c179ae\n\n|Equation|Random|SBAL|SBAL+STAP|\n|-|-|-|-|\n|Time-dependent INS|$-0.080\\pm 0.011$|$-0.081\\pm 0.013$|$-0.339\\pm 0.015$|\n\nWe have performed an experiment on a time-dependent incompressible Navier Stokes equation, simply by using the time-dependent external force in our current INS equation. The new forcing term is a sinusoidal mixture of two spatial coordinates and the temporal coordinate. The above figure and table summarize the log RMSE of Random, SBAL, and SBAL+STAP. **Our method aligns closely with time-dependent PDEs, explaining the massive gain in performance.**\n\nReferences:\n\n[1] Settles, Burr. \"Active learning literature survey.\" (2009). \\\n[2] Roy, Nicholas, and Andrew McCallum. \"Toward optimal active learning through sampling estimation of error reduction.\" ICML. Vol. 1. No. 3. 2001. \\\n[3] Musekamp, Daniel, et al. \"Active learning for neural pde solvers.\" arXiv preprint arXiv:2408.01536 (2024). \\\n[4] Shibo Li, Xin Yu, Wei Xing, Robert Kirby, Akil Narayan, and Shandian Zhe. Multi-resolution active learning of fourier neural operators. In International Conference on Artificial Intelligence and Statistics, pp. 2440–2448. PMLR, 2024. \\\n[5] Pradeep Bajracharya, Javier Quetzalcóatl Toledo-Marín, Geoffrey Fox, Shantenu Jha, and Linwei Wang. Feasibility study on active learning of smart surrogates for scientific simulations. arXiv preprint arXiv:2407.07674, 2024.", "claims_and_evidence": "The proposed algorithm relies on the heuristic of replacing simulations with predictions from the surrogate model for the skipped time steps. However, especially in the early stages, when the surrogate model may perform poorly, the sequence after the skipped time steps might not correlate with the initial condition u_0 at all. In such cases, why not query more diverse initial conditions and run fewer time steps for each scenario using the ground-truth simulator?", "methods_and_evaluation_criteria": "The evaluation metrics are well-established. For benchmarks, it is helpful to test whether the proposed methods can also be applied to PDEs where the evolution operator is time-dependent.", "theoretical_claims": "There is no theoretical guarantee provided for the proposed acquisition function. It will be helpful if the authors could include a theoretical analysis of their designed acquisition function. For instance, can they prove that the proposed acquisition function provides an optimal or near-optimal solution?", "experimental_designs_or_analyses": "I am not convinced by the authors' claim that their approach improves performance by up to five times compared to the baselines. For instance, in Figure 3, when comparing SBAL and SBAL+STAP, their performance appears similar. In most cases, SBAL is only one iteration behind SBAL+STAP. Additionally, it would be helpful if the authors can include more iterations to show the full performance until convergence. In Table 2, the results indicate that random time step selection outperforms STAP on 2 out of 5 tasks.", "relation_to_broader_scientific_literature": "There is a amount of work on batch active learning in various scenarios. This paper focuses on actively selecting the time stamps within each scenario to further improve the sample efficiency.", "essential_references_not_discussed": "The related work has been thoroughly discussed.", "other_strengths_and_weaknesses": "Please refer to the previous sections.", "other_comments_or_suggestions": "Please refer to the previous sections.", "questions_for_authors": "Please refer to the previous sections.", "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": "lDbVVpobkW", "reviewer": "Reviewer_FkHn", "summary": "The paper introduces an active learning framework STAP for surrogate modeling of PDE trajectories that selectively queries only key time steps instead of simulating entire trajectories. STAP uses a binary sampling pattern to decide which time steps to acquire via a numerical solver and which to approximate with a surrogate model. An acquisition function based on variance reduction guides the selection process, and a greedy batch acquisition algorithm is used to optimize the sampling patterns under a fixed-cost budget. Experiments on benchmark PDEs, including Burgers, KdV, Kuramoto-Sivashinsky, and Navier-Stokes equations, demonstrate that the proposed method reduces errors compared to existing active learning methods, particularly at higher quantiles.\n\n## update after rebuttal\n\nI thank the authors for the updated results. I appreciate the interesting idea presented in the paper. However, although aiming to reduce the cost of generating a PDE dataset, the method presented in this paper, in the end, makes the overall process of obtaining a surrogate model more expensive in most cases. I believe the paper can contribute validly to the field if the authors address this contradiction. I've decided to maintain my score since two specific issues regarding the evidence presented in this manuscript remain unaddressed.  \n\nTotal computational cost is larger: In the AI4PDE context, both data generation (numerical solvers) and model training/AL overhead consume computational resources. Therefore, demonstrating a reduction in total computational cost is crucial for establishing practical value. While the authors provided cost analysis upon request, the results indicated their method actually increased the total computational cost for most benchmarks presented. If the proposed AL method requires more total computation time or resources to reach a given accuracy on these problems, what is its compelling advantage over standard approaches? The paper needs to address this directly and convincingly, rather than solely relying on hypothetical scenarios.\n\nUnvalidated scalability assumption for larger-scale PDEs: The work's significance heavily relies on extrapolating its findings to more complex, high-dimensional PDEs where solver costs are presumed dominant. However, it does not consider whether the deep learning surrogate, along with the active learning strategy that depends on it, can effectively scale to these types of problems. This is a non-trivial assumption, especially in light of known limitations of deep learning, such as producing blurry predictions for very high-resolution data. Without supporting evidence, expectations regarding performance in complex regimes remain speculative.", "questions": "", "limitations": "", "rating": 2, "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": "We appreciate your thoughtful questions and feedback.\n\n> Concerns about the lack of a direct comparison of computing resources with non-active learning \n\nOur paper does not claim direct computational speedups on our benchmark PDEs; instead, it relies on the benchmark PDEs as proxies reflecting realistic, expensive simulations, a common approach in active learning studies such as AL4PDE [1].\nOur surrogate metric, the number of numerical solver-simulated timesteps, effectively represents relative computational savings in realistic, expensive simulations. A computing resources analysis would be potentially misleading due to the inherently simplified nature of our experimental datasets.\n\nhttps://anonymous.4open.science/api/repo/icml_rebuttal-E9AB/file/taylor-green.png?v=50ac6a75\n\nMany real-world PDEs inherently demand significant computational resources for numerical simulations [2][3][4]. For instance, the Taylor-Green vortex benchmark reported in [2]—standardized through the \"taubench\" work unit measure—requires between $10^4$ and $10^6$ work units per trajectory, translating to roughly 41 CPU-hours to over 4,100 CPU-hours on Dual 20-Core Intel Xeon processors (Gold 6148) [5].\nThis benchmark was conducted as a competition, in which multiple teams submitted numerical solvers. Each solver could run at various fidelity levels, with lower fidelity solutions requiring fewer computational resources at the cost of increased error tolerance. As shown in the image attached above (Figure 25(b) of [2]), even the fastest solution with the highest allowed error tolerance required approximately $10^4$ work units.\nThis substantial computational burden highlights the importance and practicality of active learning for PDEs.\n\n> Experiments with multi-step neural solvers\n\nThank you for raising this important point. Please see Common Response 1 at the bottom.\n\n> model-specific-ness of the method\n\nOur method STAP is model-agnostic, as it is applicable to any model that predicts the next PDE state given the current state (e.g. UNet). As you point out, STAP will likely select different timesteps when used with a different model. This is the desired behavior of any active learning algorithm, since its goal is to select data that are most useful to the current model at hand. Could you elaborate on why one would want an active learning strategy that selects the same data regardless of the model?\n\n> OOD data caused by surrogate models that potentially deviate from correct physical behavior\n\nThank you for raising an important point. Please see Common Response 2 in our response to reviewer f7JP.\n\n**Common Response 1. Multi-step model**\n\nhttps://anonymous.4open.science/api/repo/icml_rebuttal-E9AB/file/multistep.png?v=ab3f6da0\n\nWe have done experiments with multi-step FNO which receives N timesteps as input and outputs N timesteps. To perform STAP, we group the total number of timesteps into non-overlapping clusters of N timesteps, and perform STAP as if each cluster is one timestep with N channels.\nWe have performed two variants of experiments. In the first, we divide a timestep in our main experiment into 8 smaller timesteps, so that a total of $L$ timesteps turn into $8 L$ timesteps, and train 8-in-8-out models. In the second variant, we keep $L$ timesteps the same but train 2-in-2-out models. The figure above shows the log RMSE for Burgers, KdV, and KS of the first variant, and KS, INS and CNS of the second variant.\n\n|Equation|Random|SBAL|SBAL+STAP|\n|-|-|-|-|\n|Burgers 8L/8|$-1.670\\pm 0.0982$|$-1.893\\pm 0.053$|$-2.058\\pm 0.028$|\n|KdV 8L/8|$1.402\\pm 0.029$|$1.404\\pm 0.024$|$1.364\\pm 0.043$|\n|KS 8L/8|$1.255\\pm 0.015$|$1.232\\pm 0.012$|$1.156\\pm 0.008$|\n|KS L/2|$1.340\\pm 0.014$|$1.335\\pm 0.011$|$1.288\\pm 0.011$|\n|INS L/2|$1.124\\pm 0.017$|$1.118\\pm 0.007$|$1.081\\pm 0.012$|\n|CNS L/2|$3.593\\pm 0.023$|$3.594\\pm 0.044$|$3.42\\pm 0.050$| \n\nThe table above summarizes our results with mean log RMSE.\n\nReferences:\n\n[1] Musekamp, Daniel, et al. \"Active learning for neural pde solvers.\" arXiv preprint arXiv:2408.01536 (2024). \\\n[2] Wang, Q., Fidkowski, K., Abgrall, R., Bassi, F., Caraeni, D., Cary, A., ... & Olivier, H. (2012). “High-order CFD methods: current status and perspective.” International Journal for Numerical Methods in Fluids, 93(4), 212–232. \\\n[3] Kaneda, Y., Ishihara, T., Yokokawa, M., Itakura, K., & Uno, A. (2003). “Energy dissipation rate and energy spectrum in high resolution direct numerical simulations of turbulence in a periodic box.” Physics of Fluids, 15(2), L21–L24. \\\n[4] Heil, M. & Hazel, A. L. (2011). “Fluid–Structure Interaction in Internal Physiological Flows.” Annual Review of Fluid Mechanics, 43, 141–162.\\\n[5] Capuano, F., Beratis, N., Zhang,F., Peet, Y., Squires, K., & Balaras., E. (2023). Cost vs Accuracy: DNS of turbulent flow over a sphere using structured immersed-boundary, unstructured finite-volume, and spectral-element methods. Eur. J. Mech. B Fluids, 102:91–102.", "claims_and_evidence": "The paper claims that selective time-step acquisition using active learning improves the efficiency and accuracy of PDE surrogate modeling by reducing the expensive computation of the numerical solver. While the experimental results provide clear evidence that the proposed method outperforms other active learning methods, no cost of time and computing resources analysis against direct non-active learning are provided. The active learning in this work consists of training an ensemble of FNOs over multiple rounds. It is uncertain whether this approach saves time or computing resources compared to training a single FNO for one run on the full trajectory.\n\nThe paper also claims the cost of data acquisition from numerical solvers is extremely high. This is not well supported. According to a recent article [1], it's common for the research in the machine learning community to compare to a numerical method that is much less efficient than a SOAT method for that problem. Some of the PDE datasets are merely solved with fundamental, manually implemented Python code instead of SOAT algorithm or well-optimized software. The authors should justify that the numerical solvers used in this paper are not necessarily the best but are at least reasonably close to the SOAT solution.\n\nAlso, numerical algorithms can adjust the solving iterations to produce less accurate solutions with faster speed. As saving the cost of generating datasets for training surrogate models is the main goal of this work, it's advisable to test generating less accurate training samples from numerical solvers and compare the final accuracy and the time and computing resources cost. \n\n\n[1] McGreivy, Nick, and Ammar Hakim. \"Weak baselines and reporting biases lead to overoptimism in machine learning for fluid-related partial differential equations.\" Nature Machine Intelligence 6.10 (2024): 1256-1269.", "methods_and_evaluation_criteria": "The proposed methods and evaluation criteria make sense for the problem of surrogate modeling for PDEs.", "theoretical_claims": "No proof is provided in the paper.", "experimental_designs_or_analyses": "The experimental design was generally sound. However, one concern is that the FNO is trained using a one-input, one-output autoregressive scheme. Many previous works use a multi-step input-output approach (e.g., 10-in-10-out). While the one-step approach might achieve higher accuracy, it sacrifices prediction speed by a factor of 10. A key efficiency advantage of machine learning models over numerical solvers is they can make multi-step or skip-step predictions. The paper does not evaluate whether the proposed framework remains effective under a common multi-step configuration.", "relation_to_broader_scientific_literature": "The paper is related to AI for PDE and AI for CFD research. No apparent relevance to broader scientific literature.", "essential_references_not_discussed": "No essential references I can think of.", "other_strengths_and_weaknesses": "The concept of using selective time steps in data is interesting. The visualization of sample diversity is clear and effectively demonstrates the motivation.", "other_comments_or_suggestions": "I suggest that the authors include detailed information on computing resources and time costs for (1) the numerical solver, (2) the entire active learning procedure, and (3) the non-active learning approach. This will help readers determine whether to use active learning in their specific situations. For example, some may have a few GPUs but enough CPU cores and numerical algorithms that can effectively utilize these cores in parallel.", "questions_for_authors": "The data selection method in the framework appears to be model-specific. If FNO were replaced with a different model, given the different inductive biases, would the selected time steps change? Is it possible to develop a data selection strategy that is not model-specific so that one can compare different machine learning models without requiring each to perform its own active learning process? \n\nIn my experience, for complex physical problems, predictions from models like FNO often do not fully satisfy physical constraints. Even if these predictions are fed back into the solver, it's hard to reduce the physical residuals for just a single-step solving. Will this lead to a \"garbage in, garbage out\" scenario where, most of the time steps selected by active learning don't satisfy physics?", "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": "qBq04RnYPO", "reviewer": "Reviewer_f7JP", "summary": "This paper introduces a novel active learning (AL) framework called Selective Time-Step Acquisition for PDEs (STAP) to improve the efficiency of surrogate models for partial differential equations (PDEs). The key idea is to selectively acquire only the most informative time steps from PDE trajectories using a numerical solver while the surrogate model approximates the remaining steps. This contrasts with existing AL methods that acquire entire PDE trajectories, reducing computational costs per trajectory and allowing for exploring a more diverse set of trajectories within a given budget. The authors develop an acquisition function that estimates the utility of a set of time steps by approximating the resulting variance reduction. They demonstrate the effectiveness of STAP on benchmark PDEs like the Burgers', Korteweg-De Vries, Kuramoto-Sivashinsky, incompressible Navier-Stokes, and compressible Navier-Stokes equations. Results show that STAP significantly improves performance, reducing average error and error quantiles compared to existing methods. The method combines a numerical solver and a surrogate model to acquire data along a trajectory with reduced cost, improving sample efficiency. STAP can be seen as an add-on to existing AL methods that acquire full trajectories. The authors also explore efficient variants of STAP to reduce computational cost, and show that the success of STAP is driven by its ability to prioritize both diverse and informative time steps. The findings suggest STAP offers a more cost-efficient and accurate solution for PDE surrogate modeling with broader applicability in scientific and engineering simulations.", "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": "We appreciate your thoughtful questions and feedback.\n\n> Prone to diverging models, producing out-of-distribution inputs\n\nThank you for raising this important point. Please see Common Response 2 at the bottom.\n\n> Does not support autoregressive training techniques\n\nWe have done additional experiments with multi-step models as in [1]. Please see Common Response 1 in our response to reviewer FkHn.\n\n> What does the +- / shaded area show exactly?\n\nWe used one standard deviation. You are correct to point out that we underestimated the standard deviation. We will rerun the experiments to obtain the correct standard deviations with prediction error of all ensemble members across all seeds.\n\nhttps://anonymous.4open.science/api/repo/icml_rebuttal-E9AB/file/correct_std.png?v=273fbf72\n\nThe above figure contains the corrected log RMSE figure for some main experiments.\n\n> Retraining == ?\n\nWe are training from scratch like in most of the active learning literature [1].\n\n> Parameters of IC generators\n\nLet $U$ stand for the uniform distribution. For Burgers and KdV, the initial condition is in the form $\\sum_{i=1}^N A_i \\sin(2\\pi k_i x / L + \\phi_i)$. The amplitudes and phases are always sampled from $U([0,1])$ and $U([0,2\\pi])$. For Burgers, $N = 2$ and $k_i \\sim U(\\\\{1,2,3,4\\\\})$, and for KdV, $N=10$ and $k_i \\sim U(\\\\{1,2,3\\\\})$. For KS and INS, the initial conditions are gaussian random fields drawn from $N(0, 25(-\\Delta + 25I)^{-1})$ and $N(0, 7^{3/2}(-\\Delta + 49I)^{-2.5})$ respectively. For CNS, we first sample $\\sum_{k \\in \\\\{1,2,3,4\\\\}^3} A_k \\sin( 2\\pi k x/L + \\phi_k)$, the amplitudes and phases sampled uniformly at random for each channel ($\\rho$, $p$ and $\\mathbf{v}$). Then, we renormalize $\\rho$, $p$ to lie within $\\rho_0(1 \\pm \\Delta_\\rho)$ and $p_0(1 \\pm \\Delta_p)$, respectively, where $\\rho_0 \\sim U([0.1,10])$, $\\Delta_\\rho \\sim U([0.013, 0.26])$, $\\Delta_p \\sim U([0.04,0.8])$ and $T_0 := \\rho_0/p_0 \\sim U([0.1,10])$. The $\\mathbf{v}$ is also computed by superposing sinusoidal waves, but with amplitudes chosen so that the initial condition has the given initial Mach number $M \\sim U([0.1,1])$.\n\n> Push-forward trick\n\nWe agree that the errors seem too high. If we can’t fix the problem by the deadline, we will remove the push-forward experiment.\n\n**Common Response 2. Out-of-distribution Synthetic Inputs**\n\nhttps://anonymous.4open.science/r/icml_rebuttal-E9AB/gt_pred/ \\\nhttps://anonymous.4open.science/api/repo/icml_rebuttal-E9AB/file/kdv_energy.png?v=5d240a44\n\nReviewers have raised the concern that inaccurate surrogate models might synthesize inputs that lie far from the ground truth distribution, reducing their information gain. In fact, under limited training data, the surrogate model outputs visibly erroneous trajectories. The images in the folder of the first link are comparisons of the ground truth trajectories and predictions from surrogate models trained on 32 and 1 trajectories, respectively. The second link plots the energy of KdV states in both ground truth and predicted trajectories. We see that the surrogate model doesn't satisfy conservation of energy.\n\nhttps://anonymous.4open.science/api/repo/icml_rebuttal-E9AB/file/one_initial_train.png?v=1d82126d \\\nhttps://anonymous.4open.science/r/icml_rebuttal-E9AB/stap_pca/\n\nTo test how much this error harms STAP, we perform experiments where the initial training dataset contains 1 trajectory, compared to 32 used in our main experiments. The first link above shows the log RMSEs. To our own surprise, we find that SBAL+STAP still outperforms Random and SBAL, except for INS in the early rounds.\nThe second link is a folder that contains comparisons of FNO activation PCAs for PDE states sampled by SBAL and SBAL+STAP during the first round of active learning. Note that the PCA was fitted only to the ground truth states in random trajectories (blue points), so that the out-of-distributionness of sampled states can be properly reflected. We find that only several of the states sampled by SBAL+STAP diverge significantly from the random ground truth states. In other words, the surrogate model synthesizes erroneous inputs when looked at trajectory-wise, but they aren't necessarily out-of-distribution, hence retaining the information gain.\n\nReviewer vQAV also asked, \"why not query more diverse initial conditions and run fewer time steps\" in the earlier rounds where the surrogate model is inaccurate. The problem is that the distribution of states $u_t$ changes over time $t$. Running only up to the first few timesteps harms the model’s performance on the later timesteps, as evidenced in Appendix D.3. STAP seems to be striking the balance between sampling realistic inputs and sampling diverse timesteps $t$. How one could further improve this balance is left as a work for future research.\n\nReferences:\n\n[1] Musekamp, Daniel, et al. \"Active learning for neural pde solvers.\" arXiv preprint arXiv:2408.01536 (2024).", "claims_and_evidence": "Claim: STAP improves surrogate model performance over previous active learning methods for PDEs.\n\nEvidence: Experiments on Burgers', KdV, KS, INS, and CNS equations demonstrate that SBAL+STAP outperforms other AL baselines regarding Log RMSE (Figure 3, Table 1). \"Ratio of ∆\" values in Table 1 quantifies the improvement of SBAL+STAP over Random selection. These ratios are greater than 1 for all tested PDEs, indicating STAP provides a performance gain. Figure 4 shows that STAP improves performance on points with extreme errors (99% and 95% quantiles) and even reduces error in the middle quantiles (50%), which is rare for AL algorithms. Appendix C provides a full report of all metrics on all methods\n\nClaim: STAP achieves better performance by adaptively choosing both the frequency and specific locations of time steps to acquire.\n\nEvidence: Figure 5 shows the distribution of time steps chosen by STAP, demonstrating a tendency to acquire early time steps, with occasional selection of later time steps varying across different PDEs.  Analysis of the frequency with which STAP samples time steps reveals that it doesn't always sample time steps at an optimal frequency p to sample with. This suggests that the specific time steps sampled also matter as much as the overall frequency.\n\nClaim: STAP can be implemented efficiently without significant performance loss.\n\nEvidence: The use of two efficient variants of STAP, namely STAP MF and STAP 10, shown in Table 14.\nTable 3 summarizes the wall-clock time of each baseline method and STAP; STAP 10 incurs only a fraction of computational cost over the baseline methods.", "methods_and_evaluation_criteria": "The paper uses a well-chosen set of PDE benchmarks from the Al4PDE active learning benchmark, including Burgers', KdV, KS, INS, and CNS equations, which represent a variety of physical phenomena and complexity levels, making for a robust and generalizable evaluation of STAP. It is great that the authors use the existing AL benchmarks and it would be great if they could also publish their method on github and integrate it there with the existing framework.", "theoretical_claims": "There are no theoretical claims made in the paper.", "experimental_designs_or_analyses": "The experiments are well-designed and in line with prior work by Musekamp et al. but limited to one-step predictions.", "relation_to_broader_scientific_literature": "Active learning for ML-based PDE solvers is a recent development, with only a few recent papers having addressed this important problem. The authors make strong contributions by showing that time steps of the trajectories (rollouts) can be chosen adaptively and selectively. Compared to prior work, this is a significant advancement over the state of the art.", "essential_references_not_discussed": "None.", "other_strengths_and_weaknesses": "Strengths\nI really like the PCA visualizations of the effect. It is impressive that the methods show improvements over random sampling, even for problems where no other AL method showed an improvement!  The evaluation is strong and on multiple PDEs from AL4PDE.\n\n\nWeakness\nThe method could be prone to diverging models, producing out-of-distribution inputs to the simulator. There is no check of the inputs in terms of the uncertainty.\n\nOnly applicable to 1-step training (predicting t+1 from t), does not support autoregressive training techniques.\n\nA push-forward experiment (D.2) is a good idea to check. But the numbers seem wrong; it is hard to believe that the push-forward trick worsens things. \n\nThe connection to PDE dynamics / autoregressive behavior could have been discussed more. For example, your algorithm likely selects earlier time steps since improvements, in the beginning, will affect the rest of the rollout.  In the autoregressive setting, one might not want to select only the first steps since it will take some time for the model to diverge again and since they might be too close.", "other_comments_or_suggestions": "None.", "questions_for_authors": "What does the +- / shaded area show exactly? Confidence interval? 1 time standard deviation?  The standard deviation was probably underestimated since they took the average error of the ensemble members (ln 848).\n\nRetraining == training from scratch or fine-tuning?\n\nWhat are the parameters of the IC generator (distribution)?", "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": "bFBFfThCNz", "reviewer": "Reviewer_K9ac", "summary": "This paper introduced an active learning method for learning PDEs. The method is composed of: (1) selective time-step acquisition, where the method selects a subset of time steps for the solver to simulate while other time steps are evolved by the surrogate model, (2) an acquisition function that evaluates the variance reduction, and (3) batch acquisition. The method is evaluated on 5 PDEs including incompressible and compressible NS, KS, Burgers, and KdV equations. The method shows consistent improvement compared with other baselines.", "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": "We appreciate your thoughtful questions and feedback.\n\n> I wonder if the patterns shown in Figure 5 are consistent among different experiment runs (with different seeds)? The authors are encouraged to run multiple independent experiments, showing their patterns.\n\nhttps://anonymous.4open.science/r/icml_rebuttal-E9AB/timesteps/\n\nThank you for raising an important point. The linked folder \"timesteps\" contains sampled timesteps for four other seeds for each of the PDEs. We see that the patterns described in our paper are consistent among different seeds per task.\n\n> The method requires interleaving solver evolution with surrogate evolution, which, in my own experience, is typically engineering challenging.\n\nWe appreciate the concern about interleaving solver-based and surrogate-based time steps. In our current codebase, the only required addition is calling one step of the external solver from Python. In practice, most PDE solvers can evolve a state one step at a time through a simple function call or by using “checkpoint” files between time steps, which Python can read and write. Although some solvers are specialized or written in lower-level languages, the overhead incurred by periodic file-based communication is generally much smaller than the total cost of running a full solver trajectory, especially for problems where the solver is very expensive. By selectively acquiring just the most important steps, our approach saves substantial solver calls and overall runtime. We believe the net cost reduction outweighs the implementation complexity in most scenarios where computational demands are high. We will also make sure that our codebase is generally easy to adopt for user-defined PDE solvers.", "claims_and_evidence": "Yes. The claims that the selective time-step acquisition improves upon full trajectory acquisition is supported by convincing evidence.", "methods_and_evaluation_criteria": "Yes, they make sense", "theoretical_claims": "The paper does not have theoretical claims.", "experimental_designs_or_analyses": "The experimental design is sound. Experiments are done on incompressible and compressible NS equations which are important evaluations. Also, experiments are done comparing the method with random Bernoulli sampling, which shows that although the proposed method is on par with best Bernoulli sampling result.\n\nI wonder if the patterns shown in Figure 5 are consistent among different experiment runs (with different seeds)? The authors are encouraged to run multiple independent experiments, showing their patterns.", "relation_to_broader_scientific_literature": "The paper addresses an important problem in learning PDEs, where the time cost for running the solver is high.", "essential_references_not_discussed": "Not that I know of.", "other_strengths_and_weaknesses": "Strengths:\n\nThe paper is written in a clear way and easy to understand. The novelty is reasonable. In terms of significance, I believe the paper addresses an important problem that can be widely used in learning PDEs.\n\nWeaknesses:\n\nThe thing I worry about most is the engineering difficulty of applying the method. The method requires interleaving solver evolution with surrogate evolution, which, in my own experience, is typically engineering challenging. For example, if the solver is a well-established software which is written in a different language, it can be very hard to incorporate it with a neural network. In those case, it is probably easier if the solver runs the full trajectory. Therefore, in practice, I wonder whether the gain by selective time-step acquisition is worth the effort. Nevertheless, when incorporating the two component is not too hard, the method may be useful.\n\n## update after rebuttal:\nI have read the authors' rebuttal and mostly resolved my concerns. I remain my score.", "other_comments_or_suggestions": "N/A", "questions_for_authors": "N/A", "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: This paper introduces an active learning framework that improves sample efficiency in PDE surrogate modeling by selectively acquiring only the most informative time steps rather than the full trajectories. The method introduces a variance-reduction-based acquisition function and is evaluated on several PDE benchmarks. Results demonstrate that the method can be used as a promising add-on to existing active learning strategies.

The reviewers acknowledged the novelty and practicality of the proposed approach. They also appreciated the clear writing and thorough experiments. However, all reviewers raised concerns that should be addressed before the camera-ready version. These include the extent to which the method actually outperform existing methods, the lack of theoretical explanation for the acquisition function, and an explanation of the overall computational savings.

Given the reviews, I am on the fence about this submission and am leaning towards acceptance.
Decision: Accept (poster)