DRPruning: Efficient Large Language Model Pruning through Distributionally Robust Optimization

📄 arXiv: 2411.14055v2 📥 PDF

作者: Hexuan Deng, Wenxiang Jiao, Xuebo Liu, Jing Li, Min Zhang, Zhaopeng Tu

分类: cs.CL

发布日期: 2024-11-21 (更新: 2025-05-27)

备注: Accepted by ACL 2025 Main Conference

🔗 代码/项目: GITHUB

💡 一句话要点

提出DRPruning，通过分布鲁棒优化实现高效的大语言模型剪枝

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

关键词: 大语言模型剪枝 分布鲁棒优化 模型压缩 异构数据 领域泛化

📋 核心要点

  1. 现有大语言模型剪枝方法在不同领域表现不均，导致性能偏差，无法保证模型在异构数据上的泛化能力。
  2. DRPruning通过动态调整训练数据分布，优化模型在不同数据域上的性能，从而实现更鲁棒的剪枝。
  3. 实验证明，DRPruning在单语、多语环境下均优于同等规模模型，且对领域偏移具有更强的鲁棒性。

📝 摘要（中文）

大型语言模型（LLMs）取得了显著成果，但面临着模型规模和计算成本不断增加的挑战。结构化剪枝可以减小模型尺寸并加速推理，但通常会导致跨领域性能的不均衡下降，从而导致有偏的性能。为了解决这个问题，我们提出了DRPruning，一种在训练过程中动态调整数据分布的方法，以恢复异构和多任务数据上的平衡性能。在单语和多语环境下的实验表明，在困惑度、下游任务和指令调优方面，DRPruning优于类似大小的模型，无论是在剪枝还是持续预训练方面。进一步的分析表明了DRPruning对各种领域和分布偏移的鲁棒性。此外，DRPruning可以自动确定最佳参考损失和数据比例，表明其具有更广泛的应用潜力。代码和脚本可在https://github.com/hexuandeng/DRPruning获得。

🔬 方法详解

问题定义：论文旨在解决大语言模型结构化剪枝后，模型在不同领域或任务上的性能退化不均衡的问题。现有剪枝方法通常会导致模型在某些领域表现良好，而在其他领域表现较差，这限制了模型在实际应用中的泛化能力。

核心思路：DRPruning的核心思路是利用分布鲁棒优化（Distributionally Robust Optimization, DRO）的思想，在训练过程中动态调整数据分布，使得模型在最坏情况下的性能也能得到保证。通过这种方式，模型可以学习到对不同领域或任务都具有鲁棒性的特征表示。

技术框架：DRPruning的整体框架包括以下几个主要阶段：1) 初始模型剪枝：使用现有的结构化剪枝方法对原始大语言模型进行初步剪枝。2) 数据分布调整：根据模型在不同数据子集上的性能表现，动态调整训练数据分布。具体来说，对于模型性能较差的数据子集，增加其在训练过程中的权重；对于模型性能较好的数据子集，降低其权重。3) 模型微调：使用调整后的数据分布对剪枝后的模型进行微调，以恢复模型在各个领域或任务上的性能。

关键创新：DRPruning的关键创新在于将分布鲁棒优化引入到大语言模型剪枝中，通过动态调整数据分布来解决剪枝后模型性能不均衡的问题。与传统的剪枝方法相比，DRPruning能够更好地平衡模型在不同领域或任务上的性能，从而提高模型的泛化能力。

关键设计：DRPruning的关键设计包括：1) 损失函数：使用交叉熵损失函数作为主要的优化目标。2) 数据权重调整策略：根据模型在每个数据子集上的损失值，动态调整该子集的权重。具体来说，可以使用基于梯度的方法或者基于性能指标的方法来调整权重。3) 超参数设置：需要调整的超参数包括学习率、剪枝率、数据权重调整的幅度等。这些超参数的选择需要根据具体的任务和数据集进行调整。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

DRPruning在单语和多语环境下的实验中，均超越了同等规模的模型，在困惑度、下游任务和指令调优方面均有提升。实验结果表明，DRPruning能够有效提高剪枝后模型的性能，并增强模型对领域偏移的鲁棒性。该方法还能自动确定最佳参考损失和数据比例。

🎯 应用场景

DRPruning可应用于各种需要对大语言模型进行高效剪枝的场景，例如移动设备部署、边缘计算等资源受限环境。该方法能够提升剪枝后模型在不同领域和任务上的泛化能力，使其在实际应用中更具鲁棒性和可靠性。未来可进一步探索其在多模态大模型剪枝中的应用。

📄 摘要（原文）

Large language models (LLMs) deliver impressive results but face challenges from increasing model sizes and computational costs. Structured pruning reduces model size and speeds up inference but often causes uneven degradation across domains, leading to biased performance. To address this, we propose DRPruning, a method that dynamically adjusts the data distribution during training to restore balanced performance across heterogeneous and multi-tasking data. Experiments in monolingual and multilingual settings show that DRPruning surpasses similarly sized models in both pruning and continued pretraining over perplexity, downstream tasks, and instruction tuning. Further analysis demonstrates the robustness of DRPruning towards various domains and distribution shifts. Furthermore, DRPruning can determine optimal reference losses and data ratios automatically, suggesting potential for broader applications. Code and scripts are available at https://github.com/hexuandeng/DRPruning.