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

TL;DR

This paper introduces MoE-POT, a sparse Mixture-of-Experts Transformer for large-scale PDE pre-training that efficiently scales parameters, reduces inference costs, and improves zero-shot error on diverse PDE datasets.

Contribution

The paper presents a novel MoE architecture with dynamic expert routing and shared experts, enhancing PDE pre-training efficiency and accuracy over existing dense models.

Findings

Achieves up to 40% reduction in zero-shot error with 90M activated parameters.

Effectively captures PDE dataset features through interpretability analysis.

Reduces inference costs compared to dense models.

Abstract

Pre-training has proven effective in addressing data scarcity and performance limitations in solving PDE problems with neural operators. However, challenges remain due to the heterogeneity of PDE datasets in equation types, which leads to high errors in mixed training. Additionally, dense pre-training models that scale parameters by increasing network width or depth incur significant inference costs. To tackle these challenges, we propose a novel Mixture-of-Experts Pre-training Operator Transformer (MoE-POT), a sparse-activated architecture that scales parameters efficiently while controlling inference costs. Specifically, our model adopts a layer-wise router-gating network to dynamically select 4 routed experts from 16 expert networks during inference, enabling the model to focus on equation-specific features. Meanwhile, we also integrate 2 shared experts, aiming to capture common properties of PDE and reduce redundancy among routed experts. The final output is computed as the weighted average of the results from all activated experts. We pre-train models with parameters from 30M to 0.5B on 6 public PDE datasets. Our model with 90M activated parameters achieves up to a 40% reduction in zero-shot error compared with existing models with 120M activated parameters. Additionally, we conduct interpretability analysis, showing that dataset types can be inferred from router-gating network decisions, which validates the rationality and effectiveness of the MoE architecture.