Unisolver: PDE-Conditional Transformers Towards Universal Neural PDE Solvers

We sincerely thank Reviewer WyYo for providing a detailed review and insightful questions.

Q1: "The proposed techniques may be limited in its applicability. Though it's a good combination of existing techniques."

Thank you for the feedback. Regarding the applicability, we would like to clarify that in the $\underline{\text{Incomplete Component Scenario}}$ section, the "incomplete" setting in $\underline{\text{Figure 7}}$ actually means inference with no PDE components provided. Therefore, our method can be trained with partial PDE information, and supports inference with no PDE components, which strengthens its applicability in real-world scenarios with incomplete PDE knowledge.

We acknowledge that our method builds upon existing some existing components; however, the novelty of our method lies in the systematic integration of PDE information into neural surrogate modeling through conditional embedding, which to our knowledge has not been explored in prior works.

Q2: "Could you clarify which numerical solver was used in Table 22?"

Thank you for raising this question. In $\underline{\text{Table 22}}$, the numerical solver used for comparison is the pseudo-spectral solver adopted by Li et al. in the original FNO paper, which is not extensively optimized for the specific task. As noted in the paper you cited, FNO is reported to be up to 1,000 times faster than a pseudo-spectral solver, which is consistent with the result in $\underline{\text{Table 22}}$. We promise to cite the mentioned paper and discuss the efficiency improvement more rigorously.

Q3: "How does the neural solver handle extrapolation to future time steps?"

Thank you for your valuable suggestions. We have provided the extrapolation behavior comparison between our model and baseline models in $\underline{\text{Appendix I.3}}$. By extending the prediction horizon to time steps beyond the training range, we observe that all model's performance drops, while our method still outperforms all other baselines.

Q4: How does the model generalize, with and without fine-tuning, to new type of PDEs, especially PDEs with distinct behaviors, such as inviscid burger's equation, wave equation, or Cahn–Hilliard equation, etc.

We appreciate the reviewer's valuable suggestions on providing more PDE type generalization experiments.

Note that we have provided a new PDE generalization analysis in $\underline{\text{Appendix C.2}}$, where Unisolver trained on equations with polynomial order up to 2 demonstrates strong generalization capabilities to equations of polynomial order 3 via fine-tuning.

As per your request, we test our model's performance on 2D wave equation using 200 samples for training and 20 samples for evaluation. The 2D wave equation exibits significantly differnt behaviors from the PDEs of the training dataset. The zero-shot and fine-tuning performance of our model and a FNO model trained from scratch is shown in the table below. Relative L2 is reported.

The zero-shot performance of our model is not very impressive, showing that zero-shot generalization to PDEs with significantly different behavior is rather hard. However, when fine-tuned with only 200 samples, our model is able to achieve a relative error smaller than 1% on evaluation samples, which is far better than FNO trained from scratch and also better than Unisolver trained from scratch, demonstrating its strong generalization capability to different types of PDEs and its effectiveness in adapting to new PDEs with limited data.

We sincerely thank Reviewer GsAC for providing the insightful review and valuable suggestions.

Q1: "The claim of universality may be overstated. The equations considered remain a relatively small subset of possible PDEs." "The domain remains discretized on a grid, even in CFDBench, which further weakens the claim of universality." "Soften their claim on universality."

We thank the reviewer for providing valuable feedback.

(1) Regarding handling irregular geometry.

We acknowledge the limitation of our current method in handling irregular geometries and have discussed this in $\underline{\text{Appendix K}}$. We highlight that this limitation is shared by existing PDE foundation models such as DPOT, MPP, and Poseidon, all of which focus on data defined on regular grids. One fundamental reason behind this is the lack of suitable large-scale PDE datasets on irregular meshes. To extend Unisolver to irregular geometries, one possible approach is to replace the current canonical Transformer with geometry-general PDE models like Transolver.

(2) Regarding the potential overstatement of "universality".

We agree that the claim of "universal" may overstate the current scope of our model, as our method does not handle all possible kinds of PDEs, but rather focuses on several diverse sets of PDEs. Our definition of a "universal" neural PDE solver should be able to incorporate all possible PDE information and to generalize across the wide family of PDEs, and our method is one step beyond existing approaches by systematically encoding available PDE information. We promise to revise our title to "Unisolver: PDE-Conditional Transformers Towards Universal Neural PDE Solvers" to more accurately reflect the scope and goal of our work.

Q2: Although the paper claims to provide theoretical analysis, I did not find any substantial theoretical justification in the text.

We acknowledge the reviewer's concern. Our model design is inspired by theoretical insights on how PDE components influence the solutions, as shown in the motivating example in $\underline{\text{Section 3.1}}$. However, we would like to clarify that we do not attempt to provide formal theoretical analysis or guarantees in this work.

Q3: Suggestions on improving visualization of embedded PDE conditions.

We sincerely thank the reviewer for providing the valuable suggestions. We provide an additional visualization of the learned PDE condition embeddings, which can be found through this anonymous link: https://anonymous.4open.science/r/rebuttal-4EC7/visualization.png. The new 3D plot provides more insightful visualization of the learned embeddings, clearly reflecting how each coefficient impact the embedded PDE conditions.

Q4: Unisolver exhibits more artifacts.

We thank the reviewer for the detailed observation. $\underline{\text{Figure 14}}$ displays the error maps of each model, which is the absolute difference between the model predictions and the ground truth. Although there are visual artifacts in the error map, it is important to note that the model predictions does not display such significant artifacts. This can be demonstrated in the Full trajectory visualization in $\underline{\text{Figure 17 and 18}}$​. Actually, the absolute error is much smaller than the ground truth values, thus making the artifacts in error maps more noticeable.

We sincerely thank Reviewer 8zJe for providing valuable feedback and insightful questions.

Q1: "'Universal' is too big to say."

Thanks for this rigorous review, which is very instructive for us.

(1) We adopt "Universal" in the context of deep learning to highlight model's flexibility and broad applicability.

We appreciate the reviewer’s concern. Our use of the term "universal" is not meant to claim generality across all kinds of PDEs, but rather to highlight the flexibility and broad applicability of our model in incorporating diverse PDE information and handling large-scale PDE datasets. Similar usage also exists in the previous work as "Universal Physics Transformers" [1].

Specifically, our method allows the deep model to incorporate all available PDE information including PDE types, coefficients, boundary conditions, domain geometries and force terms. This allows the model to flexibly handle large-scale and diverse PDE datasets, as demonstrated by our experiments on three challenging benchmarks.

In particular, we have trained a single unified model for all considered 1D PDE datasets, and likewise a single model for 2D PDE datasets containing all kinds of PDE components, each covering a wide spectrum of PDE variations. The experiment setup demonstrates a certain degree of universality, as the models generalize well across diverse PDE families within their respective domains.

[1] Universal Physics Transformers: A Framework For Efficiently Scaling Neural Operators, NeurIPS 2024

(2) We will change the title to "Towards Universal Neural PDE Solvers" for scientific rigor.

Thanks for the reviewer's kind reminder, we acknowledge the potential overstatement of "universal" and we promise to revise our title to "Unisolver: PDE-Conditional Transformers Towards Universal Neural PDE Solvers" to more accurately reflect the scope and position of our work. Similar usage also exists in previous work [2]. This revision will help clarify that the goal of our work is to make progress towards more generalizable and practical neural PDE solvers.

[2] Towards Foundation Models for Scientific Machine Learning: Characterizing Scaling and Transfer Behavior, NeurIPS 2023

Q2: "For such simple examples... How do you justify?"

(1) Not all the PDE-solving tasks require large computation resources.

We apologize for the confusion. While we conduct experiments on a 32-GPU server, we would like to clarify that not all our models were trained on 32 A100 GPUs. The HeterNS dataset is a relatively small dataset and models were trained with a single GPU. For the 2D mixed PDEs dataset, we use 8 A100 GPUs to train our model, aligning with the training configuration reported in DPOT. For the 1D time-dependent PDEs dataset, our model was trained on 32 GPUs, which is justified by its large scale and high diversity, comprising over 3 million training samples in total. The computational cost of Unisolver is shown below, which can also be found in $\underline{\text{Appendix I.7}}$.

Besides, we also want to highlight that although Unisolver requires some time for training, once trained, it can generalize to new PDEs without retraining and efficiently generate solutions, as can be seen in the efficiency analysis in $\underline{\text{Appendix I.5}}$. This allows Unisolver to serve as an efficient surrogate of numerical solvers, significantly reducing the computational overhead, which further justifies the training cost.

(2) Our PDE-solving benchmarks are among the hardest ones in current research community.

Regarding the concern about simplicity, we respectfully argue that the 2D datasets are highly non-trivial, covering a wide range of complex PDE types and diverse PDE components. While the 1D dataset is relatively simpler from a numerical solving perspective, it is specifically constructed to cover a large variety of time-dependent PDEs, posing significant challenges for training and generalization from a deep learning perspective. Specifically, the 1D PDE family contains six polynomial coefficients, various viscosity terms and force terms, three types of boundary conditions, all of them can vary simultaneously, imposing intricate challenges for the model to capture the complex relationship between PDE component inputs and the corresponding solutions.

Finally, we would like to note that the datasets used in our paper are not intended as end goals, but as a foundation for systematically studying the performance of deep surrogate models across increasingly challenging PDE settings, paving the way towards more practical and powerful PDE foundation models.

We sincerely thank Reviewer Z6hm for providing valuable feedback and suggestions.

Q1: About the incomplete scenario setup and incomplete ratio.

(1) Clarify our setting.

Sorry for the confusion. We clarify the incomplete component scenario setup:

• During training, each PDE component (viscosity and force in HeterNS) is independently masked with 30% probability, resulting in 49% samples with full components.
• During evaluation, the "incomplete" setting in Figure 7 means no components available, while "complete" means full components.

The results demonstrate that our model can inference with no PDE components, and providing components boosts performance.

(2) New experiment with a larger incomplete ratio.

We perform additional ablation where each component is masked with 80% probability during training, resulting in only 4% data with full components. As shown in the results below (averaged across multiple forces), our model outperforms FNO and maintains strong performance.

No components available: (80% masked)

Full component: (80% masked)

Q2: The method assumes full knowledge of the governing equation, which may not always be realistic.

According to $\underline{\text{Incomplete component scenario section}}$, our model does not rely on complete knowledge of governing equations and can be trained with partial PDE information, supporting inference with no PDE components. This enables application to real-world settings with incomplete PDE knowledge.

Besides, we want to note that beyond real world, simulation in CAE software is also valuable where complete information is easy to obtain and Unisolver can serve as an efficient surrogate.

Q3: About contribution of LLM embedding, interpretation of learned PDE embeddings, and whether the model is simply performing pattern recognition.

As shown in $\underline{\text{Table 6}}$, incorporating LLM embeddings yields an average improvement of 5.76%, indicating they provide meaningful information. $\underline{\text{Table 7}}$ further shows that LLM embeddings outperform manually constructed symbolic embeddings on both in-distribution and downstream tasks, with over 10% improvement on Advection equation, highlighting better generalization performance of LLM embeddings and benefits beyond simple pattern recognition.

Regarding efficiency, as the LLM has been heavily optimized, generating embeddings incurs negligible computational overhead.

To improve interpretability, we provide additional visualizations: https://anonymous.4open.science/r/rebuttal-4EC7/visualization.png. We filter out equations with viscosity and force for clarity and annotate each cluster with a representative PDE formula. The visualization clearly shows the well-structured latent space of LLM embeddings.

Q4: About statement on existing methods.

We agree that "fail to incorporate all PDE information" may be too strong. We will revise the statement to "do not fully utilize all available PDE information" to soften our claim.

Q5: MPP and Poseidon should be discussed earlier; CODA , CAPE, and Zebra should be mentioned.

We will discuss MPP and Poseidon earlier in Related Works for clarity. We have already cited CAPE, and we will cite CODA and Zebra to better contextualize Unisolver’s novelty.

Q6: About distinction between neural PDE solvers and neural surrogates.

We use the term neural PDE solver to refer to PINNs and neural operators following prior works like Message Passing Neural PDE Solver and Transolver. We also acknowledge that under a narrower definition by [1], only PINNs are considered neural solvers. Considering the mixing concept, we prefer to follow the paper Message Passing Neural PDE Solver with more relative topic and maintain the "neural PDE solver".

[1] Physics-informed machine learning: A survey on problems, methods and applications.

Q7: About conditioning approach and repeating encoding.

Token repeating is an implementation trick similar to tensor broadcasting. We experiment with cross attention on HeterNS. As shown in the table below, our design outperforms cross attention in PDE information conditioning.

Q8: About difference between domain and point-wise embeddings.

The importance of domain-wise vs. point-wise components varies by dataset. In HeterNS, point-wise components (e.g., force terms) are more critical due to their strong influence on fluid patterns. For 1D and 2D mixed PDEs, both types contribute significantly in guiding the simulation.