[논문 리뷰] Choose a Transformer: Fourier or Galerkin

Shuhao Cao|arXiv (Cornell University)|2021. 05. 31.

Model Reduction and Neural Networks참고 문헌 97인용 수 58

한 줄 요약

이 논문은 연산자 학습을 위한 Fourier-type 및 Galerkin-type 자기 주의 변형(소프트맥스 없이)을 제시하고, 근사 능력을 입증하며 PDE 벤치마크에서 소프트맥스 기반 트랜스포머보다 효율성과 정확도 향상을 보여준다.

ABSTRACT

In this paper, we apply the self-attention from the state-of-the-art Transformer in Attention Is All You Need for the first time to a data-driven operator learning problem related to partial differential equations. An effort is put together to explain the heuristics of, and to improve the efficacy of the attention mechanism. By employing the operator approximation theory in Hilbert spaces, it is demonstrated for the first time that the softmax normalization in the scaled dot-product attention is sufficient but not necessary. Without softmax, the approximation capacity of a linearized Transformer variant can be proved to be comparable to a Petrov-Galerkin projection layer-wise, and the estimate is independent with respect to the sequence length. A new layer normalization scheme mimicking the Petrov-Galerkin projection is proposed to allow a scaling to propagate through attention layers, which helps the model achieve remarkable accuracy in operator learning tasks with unnormalized data. Finally, we present three operator learning experiments, including the viscid Burgers' equation, an interface Darcy flow, and an inverse interface coefficient identification problem. The newly proposed simple attention-based operator learner, Galerkin Transformer, shows significant improvements in both training cost and evaluation accuracy over its softmax-normalized counterparts.

연구 동기 및 목표

• PDE에 대한 데이터 기반 연산자 학습의 필요성과 메시 인변형(mesh-invariant), 해상도 불변 연산자의 필요성을 제시한다.
• Hilbert 공간 및 Petrov–Galerkin 원칙에 맞추어 소프트맥스 없는 자기 주의 변형을 도입한다.
• Gal­erkin Transformer와 비정규화 데이터에서도 훈련을 안정시키는 대응하는 층 정규화를 개발한다.
• 소프트맥스 기반 주의와 비교하여 PDE 벤치마크에서 훈련 효율성 및 정확도를 향상시킴을 입증한다.

제안 방법

• 소프트맥스 없이 간단한 자기 주의 변형과 메쉬 길이 가중 정규화(Galerkin-type 및 Fourier-type)를 정의한다.
• 주의를 Hilbert 공간에서 적분 변환 또는 Petrov–Galerkin 투영에 대한 행렬 없는 근사로 해석한다.
• Gal­erkin-type 주의가 n독립 근사를 갖는 학습 가능한 Petrov–Galerkin 투영에 대응한다는 것을 보이고 Céa-type 오차 한계(정리 4.3)를 제공한다.
• 학습을 안정시키고 층을 거쳐 규모를 전달하기 위해 Petrov–Galerkin 투영을 모방하는 층 정규화 체계를 제안한다.
• 잠재 기저를 피드포워드 네트워크와 위치 인코딩으로 동적으로 풍부하게 하여 연산자 반응에 적응하도록 한다.
• FT 및 GT를 소프트맥스 기반 ST와 LT 변형과 비교하여 PDE 작업에서 메모리와 속도 이점을 강조한다.

실험 결과

연구 질문

• RQ1 소프트맥스 없는 주의가 PDE 연산자 학습에서 Petrov–Galerkin 투영과 유사한 근사 능력을 달성할 수 있는가?
• RQ2 Galerkin-type 주의가 비정규화 데이터에서 안정적이고 확장 가능한 훈련을 가능하게 하며 전통적인 소프트맥스 주의에 비해 효율성을 제공하는가?
• RQ3 Fourier-type 및 Galerkin-type 주의가 벤치마크 PDE 연산자 학습 작업(점성 Burgers’, Darcy 흐름, 역 제어 계수 식별)에서 어떤 성능을 보이는가?
• RQ4 레이어 정규화 및 동적 기저 업데이트가 연산자 학습자의 정확도와 안정성 향상에 어떤 역할을 하는가?

주요 결과

• 소프트맥스 없는 Galerkin-type 주의는 Petrov–Galerkin 투영을 명시적으로 표현할 수 있으며 그 근사 능력은 시퀀스 길이에 독립적이다(정리 4.3).
• 새로운 Galerkin 투영 유형의 층 정규화가 주의 층을 통한 안정적 스케일링을 가능하게 하고 물리적 특성 추적(예: 에너지 감소)을 개선한다.
• Gal­erkin Transformer와 Fourier Transformer는 학습 메모리를 줄이고 PDE 연산자 학습 작업에서 소프트맥스 기반 상대 대비 훈련 속도를 높일 수 있다.
• viscid Burgers’ 방정식, 계면 Darcy 흐름, 역 계면 계수 식별에 대한 실험에서 GT/FT가 정확도와 효율성을 개선했다.
• Fourier-type 주의의 메모리 사용량은 소프트맥스 대비 40-50% 감소하고 Galerkin-type 주의는 일부 경우 선형 주의에 비해 상당한 속도 향상을 제공할 수 있다.
• 저자들은 결과 재현을 위한 오픈소스 PyTorch 코드를 제공한다.