Atención dispersa jerárquica hecha correctamente: hacia el modelado de contexto infinito

Resumen

El escalamiento de los modelos de lenguaje de gran escala (LLMs, por sus siglas en inglés) a contextos largos está limitado por el costo computacional cuadrático y la pobre extrapolación de longitud de la atención densa. La atención dispersa por fragmentos ofrece una alternativa prometedora, pero todos los métodos existentes quedan por detrás de la atención completa debido a su selección inexacta de fragmentos. Proponemos la Atención Dispersa Jerárquica de Hitos (HiLS, por sus siglas en inglés), un mecanismo de atención dispersa por fragmentos que aprende la selección de fragmentos de extremo a extremo bajo la pérdida de modelado de lenguaje (LM). HiLS factoriza la atención jerárquicamente: cada consulta realiza atención de forma independiente con cada fragmento recuperado para extraer información específica de ese fragmento, y las salidas resultantes se fusionan de acuerdo con las puntuaciones de recuperación de los fragmentos. Al incorporar las puntuaciones de recuperación en el cómputo de atención hacia adelante, HiLS las optimiza directamente con la pérdida de LM, lo que permite el aprendizaje de recuperación de extremo a extremo y el entrenamiento disperso nativo. Los resultados experimentales muestran que la Atención HiLS logra un rendimiento comparable, y en algunos casos superior, al de la atención completa en longitudes de contexto dentro del dominio. Mientras tanto, la Atención HiLS extrapola más de 64 veces la longitud de contexto de entrenamiento con un 90% de precisión de recuperación, superando con creces a la atención completa. Además, los modelos de atención completa existentes pueden convertirse a Atención HiLS mediante un preentrenamiento continuo ligero, preservando el rendimiento dentro del dominio mientras adquieren extrapolación de contexto ultra largo. Junto con su acceso y cómputo disperso de KV, la Atención HiLS rompe el equilibrio habitual entre eficiencia y rendimiento, permitiendo LLMs de contexto largo que son tanto más eficientes como más efectivos en tareas generales de contexto largo que sus contrapartes de atención 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.