La atención completa contraataca: transfiriendo la atención completa a dispersa en cien pasos de entrenamiento

Resumen

La inferencia de contexto largo en modelos de lenguaje de gran escala está limitada por el costo cuadrático de la atención completa. Las alternativas eficientes existentes a menudo se basan en entrenamiento disperso nativo o en la evicción heurística de tokens, generando un equilibrio no deseado entre eficiencia, costo de entrenamiento y precisión. En este trabajo, demostramos que los modelos de lenguaje con atención completa ya son intrínsecamente dispersos y pueden transformarse en modelos altamente dispersos con una adaptación mínima. Nuestro enfoque se basa en tres observaciones: (1) solo un pequeño subconjunto de cabezas de atención requiere realmente procesamiento completo de contexto largo; (2) la recuperación de largo alcance está gobernada principalmente por un subespacio de baja dimensión, lo que permite recuperar tokens relevantes de manera eficiente con un indexador de 16 dimensiones; y (3) el presupuesto de tokens útiles depende fuertemente de la consulta, lo que hace que la selección dinámica de top-p sea más adecuada que la esparsificación fija de top-k. Con base en estas ideas, proponemos RTPurbo, que retiene la caché KV completa solo para las cabezas de recuperación e introduce un indexador de tokens ligero para la atención dispersa. Al explotar la dispersión intrínseca del modelo, RTPurbo logra la esparsificación con solo unos cientos de pasos de entrenamiento. Experimentos en puntos de referencia de contexto largo y tareas de razonamiento muestran que RTPurbo preserva una precisión casi sin pérdidas mientras ofrece ganancias sustanciales de eficiencia, incluyendo una aceleración de prefill de hasta 9.36 veces en contexto de 1M y una aceleración de decodificación de aproximadamente 2.01 veces. Estos resultados sugieren que se puede obtener una inferencia dispersa robusta a partir del entrenamiento estándar de atención completa sin un costoso preentrenamiento disperso nativo.

English

Long-context inference in large language models is bottlenecked by the quadratic cost of full attention. Existing efficient alternatives often rely either on native sparse training or on heuristic token eviction, creating an undesirable trade-off among efficiency, training cost, and accuracy. In this work, we show that full-attention LLMs are already intrinsically sparse and can be transformed into highly sparse models with only minimal adaptation. Our approach is built on three observations: (1) only a small subset of attention heads truly requires full long-context processing; (2) long-range retrieval is governed primarily by a low-dimensional subspace, allowing relevant tokens to be retrieved efficiently with a 16-dimensional indexer; and (3) the useful token budget is strongly query-dependent, making dynamic top-p selection more suitable than fixed top-k sparsification. Based on these insights, we propose RTPurbo, which retains the full KV cache only for retrieval heads and introduces a lightweight token indexer for sparse attention. By exploiting the model's intrinsic sparsity, RTPurbo achieves sparsification with only a few hundred training steps. Experiments on long-context benchmarks and reasoning tasks show that RTPurbo preserves near-lossless accuracy while delivering substantial efficiency gains, including up to a 9.36times prefill speedup at 1M context and about a 2.01times decode speedup. These results suggest that strong sparse inference can be obtained from standard full-attention training without expensive native sparse pretraining.