[论文解读] Learning Operators with Coupled Attention

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

Hydrological Forecasting Using AI被引用 56

一句话总结

LOCA 提出一种新颖的算子学习框架，使用 Kernel-Coupled Attention 来学习函数空间之间的映射，具备通用逼近保证，以及在 PDE/ODE 与气候相关任务中的强数据效率、鲁棒性和泛化能力。

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.

研究动机与目标

• 激励在时空系统中学习函数空间之间的算子。
• 引入一种新颖的基于注意力的架构，用以捕捉输出函数中的相关性。
• 为所提模型提供通用逼近保证。
• 展示在真实与合成数据上的数据效率、对噪声的鲁棒性以及良好的泛化能力。

提出的方法

• 将输入函数提升到一个有限特征集 v(u)，并通过 E_phi(y)[v(u)] 在查询位置 y 处计算输出。
• 引入 Kernel-Coupled Attention (KCA) 通过一个核积分算子将不同输出查询位置的注意力权重耦合起来。
• 通过提升的输入 q_theta 和一个通用核来定义核 kappa，并进行归一化以产生耦合的注意力分布。
• 通过 D(u) 特征映射编码输入函数，然后经过通用函数 f，形成 v(u)=f(D(u))。
• 将小波散射作为可选的光谱编码器 D，用于对形变和噪声的鲁棒性。
• 通过经验风险与标准 L2 损失进行训练，并启用 Monte Carlo 或基于数值积分的核积分近似。
• 讨论实际方面：对 y 的位置编码、积分的离散化策略，以及通过梯度法优化。

实验结果

研究问题

• RQ1 LOCA 是否能够近似函数空间之间的任意连续算子（通用逼近）？
• RQ2 当输出测量数量较少时，跨输出查询位置耦合注意力是否能提升准确性和数据效率？
• RQ3 LOCA 在有噪声的输入以及分布外/泛化场景下的表现如何？
• RQ4 LOCA 与现有的算子学习方法在 ODE/PDE 与气候数据任务上的比较如何？
• RQ5 保持通用性与性能的实际实现选择有哪些？

主要发现

• 在适当假设下，LOCA 满足通用逼近性质。
• 当每个输入的输出评估数量较小时，耦合注意力能提高准确性。
• LOCA 展示对噪声的鲁棒性并降低离群值，误差集中在中位数附近。
• 实验结果在合成数据和真实数据（地表空气温度与气压）上显示出最先进的准确性和更好的泛化能力，包括超出训练数据的外推。
• 与竞争方法相比，该模型实现了高数据效率，仅需标记数据的很小一部分（6-12%）。