RAEE: A Robust Retrieval-Augmented Early Exiting Framework for Efficient Inference

📄 arXiv: 2405.15198v2 📥 PDF

作者: Lianming Huang, Shangyu Wu, Yufei Cui, Ying Xiong, Xue Liu, Tei-Wei Kuo, Nan Guan, Chun Jason Xue

分类: cs.CL

发布日期: 2024-05-24 (更新: 2024-09-20)

💡 一句话要点

提出RAEE框架，通过检索增强实现大语言模型高效且鲁棒的提前退出推理。

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 大语言模型 提前退出 检索增强 高效推理 零样本学习

📋 核心要点

  1. 现有提前退出方法依赖于训练内部分类器，训练成本高且性能提升有限，无法有效降低大语言模型的推理开销。
  2. RAEE框架将提前退出问题转化为分布预测，通过检索相似数据并利用其退出信息来指导模型，无需训练额外的分类器。
  3. 实验表明，RAEE能显著加速推理过程，并在多个下游任务上展现出强大的零样本性能，验证了其鲁棒性。

📝 摘要（中文）

本文提出了一种鲁棒的检索增强提前退出框架RAEE，旨在解决大语言模型推理计算开销大的问题。现有方法通常训练内部分类器来决定是否在中间层退出，但训练这些分类器需要大量工作，且性能提升有限。RAEE将提前退出问题建模为分布预测问题，使用相似数据的退出信息来近似该分布。该方法首先收集退出信息以构建检索数据库，然后利用检索到的相似数据的退出信息来指导骨干模型在预测的层退出。实验结果表明，RAEE能够显著加速推理，并在8个下游任务上实现鲁棒的零样本性能。

🔬 方法详解

问题定义：现有基于分类器的提前退出方法需要大量训练才能获得可比较的性能，并且没有充分利用数据之间的相似性。因此，如何设计一种无需大量训练且能有效利用数据相似性的提前退出框架，以降低大语言模型的推理开销，是本文要解决的问题。

核心思路：RAEE的核心思路是将提前退出问题建模为一个分布预测问题。具体来说，对于给定的输入，通过检索数据库中与其相似的数据，并利用这些相似数据的退出信息来近似当前输入的退出分布。然后，根据这个分布来指导模型在合适的层退出，从而实现高效的推理。

技术框架：RAEE框架主要包含以下几个阶段：1) 构建检索数据库：收集大量数据的退出信息，并将其存储在检索数据库中。退出信息包括数据本身以及模型在不同层对该数据进行推理时得到的置信度或预测结果。2) 检索相似数据：对于给定的输入，使用某种相似度度量（例如，余弦相似度）在检索数据库中找到与其最相似的若干个数据。3) 预测退出分布：根据检索到的相似数据的退出信息，计算出一个近似的退出分布。例如，可以统计相似数据在每一层退出的频率，并将该频率作为当前输入在该层退出的概率。4) 指导模型退出：根据预测的退出分布，选择一个合适的层让模型退出。例如，可以选择概率最高的层，或者根据一个预设的阈值来决定是否提前退出。

关键创新：RAEE的关键创新在于利用检索增强的方式来解决提前退出问题。与传统的基于分类器的提前退出方法相比，RAEE无需训练额外的分类器，而是通过检索相似数据并利用其退出信息来指导模型，从而降低了训练成本，并提高了模型的鲁棒性。

关键设计：RAEE的关键设计包括：1) 相似度度量：选择合适的相似度度量方法，以准确地找到与当前输入最相似的数据。2) 退出信息表示：设计一种有效的退出信息表示方法，以便能够准确地反映数据在不同层的退出情况。3) 退出分布计算：设计一种合理的退出分布计算方法，以便能够准确地预测模型在每一层退出的概率。4) 退出策略：设计一种合适的退出策略，以便能够在保证性能的前提下，尽可能地提前退出。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，RAEE在多个下游任务上实现了显著的加速效果，同时保持了与全模型相当的性能。更重要的是，RAEE在8个下游任务上展现出强大的零样本性能，证明了其鲁棒性。相比于传统的基于分类器的提前退出方法，RAEE无需额外的训练，更易于部署和应用。

🎯 应用场景

RAEE框架可应用于各种需要高效推理的大语言模型部署场景，例如移动设备、边缘计算和资源受限的环境。通过降低推理计算开销，RAEE能够提升用户体验，并降低部署成本，加速大语言模型在实际应用中的普及。

📄 摘要（原文）

Deploying large language model inference remains challenging due to their high computational overhead. Early exiting optimizes model inference by adaptively reducing the number of inference layers. Existing methods typically train internal classifiers to determine whether to exit at intermediate layers. However, such classifier-based early exiting frameworks require significant effort to train the classifiers while can only achieve comparable performance at best. To address these limitations, this paper proposes RAEE, a robust Retrieval-Augmented Early Exiting framework for efficient inference. First, this paper demonstrates that the early exiting problem can be modeled as a distribution prediction problem, where the distribution is approximated using similar data's exiting information. Then, this paper details the process of collecting exiting information to build the retrieval database. Finally, based on the pre-built retrieval database, RAEE leverages the retrieved similar data's exiting information to guide the backbone model to exit at the layer, which is predicted by the approximated distribution. Experimental results demonstrate that the proposed RAEE can significantly accelerate inference. More importantly, RAEE can also achieve a robust zero-shot performance on 8 downstream tasks.