От 4K до 400K: Расширение контекста языковых моделей с помощью активационных маяков

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

January 7, 2024

Авторы: Peitian Zhang, Zheng Liu, Shitao Xiao, Ninglu Shao, Qiwei Ye, Zhicheng Dou

cs.AI

Аннотация

Использование длинных контекстов представляет собой серьезную проблему для больших языковых моделей из-за ограниченной длины их контекстного окна. Хотя контекстное окно можно расширить с помощью тонкой настройки, это приведет к значительным затратам как на этапе обучения, так и на этапе вывода, а также негативно повлияет на исходные возможности модели. В данной работе мы предлагаем Activation Beacon, который сжимает исходные активации языковой модели в более компактные формы, что позволяет воспринимать значительно более длинный контекст при ограниченном размере окна. Activation Beacon представлен как модуль plug-and-play для языковой модели. Он полностью сохраняет исходные возможности модели на коротких контекстах, одновременно расширяя её способность обрабатывать более длинные контексты. Кроме того, он работает с короткими скользящими окнами для обработки длинного контекста, что обеспечивает конкурентоспособную эффективность по памяти и времени как на этапе обучения, так и на этапе вывода. Activation Beacon обучается на задаче авторегрессии, учитывающей смесь маяков с различными коэффициентами сжатия. Благодаря такому подходу, его можно эффективно обучить исключительно на данных с короткими последовательностями всего за 10 тысяч шагов, что занимает менее 9 часов на одной машине с 8 GPU A800. Экспериментальные исследования показывают, что Activation Beacon способен увеличить длину контекста модели Llama-2-7B в 100 раз (с 4K до 400K), одновременно демонстрируя превосходные результаты как в задачах генерации, так и в задачах понимания длинных контекстов. Наша модель и код будут доступны в репозитории 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.