将不同初始化的Transformer模型高效合并为单一多任务模型
Foldable SuperNets: Scalable Merging of Transformers with Different Initializations and Tasks
- 用冻结原模型权重训练超网，通过特征重建目标实现合并
- 在低数据场景下优于现有方法，支持多种模型规模与任务类型
- 兼具高效性与强表达力，适合资源受限的模型融合场景
近期方法致力于将相同架构、在不同任务上训练的神经网络合并为单个多任务模型。多数工作聚焦于从同一预训练网络初始化的模型合并，而本文解决更难的问题：从不同初始化、在不同任务上训练的大规模Transformer模型的合并。我们发现传统合并方法在此场景下表现崩溃，而知识蒸馏（KD）虽效果更好但成本更高且数据效率低，未充分利用原始模型权重。为此，我们提出“可折叠超网合并”（FS-Merge），通过特征重建目标训练一个包含原模型（权重冻结）的超网。训练后，将超网折叠回单个原模型大小。该方法简单、数据高效，计算成本接近KD，且理论证明其在MLP模型上表达力优于传统方法。实验显示，其在各类大小的MLP与Transformer模型上均达当前最优性能，涵盖多种任务、模态及分布偏移，尤其在低数据场景下优势显著。
原文摘要 · Abstract (English)
Recent methods aim to merge neural networks (NNs) with identical architectures trained on different tasks into a single multi-task model. While most works focus on the simpler setup of merging NNs initialized from a common pre-trained network, we target the harder problem of merging large transformers trained on different tasks from distinct initializations. We show that traditional merging methods fail catastrophically in this setup, while Knowledge Distillation (KD) achieves much better results, though at a higher cost. However, KD is data-inefficient, as it does not exploit the original models' weights. To solve this, we introduce "Foldable SuperNet Merge" (FS-Merge), which trains a SuperNet containing the original models (with frozen weights) using a feature reconstruction objective. After training, the SuperNet is folded back to the size of a single original model. FS-Merge is simple, data-efficient, has a computational cost comparable to KD, and is proven to have superior expressiveness compared to traditional merging methods on MLP models. It achieves SOTA results when tested on MLPs and transformers across various sizes, tasks, modalities, and distribution shifts, especially in low-data scenarios.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。