[논문 리뷰] Unlimiformer: Long-Range Transformers with Unlimited Length Input

Amanda Bertsch, Uri Alon외 2명|arXiv (Cornell University)|2023. 05. 02.

Handwritten Text Recognition Techniques인용 수 23

한 줄 요약

Unlimiformer는 교차 주의에 단일 k-최근접 이웃 인덱스를 추가하여 사전학습된 인코더-디코더 트랜스포머를 보강하고, 테스트 시 추가 학습 없이 무한 입력 길이를 가능하게 하며 장문서 및 책 요약 성능을 향상시킨다.

ABSTRACT

Since the proposal of transformers, these models have been limited to bounded input lengths, because of their need to attend to every token in the input. In this work, we propose Unlimiformer: a general approach that wraps any existing pretrained encoder-decoder transformer, and offloads the cross-attention computation to a single k-nearest-neighbor (kNN) index, while the returned kNN distances are the attention dot-product scores. This kNN index can be kept on either the GPU or CPU memory and queried in sub-linear time; this way, we can index practically unlimited input sequences, while every attention head in every decoder layer retrieves its top-k keys, instead of attending to every key. We evaluate Unlimiformer on several long-document and book-summarization benchmarks, showing that it can process even 500k token-long inputs from the BookSum dataset, without any input truncation at test time. We demonstrate that Unlimiformer improves pretrained models such as BART and Longformer by extending them to unlimited inputs without additional learned weights and without modifying their code. We make our code and models publicly available at https://github.com/abertsch72/unlimiformer .

연구 동기 및 목표

• 표준 컨텍스트 윈도우를 넘어선 극히 긴 입력의 처리에 동기를 부여하고 재훈련 없이 가능하게 한다.
• 모든 입력 토큰에 대한 전체 교차 주의를 대체하기 위한 일반적이고 비매개변수(non-parametric) 검색 메커니즘을 제안한다.
• 단일 k-NN 인덱스가 모든 디코더 계층과 헤드의 주의 질량(attention masses)을 근사하는 데 충분함을 보인다.
• 여러 기본 모델과 학습 방식에 걸쳐 장문서 및 책 요약 벤치마크에서 개선을 입증한다.

제안 방법

• 각 디코더 계층의 교차 주의 전에 k-NN 검색 단계를 삽입하고, 헤드당 상위 k개의 키를 선택한다.
• 중첩(overlap)이 있는 청크로 긴 입력을 인코딩하고 각 청크의 숨겨진 상태 중간 절반을 인덱싱한다.
• QWqWk^T를 이용한 헤드별 프로젝션으로 접근되는 인코더 숨겨진 상태의 단일 인덱스를 허용하도록 주의 계산을 재작성하여 계층/헤드 간에 하나의 공유 인덱스를 가능하게 한다.
• 디코딩 시간에 인덱스를 질의하여 상위-k개 검색된 키만 주의하도록 하고, 점곱 거리(dot-product distances)를 주의 점수로 사용한다.
• 메모리 한도를 16비트 숨겨진 상태로 설정하여(예: 1,000,000 토큰의 경우 2 GB) 필요에 따라 인덱스를 CPU/GPU로 오프로딩한다.
• 저비용의 테스트 시 변형들(+test Unlimiformer, +early stop w/ Unlimiformer)과 더 긴 범위의 학습 접근법(Random-encoded, Retrieval, Alternating)을 제공한다.
• LLaMA-2 및 HuggingFace Transformers와 호환되는 코드베이스와 모델 릴리스를 제공한다.

실험 결과

연구 질문

• RQ1 인코더-디코더 트랜스포머에서의 교차 주의를 테스트 시 무한대 길이의 입력을 지원하도록 k-NN 인덱스로 오프로드하는 것이 가능한가?
• RQ2 모든 디코더 계층/헤드에 걸쳐 하나의 공유된 k-NN 인덱스가 효과적인 검색을 보장하고 대부분의 주의 질량을 보존하는 데 충분한가?
• RQ3 추가 학습 파라미터 없이 기존의 사전학습된 모델을 보강하여 무제한 입력 길이를 처리할 수 있는가?
• RQ4 장거리 요약 및 관련 작업에서 Unlimiformer를 사용할 때 정확도와 계산 비용의 trade-off는 무엇인가?

주요 결과

• Unlimiformer는 추가 학습 없이 장문서 요약에서 기준 모델을 향상시킨다(예: BART_base +test Unlimiformer가 표준 파인튜닝보다 ROUGE/L 및 BERTScore가 더 높다).
• Unlimiformer를 이용한 얼리 스톱은 추가 학습 비용 없이 상당한 이득을 제공한다(예: GovReport에서 ROUGE-1이 48.7에서 51.0으로 증가).
• Unlimiformer로 학습될 때 PRIMERA와 같은 모델은 더 큰 장거리 기반 모델을 능가하거나 동등하게 만들고, Unlimiformer가 이를 더 개선할 수 있다(예: PRIMERA +test Unlimiformer가 표준 PRIMERA 대비 ROUGE/L 및 EntMent를 향상).
• BookSum 결과에서 EntMent의 이득이 나타나며(Unlimiformer+PRIMERA가 EntMent를 25.5로, 기본 PRIMERA의 11.6에서 증가).
• 검색 중심 학습 변형(Retrieval, Random-encoded, Alternating)은 데이터 세트 전반에 걸쳐 경쟁력 있는 이득을 제공하며 최적 방법은 모델과 데이터에 따라 다르다.