Learning Operators with Coupled Attention

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.

Cited by in corpus (11)

• Reliable extrapolation of deep neural operators informed by physics or sparse observations
• In-Context Operator Learning with Data Prompts for Differential Equation Problems
• SVD Perspectives for Augmenting DeepONet Flexibility and Interpretability
• Scalable Uncertainty Quantification for Deep Operator Networks using Randomized Priors
• Adaptive physics-informed neural operator for coarse-grained non-equilibrium flows
• A Mathematical Guide to Operator Learning
• Variational operator learning: A unified paradigm marrying training neural operators and solving partial differential equations
• Multi-scale Time-stepping of Partial Differential Equations with Transformers
• MCMC-Net: Accelerating Markov Chain Monte Carlo with Neural Networks for Inverse Problems
• A Deep Learning approach for parametrized and time dependent Partial Differential Equations using Dimensionality Reduction and Neural ODEs
• Boundary-to-Solution Mapping for Groundwater Flows in a Toth Basin