Scalable Model Merging with Progressive Layer-wise Distillation

本文提出ProDistill算法，通过逐层教师-学生蒸馏高效合并大型预训练模型，理论证明领域特定数据的必要性，并在视觉、语言任务上实现显著性能提升（6.14%-6.61%），展现出优越的内存和计算效率。

Model Merging, Knowledge Distillation, Few-Shot Learning, Efficiency, Large Language Model, Vision Foundation Model

Jing Xu, Jiazheng Li, Jingzhao Zhang

清华大学交叉信息研究院, 上海齐智研究所, 北京理工大学计算机学院

Generated by grok-3

Background Problem

随着大规模预训练模型在语言和视觉领域的成功，越来越多的微调模型被公开用于特定任务。然而，复杂任务往往需要整合多种能力，而传统的多任务学习需要大量数据和计算资源，模型集成则带来存储开销。模型合并（Model Merging）通过加权平均模型权重（任务向量）提供了一种高效解决方案，但合并后模型性能下降的问题依然存在，尤其是在合并模型数量增加或数据不可用时。本文从理论上证明了无数据依赖的合并算法在最坏情况下表现可能极差，强调了领域特定数据的重要性，并提出了一种少样本场景下的高效合并方法，旨在解决性能下降和资源效率问题。

Method

本文提出了ProDistill（Progressive Layer-wise Distillation），一种基于逐层教师-学生蒸馏的模型合并算法，核心思想和步骤如下：

• 核心思想：将模型合并视为知识蒸馏问题，通过逐层对齐合并模型（学生）和微调模型（教师）的中间层特征，逐步将多个微调模型的知识整合到合并模型中，避免传统端到端训练的高内存开销。
• 具体实现：
  • 使用任务向量（微调模型与预训练模型权重之差）作为合并基础，计算合并系数（元素级粒度）以加权平均任务向量。
  • 针对每一层，优化合并系数以最小化合并模型和各微调模型在该层输出特征之间的均方误差（MSE），如公式： $\min_{\lambda_1^{(l)}, \dots, \lambda_T^{(l)}} \sum_{i=1}^{T} \frac{1}{2T|\mathcal{D}_i|} \sum_{(\mathbf{z}_1, \mathbf{z}_2) \in \mathcal{D}_i^{(l-1)}} \left\| \varphi^{(l)} \left( \theta_0^{(l)} + \sum_{j=1}^{T} \lambda_j^{(l)} \circ \tau_j^{(l)}, \mathbf{z}_1 \right) - \varphi^{(l)} \left( \theta_i^{(l)}, \mathbf{z}_2 \right) \right\|^2$
  • 采用双输入设计，即合并模型和微调模型分别使用各自的前一层特征作为输入，确保逐层训练更接近全局目标。
• 关键创新：逐层训练显著降低了内存需求（仅需存储当前层的任务向量和系数），并通过实验证明其性能优于端到端训练，挑战了逐层训练性能较差的传统观点。
• 批判性思考：虽然逐层蒸馏提高了效率，但可能丢失全局信息，尤其是在深层网络中，后层依赖前层特征，若前层合并不佳，可能导致误差累积。此外，元素级合并系数虽然增强了表达能力，但在少样本场景下可能导致过拟合，论文未充分讨论这一风险。

Experiment

实验在视觉分类（Vision Transformer, ViT）、自然语言理解（BERT, RoBERTa）和自然语言生成（LLAMA2, 13B参数）任务上全面评估了ProDistill的表现：

• 数据集与设置：视觉任务使用ImageNet预训练的ViT模型，微调于8个下游数据集；NLU任务基于GLUE基准的8个任务；NLG任务合并WizardMath-13B和Llama-2-13B-Code-Alpaca模型。少样本验证集从训练集中随机采样，每任务64个样本。
• 结果：
  • 在视觉任务中，ProDistill平均性能达86.04%，比基线方法（如Task Arithmetic, AdaMerging）高出6.14%，仅比单独微调模型低4%。
  • 在NLU任务中，性能提升6.61%，在RoBERTa模型上平均达76.41%。
  • 在大模型（>10B参数）合并中，ProDistill在数学和代码任务上表现优于基线，显示出良好的扩展性。
• 效率分析：数据效率高，1-shot验证数据即可提升7%性能；计算效率上，1个epoch即可显著提升性能，10个epoch接近最终结果；内存效率上，ProDistill的GPU内存占用远低于端到端方法（如DistillMerge）。
• 合理性与不足：实验设置覆盖了多种任务和模型规模，基线选择全面（包括无训练和有训练方法），结果可视化（如t-SNE）直观展示了特征分离效果。然而，实验未充分探讨极少数据（<1-shot）或完全无数据场景下的表现，限制了方法普适性的验证。此外，逐层训练可能导致误差累积的问题未在实验中明确分析，特别是在深层模型上的潜在影响。

Further Thoughts

ProDistill的逐层蒸馏方法在资源受限环境下的应用潜力巨大，尤其是在边缘设备上部署AI模型时。然而，其对少样本数据的依赖提示我们思考是否可以结合自监督学习或无监督学习技术，进一步减少对标注数据的依赖。例如，是否可以利用预训练模型的中间层特征，通过对比学习生成伪标签用于合并？此外，逐层训练可能丢失全局信息的风险值得关注，特别是在处理长上下文依赖任务（如长文档理解）时，是否可以通过混合训练（部分层端到端，部分层逐层）来平衡效率和性能？与近期一些动态模型合并方法（如基于任务路由的Mixture-of-Experts）相比，ProDistill保持了原始模型架构的优势，但是否可以在此基础上引入动态路由机制，以进一步提升任务特定性能？这些方向可能为模型合并领域带来新的突破。