Longformer

Longformer는 긴 문서를 위한 Transformer로, 전체 셀프 어텐션을 각 토큰 주변의 슬라이딩 윈도 어텐션과 작업에 따라 선택된 소수 토큰에 대한 글로벌 어텐션으로 대체합니다. 따라서 계산 비용은 입력 길이에 따라 제곱이 아닌 선형으로 증가합니다. Allen Institute for AI의 Iz Beltagy, Matthew Peters, Arman Cohan은 2020년 논문 Longformer: 장문 Transformer에서 이 모델을 소개했습니다. BERT 및 RoBERTa 같은 인코더는 최대 512개 토큰만 읽으므로 더 긴 텍스트는 서로 간의 연결이 끊기는 청크로 나누어야 했습니다. 공개된 Longformer 모델은 한 번에 4,096개 토큰을 읽어 컨텍스트 윈도를 8배로 늘렸습니다.

Longformer 어텐션 작동 방식#

각 레이어는 하나의 어텐션 메커니즘 안에서 로컬 패턴과 글로벌 패턴을 결합합니다.

• 슬라이딩 윈도 어텐션: 각 토큰은 고정된 이웃 윈도에 주의를 기울입니다. 공개 모델에서는 양쪽으로 256개 토큰입니다.
• 확장 슬라이딩 윈도: 어텐션을 적용하는 위치 사이에 간격을 두어 계산 비용은 동일하게 유지하면서 헤드가 더 먼 위치까지 볼 수 있게 합니다. 저자들은 문자 수준 언어 모델의 상위 레이어에 있는 일부 헤드에 확장을 적용했습니다.
• 글로벌 어텐션: 선택된 토큰은 모든 위치에 주의를 기울이고 모든 위치는 별도의 쿼리, 키, 값 프로젝션을 통해 선택된 토큰에 주의를 기울입니다. 작업에 따라 선택되는 토큰이 달라집니다. 분류에는 [CLS] 토큰을, 질문 응답에는 질문 토큰을 사용합니다.

레이어를 쌓으면 쌓인 합성곱이 수용 영역을 확장하는 것처럼 윈도의 도달 범위가 넓어집니다. 즉, 합성곱 신경망(CNN)에서 윈도 크기가 w인 l개 레이어는 l × w개 토큰을 포괄합니다. 각 레이어는 n × n개가 아닌 약 n × w개의 어텐션 점수를 계산합니다. Longformer 마스크의 쌍을 계산하면 차이가 분명하게 드러납니다.

import torch

window = 512  # Longformer-base: 256 tokens on each side of every token

for n in (1024, 2048, 4096):
    pos = torch.arange(n)
    # Sliding window: each token attends to neighbors within window / 2 positions
    mask = (pos[:, None] - pos[None, :]).abs() <= window // 2
    # Global attention on token 0 ([CLS]): it sees every token and every token sees it
    mask[0, :] = True
    mask[:, 0] = True
    print(f"{n} tokens: full attention {n * n / 1e6:.1f}M scores, Longformer {mask.sum().item() / 1e6:.1f}M")

출력에는 전체 어텐션의 1.0M, 4.2M, 16.8M과 비교해 Longformer 점수가 0.5M, 1.0M, 2.0M으로 표시됩니다. 입력이 두 배가 될 때마다 Longformer의 어텐션 연산량은 두 배로 늘지만 전체 어텐션의 연산량은 네 배로 늘어납니다. 실제 구현에서는 점수의 밴드 부분만 계산하며, 이로 인해 메모리를 절약할 수 있습니다.

학습 및 결과#

Longformer는 RoBERTa 체크포인트에서 시작합니다. 저자들은 학습된 위치 임베딩 512개를 여덟 번 복사해 4,096개 위치를 처리하고, 긴 문서로 마스크 언어 모델 사전 학습을 이어 갔습니다. 파인튜닝 후 Longformer는 긴 문서 작업에서 RoBERTa를 일관되게 능가했고, WikiHop 및 TriviaQA 질의응답 벤치마크에서 당시 최고 성능을 기록했습니다. 문자 수준 언어 모델은 text8과 enwik8에서 새로운 성능을 달성했습니다.

Longformer 인코더-디코더(LED)는 텍스트 요약 및 기타 텍스트 생성 작업을 위한 디코더를 추가하며, arXiv의 전체 과학 논문을 요약하는 작업으로 성능을 입증했습니다. Hugging Face Transformers는 두 모델을 모두 구현하며 글로벌 토큰을 global_attention_mask로 받습니다. 이후 Clinical-Longformer와 같은 도메인별 버전이 등장했습니다. 이 모델은 대규모 임상 텍스트 코퍼스로 사전 학습되었으며 개체명 인식, 질문 응답, 문서 분류를 포함한 10개 작업에서 평가되었습니다.

Longformer와 관련 모델 비교#

• Transformer-XL: 긴 텍스트를 세그먼트 단위로 읽으면서 각 세그먼트가 이전 세그먼트의 캐시된 은닉 상태에 주의를 기울이도록 합니다. 이를 통해 왼쪽에서 오른쪽으로 생성할 때 재귀 방식으로 컨텍스트를 확장합니다.
• BigBird: 같은 해에 Google이 발표한 모델로, 윈도 및 글로벌 어텐션에 무작위 어텐션 연결을 추가하며 4,096개 토큰도 읽습니다.
• 희소 어텐션: Longformer가 속하는, 고정된 어텐션 패턴으로 구성된 더 넓은 계열입니다.
• Flash Attention: 온칩 메모리 타일에서 정확한 전체 어텐션을 계산하므로 메모리는 길이에 따라 선형으로 증가하지만 점수의 개수는 제곱으로 증가합니다.

비전에서의 로컬 및 글로벌 어텐션#

고해상도 이미지에서도 같은 문제가 발생합니다. 비전 Transformer(ViT)는 이미지를 수천 개의 패치 토큰으로 변환하기 때문입니다. Swin Transformer는 이동된 로컬 윈도로 이 문제에 대응하며, Ultralytics 패키지에서 실행되는 Meta의 Segment Anything Model(SAM)은 블록마다 윈도 어텐션과 글로벌 어텐션을 혼합해 이미지 분할을 수행합니다. 이 모델의 ViT-B 이미지 인코더는 12개 블록 중 8개에서 64 × 64 패치 그리드를 겹치지 않는 14 × 14 윈도로 나누어 어텐션을 적용하고, 블록 2, 5, 8, 11에서는 4,096개 패치 전체에 글로벌 어텐션을 적용합니다.