Foldable SuperNets: Scalable Merging of Transformers with Different Initializations and Tasks

TL;DR

This paper introduces FS-Merge, a data-efficient method for merging large transformers trained on different tasks and initializations into a single SuperNet, outperforming traditional methods and knowledge distillation especially in low-data scenarios.

Contribution

We propose FS-Merge, a novel SuperNet-based merging technique that is simple, data-efficient, and more expressive than traditional methods for combining diverse large transformers.

Findings

FS-Merge outperforms traditional merging methods and KD.

It achieves state-of-the-art results across various models and tasks.

FS-Merge is especially effective in low-data scenarios.

Abstract

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.