Introduction
Motivation
NLP 및 다양한 분야에서 사용되는 BERT와 같은 Transformer 기반의 모델들은 의 시간 및 공간 복잡도를 가진다는 문제점(구체적으로는 모든 query, key vector끼리의 Attention을 수행하는 Self-Attention때문)이 있습니다. 이로 인해 모델들은 input sequence의 길이(토큰의 수)가 증가할수록 복잡도가 상당히 증가합니다.
위와 같은 문제들로 인해 대부분의 Transformer 기반 모델들은 input sequence 길이에 제한을 두고 있습니다. 하지만 이러한 길이 제한은 Question & Answering, Text Summarization과 같은 길이가 긴 sequence를 처리해야 하는 task들에서는 성능 저하를 초래합니다.
Big Bird는 Transformer의 Full Self-Attention을 approximation할 수 있는 Block Sparse Attention이라는 방법론을 제안하여 위에서 언급한 문제를 해결합니다.
Contributions
- 기존 Transformer 모델보다 8배 더 긴 input sequence를 다룰 수 있는 Big Bird 모델 제안 : Transformer 기반 모델들은 일반적으로 512 tokens까지 처리가 가능하지만 Big Bird 모델은 4096 tokens을 처리 가능
- Full Self-Attention을 approximation하여 효율적인 Attention을 수행할 수 있는 Sparse Attention 방법론을 제안 : Random Attention, Window Attention, Global Attention으로 구성된 Block Sparse Attention 방법론 제안
Related Work
Long Sequence를 처리하기 위한 연구들은 다음과 같은 두 가지 방향으로 진행되어 왔습니다.
- Sequence 길이가 제한된 모델을 여러번 사용
- Full Attention을 사용하지 않고 Sparse Attention을 도입하여 연산량을 줄이려 시도
- 입력으로 들어온 모든 토큰들이 다른 모든 토큰들에 대해서 Attention되어야 하는지에 대한 의문을 품고 효율적으로 Attention을 하기 위한 방법론들을 연구
- 대표적인 연구들 : LongFormer, LinFormer, PerFormer, ReFormer, ETC(Extended Transformers Construction)
- Big Bird의 경우, 이 방법론의 일종으로 ETC 모델을 기반으로 함
Proposed Method
Big Bird Architecture
Big Bird는 Transformer에서 Full Attention을 사용했을 때의 성능을 유지하거나 넘어서면서 효율적으로 Attention을 할 수 있는 방법론을 찾는 것을 목표로 합니다. 이를 위해 Big Bird는 Transformer architecture를 기반으로 Self-Attention layer를 Full Attention이 아닌 Sparse Attention으로 대체합니다. 따라서, Big Bird의 가장 중요한 contribution인 Block Sparse Attention을 이해하는 것이 중요합니다.
Block Sparse Attention은 다음과 같은 3가지의 Attention으로 구성되면 Block Sparse Attention은 이 Attention들을 합친 것입니다.
- Random Attention : query token과 개의 무작위의 key tokens간의 Attention
- Window Attention(=Sliding Attention) : query token과 양옆 개의 key tokens간의 Attention
- Global Attention : query token과 개의 global tokens과의 Attention (global tokens은 모든 token들과 Attention을 수행)
저자들은 처음에는 Random Attention과 Window Attention만 사용하여 실험을 했으나, 다음과 같이 BERT 모델의 성능을 따라잡지는 못했다고 합니다.
이에 ETC(Extended Transformer Construction) 모델에서의 global tokens의 우수성을 참고하여 Global Attention을 도입했습니다. 이 부분에 대해서 저자들은 QA(Question Answering) task와 같이 Global Context 정보가 중요할 경우, Random Attention과 Window Attention은 이 정보를 충분히 반영하지 못한다고 설명합니다.
Big Bird에서는 다음과 같은 두 가지 방식의 global tokens을 도입하여 Global Attention을 수행했습니다.
- BigBird-ITC(Internal Transformer Construction) : 텍스트 내부의 특정 토큰들을 global token으로 정하는 방식 ex) 첫번째 token, 마지막 token
- BigBird-ETC(Extended Transformer Construction) : 텍스트 내부가 아닌 추가적인 토큰들을 도입하여 global token으로 정하는 방식 ex) BERT 모델의 [CLS] token
다음의 표는 ITC와 ETC의 차이를 보여줍니다. 표의 Attention Matrix에서 "1"은 Attention을 계산하는 부분을 의미합니다.
아래의 표는 QA(Question & Answering) task에서 다른 모델들과 BigBird-ITC와 BigBird-ETC의 성능 차이를 보여줍니다. Big Bird 모델이 다른 모델들에 비해 성능이 성능이 우수하며 ICT와 ETC의 성능은 거의 유사함을 확인할 수 있습니다.
Understanding about Block Sparse Attention with Graph Theory
논문의 저자들은 Block Sparse Attention에서 Random Attention, Window Attention, Global Attention이 필요한 이유를 Graph Tehory를 통해 설명하고 있습니다.
우선, 저자들은 query token에서 key token을 Attention하는 경우, graph에서의 edge가 있다고 가정합니다. 예를 들어, Full Attention의 경우, 다음과 같은 fully-connected graph로 나타낼 수 있습니다.
또한, Random Attention, Window Attention, Global Attention은 다음과 같이 각각 표현될 수 있습니다.
마지막으로 위의 Random Attention, Window Attention, Global Attention을 합친 Block Sparse Attention은 다음과 같이 표현 가능합니다.(이 때, Block Sparse Attention이 왜 "Block" Sparse Attention인지에 대해서는 다음 section에서 설명하도록 하겠습니다.)
이렇게 표현된 Graph들을 통해 저자들은 이론적으로 각각의 Attention이 필요한 이유를 다음과 같이 설명합니다.
- Small average path length between nodes : Complete graph에서 edge들이 특정 확률에 의해 남겨진 random graph를 생각해볼 때, 두 노드의 최단 경로(shortest path)는 node의 개수에 로그분포(logarithmic)에 비례한다는 것이 밝혀져 있습니다. 즉, random graph의 크기를 키우고 sparse하게 만들어도 shortest path가 기하급수적으로 급격히 증가하는 것이 아니고 매우 천천히 변하게 됩니다. 이러한 성질을 기반으로 random graph의 특성이 효율적인 Attention을 가능하게 해줄 것이라고 저자들은 설명하고 있습니다.
- Notion of locality : Graph Theory와 NLP 분야에서 지역성(locality)를 통해 양옆의 nodes(또는 tokens)들간의 정보가 있어서 특정 node(또는 token)의 정보를 충분히 반영할 수 있다고 설명합니다. 직관적으로 생각해볼 때, 양 옆의 tokens이 거리가 먼 tokens보다 더 중요할 것으로 생각할 수 있기 때문에 이러한 설명도 타당해 보입니다.(이는 NLP 분야에서 많이 사용되던 N-gram의 아이디어와도 유사합니다.)
- Importance of global tokens : global tokens을 추가하면 이론적으로 모든 토큰들이 global token을 통해 연결될 수 있으므로 다른 token들간의 Attention 정보가 필요할 때, global token을 통해 정보를 전달해줄 수 있습니다. 즉, 위의 그림들에서처럼 Random Attention에서 토큰 "to"와 "going"은 직접적으로 연결된 edge가 없지만 Global token인 "I"를 통해 서로의 정보를 간접적으로 전달할 수 있습니다. 이러한 이유를 들어, 저자들은 global tokens의 중요성을 강조하고 있습니다.
Implementation Details : Block Sparse Attention
Big Bird의 Sparse Attention은 Sparse Matrix가 만들어지기 때문에 GPU, TPU에서 Parallel하게 적용하기 어렵습니다. 따라서, 논문에서는 Blockifying Attention을 도입하여 Parallel Attention이 가능하도록 하였습니다.
Blockifying Attention에서는 다음과 같이 각 query, key vector를 block으로 만들고 block 단위의 행렬곱을 진행하여 Attention Score를 계산합니다. 아래에서 확인하실 수 있듯이 Blockifying Attention을 Full Attention에 대해 수행하면 대각성분들에 대해서만 Attention Score가 구해집니다.
이를 Big Bird의 Sparse Attention에 적용하기 위해서는 다음과 같이 4가지의 hyper parameter들이 도입됩니다.(아래의 표에서 는 로 쓰기도 합니다.)
가장 먼저 Window Attention(=Sliding Attention)의 경우,
번째 query block이 범위의 key block들과 Attention을 수행하는 과정은 다음과 같습니다. 아래의 그림에서 확인할 수 있듯이, 의 수 만큼(아래의 그림에서는 ) 동일한 key block들을 복사한 후에 각 block들을 오른쪽과 왼쪽으로 shift하여 이동시켜 연산하고자 하는 최종적인 key block들을 얻는 것을 확인할 수 있습니다.
이러한 과정을 통해 얻어진 key block들은 다음 그림과 같이 query block들에 대해 Blockifying Attention을 수행합니다.
다음으로 Random Attention의 경우, gather 연산을 통해 개의 random key block들을 모으고 query와의 Blockifying Attention 연산을 수행합니다.
마지막으로 Global Attention도 위의 과정과 유사하게 개의 global key block들을 모아서 Blockifying Attention을 수행합니다.
이렇게 만들어진 모든 Attention들은 GPU/TPU의 병렬 연산을 사용하기 위해 Dense Attention 구조를 만듭니다. 즉, 모든 key block들을 다음과 같이 compact tensor로 표현하여 Attention을 수행합니다.
개인적으로는 Random Attention, Window Attention, Global Attention을 합쳐 GPU/TPU 연산이 가능한 구조로 만든 Block Sparse Attention이 이 논문의 가장 큰 contribution이라고 생각했습니다.
Experiments
논문에서는 NLP Experiments와 Genomics Exmperiments를 수행했습니다. Genomics Experiments는 RNA, DNA 염기서열을 일련의 sequence 형태로 보는 것이기 때문에 본질적으로 NLP Experiments과 다르지 않습니다. 따라서, 저는 NLP Experiments 부분만 다루도록 하겠습니다.
우선, Question & Answering task에서는 Neural Questions, HotpotQA-distractor, TriviaQA-wiki, WikiHop 데이터셋들에 대해 실험을 했으며 실험 결과는 다음과 같습니다.(기본적으로 이 데이터셋들은 long sequence의 context text를 입력으로 합니다.)
아래의 표에서 확인하실 수 있듯이, LongFormer, SpanBERT 등보다 성능이 좋은 것을 확인할 수 있습니다.
다음으로 논문에서는 Document Classification task를 위해 다양한 길이와 내용을 담고 있는 데이터셋들에 대해 실험을 수행하였습니다.
그 결과, Arxiv 데이터셋과 같이 길이가 더 길면서 학습 데이터가 적은 경우에 큰 성능 향상을 보였다고 합니다.
마지막으로 Summarization task에서도 실험을 진행하였는데, 실험 결과 모든 데이터셋에서 SOTA를 달성했다고 합니다.
Conclusion
Big Bird는 Block Sparse Attention을 도입하여 효율적인 Self-Attention이 가능하도록 했습니다.
특히, Big Bird 논문은 Sparse Attention을 GPU/TPU 연산이 가능하도록 구현하는 방법론이 인상 깊었습니다.
앞으로 로민이 다루게 될 문서의 종류와 길이가 다양할 것이라고 생각하며 이로 인해 Long Sequence를 처리할 수 있는 Efficient Transformer에 대한 스터디를 하는 것이 중요하다고 생각하여 이 주제를 선택하였습니다.
Share article