Zebra: Extending Context Window with Layerwise Grouped Local-Global Attention
This addresses a critical bottleneck for LLMs in processing long text sequences, though it appears incremental as an architectural modification to existing Transformer-based approaches.
The paper tackles the challenge of extending context windows for Large Language Models by proposing Zebra, a novel architecture using grouped local-global attention layers to reduce computational complexity. Results show Zebra achieves comparable or superior performance on benchmarks while improving efficiency.
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.