# Learning Operators with Coupled Attention

arXiv: https://arxiv.org/abs/2201.01032 · HTML: https://mlsanity.com/paper/2201.01032/

## Key facts

- Cited by 90 papers in the map (data through 14 Oct 2026).
- Builds on 30 of its 72 references that are in the map.
- First posted on arXiv 4 Jan 2022.

## Builds on (most cited first, 30 of 30)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Neural Machine Translation by Jointly Learning to Align and Translate](https://mlsanity.com/paper/1409.0473/) (2014)
- [Gaussian Error Linear Units (GELUs)](https://mlsanity.com/paper/1606.08415/) (2016)
- [Fourier Neural Operator for Parametric Partial Differential Equations](https://mlsanity.com/paper/2010.08895/) (2020)
- [Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention](https://mlsanity.com/paper/2006.16236/) (2020)
- [Rethinking Attention with Performers](https://mlsanity.com/paper/2009.14794/) (2020)
- [Linformer: Self-Attention with Linear Complexity](https://mlsanity.com/paper/2006.04768/) (2020)
- [Image Transformer](https://mlsanity.com/paper/1802.05751/) (2018)
- [Neural Operator: Graph Kernel Network for Partial Differential Equations](https://mlsanity.com/paper/2003.03485/) (2020)
- [AST: Audio Spectrogram Transformer](https://mlsanity.com/paper/2104.01778/) (2021)
- [Learning the solution operator of parametric partial differential equations with physics-informed DeepOnets](https://mlsanity.com/paper/2103.10974/) (2021)
- [Nystr"omformer: A Nystr"om-Based Algorithm for Approximating Self-Attention](https://mlsanity.com/paper/2102.03902/) (2021)
- [Neural Operator: Learning Maps Between Function Spaces](https://mlsanity.com/paper/2108.08481/) (2021)
- [Multipole Graph Neural Operator for Parametric Partial Differential Equations](https://mlsanity.com/paper/2006.09535/) (2020)
- [Choose a Transformer: Fourier or Galerkin](https://mlsanity.com/paper/2105.14995/) (2021)
- [Model Reduction and Neural Networks for Parametric PDEs](https://mlsanity.com/paper/2005.03180/) (2020)
- [Multiwavelet-based Operator Learning for Differential Equations](https://mlsanity.com/paper/2109.13459/) (2021)
- [Transformer Dissection: A Unified Understanding of Transformer's Attention via the Lens of Kernel](https://mlsanity.com/paper/1908.11775/) (2019)
- [Long-time integration of parametric evolution equations with physics-informed DeepONets](https://mlsanity.com/paper/2106.05384/) (2021)
- [Improved architectures and training algorithms for deep operator networks](https://mlsanity.com/paper/2110.01654/) (2021)
- [DeepGreen: Deep Learning of Green's Functions for Nonlinear Boundary Value Problems](https://mlsanity.com/paper/2101.07206/) (2020)
- [Operator-valued Kernels for Learning from Functional Response Data](https://mlsanity.com/paper/1510.08231/) (2015)
- [Kymatio: Scattering Transforms in Python](https://mlsanity.com/paper/1812.11214/) (2018)
- [Scaling the Scattering Transform: Deep Hybrid Networks](https://mlsanity.com/paper/1703.08961/) (2017)
- [Fast PDE-constrained optimization via self-supervised operator learning](https://mlsanity.com/paper/2110.13297/) (2021)
- [Generalized Deep Image to Image Regression](https://mlsanity.com/paper/1612.03268/) (2016)
- [Training Neural Networks as Learning Data-adaptive Kernels: Provable Representation and Approximation Benefits](https://mlsanity.com/paper/1901.07114/) (2019)

## Built on by (most cited first, 30 of 90)

- [Geometry-Informed Neural Operator for Large-Scale 3D PDEs](https://mlsanity.com/paper/2309.00583/) (2023)
- [Transformer for Partial Differential Equations' Operator Learning](https://mlsanity.com/paper/2205.13671/) (2022)
- [Convolutional Neural Operators for robust and accurate learning of PDEs](https://mlsanity.com/paper/2302.01178/) (2023)
- [Poseidon: Efficient Foundation Models for PDEs](https://mlsanity.com/paper/2405.19101/) (2024)
- [MIONet: Learning multiple-input operators via tensor product](https://mlsanity.com/paper/2202.06137/) (2022)
- [Towards Foundation Models for Scientific Machine Learning: Characterizing Scaling and Transfer Behavior](https://mlsanity.com/paper/2306.00258/) (2023)
- [Physics-Informed Deep Neural Operator Networks](https://mlsanity.com/paper/2207.05748/) (2022)
- [In-Context Operator Learning with Data Prompts for Differential Equation Problems](https://mlsanity.com/paper/2304.07993/) (2023)
- [Clifford Neural Layers for PDE Modeling](https://mlsanity.com/paper/2209.04934/) (2022)
- [Physics-Informed Machine Learning: A Survey on Problems, Methods and Applications](https://mlsanity.com/paper/2211.08064/) (2022)
- [NOMAD: Nonlinear Manifold Decoders for Operator Learning](https://mlsanity.com/paper/2206.03551/) (2022)
- [A Mathematical Guide to Operator Learning](https://mlsanity.com/paper/2312.14688/) (2023)
- [Learning Deep Implicit Fourier Neural Operators (IFNOs) with Applications to Heterogeneous Material Modeling](https://mlsanity.com/paper/2203.08205/) (2022)
- [Neural Inverse Operators for Solving PDE Inverse Problems](https://mlsanity.com/paper/2301.11167/) (2023)
- [Reliable extrapolation of deep neural operators informed by physics or sparse observations](https://mlsanity.com/paper/2212.06347/) (2022)
- [Representation Equivalent Neural Operators: a Framework for Alias-free Operator Learning](https://mlsanity.com/paper/2305.19913/) (2023)
- [Blending Neural Operators and Relaxation Methods in PDE Numerical Solvers](https://mlsanity.com/paper/2208.13273/) (2022)
- [Kernel Methods are Competitive for Operator Learning](https://mlsanity.com/paper/2304.13202/) (2023)
- [Mitigating spectral bias for the multiscale operator learning](https://mlsanity.com/paper/2210.10890/) (2022)
- [SVD Perspectives for Augmenting DeepONet Flexibility and Interpretability](https://mlsanity.com/paper/2204.12670/) (2022)
- [Out-of-distributional risk bounds for neural operators with applications to the Helmholtz equation](https://mlsanity.com/paper/2301.11509/) (2023)
- [Neural Operator Learning for Long-Time Integration in Dynamical Systems with Recurrent Neural Networks](https://mlsanity.com/paper/2303.02243/) (2023)
- [On the training and generalization of deep operator networks](https://mlsanity.com/paper/2309.01020/) (2023)
- [PDE Generalization of In-Context Operator Networks: A Study on 1D Scalar Nonlinear Conservation Laws](https://mlsanity.com/paper/2401.07364/) (2024)
- [Nonlinear Reconstruction for Operator Learning of PDEs with Discontinuities](https://mlsanity.com/paper/2210.01074/) (2022)
- [Variable-Input Deep Operator Networks](https://mlsanity.com/paper/2205.11404/) (2022)
- [Semi-supervised Invertible Neural Operators for Bayesian Inverse Problems](https://mlsanity.com/paper/2209.02772/) (2022)
- [Incremental Spatial and Spectral Learning of Neural Operators for Solving Large-Scale PDEs](https://mlsanity.com/paper/2211.15188/) (2022)
- [Toward a Better Understanding of Fourier Neural Operators from a Spectral Perspective](https://mlsanity.com/paper/2404.07200/) (2024)
- [HAMLET: Graph Transformer Neural Operator for Partial Differential Equations](https://mlsanity.com/paper/2402.03541/) (2024)

## Abstract (arXiv)

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.