Constrained Edge AI Deployment: Fine-Tuning vs Distillation for LLM Compression

📄 arXiv: 2505.18166v1 📥 PDF

作者: Jacob Sander, David Moe, Achraf Cohen, Brent Venable, Venkat Dasari, Brian Jalaian

分类: cs.LG

发布日期: 2025-05-13

备注: 9 Pages, 2 Figures


💡 一句话要点

针对边缘AI部署,研究LLM压缩中微调与蒸馏的性能差异

🎯 匹配领域: 支柱二:RL算法与架构 (RL & Architecture)

关键词: 边缘AI LLM压缩 模型剪枝 知识蒸馏 自蒸馏 微调 资源受限 CommonsenseQA

📋 核心要点

  1. 边缘AI部署对LLM的计算、内存和连接性提出了严格约束,现有剪枝方法通常针对整个Transformer,复杂度高。
  2. 该研究采用简单的MLP层L2范数剪枝作为基线,重点比较交叉熵微调和KL散度自蒸馏两种再训练策略的影响。
  3. 实验表明,在相同剪枝策略下,KL散度自蒸馏在CommonsenseQA任务上与交叉熵微调性能相当甚至更好。

📝 摘要(中文)

现代基础模型通常通过结构化剪枝和再训练相结合的方式进行压缩,以满足边缘部署的严格计算、内存和连接约束。虽然最先进的剪枝方案针对整个Transformer,但我们采用了一种简单的、逐层的L2范数剪枝,仅对MLP块进行剪枝,作为固定的基线。我们的重点不是实现最大压缩,而是隔离再训练损失函数的影响:(i)使用交叉熵进行微调(L2PFT),这需要带标签的数据,以及(ii)使用KL散度进行自蒸馏,这仅利用教师模型的logits(没有标签)(L2PSD)。我们在OLMo2-7B-SFT模型上评估了这两种pipeline在CommonsenseQA上的性能,该模型适用于边缘网络中典型的间歇性或拒绝连接场景。在相同的剪枝计划下,基于KL的蒸馏在测试精度上与基于CE的微调相匹配或超过,表明即使使用基本的仅MLP剪枝,损失函数的选择也会实质性地影响资源受限环境中压缩模型的恢复。

🔬 方法详解

问题定义:论文旨在解决在资源受限的边缘设备上部署大型语言模型(LLM)时,如何有效地压缩模型并保持其性能的问题。现有方法,特别是针对整个Transformer结构的剪枝方法,计算复杂度高,且对硬件要求较高。此外,如何选择合适的再训练策略以恢复剪枝带来的性能损失也是一个关键挑战。

核心思路:论文的核心思路是简化剪枝过程,仅对Transformer中的MLP层进行L2范数剪枝,并重点研究不同的再训练损失函数对压缩模型性能的影响。通过对比交叉熵微调(需要标签数据)和KL散度自蒸馏(仅需教师模型logits),探索在资源受限环境下更有效的模型恢复方法。

技术框架:整体流程包括三个主要阶段:1) 使用OLMo2-7B-SFT模型作为初始模型;2) 对MLP层进行L2范数剪枝,形成压缩后的模型;3) 使用两种不同的再训练策略(交叉熵微调和KL散度自蒸馏)恢复模型性能。评估指标为CommonsenseQA任务的测试精度。

关键创新:该研究的关键创新在于,它强调了在边缘AI部署场景下,再训练损失函数的选择对压缩模型性能的重要性。通过对比交叉熵微调和KL散度自蒸馏,证明了即使采用简单的剪枝策略,合适的损失函数也能显著提升压缩模型的性能。

关键设计:论文采用逐层L2范数剪枝,仅针对MLP层。再训练阶段,交叉熵微调使用带标签的数据,而KL散度自蒸馏则利用教师模型的logits作为目标。实验中,使用相同的剪枝计划,对比两种再训练策略在CommonsenseQA任务上的性能表现。

🖼️ 关键图片

fig_0
fig_1

📊 实验亮点

实验结果表明,在相同的剪枝计划下,基于KL散度的自蒸馏方法在CommonsenseQA任务上的测试精度与基于交叉熵的微调方法相匹配甚至超过。这表明,即使采用简单的MLP层剪枝策略,选择合适的损失函数也能显著提升资源受限环境中压缩模型的性能。

🎯 应用场景

该研究成果可应用于各种边缘计算场景,例如智能家居、自动驾驶、工业物联网等,在这些场景中,设备通常具有有限的计算资源和网络连接。通过高效地压缩和部署LLM,可以使这些设备具备更强大的自然语言处理能力,从而实现更智能化的应用。

📄 摘要(原文)

Modern foundational models are often compressed via a combination of structured pruning and re-training to meet the strict compute, memory, and connectivity constraints of edge deployments. While state-of-the-art pruning schemes target the entire Transformer, we adopt a simple, layer-wise L2-norm pruning on only the MLP blocks as a fixed baseline. Our focus is not on achieving maximal compression, but on isolating the impact of the re-training loss function: (i) Fine-tuning with Cross- Entropy (L2PFT), which requires labeled data, versus (ii) Self-Distillation with KL-divergence, which leverages only teacher logits (no labels) (L2PSD). We evaluate both pipelines on the OLMo2- 7B-SFT model for CommonsenseQA suitable for intermittent or denied connectivity scenarios typical of edge networks. Under identical pruning schedules, KL-based distillation matches or exceeds CE fine-tuning in test accuracy, demonstrating that, even with a basic MLP-only pruning, the choice of loss function materially affects compressed model recovery in resource-constrained environments.