Title: A Plug-and-Play Query Synthesis Active Learning Framework for Neural PDE Solvers
Venue: NeurIPS 2025
Accepted: yes
Keywords: active learning

neural PDE solver

query synthesis

expected information gain
Abstract: In recent developments in scientific machine learning (SciML), neural surrogate solvers for partial differential equations (PDEs) have become powerful tools for accelerating scientific computation for various science and engineering applications. However, training neural PDE solvers often demands a large amount of high-fidelity PDE simulation data, which are expensive to generate. Active learning (AL) offers a promising solution by adaptively selecting training data from the PDE settings--including parameters, initial and boundary conditions--that are expected to be most informative to help reduce this data burden. In this work, we introduce PaPQS, a Plug-and-Play Query Synthesis AL framework that synthesizes informative PDE settings directly in the continuous design space. PaPQS optimizes the Expected Information Gain (EIG) while encouraging batch diversity, enabling model-aware exploration of the design space via backpropagation through the neural PDE solution trajectories. The framework is applicable to general PDE systems and surrogate architectures, and can be seamlessly integrated with existing AL strategies. Extensive experiments across different PDE systems demonstrate that our AL framework, PaPQS, consistently improves sample efficiency over existing AL baselines.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe work addresses the learning of an autoregressive model for time series prediction in a context where data generation is expensive. The paper investigates learning tasks on physical time series as a representative example of this problem, where a simulator generates temporal trajectories with a natural trade-off between simulation quality and speed. To improve learning efficiency, the paper proposes an active learning approach for selecting initial conditions—not randomly, but in a way that ensures the resulting trajectories most benefit the model. \n\nSpecifically, the selected initial conditions are those expected to induce the greatest change in the current model. This approach is formalized within a probabilistic Bayesian framework. A key challenge lies in identifying, in advance and without time evolution, which initial conditions are likely to contribute the most information to the existing set of trajectories. To address this, a second model, referred to as the critic, is introduced. The critic does not require autoregressive rollout and can also provide estimates when the likelihood function of the autoregressive model is inaccessible. The entire system is jointly optimized, including the parameters of the autoregressive model, the initial conditions for new trajectories, and the critic’s parameters. The method is then evaluated through a series of experiments, including several differential equations, various active learning strategies, and different model architectures.\n\nSTRENGTHS AND WEAKNESSES:\nI am uncertain about the novelty of the proposed approach and the use of critic networks in active learning. Critic networks are common in other areas to process information over time more effectively, such as reinforcement learning.\n\nMy main concern is that the overall setup seems overly complex, involving three nested non-convex optimization problems, each sensitive to initialization and hyperparameters. Even with extensive experimentation, I find it difficult to draw a conclusion about general applicability here.\n\nThe authors argue (lines 179–182) that small perturbations can lead to spikes in predictive variance due to the chaotic nature of the system and the autoregressive structure, thereby justifying their more complicated approach. But more generally, it is questionable if autoregressive setups or application of learning to chaotic system is a suitable task for approximate machine learning, because of that sensitivity. One could just as easily argue that the expected information gain is misled by small perturbations of a labeled initial condition that leads to a completely different trajectory and so drastically changes the model? \n\nI encourage the authors to explore theoretical approaches to better understand and justify their method. As noted, a purely empirical evaluation seems of limited value in this context. However, when combined with a solid theoretical analysis, a more compelling case could be made.\n\nStrengths\n\n- Relevant topic and direction of research\n\n- Many experiments and a detailed statistical evaluation\n\nWeaknesses:\n\n- Highly complex non-convex optimization\n\n- The introduction of the critic network feels abrupt—it appears in the middle of the technical derivation without sufficient context. As mentioned, the central challenge is to estimate the usefulness of a specific initial condition without evolving the system, and this is addressed using the critic. This connection should be communicated more clearly. Both the abstract and introduction should mention that a second neural network is involved and plays a key role in the method.\n\n- In my view, the derivation from the expected information gain (Equation 11) to the update equations (Equation 13) is difficult to follow. The text primarily describes what is being done, but the reasons behind the design choices are not sufficiently explained. Any additional intuition or explanation would be helpful here. It might also improve clarity to restructure Sections 3.1  (optimizer and loss term 2/regularization term)  and 3.2 (loss term 1/EIG and update equations). Also, devoting half a page to the Adam optimizer seems unnecessary, as it is a standard choice.\n\n- Equation 4, in my opinion, does not provide much insight, and the mathematical notation is somewhat confusing.\n\n- Regarding Claim 1 and Proposition 1, the statements are not mathematically precise. For example: “Assume that the dominant sources of epistemic uncertainty of the model can be captured by θ_out” is vague and would benefit from clarification. Additionally, the final step in the proof (Equation 18) should be explicitly referenced in the appendix as well. Since this step appears to be crucial, an intuitive explanation of the identity in the main text could help readers better understand the overall argument.\n\nQUESTIONS:\n- L57: A sentence in the paper ‘By explicitly targeting the reduction of the posterior uncertainty, our framework dynamically explores high-informative regions while avoiding instability-prone area’. Can you explain how to draw this conclusion? I wonder if it’s not exactly the opposite that’s the case. Shouldn’t an unstable trajectory maximize the expected information gain as processing an unstable trajectory would result in a significantly different model?\n\n- Are there any physics-specific components or assumptions in the method? If not, should the approach also be evaluated on non-physical time series tasks to better demonstrate its generality?\n\n- How does the method perform on stochastic time series?\n\nSuggestions:\n\n- As pointed out, a theoretical justification could really strengthen the paper. A study on a linear system (linear physics and linear model) may allow for a closed-form solution, which could reveal if the proposed method provides a lower variance estimator for the model parameters. \n\n- Another empirical test could be to visualize the optimization landscape. Most likely, the optimization landscape provided by the critic is smoother, especially in the beginning. Similarly, the spectrum of the Hessian indicates how fast gradient-based optimization finds solutions. This could be another strategy to gain some insights on the proposed method.\n\nMinor:\nL98: MSE loss, not RMSE\nL212: typo ‘neural’\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 3, "CLARITY": 2, "IMPACT": 3, "ORIGINALITY": 2, "DATE": "2025-06-27"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper presents PaPQS, a query synthesis approach to sample datasets for neural PDE solvers. The method uses the mine-f method to estimate the expected information gain and SWA-Gaussian to sample the posterior distribution of model parameters. Combined with a distance-based metric for ensuring diversity, the acquisition function can be optimized using gradient ascent. The experiments show that the method can improve the samples generated by other active learning strategies.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths\n1. A novel method that can be added on top of every neural PDE solver and is faster than estimating uncertainty using ensembles.\n2. One of the rare works of AL that explores the query synthesis setting.\n3. Extensive experiments with confidence intervals.\n\nWeaknesses\n1. The method needs good initialization by a different AL method. If being optimized from scratch, it does not reach the performance of other AL strategies. \n2. The discussion of related work could be improved. For example, [1] and [2]  could be mentioned as applications of AL for neural PDE solvers.  In ln. 46-48: “feature-based acquisition methods emphasize input diversity but do not account for the relationship between PDE settings and model parameters” is not substantiated. For the first part, the feature-based acquisition function focuses on the coverage in feature-space and not in the inputs, and is also able to consider informativeness, for example [3]. Could you explain the second part about the relationship between PDE settings and model parameters? Don’t feature-based methods also use the gradient w.r.t. to the parameters as the feature representation [3]?  Lastly, PDERefiner (reference 11 in the paper)  is cited in ln. 34 as an example for AL, but this paper does not discuss the matter. \n\n[1] Li, S., Yu, X., Xing, W., Kirby, R., Narayan, A., & Zhe, S. (2024, April). Multi-resolution active learning of Fourier neural operators. In International Conference on Artificial Intelligence and Statistics (pp. 2440-2448). PMLR.  \n[2] Kim, Y., Kim, H., Ko, G., & Lee, J. Active Learning with Selective Time-Step Acquisition for PDEs. In Forty-second International Conference on Machine Learning.  \n[3] Holzmüller, D., Zaverkin, V., Kästner, J., & Steinwart, I. (2023). A framework and benchmark for deep batch active learning for regression. Journal of Machine Learning Research, 24(164), 1-81.\n\nQUESTIONS:\n1. For combining SBAL and PaPQS, did you use EIG also as the acquisition function of SBAL or the original ensemble uncertainty?\n2. Why did you apply the critic T to individual states, and not the whole trajectory at once? Will this formulation still be able to consider dynamics in the information gain?\n3. Do you optimize the total number of samples to select in the given AL iteration all at once, or do you optimize within mini-batches?\nIn Eq. 8, did you use the distance in the initial parameter space and PDE parameters? If so, what distance measure did you use?\n4. Doesn’t the acceptance of only positive steps (Eq. 9) encourage getting stuck in local minima, which may have been prevented by the momentum of Adam otherwise?\n5. Minor details: In Eqs. 13 and 14, the summation index j seems not to be used within the sum. Misspelled “neural” in line 212.  In lines 133-136, you only mention batches of ICS, but shouldn’t it be batches of ICs and PDE parameters?\n\nLIMITATIONS:\nOnly one sentence, which only mentions the performance on the 2D PDE. The modest performance on a single PDE does not show a general problem for specifically multi-field PDEs; it could also be the higher number of spatial dimensions or other reasons. For example, the strong dependence of the performance on the initialization could also be mentioned.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 5, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 4, "DATE": "2025-06-28"}

comments: SUMMARY:
The paper introduces PaPQS, an active learning framework for neural partial differential equation solvers. Unlike traditional pool-based methods that are restricted to specific surrogate types, PaPQS synthesizes new initial conditions directly from the continuous input space. This approach enables finer control over informativeness and diversity in the query batch, helping avoid redundancy and instability. The framework is plug-and-play and generalizes well across multiple surrogate architectures, as demonstrated by comprehensive experiments.

STRENGTHS AND WEAKNESSES:
Strengths
- The plug-and-play design is easy to integrate and broadly applicable across different neural PDE surrogates, including U-Net, SineNet, and FNO.
- The author provides systematic and well-controlled experiments that cover key aspects of the active learning pipeline, including efficiency, architecture generalizability, and data reusability.

Weaknesses
- The method requires access to the ground-truth PDE solver, making it inapplicable in settings where only precomputed or real-world datasets are available without simulation capability.
- The performance gains are modest in several cases; for example, RMSE improvements over SBAL are minimal on Burgers, 2D-CNS, and CE equations. This experiment results cannot prove the efficacy of the proposed PaPQS method.

QUESTIONS:
- Can you compare PaPQS with non-active learning pipelines that share the same objective goal with your paper. For example, the methods that aim for faster convergence, such as those leveraging pretraining or transfer learning?

LIMITATIONS:
The authors bring up the limitations in current design choices, such as the use of fixed step sizes and the absence of physics-informed or temporally adaptive strategies.
TITLE: Official Review
CLARITY: 3
IMPACT: 2
ORIGINALITY: 3
DATE: 2025-06-28

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper presents a plug and play approach towards improving neural network based models which act as PDE solvers. These models rely on training data which can be collected based on the approach in the paper PaPQS. It has two components which consider expected information gain (EIG) and sample diversity together to provide data to train neural PDE surrogates. PsPQS can be used in conjunction with other active learning strategy. Empirical results are presented to support the claims.\n\nSTRENGTHS AND WEAKNESSES:\nStrength:\n1. The idea presented in the paper is very clean and makes a lot of sense intuitively, the authors do a good job in providing intuition for the idea.\n2. The experimental section answers a lot of interesting questions each of which have their own merit.\n3. Relevant literature has been covered.\n\nWeakness:\n1. Empirical evidence is weak to support that PaPQS has noticeable benefits.\n2. The presentation is hard to follow. I think a lot of things are mentioned in a non-standard manner. Please see questions section for more elaborate details.\n\nQUESTIONS:\n1. Fig 2. is not possible to decipher, it is clumsy and the color scheme is bad. I am not able to decode which method is doing how much better than the other. Can you make it better and more readable?\n2. Where are the tables (numerical results)? It is mentioned about Table E1, E2? I don't see them in the main paper, the hyper-links don't work. Without numerical evidence, it is hard to justify the viability of the method.\n3. Fig 3(a) I don't see too much benefit over random.\n4. Fig 4(b) Importance of the diversity component is not prominent but it has been claimed otherwise. Can authors clarify and provide more evidence?\n5. Fig 4(a) The main EIG component also doesn't seem to be independent on it's own. Can you provide more evidence?\n6. Fig 2. Is there a reason to only use PaPQS with certain AL method and not others? Where is the curve for PaPQS alone?\n\nLIMITATIONS:\nAuthors do not include any limitations of their work or any negative societal impact. I think the time complexity of the proposed approach should be discussed as it is expected to be a bottleneck when making an informed choice about the training data needed.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 2, "IMPACT": 3, "ORIGINALITY": 2, "DATE": "2025-06-30"}