De 4K a 400K: Extendiendo el contexto de los LLM con Activation Beacon

Soaring from 4K to 400K: Extending LLM's Context with Activation Beacon

January 7, 2024

Autores: Peitian Zhang, Zheng Liu, Shitao Xiao, Ninglu Shao, Qiwei Ye, Zhicheng Dou

cs.AI

Resumen

La utilización de contextos largos representa un gran desafío para los modelos de lenguaje de gran escala debido a la longitud limitada de su ventana de contexto. Aunque la ventana de contexto puede extenderse mediante ajuste fino, esto conlleva un costo considerable tanto en el tiempo de entrenamiento como en el de inferencia, y ejerce un impacto desfavorable en las capacidades originales del LLM. En este trabajo, proponemos Activation Beacon, que condensa las activaciones brutas del LLM en formas más compactas, permitiendo que perciba un contexto mucho más largo con una ventana de contexto limitada. Activation Beacon se introduce como un módulo plug-and-play para el LLM. Preserva completamente la capacidad original del LLM en contextos cortos mientras extiende la nueva capacidad para procesar contextos más largos. Además, funciona con ventanas deslizantes cortas para procesar el contexto largo, logrando una eficiencia competitiva en memoria y tiempo tanto en el entrenamiento como en la inferencia. Activation Beacon se aprende mediante la tarea de auto-regresión condicionada a una mezcla de beacons con ratios de condensación diversificados. Gracias a este enfoque, puede entrenarse eficientemente únicamente con datos de secuencias cortas en solo 10K pasos, lo que consume menos de 9 horas en una máquina con una sola GPU 8xA800. Los estudios experimentales muestran que Activation Beacon es capaz de extender la longitud de contexto de Llama-2-7B por 100 veces (de 4K a 400K), logrando al mismo tiempo un resultado superior tanto en tareas de generación como de comprensión de contextos largos. Nuestro modelo y código estarán disponibles en el repositorio BGE.

English

The utilization of long contexts poses a big challenge for large language models due to their limited context window length. Although the context window can be extended through fine-tuning, it will result in a considerable cost at both training and inference time, and exert an unfavorable impact to the LLM's original capabilities. In this work, we propose Activation Beacon, which condenses LLM's raw activations into more compact forms such that it can perceive a much longer context with a limited context window. Activation Beacon is introduced as a plug-and-play module for the LLM. It fully preserves the LLM's original capability on short contexts while extending the new capability on processing longer contexts. Besides, it works with short sliding windows to process the long context, which achieves a competitive memory and time efficiency in both training and inference. Activation Beacon is learned by the auto-regression task conditioned on a mixture of beacons with diversified condensing ratios. Thanks to such a treatment, it can be efficiently trained purely with short-sequence data in just 10K steps, which consumes less than 9 hours on a single 8xA800 GPU machine. The experimental studies show that Activation Beacon is able to extend Llama-2-7B's context length by times100 times (from 4K to 400K), meanwhile achieving a superior result on both long-context generation and understanding tasks. Our model and code will be available at the BGE repository.