PaperlayerView exact revision on arXiv

Source-linked AI summary

On Cross-Layer Alignment for Model Fusion of Heterogeneous Neural Networks

Dang Nguyen, Trang Nguyen, Khai Nguyen, Dinh Phung, Hung Bui, Nhat Ho

arXiv:2110.15538v3cs.LGcs.CV

TL;DR

Ensemble inference is costly, while layer-wise OTFusion requires networks to have the same number of layers. CLAFusion aligns and balances layers before fusion, and finetuning the fused residual networks improves accuracy on CIFAR10, CIFAR100, and Tiny-ImageNet.

  • Problem

    Ensemble inference requires storing and running multiple networks, while OTFusion cannot directly fuse networks with different numbers of layers.

    Source: §1 · p. 1, §2.3 · p. 3

  • Method

    CLAFusion uses cross-layer alignment and layer balancing to equalize network depths before applying layer-wise model fusion.

    Source: §3 · p. 4, §3.1 · p. 5

  • Results

    Finetuning the fused network achieves better RESNET accuracy on CIFAR10, CIFAR100, and Tiny-ImageNet, and the framework can improve heterogeneous model transfer performance.

    Source: §5 · p. 13

  • Takeaways & Limitations

    The fused model provides an efficient initialization for training residual networks, with potential applications in model compression and knowledge distillation.

    Source: §1 · p. 2

    Extending cross-layer alignment to multiple networks is non-trivial, and the paper offers a simple approach while leaving the problem as future work.

    Source: §3.3 · p. 9

Abstract

from arXiv · show

Layer-wise model fusion via optimal transport, named OTFusion, applies soft neuron association for unifying different pre-trained networks to save computational resources. While enjoying its success, OTFusion requires the input networks to have the same number of layers. To address this issue, we propose a novel model fusion framework, named CLAFusion, to fuse neural networks with a different number of layers, which we refer to as heterogeneous neural networks, via cross-layer alignment. The cross-layer alignment problem, which is an unbalanced assignment problem, can be solved efficiently using dynamic programming. Based on the cross-layer alignment, our framework balances the number of layers of neural networks before applying layer-wise model fusion. Our experiments indicate that CLAFusion, with an extra finetuning process, improves the accuracy of residual networks on the CIFAR10, CIFAR100, and Tiny-ImageNet datasets. Furthermore, we explore its practical usage for model compression and knowledge distillation when applying to the teacher-student setting.

1 Introduction

The introduction motivates combining pretrained networks in one compact model, then presents CLAFusion as a framework for fusing heterogeneous networks through cross-layer alignment and layer balancing. It reports improved residual-network accuracy after finetuning and potential uses in model transfer, compression, and distillation.

Motivation

  • Ensembles can improve generalization but require storing and running multiple networks, motivating a single compact network that inherits knowledge from several pretrained models.The introduction identifies knowledge distillation and model fusion as two approaches to this goal.Source: §1 · p. 1
  • Knowledge distillation transfers knowledge from large teachers to a smaller student, but requires teacher and student inference during training and often benefits from a good student initialization.A randomly initialized student often does not achieve high performance in the described setting.Source: §1 · p. 1
  • Prior neuron-alignment fusion methods such as FedMA and OTFusion are limited to networks with the same number of layers.The introduction positions CLAFusion as addressing this depth constraint.Source: §1 · p. 2

Contributions

  • CLAFusion combines cross-layer alignment, layer balancing, and layer-wise fusion to fuse heterogeneous networks; its alignment is formulated as an assignment problem solved with dynamic programming.The paper also proposes two fast layer-balancing methods and discusses extending the framework to more than two networks.Source: §1 · p. 2
  • The fused model serves as an efficient initialization for training residual networks, and CLAFusion components also apply to heterogeneous model transfer.The paper further identifies potential applications in model compression and teacher-student knowledge distillation.Source: §1 · p. 2

2 Background

The background reviews layer representations and similarity measures, neuron-alignment approaches to model fusion, and OTFusion's limits when network depths differ. It positions systematic cross-layer alignment as a response to those limits.

2.1 Layer Representation and Layer Similarity

  • Activation matrices and trained-network weight matrices are common layer representations, while Centered Kernel Alignment identifies relationships between layers of different architectures and is invariant to orthogonal transformations.Activation maps can guide knowledge distillation; weight matrices have performed well for intra-network comparisons.Source: §2.1 · pp. 2–3

2.2 Model fusion

  • Vanilla averaging works when network weights are relatively close, and neuron-alignment methods address permutation invariance but are not generalized to heterogeneous networks.The section contrasts CLAFusion's systematic alignment and layer-wise fusion with NeuralMerger's hand-crafted alignment and vector-quantized weight merging.Source: §2.2 · p. 3

2.3 Model fusion via Optimal Transport

  • OTFusion uses optimal transport to handle networks with different widths, but its layer-wise approach cannot directly fuse networks with different depths.Matching layers by index can pair mismatched layers, while removing unmatched layers can reduce the deeper network's performance and degrade the fused model.Source: §2.3 · p. 3

3 Cross-layer Alignment Model Fusion

CLAFusion aligns hidden layers across networks of different depths, balances their layer counts, and then applies layer-wise fusion. Its design combines constrained layer matching with alternative depth-balancing operations and a simple extension to multiple networks.

3.1 Cross-layer alignment

  • For same-family feed-forward ReLU models with shared input and output dimensions but different depths, CLA matches hidden layers in order and fixes endpoint matches.The paper notes that handling bias terms and batch normalization requires additional work.Source: §3 · p. 4, §3.1 · p. 5
  • Dynamic programming finds a minimum-cost, strictly increasing alignment in O(mn), skipping unmatched deeper-model layers and backtracking to recover the mapping.The first and last hidden layers are constrained to match directly.Source: §3.1 · pp. 6–7
  • For CNNs, CLAFusion aligns same-type layers separately, grouping VGG convolutions by max-pooling stages and ResNet blocks by channel-width stages.ResNet block representations use the second convolution layer’s pre-activation matrix.Source: §3.1 · p. 7

3.2 Layer balancing methods

  • Layer balancing can insert identity layers into the shallower model or merge layers in the deeper model; insertion preserves the shallower model’s information, while merging degrades the deeper model’s accuracy.Insertion yields a deeper fused model; merging yields a model with fewer layers and is slower than insertion.Source: §3.2 · pp. 7–9
  • For CNN balancing, added VGG layers use identity kernels, while added ResNet blocks use zero kernels with an identity shortcut; the described merging method does not work for CNNs.Source: §3.2 · p. 9

3.3 Extension to multiple neural networks

  • For multiple networks, CLAFusion depth-aligns each model to an estimated fused model before applying OTFusion; the paper initializes that model with the deepest network and leaves the choice unresolved generally.The multi-network assignment formulation is identified as future work.Source: §3.3 · p. 9

4 Experiments

The experiments evaluate CLAFusion for residual-network initialization, heterogeneous model transfer, teacher–student compression, and knowledge distillation. They report gains across these settings, while also examining computational costs and practical conditions.

4 Experiments

  • The experimental program also covers synthetic skill transfer, supplementary results, and computational comparisons, alongside ablations of alignment and layer-balancing choices.The main experiments use OTFusion for layer-wise fusion and defer detailed settings, assets, and compute resources to the appendix.Source: §4 · p. 9

4.1 Generate an efficient initialization

  • Finetuned CLAFusion initialization reached 93.67, 67.59, and 55.46, the highest among tested initializations across the three datasets.The experiments fuse pretrained RESNET34 and RESNET18 models before finetuning; tested baselines include finetuned models, a NET2NET depth-aligned model, and training from scratch.Source: Table 1 · p. 10, §4.1 · p. 10
  • On CIFAR100, CLAFusion exceeded finetuned pretrained models by at least 2.05 and 2.53 points when fusing 4 and 6 models, respectively.Ensemble learning used approximately 3 and 4.5 times as many resources as CLAFusion for 4- and 6-model fusion, respectively.Source: Table 2 · p. 11, §4.1 · p. 11

4.1.1 Heterogeneous model transfer

  • After finetuning, CLA with OTFusion improved over HMT alone by 0.76 on CIFAR10 and 2.68 on CIFAR100.Using CLA also outperformed the longest-chain mapping by 42.97 on CIFAR10 and 7.42 on CIFAR100.Source: Table 3 · p. 11, §4.1.1 · p. 12

4.2 Teacher-student fusion

  • On CIFAR10, the retrained fused VGG student outperformed other retrained student models at about an 11:1 teacher-to-fused-model compression ratio, with a 1.74% accuracy reduction.The experiment fused a VGG13 teacher with doubled channels and a VGG11 student with half-width channels.Source: Table 4 · p. 12, §4.2 · pp. 12–13

4.2.1 Knowledge distillation

  • Fused-model initialization achieved the best knowledge-distillation performance among tested initializations; averaged over temperatures, fused retraining scored 90.96 versus 90.73 for pretrained distillation.The best pretrained-distillation accuracy was 90.98, slightly above fused retraining, but required hyperparameter tuning; distillation epochs took 12.22 seconds versus 7.48 seconds for finetuning.Source: Table 5 · p. 12, §4.2.1 · p. 13

5 Conclusion

CLAFusion extends layer-wise fusion to heterogeneous networks through cross-layer alignment followed by layer balancing. The conclusion reports improved finetuned residual-network accuracy across three datasets and identifies model transfer, compression, and distillation as further applications.

5 Conclusion

  • CLAFusion extends layer-wise fusion to heterogeneous networks by solving cross-layer alignment and then balancing the networks' layer counts.The conclusion reports improved finetuned RESNET accuracy on CIFAR10, CIFAR100, and Tiny-ImageNet, with potential applications in model transfer, compression, and knowledge distillation.Source: §5 · p. 13

Supplement to “Model Fusion of Heterogeneous Neural Networks via Cross-Layer Alignment”

A Limitations

CLAFusion’s limitations arise from architectural constraints and the layer-balancing methods it currently supports. Its use of OTFusion also carries requirements for model compatibility and post-fusion fine-tuning.

A Limitations

  • CLAFusion does not generalize to LSTMs or Transformers because the framework lacks a reasonable way to combine their complex components.The passage identifies this as a limitation shared by weight-averaging model-fusion methods.Source: A Limitations · p. 14
  • The fused networks must belong to the same architecture family, even though they may differ in depth and width.OTFusion’s activation-based alignment requires matching activation dimensions at corresponding layers; weight-based alignment requires matching kernel sizes.Source: A Limitations · p. 14
  • Bias and batch normalization are not considered in the experiments, although the passage says they can be handled as in OTFusion.The cited discussion points to OTFusion Appendix S1 for that handling.Source: A Limitations · p. 14
  • Additional fine-tuning is needed after fusion to help the fused model surpass the original pretrained models.Source: A Limitations · p. 14
  • Merging layers degrades original-model performance and adds computational overhead; the current merging technique does not apply to convolutional layers.Source: A Limitations · p. 14

B Experiment settings

The experiments use OTFusion after CLAFusion’s alignment and layer-balancing steps, with activation-based alignment based on pre-activation matrices. The implementation and reporting conventions vary across model families and experiments.

B Experiment settings

  • All experiments apply OTFusion after alignment and layer balancing, using pre-activation matrices for activation-based neuron alignment.Source: B Experiment settings · p. 14
  • Pretrained MLP accuracy is reported at the last epoch, while VGG and RESNET use their best checkpoint; fine-tuning reports the best result across epochs.Source: B Experiment settings · p. 14
  • The implementation adapts OTFusion’s official code, uses POT for optimal transport, and uses the original paper’s source code for model transfer.Source: B Experiment settings · p. 14
  • Experiments were run on one NVIDIA A100 40GB GPU.Source: B Experiment settings · p. 16

C Ablation studies

The ablation studies assess optimal cross-layer alignment (CLA) against other strictly increasing mappings, first for fully connected networks and then for CNNs.

C Ablation studies

  • The section compares fused-model performance under the optimal CLA mapping and other strictly increasing mappings.Source: C Ablation studies · p. 16
  • For fully connected networks, the studies examine CLA’s importance and the effects of two constraints on fused-model classification accuracy.Source: C Ablation studies · p. 16
  • The section also evaluates the advantage of CLA when fusing convolutional neural networks.Source: C Ablation studies · p. 16

C.1 Fully connected neural networks

In the MLP experiments, matching both first and last hidden layers yields the best fused-model accuracy, while merging layers is slower but more accurate than adding them. The supplied CNN materials describe a VGG mapping comparison but do not report its outcome.

C.1 Fully connected neural networks

  • The MLP study uses five seeds and compares models with five versus three hidden layers, each beginning with sizes 400, 200, and 100.Model A has hidden-layer sizes 400, 200, 100, 50, 25; model B has 400, 200, 100.Source: §C.1 · p. 16
  • Across the tested MLP combinations, changing the random seed barely affects CLA mappings, and the first-layer constraint does not change them.All three representation-and-cost combinations yield the same mappings when both endpoint constraints are retained.Source: §C.1 · pp. 16–17
  • With either layer-balancing method, the best MLP fused-model accuracy occurs when both first- and last-hidden-layer constraints are imposed.Removing the last-layer constraint slightly lowers performance; mismatching the first layers causes a dramatic accuracy drop.Source: §C.1 · p. 17, Table 10 · p. 17
  • Merging layers runs slower than adding layers but achieves higher accuracy in this MLP setting.Source: §C.1 · p. 17, Table 10 · p. 17

C.2 Convolutional neural networks

  • The CNN materials specify pre-activation representations from 200 samples, a 1 minus linear CKA cost, added layers, and fine-tuning; the VGG table compares mappings.The supplied passages state the comparison’s scope but do not give a winning mapping or performance outcome.Source: §C.2 · p. 17, Table 11 · p. 18

C.2.1 VGG architecture

The VGG experiments examine cross-layer mappings between two models on CIFAR10. The optimal mapping is compared with alternative mappings before and after finetuning.

Settings and the optimal CLA

  • The experiment uses two VGG models on CIFAR10 and considers matching their convolution layers because their input, output, and fully connected layers are identical.The last three group pairs have the same number of layers, so their mappings are trivial; the first two group mappings are the ones to solve.Source: §C.2.1 · p. 18

The optimal CLA

  • The optimal mapping for the first two groups remains unchanged across random seeds.The models are divided into five groups separated by max-pooling layers.Source: §C.2.1 · pp. 18–19

The effectiveness of the optimal mapping

  • After finetuning, the optimal CLA mapping still beats other mappings, with an average accuracy of 90.64.Before finetuning, it leads other mappings by at least 2.66%.Source: §C.2.1 · p. 19

C.2.2 RESNET architecture

For RESNET, CLA aligns blocks within each stage, assigning the two RESNET18 blocks to the first and last RESNET34 blocks. The optimal mapping outperforms alternatives before finetuning and is slightly more accurate after finetuning.

C.2.2 RESNET architecture

  • Each RESNET is divided into five stages, and CLA aligns blocks within stages rather than matching individual layers.Because RESNET18 has two blocks per stage, its first block maps to the first RESNET34 block and its second to the last.Source: §C.2.2 · p. 19
  • Before finetuning, optimal CLA exceeds naive and random mappings by 3.29 and 1.68, respectively; after finetuning it is slightly more accurate than alternatives.Removing the last-layer constraint slightly reduces performance, while the three baseline mappings reach approximately 93.60% after finetuning.Source: §C.2.2 · p. 20, Table 12 · p. 20

D Additional experimental results

D.1 Skill transfer

In skill transfer, a fused model can outperform both individual models when the specialist model receives most of the weight. Across initialization and transfer experiments, CLAFusion also improves downstream results in multiple settings.

D.1 Skill transfer

  • 91.31% is the fused model’s best accuracy, above model A at 90.76% and model B at 87.58%.In four of five seeds, the fused model improves over both individual models when wA is 0.1 or 0.2.Source: §D.1 · p. 21
  • With either layer-balancing method, skill-transfer performance is best when model B’s weight proportion is 0.1; merging reaches 91.57 and 91.65, versus 91.31 for adding.The comparison averages results over five random seeds.Source: §D.1 · p. 22

D.2 Generate an efficient initialization

  • Finetuning from the fused RESNET initialization beats other initializations in all five seeds on both datasets.For four fused RESNETs, final finetuned performance rises from 93.67 to 93.76 on CIFAR10, from 67.59 to 68.92 on CIFAR100, and from 55.46 to 56.07 on Tiny-ImageNet compared with the two-model scenario.Source: §D.2 · pp. 22–23, Table 15 · p. 23
  • For six-model fusion, CLAFusion yields more favorable initialization in every run on CIFAR10, CIFAR100, and Tiny-ImageNet.The authors also report an increase in ensemble and fused-model performance as the number of pretrained models increases.Source: §D.2 · p. 23, Table 16 · p. 24

D.3 Heterogeneous model transfer

  • On CIFAR10, CLAFusion with HMT boosts target-model accuracy after transfer by at least 16.27 in seed 42 and gives the best finetuned result in four of five seeds.On CIFAR100, incorporating CLAFusion into HMT gives the best performance in all five runs.Source: §D.3 · p. 23, Table 17 · p. 24

D.4 Teacher-student fusion

In teacher-student experiments, the fused model is a productive student initialization, and distillation from it performs best across most tested temperatures. The paper also compares the resource costs of fusion, finetuning, and distillation.

D.4 Teacher-student fusion

  • Across five seeds, the fused model MF was the most productive initialization of student models in at least four.Source: §D.4 · p. 25
  • CLAFusion systematically chooses where to add identity mappings and does not enlarge student width, unlike NET2NET’s manually designed operation for a specific network type.The resulting student remains more compact than the teacher.Source: §D.4 · p. 26

D.5 Knowledge distillation

  • Distillation from the fused model gives the best accuracy in at least four of five temperatures, and its average and best performance are highest among initializations.Depth-aligned MB is the second-best initialization, followed by pretrained MB; random initialization performs much worse.Source: §D.5 · p. 26, Table 19 · p. 26

D.6 Computational efficiency of model fusion

  • Ensemble parameters rise from 32M to 96M as more models are combined, while inference time increases from 2.72 to 7.88, nearly 190%.The CLAFusion fused model remains compact regardless of the number of models fused.Source: §D.6 · p. 27, Table 20 · p. 27
  • Knowledge distillation takes at least 65% longer than finetuning with the same VGG initialization; RESNET training-time increases are 385% for MB and 243% for MF.Distillation runs forward inference through both teacher and student, whereas finetuning runs it only through the student.Source: §D.6 · p. 27, Table 21 · p. 27

Related papers

Loading 2110.15538v3…