FiRST: Finetuning Router-Selective Transformers for Input-Adaptive Latency Reduction

自回归大型语言模型（LLM）在视觉和语言处理等领域展现出卓越的性能。然而，由于通过一系列变压器层进行顺序处理，自回归解码在计算/延迟方面面临显著挑战，尤其是在移动设备和边缘设备等资源受限的环境中。现有文献中旨在通过跳过层来提高延迟的方法有两种不同的方式——1）早期退出 2）与输入无关的启发式方法，即无论输入序列如何，令牌都在预设的层退出。上述两种策略都有局限性——前者无法应用于现代框架中为加速而必需的KV缓存处理，后者未能捕捉到任务间或更广泛地说，输入序列间的层重要性的变化。为了解决这两个局限性，我们提出了FIRST算法，该算法通过使用特定于层的路由器为每个输入序列自适应地选择一组变压器层来减少推理延迟——在预填充阶段，提示决定了解码过程中将跳过的层。FIRST保留了与KV缓存的兼容性，从而实现更快的推理，同时保持质量意识。FIRST与模型无关，可以轻松应用于任何预训练的LLM。我们进一步通过结合LoRA适配器对外部数据集进行微调来提高性能，这在保持延迟优势的同时增强了任务特定的准确性。我们的方法揭示了输入自适应的重要性——事实上，不同的任务特定中间层在隐藏表示的演变中起着关键作用，具体取决于任务。广泛的实验表明，FIRST在保持与基线相当的性能的同时显著减少了延迟，使我们的方法成为低资源环境下LLM部署的有效解决方案。

该论文旨在解决自回归大型语言模型（LLM）在资源受限环境（如移动和边缘设备）中进行解码时面临的显著计算/延迟挑战。现有的早期退出和输入无关启发式方法存在局限性，前者无法处理KV缓存，后者不能捕捉层重要性的变化。

论文提出了一种名为FIRST的算法，通过使用层特定路由器为每个输入序列自适应地选择一部分Transformer层来减少推理延迟。在预填充阶段，输入序列（即提示）决定解码过程中将跳过哪些层。FIRST与KV缓存兼容，能够在保持质量的同时加速推理。此外，通过结合LoRA适配器进行微调，进一步提高了任务特定的准确性，同时保持了低延迟优势。

1. 实验设计包括广泛的实验，展示了FIRST在不同任务和数据集上的性能。2. 使用了多个标准数据集，验证了方法的有效性和通用性。3. 提供了开源代码，便于其他研究人员复现和扩展研究。4. 研究表明，输入自适应是关键，不同的中间层在不同任务中对隐藏表示的演化起着重要作用。5. 值得进一步研究的方向包括如何优化层特定路由器的设计，以及探索更多任务特定的微调策略。

1. "LayerDrop: A Probabilistic Approach to Neural Network Model Selection" - 提出了一种概率方法来选择神经网络模型中的层。 2. "DynamicAdapt: Dynamic Layer Skipping for Efficient Transformer Inference" - 探索了动态跳过Transformer层以提高推理效率的方法。 3. "Efficient Transformers: A Survey" - 对高效Transformer模型的研究进行了综述。 4. "Adaptive Inference in Neural Networks by Dynamic Halting" - 讨论了通过动态停止机制实现神经网络中的自适应推理。