BigBird: Transformers for Longer Sequences 논문 리뷰

The overview of this paper

  Transformer 기반의 BERT 같은 모델은 NLP 분야에서 가장 성공한 모델 중 하나이다. 하지만 불행하게도, Transformer 기반 모델들의 가장 큰 약점은 full attention 메커니즘 때문에 sequence의 길이에 따라 메모리 사용량이 곱절로 늘어난다는 것이다. 이를 완화하기 위해, 논문에서는 이러한 quadratic을 linear로 줄인 sparse attention 메커니즘을 사용한 Big Bird를 제안하였다. 논문에서는 BigBird가 sequence 함수의 범용적인 근사치가 되고 Turing complete하다고 말했다. 그리고 BigBird는 full attention model의 quadratic 특성을 보존하였다. 논문의 이론적 분석에서는 sparse attention 메커니즘의 부분으로써 전체 sequence를 참조하는 $O(1)$ global token을 가지는 것의 혜택을 밝혔다. 제안된 sparse attention은 이전의 모델들보다 8배 정도 더 긴 sequence를 처리할 수 있는 능력을 보여줬다. BigBird가 longer context를 처리할 수 있는 능력을 얻은 결과로 다양한 NLP task에 대해 상당히 향상된 성능을 보여줬다.
 
 

Table of Contents

1. Introduction
2. BigBird Architecture

   2-1. Random Attention
   2-2. Sliding Window Attention
   2-3. Global Attention
   2-4. BigBird Attention

3. Theoretical Results about Sparse Attention Mechanism
4. Experiments
 
 

1. Introduction

 Transformer 기반 모델인 BERT는 여러 NLP task에서 상당히 좋은 모습을 보여주고 있다. BERT의 versatility와 robustness는 Transformer가 다양한 task에 적용될 수 있었던 가장 큰 이유 중 하나이다. 그리고 이 BERT는 이전의 LSTM 같은 sequence model을 상당히 상회하는 성능을 보여줬다. Transformer에서 이뤄진 가장 획기적인 발전 중 하나는 self-attention이다. self-attention은 입력 시퀀스의 각 토큰에 대해 병렬로 평가될 수 있고, LSTM 같은 반복 신경망에 존재하는 sequential dependency를 제거하였다. 이러한 병렬 처리는 Transformer가 현대 하드웨어 가속기들의 풀 파워를 leverage함으로써, NLP model이 전례 없는 규모의 데이터셋에서 학습을 할 수 있도록 해준다. 모델은 이렇게 큰 데이터에서 pre-training을 진행하고 그 지식을 downstream task에 적용하는 방식으로 상당한 성능 향상을 이룩할 수 있었다.
 
 self-attention 메커니즘은 RNN의 약점들을 input sequence의 토큰들이 sequence의 다른 토큰들을 독립적으로 볼 수 있게 해줌으로써 극복하였다. 이러한 디자인 선택은 몇 개의 흥미로운 영향을 가지고 있다. 특히, full self-attention은 sequence length에 대해 곱절의 계산 및 메모리를 요구한다. 논문에서는 corpus가 커질 수 있지만, sequence length는 context를 여러 응용에 제공하는 것에는 매우 제한이 있다고 하였다. 일반적으로 사용 가능한 현재의 하드웨어와 모델의 크기를 사용하면 input sequence의 length가 512 token으로 조정되어야만 한다. 이러한 점은 larger context를 필요로 하는 task에 대한 직접적인 응용을 감소시킨다.
 
 하지만, self-attention과 Transformer는 매우 유용하다는 것을 알고 있다. 그렇다면 self-attention model의 어떤 면이 성능에 필수적일까? 그리고 Transformer와 비슷한 모델의 표현성에 대해서 어떻게 말할 수 있을까? 이 질문에 대한 대답은 제안된 self-attention 메커니즘이 RNN만큼 효과적이라면 디자인으로부터 명확한 대답을 내놓을 수 없다. 예를 들어, self-attention이 순열 동등처럼 sequence order를 따르지 않으면 말이다. 이러한 걱정은 Transformer가 압축된 도메인에서 모든 연속된 seq2seq 함수를 캡처할 수 있을 만큼 표현적이라는 점에서 해결되었다. 한편, full Transformer는 Turing Complete하다는 것도 증명되었다. 여기서 Turing Complete라는 것은 Turing test를 통과할 수 있다는 것을 의미한다. 그렇다면 자연스럽게 두 개의 궁금증이 떠오를 것이다.
 

• 적은 inner product를 사용해서 gully quadratic self-attention의 혜택을 성취할 수 있을까?
• sparse attention 메커니즘이 기존 네트워크의 표현성 및 유연성을 보존할 수 있을까?

 
 이 논문에서는 위의 두 개의 질문을 해결하고 longet context를 필요로 하는 task에서 성능을 향상시키는 sparse attention 메커니즘을 제공하였다. token 수에 따라 선형의 복잡도를 가지는 attention 메커니즘을 사용하는 BigBIrd를 체계적으로 개발하였다. 이러한 방법들은 그래프 sparsification method로부터 영감을 받고, 제안된 attention pattern을 형성하기 위해 full-attention이 완화될 때 Transformer의 표현력에 대한 증거가 어디에서 무너지는지 이해하였다. 이러한 이해는 BigBird를 개발할 수 있도록 도와줬고, 이는 이는 이론적으로 표현력이 풍부하고 경험적으로도 유용하다. BigBIrd는 다음 3개의 메인 파트로 이루어져 있다.
 

• $g$ global token 세트가 sequence의 모든 부분을 참조
• 모든 token은 $w$ local 근접 토큰의 세트를 참조
• 모든 token은 $r$ 램던 토큰의 세트를 참조

 
 이러한 방법은 attention 메커니즘이 더욱 긴$($8배$)$ sequence length로 scaling 될 수 있게 이끌어주었다.
 
 이 논문의 contribution을 요약하면 다음과 같다.
 

1. full Transformer의 이론적 특성을 모둔 만족. 추가 토큰을 추가하는 것이 모든 연속되는 seq2seq 함수를 오직 $O(n)$-내적으로 표현하게 하락.
2. BigBird에 의해 모델링된 늘어난 context는 여러 NLP task에서 이점을 보여줌 → SOTA 달성
3. attention 기반 모델의 새로운 응용 소개. 여기서 긴 context는 이점을 가짐.

 

 2. BigBird Architecture

 이 섹션에서는 Transformer의 각 레이어에서 input sequence $X = (x_1, \cdots, x_n) \in \mathbb{R}^{n \times d}$에서 구동되는 generalized attention 메커니즘을 사용하는 BigBird에 대해서 설명하겠다. generalized attention 메커니즘은 vertex 세트가 $[n] = {1, \cdots, n}$인 directed 그래프 $D$에 의해 묘사된다. arc$($directed 그래프의 edge 부분$)$ 세트는 attention 메커니즘이 고려하는 내적 세트를 표현하였다. $N(i)$가 $D$에서 노드 $i$의 out-neighbors 세트를 나타내도록 하게 한 다음 generalized attention 메커니즘의 $i$ 번째 출력 벡터는 다음과 같이 정의된다.
 

 
 여기서 $Q_h, K_h : \mathbb{R}^{d} \to \mathbb{R}^{m}$은 각각 query 함수와 key 함수이고, $V_h : \mathbb{R}^{d} \to \mathbb{R}^{d}$은 value 함수이다. $\sigma$는 scoring 함수이고 $H$는 head의 수를 나타낸다. 또한 $X_{N(i)}$는 모든 입력에 대해서가 아니라 오직 stacking ${x_j : j \in N(i)}$에 의해 형성된 행렬에 해당한다. 설명을 간편하게 하기 위해, 논문에서는 그래프가 sparse 일지라도 그래프 $D$의 인접 행렬 $A$에서 구동해볼 것이다. 자세히 설명하면, $A \in [0, 1]^{n \times n}, A(i, j)=1$일 때 query $i$가 key $j$를 참조하고 그렇지 않으면 0이다. 예를 들어 $A$가 1의 행렬일 때$($BERT처럼$)$, 모든 토큰이 모든 다른 토큰을 참조하기 때문에 quadratic 복잡도를 가진다. self-attention을 완전히 연결된 그래프로 보는 이러한 관점은 기존 그래프 이론을 활용하여 복잡성을 줄이는 데 도움이 된다. self-attention의 quadratice 복잡도를 줄이는 문제는 graph sparsification 문제로 볼 수 있다. 랜덤 그래프는 확장자이며 스펙트럼 속성을 포함하여 다양한 context에서 완전한 그래프를 근사화할 수 있다는 것은 잘 알려져 있다. attention 메커니즘에 대한 sparse 랜덤 그래프는 노드 간의 작은 평균 path length와 locality notion을 가져야만 한다.
 

그림 1. BigBird에서 사용되는 attention 메커니즘의 블록 생성. 하얀색은 attention의 부재를 의미한다. (a) r = 2인 random attention, (b) w = 3인 sliding window attention, (c) g = 2인 global attention, (d)는 이 모든 걸 하나로 합친 BigBirde model

 

2-1. Random Attention
 

 가장 간단한 랜덤 그래프 구조를 생각해보자. 여기서 각각의 edge는 고정 확률로 독립적으로 선택된다. 각 query가 $r$개의 임의의 key, 즉 $r$개의 임의로 선택된 key에 대해 $A(i, ·)=1$에 대해 참조하는 sparse attention 이다.
 

2-2. Sliding Window Attention
 

 BigBird의 탄생에 영감을 준 두 번째 관점은 NLP에서 대부분의 context는 상당한 양의 참조의 locality를 보여주는 데이터를 가지고 있다는 것이다. BigBird와 높은 연관성을 가지고 있는 특정 모델은 평균 최단 거리 경로와 locality의 견해 간에 좋은 밸런스를 달성하였다. 모델의 생성 프로세스는 다음과 같다. regular ring lattice, $n$ 노드와 $w$ 가중치를 연결한 그래프, 각각의 사이드에 $w/2$를 만들어서 context에서 local 구조를 캡처한다.
 
 모든 connection의 랜덤 서브셋 $k_x$는 랜덤 connection으로 대체하였다. 나머지 $(100 - k%)$의 local connection은 유지된다. context의 local 구조를 캡처하기 위해 sliding window attention을 정의해서 사용한다. 그래서 self-attention 중에 width는 $w$이고, location $i$에서 query는 $i-w/2$ 부터 $i+w/2$의 key들을 참조한다. 기호로 표현하면 $A(i, i-w/2 : i+w/2) = 1$로 그림 1의 b처럼 된다.

2-3. Global Attention

BigBird의 마지막 조각은 실험적 성능에 큰 영향을 미치는 논문의 이론적 분석에서 영감을 받았다. 논문의 이론은 시퀀스의 모든 토큰을 참조하는 토큰인 global token의 중요성을 활용하였다. 그림 1의 c를 참고하여라. 이러한 global token은 두 가지 방식으로 정의된다.

• BigBird-ITC: internal transformer construction$($ITC$)$에서는 존재하는 토큰을 global화 해서 전채 시퀀스를 참조하도록 해준다.
• BigBird-ETC: external transformer construction$($ETC$)$에서 CLS같은 추가적인 ‘global’ 토큰 추가해준다. 이 토큰은 존재하는 모든 토큰을 참조한다. 이 방식은 context를 저장할 추가적 location을 제공해준다.

2-4.BigBird Attention

BigBird에 대한 final attention은 이 3개의 특성을 모두 가지고 있다. $($그림 1의 d를 참고하여라$)$ 이때 query가 $r$ 개의 랜덤한 key를 참조한다. 각각의 query는 location의 좌우 양쪽으로 $w/2$의 토큰을 참조하고 $g$ global token을 얻게 된다.

3. Theoretical Results about Sparse Attention Mechanism

 이번 섹션에서는 sparse attention 메커니즘이 full-attention만큼 강력하고 표현력이 좋은 이유에 대해서 두 관점으로 보여주겠다.

1. sparse attention이 encoder-only에서 사용될 때 seq2seq 함수의 universal approximator가 된다.
2. sparse encoder-decoder transformer는 Turing Complete하다.

 

4. Experiments

 본 포스트에서는 genomic에 관한 부분은 다루지 않고 NLP 부분만 다루었다.

 NLP task를 위해 longer input sequence를 모델링하는 것의 혜택을 보여주었다. 이를 위해 기존의 MLM에서 시작해서 longer contiguous sequence를 활용함으로써 더 나은 contextual representation을 학습할 수 있는가를 확인하였다. 다음으로 longer sequence를 다루는 것이 TF-IDF/BM25를 사용하는 것보다 더 많은 증거들을 되찾게 해준다는 것을 보여주기 위해 지지 증거를 사용하는 QA를 이용하였다. 마지막으로 정보가 첫 512 토큰에 위치해있는지를 식별하는 long document classification에 태클을 걸었다. 아래에서는 sequence length 4096을 사용하는 BigBird에 대한 결과를 요약하였다. 그 과정에서 sequence length 외에 컴퓨팅 자원, 배치 크기, 스텝 크기 등과 같은 모든 세팅을 달리하였다.

Pretraining & MLM  MLM을 이용하여 BigBird의 base와 large 버전을 만들었다. 논문에서는 pre-training을 할 때 4개의 표준 데이터셋을 사용하였고, public RoBERTa checkpoint로부터 시작되었다. 결과를 보면 BigBird와 Longformer가 한정된 길이의 RoBERTa보다 좋은 성능을 보여주고 BigBird-ETC가 가장 좋은 성능을 보여줬다. 논문에서는 모델을 학습시킬 때 16GB의 메모리와 32-64의 배치 사이즈를 사용허였다고 말하였다. 이러한 메모리 효율성은 sparse attention의 효율적인 blocking과 sparsity 구조 덕분이다.

Question Answering$($QA$)$  논문에서는 어려운 task들인 NaturalQ, HotpotQA, TriviQA, WikiHop에 대해 테스트를 진행하였다. 그 결과는 다음의 표 1과 같다.

 

표 1. Base size model 사용해서 QA Dev 결과

 실험을 진행할 때 base model을 사용하였고 각 데이터셋에 대해 development set에서 최고의 구성을 선택하였다. 그 결과 확장된 global token을 사용하는 BigBird-ETC가 다른 모델들보다 좋은 성능을 보여준다는 것을 알 수 있었다. 따라서 이 구조를 선택하여 hidden test set에서 평가를 위한 large model을 학습시켰다.

 표 2에서 BigBird-ETC와 리더보드의 top3 모델들을 비교하였다. 그 결과 longer context를 사용하는 Longformer와 BigBird가 smaller context를 사용하는 다른 모델들보다 나은 성능을 보여줬다. 그리고 또 놀라운 것은 BigBird는 single model이지만, 다른 top3 모델들은 ensemble 모델이라는 점이다.

 

표 2. QA task의 test set에 대한 fine-tuning 결과

Classification  서로 다른 길이와 콘텐츠를 갖고 있는 다양한 문서 분류 데이터셋과 GLUE 데이터셋에 대해 테스트를 진행하였다. BERT를 따라서 논문에서는 첫 [CLS] token의 위에 cross entropy를 사용한 하나의 레이어를 사용하였다. 논문에서는 BigBird를 사용하는 것의 효과가 longer document와 적은 training example일 때 더욱 커진다고 하였다.

4-1. Encoder-Decoder Tasks

 

 encoder-decoder 세팅에서는 full self attention 때문에 quadratic 복잡도를 양쪽에서 모두 겪는다는 것을 알 수 있다. BigBird에서는 오직 encoder 측면에서만 sparse attention을 소개하였는데 이는 왜냐하면 예를 들어 요약 문제에서 출력은 200개 내지의 단어가 사용되는데 입력값으로는 3000개 내지의 단어가 들어오기 때문이다. 따라서 응용을 할 때 encoder에는 sparse attention을 사용하고 decoder에는 full self attention을 사용하는 것이 효율적이다고 밝혔다.


Summarization  document summarization은 text document의 짧고 정확한 요약을 생성하는 task이다. 논문에서는 long document의 추상적 요약에 집중하여 longer contextual encoder을 사용라는 것이 성능을 향상시킨다는 것을 보여줬다. 그에 대한 이유는 다음과 같다.

1. 가장 중요한 문장이 512개의 토큰 외의 부분이 분포되어 있을수도 있다.
2. longer document는 더 풍부한 담론 구조와 요약을 보여준다. 그러면서 더 많은 context를 관찰할 수 있다.

 pre-training은 generative task에 도움을 주기 때문에 논문에서는 base model에서 일반적 목적 MLM pre-training으로부터 warm start 뿐만 아니라 SOTA를 달성한 요약 모델의 pre-training을 활용하였다. 그 결과 longer context를 모델링하는 것이 상당한 향상을 가져온다는 것을 알 수 있었다.

 

표 3. long document에 대한 요약 ROGUE 점수

출처
https://arxiv.org/abs/2007.14062

Big Bird: Transformers for Longer Sequences

https://sh-tsang.medium.com/brief-review-big-bird-transformers-for-longer-sequences-12ccd3430e3b

Review — Big Bird: Transformers for Longer Sequences