Active Learning with Selective Time-Step Acquisition for PDEs

We appreciate your thoughtful questions and feedback.

Prone to diverging models, producing out-of-distribution inputs

Thank you for raising this important point. Please see Common Response 2 at the bottom.

Does not support autoregressive training techniques

We have done additional experiments with multi-step models as in [1]. Please see Common Response 1 in our response to reviewer FkHn.

What does the +- / shaded area show exactly?

We 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.

https://anonymous.4open.science/api/repo/icml_rebuttal-E9AB/file/correct_std.png?v=273fbf72

The above figure contains the corrected log RMSE figure for some main experiments.

Retraining == ?

We are training from scratch like in most of the active learning literature [1].

Parameters of IC generators

Let $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])$.

Push-forward trick

We agree that the errors seem too high. If we can’t fix the problem by the deadline, we will remove the push-forward experiment.

Common Response 2. Out-of-distribution Synthetic Inputs

https://anonymous.4open.science/r/icml_rebuttal-E9AB/gt_pred/
https://anonymous.4open.science/api/repo/icml_rebuttal-E9AB/file/kdv_energy.png?v=5d240a44

Reviewers 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.

https://anonymous.4open.science/api/repo/icml_rebuttal-E9AB/file/one_initial_train.png?v=1d82126d
https://anonymous.4open.science/r/icml_rebuttal-E9AB/stap_pca/

To 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. The 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.

Reviewer 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.

References:

[1] Musekamp, Daniel, et al. "Active learning for neural pde solvers." arXiv preprint arXiv:2408.01536 (2024).

We appreciate your thoughtful questions and feedback.

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.

https://anonymous.4open.science/r/icml_rebuttal-E9AB/timesteps/

Thank 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.

The method requires interleaving solver evolution with surrogate evolution, which, in my own experience, is typically engineering challenging.

We 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.

We appreciate your thoughtful questions and feedback.

Concerns about the lack of a direct comparison of computing resources with non-active learning

Our 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]. Our 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.

https://anonymous.4open.science/api/repo/icml_rebuttal-E9AB/file/taylor-green.png?v=50ac6a75

Many 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]. This 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. This substantial computational burden highlights the importance and practicality of active learning for PDEs.

Experiments with multi-step neural solvers

Thank you for raising this important point. Please see Common Response 1 at the bottom.

model-specific-ness of the method

Our 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?

OOD data caused by surrogate models that potentially deviate from correct physical behavior

Thank you for raising an important point. Please see Common Response 2 in our response to reviewer f7JP.

Common Response 1. Multi-step model

https://anonymous.4open.science/api/repo/icml_rebuttal-E9AB/file/multistep.png?v=ab3f6da0

We 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. We 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.

The table above summarizes our results with mean log RMSE.

References:

[1] Musekamp, Daniel, et al. "Active learning for neural pde solvers." arXiv preprint arXiv:2408.01536 (2024).
[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.
[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.
[4] Heil, M. & Hazel, A. L. (2011). “Fluid–Structure Interaction in Internal Physiological Flows.” Annual Review of Fluid Mechanics, 43, 141–162.
[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.

We appreciate your thoughtful questions and feedback.

Theoretical analysis of the acquisition function

Our 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.

OOD data due to poor quality of surrogate models in early stages

Thank you for raising an important point. Please see Common Response 2 in our response to reviewer f7JP.

Performance metrics

https://anonymous.4open.science/api/repo/icml_rebuttal-E9AB/file/diff.png?v=0041f607

The 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.

However, 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.

random time step selection outperforms STAP

As 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.

Experiments with more iterations

https://anonymous.4open.science/api/repo/icml_rebuttal-E9AB/file/20_rounds.png?v=7ac767fc

We 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.

Experiments with time-dependent PDEs

https://anonymous.4open.science/api/repo/icml_rebuttal-E9AB/file/time_dependent_ins.png?v=f9c179ae

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 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.

References:

[1] Settles, Burr. "Active learning literature survey." (2009).
[2] Roy, Nicholas, and Andrew McCallum. "Toward optimal active learning through sampling estimation of error reduction." ICML. Vol. 1. No. 3. 2001.
[3] Musekamp, Daniel, et al. "Active learning for neural pde solvers." arXiv preprint arXiv:2408.01536 (2024).
[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.
[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.