Research map / Architectures and efficiency / Mixture of experts
Mixture-of-Experts Operator Transformer for Large-Scale PDE Pre-Training
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes MoE-POT, a sparse mixture-of-experts neural operator (built on DPOT with CNN experts, 16 routed experts with top-4 routing plus 2 shared experts) for auto-regressive PDE pre-training across six public datasets, reporting lower zero-shot error at equal activated parameters.
Problem
Mixed-dataset pre-training of neural operators on heterogeneous PDE types causes high errors, and scaling dense models increases inference cost.
Approach
Replace the dense feed-forward layer of a Fourier-based operator transformer (DPOT) with a sparse MoE layer combining top-k routed CNN experts and always-on shared experts, so capacity scales with total parameters while activated parameters stay fixed; router gating decisions also reveal dataset identity.
Key results
- Ours-S (90M activated) reduces zero-shot L2RE on PDEBench CNS(0.1,0.01) by 57% relative to DPOT-M (122M activated), reaching 0.00959 vs 0.0224.
- After fine-tuning, MoE-POT-M achieves best results on 5 of 6 datasets, and MoE-POT-M (288M activated) at 16.6 ms single-step inference is about 68% of DPOT-L (493M) inference time of 24.3 ms.
- Router gating decisions in block 2 classify the input dataset with 97.7% accuracy (abstract reports 98%).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Increase routed experts from 16 to 32 (ablation) | MoE-POT-T · NS(1e-3) and NS(1e-5) | L2RE: NS(1e-3): 0.06920 -> 0.06680; NS(1e-5): 0.00762 -> 0.00857 | Mixed |
| Increase routed experts from 16 to 8 (ablation) | MoE-POT-T · NS(1e-3) and NS(1e-5) | L2RE: NS(1e-3): 0.06920 -> 0.06833; NS(1e-5): 0.00762 -> 0.00773 | Mixed |
| Increase Top-K from 2 to 4 (ablation) | MoE-POT-T · NS(1e-3) | L2RE: 0.06983 -> 0.06920 | Helped |
| Increase Top-K from 1 to 4 (ablation) | MoE-POT-T · NS(1e-3) | L2RE: 0.10108 -> 0.06920 | Helped |
| Set number of heads to 2 instead of 4 (ablation) | MoE-POT-T · NS(1e-3) | L2RE: 0.06920 -> 0.06748 | Helped |
| Set number of heads to 8 instead of 4 (ablation) | MoE-POT-T · NS(1e-3) | L2RE: 0.06920 -> 0.06963 | Hurt |
| Set patch size to 4 instead of 8 (ablation) | MoE-POT-T · NS(1e-3) | L2RE: 0.06920 -> 0.06226 | Helped |
| Set patch size to 16 instead of 8 (ablation) | MoE-POT-T · NS(1e-3) | L2RE: 0.06920 -> 0.08964 | Hurt |
| Use MoE-POT-M rather than DPOT-M for zero-shot pre-training | PDEBench CNS(0.1,0.01) · CNS(0.1,0.01) | L2RE: 0.0224 -> 0.00914 | Helped |
| Use MoE-POT-S rather than DPOT-M for zero-shot pre-training | SWE · SWE | L2RE: 0.00247 -> 0.00289 | Hurt |
| Use MoE-POT-M rather than DPOT-M for zero-shot pre-training | DR · DR | L2RE: 0.0288 -> 0.0300 | Hurt |
| Use MoE-POT-S rather than DPOT-T for zero-shot pre-training | CFDBench · CFDBench | L2RE: 0.00673 -> 0.00448 | Helped |
Surprising findings
- Dense FNO-M fails completely (error above 20) on some datasets during mixed pre-training, indicating training collapse under heterogeneity.
Limitations
- The mathematical basis of the router's dataset classification is not analyzed.
- Parameter-informed setting results show MoE-POT underperforming Poseidon on Wave-Layer and Wave-Gauss.
- Results rely on a fixed set of six pre-training datasets and resolution 128 with padding and masking.
Experimental setup
Datasets: FNO NS (1e-3) FNO NS (1e-5) PDEBench CNS(0.1,0.01) PDEBench SWE PDEBench DR CFDBench NS (1e-4) CNS (1,0.01) PDEArena
Benchmarks: FNO NS (1e-3) FNO NS (1e-5) PDEBench CNS PDEBench SWE PDEBench DR CFDBench PDEArena Poseidon Wave-Layer Poseidon Wave-Gauss
Metrics: L2RE inference time (ms) dataset classification accuracy
Baselines: FNO UNet FFNO GK-Transformer OFormer GNOT DPOT Geo-FNO Poseidon MPP
Builds on: DPOT Fourier neural operator mixture of experts Vision Transformer
Closest papers by meaning
- Dense Backpropagation Improves Training for Sparse Mixture-of-Experts (NeurIPS 2025 · similarity 0.89)
- DPOT: Auto-Regressive Denoising Operator Transformer for Large-Scale PDE Pre-Training (ICML 2024 · similarity 0.88)
- Scaling Beyond the GPU Memory Limit for Large Mixture-of-Experts Model Training (ICML 2024 · similarity 0.86)
- DeRS: Towards Extremely Efficient Upcycled Mixture-of-Experts Models (CVPR 2025 · similarity 0.86)
- Data-Efficient Operator Learning via Unsupervised Pretraining and In-Context Learning (NeurIPS 2024 · similarity 0.86)
- Joint MoE Scaling Laws: Mixture of Experts Can Be Memory Efficient (ICML 2025 · similarity 0.86)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.