让大模型按输入自动跳过冗余层,提速同时不丢质量。
FiRST: Finetuning Router-Selective Transformers for Input-Adaptive Latency Reduction
- 为每条输入动态选择跳过的变换层,而非固定规则。
- 在保持高质量的前提下,推理延迟显著降低。
- 适合移动端和边缘设备部署,兼容现有加速技术。
自回归大语言模型在视觉与语言处理等领域表现优异,但因逐层顺序计算导致推理延迟高,尤其在移动和边缘设备等资源受限场景下问题突出。现有跳层方法分为两类:早期退出机制(不兼容键值缓存)和基于输入无关的预设规则(忽略输入差异)。为此,本文提出FiRST,通过层级专用路由器,根据输入序列动态选择需跳过的变换层——提示阶段决定解码时跳过哪些层。该方法兼容键值缓存,实现更快速推理,且具备质量感知能力。FiRST与模型无关,可轻松部署于任意预训练大模型。实验表明,其在降低延迟的同时,在多项质量指标上优于其他跳层策略,甚至在某些情况下超越基础模型。因此,FiRST是低资源环境下大模型部署的高效可行方案。
原文摘要 · Abstract (English)
Auto-regressive Large Language Models (LLMs) demonstrate remarkable performance across different domains such as vision and language processing. However, due to sequential processing through a stack of transformer layers, autoregressive decoding faces significant computation/latency challenges, particularly in resource-constrained environments like mobile and edge devices. Existing approaches in literature that aim to improve latency via skipping layers have two distinct flavors - 1) Early exit, and 2) Input-agnostic heuristics where tokens exit at pre-determined layers irrespective of input sequence. Both the above strategies have limitations - the former cannot be applied to handle KV Caching necessary for speed-ups in modern framework and the latter does not capture the variation in layer importance across tasks or more generally, across input sequences. To address both limitations, we propose FiRST, an algorithm that reduces inference latency by using layer-specific routers to select a subset of transformer layers adaptively for each input sequence - the prompt (during the prefill stage) decides which layers will be skipped during decoding. FiRST preserves compatibility with KV caching enabling faster inference while being quality-aware. FiRST is model-agnostic and can be easily enabled on any pre-trained LLM. Our approach reveals that input adaptivity is critical - indeed, different task-specific middle layers play a crucial role in evolving hidden representations depending on tasks. Extensive experiments show that FiRST significantly reduces latency while outperforming other layer selection strategies in quality metics. It retains competitive performance to base model (without layer skipping) and in some cases, even improves upon it. FiRST is thus a promising and efficient solution for LLM deployment in low-resource environments.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。