Apple Teaches Large Models to Be Lazy: Faster Generation of the First Token While Maintaining Accuracy
内容提要
苹果和Meta AI的研究团队提出了一种名为LazyLLM的新方法，可以在保持准确度的情况下提高Llama 2预填充阶段的推理速度。LazyLLM通过渐进式剪枝和辅助缓存来实现。实验结果表明，LazyLLM在加速推理方面表现优异，同时准确度下降可以忽略不计。LazyLLM的优势包括适用范围广、无需训练、效果好。
延伸解读
为何优化TTFT至关重要
文章指出，LLM推理分为预填充和解码两阶段，预填充阶段耗时称为TTFT。以Llama 2 7B为例，TTFT的walltime是每个后续解码步骤的21倍，在LongBench上约占总生成时间的23%。因此，优化TTFT对提升整体推理效率非常关键，而现有研究多关注解码阶段，对TTFT的改进较少。
动态剪枝与静态剪枝的差异
LazyLLM采用动态剪枝，允许在不同生成步骤选择不同的token子集，即使某些token在之前步骤中被剪枝，后续仍可能被重新选用。相比之下，静态剪枝一次性对所有token进行剪枝，无法根据每一步的预测需求调整。动态剪枝能针对每个生成步骤优化下一token预测，有助于维持模型性能。
Aux Cache如何避免重复计算
由于LazyLLM在预填充阶段进行渐进式剪枝，被剪枝token的KV不会出现在后续层的KV缓存中。若在解码步骤中重新选用这些token，模型无法直接检索其KV。Aux Cache保存被剪枝token的隐藏状态，供后续迭代检索，确保每个token在每个Transformer层最多计算一次，从而避免重复计算，保证LazyLLM最慢时也比标准LLM快。
剪枝位置对性能的影响
实验表明，在同一Transformer层，剪枝越多token，模型性能下降越明显。此外，在后期层剪枝比在早期层剪枝性能更好，说明后期层对token剪枝敏感度更低。基于此，LazyLLM采用渐进式剪枝：早期层保留更多token，随着层数加深逐渐减少token数量，以平衡速度与准确度。
Q&A
LazyLLM方法是如何提高Llama 2的推理速度的？
LazyLLM通过渐进式剪枝和辅助缓存来提高Llama 2的推理速度，允许在不同生成步骤选择不同的token子集，从而减少计算量。
LazyLLM在保持准确度方面表现如何？
实验结果表明，LazyLLM在加速推理的同时，准确度下降可以忽略不计。
LazyLLM的优势是什么？
LazyLLM的优势包括适用范围广、无需训练和效果好，能够显著提升推理速度。
LazyLLM是如何处理剪枝token的？
LazyLLM在每个生成步骤使用逐层token剪枝，通过注意力图确定输入token对预测的影响，从而动态选择重要token。
LazyLLM与标准LLM相比有什么不同？
LazyLLM在推理过程中不使用所有prompt中的token，而是通过剪枝减少计算量，从而加速推理，且准确度保持较高。
LazyLLM的Aux Cache有什么作用？
Aux Cache保存被剪枝token的隐藏状态，避免在后续计算中重复计算，从而提高推理效率。
