LM-Infinite : Généralisation simple et en temps réel de la longueur pour les grands modèles de langage

LM-Infinite: Simple On-the-Fly Length Generalization for Large Language Models

August 30, 2023

Auteurs: Chi Han, Qifan Wang, Wenhan Xiong, Yu Chen, Heng Ji, Sinong Wang

cs.AI

Résumé

Ces dernières années, les modèles de langage de grande taille (LLMs) basés sur l'architecture Transformer ont connu des avancées remarquables dans divers domaines. Alors que ces LLMs sont déployés pour des tâches de plus en plus complexes, ils sont souvent confrontés à la nécessité de mener des processus de raisonnement plus longs ou de comprendre des contextes plus étendus. Dans ces situations, l'échec de généralisation en termes de longueur des LLMs sur des séquences longues devient plus prononcé. La plupart des schémas de pré-entraînement tronquent les séquences d'entraînement à une longueur fixe (comme 2048 pour LLaMa). Les LLMs ont souvent du mal à générer des textes fluides, et encore moins à accomplir des tâches en aval, après des contextes plus longs, même avec un encodage positionnel relatif conçu pour résoudre ce problème. Les solutions courantes, comme le fine-tuning sur des corpus plus longs, impliquent souvent des coûts matériels et temporels importants et nécessitent une conception minutieuse du processus d'entraînement. Pour exploiter plus efficacement la capacité de génération des LLMs existants, nous étudions théoriquement et empiriquement les principaux facteurs hors distribution (OOD) contribuant à ce problème. Inspirés par ce diagnostic, nous proposons une solution simple mais efficace pour la généralisation en longueur à la volée, LM-Infinite, qui implique uniquement un masque d'attention en forme de Lambda et une limite de distance, sans nécessiter de mise à jour des paramètres ou d'apprentissage. Nous constatons qu'elle est applicable à une variété de LLMs utilisant des méthodes d'encodage positionnel relatif. LM-Infinite est efficace sur le plan computationnel avec une complexité en temps et en espace de O(n), et démontre une fluidité et une qualité de génération constantes jusqu'à 32 000 tokens sur les ensembles de données ArXiv et OpenWebText2, avec une accélération du décodage de 2,72x. Sur des tâches en aval comme la récupération de clé, elle continue de fonctionner sur des entrées bien plus longues que les longueurs d'entraînement où les modèles classiques échouent immédiatement.

English

In recent years, there have been remarkable advancements in the performance of Transformer-based Large Language Models (LLMs) across various domains. As these LLMs are deployed for increasingly complex tasks, they often face the needs to conduct longer reasoning processes or understanding larger contexts. In these situations, the length generalization failure of LLMs on long sequences become more prominent. Most pre-training schemes truncate training sequences to a fixed length (such as 2048 for LLaMa). LLMs often struggle to generate fluent texts, let alone carry out downstream tasks, after longer contexts, even with relative positional encoding which is designed to cope with this problem. Common solutions such as finetuning on longer corpora often involves daunting hardware and time costs and requires careful training process design. To more efficiently leverage the generation capacity of existing LLMs, we theoretically and empirically investigate the main out-of-distribution (OOD) factors contributing to this problem. Inspired by this diagnosis, we propose a simple yet effective solution for on-the-fly length generalization, LM-Infinite, which involves only a Lambda-shaped attention mask and a distance limit while requiring no parameter updates or learning. We find it applicable to a variety of LLMs using relative-position encoding methods. LM-Infinite is computational efficient with O(n) time and space, and demonstrates consistent fluency and generation quality to as long as 32k tokens on ArXiv and OpenWebText2 datasets, with 2.72x decoding speedup. On downstream task such as passkey retrieval, it continues to work on inputs much longer than training lengths where vanilla models fail immediately.