Zebra : Extension de la fenêtre contextuelle par une attention locale-globale groupée par couches

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

December 14, 2023

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

cs.AI

Résumé

Cet article présente une nouvelle approche visant à améliorer les capacités des modèles de langage de grande taille (LLMs) dans le traitement et la compréhension de séquences textuelles étendues, un aspect crucial pour les applications nécessitant une compréhension approfondie et une synthèse de grands volumes d'informations. Conscients des défis inhérents à l'extension de la fenêtre contextuelle pour les LLMs, principalement basés sur l'architecture Transformer, nous proposons une nouvelle architecture de modèle, appelée Zebra. Cette architecture gère efficacement les problèmes de complexité quadratique en temps et en mémoire associés à l'attention complète dans le Transformer en utilisant des couches d'attention locale-globale groupées. Notre modèle, semblable aux rayures alternées d'un zèbre, équilibre les couches d'attention locale et globale, réduisant ainsi significativement les besoins en calcul et en mémoire. Des expériences approfondies, incluant un pré-entraînement à partir de zéro, une continuation de l'entraînement d'adaptation à des contextes longs et un réglage d'instructions longues, sont menées pour évaluer les performances de Zebra. Les résultats montrent que Zebra atteint des performances comparables ou supérieures sur des benchmarks de séquences courtes et longues, tout en améliorant l'efficacité de l'entraînement et de l'inférence.

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.