Atenção Esparsa Hierárquica Corretamente: Rumo à Modelagem de Contexto Infinito

Resumo

O escalonamento de grandes modelos de linguagem (LLMs) modernos para contextos longos é limitado pelo custo computacional quadrático e pela baixa extrapolação de comprimento da atenção densa. A atenção esparsa por blocos oferece uma alternativa promissora, mas todos os métodos existentes ficam aquém da atenção completa devido à seleção imprecisa de blocos. Propomos a Atenção Esparsa Hierárquica por Marcos (HiLS), um mecanismo de atenção esparsa por blocos que aprende a seleção de blocos de ponta a ponta sob a perda de modelagem de linguagem (LM). A HiLS fatoriza a atenção hierarquicamente: cada consulta realiza atenção independentemente com cada bloco recuperado para extrair informações específicas do bloco, e as saídas resultantes são fundidas de acordo com as pontuações de recuperação dos blocos. Ao incorporar as pontuações de recuperação no cálculo direto da atenção, a HiLS as otimiza diretamente com a perda de LM, permitindo aprendizado de recuperação de ponta a ponta e treinamento esparso nativo. Resultados experimentais mostram que a HiLS-Atenção alcança desempenho comparável e, em alguns casos, superior ao da atenção completa em comprimentos de contexto dentro do domínio. Além disso, a HiLS-Atenção extrapola mais de 64 vezes o comprimento do contexto de treinamento com 90% de precisão de recuperação, muito além da atenção completa. Ademais, modelos existentes de atenção completa podem ser convertidos para HiLS-Atenção com um pré-treinamento contínuo leve, preservando o desempenho dentro do domínio enquanto adquirem extrapolação de contexto ultra longo. Juntamente com seu acesso e computação esparsos de KV, a HiLS-Atenção quebra a usual troca entre eficiência e desempenho, possibilitando LLMs de contexto longo que são tanto mais eficientes quanto mais eficazes em tarefas gerais de contexto longo do que suas contrapartes de atenção completa.

English

Scaling modern large language models (LLMs) to long contexts is limited by the quadratic computation cost, and poor length extrapolation of dense attention. Chunk-wise sparse attention offers a promising alternative, but all existing methods fall short of full attention because of their inaccurate chunk selection. We propose Hierarchical Landmark Sparse (HiLS) Attention, a chunk-wise sparse attention mechanism that learns chunk selection end-to-end under the language-modeling (LM) loss. HiLS factorizes attention hierarchically: each query performs attention independently with each retrieved chunk to extract chunk-specific information, and the resulting outputs are fused according to chunk retrieval scores. By incorporating retrieval scores into the forward attention computation, HiLS optimizes them directly with the LM loss, enabling end-to-end retrieval learning and native sparse training. Experimental results show that HiLS-Attention achieves performance comparable to, and in some cases better than, full attention at in-domain context lengths. Meanwhile, HiLS-Attention extrapolates more than 64times the training context length with 90% retrieval accuracy, far beyond full attention. Moreover, existing full-attention models can be converted to HiLS-Attention with lightweight continued pretraining, preserving in-domain performance while acquiring ultra-long-context extrapolation. Together with its sparse KV access and computation, HiLS-Attention breaks the usual efficiency-performance trade-off, enabling long-context LLMs that are both more efficient and more effective on general long-context tasks than their full-attention counterparts.