TidalDecode: Snelle en nauwkeurige LLM-decodering met positie-persistente schaarse aandacht

TidalDecode: Fast and Accurate LLM Decoding with Position Persistent Sparse Attention

October 7, 2024

Auteurs: Lijie Yang, Zhihao Zhang, Zhuofu Chen, Zikun Li, Zhihao Jia

cs.AI

Samenvatting

Grote taalmodellen (LLM's) hebben significante vooruitgang geboekt op diverse taken binnen natuurlijke taalverwerking (NLP), waarbij modellen met lange context steeds belangrijker worden voor het verwerken van uitgebreide invoer. Echter, de groeiende omvang van de sleutel-waarde (KV) cache die nodig is voor Transformer-architecturen intensiveert de geheugenbeperkingen, vooral tijdens de decoderingsfase, wat een aanzienlijke bottleneck creëert. Bestaande schaarse aandachtsmechanismen die zijn ontworpen om deze bottleneck aan te pakken, hebben twee beperkingen: (1) ze slagen er vaak niet in om betrouwbaar de meest relevante tokens voor aandacht te identificeren, en (2) ze zien de ruimtelijke samenhang van tokenselectie over opeenvolgende Transformer-lagen over het hoofd, wat kan leiden tot prestatievermindering en aanzienlijke overhead in tokenselectie. Dit artikel introduceert TidalDecode, een eenvoudig maar effectief algoritme en systeem voor snelle en nauwkeurige decodering van LLM's door middel van positie-persistente schaarse aandacht. TidalDecode maakt gebruik van de ruimtelijke samenhang van tokens geselecteerd door bestaande schaarse aandachtsmethoden en introduceert enkele tokenselectielagen die volledige aandacht uitvoeren om de tokens met de hoogste aandachtscores te identificeren, terwijl alle andere lagen schaarse aandacht uitvoeren met de vooraf geselecteerde tokens. Deze opzet stelt TidalDecode in staat om aanzienlijk de overhead van tokenselectie voor schaarse aandacht te verminderen zonder afbreuk te doen aan de kwaliteit van de gegenereerde resultaten. Evaluatie op een diverse reeks LLM's en taken toont aan dat TidalDecode de generatieve prestaties van methoden met volledige aandacht nauw benadert, terwijl de decoderingslatentie van LLM's met maximaal 2,1x wordt verminderd.

English

Large language models (LLMs) have driven significant advancements across diverse NLP tasks, with long-context models gaining prominence for handling extended inputs. However, the expanding key-value (KV) cache size required by Transformer architectures intensifies the memory constraints, particularly during the decoding phase, creating a significant bottleneck. Existing sparse attention mechanisms designed to address this bottleneck have two limitations: (1) they often fail to reliably identify the most relevant tokens for attention, and (2) they overlook the spatial coherence of token selection across consecutive Transformer layers, which can lead to performance degradation and substantial overhead in token selection. This paper introduces TidalDecode, a simple yet effective algorithm and system for fast and accurate LLM decoding through position persistent sparse attention. TidalDecode leverages the spatial coherence of tokens selected by existing sparse attention methods and introduces a few token selection layers that perform full attention to identify the tokens with the highest attention scores, while all other layers perform sparse attention with the pre-selected tokens. This design enables TidalDecode to substantially reduce the overhead of token selection for sparse attention without sacrificing the quality of the generated results. Evaluation on a diverse set of LLMs and tasks shows that TidalDecode closely matches the generative performance of full attention methods while reducing the LLM decoding latency by up to 2.1x.