[논문 리뷰] From block-Toeplitz matrices to differential equations on graphs: towards a general theory for scalable masked Transformers

Krzysztof Choromański, Han Lin외 8명|arXiv (Cornell University)|2021. 07. 16.

Quantum optics and atomic interactions인용 수 7

한 줄 요약

이 논문은 그래프 상의 위상 조절로 마스킹을 모델링하여 트랜스포머에서 확장 가능한 마스킹 어텐션을 위한 일반적인 이론적 프레임워크를 제안한다. 블록-토플리츠 행렬 구조와 빠른 알고리즘을 통해 2차 이하의 계산을 가능하게 하며, 선형 인과 어텐션과 로그-선형 RPE-어텐션과 같은 이전 방법들을 통합하고 확장한다. d차원 RPE와 증명된 효율성을 갖춘 그래프 커널 마스킹을 위한 새로운 확장 가능한 메커니즘을 도입한다.

ABSTRACT

In this paper we provide, to the best of our knowledge, the first comprehensive approach for incorporating various masking mechanisms into Transformers architectures in a scalable way. We show that recent results on linear causal attention (Choromanski et al., 2021) and log-linear RPE-attention (Luo et al., 2021) are special cases of this general mechanism. However by casting the problem as a topological (graph-based) modulation of unmasked attention, we obtain several results unknown before, including efficient d-dimensional RPE-masking and graph-kernel masking. We leverage many mathematical techniques ranging from spectral analysis through dynamic programming and random walks to new algorithms for solving Markov processes on graphs. We provide a corresponding empirical evaluation.

연구 동기 및 목표

• 다양한 마스킹 메커니즘을 2차 복잡도 없이 트랜스포머 아키텍처에 통합할 수 있는 일반적이고 확장 가능한 이론을 개발하기 위해.
• 선형 인과 어텐션과 로그-선형 RPE-어텐션와 같은 기존의 효율적인 어텐션 메커니즘들을 단일 이론적 프레임워크 아래 통합하기 위해.
• 확장 가능한 어텐션의 범위를 d차원 위치 인코딩과 그래프 기반 커널 마스킹과 같은 새로운 설정으로 확장하기 위해.
• 스펙트럼 분석, 동적 프로그래밍, 그래프 상의 마코프 과정을 기반으로 한 수학적 기반을 마련하여 효율적인 마스킹 어텐션 계산을 가능하게 하기 위해.
• 최신 기술 기반 GNN들과의 비교를 통해 새로운 모델인 그래프 커널 어텐션 트랜스포머(GKAT)를 경험적으로 검증하기 위해.

제안 방법

• 마스킹을 비마스킹 어텐션에 대한 위상적(그래프 기반) 조절로 공식화하며, 마스크는 그래프 구조나 거리 함수로부터 유도된다.
• 마스크 행렬이 2차 이하의 행렬-벡터 곱셈을 지원하는 성질을 통해 효율적인 마스킹을 특성화함으로써, 확장 가능한 저랭크 어텐션을 가능하게 한다.
• 예를 들어 이미지나 영상와 같은 d차원 격자에서 블록-토플리츠 구조를 활용하여 1차원 시퀀스를 초월한 상대적 위치 인코딩을 일반화한다.
• 특히 최단경로 및 확산 커널에 대해 그래프 기반 마스크의 행렬-벡터 곱셈을 가속화하기 위해 푸리에 변환(FFT)과 하 Hankel 행렬 기법을 적용한다.
• 나무 구조의 그래프 상에서 하향-상향 순서의 순회를 통해 마스킹 어텐션을 효율적으로 계산하기 위해 재귀적 분할-정복 알고리즘을 설계한다.
• 패딩과 시퀀스 팩킹을 프레임워크에 통합하기 위해, 수정된 그래프 구조를 모델링하고 트리 분해 상에서 효율적인 계산을 수행한다.

실험 결과

연구 질문

• RQ1 어떤 조건에서 마스킹 메커니즘이 2차 이하의 시간 복잡도로 어텐션 메커니즘에 통합될 수 있는가?
• RQ2 구조적 행렬 성질을 활용하여 이미지 및 영상 모델링에 적합한 d차원 상대적 위치 인코딩을 어떻게 확장 가능한 형태로 만들 수 있는가?
• RQ3 최단경로 또는 확산 커널 기반의 그래프 커널 기반 어텐션 마스크는 저랭크 어텐션 설정에서 효율적으로 계산될 수 있는가?
• RQ4 기존의 효율적인 어텐션 메커니즘들인 인과 어텐션과 로그-선형 RPE는 어떻게 단일 이론적 프레임워크 아래 통합될 수 있는가?
• RQ5 이 이론에서 유도된 새로운 확장 가능한 어텐션 메커니즘인 그래프 커널 어텐션 트랜스포머(GKAT)는 기존의 GNN들을 능가할 수 있는가?

주요 결과

• 논문은 확장 가능한 마스킹 어텐션을 위한 일반 조건을 수립한다: 마스크 행렬이 2차 이하의 행렬-벡터 곱셈을 지원한다면, 마스킹된 저랭크 어텐션은 2차 이하의 시간 복잡도로 계산될 수 있다.
• 이 프레임워크는 1차원 RPE와 인과 어텐션을 d차원 RPE와 그래프 기반 마스킹의 특수 케이스로 일반화하여, 이미지나 영상과 같은 2차원 및 3차원 입력에 대한 효율적인 어텐션을 가능하게 한다.
• FFT를 통한 빠른 계산을 지원하는 새로운 d차원 RPE 메커니즘이 도입되어 시각 및 영상 작업에서의 확장 가능한 어텐션을 실현한다.
• 최단경로 거리 및 확산 커널을 사용한 그래프 커널 마스킹은 FFT와 트리 분해 기법을 통해 효율적으로 계산 가능하다고 밝혀졌다.
• 나무 구조의 그래프에 대한 제안된 알고리즘은 마스킹 어텐션 계산에 대해 $O(|V( au)| imes \text{diam}( au))$의 시간 복잡도를 달성하며, 재귀적 분해와 순환 이동 연산을 활용한다.
• 경험적 평가에서 제안된 그래프 커널 어텐션 트랜스포머(GKAT) 모델은 9개의 최신 기술 기반 GNN들을 능가하여, 제안된 프레임워크의 실용적 효용성을 검증했다.