[논문 리뷰] Landmark Attention: Random-Access Infinite Context Length for Transformers

Amirkeivan Mohtashami, Martin Jaggi|arXiv (Cornell University)|2023. 05. 25.

Topic Modeling인용 수 4

한 줄 요약

이 논문은 랜드마크 토큰을 사용해 고정 길이의 입력 블록을 검색하는 데 문을 열어주는 방식으로, Transformer 모델이 무작위 액세스 기능과 무한한 컨텍스트 길이를 갖도록 하는 Landmark Attention을 제안한다. 랜드마크 유사도를 통해 어텐션을 훈련시켜 관련 블록을 검색함으로써, 계산량과 메모리 사용량을 블록 크기와 동일한 요소로 줄일 수 있으며, 이는 예를 들어 50배의 감소를 의미한다. 이를 통해 LLaMA 7B는 32,000자 이상의 토큰에서 추론을 수행할 수 있게 되었으며, 이는 GPT-4의 컨텍스트 길이와 동일하며, Transformer-XL 수준의 성능을 유지한다.

ABSTRACT

While Transformers have shown remarkable success in natural language processing, their attention mechanism's large memory requirements have limited their ability to handle longer contexts. Prior approaches, such as recurrent memory or retrieval-based augmentation, have either compromised the random-access flexibility of attention (i.e., the capability to select any token in the entire context) or relied on separate mechanisms for relevant context retrieval, which may not be compatible with the model's attention. In this paper, we present a novel approach that allows access to the complete context while retaining random-access flexibility, closely resembling running attention on the entire context. Our method uses a landmark token to represent each block of the input and trains the attention to use it for selecting relevant blocks, enabling retrieval of blocks directly through the attention mechanism instead of by relying on a separate mechanism. Our approach seamlessly integrates with specialized data structures and the system's memory hierarchy, enabling processing of arbitrarily long context lengths. We demonstrate that our method can obtain comparable performance with Transformer-XL while significantly reducing the number of retrieved tokens in each step. Finally, we show that fine-tuning LLaMA 7B with our method successfully extends its context length capacity to over 32k tokens, allowing for inference at the context lengths of GPT-4. We release the implementation of landmark attention and the code to reproduce our experiments at https://github.com/epfml/landmark-attention/.

연구 동기 및 목표

• 긴 컨텍스트에서 자기주의 어텐션의 제곱형 메모리 및 계산 비용을 해결하기 위해.
• 반복 없이도 또는 외부 검색기 없이도 어텐션의 무작위 액세스 유연성을 유지하면서 장기간의 컨텍스트를 검색할 수 있도록 하기 위해.
• 재훈련 없이도 훈련 중에 관찰한 길이를 초월하는 임의의 컨텍스트 길이에서 추론을 가능하게 하기 위해.
• 메모리 계층 및 FAISS와 같은 데이터 구조와의 원활한 통합을 통해 대규모 컨텍스트 처리의 효율성을 높이기 위해.
• LLaMA 7B와 같은 사전 훈련된 모델을 이 메커니즘을 통해 초장거리 컨텍스트(예: 32,000토큰)로 확장하기 위해.

제안 방법

• 입력을 고정 길이의 블록으로 나누며, 각 블록은 해당 블록에 대한 어텐션의 게이트로 기능하는 랜드마크 토큰으로 표현된다.
• 추론 중에 랜드마크 토큰으로의 어텐션 점수는 어떤 블록이 검색되고 어텐션 대상이 되는지를 결정하며, 이는 이전의 어떤 블록이라도 무작위 액세스할 수 있도록 한다.
• 표준 어텐션 기법을 사용하지만, 쿼리와 랜드마크 키 벡터 간의 유사도에 따라 검색을 조건화한다.
• 이전 토큰의 캐시(키-값 캐시)를 유지하며, 랜드마크 토큰이 활성화될 때만 해당 랜드마크에 대응하는 비랜드마크 토큰만 로드한다.
• 관련 블록을 검색하기 위해 FAISS와 같은 효율적인 데이터 구조를 활용한다.
• LLaMA 7B와 같은 사전 훈련된 모델을 이 메커니즘을 사용해 미세조정함으로써, 다시 시작부터 훈련하지 않고도 컨텍스트 길이를 연장할 수 있다.

실험 결과

연구 질문

• RQ1 반복이나 외부 검색 모델에 의존하지 않고도 Transformer가 무작위 액세스 기능과 무한한 컨텍스트 길이를 갖출 수 있는가?
• RQ2 랜드마크 기반 검색이 계산량과 메모리 사용량을 크게 줄이면서도 어텐션의 유연성을 유지할 수 있는가?
• RQ3 이 방법을 기존의 대규모 모델에 적용해 초장거리 컨텍스트(예: 32,000토큰)를 처리할 수 있도록 미세조정할 수 있는가, 성능 저하 없이?
• RQ4 랜드마크 어텐션 모델의 성능이 장기간 컨텍스트에서 반복 모델인 Transformer-XL과 비교해 어떻게 되는가?
• RQ5 이 방법을 고급 데이터 구조와 통합해 장기간 추론에서 자원 사용량을 추가로 줄일 수 있는가?

주요 결과

• 랜드마크 어텐션은 모델이 훨씬 짧은 컨텍스트에서 훈련되었더라도 추가 훈련 비용 없이 임의의 컨텍스트 길이에서 추론을 가능하게 한다.
• 이 방법은 계산량과 메모리 사용량을 블록 크기와 동일한 요소로 줄이며, 예를 들어 50토큰 블록을 사용할 경우 50배의 감소를 이룬다. 이는 장기간의 컨텍스트 추론을 효율적으로 가능하게 한다.
• LLaMA 7B를 랜드마크 어텐션으로 미세조정함으로써, 최대 32,000토큰 길이의 컨텍스트에서 숨겨진 암호어를 검색하는 데 높은 정확도를 달성했으며, 이는 GPT-4의 컨텍스트 길이와 동일하다.
• 모델은 장기간의 컨텍스트 작업에서 높은 성능을 유지하며, 32,000토큰을 초월한 영역에서도 어휘 복원 정확도가 거의 완벽한 편이지만, 기본 모델은 실패한다.
• FAISS와 같은 외부 데이터 구조와의 호환성이 있으며, 이는 검색 효율성과 메모리 사용량 최적화를 추가로 가능하게 한다.
• 이 방법은 해석 가능성을 유지하며, 랜드마크로의 어텐션 점수는 어떤 블록이 접근되고 있는지 드러내어 모델의 추론 과정에 대한 통찰을 제공한다.