[论文解读] Sparse Model Soups: A Recipe for Improved Pruning via Model Averaging

Max Zimmer, Christoph Spiegel+1 more|arXiv (Cornell University)|Jun 29, 2023

Advanced Neural Network ApplicationsComputer Science被引用 3

一句话总结

本文提出稀疏模型汤（Sparse Model Soups, SMS），一种通过在微调阶段对使用不同超参数（如学习率调度和权重衰减）训练的多个剪枝模型进行平均，来增强迭代剪枝（IMP）的方法。通过确保所有模型具有相同的稀疏连接结构，SMS 在保持稀疏性的同时显著提升了泛化能力和分布外性能，且优于单个剪枝模型以及标准 IMP，同时训练时间更短。

ABSTRACT

Neural networks can be significantly compressed by pruning, yielding sparse models with reduced storage and computational demands while preserving predictive performance. Model soups (Wortsman et al., 2022) enhance generalization and out-of-distribution (OOD) performance by averaging the parameters of multiple models into a single one, without increasing inference time. However, achieving both sparsity and parameter averaging is challenging as averaging arbitrary sparse models reduces the overall sparsity due to differing sparse connectivities. This work addresses these challenges by demonstrating that exploring a single retraining phase of Iterative Magnitude Pruning (IMP) with varied hyperparameter configurations such as batch ordering or weight decay yields models suitable for averaging, sharing identical sparse connectivity by design. Averaging these models significantly enhances generalization and OOD performance over their individual counterparts. Building on this, we introduce Sparse Model Soups (SMS), a novel method for merging sparse models by initiating each prune-retrain cycle with the averaged model from the previous phase. SMS preserves sparsity, exploits sparse network benefits, is modular and fully parallelizable, and substantially improves IMP's performance. We further demonstrate that SMS can be adapted to enhance state-of-the-art pruning-during-training approaches.

研究动机与目标

• 解决在剪枝神经网络中有效结合模型平均与稀疏性的问题。
• 在不增加推理成本的前提下，提升泛化能力和分布外性能。
• 通过在多个微调模型间复用相同的稀疏连接结构，实现稀疏模型的高效、可并行训练。
• 证明在 IMP 过程中，使用不同超参数进行微调可生成适合平均的模型，同时保持稀疏性。
• 表明从平均模型初始化后续剪枝周期可进一步提升性能。

提出的方法

• 使用相同的初始权重但不同的超参数（如学习率调度和权重衰减）对多个剪枝模型进行迭代剪枝（IMP）训练。
• 通过在微调前应用剪枝，确保所有模型共享相同的稀疏连接结构，从而实现参数的直接平均。
• 通过平均多个微调后的模型参数，构建单一稀疏模型，同时保持稀疏性和推理效率。
• 从先前平均的模型开始每个后续剪枝周期，实现性能的逐步提升。
• 使用一致的微调调度（例如 ALLR）以确保模型间收敛与稳定性。
• 利用该方法的模块化与可并行特性，实现多模型扩展，且不增加推理成本。

实验结果

研究问题

• RQ1 能否在不牺牲稀疏性或增加推理成本的前提下，有效结合模型平均与剪枝神经网络中的结构化稀疏性？
• RQ2 与单个模型相比，对使用不同超参数训练的多个剪枝模型进行平均，是否能提升泛化能力和分布外性能？
• RQ3 通过从平均模型初始化后续剪枝周期，能否增强迭代剪枝（IMP）的性能？
• RQ4 在微调过程中，哪些超参数配置能生成既足够相似以支持平均，又具备足够差异性以提升泛化能力的模型？
• RQ5 微调的长度与调度方式如何影响最终稀疏模型汤的性能？

主要发现

• 在 CIFAR-100 上，稀疏模型汤（SMS）在 95% 稀疏度下达到 99.52% 的测试准确率，优于标准 IMP 的 98.79%。
• 在 97% 稀疏度下，SMS 达到 97.68% 的准确率，而 IMP 为 98.90%，表明在高稀疏度下仍保持一致的性能增益。
• 与对单个模型进行相同时间微调相比，该方法减少了训练时间，因为多个模型可并行训练且共享连接结构。
• 消融实验表明，即使仅进行一个微调周期，只要搭配合适的学习率调度（如 ALLR），也能带来显著性能提升。
• 在微调过程中使用多样化的超参数（如学习率调度和权重衰减）可生成既足够相似以支持平均，又具备足够差异性以提升泛化能力的模型。
• 从平均模型初始化下一阶段剪枝周期，可在后续微调阶段进一步提升性能，表明该方法具有级联增益效应。