포커스드 트랜스포머: 컨텍스트 스케일링을 위한 대조적 학습

Focused Transformer: Contrastive Training for Context Scaling

July 6, 2023

저자: Szymon Tworkowski, Konrad Staniszewski, Mikołaj Pacek, Yuhuai Wu, Henryk Michalewski, Piotr Miłoś

cs.AI

초록

대규모 언어 모델은 새로운 정보를 맥락적으로 통합하는 탁월한 능력을 가지고 있다. 그러나 이러한 접근법의 전체 잠재력은 종종 효과적인 맥락 길이의 제한으로 인해 제약을 받는다. 이 문제에 대한 한 가지 해결책은 (키, 값) 쌍으로 구성된 외부 메모리에 접근할 수 있는 어텐션 레이어를 부여하는 것이다. 그러나 문서의 수가 증가함에 따라 관련 키 대 비관련 키의 비율이 감소하여 모델이 비관련 키에 더 집중하게 된다. 우리는 서로 다른 의미적 값과 연결된 키가 겹쳐 구분하기 어려워지는, 이른바 '분산 문제'라는 중요한 과제를 확인했다. 이 문제를 해결하기 위해 우리는 대조 학습에서 영감을 받은 훈련 과정을 사용하는 '포커스드 트랜스포머(FoT)'라는 기법을 소개한다. 이 새로운 접근법은 (키, 값) 공간의 구조를 강화하여 맥락 길이를 확장할 수 있게 한다. 우리의 방법은 기존의 대규모 모델을 미세 조정하여 그들의 효과적인 맥락을 연장할 수 있도록 한다. 이는 3B 및 7B OpenLLaMA 체크포인트를 미세 조정한 결과로 입증된다. 우리가 'LongLLaMA'라 명명한 결과 모델들은 긴 맥락이 필요한 작업에서 진전을 보여준다. 우리는 더 나아가 LongLLaMA 모델들이 256k 맥락 길이의 패스키 검색을 능숙하게 처리함을 보여준다.

English

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.