Prompt-based Depth Pruning of Large Language Models

📄 arXiv: 2502.04348v3 📥 PDF

作者: Juyun Wee, Minjae Park, Jaeho Lee

分类: cs.CL, cs.AI

发布日期: 2025-02-04 (更新: 2025-06-12)

备注: Project: https://jwee01.github.io/PuDDing/ Code: https://github.com/tada0347/PuDDing

💡 一句话要点

提出PuDDing：一种基于Prompt路由的大语言模型动态深度剪枝方法

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

关键词: 深度剪枝 大型语言模型 Prompt路由 动态剪枝 模型压缩 推理加速

📋 核心要点

  1. 现有深度剪枝方法忽略了Transformer块重要性对任务的依赖性，导致剪枝效果受限。
  2. PuDDing通过Prompt路由动态地决定要剪枝的Transformer块，从而实现任务自适应的剪枝。
  3. 实验表明，PuDDing在常识推理任务上优于静态深度剪枝方法，并加速了推理过程。

📝 摘要（中文）

深度剪枝旨在降低大型语言模型的推理成本，通过移除不太重要的Transformer块来实现，无需特定的硬件优化。然而，我们的实验结果表明，Transformer块的重要性高度依赖于任务——对一个任务至关重要的块，在另一个任务中移除可能不会降低准确性。基于此，我们开发了一种动态深度剪枝算法，名为PuDDing（Prompt-routed Dynamic Depth Pruning），它根据输入Prompt确定要省略哪些块。PuDDing通过训练一个轻量级的路由器来预测一组选项中最佳的省略集合，而这个选项集合也是以数据驱动的方式构建的。在常识推理基准测试上的实验结果表明，PuDDing有效地加速了语言模型的推理，并且实现了比静态深度剪枝基线更好的任务性能。

🔬 方法详解

问题定义：论文旨在解决大型语言模型推理成本高昂的问题，特别是在Transformer块重要性因任务而异的情况下，如何有效地进行深度剪枝。现有静态深度剪枝方法无法根据不同任务的重要性自适应地选择剪枝的Transformer块，导致性能下降或剪枝效率不高。

核心思路：论文的核心思路是利用Prompt信息来动态地决定要剪枝的Transformer块。通过训练一个轻量级的路由器，根据输入的Prompt预测最佳的Transformer块省略集合，从而实现任务自适应的深度剪枝。这种方法允许模型在不同的任务中选择不同的Transformer块组合，以达到最佳的性能和效率。

技术框架：PuDDing框架主要包含以下几个模块：1) Transformer块省略集合构建模块：以数据驱动的方式构建一组Transformer块省略的候选集合。2) Prompt路由器：一个轻量级的神经网络，根据输入的Prompt预测最佳的Transformer块省略集合。3) 剪枝后的Transformer模型：根据Prompt路由器的预测结果，动态地选择要使用的Transformer块进行推理。整个流程是，给定一个Prompt，Prompt路由器预测一个Transformer块省略集合，然后模型只使用剩余的Transformer块进行推理。

关键创新：PuDDing的关键创新在于Prompt路由机制，它允许模型根据输入Prompt动态地选择要使用的Transformer块。与传统的静态深度剪枝方法相比，PuDDing能够更好地适应不同的任务，从而实现更高的性能和效率。此外，数据驱动的Transformer块省略集合构建方法也保证了剪枝方案的多样性和有效性。

关键设计：Prompt路由器可以使用各种轻量级的神经网络结构，例如多层感知机或循环神经网络。训练Prompt路由器的关键是定义一个合适的损失函数，以鼓励路由器选择能够最大化模型性能的Transformer块省略集合。论文中可能使用了强化学习或者直接监督学习的方法来训练Prompt路由器。Transformer块省略集合的构建可能使用了聚类或者其他数据挖掘技术，以保证集合的多样性和代表性。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

PuDDing在常识推理基准测试上取得了显著的成果，相较于静态深度剪枝方法，PuDDing在保持或提高模型性能的同时，有效地降低了推理成本。具体的性能提升幅度和推理加速比率需要在论文中查找，但总体而言，PuDDing展示了动态深度剪枝在大型语言模型上的有效性。

🎯 应用场景

PuDDing具有广泛的应用前景，可以应用于各种需要降低大型语言模型推理成本的场景，例如移动设备上的自然语言处理、边缘计算环境下的智能问答系统等。通过动态深度剪枝，PuDDing可以在保证模型性能的同时，显著降低计算资源消耗，从而使得大型语言模型能够部署在资源受限的设备上，并加速云端推理服务。

📄 摘要（原文）

Depth pruning aims to reduce the inference cost of a large language model without any hardware-specific complications, by simply removing several less important transformer blocks. However, our empirical findings suggest that the importance of a transformer block may be highly task-dependent -- a block that is crucial for a task can be removed without degrading the accuracy on another task. Based on this observation, we develop a dynamic depth pruning algorithm, coined PuDDing (Prompt-routed Dynamic Depth Pruning), which determines which blocks to omit from the model based on the input prompt. PuDDing operates by training a lightweight router to predict the best omission set among a set of options, where this option set has also been constructed in a data-driven manner. Empirical results on commonsense reasoning benchmarks demonstrate that PuDDing effectively accelerates the inference language models, and achieves better on-task performance than static depth pruning baselines.