用于LLMs的基于草稿的近似推理

由于Transformer模型计算复杂度为平方级，内存复杂度为线性级，优化长上下文大型语言模型（LLMs）的推理变得日益重要。现有的近似方法，如键值（KV）缓存丢弃、稀疏注意力以及提示压缩，通常依赖于对token或KV对重要性的粗略预测。我们提出了一种新颖的近似LLM推理框架，该框架利用小型草稿模型更准确地预测token和KV对的重要性。具体而言，我们引入了我们提出的框架的两个实例：（i）SpecKV，它利用草稿输出准确评估每个KV对的重要性，以实现更有效的KV缓存丢弃，以及（ii）SpecPC，它使用草稿模型的注意力激活来识别和丢弃不重要的提示token。据我们所知，这是首次利用草稿模型加速近似LLM推理的工作，将其效用扩展到传统的无损推测解码之外。我们通过理论和经验分析来阐述我们的方法，并展示了草稿模型和目标模型的注意力模式之间存在强相关性。在长上下文基准上的大量实验表明，我们的方法始终比现有基线实现更高的准确性，同时在内存使用、延迟和吞吐量方面保持相同的改进。我们的代码可在https://github.com/furiosa-ai/draft-based-approx-llm获取。