神经注意力搜索

我们提出了神经注意力搜索（NAtS），这是一个能够自动评估序列中每个token重要性并决定相应token是否可以在若干步骤后被丢弃的框架。这种方法可以有效减少基于transformer的模型在推理期间所需的KV缓存大小，从而降低推理成本。在本文中，我们设计了一个包含三种token类型的搜索空间： Global Tokens将被保留并被后续所有token查询； Local Tokens存活直到下一个全局token出现； Sliding Window Tokens对固定数量的后续token的推理产生影响。类似于One-Shot神经架构搜索方法，这种token类型信息可以通过可学习的注意力掩码与架构权重联合学习。在从头训练新transformer和微调现有大语言模型的实验中均表明，NAtS能够有效减少模型所需的KV缓存大小，同时保持模型的性能。