[논문 리뷰] LM-Infinite: Zero-Shot Extreme Length Generalization for Large Language Models

Chi Han, Qifan Wang외 5명|arXiv (Cornell University)|2023. 08. 30.

Topic Modeling인용 수 5

한 줄 요약

이 논문은 대규모 언어 모델에서 초장거리 일반화를 위한 간단하고 파rameter가 없는 방법인 LM-Infinite를 제안한다. 추론 중에 $λ$-형태의 어텐션 마스크와 거리 제한을 적용함으로써, 미세조정 없이도 최신의 LLM이 128k 토큰까지 유창하고 일관성 있는 텍스트를 생성할 수 있게 한다. 이로 인해 디코딩 속도가 2.72배 빨라지고, 장문 텍스트 벤치마크에서 미세조정된 모델과 유사한 성능을 달성한다.

ABSTRACT

Today's large language models (LLMs) typically train on short text segments (e.g., <4K tokens) due to the quadratic complexity of their Transformer architectures. As a result, their performance suffers drastically on inputs longer than those encountered during training, substantially limiting their applications in real-world tasks involving long contexts such as encoding scientific articles, code repositories, or long dialogues. Through theoretical analysis and empirical investigation, this work identifies three major factors contributing to this length generalization failure. Our theoretical analysis further reveals that commonly used techniques like truncating the attention window or relative positional encodings are inadequate to address them. Answering these challenges, we propose LM-Infinite, a simple and effective method for enhancing LLMs' capabilities of handling long contexts. LM-Infinite is highly flexible and can be used with most modern LLMs off-the-shelf. Without any parameter updates, it allows LLMs pre-trained with 2K or 4K-long segments to generalize to up to 200M length inputs while retaining perplexity. It also improves performance on downstream tasks such as Passkey Retrieval and Qasper in the zero-shot setting. LM-Infinite brings substantial efficiency improvements: it achieves 2.7x decoding speed up and 7.5x memory saving over the original model. Our codes are released at \url{https://github.com/Glaciohound/LM-Infinite}.

연구 동기 및 목표

• 모델의 훈련 길이를 초월해 긴 시퀀스를 생성할 때 발생하는 길이 일반화 실패 문제를 해결한다.
• 장문 시퀀스에서 성능 저하를 유발하는 분포 외 요인(미래의 거리, 과도하게 많은 어텐션을 받는 토큰 수, 암묵적으로 인코딩된 위치 정보)을 규명하고 진단한다.
• 기존에 사전 훈련된 LLM이 파rameter 업데이트나 미세조정 없이도 초장거리로의 일반화를 유창하게 수행할 수 있도록 즉시 적용 가능한 솔루션을 개발한다.
• 다양한 LLM 아키텍처와 장문 텍스트 데이터셋에서 일관된 생성 품질과 높은 계산 효율성을 달성한다.

제안 방법

• 각 쿼리가 처리하는 토큰 수를 제한하는 $λ$-형태의 어텐션 마스크를 도입하여 계산 부담을 줄이고 분포 외 어텐션 패턴을 완화한다.
• 상대적 위치 인코딩에서 고려하는 최대 상대적 거리를 제한하는 거리 제한을 적용하여 추론 중에 발생할 수 있는 새로운 거리에 대비한다.
• 사전 훈련된 모델의 재학습, 미세조정, 파rameter 업데이트 없이 추론 단계에서만 적용한다.
• 장문 시퀀스에 대한 전체 자기어텐션을 피하기 위해 $O(n)$ 시간 및 공간 복잡도를 유지함으로써 효율적인 디코딩을 가능하게 한다.
• RoPE 및 Alibi와 같은 기존의 상대적 위치 인코딩 기반 기법과의 통합을 통해 다양한 LLM 간 광범위한 호환성을 확보한다.
• 디코딩 중 실시간으로 적용하여 아키텍처 변경 없이도 동적 장문 처리가 가능하다.

Figure 1: Diagnosis of three OOD factors in LLMs.

실험 결과

연구 질문

• RQ1 LLM이 훈련 길이를 초월해 생성할 때 발생하는 길이 일반화 실패의 주요 분포 외 요인은 무엇인가?
• RQ2 간단하고 파rameter가 없는 방법이 미세조정 없이도 장문 텍스트 생성의 유창성과 일관성을 향상시킬 수 있는가?
• RQ3 $λ$-형태의 어텐션 마스크와 거리 제한이 128k 토큰까지의 시퀀스에서 생성 품질을 얼마나 잘 유지할 수 있는가?
• RQ4 이러한 방법은 장문 텍스트 작업에서 잘라내기 기반 모델과 미세조정된 모델에 비해 효율성과 성능 면에서 어떻게 비교되는가?
• RQ5 이 방법은 상대적 위치 인코딩을 사용하는 다양한 LLM 아키텍처(예: LLaMA, Llama-2, MPT, GPT-J)에 대해 일반화 가능한가?

주요 결과

• LM-Infinite는 ArXiv의 32k 토큰 시퀀스에서 LLaMA와 Llama-2가 각각 BLEU 점수 19.7과 18.4를 기록하며, 미세조정된 모델과 유사한 수준의 유창성과 일관성을 유지한다.
• OpenWebText2에서 Llama-2 + LM-Infinite는 16k 토큰에서 BLEU 점수 9.6을 기록했고, 기반 Llama-2 모델은 난잡한 텍스트(bleu = 0.0)를 생성했다.
• 32k 길이의 시퀀스에서 디코딩 속도가 2.72배 빨라지고, 인코딩 속도가 3.16배 빨라져 추론 효율성이 크게 향상되었다.
• 이 방법은 잘라내기 기반 모델보다 약 5점의 BLEU 점수 우수성을 기록했으며, 유사한 계산 비용을 사용했고, 동일한 BLEU 점수를 달성하기 위해 계산량의 25%만으로도 가능했다.
• ArXiv와 OpenWebText2 양쪽에서 128k 토큰까지 일관된 생성 품질을 유지하여, 강력한 제로샷 일반화 능력을 입증했다.
• LLaMA, Llama-2, MPT-7B, GPT-J 등 다양한 최신 LLM 아키텍처에 대해 상대적 위치 인코딩을 사용하는 경우에 모두 일반적으로 적용 가능하며, 모델별 맞춤형 적응이 필요하지 않다.

Figure 2: (a) \model is a plug-and-play solution for various LLMs, consisting of a $\Lambda$ -shaped mask and a distance constraint during attention. (b) We also provide a conceptual model for understanding how relative position encoding works.