Zebra: Расширение контекстного окна с помощью послойно сгруппированного локально-глобального внимания

Zebra: Extending Context Window with Layerwise Grouped Local-Global Attention

December 14, 2023

Авторы: Kaiqiang Song, Xiaoyang Wang, Sangwoo Cho, Xiaoman Pan, Dong Yu

cs.AI

Аннотация

В данной статье представлен новый подход для расширения возможностей крупных языковых моделей (LLM) в обработке и понимании длинных текстовых последовательностей, что является критически важным аспектом для приложений, требующих глубокого анализа и синтеза больших объемов информации. Учитывая присущие трудности в расширении контекстного окна для LLM, основанных преимущественно на архитектуре Transformer, мы предлагаем новую архитектуру модели, названную Zebra. Эта архитектура эффективно решает проблемы квадратичной временной и пространственной сложности, связанные с полным вниманием в Transformer, за счет использования группированных локально-глобальных слоев внимания. Наша модель, подобно чередующимся полосам зебры, балансирует между локальными и глобальными слоями внимания, значительно снижая вычислительные требования и потребление памяти. Для оценки производительности Zebra проведены комплексные эксперименты, включая предварительное обучение с нуля, продолжение адаптационного обучения для длинных контекстов и тонкую настройку на длинных инструкциях. Результаты показывают, что Zebra демонстрирует сопоставимую или превосходящую производительность на бенчмарках как для коротких, так и для длинных последовательностей, одновременно повышая эффективность обучения и вывода.

English

This paper introduces a novel approach to enhance the capabilities of Large Language Models (LLMs) in processing and understanding extensive text sequences, a critical aspect in applications requiring deep comprehension and synthesis of large volumes of information. Recognizing the inherent challenges in extending the context window for LLMs, primarily built on Transformer architecture, we propose a new model architecture, referred to as Zebra. This architecture efficiently manages the quadratic time and memory complexity issues associated with full attention in the Transformer by employing grouped local-global attention layers. Our model, akin to a zebra's alternating stripes, balances local and global attention layers, significantly reducing computational requirements and memory consumption. Comprehensive experiments, including pretraining from scratch, continuation of long context adaptation training, and long instruction tuning, are conducted to evaluate the Zebra's performance. The results show that Zebra achieves comparable or superior performance on both short and long sequence benchmarks, while also enhancing training and inference efficiency.