Zebra: Extensión de la Ventana de Contexto con Atención Local-Global Agrupada por Capas

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

December 14, 2023

Autores: Kaiqiang Song, Xiaoyang Wang, Sangwoo Cho, Xiaoman Pan, Dong Yu

cs.AI

Resumen

Este artículo presenta un enfoque novedoso para mejorar las capacidades de los Modelos de Lenguaje de Gran Escala (LLMs) en el procesamiento y comprensión de secuencias de texto extensas, un aspecto crítico en aplicaciones que requieren una comprensión profunda y síntesis de grandes volúmenes de información. Reconociendo los desafíos inherentes en la extensión de la ventana de contexto para los LLMs, principalmente basados en la arquitectura Transformer, proponemos una nueva arquitectura de modelo, denominada Zebra. Esta arquitectura gestiona eficientemente los problemas de complejidad cuadrática en tiempo y memoria asociados con la atención completa en el Transformer mediante el uso de capas de atención local-global agrupadas. Nuestro modelo, similar a las rayas alternas de una cebra, equilibra capas de atención local y global, reduciendo significativamente los requisitos computacionales y el consumo de memoria. Se llevan a cabo experimentos exhaustivos, que incluyen preentrenamiento desde cero, continuación de entrenamiento de adaptación de contexto largo y ajuste de instrucciones largas, para evaluar el rendimiento de Zebra. Los resultados muestran que Zebra logra un rendimiento comparable o superior en puntos de referencia de secuencias cortas y largas, al mismo tiempo que mejora la eficiencia en el entrenamiento y la inferencia.

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.