突破LLM长度限制：Attention Sinks如何实现无限文本生成而不牺牲性能

【免费下载链接】attention_sinksExtend existing LLMs way beyond the original training length with constant memory usage, without retraining项目地址: https://gitcode.com/gh_mirrors/at/attention_sinks

你是否遇到过大型语言模型（LLM）在处理长文本时突然"失忆"或生成混乱内容的情况？🤔 这通常是因为传统LLM在生成超过其训练长度的文本时，会出现内存线性增长和性能急剧下降的问题。而Attention Sinks技术正是解决这一痛点的革命性方案！它能让现有LLM突破原始训练长度限制，以恒定内存占用实现无限文本生成，且无需重新训练模型。

什么是Attention Sinks？

Attention Sinks是基于论文《Efficient Streaming Language Models with Attention Sinks》实现的开源项目，它通过修改滑动窗口注意力机制，让LLM能够无限生成流畅文本。核心原理是保留少量初始"注意力汇点"（attention sink）标记和最近的上下文窗口，同时丢弃中间标记，从而在保持恒定内存使用的同时维持生成质量。

三种方案对比

传统LLM处理长文本时主要面临两种困境：

• 普通transformers：内存使用随文本长度线性增长，超过训练长度后性能严重下降
• 窗口注意力：内存使用恒定但会丢失早期上下文，很快就会生成无意义内容

而Attention Sinks则结合了两者的优势：

• 恒定内存占用（通过4个注意力汇点标记+1020个最近标记的窗口机制）
• 长期保持生成流畅性，即使处理数百万标记也不会"失忆"

惊人的性能表现

困惑度对比实验

以下是不同模型在三种方案下的困惑度（perplexity）对比，困惑度越低表示生成质量越好：

Llama-2-7b-hf与Falcon-7B模型表现

• transformers：VRAM使用线性增长，超过训练长度后性能崩溃
• windowed：VRAM使用恒定，但上下文窗口滑动后性能立即下降
• attention_sinks：VRAM使用恒定，且困惑度始终保持稳定

MPT-7B与Pythia-6.9B模型表现实验结果表明，Attention Sinks在各种模型架构上都能保持稳定的低困惑度，而其他方案在长文本生成中很快就会失效。

无限生成流畅度测试

在Llama 2 7B模型上进行的10,000标记无限生成测试中：

• transformers方案日志：约1900标记后失去流畅性，开始生成乱码如🤖🧠👨‍���������������������
• 窗口注意力方案日志：约1000标记后失去流畅性，生成大量换行和无意义字符OOOMMO̶OANOOAMOO̶OMMO
• Attention Sinks方案日志：全程保持流畅，完成全部10k标记测试✅

快速开始：5分钟上手Attention Sinks

安装步骤

安装Attention Sinks非常简单，只需一行命令：

pip install attention_sinks

基本使用方法

使用Attention Sinks与标准transformers API几乎相同，唯一区别是从attention_sinks导入模型类：

from attention_sinks import AutoModelForCausalLM model = AutoModelForCausalLM.from_pretrained( "mosaicml/mpt-7b", device_map="auto", torch_dtype=torch.float16, # Attention Sinks特有参数 attention_sink_size=4, # 注意力汇点标记数量，默认4 attention_sink_window_size=252 # 滑动窗口大小，默认1020 )

支持的模型

Attention Sinks目前支持多种主流LLM架构：

• Llama系列（包括Llama 2）
• Mistral
• Falcon
• MPT
• GPTNeoX (Pythia)
• GPT-J
• Qwen
• StableLM_epoch
• BTLM
• Yi

高级应用：构建流式对话系统

Attention Sinks特别适合构建需要长时间运行的对话式AI系统。通过demo/streaming.py脚本，你可以实现多轮对话而不会丢失上下文或耗尽内存。

以下是构建流式对话系统的核心步骤：

1. 加载模型时配置适当的注意力汇点大小和窗口大小
2. 保存每轮对话生成的past_key_values作为历史上下文
3. 在后续对话中传入历史past_key_values以保持上下文连贯

这种方法已在MT-Bench基准测试中证明能显著提高多轮对话的流畅性，尤其对于Llama-2-7b-chat、Mistral-7B-Instruct等对话模型效果显著。

常见问题解答

Attention Sinks真的能处理无限长文本吗？

是的！它通过保留最近标记和注意力汇点，在恒定内存占用下实现无限文本生成。实验表明模型困惑度在400万标记后仍然保持稳定。

这是否扩展了LLM的上下文窗口？

不是。上下文窗口仍受原始预训练限制。Attention Sinks的优势在于能在不刷新缓存的情况下持续生成，而不是增加模型能理解的上下文长度。

最适合的使用场景是什么？

Attention Sinks最适合流式应用，如：

• 多轮对话系统
• 实时日志分析
• 持续运行的AI助手
• 需要长时间生成的内容创作

与其他上下文扩展方法有冲突吗？

没有！Attention Sinks可以与LongChat等上下文扩展方法结合使用，进一步提升长文本处理能力。

性能基准测试

项目提供了完整的基准测试工具，你可以通过benchmark/perplexity.py脚本测试不同模型在各种配置下的表现：

python benchmark/perplexity.py --experiment attention_sinks --model_name_or_path meta-llama/Llama-2-7b-hf

测试结果会生成CSV文件，可通过benchmark/plot_perplexity.py可视化：

python benchmark/plot_perplexity.py --features perplexity vram --output_dir benchmark/outputs_llama_2_7b

总结

Attention Sinks为LLM的长文本生成提供了突破性解决方案，它以最小的改动实现了：

• 🔄 无限文本生成能力
• 📈 恒定内存占用
• 🚀 不牺牲生成质量和性能
• 🔌 与现有transformers API无缝集成

无论你是构建对话机器人、内容生成工具还是需要处理长文本的应用，Attention Sinks都能帮助你突破LLM的长度限制，开启无限可能！

要开始使用，只需执行：

git clone https://gitcode.com/gh_mirrors/at/attention_sinks cd attention_sinks pip install .

查看项目demo目录获取更多使用示例，或参考CHANGELOG.md了解最新功能更新。

【免费下载链接】attention_sinksExtend existing LLMs way beyond the original training length with constant memory usage, without retraining项目地址: https://gitcode.com/gh_mirrors/at/attention_sinks