自动识别可删的序列标记,减少大模型推理时的缓存占用。
Neural Attention Search
- 设计三种令牌类型,动态决定保留或丢弃
- 实验显示缓存大小显著降低,性能不变
- 适合需优化推理效率的大模型部署场景
我们提出神经注意力搜索(NAtS),一种自动评估序列中每个标记重要性并判断其是否可在后续步骤中被丢弃的框架。该方法能有效降低基于Transformer模型在推理时所需的键值(KV)缓存大小,从而减少推理开销。本文设计了一个包含三种令牌类型的搜索空间:(i) 全局令牌将被保留并供后续所有令牌查询;(ii) 局部令牌仅在下一个全局令牌出现前存活;(iii) 滑动窗口令牌对接下来固定数量的后续令牌产生影响。类似于单次神经架构搜索,令牌类型信息可通过可学习的注意力掩码与网络权重联合学习。在从头训练新Transformer及微调现有大型语言模型的实验中,NAtS均能高效降低模型所需的KV缓存大小,同时保持模型性能。
原文摘要 · Abstract (English)
We present Neural Attention Search (NAtS), a framework that automatically evaluates the importance of each token within a sequence and determines if the corresponding token can be dropped after several steps. This approach can efficiently reduce the KV cache sizes required by transformer-based models during inference and thus reduce inference costs. In this paper, we design a search space that contains three token types: (i) Global Tokens will be preserved and queried by all the following tokens. (ii) Local Tokens survive until the next global token appears. (iii) Sliding Window Tokens have an impact on the inference of a fixed size of the next following tokens. Similar to the One-Shot Neural Architecture Search approach, this token-type information can be learned jointly with the architecture weights via a learnable attention mask. Experiments on both training a new transformer from scratch and fine-tuning existing large language models show that NAtS can efficiently reduce the KV cache size required for the models while maintaining the models' performance.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。