A Atenção Completa Contra-Ataca: Transferindo Atenção Completa para Atenção Esparsa em Cem Etapas de Treinamento

Resumo

A inferência de contexto longo em modelos de linguagem de grande escala é limitada pelo custo quadrático da atenção total. Alternativas eficientes existentes frequentemente dependem de treinamento esparso nativo ou de evicção heurística de tokens, criando uma troca indesejável entre eficiência, custo de treinamento e precisão. Neste trabalho, mostramos que LLMs de atenção total já são intrinsecamente esparsos e podem ser transformados em modelos altamente esparsos com apenas adaptação mínima. Nossa abordagem é baseada em três observações: (1) apenas um pequeno subconjunto de cabeças de atenção realmente requer processamento completo de contexto longo; (2) a recuperação de longo alcance é governada principalmente por um subespaço de baixa dimensão, permitindo que tokens relevantes sejam recuperados eficientemente com um indexador de 16 dimensões; e (3) o orçamento de tokens úteis é fortemente dependente da consulta, tornando a seleção dinâmica top-p mais adequada do que a esparsificação fixa top-k. Com base nesses insights, propomos o RTPurbo, que retém o cache KV completo apenas para cabeças de recuperação e introduz um indexador de tokens leve para atenção esparsa. Ao explorar a esparsidade intrínseca do modelo, o RTPurbo alcança esparsificação com apenas algumas centenas de etapas de treinamento. Experimentos em benchmarks de contexto longo e tarefas de raciocínio mostram que o RTPurbo preserva precisão quase sem perdas, ao mesmo tempo que proporciona ganhos substanciais de eficiência, incluindo um aumento de velocidade de preenchimento de até 9,36× em contexto de 1M e cerca de 2,01× na decodificação. Esses resultados sugerem que uma inferência esparsa robusta pode ser obtida a partir de treinamento padrão de atenção total, sem a necessidade de pré-treinamento esparso nativo caro.

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.