张量 Transformer 解码中 KV-Cache 压缩的时间与内存权衡
机器学习
2025-03-28 v2 人工智能
计算复杂性
计算与语言
摘要
张量版 Transformer 中的键值（KV）缓存在推理过程中构成了显著瓶颈。已有工作分析了标准注意力机制中的空间复杂度基本障碍 [Haris and Onak, 2025]，而我们的工作将空间复杂度障碍的结果推广至张量注意力版本。我们的理论贡献依赖于从通信复杂度的归约，推导出当 时张量结构注意力机制的内存下界。此外，我们引入了两种类型的张量注意力缓存，并给出了两种场景下的时间与内存权衡。总体而言，我们的工作为理解张量注意力解码中 KV-Cache 压缩的时间与内存权衡提供了理论基础，并为开发更具内存效率的张量注意力 Transformer 架构提供了更多视角。
引用
@article{arxiv.2503.11108,
title = {Time and Memory Trade-off of KV-Cache Compression in Tensor Transformer Decoding},
author = {Yifang Chen and Xiaoyu Li and Yingyu Liang and Zhenmei Shi and Zhao Song and Yu Tian},
journal= {arXiv preprint arXiv:2503.11108},
year = {2025}
}