带有注意力汇的高效流式语言模型

[论文] [幻灯片][视频]

https://github.com/mit-han-lab/streaming-llm/assets/40906949/2bd1cda4-a0bd-47d1-a023-fbf7779b8358

简而言之

我们能够在不牺牲效率和性能的情况下，将大语言模型部署用于无限长度的输入。

新闻

• [2024/02] StreamingLLM 被 MIT 新闻以亮点形式报道：MIT News: 一种让 AI 聊天机器人全天候对话而不崩溃的新方法!
• [2024/01] HPC-AI Tech 的 SwiftInfer 已集成 StreamingLLM，支持 LLM 推理处理无限长度的输入。
• [2024/01] NVIDIA 的 TensorRT-LLM 已集成 StreamingLLM！
• [2023/12] CMU、UW 和 OctoAI 将 StreamingLLM 集成到他们的系统中，使得在 iPhone 上也能实现无尽且高效的 LLM 生成，详情请见 X。
• [2023/12] HuggingFace Transformers 已合并 StreamingLLM 的 PR。
• [2023/10] StreamingLLM 已被 Intel Extension for Transformers 集成。
• [2023/10] 第三方实现 Attention Sinks（由 tomaarsen 提供）使更多 HuggingFace 的 LLM 支持 StreamingLLM。

摘要

在多轮对话等需要长时间交互的流式应用中部署大型语言模型（LLMs）迫在眉睫，但同时也面临两大挑战。首先，在解码阶段，缓存先前 token 的 Key 和 Value 状态（KV）会占用大量内存。其次，主流 LLM 很难泛化到超过其训练序列长度的文本。窗口注意力机制只缓存最近的 KV，看似自然的选择——但我们发现，当文本长度超过缓存大小时，这种方法就会失效。我们观察到一种有趣的现象，即“注意力汇”，保留初始 token 的 KV 可以显著恢复窗口注意力的效果。本文首先证明，注意力汇的出现是因为初始 token 会形成一个“汇”，即使它们在语义上并不重要，也会产生很强的注意力得分。基于上述分析，我们提出了 StreamingLLM，这是一种高效的框架，无需任何微调即可使使用有限长度注意力窗口训练的 LLM 泛化到无限长度的序列。我们表明，StreamingLLM 能够让 Llama-2、MPT、Falcon 和 Pythia 等模型稳定高效地进行语言建模，支持高达 400 万个 token 甚至更多的输入。此外，我们还发现，在预训练阶段添加一个占位符 token 作为专门的注意力汇，可以进一步提升流式部署的效果。在流式场景下，StreamingLLM 相比滑动窗口重新计算的基线方法，速度最高可提升 22.2 倍。

使用方法

环境设置

conda create -yn streaming python=3.8
conda activate streaming

pip install torch torchvision torchaudio
pip install transformers==4.33.0 accelerate datasets evaluate wandb scikit-learn scipy sentencepiece

python setup.py develop

运行流式 Llama 聊天机器人

CUDA_VISIBLE_DEVICES=0 python examples/run_streaming_llama.py  --enable_streaming

常见问题解答

1. “处理无限长度输入”对 LLM 有何意义？

   处理无限长度文本对 LLM 来说是一个挑战。尤其是存储所有之前的 Key 和 Value (KV) 状态会消耗大量内存，而且模型可能难以生成超出其训练序列长度的文本。StreamingLLM 通过仅保留最近的 token 和注意力汇，丢弃中间的 token 来解决这一问题。这样，模型可以在不重置缓存的情况下，基于最近的 token 生成连贯的文本——这是以往方法所不具备的能力。

2. LLM 的上下文窗口是否被扩展了？

   没有。上下文窗口保持不变。StreamingLLM 只保留最近的 token 和注意力汇，丢弃中间的 token。这意味着模型只能处理最新的 token。上下文窗口仍然受限于其最初的预训练设置。例如，如果 Llama-2 是用 4096 个 token 的上下文窗口进行预训练的，那么在 Llama-2 上使用 StreamingLLM 时，最大缓存大小仍然是 4096。

3. 我可以将一本书这样的长文本输入到 StreamingLLM 中进行摘要吗？

   虽然您可以输入一段较长的文本，但模型只会识别最近的 token。因此，如果输入的是一本书，StreamingLLM 可能只会总结最后几段，这可能并不太有意义。正如前面强调的那样，我们既没有扩展 LLM 的上下文窗口，也没有增强其长期记忆能力。StreamingLLM 的优势在于，它能够基于最近的 token 生成流畅的文本，而无需刷新缓存。

4. StreamingLLM 的理想应用场景是什么？

   StreamingLLM 针对流式应用进行了优化，例如多轮对话。它非常适合那些需要模型持续运行，同时又不需要大量内存或依赖过往数据的场景。例如，基于 LLM 的日常助手。StreamingLLM 可以让模型持续工作，根据最近的对话内容做出响应，而无需刷新缓存。相比之下，早期的方法要么会在对话长度超过训练长度时重置缓存（从而丢失最近的上下文），要么需要从最近的文本历史中重新计算 KV 状态，这会非常耗时。

5. StreamingLLM 与近期关于上下文扩展的研究有何关系？

   StreamingLLM 与近期的上下文扩展方法是正交的，并且可以与这些方法集成。在 StreamingLLM 的语境下，“上下文扩展”指的是使用更大的缓存来存储更多最近的 token。有关实际演示，请参阅我们论文中的图 9，其中我们展示了如何将 StreamingLLM 应用于 LongChat-7B-v1.5-32K 和 Llama-2-7B-32K-Instruct 等模型。

待办事项

我们将按以下顺序发布代码和数据，请继续关注！

• 发布 StreamingLLM 的核心代码，包括 Llama-2、MPT、Falcon 和 Pythia。
• 发布困惑度评估代码。
• 发布流式 Llama 聊天机器人演示。
• 发布 StreamEval 数据集及评估代码。

引用

如果您认为 StreamingLLM 对您的项目或研究有用或相关，请引用我们的论文：

@article{xiao2023streamingllm,
        title={Efficient Streaming Language Models with Attention Sinks},
        author={Xiao, Guangxuan and Tian, Yuandong and Chen, Beidi and Han, Song and Lewis, Mike},
        journal={arXiv},
        year={2023}
        }

StreamingLLM 快速上手指南

StreamingLLM 是一个高效框架，允许大语言模型（LLM）在无需微调的情况下处理无限长度的输入流。它通过保留初始 token（Attention Sinks）和最近 token 的 KV 状态，解决了长对话中显存占用过高及模型无法泛化到超长序列的问题，特别适用于多轮对话等流式应用场景。

环境准备

• 操作系统: Linux (推荐) 或 macOS
• Python 版本: 3.8+
• 硬件要求: NVIDIA GPU (需支持 CUDA)
• 前置依赖: Conda (推荐用于环境管理)

安装步骤

1. 创建并激活虚拟环境

   conda create -yn streaming python=3.8
   conda activate streaming
   

2. 安装深度学习基础库 (注：如需国内加速，可添加 -i https://pypi.tuna.tsinghua.edu.cn/simple 参数)

   pip install torch torchvision torchaudio
   

3. 安装项目依赖

   pip install transformers==4.33.0 accelerate datasets evaluate wandb scikit-learn scipy sentencepiece
   

4. 安装 StreamingLLM 克隆仓库后进入目录，执行开发模式安装：

   python setup.py develop
   

基本使用

以下命令将启动一个支持无限长度输入的 Llama 聊天机器人示例。请确保至少有一块可用的 GPU。

CUDA_VISIBLE_DEVICES=0 python examples/run_streaming_llama.py --enable_streaming

运行后，模型将仅缓存最近的 token 和初始的 Attention Sinks，从而在长对话中保持稳定的生成速度和较低的显存占用。