[논문리뷰] Longformer: The Long-Document Transformer

**논문: https://arxiv.org/abs/2004.05150

**깃헙: https://github.com/allenai/longformer

 

---

MOTIVATION

Transformer의 self-attention은 인풋 시퀀스 길이의 제곱에 비례하여 메모리와 계산량이 늘어난다. 이는 그동안 길이가 긴 텍스트를 인풋으로 받는 NLU 태스크를 처리하는 데에 병목이 되었다. 긴 컨텍스트에 대한 질의응답과 같이 NLU에는 긴 시퀀스를 처리해야 하는 문제들이 존재하고, BERT의 경우 다음과 같은 방법들을 통해 Transformer의 512토큰 제한 문제를 우회했다:

   1) 512 글자 제한에 맞도록 인풋 문서 일부분을 버리기

   2) stride를 주며 독립적으로 문서의 조각들을 처리

   3) 관련 있는 문서 조각을 받아와 모델 추론을 수행하는 두 단계의 추론 수행

하지만 위의 세 가지 방법들은 모두 하나로 된 텍스트를 쪼개서 처리하는 데에서 오는 정보 손실을 감수해야 한다.

 

본 논문에서 제안하는 <Longformer>는 

- 시퀀스 길이에 선형적으로 비례하는 메모리 & 계산량 스케일의 attention 매커니즘 제안하였고

- 이 방법은 기존의 self-attention을 즉각적으로 대체할 수 있다. (drop-in 방식으로 대체 가능)

 

Transformer에서 처리할 수 있는 인풋 길이를 늘리고자 하는 시도는 기존에도 있었으나, 대부분의 실험이 Language modeling의 성능 향상에 국한되었다는 아쉬움이 있었다. 그러나 Longformer에서는 제안한 매커니즘을 기계독해(TriviaQA), 문서 분류, 일관성 모델링 태스크와 같은 다운스트림 NLU 태스크에 적용해 실험하였고, 그 결과 RoBERTa보다 일관적으로 좋은 성능을 보이며 데이터셋 중 두 개에 대해서는 SOTA 성능을 보이기까지 하였다. 

---

모델

Longformer에서 제안한 attention은 기존의 full-attention의 self-attention 행렬을 특정 패턴에 따라 sparse하게 만들어서 사용한다. 논문에서 사용한 attention 패턴은 다음의 세 가지이다. 

 

 

Sliding Window (b, 이하 '윈도우 방식'으로 표기)

토큰 근처의 w개의 고정된 크기의 윈도우에 대해서만 attention 값을 계산하고, 윈도우를 벗어나 멀리 떨어진 토큰에 대해서는 attend하지 않는 방식. 이러한 레이어를 여러 층 쌓게 되면 CNN과 비슷하게 receptive field가 커지고, 결국 넓은 context의 정보를 가지는 representation을 생성할 수 있다. l 층의 레이어를 쌓았을 때, 마지막 레이어에서 토큰에 대한 receptive field는 l x w이다. 

 

윈도우 방식의 패턴에서는 특정 단어 앞 뒤 1/2*w 만큼의 토큰에 대해 attention을 계산하고, 따라서 인풋 시퀀스 길이 n에 대한 계산 복잡도는 O(n x w), 즉 선형 복잡도를 가지게 된다. 이 때 이 방식이 효율적이기 위해서는 시퀀스 길이 n에 비해 윈도우 크기 w가 작은 수여야 한다는 제약이 있다.

 

 

Dilated Sliding Window (c, 이하 '확장된 윈도우 방식'으로 표기)

그림과 같이 d개씩 토큰을 건너뛰며 윈도우를 듬성듬성하게 확장하여 구성하는 방법이다. 이 경우, l개 층을 쌓았을 때 receptive field의 크기는 l x d x w으로, d를 작게 설정하더라도 기본적인 윈도우 방식보다 훨씬 넓은 영역에 대한 representation을 생성할 수 있다.

 

특히 Transformer에서는 multi-head attention을 사용하는데, 각각의 head에 대해 다른 dilation 수치를 적용하는 것이 모델 성능 향상에 도움이 되었다. 어떠한 head는 보다 지역적인 정보를 모으는 데에 집중하게 되고, 다른 head는 더 긴 문맥 정보를 처리하는 데에 집중할 수 있게 되기 때문이다. 

 

 

Global Attention (d)

텍스트 분류, 질의응답과 같은 태스크에서는 [CLS], [SEP]과 같은 스페셜 토큰을 사용한다. 이러한 토큰의 경우 전체적인 문맥에 대해 attend하도록 하는 것이 유리하고, 따라서 논문에서는 미리 선택해 둔 스페셜 토큰의 위치에 대해서는 global attention을 적용하여 전체 토큰에 대해 attend할 수 있도록 하였다. 

 

스페셜 토큰의 개수는 문맥 길이에 비해 매우 적기 때문에 몇몇 토큰에 대해 global attention을 적용하는 것은 계산에 있어 큰 부담이 되지 않는다. 그러나 이러한 방법을 통해 모델은 전체 시퀀스에 대한 문맥적인 representation을 생성하는 법을 학습하게 되고, 이는 다운스트림 태스크를 수행하는 데에 있어 유리하게 작용한다.

 

Global attention과 윈도우 방식은 기본적으로 attention 계산을 위해 필요한 Q, K, V에 대한 웨이트를 학습해야 한다는 점에서 본질적으로 다르지 않다. 하지만 본 논문에서는 윈도우 방식에 사용한 웨이트와 global attention에 사용한 웨이트를 별도로 두었다. (지역적인 정보를 모으는 역할을 하는 아이와 전체 문맥 정보를 가지고 와야 하는 아이를 분리한 것으로 생각됨) 이때 global attention의 웨이트를 윈도우 방식에서 사용하는 웨이트와 같게 초기화 하는 것이 성능에 좋았다고 한다.

 

Attention 계산을 위해 필요한 CUDA 커널

Sliding window 방식의 attention을 구현하려면, banded matrix 곱셈(대각선 일부를 제외한 모든 output을 0으로 만드는 연산)을 구현해야 한다. 그런데 현존하는 딥러닝 라이브러리(Pytorch, TensorFlow)는 이러한 연산을 지원하지 않고, 나이브한 방식으로 이 곱셈을 구현하면 연산이 매우 느려진다. 

 

그래서, 논문의 저자는 CUDA 커널을 커스터마이즈하여 효율적으로 attention 연산이 가능하게 했다!! (대단 ㅠㅠ)

 

Tensor Virtual Machine (TVM)

- TVM을 이용해 커스텀 CUDA 커널을 구현

- TVM은 함수에 대한 high-level description을 최적화된 device-specific한 코드로 바꿔주는 딥러닝 컴파일러. 

- 이를 이용함으로써 banded matrix 곱셈을 GPU 쓰래드에 병력 계산까지 가능한 CUDA 코드로 바꿨다. 

- 커스터마이즈한 Longformer는 최적화된 full self attention과 속도가 거의 비슷하면서도 메모리 사용이 매우 효율적이다!

 

결론적으로 모델은 ...

- Autoregressive Language Modeling 셋팅에서 모델을 개발하고 평가함

 

- 확장된 윈도우 방식(dilated sliding window)의 attention 사용

   > 낮은 layer에서는

       a. 작은 윈도우 사용 : 지역적인 정보를 모델링 & 계산 효율성 확보

       b. dilate하지 않음 : 즉각적으로 지역 정보를 활용

   > 높은 layer에서는

       a. 큰 윈도우 사용 : 인풋에 대한 high-level representation 얻을 수 있도록 함

       b. 2개의 head에 대해서만 dilate : 지역적인 정보를 희생하지 않으면서도 먼 거리에 있는 토큰에도 attend할 수 있음

 

LM 태스크

- 이상적으로는 최대한 큰 윈도우 크기를 사용해 GPU에 올릴 수 있는 최대 길이 시퀀스에 대해 훈련하는 것이 좋다. 하지만 고차원의 문맥 정보를 활용하기에 앞서 일단 지역적인 정보를 배우기 위해서는 다량의 그라이언트 업데이트가 필요하다. 따라서 저자는 학습 단계에 따라 attention 윈도우의 크기와 인풋 시퀀스의 길이를 점진적으로 증가시키는 전략을 선택하였다.

 

- 구체적으로 첫 번째 단계에서는 짧은 시퀀스에 대해 작은 윈도우로 학습을 시작하여, 이후 훈련 단계에 따라 윈도우 크기와 시퀀스 길이를 두 배로 늘리면서 learning rate는 반감하는 전략을 취했다. 이로써 가장 학습 시간이 오래 드는 '긴 시퀀스에 대해 큰 윈도우 사이즈로 학습하는' 부분을 마지막 단계에 두고, 전체적인 학습을 빠르게 만들 수 있었다.

 

- 최종적으로는 2,048개의 토큰에서 시작하여 5개의 훈련 단계를 거쳐 마지막에는 23,040개의 토큰까지 학습하게 된다.

 

주요 Hyper-parameter

• relative position embedding w/ sinusoidal weight 사용 (TransformerXL과 동일)
• small모델 (12레이어, 512 히든 사이즈) / large모델 (30레이어, 512 히든 사이즈)에 대해 실험
• mixed precision training (16&32) / attention 계산은 늘 fp32로 해서 불안정성을 피함
• gradient checkpointing을 사용해 48GB RTX8000 GPU에서 실험 진행

Language Modeling 실험 결과 (enwik8)

- Transformer XL보다 성능이 좋고 Sparse Transformer 성능과 비슷하며, 파라미터 개수가 두배에 달하는 최근 모델에는 성능이 못 미친다.

- 밑단 레이어에서 위로 갈수록 윈도우 크기를 늘려가는(32->512) 전략이 가장 성능이 좋았다.

- 상위 레이어에서 두 개의 head에 dilation을 적용했을 때 BPC=1.20으로 가장 우수했다. 

 

다운스트림 태스크

- 최신 SOTA논문들과 마찬가지로 다운스트림 태스크 수행을 위해 Longformer을 다량의 코퍼스에 사전학습한 후 텍스트 분류, QA 등 여섯 가지 하위 태스크에 fine-tuning하는 전략을 선택하여 성능을 높인다. 이렇게 학습한 Longformer은 최대 4,096토큰까지 처리할 수 있다. (BERT의 8배 길이)

 

- 성능을 위해 RoBERTa 체크포인트에서 attention부분만 Longformer의 패턴을 끼워넣어 이어 학습함

   > 윈도우 방식(sliding window, w=512)의 attention 패턴을 선택 -> RoBERTa와 계산량은 똑같아짐

   > Position Embedding은 RoBERTa가 학습한 weight를 활용하기 위해 512토큰 임베딩을 copy하는 방식을 선택. 이는 BERT의 attention head가 지역적인 정보에 특히 attend하도록 학습한다는 분석에서 착안하여 일부 점프 구간을 제외하고는 기존에 학습한 local 구조를 유지하기 위함임. 이 방법은 매우 간단하면서도 효율적인 것으로 나타남.

   > 결과적으로 RoBERTa에서 이어서 학습한 Longformer은 적은 그라디언트 업데이트로도 빠르게 수렴하였다.

 

- RoBERTa의 base와 large 모델을 활용해 배치 사이즈 64로 시퀀스 길이 4,096의 인풋에 대해  65K번의 그라디언트 업데이트

- 최대 학습률은 3e-5로, 500스텝간은 linear warmup을 거친 후 3 polynomial decay 룰에 의해 점차적으로 감소

 

결과

 

 

특히 긴 문단을 다뤄야 하는 QA, WikiHop 태스크와 TriviaQA에서는 SOTA를 달성

 

---

**깃허브 링크

https://github.com/allenai/longformer

allenai/longformer