Atenção Delta: Inferência Rápida e Precisa de Atenção Esparsa por Correção Delta

Delta Attention: Fast and Accurate Sparse Attention Inference by Delta Correction

May 16, 2025

Autores: Jeffrey Willette, Heejun Lee, Sung Ju Hwang

cs.AI

Resumo

O mecanismo de atenção de um transformer possui complexidade quadrática, resultando em altos custos de inferência e latência para sequências longas. No entanto, as matrizes de atenção são majoritariamente esparsas, o que implica que muitas entradas podem ser omitidas do cálculo para uma inferência eficiente. Métodos de inferência de atenção esparsa visam reduzir esse custo computacional; entretanto, eles também vêm acompanhados de uma degradação de desempenho problemática. Descobrimos que uma razão para essa degradação é que o cálculo esparso induz uma mudança distribucional nas saídas de atenção. Essa mudança distribucional faz com que as consultas no tempo de decodificação falhem em se alinhar adequadamente com as chaves apropriadas da etapa de preenchimento, levando a uma queda no desempenho. Propomos um procedimento simples, novo e eficaz para corrigir essa mudança distribucional, aproximando a distribuição das saídas de atenção esparsa daquela da atenção quadrática. Nosso método pode ser aplicado sobre qualquer método de atenção esparsa e resulta em um aumento médio de desempenho de 36 pontos percentuais, recuperando 88% da precisão da atenção quadrática no benchmark RULER de 131K quando aplicado sobre a atenção de janela deslizante com tokens de sumidouro, enquanto adiciona apenas uma pequena sobrecarga. Nosso método pode manter aproximadamente 98,5% de esparsidade em relação à atenção quadrática completa, tornando nosso modelo 32 vezes mais rápido que o Flash Attention 2 ao processar preenchimentos de 1 milhão de tokens.

English

The attention mechanism of a transformer has a quadratic complexity, leading to high inference costs and latency for long sequences. However, attention matrices are mostly sparse, which implies that many entries may be omitted from computation for efficient inference. Sparse attention inference methods aim to reduce this computational burden; however, they also come with a troublesome performance degradation. We discover that one reason for this degradation is that the sparse calculation induces a distributional shift in the attention outputs. The distributional shift causes decoding-time queries to fail to align well with the appropriate keys from the prefill stage, leading to a drop in performance. We propose a simple, novel, and effective procedure for correcting this distributional shift, bringing the distribution of sparse attention outputs closer to that of quadratic attention. Our method can be applied on top of any sparse attention method, and results in an average 36%pt performance increase, recovering 88% of quadratic attention accuracy on the 131K RULER benchmark when applied on top of sliding window attention with sink tokens while only adding a small overhead. Our method can maintain approximately 98.5% sparsity over full quadratic attention, making our model 32 times faster than Flash Attention 2 when processing 1M token prefills.