Zebra: Estensione della Finestra di Contesto con Attenzione Locale-Globale a Gruppi Stratificati

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

December 14, 2023

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

cs.AI

Abstract

Questo articolo introduce un approccio innovativo per potenziare le capacità dei Large Language Models (LLM) nell'elaborazione e comprensione di sequenze testuali estese, un aspetto cruciale per applicazioni che richiedono una profonda comprensione e sintesi di grandi volumi di informazioni. Riconoscendo le sfide intrinseche nell'estendere la finestra contestuale per i LLM, principalmente basati sull'architettura Transformer, proponiamo una nuova architettura di modello, denominata Zebra. Questa architettura gestisce in modo efficiente i problemi di complessità quadratica in termini di tempo e memoria associati all'attenzione completa nel Transformer, impiegando strati di attenzione locale-globale raggruppati. Il nostro modello, simile alle strisce alternate di una zebra, bilancia strati di attenzione locale e globale, riducendo significativamente i requisiti computazionali e il consumo di memoria. Sono stati condotti esperimenti completi, tra cui pretraining da zero, continuazione dell'adattamento a contesti lunghi e tuning su istruzioni lunghe, per valutare le prestazioni di Zebra. I risultati dimostrano che Zebra raggiunge prestazioni comparabili o superiori su benchmark sia per sequenze brevi che lunghe, migliorando al contempo l'efficienza durante l'addestramento e l'inferenza.

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.