[论文解读] SkipDecode: Autoregressive Skip Decoding with Batching and Caching for Efficient LLM Inference

Luciano Del Corro, Allie Del Giorno+4 more|arXiv (Cornell University)|Jul 5, 2023

Topic Modeling被引用 7

一句话总结

SkipDecode 是一种新颖的自回归大语言模型推理中的逐 token 早期退出方法，通过在序列位置上使用单调递减的退出点，实现了高效的批处理和 KV 缓存。该方法通过跳过较低层并集中计算资源于高层，实现了 2 倍至 5 倍的速度提升，同时性能损失可忽略不计，确保与实际部署流水线的兼容性。

ABSTRACT

Autoregressive large language models (LLMs) have made remarkable progress in various natural language generation tasks. However, they incur high computation cost and latency resulting from the autoregressive token-by-token generation. To address this issue, several approaches have been proposed to reduce computational cost using early-exit strategies. These strategies enable faster text generation using reduced computation without applying the full computation graph to each token. While existing token-level early exit methods show promising results for online inference, they cannot be readily applied for batch inferencing and Key-Value caching. This is because they have to wait until the last token in a batch exits before they can stop computing. This severely limits the practical application of such techniques. In this paper, we propose a simple and effective token-level early exit method, SkipDecode, designed to work seamlessly with batch inferencing and KV caching. It overcomes prior constraints by setting up a singular exit point for every token in a batch at each sequence position. It also guarantees a monotonic decrease in exit points, thereby eliminating the need to recompute KV Caches for preceding tokens. Rather than terminating computation prematurely as in prior works, our approach bypasses lower to middle layers, devoting most of the computational resources to upper layers, allowing later tokens to benefit from the compute expenditure by earlier tokens. Our experimental results show that SkipDecode can obtain 2x to 5x inference speedups with negligible regression across a variety of tasks. This is achieved using OPT models of 1.3 billion and 6.7 billion parameters, all the while being directly compatible with batching and KV caching optimization techniques.

研究动机与目标

• 解决现有逐 token 早期退出策略在自回归大语言模型推理中与批处理和 KV 缓存不兼容的问题。
• 克服推理必须等待批次中最后一个 token 退出的瓶颈，从而限制了潜在的速度提升。
• 通过在批次中所有 token 的每个序列位置建立统一的退出点，实现可预测且可控的计算预算。
• 利用序列中后续 token 更具可预测性的观察，降低对下游 token 的计算需求。
• 通过战略性地将计算资源分配给高层 Transformer 层，在显著降低推理延迟的同时保持高质量的生成结果。

提出的方法

• 在序列位置上引入单调递减的退出策略，即早期 token 使用更多层，而后期 token 使用更少层，基于损失衡量的不确定性递减。
• 为批次中所有 token 在每个序列位置定义单一共享的退出点，避免在后期 token 退出时重新计算 KV 缓存。
• 对后期 token 跳过低层和中间层的 Transformer 层，将计算资源重新定向至高层，以利用早期 token 的计算结果。
• 使用线性衰减函数定义序列中退出位置，确保单调性并避免缓存的 key-value 状态重新计算。
• 与标准的批处理和 KV 缓存优化无缝集成，支持在真实系统中的实际部署。
• 通过为每个序列位置设置最大退出层，确保计算预算可控，避免不可预测的全网络计算。

(a) OPT-350m on OpenWebText . Average loss per token position shows a strong monotonically-decreasing trend for general text.

实验结果

研究问题

• RQ1 能否使逐 token 早期退出策略与自回归大语言模型推理中的批处理和 KV 缓存兼容？
• RQ2 在序列位置上采用单调递减的退出策略是否能在保持生成质量的同时降低计算成本？
• RQ3 当后期 token 在早期 token 之后退出时，能否在不重新计算 KV 缓存的情况下应用早期退出？
• RQ4 在受控计算预算下，SkipDecode 的速度-性能权衡与现有方法相比如何？
• RQ5 该方法能否在不牺牲推理效率的前提下，支持依赖批处理和 KV 缓存的实际部署流水线？

主要发现

• SkipDecode 在 OPT-1.3B 和 OPT-6.7B 模型上，于多个生成基准测试中实现了相对于完整推理 2 倍至 5 倍的速度提升。
• 该方法保持了与完整模型几乎相同的生成质量，即使在高速度提升因子下性能损失也微乎其微。
• 通过强制退出点单调递减，SkipDecode 消除了早期 token KV 缓存重新计算的需求，从而实现了高效的批处理。
• 该方法与批处理和 KV 缓存完全兼容，克服了先前早期退出方法的主要局限。
• 退出位置的线性衰减策略表现有效且稳健，在初步实验中优于幂律等替代函数。
• Oracle 实验确认，观察到的损失趋势（按 token 位置递减）支持早期、深层计算的设计选择。

(b) OPT-350m (finetuned) on Reddit-TLDR . Average loss per token position. Decreasing trend but with a different function.