Locret：在消费级设备上通过可训练的保留头实现长上下文LLM推理中的高效驱逐
计算与语言
2025-01-31 v2
摘要
扩展大型语言模型（LLM）的输入上下文长度会显著增加计算成本和内存开销，以维护注意力键值（KV）缓存。现有的KV缓存压缩方法存在压缩效率低下和内存减少效果有限的问题，使得LLM难以在消费级设备上进行长上下文推理，尤其是在推理长上下文流式输入时。这些障碍阻碍了消费级设备支持更复杂的应用，对LLM的普及构成了挑战。为解决这一问题，我们提出了Locret，这是第一个创建与分块预填充兼容的驱逐策略的框架。通过可学习的保留头评估KV缓存单元的因果重要性，Locret能够精确驱逐缓存单元，从而实现高效的长上下文推理。在我们广泛的实证研究中，Locret在内存效率和生成质量方面优于近期流行且有竞争力的方法——Locret实现了高达20倍的KV缓存压缩比，且性能损失小于10%。此外，Locret在单个NVIDIA 4090 GPU上实现了128K+的长上下文推理，且不损害生成质量，仅需不到1 GPU小时的额外训练时间。
关键词
引用
@article{arxiv.2410.01805,
title = {Locret: Enhancing Eviction in Long-Context LLM Inference with Trained Retaining Heads on Consumer-Grade Devices},
author = {Yuxiang Huang and Binhang Yuan and Xu Han and Chaojun Xiao and Zhiyuan Liu},
journal= {arXiv preprint arXiv:2410.01805},
year = {2025}
}
备注
Preprints