LM-Infinite: Generalização Simples de Comprimento em Tempo Real para Modelos de Linguagem de Grande 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

Resumo

Nos últimos anos, houve avanços notáveis no desempenho de Modelos de Linguagem de Grande Escala (LLMs) baseados em Transformers em diversos domínios. À medida que esses LLMs são implantados para tarefas cada vez mais complexas, eles frequentemente enfrentam a necessidade de realizar processos de raciocínio mais longos ou compreender contextos maiores. Nessas situações, a falha de generalização de comprimento dos LLMs em sequências longas torna-se mais evidente. A maioria dos esquemas de pré-treinamento trunca as sequências de treinamento para um comprimento fixo (como 2048 para LLaMa). Os LLMs frequentemente lutam para gerar textos fluentes, muito menos realizar tarefas subsequentes, após contextos mais longos, mesmo com codificação posicional relativa, que foi projetada para lidar com esse problema. Soluções comuns, como ajuste fino em corpora mais longos, muitas vezes envolvem custos assustadores de hardware e tempo e exigem um design cuidadoso do processo de treinamento. Para aproveitar de forma mais eficiente a capacidade de geração dos LLMs existentes, investigamos teoricamente e empiricamente os principais fatores fora da distribuição (OOD) que contribuem para esse problema. Inspirados por esse diagnóstico, propomos uma solução simples, mas eficaz, para generalização de comprimento em tempo real, o LM-Infinite, que envolve apenas uma máscara de atenção em forma de Lambda e um limite de distância, sem exigir atualizações de parâmetros ou aprendizado. Descobrimos que ele é aplicável a uma variedade de LLMs que utilizam métodos de codificação posicional relativa. O LM-Infinite é computacionalmente eficiente, com tempo e espaço O(n), e demonstra fluência e qualidade de geração consistentes para até 32 mil tokens nos conjuntos de dados ArXiv e OpenWebText2, com uma aceleração de decodificação de 2,72x. Em tarefas subsequentes, como recuperação de chave de acesso, ele continua a funcionar em entradas muito mais longas do que os comprimentos de treinamento, onde os modelos convencionais falham imediatamente.

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.