Descubriendo las Joyas en las Primeras Capas: Acelerando LLMs de Largo Contexto con Reducción de Tokens de Entrada de 1000 veces

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

Resumen

Los Modelos de Lenguaje de Gran Tamaño (LLMs) han demostrado capacidades notables en el manejo de entradas de contexto largo, pero esto conlleva un aumento en los recursos computacionales y la latencia. Nuestra investigación introduce un enfoque novedoso para el cuello de botella del contexto largo con el fin de acelerar la inferencia de LLM y reducir el consumo de memoria de la GPU. Nuestra investigación demuestra que los LLM pueden identificar tokens relevantes en las capas iniciales antes de generar respuestas a una consulta. Aprovechando esta perspicacia, proponemos un algoritmo que utiliza las capas iniciales de un LLM como filtros para seleccionar y comprimir tokens de entrada, reduciendo significativamente la longitud del contexto para el procesamiento posterior. Nuestro método, GemFilter, muestra mejoras sustanciales tanto en velocidad como en eficiencia de memoria en comparación con técnicas existentes, como la atención estándar y SnapKV/H2O. Destacadamente, logra una aceleración de 2.4 veces y una reducción del 30\% en el uso de memoria de la GPU en comparación con los métodos de estado del arte. La evaluación en la tarea de "Encontrar una Aguja en un Pajar" muestra que GemFilter supera significativamente a la atención estándar, SnapKV y demuestra un rendimiento comparable en el desafío LongBench. GemFilter es simple, no requiere entrenamiento y es ampliamente aplicable en diferentes LLMs. Esencialmente, proporciona interpretabilidad al permitir a los humanos inspeccionar la secuencia de entrada seleccionada. Estos hallazgos no solo ofrecen beneficios prácticos para la implementación de LLM, sino que también mejoran nuestra comprensión de los mecanismos internos de los LLM, allanando el camino para futuras optimizaciones en el diseño e inferencia de LLM. Nuestro código está disponible en 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.