KV预测提高首个token的生成速度
计算与语言
2024-10-14 v1 人工智能
摘要
基于Transformer的语言模型推理以提示处理步骤开始。在此步骤中，模型生成第一个输出token并存储用于后续生成步骤的KV缓存。当提示长度或批量大小增加时， billion参数模型在边缘设备上的提示处理步骤可能耗时数十秒甚至更长。这会通过显著增加模型输出的延迟来降低用户体验。为减少生成第一个输出（即"time to first token"，简称TTFT）所需的时间，我们引入一种新方法，称为KV预测。在我们的方法中，使用小型辅助模型处理提示并生成基础模型所需KV缓存的近似值。然后，使用该近似KV缓存进行自回归生成，而无需再查询辅助模型。我们展示，在TTFT FLOPs预算范围内与基线方法相比，我们的方法在效率-准确性之间实现了帕累托最优的权衡。在TriviaQA上，我们在各种TTFT FLOPs预算下均实现了15%-50%的相对准确性提升。在固定TTFT FLOPs预算下，我们还在HumanEval Python代码完成任务上实现了最高30%的准确性提升。此外，我们在Apple M2 Pro CPU上对模型进行基准测试，证明了我们的FLOPs改进在硬件上 translates to TTFT加速。我们在https://github.com/apple/corenet/tree/main/projects/kv-prediction 上发布了代码。
关键词
引用
@article{arxiv.2410.08391,
title = {KV Prediction for Improved Time to First Token},
author = {Maxwell Horton and Qingqing Cao and Chenfan Sun and Yanzi Jin and Sachin Mehta and Mohammad Rastegari and Moin Nabi},
journal= {arXiv preprint arXiv:2410.08391},
year = {2024}
}