Multi-Level Monte Carlo Training of Neural Operators

TL;DR

This paper introduces a Multi-Level Monte Carlo training method for neural operators that reduces computational costs by leveraging multi-resolution data, maintaining high accuracy for PDE-related operator learning.

Contribution

It presents a novel MLMC-based training framework applicable to various neural operator architectures, improving efficiency over traditional single-resolution methods.

Findings

Enhanced computational efficiency demonstrated across multiple models

Existence of a Pareto curve between accuracy and computational time

Effective use of multi-resolution data for training neural operators

Abstract

Operator learning is a rapidly growing field that aims to approximate nonlinear operators related to partial differential equations (PDEs) using neural operators. These rely on discretization of input and output functions and are, usually, expensive to train for large-scale problems at high-resolution. Motivated by this, we present a Multi-Level Monte Carlo (MLMC) approach to train neural operators by leveraging a hierarchy of resolutions of function discretization. Our framework relies on using gradient corrections from fewer samples of fine-resolution data to decrease the computational cost of training while maintaining a high level accuracy. The proposed MLMC training procedure can be applied to any architecture accepting multi-resolution data. Our numerical experiments on a range of state-of-the-art models and test-cases demonstrate improved computational efficiency compared to traditional single-resolution training approaches, and highlight the existence of a Pareto curve between accuracy and computational time, related to the number of samples per resolution.