Zebra: Erweiterung des Kontextfensters durch schichtenweise gruppierte lokale-globale Aufmerksamkeit

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

December 14, 2023

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

cs.AI

Zusammenfassung

Dieses Papier stellt einen neuartigen Ansatz vor, um die Fähigkeiten von Large Language Models (LLMs) bei der Verarbeitung und dem Verständnis umfangreicher Textsequenzen zu verbessern, einem entscheidenden Aspekt in Anwendungen, die ein tiefes Verständnis und die Synthese großer Informationsmengen erfordern. In Anerkennung der inhärenten Herausforderungen bei der Erweiterung des Kontextfensters für LLMs, die hauptsächlich auf der Transformer-Architektur basieren, schlagen wir eine neue Modellarchitektur vor, die als Zebra bezeichnet wird. Diese Architektur bewältigt effizient die quadratischen Zeit- und Speicherkomplexitätsprobleme, die mit der vollständigen Aufmerksamkeit im Transformer verbunden sind, indem sie gruppierte lokale-globale Aufmerksamkeitsschichten einsetzt. Unser Modell, ähnlich den abwechselnden Streifen eines Zebras, balanciert lokale und globale Aufmerksamkeitsschichten und reduziert dadurch erheblich den Rechenaufwand und den Speicherverbrauch. Umfassende Experimente, einschließlich Pretraining von Grund auf, Fortsetzung des Trainings zur Anpassung an lange Kontexte und Feinabstimmung langer Instruktionen, werden durchgeführt, um die Leistung von Zebra zu bewerten. Die Ergebnisse zeigen, dass Zebra vergleichbare oder überlegene Leistung sowohl bei kurzen als auch bei langen Sequenzbenchmarks erzielt und gleichzeitig die Effizienz des Trainings und der Inferenz verbessert.

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.