Hugging Face Reads: Long-range Transformers

TL;DR

Hugging Face는 표준 Transformer의 이차적 메모리 및 시간 복잡도 병목 현상을 해결하기 위해 Longformer, Compressive Transformer, Linformer, Performer의 네 가지 주요 아키텍처 접근 방식을 분석했습니다. 이러한 방법들은 문서 수준의 NLP, 음성 및 단백질 모델링에 필수적인 기존 512 또는 1024 토큰 제한을 훨씬 넘어서는 시퀀스 처리를 가능하게 합니다.

이차적 병목 현상 극복하기

표준 Transformer의 self-attention은 시퀀스 길이 $n$에 따라 이차적으로($O(n^2)$) 확장되므로, 긴 문서에 대해 상당한 메모리 및 연산 병목 현상을 생성합니다. 이를 해결하기 위해 연구자들은 이 복잡도를 선형($O(n)$)으로 줄이는 것을 목표로 하는 "Efficient Transformers"를 개발했습니다. 이러한 접근 방식은 일반적으로 사용자 정의 attention pattern, recurrence, low-rank approximation, kernel approximation의 네 가지 범주로 나뉩니다.

장거리 모델링을 위한 아키텍처 접근 방식

Longformer: Custom Attention Patterns

Longformer는 표준 self-attention을 windowed (local) 및 global attention의 조합으로 대체하여 시퀀스 길이에 따라 선형적으로 확장될 수 있도록 합니다.

  • Mechanism: autoregressive language modeling을 위해 dilated windowed self-attention을 사용하고, encoder pre-training을 위해 local windowed 및 global bi-directional attention의 혼합을 사용합니다. Global attention은 task-specific token(예: [CLS] token 또는 QA의 question token)에 적용되어 정보가 전체 시퀀스에 걸쳐 흐를 수 있도록 합니다.
  • Key Advantage: self-attention layer는 drop-in replacement이므로, 비용이 많이 드는 처음부터의 pre-training 없이도 pre-trained checkpoint를 장거리 입력에 맞게 조정할 수 있습니다.
  • Trade-off: sliding window attention은 특정 하드웨어(예: TPU)에서 느려질 수 있는 indexing operation에 의존합니다.

Compressive Transformer: Recurrence

Transformer-XL을 기반으로 구축된 Compressive Transformer는 버려질 수 있는 과거의 activation을 저장하기 위해 압축된 메모리를 도입합니다.

  • Mechanism: compression function(예: max/mean pooling 또는 1D convolution)을 사용하여 과거의 activation을 $c$ 배만큼 압축합니다. 이를 통해 모델이 최근 토큰의 regular memory와 훨씬 오래된 토큰의 compressed memory 모두에 attention을 줄 수 있게 합니다.
  • Key Advantage: enwik8 및 WikiText-103과 같은 장거리 language modeling benchmark에서, 특히 긴 거리에 걸쳐 나타나는 희귀 단어에 대해 perplexity를 크게 개선합니다.
  • Trade-off: Training은 취약하며, effective batch size를 점진적으로 늘리는 특수한 optimization schedule이 필요합니다.

Linformer: Low-Rank Approximations

Linformer는 self-attention matrix가 low-rank라는 관찰을 바탕으로 시퀀스 길이를 더 작은 차원으로 투영(projection)하여 복잡도를 줄입니다.

  • Mechanism: Johnson-Lindenstrauss lemma를 활용하여, Linformer는 attention context matrix의 low-rank decomposition을 학습합니다. 이는 $n imes n$ matrix가 계산되거나 저장되지 않도록 보장합니다.
  • Key Advantage: Inference speed(time-clock)는 시퀀스 길이 증가에 영향을 받지 않으며, convergence speed는 표준 Transformer와 비교하여 안정적으로 유지됩니다.
  • Trade-off: decomposition은 training 시 결정된 고정된 context length를 위해 설계되었으며, 적응 과정 없이 더 긴 시퀀스에 일반화되지 않습니다.

Performer: Kernel Approximations

Performer는 sparsity나 low-rank prior를 사용하지 않고 softmax attention kernel을 근사사(approximation)하기 위해 FAVOR+ (Fast Attention Via Orthogonal Random positive features) 알고리즘을 사용합니다.

  • Mechanism: random feature map을 사용하여 softmax function을 근사하여, query multiplication 이전에 $K imes V$ matrix multiplication을 수행할 수 있게 합니다. 이는 $n imes n$ attention matrix의 계산을 효과적으로 우회합니다.
  • Key Advantage: attention matrix의 구조에 대해 어떠한 가정을 하지 않기 때문에, 음성 및 단백질 시퀀스를 포함한 다양한 modality에 매우 높은 적용성을 가집니다.
  • Trade-off: 작은 approximation error가 여러 Transformer layer를 통해 전파될 수 있으며, 이는 pre-trained network의 fine-tuning 안정성에 영향을 미칠 수 있습니다.

비교 분석 및 토론

Inductive Biases and Trade-offs

아키텍처의 선택은 작업의 특정 요구 사항과 사용 가능한 데이터에 따라 달라집니다:

  • Longformer vs. Linformer: Longformer는 fixed sparse pattern을 사용하며, Linformer는 low-rank factorization을 학습합니다. Longformer는 일반적으로 Linformer보다 효율성이 낮지만 시퀀스 길이에 대해 더 유연합니다(단, Linformer는 training context length에 제한됩니다).
  • Performer: Kernel 자체를 근사함으로써 두 방식 모두와 차 differs, attention matrix가 sparse하거나 low-rank라는 가정을 하지 않는 다재능적인 drop-in replacement입니다.

Positional Embeddings

Positional encoding은 장거리 효율성의 핵심 요소입니다:

  • Relative Positional Embeddings: Transformer-XL 및 Compressive Transformers에서 사용됩니다. 이들은 보이지 않은 시퀀스 길이로 쉽게 확장 가능하지만, 연산 비용이 높습니다.
  • Absolute Positional Embeddings: Longformer 및 Linformer에서 사용됩니다. 이들은 연산 효율성이 더 높지만, training 시 보았던 시퀀스보다 더 긴 시퀀스에 대해서는 유연성이 떨어집니다.
  • Position-Infused Attention: Shortformer에서 도입된 대안으로, token embedding 대신 query와 key에 positional information을 추가합니다.

Training Strategies

Shortformer 및 Longformer와 같은 모델의 증거는 짧은 시퀀스에서 training을 진행하고 길이를 점진적으로 늘려가는 것이 더 빠른 training과 더 강력한 downstream performance를 이끌어낸다는 것을 시사합니다. 이는 모델이 데이터의 가 spurious correlations에 대한 의존성을 방지합니다.

Sources