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.

Learning Deep Implicit Fourier Neural Operators (IFNOs) with Applications to Heterogeneous Material Modeling

“…When comparing the prediction accuracy from different models, similar to the previous examples, the FNO suffers from overfitting and vanishing gradient issues when L > 2, especially in the original (more noisy) dataset. This finding is consistent with the results reported in [68,103], where the performance of the FNOs was found to be deteriorated on noisy datasets. In contrast, the accuracy of the IFNOs monotonically improves with the increase of L. Both neural operator models outperforms the conventional constitutive modeling approaches by around one order of magnitude.…”

Physics-Informed Deep Neural Operator Networks

“…In the regularized cavity flow problem presented in [17], an additional trigonometric input function was included in the branch network to incorporate periodic boundary conditions. Similarly, an additional input function is included in [20], which is created by averaging a feature embedded in the inputs of the branch network over probability distributions that depend on the corresponding query locations of the output function, employing the kernelcoupled attention mechanism, which allows the operator to accurately model correlations between the query locations of the output functions.…”

Physics informed token transformer for solving partial differential equations

“…Guo et al [29] introduces attention as an instance-based learnable kernel for direct sampling method and demonstrates superiority on boundary value inverse problems. Learning operators with coupled attention [32] uses attention weights to learn correlations in the output domain and enables sample-efficient training of the model. General neural operator transformer for operator learning [25] proposes a heterogeneous attention architecture that stacks multiple cross-attention layers and uses a geometric gating mechanism to adaptively aggregate features from query points.…”