Mixture-of-Experts Operator Transformer for Large-Scale PDE Pre-Training

Dear Reviewer,

Thank you for your incredibly thorough and insightful review. We are grateful for your feedback and for acknowledging our work's motivation and novelty. Your critiques regarding clarity, evaluation, and literature context are invaluable and provide a clear roadmap for improving our paper. We sincerely apologize for these shortcomings and are committed to a full revision. Below are our responses and new experimental results.

---

On Literature Review, Novelty, and Comparison with Poseidon

Our initial choice of DPOT as a baseline was threefold: to ensure clear attribution of gains from our MoE module; the divergent experimental setups (our auto-regressive, "trajectory-to-trajectory" paradigm differs significantly from Poseidon's excellent parameter-informed approach); and the fact that Poseidon was contemporaneous work released after our experiments began.

• We agree that testing on challenging, unseen tasks is essential. To address this, we conducted new experiments on the Wave-Layer and Wave-Gauss datasets you proposed.

• These results compellingly show that each model excels in its intended paradigm. This highlights that they are both powerful tools with different, valuable scopes of application.

• Commitment: We commit to including a thorough discussion of Poseidon and other related MoE works in our literature review. We will also incorporate these new experimental results to provide a more complete and honest picture of our model's performance and positioning within the field.

On Task Definition, Evaluation, and Clarity (Q4, Q6, Q7, Q17, Q18, Q23)

Your critiques on our operator learning task definition, limited evaluation, and writing clarity are spot on. We sincerely apologize for these shortcomings and are committed to a full revision, including Section 3, to ensure clarity and logical flow.

• Task Definition & Motivation: You are correct that our "trajectory-to-trajectory" task differs from the classic formulation. Our setup, following DPOT, is motivated by real-world scenarios (e.g., weather forecasting) where explicit PDE parameters are unknown and must be inferred from observed data. While parameter-informed approaches like Poseidon are powerful, our core contribution is introducing the MoE architecture to this specific auto-regressive paradigm.
• Dataset Choice & Heterogeneity: We chose our datasets for fair comparison with DPOT. Our central argument is that the key challenge is not single-task difficulty but handling the heterogeneity of mixed datasets. This "knowledge conflict" is where dense models fail, as shown in Figure 2, validating the need for our architecture. Even for tasks with similar equations like PDEArena, our router distinguishes them with 100% accuracy, proving it captures significant distributional differences beyond just the equation type.

On Dense Model Challenges & MoE's Rationale (Q1, Q3)

The performance collapse of dense models like FNO on mixed datasets is due to "negative transfer," where a single shared kernel struggles to learn conflicting physical laws. To validate this, we tested a dense DPOT-S model as more heterogeneous datasets were added. The results below show that adding more data often degrades DPOT's performance, while our MoE-POT remains superior. This confirms that a sparse architecture is crucial for handling knowledge conflict in diverse PDE pre-training.

On Interpretability with Unseen Datasets (Q36)

To stress-test our router's generalization, we tested its ability to classify data from two challenging, unseen downstream tasks (Wave-Layer and Wave-Gauss).

Furthermore, the router exhibited distinct and highly confident selection patterns for these new tasks. For example, in Block-2, Wave-Layer data triggered experts 1 (79%), 9 (31%), 11 (100%), 13 (100%), and 15 (100%). The router not only classified them perfectly but also exhibited distinct and confident expert activation patterns. This proves the router learns fundamental physical dynamics rather than simply memorizing the training set, and that the experts have learned specialized knowledge.

On the Emergence of Classification Accuracy during Training (Q37)

To show how the router's classification ability arises, we tracked its accuracy during training. The results confirm this is an emergent property: accuracy starts low but steadily improves, eventually reaching 100%. This aligns with the intuition that the model learns to separate data distributions to optimize overall performance.

Responses to Specific Questions

On Experimental Setup

• Q1: FNO's performance collapses due to "negative transfer" from its shared kernel learning conflicting PDE patterns. DPOT mitigates this but doesn't fully solve it, as supported by our new experiments showing its performance wavers when more datasets are added.
• Q2: Noise is injected into all input frames, with variance proportional to the input's norm and applied uniformly across time.
• Q32: We followed DPOT's preprocessing, normalizing spatial and temporal domains; details will be added to the appendix.
• Q33: A fixed input length of T=10 was used for all experiments, consistent with the source datasets.
• Q34: Yes, we use auto-regressive rollouts. The reported L2RE is the average error over the entire sequence. We will add a requested error accumulation plot comparing MoE-POT and DPOT to the final paper.
• Q39, Q40: The pre-training dataset has ~30,000 samples; each downstream task has 2,000. We will provide a detailed table.

On Model & Methods

• Q8: You are correct; the output dimension after patchification should be $(H/P)\times(W/P)\times d$. We will correct this typo.
• Q9-11: $W_t$ is a learnable MLP for positional encoding, $W_t$ is a learnable linear transformation, and $\gamma$ are fixed, non-learnable Fourier features.
• Q12: You are correct; "Fourier Attention" is a misnomer. We will correct this to "Fourier Layer" throughout the paper.
• Q13: The grouping is performed on the channel dimension, where h is the number of heads.
• Q14: We chose lightweight CNNs for the experts as they effectively preserve spatial structure and proved more stable than alternatives like FNO in our early experiments.
• Q15: Batching is over samples, so x in Eq. (11) represents a single sample in the batch B. We will correct this confusing notation.
• Q16: e_i and Importance(e_i) are not the same. e_i is the i-th expert network, while Importance(e_i) is a scalar value—the total routing weight for that expert in a batch—used for the load-balancing loss.

On Results & Inference

• Q5: The per-step millisecond difference aggregates to a significant advantage in long rollouts or optimization loops. Our point is MoE's higher capacity and accuracy for a comparable computational cost.
• Q19, Q34: L2RE is the average error over the full auto-regressive rollout, not a single step.
• Q21, Q28: The values (0.1, 0.01) are the physical coefficients $\zeta$ and $\eta$ in the CNS equations; we will define all such constants in the appendix.
• Q22: During fine-tuning, we freeze the router and update only the expert networks to preserve the routing strategy while adapting expert knowledge.
• Q24: We agree the slope in Fig. 4(b) can be misleading. Our claim is about efficiency: better performance for a given number of activated parameters. We will rephrase this for clarity.
• Q25, Q26: We will add supplementary figures for all tasks and clearly define Y (mean expert preference) and I (input's expert preference) in the appendix.
• Q27, Q29-31: All requested details on T, u, D, R, initial conditions, and geometry will be added to the appendix.

---

Thanks again

We sincerely thank you again for your constructive and detailed feedback, which has helped us identify clear pathways to improve our paper. Should you have any further questions or require additional discussion, please don't hesitate to reach out. If we have adequately addressed your concerns, we would be grateful for your consideration in adjusting your evaluation score accordingly.

Dear Reviewer,

Thank you for your prompt reply and for giving us the opportunity to provide further clarification. We sincerely appreciate you taking the time to engage with our work so deeply. Due to the word and time constraints of the initial rebuttal, some experimental details were not fully elaborated. We are happy to provide more detailed clarification on these points below.

---

On the MoE-POT and Poseidon Comparison (Q1)

We ran comparisons under two distinct settings. We apologize that we only reported partial results in our initial response due to time limitations. Here are the complete details.

1. Poseidon's Setting

To align with the setup in the Poseidon paper, we modified MoE-POT's input/output to match theirs. We trained both models on 2,000 trajectories for 200 epochs.

2. Our Setting

We used the official open-source Poseidon repository to train their model under our loss function and experimental setting. In our initial experiment, we overlooked that the default training script for Poseidon uses only 128 trajectories, which led to an unfair comparison. We sincerely apologize for this oversight. We have now re-run the experiments with a full and fair comparison, including results on MoE-POT with varying numbers of trajectories. All models were trained for 200 epochs.

These corrected and expanded results provide a much clearer picture of performance under different data constraints.

On Performance with Increasing Dataset Heterogeneity (Q2)

To ensure a fair comparison, we selected all datasets from the pre-training and downstream task sets of the original DPOT paper. Specifically:

• 6 Datasets (Our paper's pre-training set): NS(1e-5), NS(1e-3), CNS(0.1, 0.01), SWE, DR, CFDBench.
• 10 Datasets (Added 4): NS(1e-4), CNS(1, 0.1), NavierStokes-2D (from PDEArena), NavierStokes-2D-conditioned (from PDEArena).
• 12 Datasets (Added 2 more): CNS(1, 0.01), CNS(0.1, 0.1).

For each experiment (6, 10, and 12 datasets), we pre-trained both DPOT-S and MoE-POT-S completely from scratch on the specified data mixture.

As the table shows, when the number of datasets increases, DPOT’s error rises noticeably on the SWE and DR datasets, demonstrating the negative transfer effect. In contrast, MoE-POT's performance remains highly stable, showcasing its superior ability to handle knowledge conflict from heterogeneous sources.

On Fine-Tuning with Limited Data Samples (Q3)

This is indeed a critical test for any foundation model. We evaluated our MoE-POT-S model on the Wave-Layer downstream task by fine-tuning it with a limited number of training trajectories. The table below shows the L2RE for various combinations of fine-tuning epochs and training trajectory counts.

These results confirm that MoE-POT maintains strong performance and learns effectively even when fine-tuned on a very small number of data samples, a key advantage for practical applications.

---

Thanks Again

We will integrate all of these new experiments and more detailed analyses into the final version of our paper. We will also ensure that Poseidon and other related algorithms are thoroughly discussed and cited in the related work section.

We sincerely hope these clarifications have adequately addressed your remaining concerns. We are grateful for your constructive feedback, which is truly invaluable in improving our work. If we have successfully addressed your points, we would be deeply grateful for your consideration in adjusting your evaluation score accordingly.

Dear Reviewer,

Thank you for your prompt and detailed follow-up. We appreciate you providing us with the opportunity to offer further clarifications. Your insightful questions are instrumental in strengthening our work, and we are happy to provide the additional results and details you requested.

---

On the Comparison with Poseidon (Q1)

Your question regarding the discrepancy between our reproduced results for Poseidon and those published in its original paper is indeed critical. After a careful investigation prompted by your question, we have identified a key factor causing this inconsistency: a difference in the evaluation protocol.

In Poseidon's official code repository, the Wave-Layer and Wave-Gauss tasks are evaluated at timestep 14. In contrast, our previous experiments followed a uniform standard used for other tasks in our framework, evaluating at the final timestep (T=20). As error accumulates over time, evaluating at the final timestep naturally yields higher error values. Furthermore, we used the L2RE norm for our error evaluation, whereas the original Poseidon paper uses the L1 norm.

To provide a completely fair evaluation that is directly comparable to the published literature, we have re-run the fine-tuning and testing for our MoE-POT model by strictly adhering to Poseidon's evaluation standard (evaluating at timestep 14 using the L1 norm**)**. All models were fine-tuned for 200 epochs.

The table below presents the performance of MoE-POT-S on a varying number of training trajectories under the Poseidon setting (L1 norm, evaluated at T=14):

As shown, under Poseidon's specific setting, its performance is superior to MoE-POT. Poseidon is an outstanding work that has had a significant positive impact on the AI4PDE field. However, as our paper's primary focus is on the DPOT-style pre-training and fine-tuning paradigm, we do not intend to claim superiority over Poseidon in its native setting. Our goal with these experiments is simply to provide the direct comparison you requested.

On the DPOT Training Protocol (Q2)

Yes, your understanding is correct.

While we are aware that the original DPOT paper reported results on 12 datasets, we retrained the DPOT-S model entirely from scratch on our mixed datasets containing 6, 10, and 12 datasets, respectively. This was done to ensure a controlled and fair comparison, isolating the architectural differences between DPOT and MoE-POT in handling "knowledge conflict" by training both models on the exact same data mixtures of progressively increasing heterogeneity.

Furthermore, all pre-trained models in our paper's main experiments were trained from scratch to fairly compare the performance of different architectures under identical conditions. The training scripts for both DPOT and MoE-POT in our work are based on the official DPOT repository, with the only difference being the model architecture itself. Our code has been submitted in the supplementary material for full transparency.

Below is the updated table, which includes data from the original DPOT paper for reference. We have corrected a minor data misalignment from our previous response—please consider this version definitive.

Both our re-implemented DPOT results and the original DPOT paper's data support the same conclusion: MoE-POT achieves better performance than DPOT, and critically, its performance remains robust as dataset heterogeneity increases.

Dear Reviewer,

Thank you for your detailed and insightful review. We are very grateful for your positive feedback on our paper's motivation and the novelty of our approach. Your comments are exceptionally helpful, and we agree that they point to clear ways we can improve the paper. We have carefully considered all your questions and provide the following responses, including new experimental results, which we hope will address your concerns.

---

On Clarity and Experimental Details

We completely agree that presenting experimental details with utmost clarity is crucial. Thank you for these excellent suggestions. We have uploaded our code for full transparency.

• Experimental Setup & Data Splits: As stated on Page 6 (Line 186), our datasets are from three widely-used public benchmarks: FNO, PDEBench, and CFDBench. We used the standard, publicly available train/validation/test splits and sample sizes for all benchmarks to ensure fair and reproducible comparisons, following the standard procedures of prior works like DPOT and FNO.
• Commitment to Improvement: We recognize that these details could be presented more clearly. In the final version, we will:
  • Add a comprehensive summary table in the appendix, detailing the train/val/test splits, sample counts, fine-tuning procedures, and data sources for all pre-training and downstream datasets. We will also include quantitative metrics of task difficulty (e.g., Reynolds number for NS equations) and discuss their implications for generalization.
  • Revise Figure 2 (right) into a single, unified plot with grouped bar charts. This will allow for a direct and easy comparison of expert activation patterns across different datasets, as you wisely suggested.
  • Improve Table 1's presentation by adding a summary figure in the main text that plots performance against the number of fine-tuning samples for MoE-POT vs. DPOT, visually highlighting the value added by our MoE architecture. The full table will be retained in the appendix.
  • Enlarger all figure labels and ticks to improve readability.

On Performance Improvement and Multi-Step Rollout Accuracy

You have raised a crucial point about the significance of our performance gains and the fairness of our comparisons. We would like to offer a detailed clarification.

• Reported Errors are for Multi-Step Rollouts: We would like to clarify a critical point regarding our evaluation: the reported errors in our tables are not for single-step prediction but are the average errors over long-term, auto-regressive rollouts. Our experimental setup follows the standard protocol used in DPOT and FNO. For instance, on the PDEBench-SWE dataset, we input the first 10 frames and auto-regressively predict the subsequent 90 frames, with the reported error being the average over these 90 predicted steps.
• The Value of "Marginal" Error Reduction: For spatio-temporal dynamics, even a small reduction in single-step prediction error is vital.
  • Significant Relative Improvement: As shown in Tables 1 and 2, our model achieves a substantial relative performance improvement on multiple tasks. In the field of neural operators, where current SOTA precision (1e-2 to 1e-4) still lags behind traditional solvers (<1e-7), every bit of error reduction is a meaningful step forward.
  • Importance of Compounded Error: A significant reduction in single-step error (e.g., 50% relative improvement) drastically suppresses the accumulation of errors during long-term rollout predictions. This leads to far more stable and accurate simulations in practical applications, directly addressing your concern about multi-step performance.
• Fair Comparison with DPOT:
  • Fine-tuning Details: As described on Page 7 (Lines 221-222) and Page 8 (Line 240), our fine-tuning protocol was applied consistently across all pre-trained models (DPOT and ours) to ensure a fair comparison. For Table 1, models were fine-tuned for 200 epochs; for Table 2, 500 epochs. The training set of the corresponding dataset was used for fine-tuning.
  • Convergence and Architectural Advantage: In all fine-tuning experiments, both MoE-POT and DPOT models reached convergence (typically before 180 epochs for Table 1 tasks and before 400 epochs for Table 2 tasks). Therefore, there was no possibility of DPOT "catching up" with more training. We argue the performance gap stems from a fundamental architectural advantage. As a dense model, DPOT's capacity is limited by its activated parameters. MoE-POT, through sparse activation, leverages a much larger total parameter count, giving it a higher capacity to absorb heterogeneous knowledge more effectively.

On Inference Efficiency and DPOT-L Comparison

• Advantage in Inference Efficiency: You correctly observe that MoE-POT-M and DPOT-M have similar inference times. This is precisely where our model's advantage lies. As shown in Table 3, MoE-POT-M (288M activated, 489M total params) has nearly the same inference time (16.6ms) as DPOT-M (158M activated, 16.7ms). However, as shown in Table 1, MoE-POT-M is far more accurate for this same inference cost. This demonstrates a significantly better efficiency-to-performance ratio and highlights the scalability advantage of the MoE architecture.
• DPOT-L Experimental Results: Due to computational constraints, we were unable to complete the full pre-training for a corresponding MoE-POT-L model. To address your question, we provide the zero-shot results for the official DPOT-L model below.

These results show that MoE-POT-M achieves comparable or superior performance to DPOT-L on 4 out of 6 datasets, using only about half the activated parameters (and thus, computational cost). This strongly supports our claims. We will add a more detailed comparison with DPOT-L in the final version.

On Interpretability and the MoE Mechanism

Thank you for appreciating the interpretability analysis and for pushing us to explore it further. These are excellent suggestions.

• Role of Shared Experts: The two shared experts are, by design, activated for every input. In our initial experiments, models without shared experts often ended up training unstable.
• Stress-Testing with Out-of-Distribution Data: In our original analysis (Appendix B.3), we used the training set to compute the mean expert preference vector for each dataset and the test set to calculate classification accuracy. To address your excellent question about OOD data, we conducted a new experiment testing the classification accuracy on two challenging downstream tasks from the Poseidon paper [1] that were not seen during pre-training.

The router achieves 100% classification accuracy on these unseen datasets. Furthermore, in Block-2, the expert network ratios for the Wave-Layer dataset were: Expert 1 (79%), 9 (31%), 11 (100%), 13 (100%), 15 (100%), others (0%). This strongly demonstrates that the gating network learns to distinguish fundamental properties of different PDEs and generalizes this capability to new, unseen tasks.

Why Not One-Hot Encode?:

• Our work follows the experimental paradigm and dataset of DPOT and FNO, which aims to simulate real-world scenarios (e.g., weather forecasting) where only solution trajectories are available, and the exact PDE parameters are unknown. The goal is for the model to implicitly infer the system's dynamics from observed data, which is precisely the role our router network successfully plays. While parameter-informed modeling via encoding is a valid and important research direction, our core contribution is the effective introduction of the MoE structure into this auto-regressive pre-training paradigm.

On Minor Questions

• FNO with MoE: During our research, we explored numerous combinations, including MoE with FNO and FFNO. These attempts were unsuccessful, leading to unstable training and large errors that could not be resolved by fine-tuning. The current MoE-POT architecture was the final stable and effective design that emerged from this rigorous process.
• Table 1 Caption: Thank you for pointing this out. The "-" indicates that the model's L2RE on that task exceeded 20, signifying that training diverged or failed completely. We will clarify this in the caption.
• Wp in Eqn. (3): You are correct. 'n' is the feature dimension of the positional encoding, which we typically set to be the same as the number of channels C. We will explicitly define this in the final version.

References:

[1] Poseidon: Efficient foundation models for pdes, NeurIPS 2024.

---

Thanks again

We sincerely thank you again for your constructive and detailed feedback, which has helped us identify clear pathways to improve our paper. Should you have any further questions or require additional discussion, please don't hesitate to reach out. If we have adequately addressed your concerns, we would be grateful for your consideration in adjusting your evaluation score accordingly.

Dear Reviewer,

Thank you for your prompt follow-up and for giving us the opportunity to provide further clarification. We apologize that our previous response was not entirely clear due to length constraints. We are grateful for the chance to address your remaining questions in more detail.

---

On the Nature of Rollout Errors (Q1)

To clarify directly: all results reported in our paper, including those in Table 1, are indeed multi-step rollout errors. We apologize if our previous phrasing caused any confusion and would like to explain our evaluation protocol precisely.

Our methodology, consistent with the DPOT benchmark, involves an auto-regressive rollout. For example, on the PDEBench-SWE dataset, we input the solution function images for frames 1, 2, ..., 10. We have the model predict frame 11. We then use frames 2, 3, ..., 10, 11 (the newly predicted frames) to predict frame 12. We then use frames 3, 4, ..., 10, 11 (the newly predicted frames), 12 (the newly predicted frames) to predict frame 13. This cycle continues until we reach frame 100. We then use the average L2RE error over frames 11, 12, ..., 100 as the model's error on the current dataset. This evaluation method is used in all experiments in the paper (including all tables), and is also the evaluation method used in the DPOT experiments.

We believe this process aligns with what you refer to as "multi-step rollout", where each prediction in the sequence is a "single-step prediction".You are absolutely right that a small reduction in single-step error leads to a significant reduction in the cumulative multi-step error. This effect is very pronounced because errors compound with each new prediction. To illustrate this with the SWE dataset:

As the table shows:

1. The error for both models increases with more rollout steps, confirming the effect of error accumulation.
2. The performance gap between DPOT-S and MoE-POT-S widens from a ~25% relative difference at frame 50 to a ~45% difference at frame 100. This super-linear divergence in error underscores the critical impact of reducing single-step error.

Your suggestion is invaluable for improving the paper's clarity. We commit to explicitly stating this evaluation protocol in the main text and table captions in the final version. We will also add this detailed breakdown of per-frame error analysis to the appendix to avoid any ambiguity for future readers.

---

On Fine-tuning Performance vs. Number of Examples (Q2)

We sincerely apologize for misinterpreting your previous question. Your intuition is spot on: increasing the number of fine-tuning samples can indeed allow a model like DPOT to eventually match the performance of MoE-POT.

To address your question directly, we conducted new experiments to illustrate this relationship. We tested performance on two tasks: a downstream task closely related to our pre-training data (NS 1e-4) and an out-of-distribution task from Poseidon [1] (Wave-Layer).

• Experiment 1: In-Distribution Fine-tuning on NS (1e-4) (500 Epochs)

• Experiment 2: Out-of-Distribution Fine-tuning on Wave-Layer (500 Epochs)

As demonstrated in the experiments above, and as you might expect, model performance improves on downstream tasks as the number of samples increases. Therefore, DPOT with more samples will outperform MoE-POT. In this scenario, a high-performing model requires less data to achieve similar performance, or achieves better performance with the same number of samples. Therefore, in all downstream task experiments in this article, we use a fixed number of samples.

We are very grateful for this insightful question and will incorporate these new results and analysis into the final version of our paper.

References:

[1] Pdebench: An extensive benchmark for scientific machine learning, NIPS 2022.

---

We sincerely thank you once again for your detailed and constructive engagement with our work. We hope these detailed explanations fully address your remaining concerns. Should you have any further questions, please do not hesitate to let us know. We remain available for any further discussion. We are very grateful for your time and guidance, and we would be deeply appreciative if you would consider our clarifications in your final evaluation of our work.

Dear Reviewer,

Thank you for your thoughtful and constructive feedback. We appreciate your time and have carefully considered your suggestions. We hope the following clarifications address your concerns and further highlight the significance of our work.

---

On the Novelty and Uniqueness of MoE-POT (Weakness 1 & Questions 3)

Our core innovation is not merely substituting a layer but systematically addressing the critical bottleneck of data heterogeneity in large-scale PDE pre-training. MoE-POT is presented as a well-validated and highly interpretable solution to this specific, previously underexplored challenge.

• Problem-Driven Innovation: Prior works have overlooked the "knowledge conflict" that arises from mixing diverse PDEs. Our experiments show this is a critical issue (Figure 2, left), as dense models suffer a drastic performance drop when trained on heterogeneous datasets.
• Unique and Validated Architectural Solution: The uniqueness of MoE-POT stems from its "divide-and-conquer" strategy, enabled by the combination of routed and shared experts. The routed experts specialize in the unique characteristics of specific PDEs, while the shared experts capture common physical principles. Crucially, we validate this mechanism's effectiveness through interpretability analysis (Figure 2, right; Figure 4c), demonstrating that the model learns to dynamically select experts based on the PDE type with an accuracy of up to 98%.
• Rigorous Experimental Exploration: The current MoE-POT architecture is the result of extensive experimentation. We explored numerous combinations of MoE with other neural operators (including MPP, FNO, and FFNO) and various MoE structures (e.g., without shared experts, using MLP-based experts instead of CNNs). These alternative designs frequently failed, resulting in models with substantial errors on certain datasets that could not be rectified through fine-tuning. The final version of MoE-POT is the stable and effective architecture that emerged from this rigorous process.

Therefore, the novelty of MoE-POT is not in an isolated architectural tweak but in its systematic demonstration that a sparse expert architecture is an effective and interpretable solution to the data heterogeneity challenge in PDE pre-training, supported by comprehensive evidence that prior models have not explored.

On the Cost-Benefit Analysis and Comparison with DPOT-L (Weakness 2 & Question 1)

Thank you for your detailed questions regarding the model's efficiency and training costs. We would like to clarify the cost-benefit trade-off of the MoE architecture and address the comparison with DPOT models.

• Clarifying the Comparison: You rightly point out the comparison between our MoE-POT-M (489M total, 288M activated) and DPOT-L (493M total & activated). Our key argument is that the MoE architecture decouples total parameters (storage cost) from activated parameters (computational cost), achieving superior efficiency.
• Greater Parameter Efficiency & Faster Inference: Our MoE-POT-S (90M activated) outperforms the larger DPOT-M (122M activated) on several tasks. Furthermore, MoE-POT-M (489M total) has an inference time comparable to the much smaller DPOT-M and is 32% faster than the similarly sized DPOT-L, a critical advantage for iterative PDE solving.
• Training Overhead and DPOT-L Results: We acknowledge that the large size of the L-series models made it challenging to complete the full pre-training for MoE-POT-L within our available computational resources, which is why a direct comparison was not included in the main table. We also recognize that MoE requires more memory to store the total parameters. However, we argue this is a highly valuable trade-off for achieving superior model performance, stronger generalization, and critically, higher inference efficiency. This rationale is widely accepted in other fields, such as large language models. In practical, large-scale deployments, memory usage can be optimized by loading shared expert networks that are used across multiple parallel models.

To further address your concerns, we provide supplementary experimental results (L2RE) for DPOT-L below:

As the table shows: 1) DPOT's performance improves marginally when scaling from 122M to 493M parameters. 2) When comparing DPOT-L (493M activated) with MoE-POT-M (288M activated), DPOT-L performs best on two datasets, while MoE-POT is superior on the other four. This indicates that MoE-POT achieves comparable, and in many cases better, performance with nearly half the activated parameters, which aligns with our claims. We will clarify this cost-benefit analysis and incorporate more detailed results for L-sized models in the final version of the paper.

On the Lack of Theoretical Analysis (Weakness 3)

We fully agree with your assessment and have candidly noted this in Section 6, "Limitations and Conclusions". The current work focuses on empirically demonstrating the feasibility and effectiveness of the MoE architecture for PDE pre-training. Our results, particularly the 98% dataset classification accuracy, provide strong empirical evidence for the mechanism that the router network can identify PDE types. Investigating the underlying mathematical principles and theoretical guarantees is a valuable and important direction for future research, which we are excited to pursue.

On Effectively Handling Diverse PDE Datasets (Question 2)

Thank you for this insightful question. We believe the existing evidence in our paper robustly demonstrates that MoE-POT handles heterogeneous datasets more effectively than dense models like DPOT, due to its inherent architectural advantages.

• Fundamental Limitation of Dense Models: Our work highlights a core weakness of dense models. As shown in Figure 2 (left), when training on PDEs with disparate properties, the performance of dense model FNO degrades sharply, with errors increasing by up to 5000%.
• MoE-POT's Architectural Advantage: MoE-POT is explicitly designed to resolve this "knowledge conflict," and its router network intelligently identifies different PDE tasks and assigns them to specialized expert networks. Our interpretability analysis provides direct evidence for this mechanism: the router's selection weights alone can infer the input data's PDE type with 97.7% accuracy. This confirms that MoE-POT can effectively distinguish and process heterogeneous data in a way that dense architectures like DPOT cannot.
• Superior Performance as a Result: This architectural superiority translates directly into leading performance. In the zero-shot comparison in Table 1, our MoE-POT-S (90M activated params) reduces the error on the PDEBench-CNS dataset by 57% compared to the larger DPOT-M (122M activated params) after pre-training on six heterogeneous datasets. This clearly shows MoE-POT learns more effectively from diverse data with lower computational cost.

To further address your concern and illustrate the relationship between the number of datasets and performance for a dense model, we conducted the following supplementary experiment with DPOT-S, measuring its performance on the six core datasets as we increased the total number of pre-training datasets.

These results show that: 1) Regardless of the number of training datasets, DPOT-S does not match the performance of MoE-POT-S. 2) For a fixed-capacity dense model, simply increasing the number of training datasets does not guarantee better performance. In fact, performance degraded on four of the six datasets when scaling from 6 to 12. This reinforces our central claim that mixing disparate PDE datasets in a dense model leads to performance degradation and highlights the rationale for our MoE architecture. We will add a more detailed analysis and these experiments to the final version.

---

Thanks again

We sincerely thank you again for your constructive and detailed feedback, which has helped us identify clear pathways to improve our paper. Should you have any further questions or require additional discussion, please don't hesitate to reach out. If we have adequately addressed your concerns, we would be grateful for your consideration in adjusting your evaluation score accordingly.

To provide more definitive empirical evidence for our claims and to comprehensively address your question, we conducted a rigorous set of experiments comparing MoE-POT-S and the dense DPOT-S model as the heterogeneity of the pre-training data was increased. Due to the time constraints of the rebuttal period, these experiments were not yet complete at the time of our initial response. We are now pleased to share the finalized results. Our conclusions remain unchanged and are, in fact, strengthened by this new data.

Experimental Setup

For full transparency and fairness, we selected all datasets from the original DPOT paper's pre-training and downstream task sets. The specific mixtures were as follows:

• 6 Datasets (This paper's pre-training set):
  • NS(1e-5); 2. NS(1e-3); 3. CNS(0.1, 0.01); 4. SWE; 5. DR; 6. CFDBench
• 10 Datasets (Added 4 diverse tasks):
  • The 6 datasets above, plus: 7. NS(1e-4); 8. CNS(1, 0.1); 9. NavierStokes-2D (from PDEArena); 10. NavierStokes-2D-conditioned (from PDEArena)
• 12 Datasets (Added 2 more diverse tasks):
  • The 10 datasets above, plus: 11. CNS(1, 0.01); 12. CNS(0.1, 0.1)

Crucially, for each configuration (6, 10, and 12 datasets), both DPOT-S and MoE-POT-S were pre-trained entirely from scratch on the specified data mixture; this was not a continuation of a previous training run.

Complete Experimental Results (L2RE)

Below is the updated table, which includes data from the original DPOT paper for reference. We have corrected a minor data misalignment from our previous response—please consider this version definitive.

Analysis

These comprehensive results lead to two clear conclusions:

1. Superior Performance of MoE-POT: Regardless of the number of training datasets, MoE-POT-S consistently outperforms DPOT-S on a majority of the evaluation tasks.
2. Robustness to Heterogeneity**:** For a fixed-capacity dense model like DPOT-S, simply increasing the number and diversity of training datasets does not guarantee better performance. Notably, as the dataset mixture grows more heterogeneous (from 6 to 12), the performance of DPOT-S degrades significantly on the SWE and DR datasets. In stark contrast, MoE-POT-S maintains highly stable.

This experiment robustly reinforces our central claim: mixing disparate PDE datasets in a dense model leads to performance degradation due to "knowledge conflict," and it powerfully highlights the rationale for our sparse MoE architecture, which is specifically designed to overcome this challenge. We will add a more detailed analysis and these complete experimental results to the final version of the paper.

Dear Reviewer,

Thank you for your thorough review and positive assessment of our work. We are grateful for your insightful feedback and for recognizing the novelty and promise of applying a Mixture-of-Experts (MoE) architecture to PDE solvers. Your questions have been instrumental in helping us clarify our contributions. We have carefully considered your points and provide the following detailed responses.

---

On the Choice of Baselines (DPOT vs. Poseidon/MPP) (Weakness 1)

Thank you for this valuable suggestion. We chose DPOT as our primary multi-physics baseline for three main reasons:

• Clear Attribution of Gains: Our MoE-POT architecture is built upon the DPOT model, where we specifically replace the dense feed-forward network with our sparse expert layer. This design allows us to cleanly attribute performance differences to the advantages conferred by the MoE architecture itself, providing a direct and compelling validation of our core method. Comparing with structurally disparate models like Poseidon would make it difficult to isolate the source of any performance gains.
• Differences in Experimental Setup: Our work, following DPOT, uses an auto-regressive paradigm: predicting the next frame from a sequence of previous frames. In contrast, models like MPP and Poseidon have different setups. For instance, Poseidon takes problem parameters, boundary conditions, and initial conditions as input to predict a future state. As noted in the original DPOT paper, DPOT outperformed MPP in most scenarios. Furthermore, Poseidon requires PDE parameters as input, which are not available in the public datasets we use (FNO, PDEBench, CFDBench), making a direct comparison infeasible under our experimental protocol.
• Contemporaneous Work: Poseidon is an excellent and cutting-edge model. At the time we began our experimental cycle, this work had not yet been publicly released and was developed concurrently with our own. Therefore, we were unable to include it in our initial baseline comparisons.

We agree that comparison with the latest state-of-the-art models is crucial. To address your concern, we conducted supplementary fine-tuning experiments on two challenging downstream tasks from the Poseidon paper, Wave-Layer and Wave-Gauss, using the $L_2$ relative error metric. We considered two settings:

1. Our Setting (Auto-regressive, no parameters):

2. Poseidon's Setting (Parameter-informed):

These results indicate that each model excels in its native experimental setting, suggesting they have different scopes of application rather than one being definitively superior. Due to the time constraints of the rebuttal period, we could not complete an exhaustive comparison. However, if the paper is accepted, we commit to including a more thorough comparison with Poseidon and other recent SOTA models in the final version, along with citations to these important works.

On Comparison with Advanced FNO Variants (Weakness 2)

Thank you for your close reading. We completely agree that a comparison limited to the original FNO would be insufficient. In fact, in Table 1 of our paper, we have already conducted a comprehensive comparison against a range of more advanced single-task baselines. This includes not only the FFNO you mentioned but also GK-Transformer, OFormer, and GNOT. The results demonstrate that our fine-tuned MoE-POT model outperforms these state-of-the-art single-task baselines on the vast majority of datasets, validating the effectiveness of our pre-training framework.

On the Difficulty of Pre-training Datasets (Weakness 3)

We appreciate your concern regarding the choice of datasets. We selected them based on the following considerations:

• Standard and Fair Benchmarking: The six datasets we use are sourced from three widely-accepted public benchmarks: FNO, PDEBench, and CFDBench. They are also the datasets used for pre-training by DPOT. Using these standard benchmarks ensures that our results are fair, reproducible, and comparable with existing and future research.
• The Core Challenge is "Heterogeneity," Not "Single-Task Difficulty": The central thesis of our paper is to solve the performance bottleneck that arises when a single model must handle a heterogeneous mixture of PDEs. As we explicitly demonstrate in Figure 2 (left), even on these individually "less difficult" datasets, mixing them during training causes the performance of a traditional dense model to collapse. Therefore, the advantage of our architecture is its ability to effectively manage data diversity and heterogeneity, a property that is best demonstrated through the diverse dataset combination we selected.

On Fine-Tuning Epochs and the Overall Compute Budget (Weakness 4)

This is a very insightful and fair point. We offer two clarifications:

• Relative Cost of Fine-tuning: Our pre-training phase runs for 1000 epochs, making the 200-epoch fine-tuning a reasonable fraction of the total training. More importantly, the value of fine-tuning lies in its ability to rapidly adapt general knowledge to specific tasks. The results in Table 2 clearly show that pre-trained models ("w/ Pre-train") significantly outperform models trained from scratch ("w/o Pre-train") on all downstream tasks, proving the immense value and knowledge transfer capability of pre-training.
• Model Capacity vs. Training Duration: The idea of aligning the total compute budget is an excellent one. However, we argue that smaller models like FFNO, even if granted equivalent training time, would quickly hit a performance ceiling due to their limited model capacity, potentially even degrading from overfitting. Indeed, all baseline results in our paper represent the best performance achieved during their training until full convergence. In contrast, our MoE architecture's success stems from having a very large total parameter count (ensuring high capacity) while keeping computational costs low via sparse activation. Our scaling experiments (Sec 5.3) further support the conclusion that larger model capacity leads to better performance.

Clarification of "Feature Competition" and "Knowledge Illusion" (Question)

Thank you for pointing out this imprecise phrasing. These terms were intended as intuitive descriptors for the well-established phenomenon known as "Negative Transfer" in multitask learning.

• "Feature Competition" refers to the scenario where a dense model, forced to learn fundamentally different tasks (e.g., fluid dynamics of NS equations vs. reaction-diffusion of DR equations) with a single set of shared parameters, experiences conflicting gradient updates. This "tug-of-war" between tasks prevents the model from fully learning any single task, thereby harming overall performance.
• "Knowledge Illusion" describes the outcome of this phenomenon: the model appears to have learned something about all tasks, but this knowledge is shallow and compromised, failing to capture the core physics of each task and leading to poor generalization.

We acknowledge that we did not provide direct citations for these descriptive terms. In the final version of the paper, we will replace them with the more standard academic terms "negative transfer" or "inter-task interference" and cite foundational multitask learning literature, such as [1], to make our argument more rigorous.

References:

[1] Multitask learning, Caruana, 1997

---

Thanks again

We sincerely thank you again for your constructive and detailed feedback, which has helped us identify clear pathways to improve our paper. Should you have any further questions or require additional discussion, please don't hesitate to reach out. If we have adequately addressed your concerns, we would be grateful for your consideration in adjusting your evaluation score accordingly.

We agree that a robust comparison with state-of-the-art models like Poseidon is crucial for positioning our work. Due to the time constraints of the rebuttal period, our initial experiments were incomplete. We have now finalized our supplementary experiments comparing MoE-POT with Poseidon, and we are happy to report that our initial conclusions remain consistent.

---

We conducted comparisons under two distinct settings to ensure a comprehensive evaluation:

1. Poseidon's Setting

To align with the methodology in the Poseidon paper, we modified MoE-POT’s architecture to match their input-output structure (i.e., taking problem parameters and initial conditions as input to predict a specific future frame). For this comparison, we fine-tuned each model on 2,000 trajectories for 200 epochs.

2. Our Setting

For this setting, we used the official open-source Poseidon repository and adapted its loss calculation to fit our auto-regressive framework. All models were trained for 200 epochs.

In our initial analysis, we overlooked that the official Poseidon training script defaults to using only 128 trajectories, which resulted in an unfair comparison. We sincerely apologize for this oversight. We have now included a more complete and equitable comparison, detailing the performance of MoE-POT-S with varying numbers of training trajectories to demonstrate its effectiveness in low-data regimes.

These updated and expanded results compellingly show that each model excels within its native experimental setting. This suggests they are powerful tools with different, valuable scopes of application rather than one being definitively superior across all paradigms. Due to the time constraints of the rebuttal period, we could not complete an exhaustive comparison. However, if the paper is accepted, we commit to including a more thorough comparison with Poseidon and other recent SOTA models in the final version, along with citations to these important works.

Dear Reviewer wk8j,

Thank you very much for taking the time to read our rebuttal and for promptly acknowledging it. We truly appreciate your engagement with our work during the review process.

If you have any further thoughts, questions, or suggestions---either regarding our current response or the broader direction of the work---we would be genuinely grateful to hear them. We highly value your perspective and would welcome any opportunity for continued discussion or clarification.

Best,

Authors