Hyper-Representations for Pre-Training and Transfer Learning
Abstract
Learning representations of neural network weights given a model zoo is an emerging and challenging area with many potential applications from model inspection, to neural architecture search or knowledge distillation.
Recently, an autoencoder trained on a model zoo was able to learn a hyper-representation , which captures intrinsic and extrinsic properties of the models in the zoo.
In this work, we extend hyperrepresentations for generative use to sample new model weights as pre-training.
We propose layerwise loss normalization which we demonstrate is key to generate high-performing models and a sampling method based on the empirical density of hyper-representations.
The models generated using our methods are diverse, performant and capable to outperform conventional baselines for transfer learning.
Our results indicate the potential of knowledge aggregation from model zoos to new models via hyper-representations thereby paving the avenue for novel research directions.
1. Introduction
Pre-training is a key element to transfer learning.
What if we could encode the knowledge of an entire population of neural networks for transfer learning to a new domain or dataset?
Recent work (Unterthiner et al., 2020; Sch ¨ urholt et al., 2021; Martin et al., 2021) demonstrated the capability to predict properties of individual neural networks captured by such population of neural networks (often also referred to as model zoo ).
In particular, (Sch¨urholt et al., 2021) proposed to learn a lower-dimensional representation of a model zoo able to capture the underlying manifold of all neural network models populating the model zoo.
In their work, they trained so called hyper-presentations with a transformer autoencoder directly from the weights of a zoo and used them to predict several model properties such as accuracy, hyperparameters or architecture configurations of individual neural networks.
Instead of using hyper-presentations for model properties prediction only, in this work, we propose to exploit the knowledge encoded in hyper-representations as pretraining for subsequent transfer learning.
We train hyperrepresentations of model zoos and use them in combination with the decoder as a generative model i.e., to sample model weights as pre-trained models in a single forward pass through the decoder.
To that end, we introduce layerwise loss normalization improving the quality of decoded neural network weights significantly and demonstrate that conditioning the proposed sampling methods on particular properties of the topology of the hyper-representation makes a difference for transfer learning.
Our results show that the proposed approach is potent enough to be used as another form of pre-training able to out-perform conventional baselines in transfer learning.
Previous work on generating model weights proposed (Graph) HyperNetworks (Ha et al., 2016; Zhang et al., 2020;
Knyazev et al., 2021), Bayesian HyperNetworks (Deutsch, 2018), HyperGANs (Ratzlaff and Fuxin, 2019) and
HyperTransformers (Zhmoginov et al., 2022) for neural architecture search, model compression, ensembling, transferor meta-learning.
These methods learn representations from the images and labels of the target domain.
In contrast, our approach only uses model weights and does not need access to the underlying data samples and labels rendering it more compact to the original data.
In addition to the ability to generate novel and diverse model weights, compared to previous works our approach (a) can generate novel weights conditionally on model zoos for unseen datasets and (b) can be conditioned on the latent factors of the underlying hyper-representation.
Notably, both (a) and (b) can be done without the need to retrain hyper-representations.
The results suggest our approach (Figure 1) to be a promising step towards the use of hyper-representation as generative models able to encapsulate knowledge of model zoos for transfer learning.
2. Hyper-Representation Training
The first stage of our method that corresponds to learning a hyper-representation of a population of neural networks, called a model zoo (Sch ¨ urholt et al., 2021).
In this context, a model zoo consists of models of the same architecture trained on the same task such as CIFARimage classification (Krizhevsky, 2009).
Specifically, a hyper-representation is learned using an autoencoder ˆ w i
= h ( g ( w i )) on a zoo of M models { w i } M 1 , where w i is the flattened vector of dimension N of all the weights of the i -th model.
The encoder g compresses vector w i to fixedsize hyper-representation z i
= g ( w i ) of lower dimension.
The decoder h decompresses the hyper-representation to the reconstructed vector ˆ w i .
Both encoder and decoder are built on a self-attention blocks.
The samples from model zoos are understood as sequences of convolutional or fully connected neurons.
Each of the these is encoded as a token embedding and concatenated to a sequence.
The sequence is passed through several layers of multi-head self-attention.
Afterwards, a special compression token summarizing the entire sequence is linearly compressed to the bottleneck.
The output is fed through a tanh-activation to achieve a bounded latent space
z i for the hyper-representation.
The decoder is symmetric to the encoder, the embeddings are linearly decompressed from hyper-representations z i and position encodings added.
Training is done in a multi-task fashion, minimizing the composite loss L = [Equation] +
[Equation] c , where L c is a contrastive loss and L MSE is a weight reconstruction loss (see details in (Sch¨urholt et al., 2021)).
We can write the latter in a layer-wise way to facilitate our discussion in § 3.1:
L MSE =
[Equation]
MN (cid:88) M i =
[Equation]
(cid:88) L
[Equation]
where ˆ w ( l ) i , w ( l ) i are reconstructed and original weights for the l -th layer of the i -th model in the zoo; L is the total number of layers.
The contrastive loss L c leverages two types of data augmentation at train time to impose structure on the latent space: permutation exploiting inherent symmetries of the weight space and random erasing.
3. Methods 3.1. Layer-Wise Loss Normalization
We observed that hyper-representations as proposed by (Sch ¨ urholt et al., 2021) decode to dysfunctional models, with performance around random guessing.
To alleviate that, we propose a novel layer-wise loss normalization (LWLN), which we motivate and detail in the following.
Due to the MSE training loss, the reconstruction error can generally be expected to be uniformly distributed over all weights and layers of the weight vector w .
However, the weight magnitudes of many of our zoos are unevenly distributed across different layers.
In these zoos, the even distribution of reconstruction errors lead to undesired effects.
Layers with broader distributions and large weights are reconstructed well, while layers with narrow distributions and small weights are disregarded.
The latter can become a weak link in the reconstructed models, causing performance to drop significantly down to random guessing.
The top row of Figure 2 shows an example of a baseline hyperrepresentation learned on the zoo of SVHN models (Netzer et al., 2011).
Common initialization schemes (He et al., 2015; Glorot and Bengio, 2010) produce distributions with different scaling factors per layer, so the issue is not an artifact of the zoos, but can exist in real world model populations.
Similarly, recent work on generating models normalizes weights to boost performance (Knyazev et al., 2021).
In order to achieve equally accurate reconstruction across the layers, we introduce a layer-wise loss normalization (LWLN) with the mean µ l and standard deviation σ l of all weights in layer l estimated over the train split of the zoo:
[Equation]
3.2. Sampling from Hyper-Representations
We introduce methods to draw diverse and high-quality samples z ∗ ∼ p ( z ) from the learned hyper-representation space to generate model weights w ∗ =
h ( z ∗ ) .
Assuming the smoothness and robustness of the learned hyperrepresentation space, we follow (Liu et al., 2019; Guo et al., 2019; Ghosh et al., 2020) in estimating the density to draw samples from a regularized autoencoder.
We propose to draw samples from the density of the train set in representation space.
We use the embeddings of the train set as anchor samples { z i } to model the density.
The dimensionality D of hyper-representations z in (Sch ¨ urholt et al., 2021), as well as in our work, is relatively high due to the challenge of compressing weights w .
We make a conditional independence assumption to facilitate sampling: p ( z ( j ) | z ( k ) , w ) =
p ( z ( j ) | w ) , where z ( j ) is the j -th dimensionality of the embedding z .
To model the distribution of each j -th dimensionality, we choose kernel density estimation (KDE), as it is a powerful yet simple, non-parametric and deterministic method with a single hyperparameter.
We fit a KDE to the M anchor samples { z ( j ) i } Mi =1 of each dimension j , and draw samples z ( j ) from that distribution:
z ( j ) ∼ p
( z ( j ) ) =
1 Mh (cid:80) Mi =1 K ( z ( j ) − z
( j ) i h ) , where K ( x ) =
(2 π ) − 1 / 2 exp ( − x 2 2 ) is the Gaussian kernel and h is a bandwidth hyperparameter.
The samples of each dimension z ( j ) are concatenated to form samples z ∗ =
[ z (1) , z (2) , · · · , z ( D ) ] .
This method is denoted as S KDE .
We observe that many anchor samples from { z i } correspond to the models with relatively poor accuracy (Figure 2). To improve the quality of sampled weights, we consider the variants of S KDE using only those embeddings of training samples corresponding to the top 30% performing models.
We denote that sampling method as S KDE30 .
4. Experiments 4.1. Experimental Setup
We train and evaluate our approaches on four image classification datasets: MNIST (Lecun et al., 1998), SVHN (Netzer et al., 2011), CIFAR-10 (Krizhevsky, 2009) and
STL10 (Coates et al., 2011).
For each dataset, there is a model zoo that we use to train an autoencoder.
Model zoos: In each image dataset, a zoo contains
M = 1000 convolutional networks of the same architecture with three convolutional layers and two fully-connected layers ( L = 5 ).
Varying only in the random seeds, all models of the zoo are trained for 50 epochs with the same hyperparameters following (Sch¨urholt et al., 2021).
To integrate higher diversity in the zoo, initial weights are uniformly sampled from a wider range of values rather than using welltuned initializations of (Glorot and Bengio, 2010;
He et al., 2015).
Each zoo is split in the train (70%), validation (15%) and test (15%) splits.
To incorporate the learning dynamics, we train autoencoders on the models trained for 21epochs following (Sch ¨ urholt et al., 2021).
Here the models have already achieved high performance, but have not fully converged.
The development in the remaining epochs of each model is treated as hold-out data to compare against.
We use the MNIST and SVHN zoos from (Sch¨urholt et al., 2021) and based on them create the CIFAR-10 and STLzoos.
Details on the zoos can be found in Appendix A.
Self-Supervised Hyper-Representation Training: We train separate hyper-representations on each of the model zoos and use the checkpoint with lowest reconstruction error on the validation set.
Using the proposed sampling methods (§ 3.2), we generate new embeddings and decode them to weights.
We evaluate sampled populations as initializations (epoch 0) and by fine-tuning for up to 25 epochs.
We distinguish between in-dataset and transfer-learning.
For in-dataset, the same image dataset is used for training and evaluating our hyper-representations and baselines.
For transfer-learning, pre-trained models B F and the hyperrepresentation model are trained on a source dataset.
Subsequently, the pre-trained models B F and the samples S KDE are fine-tuned on the target domain.
The baseline ( B T ) contains models trained from scratch on the target domain.
Baselines: As the first baseline, we consider the autoencoder of (Sch¨urholt et al., 2021), which is same as ours but without the proposed layer-wise loss-normalization (LWLN, § 3.1).
We combine this autoencoder with the S KDE30 sampling method and, hence, denote it as B KDE 30 .
We consider two other baselines based on training models with stochastic gradient descent (SGD): training from scratch on the target classification task B T , and training on a source followed by fine-tuning on the target task B F .
The latter remains one of the strongest transfer learning baselines (Chen et al., 2019;
[Table]
Dhillon et al., 2019; Kolesnikov et al., 2020).
4.2. Results
Evaluation of layer-wise loss normalization: We compare
S KDE30 that is based on our autoencoder with layer-wise loss normalization (LWLN) to the baseline autoencoder using the same sampling method ( B KDE 30 ) without fine-tuning.
On all datasets except for MNIST, S KDE30 considerably outperform B KDE 30 with the latter performing just above 10% (random guessing), see Table 1 (rows with epoch 0).
We attribute the success of LWLN to two main factors.
First,
LWLN prevents the collapse of reconstruction to the mean (compare Figure 2-top to Figure 2-bottom).
Second, by fixing the weak links, the reconstructed models perform significantly better.
We also evaluated input-to-output feature normalization, s.t. encoder and decoder operate on normalized weights, but empirically found it did not work as well as a normalization just for the loss.
Sampling for in-dataset fine-tuning: When fine-tuning, our S KDE 30 and baseline B KDE 30 appear to gradually converge to similar performance (Table 1).
While unfortunate, this result aligns well with previous findings that longer training and enough data make initialization less important (Mishkin and Matas, 2015;
He et al., 2019; Rasmus et al., 2015).
Comparing S KDE and S KDE 30 , we observe that conditioning the samples on the better models in the zoo improves the performance.
We also compare S KDE and S KDE 30 to training models from scratch ( B T ).
On all four datasets, both ours and the baseline hyper-representations outperform B T when generated weights are fine-tuned for the same number of epochs as B T .
Notably, on MNIST and SVHN generated weights fine-tuned for 25 epochs are even better than B T run for 50 epochs.
The comparison to B T trained for 50 epochs on the image dataset is interesting, since the hyper-representations were trained on model weights trained for up to 25 epochs, and so their overall training epochs on the image dataset is equal.
On CIFARand STL-10, all populations are limited by the architecture and saturate below 50 and 40 % accuracy.
These findings show that the models initialized with generated weights can learn faster and in some cases achieve higher performance in 25 epochs than B T in 50 epochs.
Sampling for Cross-dataset Initialization We investigate the effectiveness of our method in a transfer-learning setup across image datasets.
Here, a zoo is trained on a source dataset, e.g., SVHN.
A hyper-representation is trained on that zoo and models are generated from it.
These models are transferred to a target dataset, e.g., MNIST.
We report transfer learning results from SVHN to MNIST and from STL-10 to CIFAR-10 as two representative scenarios.
Results on all datasets can be found in Appendix B.
Note that for STL-10 to CIFAR-10 the performance of all methods saturate quickly due to the limited capacity of models in the zoo making further improvements challenging
[Table]
In transfer learning from SVHN to MNIST, the sampled populations on average learn faster and achieve significantly higher performance than the B T baseline and generally compares favorably to B F (Table 2).
In the STL-10 to
CIFAR-10 experiment, all populations appear to saturate with only small differences in their performances (Table 2).
We found that all datasets are useful sources for all targets (see Appendix B).
This might be explained by the ability of hyper-representations to capture a generic inductive prior useful across different domains.
5. Conclusion
In this paper, we propose to use hyper-representations as pre-training for transfer learning.
We extend the training objective of hyper-representations by a novel layer-wise loss normalization which is key to the capability of generating functional models.
Our method allows us to generate populations of model weights in a single forward pass.
We evaluate sampled models both in-dataset as well as in transfer learning and find them capable to outperform both models trained from scratch, as well as pre-trained and fine-tuned models.
Our work might serve as a building block for transfer learning from different domains, meta learning or continual learning.
A. Model Zoo Details
All model zoos share one general CNN architecture, outlined in Table A.
The hyperparameter choices for each of the population are listed in Table A.
The hyperparameters are chosen to generate zoos with smooth, continuous development and spread in performance.
The models in the zoos trained on MNIST, SVHN and USPS have one input channel andlearnable parameters.
Due to the three input channels in the CIFAR-10 and STL-10 zoos, their models have 2864 learnable parameters.
[Table]
[Table]
B. Cross Dataset Initialization Results
In Table 5 we report the remaining results on cross dataset initialization.
On MNIST to SVHN, the sampled population outperforms both baselines.
On CIFAR-10 to STL-10, much like the other way around, all populations saturate.
[Table]
Limitations of Zoos with Small Models To thoroughly investigate different methods and make experiments feasible, we chose to use the model zoos of the same scale as in (Sch ¨ urholt et al., 2021).
While on MNIST and SVHN, the architectures of such model zoos allowed us to achieve a reasonably high performance, on CIFAR-10 and STL-10, the performance of all populations is limited by the low capacity of the models architecture.
The models saturate at around 50% and 40% accuracy, respectively.
We hypothesize that due to the high remaining loss, the weight updates are correspondingly large without converging or improving performance.
This may cause the weights to contain relatively little signal and high noise.
Indeed, learning hyper-representations on the CIFAR-10 and STL-10 zoos was difficult and never reached similar reconstruction performance as in the MNIST and SVHN zoos.
That in turn additionally limits the performance of the sampled populations.
Future work will therefore focus on architectures with higher capacity.