Unisolver: PDE-Conditional Transformers Are Universal PDE Solvers

Q4: "I am wondering if a simple manual encoding could be enough for a similar performance. This question also relates to the "Random vector" ablation study in Table 7. How did the random vectors get assigned for each PDE?"

Sorry for the confusion. The baseline in $\underline{\text{Table 7}}$ using "random vector" encoding actually uses orthogonal random vectors. It is equivalent to a baseline using "one-hot" encoding to indicate which PDE is currently used. We have updated the description of this baseline $\underline{\text{in Lines 430 and 442 of the revised paper}}$.

Moreover, following your request, we conduct additional ablation studies by manually constructing a one-hot vector for each PDE term and combining them to represent a full PDE. The results are shown in the table below. We observe that the LLM embedding case consistently outperforms the manual encoding case in both in-distribution tests and four zero-shot generalization settings, showing a 4.23% average improvement. Although the manual encoding aims to preserve the mathematical structure of the PDE as much as possible, the handcrafted features struggle to perfectly capture the mathematical structure provided by LLMs visualized in $\underline{\text{Figure 7(b) of the original submission}}$, leading to a decrease in performance.

Further, using an LLM to encode the symbolic information allows for more flexible adaptation to more diverse sets of PDEs. For example, when we want to extend the 1D PDE dataset to contain non-polynomial functions like sinusoidal terms. This flexibility is particularly advantageous, as it reduces the manual effort associated with designing unique embeddings for each new PDE component, enabling broader adaptability and scalability across various PDE applications. We have included this ablation study in $\underline{\text{Appendix B in the revised paper}}$.

We would like to sincerely thank Reviewer fFn6 for providing a detailed review and insightful questions.

Q1: "The current experiments seemed to fix the temporal evolution up to 20 steps. Experiments and analysis for long trajectory prediction will be invaluable."

Thank you for this valuable suggestion. In our experiments, the 1D time-dependent PDEs benchmark indeed includes 100 time steps within the range [0, 1]. On this benchmark, Unisolver utilizes an implicit neural representation to decode the solution at any spatial-temporal coordinates within a fixed domain: $(x,t)\in[-1,1]\times[0,1]$.

While on the HeterNS benchmark, the number of evolution steps is limited to 20, the actual time interval between 2 consecutive frames is 1 second which is a substantial temporal gap. A numerical solver would require approximately 10,000 time-stepping iterations to achieve precise one-second solutions.

Further, as per your request, we extend the temporal evolution steps of HeterNS to 30 steps, corresponding to 30 seconds of complex fluid dynamics, and report the zero-shot performance comparison between Unisolver and the baselines in the table below. We present the performance on the subdataset with a viscosity coefficient of $\nu = 1 \times 10^{-5}$ and a force coefficient of $\omega = 2$. This is a particularly challenging task, as these models have never seen such long trajectories in the training data (at most 20 seconds). Despite this, Unisolver achieves the best performance. We have included this experiment in $\underline{\text{Appendix H.3 of the revised paper}}$.

Q2: "It will be particularly useful to have actual codes because of the combinatory design of the proposed model."

We sincerely thank you for providing this helpful suggestion. To facilitate the understanding of Unisolver's design, we have included an example inference code for the 1D time-dependent PDEs benchmark in the $\underline{\text{supplementary material}}$, along with the checkpoint of our largest Unisolver model, which was trained for 3000 A100 hours on 3 million samples. You can modify the PDE configurations like PDE coefficients, external forces and boundary conditions and inference with Unisolver by yourself. Furthermore, we promise to make both training and inference code public upon acceptance.

Q3: "What will happen if the LLM model (i.e., conditioning equation symbols) is omitted in a single PDE dataset?"

Thank you for this insightful question. In the HeterNS benchmark, we do omit the equation symbol embeddings, as the entire dataset is based on a single PDE with varying coefficients and external forces. In this case, the equation symbol embeddings are unnecessary because all samples share the same underlying PDE structure. The LLM embeddings are helpful when the dataset contains multiple types of PDEs by allowing the model to distinguish between different PDEs more effectively. Consequently, only the viscosity coefficients and external forces are inputs to the model on the HeterNS dataset.

We would like to sincerely thank Reviewer awFj for providing a detailed review and insightful questions.

Q1: "It is suggested to remove these proofs and cite a textbook. Further, I think it is unnecessary to use Theorem 1 to introduce the concepts of domain-wise and point-wise PDE components."

Thanks for your valuable suggestion. As per your request, we have removed the proofs in the appendix and cited the proof from a textbook instead. Additionally, we have revised $\underline{\text{Section 3.1}}$ to use the vibrating string equation solely as a motivating example for introducing the concepts of domain-wise and point-wise PDE components. We believe that the $\underline{\text{revised paper}}$ delivers a more compact and direct description of our motivation.

Q2: "Please explain why you choose to predict solutions in a data-driven manner instead of using numerical solutions and PINNs, which both can take full advantage of complete PDE components. Especially, please disucss the pros and cons of Unisolver vs PINNs in terms of accuracy and computational efficiency."

Firstly, as demonstrated in our experiments in "Incomplete component scenario" $\underline{\text{in Line 496 of Section 4.4}}$, Unisolver exhibits strong performance even when some PDE components are incomplete. In contrast, PINNs and traditional numerical methods are inapplicable under such circumstances.

Further, regarding accuracy, it is well known that PINNs "often face severe difficulties and even fail to tackle problems whose solution exhibits highly nonlinear, multi-scale, or chaotic behavior"[1], and the optimization process of PINNs often faces thorny challenges [2], even for a simple steady Navier-Stokes Equation [3]. The datasets used in our work, such as HeterNS, can present significant challenges for PINNs due to their high non-linearity and chaotic nature.

In terms of computational efficiency, as stated in $\underline{\text{Lines 35-36, 44-47 of main text}}$, PINNs and numerical methods also struggle to generalize to new PDEs with varying PDE components like coefficients or boundary conditions, requiring retraining for each new PDE. In contrast, although Unisolver requires some time for data collection and training, once the model is well trained, it can generalize to new PDEs without retraining and generate a solution within a second. The efficiency benefit is the most widely acknowledged advantage of Neural Operator than PINNs.

Moreover, we provide the inference times and memory consumption for each model to predict a single frame on the HeterNS dataset, along with the calculation time and memory consumption for the numerical solver which is used to generate the dataset, as summarized in $\underline{\text{Appendix H.5 in the revised paper}}$. Unisolver is 1,000+ times faster than the numerical solver.

[1] Respecting causality for training physics-informed neural networks. Computer Methods in Applied Mechanics and Engineering 2024.

[2] RoPINN: Region Optimized Physics-Informed Neural Networks. NeurIPS 2024.

[3] PINNsFormer: A Transformer-Based Framework For Physics-Informed Neural Networks, ICLR 2024

Q3: "Traditional operator learning method such as FNO train neural networks to generalize across different configuratons (such as viscosity coefficients and force terms) of the same type of PDE."

We respectfully point out that this comment exists some misunderstandings of FNO. It is worth noting that in the original FNO paper, the model was not designed to handle PDEs with different configurations, but rather train separate models for every single configuration. Specifically, the original NS dataset proposed by FNO test the model's capability to generalize to unseen initial conditions under a fixed PDE configuration.

In contrast, Unisolver handles a more general generalization task, requiring the model to generalize not only to unseen initial conditions, but also to other PDE components like the coefficients and external forces. The HeterNS benchmark is an extension of the original NS benchmark to test the models' capability to generalize across different configurations within the same PDE type, while the 1D time-dependent PDEs and 2D mixed PDEs benchmark require the model to generalize across different types of PDEs.

To be more clear, we have listed the comparison between settings of Unisolver and FNO as follows and add $\underline{\text{Figure 2 in revised paper}}$ for clarity.

Q4: "Did different types of PDEs help each other during prediction? If possible, please give an ablation study to show the transfer learning capability of Unisolver among different types of PDEs, and justify the integration of a wide range of PDEs instead of treating them independently."

As per your request, we design an ablation study to evaluate the effect of joint training on the 1D time-dependent PDE benchmark. As stated in $\underline{\text{Appendix E of the original submission}}$, the general equation formulations used in this benchmark include two polynomials, $f_0$ and $f_1$, both with a maximum order of 3. We construct three distinct sub-datasets, each with 10,000 samples, to test the impact of joint training. The polynomials in each dataset are fixed to orders of 1, 2, and 3, respectively, ensuring that the PDEs contained in these three datasets do not overlap. For instance, in the dataset with polynomials of order 3, only $c_{03}$ and $c_{13}$ are non-zero terms, while $c_{01}$, $c_{02}$, $c_{11}$ and $c_{12}$ are fixed to zero.

We conduct both joint training and independent training for 500 epochs on these 3 subdatasets. The results are shown in the table below.

It can be shown that joint training significantly boosts performance on three datasets, even though the 3 types of PDEs do not overlap, demonstrating that these PDEs do "help each other" during training. We have included this experiment in $\underline{\text{Appendix C in the revised paper}}$.

As for the transfer learning capability of Unisolver, we have provided a equation generalization experiment in $\underline{\text{Appendix G.3 of original submission}}$. The results demonstrate that Unisolver has strong generalization capabilities to unseen PDEs, even when the downstream PDEs are entirely new.

Q5: "Please explain why Unisolver works much better than PINO in table 3 in terms of relative error. Can we improve the accuracy of PINO if we weight the physics-informed loss more heavily?"

Thanks for pointing this out. Firstly, we would like to detail that PINO in Table 3 is essentially a FNO model augmented with a physics-informed loss, as we stated in $\underline{\text{Lines 319-321 of main text}}$. It is worth noting that although physics-informed loss proposed by PINO seems to help the model training, in pratice, this loss cannot be accurately calculated in experimented data due to the finite spatiotemporal resolution.

The HeterNS benchmark, as well as the original NS benchmark in the FNO paper, employs a coarse spatial resolution of $64 \times 64$ and a long temporal interval of 1 second. Meanwhile, the official implementation of PINO utilizes a finite difference method to compute the temporal derivatives, which cannot produce an accurate estimation of derivatives given the long temporal interval between two consecutive solution frames.

We try to raise the weight of the physics-informed loss for PINO, and report the results below. The results show that the performance of PINO gets worse if we raise the weight of the physics-informed loss.

Q6: "Please give the computational cost to generate the datasets, including hardware used and simulation time, especially for HeterNS."

Note that the 1D time-dependent PDEs and 2D mixed PDEs benchmark are not proposed by us and we cannot find the exact data generation time in public resources. However, we contacted the authors of PDEformer by email, and they said that it took about 2 weeks to generate "1D time-dependent PDEs" with 96 CPU cores. As for the 2D mixed PDEs benchmark, it is collected by DPOT with several data sources, which is hard to calculate the extract generation time due to inconsistent experiment environments.

As for the HeterNS benchmark, we use a single A100 GPU to generate all the training and test samples. The total simulation time is about 35 hours for generating the 15000 training samples. Again, we want to highlight that the inference efficiency of Unisolver should not be underestimated, which does not require retraining for new tasks, leading to an efficient surrogate model for traditional numerical methods.

Q7: "Why use the LLM embedding of PDE symbol instead of Latex code? It is better to give an intuitive example to show the additional useful information provided by LLM beyond equation symbols."

Unlike numerical information such as the PDE coefficients and external forces, the LaTex codes of the PDE symbols are essential strings that cannot be feed into the neural network directly. Therefore, we use an LLM to encode the LaTeX codes into a fixed-dimensional vector.

Meanwhile, $\underline{\text{Figure 6(b)}}$ shows the PCA visualization of the LLM embeddings. We can observe that the LLM is able to encode PDEs with similar complexities into similar embeddings, which intuitively shows that the equation symbol embeddings capture prior mathematical information beyond the literal symbols.

Thank you again for your thoughtful feedback. We noticed that you have rated our paper with a score of 3, and assigned a confidence level of 5, which indicates a rather negative evaluation. If you have any further questions or require additional clarification, we would be pleased to offer additional explanations.

Dear Reviewer awFj,

We kindly remind you that there are only two days left of the two-week Reviewer-author discussion period. So please kindly let us know if our response has addressed your concerns. We will be happy to deal with any further issues/questions.

We made every effort to address the concerns you suggested:

• The vibrating string equation in $\underline{\text{Section 3.1}}$ is now used solely as a motivating example to introduce the concepts of domain-wise and point-wise PDE components. We also remove the proof in $\underline{\text{Appendix E}}$ and cite standard references.
• We elaborate on the advantages of Unisolver over PINNs and numerical solvers, especially in terms of generalizability and computational efficiency. We also include inference time and memory consumption comparisons between Unisolver and numerical solvers in $\underline{\text{Appendix H.5 of the revised paper}}$.
• We add an ablation study on the 1D time-dependent PDE benchmark in $\underline{\text{Appendix C.1}}$ to demonstrate that joint training on different PDEs significantly improves performance. We also provide Unisolver's capability to generalize to entirely new PDEs in $\underline{\text{Appendix C.2}}$.
• We provide detailed explanations on PINO's performance and add an experiment that raises the weight of the physics-informed loss following your suggestion.
• We provide detailed generation costs of the HeterNS dataset.
• We explain the reason of using an LLM to encode the symbolic information of PDEs.

These updates have already been incorporated into the $\underline{\text{revised paper}}$. Inference code is also included in $\underline{\text{supplementary material}}$.

We hope our responses and revisions have addressed all your concerns effectively.

Thanks again for your dedication to reviewing our paper. Looking forward to your reply.

We would like to sincerely thank Reviewer 1Nnh for providing a detailed review and insightful questions.

Q1: "The components defining the behavior of a PDE are already known and not a "key finding," as stated in the abstract."

Many thanks for your feedback. We agree that the components defining the behavior of a PDE are already known. We have removed the proofs in the appendix and revised the $\underline{\text{Abstract and Section 3.1}}$ to use the vibrating string equation solely as a motivating example. We believe that the $\underline{\text{revised paper}}$ delivers a more compact and clear description of our motivation.

Q2: "Limited methodological novelty (but a good combination of existing components)."

Thanks for your appreciation of our overall design. It is worth noticing that while our approach leverages existing components, the design is non-trivial. We carefully consider each component and develop tailored embedding strategies to capture their unique characteristics effectively.

Furthermore, we are the first to clearly define a method for integrating PDE information into deep models via conditional modeling. This comprehensive approach leads to highly effective performance as demonstrated by the extensive experimental results.

Q3: "No evaluation of inference times and memory consumption compared to the baselines."

As per your request, we provide the inference times and memory consumption for each model to predict a single frame on the HeterNS dataset, along with the calculation time and memory consumption for the numerical solver which is used to generate the dataset to calculate the next frame, as summarized in the table below. The results are measured on an A100 GPU with a batch size of 1. Unisolver demonstrates comparable inference speed to FNO, while consuming less memory. Besides, all neural PDE solvers are approximately 1,000 times faster than the numerical solver. We have included the computational cost comparison in $\underline{\text{Appendix H.5 of the revised paper}}$.

Q4: "The ablation shows that the performance drops when not using the LLM by comparing it to a random vector encoding, but it is not shown that this is due to prior mathematical knowledge of the LLM."

We appreciate your thoughtful feedback. To further verify the role of LLM embeddings in encoding PDE information, we conduct two additional ablation experiments. In the first experiment, the LLM only encodes the number of non-zero terms in the 1D PDE. In the second experiment, the LLM encodes the "wrong" PDE information. Specifically, we replace “*” with “/” and adjust polynomial orders to their reciprocals. For example, the original latex code $u\_t + c\_{01} * u + c\_{02} * u^2 + s(x) + (c\_{11} * u + c\_{13} * u^3)\_x = 0$ is transformed into $u\_t + c\_{01} / u + c\_{02} / u^{1/2} + s(x) + (c\_{11} / u + c\_{13} / u^{1/3})\_x = 0$. The results of these two ablation studies are shown below.

The results indicate that the model indeed obtains additional information beyond merely the count of non-zero terms from the LLM embeddings. Moreover, embedding "wrong" mathematical information generally leads to a decline in performance, highlighting the importance of accurately embedding the PDE information. While we cannot definitely claim that the LLM "understands" mathematical knowledge, we can confirm that the use of LLM enables us to encode useful mathematical information into deep representations. Furthermore, we also update the corresponding "mathematical knowledge" to "mathematical information" in the $\underline{\text{revised paper}}$ for scientific rigor. Meanwhile, the additional ablation studies have also been included in $\underline{\text{Appendix B of the revised paper}}$.

Q5: "Please clarify how you performed the random vector encoding."

Wouldn't a one-hot encoding of which PDE version is currently used be a better baseline?

Sorry for the confusion. The baseline using "random vector" encoding actually uses orthogonal random vectors. It is equivalent to a baseline using "one-hot" encoding to indicate which PDE is currently used. We have updated the description of random vector to orthogonal random vectors in $\underline{\text{Line 430 and 442 in the revised paper}}$ accordingly.

Q6: "You introduce a binary mask to notice the model which component is active. Shouldn’t the LLM encoding of the equation handle this?"

Thank you for your thoughtful feedback. In the data pre-processing stage, if a component does not exist in a certain dataset, we set this component to zero by default. However, even a zero-valued component might also hold its physical significance—for example, a zero-valued external force may convey specific physical meaning. Therefore, we include an additional binary mask to indicate to the model whether a component is an actual PDE term or simply set to zero by default. While the LLM embedding can provide some indication of this information, it does not serve as the input to the encoders of other components. This binary mask, however, aids the encoder's learning and further clarifies the information without introducing significant computational overhead. Moreover, we update the description in $\underline{\text{Line 240 in the revised paper}}$.

Q7: "How costly is using an LLM to encode the equation?"

Thank you for your insightful question regarding the computational overhead of generating LLM embeddings. It takes about 450 seconds to get the embeddings for all 576 types of PDEs using the Llama3-8B model on a single A100 GPU.

Note that large language models have been extensively optimized for engineering applications by extensive researchers and companies, which allows for efficient processing within reasonable time constraints.

Q8: "Is the model trained with one-step predictions or autoregressive rollouts?"

For the HeterNS and the 2D mixed PDEs benchmark, following previous methods FNO and DPOT, we train the model using one-step predictions and test the model in an autoregressive manner. For the 1D time-dependent PDEs benchmark, we train the model to predict the solution at specific spatial-temporal coordinates through an INR following PDEformer.

Moreover, we update the description in $\underline{\text{Line 186, 214 and 230 in the revised paper}}$.

Q9: "How are the inputs aggregated in the “deep condition consolidation”?"

Sorry for the confusion. In this paper, we simply add the deep conditions within the same class. And we have updated the description in $\underline{\text{Line 256 in the revised paper}}$.

Q10: "Please clarify in the paper on all figures that the numbers refer to the relative L2, especially since you switch between relative L2 and relative L2 x 10^-2."

Thank you for pointing this out. We have updated the corresponding figures and tables to clearly indicate the metrics we use $\underline{\text{in the revised paper}}$.

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

Clarify some misunderstandings in the review.

First, we would like to clarify a few points that may have led to some misunderstandings to ensure that our work is accurately represented.

(1) About Method: Unisolver is not an LLM based PDE solver.

You commented in your summary of our paper that “In this paper, they add symbolic representation of the PDE operator to an LLM based PDE solver.”

In fact, we don't propose an LLM based PDE solver. We employ an LLM as a frozen encoder to get the symbolic embedding of a PDE. The frozen embedding is then fed into our model to inform the symbolic information of the PDE. The coefficients, boundary conditions, external forces and geometry information are all represented as scalars or fields and encoded by our proposed Universal Component Embedding approach. Thus, the main body of Unisolver is a conditional Transformer, where LLM just performs as an embedding module. To make this issue clear, we have revised the paper by updating the model architecture in $\underline{\text{Figure 3 of the revised paper}}$.

(2) About Topic: Neural PDE solvers hold significant advantage in efficiency than traditional numerical methods.

You commented that "PDE solvers are a machine learning approach to a problem where, if the inputs are properly specified, and the PDEs is well-posed, there is no need for machine learning".

We respectfully disagree with your opinion. In comparison to traditional solvers, Neural PDE solvers have shown great advantages in both efficiency and generalizability. For example, FNO can "learn an entire family of PDEs, in contrast to classical methods which solve one instance of the equation" [1]. After training on a wide range of PDEs, Unisolver is able to generalize to PDEs with varying PDE components. Although the training time of Unisolver is relatively long, we can generate the solution on unseen PDEs much faster than traditional numerical PDE solvers. The inference efficiency benefit should not be underestimated. Please think about GPT-4, the training period can be very long, while once the model is well trained it can support very fast and parallel inference.

Moreover, we provide the inference times and memory consumption for each model to predict a single frame on the HeterNS dataset, along with the calculation time and memory consumption for the numerical solver which is used to generate the dataset to calculate the next frame, as summarized in $\underline{\text{Appendix H.5 of the revised paper}}$. Unisolver is 1,000+ times faster than the numerical solver.

[1] Fourier Neural Operator for Parametric Partial Differential Equations. ICLR 2020.

(3) About Related Work: The most widely acknowledged advantage of Neural Operators lies in their efficiency, rather than in their ability to operate without requiring a specific equation.

You commented that "This paper critiques existing methods as not taking full advantage of the data, but this is a misrepresentation of the learning problem for previous methods, which are expressly designed to solve the PDE without knowing the form of the equation."

Actually, instead of the ability to solve PDE without exact equations, the most acknowledged benefit of Neural Operators is their efficiency in solving new tasks. We respectively quote some sentences in the previous papers, such as "As a result, the learning-based method can be orders of magnitude faster than the conventional solvers." in FNO (ICLR 2021) and "Serving as surrogate models, neural operators can be more than five orders of magnitude faster than traditional numerical solvers" in DPOT (ICML 2024). Under this spirit, Unisolver is one step towards making the neural PDE solvers applicable to more diverse scenarios, which is of great value in practice.

Further, all the previous methods cited in this paper (e.g., FNO, DPOT, PDEformer,f and ICON) adopt the simulation data for experiments and pursue an efficient surrogate for numerical methods. Since the equation information is just off the shelf in this context, we believe that Unisolver's setting is natural and reasonable.

(4) About Experiments: Many of our experiments follows previously used benchmarks and metrics, and directly compare with officially reported results of baselines.

You commented that "There are not standard datasets with standard performance measures: these are new datasets defined by the authors, and performance measures and baselines defined by the authors."

We respectfully disagree with your opinion. Firstly, we propose the HeterNS dataset by extending the standard Navier-Stokes dataset proposed in FNO with broader viscosity coefficients and external forces. The other two datasets, namely the 1D time-dependent PDEs and the 2D mixed PDEs, are proposed by PDEformer and collected by DPOT respectively. Secondly, the performance measurement metric—Relative L2 we used is widely recognized within the Neural PDE solver field. This metric is considered standard for assessing model performance, as used by FNO, PDEformer and DPOT. Lastly, we compare Unisolver against the official implementation of each baseline in the HeterNS benchmark. The only change we make to these baselines is to augment these models with PDE components to ensure that they get the same input as Unisolver. As for the 1D time-dependent PDEs and 2D mixed PDEs benchmarks, we compare Unisolver against PDEformer and DPOT in a head-to-head manner without making any changes to them.

We hope the above clarifications can help the reviewer understand the value of our work better. Next, we will answer the detailed questions.

Q1: "It was not clear exactly what the PDE was trained on, in for example, Table 2. " "Inputs, outputs, and data need to be specified."

We have provided detailed descriptions of the PDEs we trained on in $\underline{\text{Appendix E of original submission}}$, including the the analytical form of each equation. Following your suggestion, we have also included a new figure $\underline{\text{Figure 2 of main text}}$ to illustrate the input and output of our model.

Q2: "Also the modification to the baselines is not clearly described." "These are no longer baselines, but new methods."

We appreciate your concerns regarding our comparisons with the baselines. Actually, we have already provided implementation details of each baseline in $\underline{\text{Appendix F of original submission}}$. We would like to clarify that our comparisons are indeed fair. In the HeterNS benchmark, we input external force and coefficient information in numerical form rather than as text to Unisolver. For consistency, we chose to concatenate this information with the baseline input, allowing for a meaningful and fair comparison.

Additionally, we run FNO without concatenating the PDE components and report these results below, demonstrating that FNO’s performance declines without access to this information. That is why we "augment these baselines" in our original submission instead of directly using their original design.

Q3: "Motivated by the mathematical structure of PDEs, we define the concept of complete PDE components, classify them into domain-wise and point-wise categories, and derive a decoupled conditioning mechanism for introducing physics information into PDE solving. This is a concept completely original to the paper, which ignores existing expertise on PDE operators, and is not clearly explained. It's also not clear how it corresponds to what they do in the paper."

Firstly, We would like to clarify that the concept of “Complete PDE Component” is specific to conditional deep neural networks and is not intended to serve as a formal mathematical theorem. In the context of deep learning, it is widely recognized that providing models with fine-grained control conditions can lead to more deterministic and desired results (See [1,2]). Our goal here is to define the conditions necessary for a Neural PDE solver to function effectively.

Additionally, we categorize these PDE components into domain-wise and point-wise components also in the context of deep learning. Domain-wise components such as coefficients and boundary condition types are generally represented scalars or fixed-dimensional vectors that are uniform across the domain, and point-wise components are generally functions corresponding to specific spatial and temporal coordinates. To integrate these components effectively, we design tailored embedding strategies that incorporate them as deep conditions in Unisolver.

We have highlighted the deep learning context in the $\underline{\text{revised paper}}$.

[1] Adding Conditional Control to Text-to-Image Diffusion Models, ICCV 2023

[2] Scalable Diffusion Models with Transformers, ICCV 2023

Q4: "Theorem 1: This is not a theorem." "The solution expressed in the theorem is particular to the equation, and there do not exist similar explicit solutions for most PDEs."

Thanks for your rigorous review. We acknowledge that this particular PDE simply serves as a motivating example to show the concept of PDE components. We categorize these PDE components in the context of deep learning. The model architecture is motivated by the categorization and also builds upon empirical experiments. We have modified this part in the $\underline{\text{revised paper}}$.

Q5: "All the equation prompts had constant coefficients. How would you prompt a two dimensional equation, with variable coefficients?"

As stated in $\underline{\text{Equation formulation paragraph in Line 233 of the original submission}}$, we don't use an LLM to encode numerical information such as coefficients. The LLM is only responsible for embedding the symbolic information of the PDE. For your further information, it may be hard for LLM to directly solve a PDE. What we want to do in Unisolver is to utilize LLM's basic knowledge in understanding PDEs (Please see visualization in $\underline{\text{the end of page 9}}$).

For a two dimensional equation with variable coefficients, we would use patch embedding to embed the coefficients into deep conditions.

Q6: "How do you distinguish between the exact PDE and incomplete knowledge of the PDE?"

Again, As stated in $\underline{\text{Equation formulation paragraph}}$, we would like to clarify that our LLM embeddings only encodes the Latex code of the PDE, which does not include the numerical part of the PDE. This means that not all PDE information is captured within these symbols. We believe that LLMs, while capable of general mathematical understanding, cannot address all aspects of mathematical problems comprehensively and accurately. The embeddings primarily capture a general understanding relevant to a class of PDEs.

To address incomplete knowledge scenarios, we conduct experiments with “incomplete” PDE information in "Incomplete component scenario" section $\underline{\text{in Line 496 of original submission}}$, where certain components were omitted and replaced with a learnable token to represent the unknown PDE component. These experiments demonstrate that Unisolver still performs well without complete information of the PDE.

Q7: "There is no universal PDE solver, since PDEs do not universally have solutions. Some PDE operators are well-posed, which means there existing unique solutions which are stable under certain types of perturbations of the data."

Thanks for your valuable question. We have revised our title to Universal Neural PDE Solver to moderate our claim. In the context of deep learning, "Universal" means that the deep model can solve more diverse distribution (Please see [1]).

Besides, it is important to clarify that our model is not intended to solve every possible PDE. Equations that are not well-posed cannot be solved without sufficient training data. However, our paper focuses on one of the most prominent and cutting-edge problems in the community: whether a trained Neural PDE solver can generalize to more diverse PDEs, which is also the focus of PDEformer and DPOT. To further illustrate our approach, we have modified $\underline{\text{Figure 2,3 in the revised paper}}$ to highlight our definition of "universal" in this context.

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

Q8: "This is an imprecise characterization. There are multiple formulations of PDE problems: these include: (i) learning a solution operator, based on a dataset of solutions on grids, and (ii) learning a single PDE solution, based on incomplete data, as well as knowledge of the PDEs, as well as (iii) learning the PDE operator from sample(s) of (possibly incomplete data."

The categorization of Neural PDE solvers into pure data-driven methods and physics-informed methods is well accepted [1] [2]. Data-driven methods utilize data from numerical solvers or existing experiments, while physics-informed methods only depend on the PDE constraints. A representative example of data-driven methods is the class of neural operators, while a representative example of physics-informed methods is the class of PINNs.

It would be highly appreciated if you provided references that support your mentioned specific categorization.

[1] Physics-informed neural operator for learning partial differential equations. ACM/JMS Journal of Data Science 2024

[2] Transolver: A fast transformer solver for pdes on general geometries. ICML 2024

Q9: "What exactly is the training data, expressed in terms of the solutions to which PDEs? Was it all the PDE solutions, were the text prompts included with the data? What is the test data? An initial condition? Along with the text of the PDE? Please describe what were the inputs to the training dataset, what were the inputs to the model (besides the latex text prompt for the PDE) and the models trained on a dataset of PDEs for each of the three main cases?"

We have provided detailed descriptions of training data in $\underline{\text{Appendix E of the original submission}}$. To be more clear, we have also included or modified figures in the main text ($\underline{\text{Figure 2,3 in revised paper}}$) to illustrate the input and output of our model more clearly.

Q10: "what is meant by "favorable generalizability"?"

Favorable generalizability means the empirical improvement on the standard benchmarks we considered, where our model significantly outperforms numerous baseline models in a zero-shot generalization setting. For detailed information on the benchmark design and baseline implementations, please refer to $\underline{\text{Appendix E and F of the original submission}}$.

Q11: "This is not clear: what are 'point-wise deep conditions'"

As stated in $\underline{\text{Sec 3.2 of the original submission}}$, we provide a detailed explanation of how to embed all considered PDE components into deep PDE conditions. The point-wise deep conditions are obtained by embedding each point-wise component and then aggregating them, as illustrated in $\underline{\text{Figure 3 of the original submission}}$, which is relatively clear in the context of deep learning.

Again, we have to point out that Unisolver is not an LLM based solver. The LLM serves as a frozen encoder to encode the symbolic information of the PDE, while the numerical part of the PDE are encoded either by an MLP layer or a Patch-Embedding module.

Dear Reviewer uAAv,

We kindly remind you that there are only two days left of the two-week Reviewer-author discussion period. So please kindly let us know if our response has addressed your concerns. We will be happy to deal with any further issues/questions.

We respectfully clarify some misunderstandings about our work:

• We clarify that Unisolver is not an LLM-based PDE solver, as the LLM serves only as a frozen encoder for symbolic embedding.
• We emphasize the efficiency and generalizability advantages of neural PDE solvers over traditional numerical methods.
• We highlight that in recent papers, the most widely acknowledged advantage of Neural Operators lies in their efficiency, rather than in their ability to operate without requiring a specific equation.
• We clarify that the datasets and performance metrics we use are all standard and that we make fair comparisons with baselines.

We also made every effort to address the concerns you suggested:

• We have revised our title to Universal Neural PDE Solver to moderate our claim.
• We include a new figure $\underline{\text{Figure 2}}$ and revise $\underline{\text{Figure 3 of main text}}$​ to better illustrate the problem setup of our work and the input and output of Unisolver.
• We emphasize that the concept of "Complete PDE Component" is specific to conditional deep neural networks and the categorization of domain-wise and point-wise components is also in the context of deep learning.
• We have revised $\underline{\text{Section 3.1}}$ to use the vibrating string equation solely as a motivating example for introducing the concepts of domain-wise and point-wise PDE components.

These updates have already been incorporated into the $\underline{\text{revised paper}}$. We hope our responses and revisions have addressed all your concerns effectively.

Thanks again for your dedication to reviewing our paper. Looking forward to your reply.

Dear reviewer uAAv,

Thanks for your time in reviewing our paper.

As the discussion period is nearing its end, we are eager to know whether our response has addressed your concerns about the problem setup and the value of machine learning for "properly specified inputs and well-posed PDEs."

We sincerely look forward to your feedback.

Best regards,

Authors

Dear reviewer uAAv,

We sincerely thank you once again for your time reviewing our work.

Since there are only 9 hours left before the deadline for the reviewer-author discussion, we would greatly appreciate your feedback on whether our responses have adequately addressed your concerns.

For clarity, we highlight some key points from our response here:

• Our topic is meaningful and practical, which is also actively explored by NVIDIA. Even when "the inputs ... well-posed", Neural PDE Solvers hold unique advantages in efficiency. Specifically, Unisolver is over 1,000 times faster than traditional numerical solvers, see our previous response. This is why some top companies, such as NVIDIA Omniverse Blueprints, still work on Neural PDE Solvers for pure simulation scenarios with complete PDE information. Thus, we respectfully disagree with your opinion that "if the inputs are properly specified, and the PDE is well-posed, there is no need for machine learning".

• The problem setup is clear in our paper with all the input-output well defined. Our problem definition has been clearly presented in $\underline{\text{Figure 2 and the problem setup paragraph in Section 3 in the revised paper}}$. Besides, our paper also includes all the inputs and outputs of each benchmark. Thus, we think the problem setup is crystal clear. We do hope you can refer to $\underline{\text{Page 3-4}}$ of our revised paper.

Based on the previous discussion, we think there are some foundational misunderstandings in your review. Since it is nearing the end of the rebuttal period, we do hope you can carefully review our clarification and reconsider your score.

Best regards,

Authors