Zebra: Extensão da Janela de Contexto com Atenção Local-Global Agrupada por Camadas

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

Resumo

Este artigo introduz uma abordagem inovadora para aprimorar as capacidades dos Modelos de Linguagem de Grande Escala (LLMs) no processamento e compreensão de sequências extensas de texto, um aspecto crítico em aplicações que exigem profunda compreensão e síntese de grandes volumes de informação. Reconhecendo os desafios inerentes à extensão da janela de contexto para LLMs, principalmente construídos sobre a arquitetura Transformer, propomos uma nova arquitetura de modelo, denominada Zebra. Essa arquitetura gerencia de forma eficiente os problemas de complexidade quadrática de tempo e memória associados à atenção completa no Transformer, empregando camadas de atenção local-global agrupadas. Nosso modelo, semelhante às listras alternadas de uma zebra, equilibra camadas de atenção local e global, reduzindo significativamente os requisitos computacionais e o consumo de memória. Experimentos abrangentes, incluindo pré-treinamento do zero, continuação do treinamento de adaptação de contexto longo e ajuste fino de instruções longas, são conduzidos para avaliar o desempenho do Zebra. Os resultados mostram que o Zebra alcança desempenho comparável ou superior em benchmarks de sequências curtas e longas, ao mesmo tempo em que melhora a eficiência de treinamento e inferência.

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.