Descobrindo as Joias nas Camadas Iniciais: Acelerando LLMs de Longo Contexto com Redução de 1000x nos Tokens de Entrada

Discovering the Gems in Early Layers: Accelerating Long-Context LLMs with 1000x Input Token Reduction

September 25, 2024

Autores: Zhenmei Shi, Yifei Ming, Xuan-Phi Nguyen, Yingyu Liang, Shafiq Joty

cs.AI

Resumo

Os Modelos de Linguagem de Grande Escala (LLMs) demonstraram capacidades notáveis no tratamento de entradas de contexto longo, mas isso vem com o custo de recursos computacionais e latência aumentados. Nossa pesquisa introduz uma abordagem inovadora para o gargalo do contexto longo, a fim de acelerar a inferência de LLM e reduzir o consumo de memória da GPU. Nossa pesquisa demonstra que os LLMs podem identificar tokens relevantes nas camadas iniciais antes de gerar respostas a uma consulta. Aproveitando essa percepção, propomos um algoritmo que utiliza as camadas iniciais de um LLM como filtros para selecionar e comprimir tokens de entrada, reduzindo significativamente o comprimento do contexto para processamento subsequente. Nosso método, GemFilter, demonstra melhorias substanciais tanto em velocidade quanto em eficiência de memória em comparação com técnicas existentes, como atenção padrão e SnapKV/H2O. Notavelmente, ele alcança um aumento de velocidade de 2,4 vezes e uma redução de 30% no uso de memória da GPU em comparação com métodos de referência. A avaliação na tarefa Needle in a Haystack mostra que o GemFilter supera significativamente a atenção padrão, SnapKV, e demonstra desempenho comparável no desafio LongBench. O GemFilter é simples, não requer treinamento e é amplamente aplicável em diferentes LLMs. Crucialmente, ele fornece interpretabilidade ao permitir que os humanos inspecionem a sequência de entrada selecionada. Essas descobertas não apenas oferecem benefícios práticos para a implementação de LLMs, mas também aprimoram nossa compreensão dos mecanismos internos dos LLMs, abrindo caminho para mais otimizações no design e inferência de LLMs. Nosso código está disponível em https://github.com/SalesforceAIResearch/GemFilter.

English

Large Language Models (LLMs) have demonstrated remarkable capabilities in handling long context inputs, but this comes at the cost of increased computational resources and latency. Our research introduces a novel approach for the long context bottleneck to accelerate LLM inference and reduce GPU memory consumption. Our research demonstrates that LLMs can identify relevant tokens in the early layers before generating answers to a query. Leveraging this insight, we propose an algorithm that uses early layers of an LLM as filters to select and compress input tokens, significantly reducing the context length for subsequent processing. Our method, GemFilter, demonstrates substantial improvements in both speed and memory efficiency compared to existing techniques, such as standard attention and SnapKV/H2O. Notably, it achieves a 2.4times speedup and 30\% reduction in GPU memory usage compared to SOTA methods. Evaluation on the Needle in a Haystack task shows that GemFilter significantly outperforms standard attention, SnapKV and demonstrates comparable performance on the LongBench challenge. GemFilter is simple, training-free, and broadly applicable across different LLMs. Crucially, it provides interpretability by allowing humans to inspect the selected input sequence. These findings not only offer practical benefits for LLM deployment, but also enhance our understanding of LLM internal mechanisms, paving the way for further optimizations in LLM design and inference. Our code is available at https://github.com/SalesforceAIResearch/GemFilter.