PaperlayerView exact revision on arXiv

Source-linked AI summary

Learning Operators with Coupled Attention

Georgios Kissas, Jacob Seidman, Leonardo Ferreira Guilhoto, Victor M. Preciado, George J. Pappas, Paris Perdikaris

arXiv:2201.01032v1cs.LGcs.AIphysics.comp-ph

TL;DR

Fixed-grid image regression can require a newly tuned model for queries outside its training grid. LOCA couples output-query-dependent attention distributions to model output correlations; the authors report data efficiency, robustness, and generalization, alongside an approximation guarantee under stated conditions.

  • Problem

    Fixed-grid image regression requires a newly built and tuned model to predict at output locations outside its training grid.

    Source: §1 · p. 2

  • Method

    LOCA averages input-function feature embeddings with output-query-dependent distributions coupled through an integral kernel to model correlations between query scores.

    Source: §8 · p. 18

  • Results

    The authors report high performance with 6–12% of labeled data, robustness to noise and initialization randomness, and generalization on climate and synthetic out-of-distribution tasks.

    Source: §1 · p. 3

  • Takeaways & Limitations

    Under the stated conditions, the architecture can approximate any functional from the specified compact input-function space to R.

    Source: §4 · p. 8

    Numerically approximating KCA integrals can be costly: Monte Carlo has O(P^2) complexity, while quadrature has O(PQ + Q^2).

    Source: §8 · p. 18

Abstract

from arXiv · show

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.

1 Introduction

The introduction motivates learning operators directly between function spaces and proposes an attention-based framework that couples output-query weights to model output correlations. It presents theoretical guarantees and empirical claims of data efficiency, robustness, and generalization.

1 Introduction

  • Operator learning models functional responses in physics and engineering and can serve as fast surrogates for traditional numerical solvers.The introduction gives differential-equation systems as a motivating setting.Source: §1 · p. 1
  • Function-space operator models avoid a fixed-grid limitation: predicting at an output location outside the training grid with image regression requires building and tuning a new model.Modeling before discretization permits greater flexibility in point-measurement locations.Source: §1 · p. 2
  • Prior operator-learning approaches include feed-forward methods such as DeepONet [23] and graph-based methods [31] [32] [33], while Fourier Neural Operator and DeepONet have universal-approximation guarantees [21].Other cited approaches include Green’s-function, multi-wavelet, and random-feature methods, for which the passage reports no approximation-power guarantees [34, 35].Source: §1 · pp. 2–3
  • LOCA couples attention distributions across output query locations with an integral transform to model output-function correlations, and the paper proves a universal approximation property under certain assumptions.The mechanism is called Kernel-Coupled Attention.Source: §1 · p. 3
  • The authors report high performance with 6–12% of labeled data, robustness to input noise and initialization randomness, and generalization on climate and synthetic out-of-distribution tasks.They also report test errors concentrated around the median with fewer outliers than competing methods.Source: §1 · p. 3

2 Problem Formulation

The section formulates supervised operator learning from paired input and output functions: learn an operator approximating the unknown ground-truth map. It frames this function-space formulation as applicable to differential-equation responses and flexible measurement locations.

2 Problem Formulation

  • Given paired input and output functions generated by an unknown operator G, the task is to learn an operator F that approximates G.The inputs and outputs belong to spaces of continuous functions over their respective domains.Source: §2 · pp. 3–4
  • The formulation covers differential-equation solution maps and dynamical-system trajectories, where an approximate operator can provide computational speedups and efficient input sensitivities.The examples use an initial condition or a forcing term as the input function.Source: §2 · p. 4
  • Although measurements may be available only at finitely many locations, the model is formulated for functions before discretization, supporting operator learning across multiple output resolutions.Discretization is considered at implementation time.Source: §2 · p. 4

3 Proposed Model: Learning Operators with Coupled Attention (LOCA)

LOCA predicts output functions by averaging features of the input function with query-dependent attention weights, then couples those weights across output query locations using an integral kernel.

3.1 The Attention Mechanism

  • LOCA adapts attention to operator learning by averaging input-function features with distributions that depend on each output query location.The model lifts input functions to feature vectors and averages their rows over the query-dependent distributions.Source: §3.1 · p. 5

3.2 Kernel-Coupled Attention Weights

  • Kernel-Coupled Attention integrates proposal scores against a kernel to relate query distributions, and outperforms uncoupled attention when output evaluations per example are few.The kernel’s form determines which query locations are coupled; a narrow-bandwidth local kernel couples nearby locations.Source: §3 · p. 4, §3.2 · p. 5

3.3 Formulation of the Coupling Kernel

  • The coupling kernel is formed by lifting query locations and applying a universal kernel, with tunable parameters learning similarities across the output domain.The text gives a Gaussian RBF example and normalizes the kernel output to maintain the proposal score function’s relative scale.Source: §3.3 · p. 5

3.4 Input Function Feature Encoding

  • The input feature representation composes a map D from the input function to a finite-dimensional vector with a universal function approximator f.Using point evaluations at fixed locations for D prevents accepting measurements elsewhere or varying input resolution across passes; Fourier or wavelet projections are alternatives.Source: §3.4 · p. 6

3.5 Model Summary

  • The forward pass combines the input feature encoder v, score function g, softmax σ, and coupling kernel κ to produce output values.Figure 2 depicts the feature representation averaged with query-dependent attention weights to evaluate the solution at each query location.Source: Figure 2 · p. 6, §3.5 · p. 7

4 Theoretical Guarantees of Universality

The section establishes conditions under which LOCA’s normalization, kernel coupling, and spectral input encoding preserve universal approximation. It also gives an example using Fourier modes for C1 input functions.

Definition of universality

  • Universality means density among operators under the supremum norm on the stated compact spaces.The definition requires approximation of every continuous operator to arbitrary accuracy.Source: §4 · p. 7

Normalization and kernel coupling

  • LOCA remains universal after softmax normalization, and kernel coupling preserves this guarantee when the kernel conditions are satisfied.The kernel integral operator preserves density; the coupling argument uses a dense RKHS image to approximate continuous output-domain weights.Source: §4 · pp. 7–8

Kernel conditions

  • The coupling kernel is symmetric and positive definite, and it defines a universal RKHS when q is injective.These are the stated conditions supporting the kernel-coupling universality result.Source: §4 · p. 8

Spectral encoding

  • Spectral encoding preserves universality when basis projections converge uniformly over the compact input set.Under that condition, the proposition guarantees approximation of continuous functionals using a sufficiently large finite projection.Source: §4 · p. 8

Example

  • For a compact C1 input set, finite Fourier-mode projections can support approximation of any functional from that set to R.The section gives this as an example of the architecture’s expressive capacity.Source: §4 · p. 8

5 Implementation Aspects

The implementation combines numerical integration, query-location encodings, and a scattering-based input representation, then trains the model by stochastic gradient descent and evaluates predictions on super-resolution grids.

5.1 Computation of the Kernel Integrals

  • Monte Carlo estimates the kernel integrals from P sampled domain points, while Gauss-Legendre quadrature offers an accurate, efficient alternative in low dimensions.The normalized kernel cancels the domain-volume factor; query locations may also be among the Monte Carlo points.Source: §5.1 · pp. 8–9
  • For translation-invariant kernels, Fourier-domain integration is an option, but efficient implementations generally require equally spaced grids.The authors retain arbitrary query locations and instead switch between Monte Carlo and quadrature by problem.Source: §5.1 · p. 9

5.2 Positional Encoding of Output Query Locations

  • Query locations are encoded by physical coordinates rather than their index in a list, which generally lacks physical meaning.The encoding extends positional features across multiple spatial dimensions.Source: §5.2 · p. 9

5.3 Wavelet Scattering Networks as a Spectral Encoder

  • The scattering transform is motivated by Fourier encoders’ limited robustness to input noise and small deformations.It alternates wavelet convolutions and complex modulus operations; the cited result states Lipschitz continuity to small deformations, unlike Fourier-transform modulus.Source: §5.3 · p. 9
  • The scattering encoder uses a finite set of paths controlled by maximum scale J, rotation-group size L, and maximum path length m0.This practical encoder is not necessarily covered by Proposition 4.3, although it gives the best empirical performance reported here.Source: §5.3 · p. 10

5.4 Loss Function and Training

  • Training minimizes empirical risk with stochastic gradient descent, using reverse-mode automatic differentiation for parameter gradients.The described loss is the simplest choice; weighted mean-square or relative L2 losses are also mentioned as alternatives.Source: §5.4 · p. 10

5.5 Implementation Overview

  • The implementation preprocesses inputs and coordinates, trains on mini-batches, then makes one-shot super-resolution predictions evaluated by relative L2 error.The forward pass combines the coupling kernel, numerical KCA approximation, input features, and expectation to produce predictions.Source: §5.5 · p. 10

6 Connections to Existing Operator Learning Methods

LOCA shares a final feature-weighted evaluation step with DeepONet but adds normalized, kernel-coupled attention and a different input representation. Its connections to neural operators and transformer variants clarify these architectural distinctions.

6.1 DeepONets

  • DeepONet is recovered from LOCA by removing attention normalization and kernel coupling, and replacing scattering features with fixed-sensor samples.The input feature map corresponds to DeepONet’s branch network and the query-dependent distribution to its trunk network.Source: §6.1 · p. 11

6.2 Neural Operators

  • A neural-operator layer can resemble DeepONet under specified zero-weight, tiled-input, separable-kernel, and output-summing choices.LOCA connects transitively through DeepONet; its scattering component can also be viewed as fixed-weight neural-operator layers.Source: §6.2 · p. 11

6.3 Other Attention-Based Architectures

  • Unlike Galerkin/Fourier and neural-operator transformers, LOCA attends over finite input features and couples attention weights across the output domain.The kernel acts on an output-domain score function in LOCA, rather than directly on internal input and output functions as in neural operators.Source: §6.3 · pp. 11–12, §6.2 · p. 11

7 Experimental Results

The section compares LOCA with FNO and DON across operator-learning benchmarks, assessing sparse output labels, noisy data, initialization variability, and distribution shifts.

7.1 Data Efficiency

  • 6% of available output measurements lets LOCA achieve lower Darcy errors than DON and FNO trained with up to 100%; it also performs well at 1.5%.The comparison uses relative L2 error over the test set.Source: §7.1 · p. 13, Figure 4 · p. 14
  • With 1.5% labeled data and 200 training examples, removing KCA destabilizes training and raises test error, while LOCA with KCA performs well.Table 1 reports Darcy relative L2 error statistics with and without KCA.Source: Table 1 · p. 14, §7.1 · p. 14
  • At 2.5% labeled data, LOCA outperforms DON and FNO on shallow-water height ρ, while velocity median errors are similar and LOCA has smaller spread and fewer outliers.The velocity components are v1 and v2.Source: §7.1 · p. 13, Figure 5 · p. 15

7.2 Robustness

  • Across 10 random initializations with 7% labeled outputs, LOCA has a smaller maximum-error spread than competitors for v1 and performance similar to FNO for v2.This test uses clean inputs and noisy outputs.Source: §7.2 · p. 15, Figure 7 · p. 17

7.3 Generalization

  • With test Gaussian-process length-scale fixed at 0.1, LOCA outperforms competitors across training length-scales 0.1–0.9.The test length-scale is fixed while the training length-scale varies across nine datasets.Source: §7.3 · p. 17
  • With 4% of pressure measurements for training, LOCA achieves the lowest error and a small error spread when predicting 2005–2010 pressure from 2000–2005 temperature data.The experiment evaluates climate prediction on a later time period than the training data.Source: §7.3 · p. 17, Figure 8 · p. 18
  • Across input distributions spanning 2 orders of magnitude in length-scale and 4 in amplitude, LOCA is more accurate and has a smaller error spread than DON and FNO.FNO's median is close to LOCA's, but it has many high-error outliers.Source: §7.3 · pp. 17–18

8 Discussion

LOCA averages input-function feature embeddings at output queries and couples query weights through Kernel-Coupled Attention. The discussion reports robustness to noisy inputs, generalization across input distributions, and extrapolation on a climate functional-regression task.

8 Discussion

  • Kernel-Coupled Attention integrates query scores against a similarity kernel, explicitly modeling correlations between output query locations.LOCA averages input-function feature embeddings using probability distributions that depend on output query locations.Source: §8 · p. 18
  • A wavelet-scattering feature encoder is used for robustness to small input deformations and noise, a property the experiments support.Source: §8 · p. 18
  • Experiments show generalization across varying functional-input distributions and extrapolation on a functional-regression task with global climate data.Source: §8 · p. 18
  • Approximating KCA integrals costs O(P^2) with Monte Carlo or O(PQ + Q^2) with quadrature, and high-dimensional domains make integration increasingly costly.The authors suggest parallelization and kernel approximations as future ways to improve scaling.Source: §8 · pp. 18–19
  • Figure 9 compares model error boxplots across increasingly out-of-distribution training sets using all available output measurements; Figure 10 shows test errors for inputs with varied lengthscales and amplitudes.Figure 10 also shows sample input and output functions.Source: Figure 9 · p. 19, Figure 10 · p. 20
  • A layered LOCA variant did not significantly improve performance in the reported experiments, while a context-vector extension could model output uncertainty.Source: §8 · p. 19
  • Operator-learning models can speed repeated PDE solves and sensitivity evaluation, supporting design and gradient-based control optimization.The discussion notes that sensitivities can be computed through automatic differentiation.Source: §8 · pp. 19–20

A Nomenclature

B Proof of Theorem 4.1

The proof approximates the compact image of the operator with a finite-dimensional subspace, then expresses the approximation through averages of input functionals.

B Proof of Theorem 4.1

  • A compact subset of a Banach space can be approximated arbitrarily closely by a finite-dimensional subspace.Source: B Proof of Theorem 4.1 · p. 24
  • Because the operator is continuous on a compact input set, its image is compact, so the finite-dimensional approximation lemma applies to that image.Source: B Proof of Theorem 4.1 · p. 24
  • The resulting approximation is rewritten as a vector of averages of modified functionals, which form the coordinates of the input feature vector.The construction separates positive and negative parts and creates coordinate-specific functions for each output dimension.Source: B Proof of Theorem 4.1 · pp. 24–25

C Proof of Proposition 4.1

The proof establishes density by using continuity to transfer dense sets through maps and relating the images of the coupling-kernel operator and its square root.

C Proof of Proposition 4.1

D Proof of Proposition 4.2

The proof establishes symmetry and positive semi-definiteness of the coupling kernel, then uses feature-map injectivity to obtain universality.

D Proof of Proposition 4.2

E Proof of Proposition 4.3

The appendix gives a compactness-and-density argument for Proposition 4.3, then introduces experimental supplements and analyzes the computational cost of LOCA’s KCA integrals.

E Proof of Proposition 4.3

  • Uniform continuity on compact U bounds the change in h, and finite-dimensional projections approximate each u closely enough to transfer the approximation to h.The proof then uses density of A_d in C(R^d, R^n) to approximate h composed with r within ε/2.Source: E Proof of Proposition 4.3 · p. 26

F Supplementary Information for Experiments

F.1 Computational Complexity

  • For P query locations and Q integration nodes, KCA kernel evaluation costs PQ + Q^2 operations, with the k(z, z) computation shared across queries.For Monte Carlo integration, Q = P and y = z, reducing the computation to Q^2.Source: §F.1 · pp. 26–27
  • Monte Carlo integration performs P^2 computations once, with cost scaling exponentially in P; Gauss-Legendre can cost less when Q is small.The comparison is stated for the integral-computation methods used in the KCA mechanism.Source: §F.1 · pp. 26–27

F.2 Architecture Choices and Hyper-parameter Settings

This appendix records model settings, computational costs, comparison metrics, and experiment details. It describes the antiderivative data and evaluation setups, including tests across input length scales.

F.2.2 DON

  • DON’s benchmark settings specify branch and trunk architectures and training parameters; the implementation also uses harmonic feature expansions and output normalization.The reported modifications follow improvements proposed in Lu et al. [64].Source: §F.2.2 · p. 27, Table 8 · p. 28, Table 9 · p. 29

F.2.3 FNO

  • FNO benchmark configurations and training parameters are provided in Tables 10 and 11.The appendix describes architecture choices, dataset sizes, evaluation points, batch size, and training epochs.Source: §F.2.3 · p. 27, Table 10 · p. 29, Table 11 · p. 29

F.3 Computational Cost

  • Training wall-clock times for all methods are in the same order of magnitude; the Darcy comparison uses P = 1024 output measurements per example.All models were trained on one NVIDIA RTX A6000 GPU.Source: §F.3 · p. 27, Table 13 · p. 29

F.2.1 LOCA

  • LOCA’s benchmark-specific architecture choices specify g, q, φ(y), and v, alongside model parameters and the kernel-integral computation method.Tables 5–7 cover training and test sizes, evaluation points, latent and positional dimensions, training settings, scattering-network parameters, and integration methods.Source: §F.2.1 · p. 27, Table 4 · p. 28, Table 5 · p. 28

F.3 Computational Cost

  • The appendix reports model parameter counts by benchmark in Table 12 and training wall-clock times in Table 13.The time comparison is for training on a single NVIDIA RTX A6000 GPU.Source: §F.3 · p. 27, Table 12 · p. 29, Table 13 · p. 29

F.4 Comparison Metrics

  • Test accuracy is assessed with relative L2 error over full-resolution reconstructions, with medians, quantiles, and outliers reported across test examples.The metric compares each predicted solution with its full-resolution ground truth.Source: §F.4 · p. 30

F.5 Experiments

  • The experiment supplement covers antiderivative, Darcy flow, Mechanical MNIST, shallow-water, and climate-modeling scenarios.It provides further details about the experiments discussed in Section 7.Source: §F.5 · p. 30

F.5.1 Antiderivative

  • For antiderivative training, forcing functions are sampled from a Gaussian-process prior and measured at 500 points, then numerically integrated into 100 output measurements used to train operator-learning models.In the out-of-distribution setup, training length scales span 0.1–0.9 in nine cases spaced by 0.1, with 1,000 solutions in each training and test set.Source: §F.5.1 · p. 30

F.5.2 Darcy Flow

The supplementary experiments describe four operator-learning benchmarks and show visual comparisons of model predictions with reference solutions. The reported test examples include unseen inputs for Darcy flow and shallow-water dynamics.

F.5.2 Darcy Flow

  • For three randomly selected Darcy-flow inputs absent from training, the model’s predictions visually agree well with the reference solutions.Figure 11 shows each input, prediction, ground truth, and absolute error.Source: §F.5.2 · p. 31, Figure 11 · p. 31

F.5.3 Mechanical MNIST

  • Mechanical MNIST evaluates later-time displacements from initial deformations, using 70,000 finite-element simulations and 60,000 training versus 10,000 test digits.Figures 12 and 13 compare test predictions with ground truth; Figure 13 shows initial and final pixel positions and point-wise error.Source: §F.5.3 · pp. 31–32, Figure 12 · p. 33, Figure 13 · p. 33

F.5.4 Shallow Water Equations

  • For a held-out initial condition, the model provides favorable shallow-water solutions at five times from t = 0.11s to 0.31s.Figures 14–18 compare the initial conditions, predicted and ground-truth parameters, and point-wise absolute error.Source: §F.5.4 · p. 34, Figure 14 · p. 35, Figure 15 · p. 36

F.5.5 Climate Modeling

  • The climate benchmark learns a black-box operator mapping surface air temperature to surface air pressure; daily test data cover 2005 to 2010 (Ntest = 1825, excluding leap years) on a 72 × 72 grid.Figure 19 shows the input temperature, output prediction, ground truth, and absolute error.Source: §F.5.5 · p. 34, Figure 19 · p. 39
Loading 2201.01032v1…