M ODEL F USION OF H ETEROGENEOUS N EURAL N ETWORKS VIA C ROSS -L AYER A LIGNMENT
Anonymous authors
Paper under double-blind review
A BSTRACT
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 efﬁciently 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 synthetic experiments indicate that the fused network from CLAFusion achieves a more favorable performance compared to the individual networks trained on heterogeneous data without the need for any retraining.
With an extra ﬁnetuning process, it improves the accuracy of residual networks on the CIFAR10 dataset.
Finally, we explore its application for model compression and knowledge distillation when applying to the teacher-student setting.
I NTRODUCTION
The ubiquitous presence of deep neural networks raises the question, “Can we combine the knowledge of two and more deep neural networks to improve performance?"
As one of the earliest successful methods, ensemble learning aggregates the output over a collection of trained models to improve the generalization ability (Hansen & Salamon, 1990; Zhou, 2012).
However, it is expensive in terms of computational resources because it requires storing and running many trained neural networks during inference.
To overcome the previous hardness, we need to ﬁnd a single neural network that can inherit the knowledge from multiple pre-trained neural networks and have a small size at the same time.
To our best knowledge, there are two popular approaches for this purpose: knowledge distillation and model fusion.
The ﬁrst approach, knowledge distillation , is a machine learning technique that transfers knowledge from large networks (teachers) into a smaller one (student) (Hinton et al., 2015).
The smaller network is trained by the task-speciﬁc loss and additional distillation losses that encourage the student networks to mimic the prediction of the teachers.
The distillation process, however, is computationally expensive because it trains the student network from scratch.
In addition, at each training epoch knowledge distillation runs forward inference for all networks and demands sharing the large training data.
In some situations sharing data is prohibited due to critical issues such as data security or network bandwidth constraints.
The second line of works is called model fusion which is the problem of merging a collection of pre-trained networks into a uniﬁed network (Wang et al., 2020; Singh & Jaggi, 2020).
The simplest technique in model fusion is vanilla averaging (Utans, 1996; Smith & Gashler, 2017), which computes the weighted average of pre-trained network parameters without the need for retraining.
Their methods fuse neural networks with the same architecture without considering the permutation invariance nature of neural networks.
To mitigate this problem, the idea of solving a neuron alignment problem before applying weight averaging is discussed in Ashmore & Gashler (2015);
Li et al.
Recently, there are two concurrent works that beneﬁt from the idea of neuron alignment.
FedMA (Wang et al., 2020) formulated and solved the assignment problem to ﬁnd a permutation matrix that undoes the permutation of weight matrices.
While OTFusion (Singh & Jaggi, 2020) viewed the problem through the lens of optimal transport (Monge, 1781; Kantorovich, 2006) and utilized the transport map to align the weight matrices.
Though both methods can fuse multiple neural networks, their applications are limited to networks with the same number of layers.
Contributions.
In this paper, we propose a model fusion framework for fusing heterogeneous neural networks, namely, unequal width and unequal depth neural networks, which is named as Cross-Layer
Alignment Fusion (CLAFusion).
CLAFusion consists of three parts: cross-layer alignment, layer balancing method, and layer-wise model fusion method.
We ﬁrst formulate the cross-layer alignment as an assignment problem using layer representation and layer similarity, then propose a dynamic programming-based algorithm to solve it.
Next, we present two natural and fast methods to balance the number of layers between two networks.
Finally, our experiments demonstrate the efﬁciency of CLAFusion on three different setups.
In skill transfer, our framework successfully fuses two heterogeneous neural networks trained on heterogeneous data and improves the performance over the individual networks without retraining.
In addition, the fused model from CLAFusion serves as an efﬁcient initialization when training residual networks.
Furthermore, CLAFusion shows potential applications for model compression and knowledge distillation in the teacher-student setting.
B ACKGROUND
In this section, we ﬁrst recall the layer representation and layer similarity that have been widely used in applications.
Then, we review available model fusion methods and their limitations.
Finally, we discuss OTFusion and the challenge of applying this framework to heterogeneous neural networks.
2. L AYER R EPRESENTATION AND L AYER S IMILARITY
A common layer representation is the matrix of activations (Kornblith et al., 2019).
Activation matrix, also known as activation map, can be used as a type of knowledge to guide the training in knowledge distillation (Gou et al., 2021).
Another choice is the weight matrix of the trained neural networks, which has been shown to perform well on intra-network comparison tasks Neill et al.
The neural network exhibits the same output when its neurons in each layer are permuted, this is so-called the permutation invariance nature of neural network parameters.
Due to the permutation invariance property, a meaningful layer similarity index should be invariant to orthogonal transformations (Kornblith et al., 2019).
Their proposed Centered Kernel Alignment (CKA) effectively identiﬁes the relationship between layers of different architectures.
2. M ODEL FUSION
Despite showing good empirical results, vanilla averaging (Utans, 1996; Smith & Gashler, 2017) only works in the case when the weights of individual networks are relatively close in the weight space.
As an effective remedy, several works on model fusion considered performing permutation of neurons in hidden layers before applying vanilla averaging.
FBA-Wagging (Ashmore & Gashler, 2015), Elastic Weight Consolidation (Leontev et al., 2019), and FedMA (Wang et al., 2020) formulate the neuron association problem as an assignment problem and align the neurons.
Nevertheless, those variants are not generalized to heterogeneous neural networks.
There is limited literature devoted to the discussion of fusing heterogeneous neural networks.
NeuralMerger
(Chou et al., 2018) is one of the few attempts to deal with this setting.
However, there are two key differences in NeuralMerger from our CLAFusion.
First, their cross-layer alignment is hand-crafted and dedicated to speciﬁc neural network architectures.
Secondly, they decompose weights into lower dimensions and use vector quantization to merge the weights of two networks.
On the other hand, we introduce a systematic way to solve the cross-layer alignment problem, and combining weights is done using the layer-wise model fusion method.
2. M ODEL FUSION VIA O PTIMAL T RANSPORT
Recently, Singh & Jaggi (2020) proposed OTFusion, which solves the neuron association problem based on free-support barycenters (Cuturi & Doucet, 2014).
OTFusion leads to a noticeable improve- [Figure] ment over vanilla averaging when all individual networks have the same architecture.
One advantage of OTFusion over other vanilla averaging-based model fusion methods is that it can fuse networks with different widths (i.e., number of neurons) thanks to the nature of optimal transport.
The challenge of applying OTFusion to heterogeneous neural networks.
OTFusion, however, is still a layer-wise approach that cannot be directly employed in heterogeneous neural networks setting.
We need to ﬁnd two corresponding layers, matching their neurons before averaging weight matrices.
There are two challenges to this scheme.
Firstly, the one-to-one mapping between layers is not available in advance.
A naive one-to-one mapping of layers with the same index as the layer-wise approach may cause a mismatch when the numbers of layers are different.
For example, we analyze two VGG conﬁgurations (Simonyan & Zisserman, 2014, Table 1) with different depths: VGG and VGG13.
The second convolution layer of VGG11 has 128 channels and an input image of size 112 × 112 .
While that of VGG13 has 64 channels and an input image of size 224 × 224 .
Secondly, assume that we have successfully found a one-to-one mapping, there still necessitates a special treatment for remaining layers that has no counterparts.
Simply removing those layers has an adverse effect on the deeper network, this may cumulatively degrade the performance of the fused model.
[Equation]
To overcome the challenges of OTFusion, in this section, we propose Cross-layer Alignment Model
Fusion (in short, CLAFusion ) framework for fusing heterogeneous neural networks.
Notation.
We deﬁne (cid:74) n (cid:75) as a set of integers { 1 , . . . , n } and I n as an identity matrix of size n × n .
We use A, B to denote the individual models and
F to denote the fused model.
l indicates the layer index.
Layer l of model A has p ( l ) A neurons and an activation function f ( l ) A .
The pre-activation and activation vector at layer l of model A are denoted as z ( l ) A , x ( l ) A ∈ R p ( l ) A , respectively.
The weight matrix between layers l and l − 1 of model
A is W ( l,l − 1) A .
The following equations hold between two consecutive layers of model A (we omit the bias terms here).
[Equation]
We stack the pre-activation vector over t samples to form a t − row pre-activation matrix.
Let Z ( l ) A ∈
R t × p ( l ) A denote a matrix of pre-activations at layer l of model
A for t samples, and Z ( l ) B ∈ R t × p ( l ) B denote a matrix of pre-activations at layer l of model B for the same t samples.
The activation matrices X ( l ) A ,
X B ( l ) ∈ R t × p ( l ) B are obtained by applying the corresponding activation functions element-wise to the pre-activation matrices.
Problem setting.
Hereafter, we consider fusing two feed-forward networks A and B of the same architecture family .
Two networks have the same input and output dimension but a different number of hidden layers.
In each network, the hidden layer has an activation function of ReLU, which is the general setting in modern architectures.
Additional work is required to handle bias terms and the batch normalization layer properly so we leave them for future work.
Let m and n be the number of hidden layers of models A and B, respectively.
Taking the input and output layers into account, the
[Figure] [Figure] [Figure] numbers of layers are m + 2 and n + 2 , respectively.
Without loss of generality, assume that m ≥ n .
The layer index of model A and B are { 0 , 1 , . . .
, m + 1 } and { 0 , 1 , . .
. , n + 1 } , respectively.
General strategy.
Our framework has three components, which are illustrated in Figure 1.
The ﬁrst part is a cross-layer alignment which is a one-to-one mapping from hidden layers of model B to hidden layers of model A.
The second part is a layer balancing method that equalizes the number of layers of two models based on the cross-layer alignment.
The last part is a layer-wise model fusion method that is applied to same-sized models.
The third part of our framework adopts any model fusion methods that can fuse two neural networks with the same number of layers.
Next, we give more details about the ﬁrst and second parts as well as provide some concrete methods.
3. C ROSS - LAYER ALIGNMENT
Cross-layer alignment (CLA) is a nontrivial problem that also arises in the context of the feature-based knowledge distillation approach (Zagoruyko & Komodakis, 2016; Tung & Mori, 2019; Chen et al., 2020).
CLA requires ﬁnding a one-to-one function that matches similar layers between two models.
Out of a few model fusion approaches considering CLA, NeuralMerger is based on the hand-crafted layer association, which causes a lack of generalization.
To address this matter, we introduce an efﬁcient method for solving the CLA problem.
Because the input and output layers of the two models are identical in both dimension and functionality, we only align their hidden layers.
In addition, the ﬁrst and last hidden layers usually play a critical role in the model performance (Zhang et al., 2019).
Therefore, we directly match the ﬁrst and last hidden layers of the two models.
The importance of this constraint in CLA is studied in Appendix C.
Furthermore, two networks are feed-forward, so the mapping is necessary to keep the sequential order of layers.
Taking all into consideration, we formulate the CLA problem as follows.
Deﬁnition 1.
Assume that we have two layer representations for hidden layers of two models as,
L A = { L (1) A , . . .
, L ( m ) A } and L B = { L (1) B , . . .
, L ( n ) B } .
Let C i,j denote the cost matrix between them, i.e., C i,j = d ( L ( i ) A , L ( j ) B ) where d is a layer dissimilarity.
The optimal CLA is a strictly increasing mapping a : (cid:74) n (cid:75) (cid:55)→ (cid:74) m (cid:75) satisfying a (1)
= 1 , a ( n ) = m and minimizing (cid:80) ni =1 C a ( i ) ,i .
Discussion.
The above problem is a special case of the (linear) unbalanced assignment problem (Ramshaw & Tarjan, 2012).
The condition of strictly increasing ensures the sequential order of layers when matching.
If we have two increasing layer representations in 1-D, the problem can be interpreted as the Partial Optimal Transport in 1-D problem with uniform weights.
It can be solved efﬁciently in O ( mn ) using the proposed method in Bonneel & Coeurjolly (2019).
However, the increasing constraint on layer representations is quite strict and layer representations might not be in 1-D in general.
On the other hand, the unbalanced assignment problem can be solved using the generalization of the Hungarian algorithm (Kuhn, 1955).
The time complexity of the Hungarian algorithm in our CLA problem is O ( mn 2 ) . Algorithm Input: C = [ C i,j ] i,j
S ( i, i ) ← (cid:80) il =1 C l,l , i ∈ (cid:74) n (cid:75) S (0 , j ) ← 0 , j ∈ (cid:74) m (cid:75) S ( i, 0) ← 0 , i ∈ (cid:74) n (cid:75) for i = 1 to n do for j = i + 1 to m do S ( i, j ) ← min { S ( i, j − 1) , S ( i − 1 , j − 1) + C j,i } end for end for a (1) ← 1; a ( n ) ← m ; i ← n − 1; j ← m − 1 while i ≥ 2 do while j ≥ i + 1 and S ( i, j ) = S ( i, j − 1) do j ← j − 1 end while a ( i ) = j ; i ← i − 1; j ← j − 1 end while
Output: a
Cross-layer alignment algorithm.
We propose an efﬁcient algorithm based on dynamic programming to solve the CLA problem.
Algorithm 1 details the pseudo-code for our algorithm. Given the cost matrix, the time complexity of Algorithm 1 is only O ( mn ) .
Note that the optimal alignment from Deﬁnition 1 is not necessarily unique so we choose one which is obtained by backtracking.
Layer representation. An important process of computing the cost matrix is to design an appropriate representation for the hidden layer. We list three possible layer representations for hidden layers of two models. The ﬁrst representation is the number of neurons in each layer.
This encourages two layers that have a similar number of neurons to match with each other.
The second choice is the activation matrix which is widely used for cross-layer comparison.
We can also use the preactivation matrix instead of the activation matrix.
The third alternative is the weight matrix of the pre-trained model.
Next, we discuss the choice of the cost function (layer dissimilarity).
For the ﬁrst representation, either the Euclidean distance or its squared version is sufﬁcient.
To measure the discrepancy between two activation matrices we choose the dissimilarity index based on CKA while another option is the Wasserstein distance.
Finally, the cosine distance, Kullback–Leibler divergence, and Wasserstein distance can be used in the weight matrix representation (Neill et al., 2020).
CLA for Convolutional neural network (CNN).
Following NeuralMerger, we do not match fully connected layers and convolution layers.
It is natural that a convolution layer has different functionality from a fully connected layer.
In addition, there is a lack of an appropriate way to combine the weights of a fully connected layer and a convolution layer.
For the rest of this paper, we consider CNN architecture which consists of consecutive convolution layers followed by fully connected layers such as VGG (Simonyan & Zisserman, 2014) and RESNET (He et al., 2016).
The key idea is to solve the CLA for same-type layers separately then combine two mappings.
Note that the combination of two strictly increasing mappings is also a strictly increasing mapping, thus it satisﬁes the strictly increasing requirement.
We further break the convolution layers part into smaller groups and ﬁnd the mapping for each pair of groups.
In particular, we divide into 5 groups separated by the max-pooling layers in VGG.
In RESNET, we group consecutive blocks with the same number of channels, also known as stage, and in each stage, align blocks instead of layers.
For the block representation, we choose the (pre-)activation matrix of the second convolution layer in each block.
3. L AYER BALANCING METHODS
As mentioned above, some layers in model A may have no counterparts in model B in the optimal mapping.
Naturally, we have two opposite directions: reduce layers in model A or add layers into model B. Assume that we have already balanced the number of layers up to layer l of model
B. At layer l + 1 of model B, there are two possibilities.
If a ( l + 1) − a ( l ) = 1 , the layer-wise fusion method is ready to apply up to layer l + 1 .
If a ( l + 1) − a ( l ) > 1 , we can either merge layers between [Figure] a ( l ) and a ( l + 1) of model A or add layers between l and l + 1 of model
B. Next, we discuss two approaches in the case a ( l + 1) − a ( l ) = 2 .
When a ( l + 1) − a ( l ) > 2 , we can repeat the same approach a ( l + 1) − a ( l ) − 1 times.
Add layers into model
B. We add a layer l (cid:48) between layer l and l + 1 of model
B. The new weight matrices are deﬁned as W ( l (cid:48) ,l ) B ← I p ( l [Equation] B and W [Equation]
[Equation] W [Equation] ∈ R
[Equation]
× p [Equation] l ) B .
The new activation function is deﬁned as f ( l (cid:48) )
B ← f ( a ( l )+1) A , which is ReLU.
Because x ( l ) B (cid:23) 0 for all l ∈ (cid:74) n (cid:75) , from Equation 1 we have
[Equation]
Therefore, the information of model B remains unchanged after adding layer l (cid:48) .
Note that the new layer is just an identity mapping, which is a trick that has been used in RESNET and NET2NET (Chen et al., 2015).
Also discussed in NET2NET, the adding layers method is a function-preserving transformation that allows us to generate a valuable initialization for training a larger network.
We later examine this hypothesis in our experiments.
Merge layers in model A.
We merge layer a ( l )+1 into layer a ( l ) of model
A by directly connecting layer a ( l ) to layer a ( l + 1) .
Because f ( a ( l )+1)
A is ReLU, the new weight matrix can be written as
[Equation]
where D ( a ( l )+1)
A ∈ R p ( a ( l )+1)
A × p ( a ( l )+1)
A is an input-dependent diagonal matrix with 0s and 1s on its diagonal.
The i th entry in the diagonal has a value of 1 if the i th entry of z ( a ( l )+1)
A is positive andotherwise.
Because the actual sign of neuron i at layer a ( l ) + 1 varies by the input, we provide a simple estimation.
Given that the neuron i ∈ (cid:74) p ( a ( l )+1) A (cid:75) has a pre-activation vector over t samples as z .
,i ∈ R t , which is the i th column of the pre-activation matrix Z ( a ( l )+1) A .
We estimate the sign of neuron i using either sign of sum (sgn (cid:80) tj =1 z j,i ) or sign of majority (i.e., 1 if at least t / 2 samples have positive activation values, 0 otherwise).
Pros and cons of balance methods.
An advantage of merging layers is that the fused model has fewer layers.
However, merging layers degrades the accuracy of model A and it is slower than adding layers method that does not involve any complex calculations.
On the ﬂip side, adding layers does not affect the accuracy of model B but results in a deeper fused model.
Surprisingly, merging layers shows comparatively better performance than adding layers in our experiments on MLP.
Performance comparison between the two methods will be provided in Appendix C and D.1.
Note that we can use more sophisticated model compression methods instead of the merging layers method.
Similarly, it is possible to replace the adding layers method with a more advanced network expanding technique (Wei et al., 2016).
Here, we just introduce two natural and fast ways to demonstrate our framework.
Balancing the number of layers for CNN.
The same merging method does not work in the case of CNN.
On the other hand, the adding layers method can be easily applied for convolution layers. [Figure]
G l o b a l t e s t a cc u r a c y M F M A M B For VGG, we can set all ﬁlters of a new convolution layer to identity kernels.
For RESNET, we add a new block in which all ﬁlters of two convolution layers become zero kernels while the short-cut connection remains as the identity mapping.
A block without a short-cut connection can be simply replaced by two identity convolution layers as in VGG.
E XPERIMENTS
Outline.
Our method is ﬁrst applied for two MLPs trained on the synthetic MNIST dataset (LeCun et al., 1998).
We then fuse a RESNET34 and a RESNET18 trained on CIFAR10 dataset (Krizhevsky et al., 2009).
The last one is a teacher-student setting in which VGG architecture is utilized.
The detailed settings including training hyperparameters, used assets, and computational resources are speciﬁed in Appendix B. We further conduct ablation studies to validate the effectiveness of crosslayer alignment and compare two layer balancing methods in Appendix C.
The experimental results for different seeds can be found in Appendix D.
4. S KILL TRANSFER
Settings.
In skill transfer, we aim to obtain a single model that can inherit both overall and specialized skills from the two individual models.
We adopt the same heterogeneous data-split technique as in (Singh & Jaggi, 2020, Section 5.1) for two MLPs A and B.
The MNIST dataset is separated into two datasets.
One dataset has all images of label 4 and 10% images of other labels while the other contains the rest.
We train a MLP with 3 hidden layers of size 400 , 200 , 100 for model A and a MLP with 4 hidden layers of size 400 , 200 , 100 , 50 for model B. The pre-activation matrix is obtained by running inference for 400 samples. We use the number of neurons as layer representation and adding layers as the balancing method.
The weight proportion of model B, which is referred to as w B , is swept over 11 different values: w B = { 0 . 0 , 0 . 1 , . . .
Results.
Figure 4 shows that when the weight proportion of model B is small ( 0 . 1 or 0 . 2 ), the fused model improves over both individual models.
The reason is that the specialist model A has higher accuracy than the generalist model
B.
When distributing the weight proportion fairly ( 0 . 5 ), the fused model performs worst due to the effect of heterogeneous training data.
The fused model has the best accuracy of 91 .
31% , which is higher than those of models
A ( 90 .
76% ) and B ( 87 .
This claims that the knowledge has been transferred successfully to the fused model without retraining.
4. G ENERATE AN EFFICIENT INITIALIZATION
Settings. Following the experiments conducted in (Singh & Jaggi, 2020, Section 5.3), CLAFusion can be used to generate an initialization when training the larger neural network.
We apply CLAFusion to pre-trained RESNET34 and RESNET18 on CIFAR10 dataset.
Layer representation is the preactivation matrix of 200 samples while the cost function is calculated by subtracting the linear CKA from 1.
After fusing, we also perform ﬁnetuning to improve the accuracy as observed in previous
[Table] [Table] model fusion works (Chou et al., 2018; Singh & Jaggi, 2020). The ﬁnetuning hyperparameters, which are adopted from (Singh & Jaggi, 2020, Section S3.1.3), are reported in Table
Baselines.
The idea of transferring the knowledge from a pre-trained network to a deeper network before continuing training is similar to NET2NET.
The term M B depth-aligned is referred to the model state after adding layers into the shallower model but before fusing with the deeper model.
We use NET2NET to generate a RESNET34 model ( M B depth-aligned) which preserves the performance of M B .
Note that M B depth-aligned has the same accuracy as the smaller model and the same architecture as the fused model.
To compare with our method, we also ﬁnetune the pre-trained models, M B depth-aligned, and train a RESNET34 from scratch.
As a reference, we report the result of ensemble learning that calculates the average predictions over all individual models.
Results.
As can be seen from Table 1 retraining from the fused model gains an accuracy of 93 .
60 , which is the highest among all initializations.
This demonstrates the advantage of CLAFusion over NET2NET when initializing a large model from a pre-trained small model.
Our method allows combining the knowledge from the pre-trained large model to generate a better initialization, rather than solely relying on the knowledge of the small model.
Although the result is still lower than the accuracy of ensemble learning, the fused model almost halves the computational resources.
Settings.
Heterogeneous model transfer (Wang et al., 2021) is a branch of model transfer that deals with heterogeneous neural networks.
It is contrasted to homogeneous model transfer, also known as transfer learning, in which the same network architecture is used in both the pre-training and ﬁnetuning phase.
In this experiment, we apply CLA and OTFusion to the heterogeneous model transfer method (in short, HMT).
We transfer the pre-trained RESNET18 to the pre-trained RESNET34 using different methods, then ﬁnetune and compare their performance.
We use the same pre-trained models, layer representation, layer dissimilarity, and ﬁnetuning hyperparameters as in Section 4.2. Baselines. The heterogeneous model transfer method (HMT) in the original paper (Wang et al., 2021) is used. Model parameters of the pre-trained RESNET18 are transferred to the pre-trained RESNET34.
The weights that are not transferred are set to the weights of the pre-trained RESNET34.
In the second method (HMT + Naive avg), after transferring we apply naive averaging to the transferred
RESNET34 and the pre-trained RESNET34. Naive averaging is replaced with OTFusion in the third method (HMT + OTFusion). Different from the third method, the last method utilizes the optimal mapping from the CLA problem instead of the longest chain for layer-to-layer transfer.
Results. Table 2 demonstrates the average performance of different transfer methods. The second row reports the model accuracy before ﬁnetuning while the third row shows the ﬁnal performance.
Using
CLA improves the performance of model transfer signiﬁcantly.
After ﬁnetuning, the combination of CLA and OTFusion leads to an improvement of 0 .
77 compared to using HMT only.
4. T EACHER - STUDENT FUSION
Settings.
In this experiment, we transfer the knowledge from the pre-trained teacher model to the student model.
We train two VGG models on CIFAR10 dataset: model A has VGG13 architecture
[Table] [Table]
with a double number of neurons at each layer while model B has VGG11 architecture with a half number of neurons at each layer. After fusing the teacher and student models, we retrain the fused model and compare it to the retraining result of the student model.
The layer measure and the layer dissimilarity are identical to Section 4.2. We adopt the best hyperparameters as reported in (Singh & Jaggi, 2020, Section S11) for retraining. All ﬁnetuning hyperparameters are reported in Table 10.
Results.
The classiﬁcation accuracies are summarized in Table 3.
After retraining, the fused model yields better performance than retraining other student models.
The compression ratio between the fused model and the teacher is about 11 at a cost of reducing 1 . 74% accuracy.
This suggests that
CLAFusion can act as a network pruning method to compress a heavy model into a lightweight one.
Settings.
As discussed in OTFusion, the fused model is an efﬁcient initialization for knowledge distillation.
To examine the same application of CLAFusion, we perform knowledge distillation to the above pre-trained VGG models.
We employ the method in Hinton et al.
(2015) to match the logit distribution of the student model to the teacher model.
For initialization, we consider ﬁve different choices for the student model: (a) randomly initialized M B , (b) M B , (c) M F , (d) M B depth-aligned, (e) randomly initialized M F .
We sweep over a set of hyperparameters to choose the best combination that maximizes the accuracy of the student model. The sets of hyperparameters for distillation also follows the setting in (Singh & Jaggi, 2020, Section S12): temperature T = { 20 , 10 , 8 , 4 , 1 } and loss-weight factor γ = { 0 . 05 , 0 . 1 , 0 . 5 , 0 . For a fair comparison, the hyperparameters for training student models are identical to ﬁnetune hyperparameters in Section 4.3. Results.
The results for retraining and knowledge distillation are reported in Table 4, further results can be found in Table 14.
The ﬁrst row shows the best performance for all combinations of hyperparameters.
The second row is the mean of the best accuracies obtained at different temperatures.
Using the fused model as an initialization achieves the best performance among different choices of initializations.
It showcases the application of CLAFusion as an efﬁcient initialization for knowledge distillation.
If averaging over different temperatures, retraining the fused model (90 . 96) works better than vanilla distilling from the teacher model into the original student architecture (88 .
Even the best accuracy of vanilla distillation (90 . 98) is slightly higher than that of retraining the fused model but comes with the cost of hyperparameter tuning.
It further strengthens the claim that CLAFusion in combination with ﬁnetuning can act as a model compression method.
C ONCLUSION
In the paper, we have presented a framework for model fusion.
Our CLAFusion extends layer-wise model fusion methods to the setting of heterogeneous neural networks by solving a cross-layer alignment problem, followed by a layer balancing step.
CLAFusion has been successfully applied in the setting of heterogeneous data as a skill transfer method.
In addition, ﬁnetuning the fused network from CLAFusion achieves a better accuracy for RESNET trained on CIFAR10 dataset.
Furthermore, it shows potential applications for model compression and knowledge distillation.
Reproducibility Statement: We provide our source codes for all experiments in the paper in the supplementary of the paper.
The details of experimental settings, computational infrastructure, and other used public libraries are given in Appendix B. Neural network architectures and ﬁnetuning process of all experiments in the main text can be found in Appendix D. We also provide links to pre-trained models for each experiment in the README ﬁle of our source code.
Ethics Statement: We expect that there are no negative social impacts from our work.
Supplement
to “Model Fusion of Heterogeneous Neural Networks via Cross-Layer Alignment”
In this supplementary material, we ﬁrst present some thoughts on the application of CLAFusion to multiple models scenario in Appendix A.
The common settings, used assets, and computational resources are described in Appendix B.
Next, we study the importance of constraints on the ﬁrst and last layers as well as the efﬁciency of CLA in Appendix C.
Finally, we provide additional results of our experiments in Appendix D.
A CLAF USION FOR MULTIPLE NEURAL NETWORKS .
The cross-layer alignment problem for multiple neural networks is a non-trivial task.
It can be formulated as a multi-index assignment problem (Spieksma, 2000; Huang & Lim, 2003).
In addition, it has some connections to the multi-marginal optimal partial optimal transport (Figalli, 2010; Kitagawa & Pass, 2015).
It is an interesting direction for future work.
Nevertheless, there are two naive approaches to extend CLAFusion to the case of multiple networks.
Consider K pre-trained models { M i } Kk =1 .
In the ﬁrst approach, we can apply CLAFusion for two models iteratively K − 1 times to fuse K models.
In each iteration, we fuse a pair of models and replace them with their fused model, thus the number of models reduces by 1 .
The second approach is similar to what they did in OTFusion (Singh & Jaggi, 2020).
Starting with an estimation of the fused model M F .
We apply the ﬁrst and second parts of CLAFusion K times to align K pre-trained models with respect to the fused model.
Finally, we average the weights of those aligned networks to produce the ﬁnal weights for the fused model.
The choice of M F plays an important role in this approach.
However, it is unclear how it is chosen in OTFusion. Results.
We tried two naive approaches to fuse multiple networks.
Unfortunately, we could not achieve a desirable outcome.
As for the ﬁrst approach, it cumulatively deteriorates the performance of the fused networks after each iteration. For the second approach, we could not achieve a good initialization for the fused model as in OTFusion. It is important to stress that our settings are more difﬁcult than previous model fusion methods because we consider two networks with different numbers of layers. B E XPERIMENT SETTINGS
Following OTFusion [Table] the layer-wise model fusion is computed using the activation-based alignment strategy in which the pre-activation matrix instead of the activation matrix is used for the neuron measure.
The hyperparameters for pre-trained models are summarized in Table 5.
Without any further speciﬁcation, the random seed for training or retraining is set to its default value.
The accuracy of pre-trained MLP is the accuracy of the last epoch.
While the accuracies of VGG and RESNET are reported as the best performing checkpoint.
For all ﬁnetuning experiments, we always chose the best record among all epochs.
[Table]
[Table]
Open source code.
We adapt the ofﬁcial implementation of OTFusion 1 (Singh & Jaggi, 2020) in our implementation.
For computing the optimal transport, we use Python Optimal Transport (POT) library 2 (Flamary et al., 2021).
For model transfer, the source code of the original paper 3 (Wang et al., 2021) is utilized.
Computational resources.
All experiments are done on 1 Tesla V100 GPU.
C A BLATION STUDIES
Settings.
We compare the performance of different combinations of layer representation and layer balancing methods in our framework.
To prove the efﬁciency of the CLA step, we compare the result of fusing two models using different mappings.
We also consider removing the constraint on the ﬁrst and last hidden layers.
We train two pairs of MLP in 5 different seeds.
In both two pairs, model A has 5 hidden layers of size 400, 200, 100, 50, 25.
Model B has 2 hidden layers of size 400, 100 in the ﬁrst pair while it has 3 hidden layers of size 400, 200, 100 in the second pair.
Results of the optimal mappings.
The cross-layer mappings obtained using Algorithm 1 with different combinations are given in Table 6.
According to Deﬁnition 1 in Section 3.1, the optimal mapping is strictly increasing mapping a : (cid:74) n (cid:75) (cid:55)→ (cid:74) m (cid:75) where n and m are the number of layers of the shallower and deeper networks, respectively.
In this case, we have m = 5 , n = 2 for the ﬁrst pair, and n = 3 for the second pair.
The optimal mapping a is represented as [ a (1) , a (2) , . .
. , a ( n )] .
For example, [1, 2, 5] means that a (1) = 1 , a (2) = 2 , and a (3)
= 5 . We observe that changing the random seed barely affects the result of CLA in this setting.
When keeping both ﬁrst and last layer constraints, all three combinations yield the same mappings in both pairs.
For the second pair, all three combinations share the same mappings in 4 out of 5 seeds even if none of the two constraints are placed (When the random seed is 2, the optimal mapping for pre-activation matrix + 1 - linear CKA is [1, 2, 5] instead of [1, 2, 3]).
On the other hand, removing both constraints in the ﬁrst pair makes each combination have a different mapping.
The effectiveness of the optimal mapping.
We fuse each pair of MLPs for all 10 possible mappings between two models A and B. For the layer balancing method, we try both adding layers and merging layers (the sign of sum estimation).
The average performance of pre-trained models and the fused model across 5 random seeds are summarized in Table 7.
In three out of four scenarios, the best accuracy of the fused model is achieved when imposing both the ﬁrst and last layer constraints.
Removing the last layer constraint decreases the performance of the fused model.
When removing the ﬁrst layer constraint, the accuracy drops even more signiﬁcantly.
Because an input image may consist of negative cell values due to normalization, adding a new layer as the ﬁrst hidden layer does not maintain the accuracy of the shallower model.
Therefore, it is necessary to match the ﬁrst and last hidden layers of the two models.
In addition, the mappings obtained from the CLA step generally result in higher accuracy than other mappings, proving the efﬁciency of the CLA.
Comparing between layer balancing methods, the merging layers method runs slower but leads to a higher accuracy than the adding layers in both pairs. [Table] [Table] [Table]
D A DDITIONAL EXPERIMENT RESULTS D. S KILL TRANSFER
Results.
The results of skill transfer for 5 different random seeds are reported in Table 8.
The best accuracy for each seed is written in bold font.
The special list model A generally has higher accuracy than the generalist model B seeds because it has been trained on all 10 labels.
In 4 out of 5 seeds, the fused model improves over both individual models when the weight proportion of model B is small ( w B = { 0 . 1 , 0 . 2 } ).
Comparison between two layer balancing methods.
For comparison purposes, we further conduct the same procedure but replace adding layers by merging layers (the sign of sum estimation).
Similar to the previous experiment, the fused model attains the best performance when the weight proportion of model B is small ( w B = { 0 . 1 , 0 . 2 } ).
Figure 5 illustrates the average performance of two layer balancing methods on the skill transfer task.
When averaging over 5 random seeds, both methods
[Figure]
G l o b a l t e s t a cc u r a c y
M F M A
M B aligned 0 .0 .0 .0 .0 .1 .
Weight proportion of model B 72 .
G l o b a l t e s t a cc u r a c y
M F M A
M B aligned
[Figure]
[Figure] [Figure]
[Table]
perform best if w B = 0 . The best accuracy of merging layers (91 . 57) is slightly higher than that of adding layers (91 . 31) even though the accuracy of M B depth-aligned drops from 87 .
58 to 82 . 10 .
D. G ENERATE AN EFFICIENT INITIALIZATION
Results.
Finetuning from the fused model always improves the accuracy of the pre-trained models as illustrated in Table 11.
In comparison with other initializations, it yields the best performance inout of 5 seeds.
M B depth-aligned, which is the result of NET2NET (Chen et al., 2015) operation, only performs better than ﬁnetuning from scratch and be comparable to that of HMT.
Starting from the results of NET2NET and HMT are even worse than that from the smaller pre-trained model.
(Note that M B → M A divergences with LR = 0 .
so its accuracy is 10, which equals to random guessing.)
[Table]
[Table]
D. H ETEROGENEOUS MODEL TRANSFER
Results.
Table 11 details the results of 4 different model transfer methods.
Combining CLA and
OTFusion with HMT always results in a great boost (at least 17 . 73 ) in the accuracy of the transferred model.
After ﬁnetuning, it gives the best performance in 3 out of 5 random seeds.
D. T EACHER - STUDENT FUSION
Results.
Table 13 illustrates the ﬁnetuning result of teacher and student models across 5 different seeds.
The fused model M F is the most productive initialization of student models in 4 out of 5 seeds.
Retraining the fused model always yields higher accuracy than continuing training the student model.
Discussion
At ﬁrst glance, our approach may resemble the NET2NET operations.
Although both approaches increase the depth of the student network and use the deeper student as a good initialization, there are two major differences between ours and NET2NET.
Firstly, we present a systematic way of ﬁnding the location to add the identity mappings for different types of network architectures while NET2NET is manually designed for a speciﬁc type of network.
Secondly, we do not enlarge the width of the student network (equivalently, use only NET2DEEPNET operation) so that the student network remains more compact than the teacher network.
D. K NOWLEDGE DISTILLATION
Results.
The results of distilling the teacher model into the student model are given in Table 14.
Knowledge distillation from the fused model leads to the best accuracy in 4 out of 5 temperatures.
In addition, both its average and best performance are the highest among all initializations.
M B depth-aligned, which is obtained from the NET2NET operation, is the second-best choice, followed by the pre-trained student model M B .
Naive knowledge distillation into either random initialized student models, however, yields a much lower accuracy.
This suggests that CLAFusion can serve as an efﬁcient initialization for knowledge distillation. [Table]