LM-Infinite: Generalización Simple de Longitud en Tiempo Real para Modelos de Lenguaje a Gran Escala

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

August 30, 2023

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

cs.AI

Resumen

En los últimos años, se han logrado avances notables en el rendimiento de los Modelos de Lenguaje de Gran Escala (LLMs) basados en Transformers en diversos dominios. A medida que estos LLMs se despliegan para tareas cada vez más complejas, a menudo enfrentan la necesidad de llevar a cabo procesos de razonamiento más largos o comprender contextos más amplios. En estas situaciones, el fallo en la generalización de longitud de los LLMs en secuencias largas se vuelve más evidente. La mayoría de los esquemas de preentrenamiento truncan las secuencias de entrenamiento a una longitud fija (como 2048 para LLaMa). Los LLMs suelen tener dificultades para generar textos fluidos, y mucho menos para realizar tareas posteriores, después de contextos más largos, incluso con codificación posicional relativa que está diseñada para abordar este problema. Soluciones comunes como el ajuste fino en corpus más largos a menudo implican costos de hardware y tiempo abrumadores y requieren un diseño cuidadoso del proceso de entrenamiento. Para aprovechar de manera más eficiente la capacidad de generación de los LLMs existentes, investigamos teórica y empíricamente los principales factores fuera de distribución (OOD) que contribuyen a este problema. Inspirados por este diagnóstico, proponemos una solución simple pero efectiva para la generalización de longitud sobre la marcha, LM-Infinite, que implica solo una máscara de atención en forma de Lambda y un límite de distancia, sin necesidad de actualizaciones de parámetros o aprendizaje. Encontramos que es aplicable a una variedad de LLMs que utilizan métodos de codificación posicional relativa. LM-Infinite es computacionalmente eficiente con tiempo y espacio O(n), y demuestra fluidez y calidad de generación consistentes hasta 32k tokens en los conjuntos de datos ArXiv y OpenWebText2, con una aceleración de decodificación de 2.72x. En tareas posteriores como la recuperación de claves, continúa funcionando en entradas mucho más largas que las longitudes de entrenamiento donde los modelos convencionales fallan inmediatamente.

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.