Deploying Large Language Models (LLMs) in streaming applications such as multi-round dialogue, where long interactions are expected, is urgently needed but poses two major challenges. Firstly, during the decoding stage, caching previous tokens' Key and Value states (KV) consumes extensive memory. Secondly, popular LLMs cannot generalize to longer texts than the training sequence length. Window attention, where only the most recent KVs are cached, is a natural approach -- but we show that it fails when the text length surpasses the cache size. We observe an interesting phenomenon, namely attention sink, that keeping the KV of initial tokens will largely recover the performance of window attention. In this paper, we first demonstrate that the emergence of attention sink is due to the strong attention scores towards initial tokens as a ``sink'' even if they are not semantically important. Based on the above analysis, we introduce StreamingLLM, an efficient framework that enables LLMs trained with a finite length attention window to generalize to infinite sequence lengths without any fine-tuning. We show that StreamingLLM can enable Llama-2, MPT, Falcon, and Pythia to perform stable and efficient language modeling with up to 4 million tokens and more. In addition, we discover that adding a placeholder token as a dedicated attention sink during pre-training can further improve streaming deployment. In streaming settings, StreamingLLM outperforms the sliding window recomputation baseline by up to 22.2x speedup. Code and datasets are provided at https://github.com/mit-han-lab/streaming-llm.
翻译:将大型语言模型(LLMs)部署到多轮对话等需要长交互的流式应用中具有迫切需求,但面临两大挑战:首先,解码阶段缓存历史token的键值状态(KV)会消耗大量内存;其次,现有LLM无法泛化至超过训练序列长度的文本。窗口注意力机制(仅缓存最近的KVs)是自然之选——但我们的研究表明,当文本长度超过缓存容量时该方法会失效。我们观察到名为"注意力汇聚"的有趣现象:保留初始token的KV能大幅恢复窗口注意力的性能。本文首先证明注意力汇聚的产生源于模型对初始token的强注意力得分,即使这些token不承载语义重要性,它们仍作为"汇聚点"存在。基于上述分析,我们提出StreamingLLM——一个无需微调即可使有限注意力窗口训练的LLM泛化至无限序列长度的高效框架。实验表明,StreamingLLM能支撑Llama-2、MPT、Falcon及Pythia等模型对高达400万以上token进行稳定高效的语言建模。此外,我们发现预训练阶段添加占位token作为专用注意力汇聚点可进一步优化流式部署。在流式场景下,StreamingLLM相比滑动窗口重计算基线实现最高22.2倍加速。代码和数据集已开源至https://github.com/mit-han-lab/streaming-llm。