PaperlayerView exact revision on arXiv

Source-linked AI summary

ZipIt! Merging Models from Different Tasks without Training

George Stoica, Daniel Bolya, Jakob Bjorner, Pratik Ramesh, Taylor Hearn, Judy Hoffman

arXiv:2305.03053v3cs.CVcs.LG

TL;DR

The paper asks how independently initialized models trained on disjoint tasks can be combined into one multitask model without retraining. ZipIt! addresses this with within-model feature merging and partial zipping, and significantly outperforms prior work across difficult settings, with reported improvements of up to 20% and over 15% for the two strategies.

  • Problem

    Merging independently initialized models trained on completely disjoint tasks without additional training is difficult because prior methods do not fully exploit similarities or account for dissimilarities.

  • Method

    ZipIt! is a graph-based framework that zips correlated features within and across same-architecture models and can partially zip them to retain separate task-specific heads.

  • Results

    ZipIt! significantly outperforms prior work across several difficult settings, with within-model merging improving accuracy by up to 20% and partial zipping by over 15%.

  • Takeaways & Limitations

    Merging models trained on disjoint tasks becomes more feasible without retraining when the method separately exploits shared features and preserves task-specific structure.

  • Takeaways & Limitations

    The general local zip operation applies to layers with weight matrices and cannot be directly inserted into layers such as BatchNorm or ReLU.

Abstract

from arXiv · show

Typical deep visual recognition models are capable of performing the one task they were trained on. In this paper, we tackle the extremely difficult problem of combining distinct models with different initializations, each solving a separate task, into one multi-task model without any additional training. Prior work in model merging permutes one model to the space of the other then averages them together. While this works for models trained on the same task, we find that this fails to account for the differences in models trained on disjoint tasks. Thus, we introduce "ZipIt!", a general method for merging two arbitrary models of the same architecture that incorporates two simple strategies. First, in order to account for features that aren't shared between models, we expand the model merging problem to allow for merging features within each model by defining a general "zip" operation. Second, we add support for partially zipping the models up until a specified layer, naturally creating a multi-head model. We find that these two changes combined account for 20-60% improvement over prior work, making it more feasible to merge models trained on disjoint tasks without retraining.

1 INTRODUCTION

ZipIt! addresses merging independently initialized models trained on separate tasks without additional training. It combines within-model feature merging with partial zipping to create multitask models and outperforms prior approaches.

  • Existing model-merging methods combine models trained on the same task, whereas ZipIt! targets models trained on completely separate tasks.
  • Permutation alone assumes most features across models are correlated, an assumption that may fail for models trained on different tasks.ZipIt! instead supports zipping correlated features within and across models.
  • Up to 20% accuracy improvement over permutation-based approaches comes from allowing correlated features to merge within each model.
  • Partial zipping merges layers only up to a specified point, then preserves the original networks’ remaining layers as separate heads.Depending on task difficulty, this can improve accuracy by over 15% while keeping most layers merged.
  • ZipIt! merges same-architecture models trained on different tasks into one multitask model without retraining, using graph-based merging and unmerging.The method is evaluated on disjoint CIFAR and ImageNet categories and on several independent datasets.

2 RELATED WORK

Model merging grew from averaging related checkpoints to aligning differently initialized models through permutations and feature correlations. These approaches remain centered on models trained for the same task, while ZipIt! extends the setting to disjoint tasks.

  • Model Soups average models finetuned from the same pretrained initialization to improve accuracy and robustness.
  • Mode-connectivity methods seek low-loss paths or interpolations between models, motivating weight averaging when models occupy compatible basins.
  • Prior merging methods primarily combine models trained on the same task, whereas ZipIt! targets separately trained models with different tasks and label sets.
  • Git Re-Basin aligns differently initialized models trained on the same data, while REPAIR adds parameters and adjusts batch normalization.REPAIR improves Git Re-Basin’s accuracy according to the related-work discussion.

3 BACKGROUND AND MOTIVATION

Weight interpolation can merge models in a shared loss basin, and permutation can sometimes align differently initialized models to such a basin. For disjoint tasks, however, task-specific basins make permutation-based interpolation perform poorly.

  • Models finetuned from the same checkpoint can often be merged by linearly interpolating their weights because they occupy the same loss basin.
  • Models with different initializations often occupy different modes, so direct interpolation typically produces random accuracy.
  • Permutation can align differently initialized models into a common basin by exploiting internal symmetries that preserve network outputs.
  • For models trained on different tasks, the optimal permutation may remain favorable for one task but outside a basin for the other, making interpolation worse than either original model.

4 ZIPIT!

ZipIt! merges corresponding layers by matching and combining redundant features within and across models, while propagating merge and unmerge operations through intervening modules. It also supports partial zipping, preserving later task-specific layers as separate heads.

  • Feature zipping: ZipIt! combines features within each model and across models, unlike permutation-based merging, to better retain performance when task-specific features differ.The method merges each layer’s feature spaces and can combine features from the same model or from both models.
  • Feature zipping: Feature similarity is measured by pairwise activation correlations over a small unlabeled image set, including all features in the concatenated model space.Greedy matching then pairs the highest-correlation features without replacement.
  • Layer-wise zip operation: The merge matrix M averages matched features into a shared output space, while the unmerge matrix U maps that space back for compatibility with subsequent layers.For a match between two features, the corresponding row of M contains 1/2 at both feature indices.
  • Layer-wise zip operation: ZipIt! fuses merge and unmerge operations into corresponding layer parameters and propagates them through parameter-free or element-wise modules and skip connections.Merge and unmerge matrices are fused when a layer with weights is reached, while propagation handles modules such as ReLU, BatchNorm, and skip connections.
  • Partial Zip: Partial zipping merges only early layers and leaves later layers unmerged, producing a multi-head model when output spaces are incompatible or full merging would lose accuracy.The latest unmerge operation aligns inputs to the first unmerged layer in each original network.
  • Extensions: The same-model budget β controls the percentage of merged features that may come from within-model merges, with β = 0 recovering the across-model-only operation.Each model receives an equal portion of the within-model merge budget.

5 RESULTS

ZipIt! is evaluated on disjoint class splits, ImageNet subsets, completely different datasets, and different output modalities. Across these settings, partial zipping consistently improves performance over permutation-based merging and can approach ensemble accuracy without retraining.

  • CIFAR results: ZipIt! improves joint CIFAR-10 accuracy by 32.9% over Git Re-Basin and 20.7% over the Permute baseline.
  • CIFAR results: 83.8% joint accuracy is achieved on CIFAR-10 when ZipIt! zips 13 of 20 layers, only 3.6% behind the ensemble.
  • CIFAR results: ZipIt!13/20 exceeds Git Re-Basin by 29.2% on CIFAR-100 joint accuracy and remains 3.3% below the ensemble.
  • ImageNet-1k results: On ImageNet-1k, Git Re-Basin reaches 3.1% joint accuracy, whereas all-layer Permute and ZipIt! reach 8.6%.ZipIt!10/50 reaches close to ensemble accuracy while saving FLOPs.
  • Multi-dataset merging: Partial zipping improves multi-dataset merging by up to 33% for pairs and 50% for all four models over the Permute baseline.
  • Multiple output modalities: ZipIt! merges a DeeplabV3 segmentation backbone with an ImageNet classifier, retaining good performance on both modalities with half the layers merged.

6 ANALYSIS

The analysis studies when within-model merging and model scale matter. Same-model merges help substantially on CIFAR, while increased capacity enables ZipIt! to approach ensemble performance.

  • Merging within models: On CIFAR-100, the optimal within-model merge budget exceeds 0.8, whereas ImageNet shows fewer redundant features within each model.
  • Model scale: When ResNet-20 width is below 4×, ZipIt! and Permute perform identically; with greater scale, ZipIt! trends toward the 75% ensemble upper bound.
  • Matching algorithm: Allowing both permutations and within-model feature merging improves CIFAR-10 joint accuracy by up to 21.2% over identity or permutation matching.The greedy graph-based approach is nearly as accurate as the optimal algorithm while being two orders of magnitude faster.

7 CONCLUSION

ZipIt! addresses model merging across disjoint tasks without additional training by exploiting both cross-model and within-model feature relationships. The paper reports substantial gains over prior methods across difficult settings and provides algorithmic and experimental details for reproducibility.

  • ZipIt! targets models trained on completely disjoint tasks, where prior merging methods underperform because they neither exploit similarities nor account for dissimilarities.
  • The authors state that code, experiments, baselines, algorithm details, experimental details, and a proof of Theorem 1 will be released or included.
  • Feature correlations decrease through successive partial-zipping locations, consistent with deeper layers encoding increasingly task-specific features.

B DATA USAGE

ZipIt! needs only a small activation sample, while augmentation matters more than dataset size and layer propagation includes a nonlinear-feature approximation.

  • Data Usage: Activation matching uses the full CIFAR training set, 1% of ImageNet, or the smallest dataset size for the multi-dataset experiment.The same training augmentations are used for activation computation.
  • Data Usage: 100 images on CIFAR-100 and 640 images on ImageNet match results obtained from much larger activation datasets.These correspond to 2 batches for CIFAR-100 and 0.05% of ImageNet, versus 50,000 and 64,048 images respectively.
  • Data Usage: Data augmentation adds almost 4% average task accuracy on CIFAR-100 but slightly hurts ImageNet performance within variance.The authors nevertheless use augmentation consistently in experiments.
  • Partial Zip: Leaving some layers unzipped recovers substantial performance while retaining most of the merged model.This is the purpose of varying partial zip depth.
  • Zip Propagation: ReLU layers require no merge operation, but carrying merge matrices through them approximates a linear merge on nonlinear features.The passage identifies this approximation as a potential direction for future work.

D CROSS ENTROPY ON CIFAR

With cross-entropy training, fully merging models with incompatible output spaces performs poorly; partial zipping substantially mitigates this issue across ResNet and VGG experiments.

  • Training Objective: CLIP loss aligns output spaces, which is necessary for prior methods that merge the entire network to obtain good accuracy.The main CIFAR models therefore use CLIP embeddings of class names as targets.
  • CIFAR-100: ZipIt! with partial zipping recovers much of the performance lost by fully merging CIFAR-100 models under cross-entropy.The accompanying table compares ZipIt! against baselines using 16×-width ResNet-20 models.
  • Cross Entropy: Partial zipping is extremely important with cross-entropy because later layers become incompatible when models have different output spaces.Fully merging the network fails in this setting, whereas leaving later layers task-specific can mitigate the problem.
  • VGG: ZipIt! outperforms all baselines on the smaller VGG11 CIFAR-10 setting, while partial zipping is important for matching original-model accuracy.The experiment uses cross-entropy and a 1×-width model, limiting the available capacity.

E IMAGENET WITH 1.5X WIDTH

Increasing ResNet-50 width improves ZipIt! substantially on the challenging ImageNet-1k setting, especially when only part of the network is zipped.

  • Motivation: The experiment tests whether ZipIt!’s strong width scaling on CIFAR-100 also appears on the harder ImageNet-1k task.The comparison uses original 1× and expanded 1.5×-width ResNet-50 models.
  • Width Comparison: 49.0% joint accuracy replaces 33.1% when merging 37 layers as width increases from 1× to 1.5×.The corresponding 1× and 1.5× models use cross-entropy, for which fully zipped models perform poorly.
  • Width Comparison: 64.1% becomes 55.8% for 22 merged layers, and 66.8% becomes 60.9% for 10 merged layers, with the latter within 1% of the ensemble.These gains show that extra width helps across partial-zipping depths.

F MERGING MODELS WITH DIFFERENT OUTPUT MODALITIES

ZipIt! merges models for semantic segmentation and image classification by combining features across and within models, with partial merging preserving task-specific layers.

  • Experimental Setting: The experiment combines different initializations, datasets, tasks, and output modalities using a DeepLabV3 ResNet-50 backbone and an ImageNet ResNet-50.The segmentation backbone was additionally finetuned on Pascal VOC after ImageNet pretraining.
  • Different Output Modalities: ZipIt! merges nearly half of the ResNet-50 layers between a Pascal VOC segmentation model and an ImageNet-1k classifier without training.Performance is evaluated using Pascal VOC mIoU and ImageNet-1k accuracy.
  • Theory: Allowing feature merges within each model yields a tighter bound than permutation-only merging in the stated theoretical setting.The paper introduces and proves a theorem extending prior bounds beyond permutations.
  • Background: The prior theorem states that sufficiently wide randomly initialized two-layer networks are likely linearly mode connected after a suitable permutation.ZipIt! generalizes the function-invariant transformations used for merging beyond permutation matrices.
  • Theory: The theoretical extension permits arbitrary reducibility between network parameter sets and relaxes uniformity assumptions while retaining the same bounds.It also does not require models to lie on the same expansion manifold.

G.2 A THEORETICAL RESULT

Theorem 1 extends prior permutation-based bounds by allowing within-model reductions and feature transformations, yielding a tighter barrier bound. The result extends beyond uniform random initializations under boundedness and centering conditions.

  • Theorem 1 yields a strictly tighter bound than Theorem 3.1 when transformations and within-model reduction are allowed, and matches it when r = h.The proof adds the reducibility property to the prior derivation.
  • Transformations T and T′ reduce each model to widths r and r′, then expand and permute them back to width h before interpolation.The transformations are composed as permutation, expansion, and reduction operations.
  • Within-model redundancy can produce a strictly lower barrier and approach zero barrier faster as h increases than permutation-only merging.The stated theoretical result explicitly holds for random initializations and is supported by experiments.
  • The bound depends on the unmatched-row count n = (r + r′) − h, which is no larger than h and therefore tightens the corresponding size bound.The derivation replaces h with n in the concentration and cardinality bounds.
  • The same bounds extend to arbitrary continuous distributions when weights are bounded and output parameters are centered, so uniform sampling is unnecessary.The extension assumes the relevant distributions satisfy the stated boundedness and centering conditions.

H EXPERIMENTS IN SETTINGS OF CONCURRENT WORKS

The concurrent-work comparison evaluates ZipIt! without REPAIR or additional training, making the reported results lower bounds for those stronger settings. ZipIt!19/20 nearly matches ensemble task-average accuracy while using half the FLOPs and outperforming the nearest baseline.

  • All Table 10 results are lower bounds because they exclude REPAIR and training, unlike the concurrent comparison setting.REPAIR adds new parameters and substantially improves merging methods.
  • ZipIt!19/20 reaches near-ensemble performance with half the FLOPs and vastly outperforms the nearest baseline.Partially merging brings ZipIt! even closer to the ensemble.
  • ZipIt!19/20 nearly matches the ensemble “Task Avg.” accuracy without requiring any training.Git Re-Basin performs especially poorly in this setting.
Loading 2305.03053v3…