LazyLLM: Dynamic Token Pruning for Efficient Long Context LLM Inference

transformer-based大型语言模型的推断包括两个连续的阶段：1）预填充阶段计算提示的KV缓存并生成第一个令牌，2）解码阶段生成后续令牌。对于长提示，必须在预填充阶段为所有令牌计算KV缓存，这可能会显著增加生成第一个令牌所需的时间。因此，预填充阶段可能成为生成过程的瓶颈。一个未解决的问题是，是否所有提示令牌都对生成第一个令牌至关重要。为了回答这个问题，我们引入了一种新方法LazyLLM，它能够有选择地计算在预填充和解码阶段对下一个令牌预测重要的令牌的KV。与一次性修剪提示的静态修剪方法不同，LazyLLM允许语言模型在不同的生成步骤中动态选择来自上下文的不同令牌子集，即使它们在以前的步骤中可能已被修剪。在各种任务的标准数据集上进行的广泛实验表明，LazyLLM是一种通用方法，可以与现有的语言模型无缝集成，以显著加速生成而无需微调。例如，在多文档问答任务中，LazyLLM加速了LLama 27B模型的预填充阶段2.34倍，同时保持准确性。

LazyLLM论文试图解决生成长文本任务时，预填充阶段的速度瓶颈问题，提出了一种新方法来加速生成过程。

LazyLLM的关键思路是在预填充和解码阶段中，仅选择与下一个预测标记相关的重要令牌来计算KV缓存，从而加速生成第一个标记。

该论文提出的方法是通用的，可以与现有的语言模型无缝集成，而不需要微调。在多文档问答任务中，LazyLLM加速了LLama 2 7B模型的预填充阶段2.34倍，同时保持准确性。论文还提供了实验细节和数据集信息，并公开了代码。

在相关研究中，一些工作旨在通过对模型进行剪枝来加速生成过程，如Sparse Transformers和Adaptive Computation Time。