Nyströmformer: Nyström 방법을 이용한 선형 시간 및 메모리에서의 self-attention 근사

TL;DR

Nyströmformer는 표준 self-attention 메커니즘을 근사하여 전통적인 2차 복잡도 $O(n^2)$ 대신 선형 시간 및 메모리 복잡도 $O(n)$를 달성하는 효율적인 Transformer 변형입니다. 이를 통해 모델은 훨씬 더 긴 입력 시퀀스를 처리하면서도 NLP 및 컴퓨터 비전 작업 전반에서 경쟁력 있는 성능을 유지할 수 있습니다.

행렬 근사를 위한 Nyström 방법

Nyström 방법은 $m$개의 행과 열을 샘플링함으로써 큰 행렬 $P^{n \times n}$을 근사합니다. 이러한 샘플들을 블록 행렬로 배열하면 원래 행렬은 네 개의 서브매트릭스로 나뉩니다: $A_{P}$ ($m \times m$), $B_{P}$ ($m \times (n - m)$), $F_{P}$ ($(n - m) \times m$), 그리고 $C_{P}$ ($(n - m) \times (n - m)$).

$A_{P}$, $B_{P}$, $F_{P}$의 항목들은 샘플링을 통해 알려지므로, 알 수 없는 서브매트릭스 $C_{P}$는 다음과 같이 추정될 수 있습니다:

$$C_{P} = F_{P} A_{P}^{+} B_{P}$$

여기서 $+$는 Moore-Penrose 의사역을 나타냅니다. 결과적인 Nyström 근사 $ \hat{P}$는 세 개의 행렬 곱으로 표현되어 전체 행렬을 계산하는 비용을 크게 줄입니다.

Nyström을 Self-Attention에 적용하기

표준 self-attention은 소프트맥스 행렬 $S = \text{softmax}(\frac{QK^T}{\sqrt{d}})$에 의존합니다. $S$에 Nyström 방법을 직접 적용하는 것은 불가능합니다. 왜냐하면 소프트맥스 행렬의 단일 열을 계산하려면 행별 소프트맥스 분모를 구하기 위해 모든 다른 열에 대한 정보가 필요하기 때문입니다.

이를 해결하기 위해 Nyströmformer는 결과 소프트맥스 행렬에서 샘플링하는 대신 쿼리($Q$)와 키($K$)에서 "랜드마크"(Nyström 포인트)를 샘플링합니다. 모델은 쿼리 랜드마크 $\tilde{Q}$와 키 랜드마크 $\tilde{K}$를 기반으로 세 개의 행렬을 정의합니다:

  • $\tilde{F} = \text{softmax}(\frac{Q\tilde{K}^T}{\sqrt{d}})$
  • $\tilde{A} = \text{softmax}(\frac{\tilde{Q}\tilde{K}^T}{\sqrt{d}})^{+}$
  • $\tilde{B} = \text{softmax}(\frac{\tilde{Q}K^T}{\sqrt{d}})$

근사된 소프트맥스 행렬 $\hat{S}$는 다음과 같이 계산됩니다 $\hat{S} = \tilde{F} \tilde{A} \tilde{B}$. 이 근사를 값($V$)과 곱함으로써 모델은 전체 $QK^T$ 곱을 전혀 계산하지 않고도 self-attention을 선형적으로 근사합니다.

랜드마크 선택 및 구현

Nyströmformer는 세그먼트 평균을 사용해 $\tilde{Q}$와 $\tilde{K}$를 구성합니다. $n$개의 토큰을 $m$개의 세그먼트로 나누고, 각 세그먼트의 평균을 랜드마크로 사용합니다. 연구에 따르면 $n = 4096$ 또는 $8192$와 같은 긴 시퀀스 길이에서도 32개 또는 64개의 랜드마크만 사용해도 경쟁력 있는 성능을 제공한다는 것이 밝혀졌습니다.

구현 세부 사항

  • Depthwise Convolution: 아키텍처는 값에 추가되는 스킵 연결로 1D depthwise convolution (DConv)을 포함합니다.
  • Complexity: 이 접근법은 $O(n^2)$ 복잡도를 성공적으로 회피하여 긴 시퀀스에서도 효율성을 확보합니다.

Hugging Face에서의 사용 가능성 및 활용

Nyströmformer는 Hugging Face Transformers 라이브러리를 통해 Masked Language Modeling (MLM) 용으로 제공됩니다. 시퀀스 길이에 따라 네 개의 체크포인트가 제공됩니다: nystromformer-512, nystromformer-1024, nystromformer-2048, 그리고 nystromformer-4096.

사용자는 NystromformerConfig의 num_landmarks 파라미터를 사용해 랜드마크 수 $m$을 제어할 수 있습니다. 모델은 NystromformerForMaskedLM 클래스를 사용하거나 fill-mask와 같은 작업을 위해 고수준 pipeline API를 통해 배포할 수 있습니다.

Sources