Subindo de 4K para 400K: Expandindo o Contexto de LLMs com o 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

Resumo

A utilização de contextos longos representa um grande desafio para modelos de linguagem de grande escala devido ao comprimento limitado de sua janela de contexto. Embora a janela de contexto possa ser estendida por meio de ajuste fino, isso resultará em um custo considerável tanto no tempo de treinamento quanto no de inferência, além de impactar negativamente as capacidades originais do LLM. Neste trabalho, propomos o Activation Beacon, que condensa as ativações brutas do LLM em formas mais compactas, permitindo que ele perceba um contexto muito mais longo com uma janela de contexto limitada. O Activation Beacon é introduzido como um módulo plug-and-play para o LLM. Ele preserva totalmente a capacidade original do LLM em contextos curtos, ao mesmo tempo em que estende a nova capacidade de processar contextos mais longos. Além disso, ele funciona com janelas deslizantes curtas para processar o contexto longo, alcançando uma eficiência competitiva de memória e tempo tanto no treinamento quanto na inferência. O Activation Beacon é aprendido pela tarefa de autoregressão condicionada a uma mistura de beacons com taxas de condensação diversificadas. Graças a esse tratamento, ele pode ser treinado de forma eficiente apenas com dados de sequências curtas em apenas 10 mil passos, consumindo menos de 9 horas em uma única máquina com 8 GPUs A800. Os estudos experimentais mostram que o Activation Beacon é capaz de estender o comprimento de contexto do Llama-2-7B em 100 vezes (de 4K para 400K), ao mesmo tempo em que alcança um resultado superior em tarefas de geração e compreensão de contexto longo. Nosso modelo e código estarão disponíveis no repositório 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.