arXiv:2502.13251v2 类别: 交叉学科 摘要: 我们提出了一种名为 Neural Attention Search (NAtS) 的框架，该框架能够自动评估序列中每个令牌的重要性，并确定在几轮后是否可以删除相应的令牌。此方法可以有效地减少变压器模型在推理过程中的 KV 缓存大小需求，从而降低推理成本。在本文中，我们设计了一个包含三种令牌类型的空间：(i) 全局令牌将被后续所有令牌保留并查询。(ii) 局部令牌存活直到下一个全局令牌出现。(iii) 滑动窗口令牌对固定大小的后续令牌的推理产生影响。类似于 One-Shot 神经架构搜索方法，这种令牌类型信息可以通过可学习的注意力掩码与架构权重联合学习。从头训练一个新的变压器和 fine-tune 现有的大型语言模型的实验表明，NAtS 可以有效地减少模型所需的 KV 缓存大小，同时保持模型的性能。