Дельта-внимание: Быстрый и точный вывод разреженного внимания с помощью коррекции дельты

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

May 16, 2025

Авторы: Jeffrey Willette, Heejun Lee, Sung Ju Hwang

cs.AI

Аннотация

Механизм внимания трансформатора обладает квадратичной сложностью, что приводит к высоким затратам на вывод и задержкам при обработке длинных последовательностей. Однако матрицы внимания в основном являются разреженными, что подразумевает возможность исключения многих элементов из вычислений для эффективного вывода. Методы разреженного внимания направлены на снижение этой вычислительной нагрузки, но они также сопровождаются проблематичным снижением производительности. Мы обнаружили, что одной из причин этого снижения является то, что разреженные вычисления вызывают сдвиг распределения в выходных данных внимания. Этот сдвиг распределения приводит к тому, что запросы на этапе декодирования не могут эффективно сопоставляться с соответствующими ключами из этапа предварительного заполнения, что вызывает падение производительности. Мы предлагаем простую, новую и эффективную процедуру для коррекции этого сдвига распределения, приближая распределение выходных данных разреженного внимания к распределению квадратичного внимания. Наш метод может быть применен поверх любого метода разреженного внимания и приводит к увеличению производительности в среднем на 36 процентных пунктов, восстанавливая 88% точности квадратичного внимания на бенчмарке 131K RULER при применении поверх скользящего окна внимания с токенами-приемниками, добавляя при этом лишь небольшой накладной расход. Наш метод позволяет сохранить примерно 98,5% разреженности по сравнению с полным квадратичным вниманием, что делает нашу модель в 32 раза быстрее, чем Flash Attention 2 при обработке предварительных заполнений на 1 миллион токенов.

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.