Understanding BigBird's Block Sparse Attention

BigBird는 표준 Transformer 모델의 $O(n^2)$ 시간 및 메모리 복잡성을 해결하여 최대 4096 토큰까지의 시퀀스를 처리할 수 있게 합니다. 전체 어텐션을 블록 희소 어텐션 메커니즘으로 교체함으로써, BigBird는 긴 문서 요약 및 긴 컨텍스트를 갖는 질문 응답과 같은 장기 시퀀스 작업에서 최첨단 결과를 달성합니다.

Block Sparse Attention Mechanism

BigBird의 블록 희소 어텐션은 효율성을 위해 설계된 BERT 전체 어텐션의 근사치이며, 우수성을 목표로 하지 않습니다. 쿼리 토큰이 주목해야 하는 토큰 수를 세 가지 서로 다른 유형의 어텐션을 결합함으로써 줄입니다:

  • Sliding Attention: 토큰이 인접한 이웃 토큰에 주목합니다. 이는 시퀀스 내 단어들이 일반적으로 인접 토큰에 강하게 의존하기 때문에 지역 의존성을 포착합니다.
  • Global Tokens: 소수의 토큰이 시퀀스의 모든 다른 토큰에 주목하고, 모든 다른 토큰도 이 토큰들을 주목합니다. 이는 슬라이딩 어텐션만으로는 어려운 정보를 시퀀스 전체에 더 빠르게 전달할 수 있게 합니다.
  • Random Tokens: 몇몇 토큰을 무작위로 선택하여 다른 토큰에 주목하게 함으로써, 시퀀스 그래프에서 먼 노드 간 정보 전달 거리를 추가로 줄입니다.

전체 어텐션은 단일 레이어에서 어떤 두 토큰 간에도 정보를 전달할 수 있지만, 블록 희소 어텐션은 정보를 시퀀스 전체에 전달하기 위해 여러 레이어가 필요할 수 있습니다. 전역 및 무작위 연결의 조합은 몇 개의 레이어만으로도 정보가 빠르게 전달될 수 있도록 보장합니다.

Technical Implementation

GPU와 TPU에서 효율적으로 구현하기 위해 BigBird는 블록 기반 접근 방식을 사용합니다. 시퀀스를 크기 $b$인 블록으로 나눕니다.

Attention Calculation by Block

  • First and Last Blocks: 첫 번째 블록 ($q_1$)과 마지막 블록 ($q_n$)은 시퀀스의 모든 다른 토큰에 대해 일반 어텐션 연산을 수행합니다.
  • Middle Blocks: 블록 $q_{3:n-2}$에 속한 토큰에 대해 모델은 전역, 슬라이딩, 무작위 키를 수집하고 선택된 키에 대해서만 어텐션을 계산합니다.
  • Boundary Blocks: 블록 $q_2$와 $q_{n-1}$은 특정 키 집합(첫 번째 블록, 마지막 블록, 인접 슬라이딩 블록 포함)을 수집하여 희소 구조를 유지합니다.

Complexity Comparison

시퀀스 길이가 증가함에 따라 BigBird는 계산 부담을 크게 줄입니다. BERT의 2차 스케일링에 비해 BigBird는 선형으로 확장됩니다:

Attention Type Sequence Length Time & Memory Complexity
Original Full (BERT) 512 $T$
Original Full (BERT) 1024 $4 \times T$
Original Full (BERT) 4096 $64 \times T$
Block Sparse (BigBird) 1024 $2 \times T$
Block Sparse (BigBird) 4096 $8 \times T$

Training Strategies: ITC vs. ETC

BigBird는 두 가지 다른 전략으로 학습될 수 있습니다:

  • Internal Transformer Construction (ITC): 소수의 전역 토큰을 사용합니다. 계산량이 적으며 많은 작업에 충분합니다.
  • Extended Transformer Construction (ETC): 추가 전역 토큰을 도입합니다. 전체 질문을 컨텍스트 토큰이 전역적으로 주목해야 하는 질문 응답과 같은 작업에 특히 유용합니다.

Integration with Hugging Face Transformers

BigBird는 transformers 라이브러리에서 BigBirdModel을 통해 제공됩니다. 주요 구현 세부 사항은 다음과 같습니다:

  • Sequence Length Requirements: 시퀀스 길이는 block_size의 배수여야 합니다. Hugging Face는 자동으로 시퀀스를 블록 크기의 가장 작은 배수로 패딩합니다.
  • Automatic Fallback: 시퀀스 길이가 필요한 전역, 무작위, 슬라이딩 토큰을 지원하기에 너무 짧거나 모델이 디코더(BigBirdForCausalLM)로 사용될 경우, 라이브러리는 자동으로 attention_type을 original_full로 전환합니다.
  • Recommended Usage: 저자들은 1024 토큰보다 짧은 시퀀스에 대해 attention_type="original_full"을 설정할 것을 권장합니다.

사용 가능한 체크포인트에는 bigbird-roberta-base, bigbird-roberta-large, bigbird-base-trivia-itc가 포함됩니다.

Sources