[논문 리뷰] Learning Operators with Coupled Attention

Georgios Kissas, J. S. Seidman외 4명|arXiv (Cornell University)|2022. 01. 04.

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.

연구 동기 및 목표

• 시공간 시스템을 위한 함수 공간 간 연산자 학습의 필요성을 제기한다.
• 출력 함수 간의 상관관계를 포착하는 새로운 주의(attention) 기반 아키텍처를 도입한다.
• 제안된 모델에 대한 보편 근사 보장을 제공한다.
• 실제 데이터와 합성 데이터 모두에서 데이터 효율성, 잡음에 대한 강건성, 그리고 우수한 일반화를 시연한다.

제안 방법

• 입력 함수들을 유한 특징 집합 v(u)로 올리고, 쿼리 위치 y에서의 출력을 E_phi(y)[v(u)]를 통해 계산한다.
• Kernel-Coupled Attention(KCA)를 도입하여 커널 적분 연산자를 사용해 서로 다른 출력 쿼리 위치 간에 주의 가중치를 연결한다.
• 리프팅된 입력 q_theta와 보편 커널을 통해 커널 kappa를 정의하고, 이를 정규화하여 결합된 attention 분포를 생성한다.
• 입력 함수를 D(u) 특징 맵을 거쳐 보편 함수 f를 적용하여 v(u)=f(D(u))를 형성한다.
• 변형과 잡음에 대한 강건성을 위해 D로서 웨이브렛 산란을 선택적으로 스펙트럴 인코더로 사용한다.
• 표준 L2 손실을 이용한 경험적 위험 하에서의 학습을 제공하고, 커널 적분에 대해 몬테 카를로(Monte Carlo) 또는 사분적(quadrature) 근사를 가능하게 한다.
• y에 대한 위치 인코딩, 적분의 이산화 전략, 및 그래디언트 기반 방법을 통한 최적화 등 실용적 측면을 논의한다.

실험 결과

연구 질문

• RQ1 LOCA가 함수 공간 간의 모든 연산자를 근사할 수 있는가(보편 근사성)?
• RQ2 출력 쿼리 위치 간의 주의 결합이 출력 측정 수가 작을 때 정확도와 데이터 효율성을 개선하는가?
• RQ3 노이즈가 있는 입력과 분포 밖/일반화 시나리오에서 LOCA는 어떻게 작동하는가?
• RQ4 LOCA는 ODE/PDE 및 기후 데이터 작업에서 기존의 연산자 학습 방법과 어떻게 비교되는가?
• RQ5 보편성과 성능을 유지하는 실용적 구현 선택은 무엇인가?

주요 결과

• LOCA는 적절한 가정 하에 보편 근사 성질을 만족한다.
• 출력 평가 수가 작을 때 Coupled attention이 정확도를 향상시킨다.
• LOCA는 노이즈에 대한 강건성과 이상치 감소를 보여주며, 오차가 중앙값 주변으로 집중된다.
• 합성 데이터와 실제 데이터(Earth surface air temperature and pressure)에서 상태-최첨단 정확도와 일반화, 특히 학습 데이터 외의 외삽 가능성이 향상되었다.
• 모델은 데이터 효율이 높아, 경쟁 방법에 비해 표시된 데이터의 소량(6-12%)만으로도 달성된다.