Dai 4K ai 400K: Estendere il Contesto degli LLM con Activation Beacon

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

January 7, 2024

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

cs.AI

Abstract

L'utilizzo di contesti lunghi rappresenta una grande sfida per i grandi modelli linguistici a causa della loro limitata lunghezza della finestra contestuale. Sebbene la finestra contestuale possa essere estesa attraverso il fine-tuning, ciò comporterà un costo considerevole sia in fase di addestramento che di inferenza, e avrà un impatto sfavorevole sulle capacità originali del LLM. In questo lavoro, proponiamo Activation Beacon, che condensa le attivazioni grezze del LLM in forme più compatte in modo che possa percepire un contesto molto più lungo con una finestra contestuale limitata. Activation Beacon è introdotto come un modulo plug-and-play per il LLM. Preserva completamente le capacità originali del LLM sui contesti brevi estendendo al contempo la nuova capacità di elaborare contesti più lunghi. Inoltre, funziona con finestre scorrevoli brevi per elaborare il contesto lungo, ottenendo un'efficienza competitiva in termini di memoria e tempo sia in fase di addestramento che di inferenza. Activation Beacon viene appreso attraverso il task di auto-regressione condizionato su una miscela di beacon con rapporti di condensazione diversificati. Grazie a questo approccio, può essere addestrato in modo efficiente utilizzando esclusivamente dati a sequenza breve in soli 10K passi, che richiedono meno di 9 ore su una singola macchina con 8 GPU A800. Gli studi sperimentali dimostrano che Activation Beacon è in grado di estendere la lunghezza contestuale di Llama-2-7B di 100 volte (da 4K a 400K), ottenendo al contempo un risultato superiore sia nei task di generazione che di comprensione di contesti lunghi. Il nostro modello e il codice saranno disponibili nel repository 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.