arXiv:2502.13251v1 Announce Type: cross 摘要：我们提出了一种名为神经注意力搜索（NAtS）的框架，该框架能够自动评估序列中每个令牌的重要性，并确定在经过若干步之后，相应的令牌是否可以被丢弃。这种方法可以有效地减少基于Transformer的模型在推理过程中所需的KV缓存大小，从而降低推理成本。在本文中，我们设计了一个包含三种类型的令牌的空间：（i）全局令牌将被所有后续令牌保留并查询。（ii）局部令牌存活到下一个全局令牌出现为止。（iii）滑动窗口令牌对固定大小的后续令牌的推理产生影响。类似于One-Shot神经架构搜索方法，这些令牌类型的信息可以通过可学习的注意力掩码与架构权重联合学习。在从头训练一个新Transformer和微调现有的大型语言模型的实验中都表明，NAtS可以有效地减少模型所需的KV缓存大小，同时保持模型的性能。