ReplaceMe: Network Simplification via Depth Pruning and Transformer Block Linearization

📄 arXiv: 2505.02819v3 📥 PDF

作者: Dmitriy Shopkhoev, Ammar Ali, Magauiya Zhussip, Valentin Malykh, Stamatios Lefkimmiatis, Nikos Komodakis, Sergey Zagoruyko

分类: cs.CL, cs.AI, cs.LG

发布日期: 2025-05-05 (更新: 2025-06-20)

🔗 代码/项目: GITHUB


💡 一句话要点

提出ReplaceMe,一种免训练的深度剪枝方法,通过线性化Transformer块简化网络。

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

关键词: 深度剪枝 Transformer模型 模型压缩 免训练 线性化 大型语言模型 网络简化

📋 核心要点

  1. 现有Transformer模型剪枝方法通常需要额外的训练或微调,计算成本高昂且流程复杂。
  2. ReplaceMe通过线性变换近似被剪枝的Transformer块,并将其无缝集成到剩余网络中,无需额外训练。
  3. 实验表明,ReplaceMe在多个LLM上实现了显著的剪枝比例,同时保持了较高的性能,优于其他免训练方法。

📝 摘要(中文)

本文提出ReplaceMe,一种广义的免训练深度剪枝方法,它能有效地用线性运算替换Transformer块,同时在低压缩率下保持高性能。与需要额外训练或微调的传统剪枝方法不同,我们的方法只需要一个小的校准数据集来估计一个线性变换,该变换近似于被剪枝的块。估计的线性映射可以无缝地与剩余的Transformer块合并,无需任何额外的网络参数。实验表明,ReplaceMe始终优于其他免训练方法,并且与涉及大量再训练/微调和架构修改的最新剪枝方法相比,仍然具有很强的竞争力。应用于多个大型语言模型(LLM)时,ReplaceMe实现了高达25%的剪枝,同时在开放基准测试中保留了大约90%的原始模型性能——无需任何训练或修复步骤,从而最大限度地减少了计算开销。

🔬 方法详解

问题定义:论文旨在解决大型Transformer模型剪枝过程中,需要大量训练或微调带来的计算开销和时间成本问题。现有剪枝方法通常需要对模型进行再训练,以恢复剪枝造成的性能损失,这限制了剪枝方法的效率和应用范围。

核心思路:ReplaceMe的核心思路是用一个线性变换来近似被剪枝的Transformer块的功能。通过这种方式,可以移除整个Transformer块,而无需重新训练模型。该线性变换通过一个小的校准数据集进行估计,从而最小化了对原始模型性能的影响。

技术框架:ReplaceMe的整体流程包括以下几个步骤:1) 选择要剪枝的Transformer块;2) 使用校准数据集,通过最小化剪枝前后块的输出差异,估计一个线性变换来近似被剪枝的块;3) 将估计的线性变换与剩余的Transformer块合并,从而创建一个更小的模型。该过程无需任何训练或微调。

关键创新:ReplaceMe的关键创新在于提出了一种免训练的深度剪枝方法,它通过线性化Transformer块来简化网络,避免了传统剪枝方法中耗时的再训练过程。与现有方法的本质区别在于,ReplaceMe不需要任何梯度更新,而是直接通过线性变换来近似被剪枝的模块。

关键设计:ReplaceMe的关键设计包括:1) 使用最小二乘法来估计线性变换的参数,目标是最小化剪枝前后Transformer块输出的均方误差;2) 校准数据集的选择,需要具有代表性,能够覆盖模型在实际应用中可能遇到的输入分布;3) 线性变换与相邻Transformer块的无缝集成,确保剪枝后的模型结构与原始模型兼容。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

ReplaceMe在多个大型语言模型上进行了实验,结果表明,该方法可以在实现高达25%的剪枝比例的同时,保留大约90%的原始模型性能。与需要大量再训练/微调的最新剪枝方法相比,ReplaceMe具有很强的竞争力,并且优于其他免训练方法。例如,在某些基准测试中,ReplaceMe在保持相似性能的同时,显著减少了模型的参数数量。

🎯 应用场景

ReplaceMe可应用于各种大型语言模型(LLM)的压缩和加速,尤其是在资源受限的环境中,如移动设备或边缘计算平台。该方法能够降低模型的计算复杂度和存储需求,从而提高推理速度和降低功耗。此外,ReplaceMe还可以用于模型蒸馏,将大型模型的知识迁移到小型模型中。

📄 摘要(原文)

We introduce ReplaceMe, a generalized training-free depth pruning method that effectively replaces transformer blocks with a linear operation, while maintaining high performance for low compression ratios. In contrast to conventional pruning approaches that require additional training or fine-tuning, our approach requires only a small calibration dataset that is used to estimate a linear transformation, which approximates the pruned blocks. The estimated linear mapping can be seamlessly merged with the remaining transformer blocks, eliminating the need for any additional network parameters. Our experiments show that ReplaceMe consistently outperforms other training-free approaches and remains highly competitive with state-of-the-art pruning methods that involve extensive retraining/fine-tuning and architectural modifications. Applied to several large language models (LLMs), ReplaceMe achieves up to 25% pruning while retaining approximately 90% of the original model's performance on open benchmarks - without any training or healing steps, resulting in minimal computational overhead (see Fig.1). We provide an open-source library implementing ReplaceMe alongside several state-of-the-art depth pruning techniques, available at https://github.com/mts-ai/ReplaceMe.