Maximal Update Parametrization and Zero-Shot Hyperparameter Transfer for Fourier Neural Operators

TL;DR

This paper introduces $Transfer-FNO, a zero-shot hyperparameter transfer method for Fourier Neural Operators, enabling efficient scaling to complex PDEs by transferring hyperparameters from smaller models, validated through extensive experiments.

Contribution

We develop a $Transfer-FNO technique based on $P that allows hyperparameters to be transferred across models with different Fourier modes, reducing tuning costs for large FNOs.

Findings

Transfer-FNO reduces hyperparameter tuning costs for large models.

It maintains or improves accuracy when scaling to complex PDEs.

Empirical validation across various PDEs supports effectiveness.

Abstract

Fourier Neural Operators (FNOs) offer a principled approach for solving complex partial differential equations (PDEs). However, scaling them to handle more complex PDEs requires increasing the number of Fourier modes, which significantly expands the number of model parameters and makes hyperparameter tuning computationally impractical. To address this, we introduce $\mu$Transfer-FNO, a zero-shot hyperparameter transfer technique that enables optimal configurations, tuned on smaller FNOs, to be directly applied to billion-parameter FNOs without additional tuning. Building on the Maximal Update Parametrization ($\mu$P) framework, we mathematically derive a parametrization scheme that facilitates the transfer of optimal hyperparameters across models with different numbers of Fourier modes in FNOs, which is validated through extensive experiments on various PDEs. Our empirical study shows that Transfer-FNO reduces computational cost for tuning hyperparameters on large FNOs while maintaining or improving accuracy.