Flexformer: Flexible Linear Transformer with Learnable Attention Kernel

기술 요약: Flexformer

문제 정의

트랜스포머 모델은 장거리 의존성을 포착하기 위해 도트 곱 어텐션(dot-product attention) 메커니ме리즘에 의존하지만, 이 메커니즘은 시퀀스 길이 $N$에 대해 이차 시간 및 공간 복잡도($O(N^2)$)를 유발합니다. 이러한 병목 현상은 긴 시퀀스에 대한 확장성을 심각하게 제한합니다. 커널 기반의 선형 어텐션 방식들은 특징 맵(feature maps)을 통해 어텐션 메커니즘을 근사함으로써 복잡도를 선형($O(N)$)으로 줄이지만, 기존 방식들은 효율성과 표현력 사이에서 트레이드오프를 겪습니다:

• 고정된 커널 (Fixed Kernels): Performer나 RFA와 같은 방법들은 소프트맥스 커널을 근사하기 위해 무작위 푸리에 특징(random Fourier features, RFF)에 의존합니다. 그러나 이들은 커널을 고정된 것으로 취급하며, 소프트맥스가 최적이라고 가정합니다. 이는 모든 시나리오에서 소프트맥스가 반드시 최적인 것은 아니라는 점을 간과합니다.
• 약하게 학습 가능한 커널 (Weakly Learnable Kernels): 최근 연구들(예: Hedgehog, Polaformer)은 소프트맥스의 저엔트로피(low-entropy) 특성을 모방하기 위해 학습 가능한 매핑과 특정 함수(지수 함수, 거듭제곱 함수)를 도입했습니다. 하지만 이러한 방법들은 학습된 커널 패밀리가 광범위한 어텐션 패턴을 수용할 만큼 충분히 표현력이 있는지 보장하지 못하며, 이는 잠재적으로 성능 저하를 초래할 수 있습니다.

핵심 과제는 선형 복잡도를 유지하면서도 고정된 커널 및 약하게 학습 가능한 베이스라인보다 뛰어난 성능을 제공할 수 있는 유연하고 학습 가능한 어텐션 커널을 구축하는 것입니다.

방법론

저자들은 **무작위 푸리에 특징(RFF)**을 사용하여 데이터 주도 방식으로 어텐션 커널을 학습하는 선형 트랜스포머 프레임워크인 Flexformer를 제안합니다.

이론적 기초

이 방법은 스펙트럼 표현 이론(spectral representation theory)에 기반을 둡니다.

1. 정상 커널 (Stationary Kernels): 보크너 정리(Bochner's theorem, 정리 3.1)에 따라, 정상 커널 $k(x-y)$는 스펙트럼 밀도 $p(\omega)$에 대응합니다. 표준 RFF는 이를 고정된 분포(예: 가우시안 분포를 사용하는 소프트맥스)로부터 주파수 $\omega$를 샘플링하여 근사합니다.
2. 비정상 커널 (Non-Stationary Kernels): 야글로프 정리(Yaglom theorem, 정리 3.2)에 따라, 비정상 커널은 결합 스펙트럼 밀도를 통해 표현될 수 있습니다.

Flexformer 아키텍처

Flexformer는 주파수를 고정된 분포에서 샘플링하는 대신, 스펙트럼 주파수를 엔드투엔드(end-to-end)로 최적화되는 학습 가능한 파라미터로 취급합니다.

• Flexformer$_s$ (Stationary): 주파수 $\{\omega_i\}$는 학습 가능한 파라미터입니다. 이를 통해 모델은 데이터로부터 최적의 스펙트럼 밀도 $p(\omega)$를 직접 학습할 수 있습니다. 특징 맵은 다음과 같이 구성됩니다:
  $$\tilde{\phi}_n(x) = \exp\left(\frac{\|x\|^2}{2\sqrt{d}}\right) \phi_n(x)$$
  여기서 $\phi_n(x)$는 학습 가능한 주파수의 코사인 및 사인 항으로 구성됩니다. 이 공식은 소프트맥스 커널의 편향되지 않은 추정치(unbiased estimate)를 엄격하게 포함하므로, 필요할 경우 소프트맥스 어텐션을 복구할 수 있음을 보장합니다.

• Flexformer$_n$ (Non-Stationary): 표현력을 높이기 위해, 저자들은 야글로프 정리를 사용하여 비정상 커널으로 프레임워크를 확장합니다. 여기에는 주파수 쌍 $\{(\omega_{1i}, \omega_{2i})\}$와 스케일링 파라미터 $\tau$(고정된 $2\sqrt{d}$를 대체)를 학습하는 과정이 포함됩니다. 특징 맵은 다음과 같습니다:
  $$\tilde{\phi}_n(x) = \exp\left(\frac{\|x\|^2}{\exp(\tau)}\right) \phi_n(x)$$
  비정상 변형은 이론적으로 정상 변형보다 엄격하게 더 높은 표현력을 가집니다.

복잡도

두 변형 모두 시퀀스 길이 $N$에 대해 선형 시간 복잡도 $O(Ndd')$및 **선형 공간 복잡도**$O(Nd + Nd' + dd')$를 유지합니다. 여기서$d$는 은닉 차원이고$d'$는 특징 맵의 차원입니다.

주요 기여

1. 학습 가능한 커널 프레임워크: 무작위 푸리에 특징 프레임워크 내에서 학습 가능한 스펙트럼 주파수를 활용하여, 모델이 데이터로부터 광범위한 어텐션 커널 패밀리를 학습할 수 있게 하는 Flexformer를 제안합니다.
2. 이론적 표현력: Flexformer 커널 패밀리가 소프트맥스 커널을 엄격하게 포함하며(정상 변형의 경우), 비정상 변형의 경우 엄격하게 더 높은 표현력을 제공함을 입증하여, 고정되거나 경험적으로 설계된 커널의 한계를 해결합니다.
3. 증류 및 전이 가능성: Flexformer가 사전 학습된 트랜스포머로부터 효과적으로 증류되어 소프트맥스 어텐션 성능을 회복할 수 있음을 보여주며, 강력한 교차 도메인 커널 전이성을 나타냅니다.

실험 결과

저자들은 언어 모델링, 시퀀스 분류 및 증류 작업에서 Flexformer를 평가했습니다.

• Long-Range Arena (LRA): Flexformer는 모든 작업(ListOps, Text, Retrieval, Image, Pathfinder)에서 가장 높은 평균 정확도를 달성했습니다. 특히, Flexformer$_n$은 문서 검색(document retrieval), 이미지 분류(image classification), Pathfinder에서 최고의 결과를 얻었으며, 기존의 가장 우수한 선형 어텐션 베이스라인보다 상대적으로 4.4% 높은 평균 정확도를 기록했습니다.
• 자기회귀 언어 모델링 (WikiText-103):
  • 소규모 모델(41M 파라미터)에서 Flexformer 변형들은 모든 다른 선형 어텐션 베이스라인을 능가했으며, 바닐라 트랜스포머의 성능에 근접했습니다.
  • 대규모 모델(247M 파라미터)에서 Flexformer$_n$은 바닐라 트랜스포머를 넘어섰으며, 이는 학습 가능한 커널이 증가된 모델 용량으로부터 큰 이득을 얻음을 시사합니다.
• 효율성: 문서 검색 작업에서 Flexformer는 바닐라 트랜스포머와 비교하여 2.6배의 훈련 속도 향상과 84%의 메모리 사용량 감소를 달려냈으며, 시간 및 메모리 측면에서 선형 스케일링을 유지했습니다.
• 증류 (GLUE 벤치마크): 사전 학습된 RoBERTa 모델을 증류할 때, Flexformer$_n$은 소프트맥스 어텐션 동작을 거의 완벽하게 복구하였으며, 학습 불가능한 선형 방식(Performer, RFA)을 능가하고 심지어 CoLA 데이터셋에서는 표준 소프트맥스보다 뛰어난 성능을 보였습니다.
• 커널 전이 가능성: 교차 도메인 실험에서, 한 GLUE 데이터셋에서 학습된 Flexformer$_n$의 커널은 Hedgehog에서 학습된 커널보다 다른 데이터셋으로 더 효과적으로 전이되었으며, 도메인 외(out-of-domain) 작업에서 더 적은 성능 저하를 보였습니다.

의의 및 주장

본 논문은 Flexformer가 긴 시퀀스 모델링을 위한 효과적이고 확장 가능한 소프트맥스 어텐션의 대안을 제공한다고 주장합니다. 이 연구의 주요 의의는 선형 어텐션의 효율성과 표현력 사이의 긴장을 해결하는 데 있습니다:

• 소프트맥스가 최적이라는 가정에서 벗어나 데이터로부터 직접 커널을 학습합니다.
• 경험적인 휴리스틱에 의존하는 대신, 스펙트럼 표현 이론(spectral representation theory)을 통해 유연한 커널을 구축하는 원칙적인 접근 방식을 제공합니다.
• 선형 어텐션이 모델 규모를 키우거나 증류할 때, 계산 효율성을 유지하면서도 이차 어텐션의 성능과 일치하거나 이를 능가할 수 있음을 입증합니다.

한계점: 저자들은 Flexformer가 학습된 어텐션 커널이 양의 정부호(positive definite)임을 가정한다는 점을 언급했습니다. 양의 정부호 커널이 어텐션에 있어 이론적으로 최적의 선택인지 여부는 향려 연구를 위한 미해결 과제로 남아 있습니다.