FreqMoE: Dynamic Frequency Enhancement for Neural PDE Solvers

TL;DR

FreqMoE introduces a progressive training framework that enhances neural PDE solvers by efficiently modeling high-frequency signals, leading to significant accuracy improvements and parameter reduction.

Contribution

It presents a novel frequency-dependent mixture of experts approach that extends low-frequency learned weights to high frequencies, improving efficiency and performance in neural PDE solvers.

Findings

Achieves up to 16.6% accuracy improvement over dense FNO.

Uses only 2.1% of parameters compared to dense FNO.

Demonstrates stability and generalization across various PDEs and grid structures.

Abstract

Fourier Neural Operators (FNO) have emerged as promising solutions for efficiently solving partial differential equations (PDEs) by learning infinite-dimensional function mappings through frequency domain transformations. However, the sparsity of high-frequency signals limits computational efficiency for high-dimensional inputs, and fixed-pattern truncation often causes high-frequency signal loss, reducing performance in scenarios such as high-resolution inputs or long-term predictions. To address these challenges, we propose FreqMoE, an efficient and progressive training framework that exploits the dependency of high-frequency signals on low-frequency components. The model first learns low-frequency weights and then applies a sparse upward-cycling strategy to construct a mixture of experts (MoE) in the frequency domain, effectively extending the learned weights to high-frequency regions. Experiments on both regular and irregular grid PDEs demonstrate that FreqMoE achieves up to 16.6% accuracy improvement while using merely 2.1% parameters (47.32x reduction) compared to dense FNO. Furthermore, the approach demonstrates remarkable stability in long-term predictions and generalizes seamlessly to various FNO variants and grid structures, establishing a new ``Low frequency Pretraining, High frequency Fine-tuning'' paradigm for solving PDEs.