[Paper Review] Learning Operators with Coupled Attention

Georgios Kissas, J. S. Seidman+4 more|arXiv (Cornell University)|Jan 4, 2022

Hydrological Forecasting Using AI56 citations

TL;DR

LOCA proposes a novel operator-learning framework that uses Kernel-Coupled Attention to learn mappings between function spaces, with universal approximation guarantees and strong data efficiency, robustness, and generalization in PDE/ODE and climate-related tasks.

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.

Motivation & Objective

• Motivate learning operators between function spaces for spatio-temporal systems.
• Introduce a novel attention-based architecture that captures correlations in output functions.
• Provide universal approximation guarantees for the proposed model.
• Demonstrate data efficiency, robustness to noise, and good generalization on real and synthetic data.

Proposed method

• Lift input functions to a finite feature set v(u) and compute outputs at query locations y via E_phi(y)[v(u)].
• Introduce Kernel-Coupled Attention (KCA) to couple attention weights across different output query locations using a kernel integral operator.
• Define the kernel kappa via a lifted input q_theta and a universal kernel, normalizing to produce coupled attention distributions.
• Encode input functions through a D(u) feature map followed by a universal function f, forming v(u)=f(D(u)).
• Use wavelet scattering as an optional spectral encoder D for robustness to deformations and noise.
• Provide training via empirical risk with standard L2 loss and enable Monte Carlo or quadrature-based integral approximations for the kernel integrals.
• Discuss practical aspects: positional encodings for y, discretization strategies for integrals, and optimization via gradient-based methods.

Experimental results

Research questions

• RQ1 Can LOCA approximate any continuous operator between function spaces (universal approximation)?
• RQ2 Does coupling attention across output query locations improve accuracy and data efficiency when the number of output measurements is small?
• RQ3 How does LOCA perform under noisy inputs and in out-of-distribution/generalization scenarios?
• RQ4 How does LOCA compare with existing operator-learning methods on ODE/PDE and climate-data tasks?
• RQ5 What are the practical implementation choices that preserve universality and performance?

Key findings

• LOCA satisfies a universal approximation property under appropriate assumptions.
• Coupled attention improves accuracy when the number of output evaluations per input is small.
• LOCA demonstrates robustness to noise and reduces outliers, with errors concentrated around the median.
• Empirical results show state-of-the-art accuracy and better generalization, including extrapolation beyond training data, on synthetic and real data (Earth surface air temperature and pressure).
• The model achieves high data efficiency, requiring only a small fraction (6-12%) of labeled data compared to competing methods.