Learning Operators with Coupled Attention

TL;DR

LOCA is a novel operator learning method using coupled attention mechanisms to effectively model complex systems, demonstrating state-of-the-art accuracy and robustness in various scientific applications.

Contribution

The paper introduces LOCA, a new operator learning architecture with coupled attention, providing theoretical guarantees and improved performance over existing methods.

Findings

Achieves state-of-the-art accuracy in PDE-based operator learning

Demonstrates robustness to noisy input data

Maintains low error spread in out-of-distribution predictions

Abstract

Supervised operator learning is an emerging machine learning paradigm with applications to modeling the evolution of spatio-temporal dynamical systems and approximating general black-box relationships between functional data. We propose a novel operator learning method, LOCA (Learning Operators with Coupled Attention), motivated from the recent success of the attention mechanism. In our architecture, the input functions are mapped to a finite set of features which are then averaged with attention weights that depend on the output query locations. By coupling these attention weights together with an integral transform, LOCA is able to explicitly learn correlations in the target output functions, enabling us to approximate nonlinear operators even when the number of output function in the training set measurements is very small. Our formulation is accompanied by rigorous approximation theoretic guarantees on the universal expressiveness of the proposed model. Empirically, we evaluate the performance of LOCA on several operator learning scenarios involving systems governed by ordinary and partial differential equations, as well as a black-box climate prediction problem. Through these scenarios we demonstrate state of the art accuracy, robustness with respect to noisy input data, and a consistently small spread of errors over testing data sets, even for out-of-distribution prediction tasks.