LazyLLM：高效长上下文 LLM 推理的动态 token 剪枝
计算与语言
2024-07-22 v1 人工智能
机器学习
摘要
基于 transformer 的大语言模型推理包含两个顺序阶段：1）prefilling 阶段计算提示词的 KV 缓存并生成第一个 token，2）解码阶段生成后续 token。对于长提示词，prefilling 阶段必须计算所有 token 的 KV 缓存，这会显著增加生成第一个 token 所需的时间。因此，prefilling 阶段可能成为生成过程的瓶颈。一个开放问题是，是否所有提示 token 都对于生成第一个 token 至关重要。为了回答此问题，我们引入一种新方法 LazyLLM，在 prefilling 和解码阶段选择对下一个 token 预测重要的 token 进行 KV 计算。与静态剪枝方法一次性剪枝提示不同，LazyLLM 允许语言模型在不同的生成步骤中从上下文中动态选择不同的 token 子集，尽管它们可能在之前的步骤中被剪枝。广泛的实验在各种任务的标准数据集上表明，LazyLLM 是一种通用方法，可无缝集成到现有语言模型中，以显著加快生成速度而无需微调。例如，在多文档问答任务中，LazyLLM 将 LLaMA 2 7B 模型的 prefilling 阶段加快 2.34 倍，同时保持准确性。
引用
@article{arxiv.2407.14057,
title = {LazyLLM: Dynamic Token Pruning for Efficient Long Context LLM Inference},
author = {Qichen Fu and Minsik Cho and Thomas Merth and Sachin Mehta and Mohammad Rastegari and Mahyar Najibi},
journal= {arXiv preprint arXiv:2407.14057},
year = {2024}
}