De 4K à 400K : Étendre le contexte des LLM avec Activation Beacon

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

January 7, 2024

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

cs.AI

Résumé

L'utilisation de contextes longs représente un défi majeur pour les grands modèles de langage en raison de la longueur limitée de leur fenêtre contextuelle. Bien que cette fenêtre puisse être étendue par un ajustement fin, cela entraînerait un coût considérable tant au niveau de l'entraînement que de l'inférence, tout en ayant un impact défavorable sur les capacités originales du modèle. Dans ce travail, nous proposons **Activation Beacon**, qui condense les activations brutes du modèle en des formes plus compactes, permettant ainsi de percevoir un contexte beaucoup plus long avec une fenêtre contextuelle limitée. Activation Beacon est introduit comme un module plug-and-play pour le modèle. Il préserve entièrement les capacités originales du modèle sur des contextes courts tout en étendant ses nouvelles capacités pour traiter des contextes plus longs. De plus, il fonctionne avec des fenêtres glissantes courtes pour traiter le contexte long, ce qui permet d'atteindre une efficacité compétitive en termes de mémoire et de temps, tant à l'entraînement qu'à l'inférence. Activation Beacon est appris par une tâche d'autorégression conditionnée sur un mélange de balises avec des ratios de condensation diversifiés. Grâce à cette approche, il peut être efficacement entraîné uniquement avec des données de séquences courtes en seulement 10 000 étapes, ce qui consomme moins de 9 heures sur une machine équipée d'une seule carte GPU 8xA800. Les études expérimentales montrent qu'Activation Beacon est capable d'étendre la longueur contextuelle de Llama-2-7B par un facteur de 100 (de 4K à 400K), tout en obtenant des résultats supérieurs sur les tâches de génération et de compréhension de contextes longs. Notre modèle et notre code seront disponibles dans le dépôt 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.