[논문 리뷰] Focused Transformer: Contrastive Training for Context Scaling

Szymon Tworkowski, Konrad Staniszewski외 4명|arXiv (Cornell University)|2023. 07. 06.

Topic Modeling인용 수 13

한 줄 요약

논문은 Focused Transformer(FoT)를 도입하며, 기억(attention) 기반 대조학습으로 컨텍스트 길이를 확장하고, 미세조정된 OpenLLaMA 모델(OpenLLaMA를 LongLLaMA로)로 최대 256k 토큰을 처리하고 장기 컨텍스트 작업을 개선합니다.

ABSTRACT

Large language models have an exceptional capability to incorporate new information in a contextual manner. However, the full potential of such an approach is often restrained due to a limitation in the effective context length. One solution to this issue is to endow an attention layer with access to an external memory, which comprises of (key, value) pairs. Yet, as the number of documents increases, the proportion of relevant keys to irrelevant ones decreases, leading the model to focus more on the irrelevant keys. We identify a significant challenge, dubbed the distraction issue, where keys linked to different semantic values might overlap, making them hard to distinguish. To tackle this problem, we introduce the Focused Transformer (FoT), a technique that employs a training process inspired by contrastive learning. This novel approach enhances the structure of the (key, value) space, enabling an extension of the context length. Our method allows for fine-tuning pre-existing, large-scale models to lengthen their effective context. This is demonstrated by our fine-tuning of $3B$ and $7B$ OpenLLaMA checkpoints. The resulting models, which we name LongLLaMA, exhibit advancements in tasks requiring a long context. We further illustrate that our LongLLaMA models adeptly manage a $256 k$ context length for passkey retrieval.

연구 동기 및 목표

• 멀티-문서 주의에서 컨텍스트 확장을 제한하는 산만 문제를 식별한다.
• 키-값 공간을 형성하여 처음부터 재학습하지 않고도 FoT로 컨텍스트 길이를 확장한다.
• 기존 모델(OpenLLaMA)을 실용적으로 미세조정하여 더 긴 컨텍스트를 가진 LongLLaMA를 생성하는 것을 시연한다.
• 패스키 검색, QA, 페무샷 작업에서 장기 컨텍스트 능력을 실험적으로 평가하고 학습 길이를 넘어 외삽(extrapolation)을 보인다.

제안 방법

• 추론 중에 kNN을 통해 추가( key, value ) 컨텍스트를 검색하는 메모리 어텐션 레이어를 추가한다.
• 대조 학습에서 영감을 얻은 crossbatch 절차로 메모리 어텐션을 양성(현재/이전 컨텍스트)과 음성(다른 문서)에 노출시키며 학습한다.
• 메모리 계층에서( LongLLaMA를 제외하고) 위치 인코딩을 제거하여 외삽을 가능하게 한다; 정확한 kNN 검색을 위해 FAISS를 사용한다.
• FoT로 OpenLLaMA-3B 및 OpenLLaMA-7B를 미세조정하여 LongLLaMA 모델을 생성하고, 패스키 검색(passkey retrieval), Qasper, TREC, WebQS 및 장기 컨텍스트 데이터셋(PG-19, arXiv, GitHub, Isabelle)을 평가한다.
• FoT가 키-값 공간의 구조를 개선하여 전체 재학습 없이도 확장된 컨텍스트에 주의를 가능하게 함을 시연한다.
• FoT 기반 모델은 학습 한계를 넘어 컨텍스트 길이를 외삽할 수 있음을 보여준다(예: 패스키 검색에서 256k 토큰).

실험 결과

연구 질문

• RQ1 FoT가 많은 문서에 걸쳐 트랜스포머 컨텍스트를 확장할 때 산만 문제를 완화할 수 있는가?
• RQ2 FoT가 그러한 길이의 학습 없이도 훨씬 더 긴 컨텍스트로 외삽을 가능하게 하는가?
• RQ3 더 긴 컨텍스트에서도 FoT로 학습된 모델이 퍼플렉시티를 유지하거나 향상시키며 페워샷 성능을 보이는가?
• RQ4 패스키 검색 및 문서 QA와 같은 장기 컨텍스트 작업에서 실용적 이점은 무엇인가?
• RQ5 다른 모델 크기 및 데이터셋(PG-19, arXiv, GitHub, Isabelle, Qasper, TREC, WebQS)에서 FoT의 성능은 어떠한가?

주요 결과

• FoT는 crossbatch 대조 학습을 통해 키-값 공간을 형성함으로써 산만 문제를 완화한다.
• LongLLaMA-3B 및 LongLLaMA-7B는 256k 토큰까지 외삽하며, 패스키 검색에서 100k에서 94.5%의 정확도, 256k에서 73%를 달성한다.
• LongLLaMA 모델은 더 긴 컨텍스트에서 TREC 및 WebQS에서 정확도 향상을 보인다.
• FoT는 표준 미세조정 베이스라인과 달리 학습 컨텍스트 길이를 넘어 외삽을 가능하게 한다.
• PG-19, arXiv, GitHub, Isabelle에서 FoT의 더 긴 컨텍스트로 퍼플렉시티가 개선되며, 베이스라인 대비 최대 64k 컨텍스트 이득을 보인다.
• FoT는 아키텍처를 변경하지 않고 기존 모델을 미세조정하여 컨텍스트를 확장할 수 있는 간단한 플러그앤플레이 확장 기능이다.