Efficient Sparse Attention needs Adaptive Token Release

Anonymous

Abstract: In recent years, Large Language Models (LLMs) have demonstrated remarkable capabilities across a wide array of text-centric tasks. However, their `large' scale introduces significant computational and storage challenges, particularly in managing the key-value states of the transformer, which limits their wider applicability.Therefore, we propose to adaptively release resources from caching and rebuild the necessary key-value states. Particularly, we accomplish this by a lighting controller module to approximate an ideal top-$K$ sparse attention. This module retains the tokens with the highest top-$K$ attention weights and simultaneously rebuilds the discarded tokens, which may become essential for future decoding.Comprehensive experiments in natural language generation and natural language modeling task reveal that our method is not only competitive with full attention in terms of performance but also achieves a significant throughput improvement of up to $\textbf{221.8\%}$.The code for replication is available on the https://anonymous.4open.science/r/ADORE-5384.

Paper Type: long

Research Area: Generation

Contribution Types: Model analysis & interpretability, Approaches low compute settings-efficiency

Languages Studied: English