Self-Supervised Representation Learning on Neural Network Weights for Model Characteristic Prediction

1 Introduction . This work investigates populations of Neural Network ( NN ) models and aims to learn representations of them using Self-Supervised Learning . Within NN populations , not all model training is successful , i.e. , some overfit and others generalize . This may be due to the non-convexity of the loss surface during optimization [ Goodfellow et al. , 2015 ] , the high dimensionality of the optimization space , or the sensitivity to hyperparameters [ Hanin and Rolnick , 2018 ] , which causes models to converge to different regions in weight space . What is still not yet fully understood , is how different regions in weight space are related to model characteristics . Previous work has made progress investigating characteristics of NN models , e.g by visualizing learned features [ Zeiler and Fergus , 2014 , Karpathy et al. , 2015 ] . Another line of work compares the activations of pairs of NN models [ Raghu et al. , 2017 , Morcos et al. , 2018 , Kornblith et al. , 2019 ] . Both approaches rely on the expressiveness of the data , and are , in the latter case , limited to two models at a time . Other approaches predict model properties , such as accuracy , generalization gap , or hyperparameters from the margin distribution [ Yak et al. , 2019 , Jiang et al. , 2019 ] , graph topology features [ Corneanu et al. , 2020 ] or eigenvalue decomposition of the weight matrices [ Martin and Mahoney , 2019 ] . In a similar direction , other publications propose to investigate populations of models and to predict properties directly from their weights or weight statistics in a supervised way [ Unterthiner et al. , 2020 , Eilertsen et al. , 2020 ] . However , these manually designed features may not fully capture the latent model characteristics embedded in the weight space . 1https : //github.com/HSG-AIML/NeurIPS_2021-Weight_Space_Learning 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . Therefore , our goal is to learn task-agnostic representations from populations of NN models able to reveal such characteristics . Self-Supervised Learning ( SSL ) is able to reveal latent structure in complex data without the need of labels , e.g. , by compressing and reconstructing data [ Goodfellow et al. , 2016 , Kingma and Welling , 2014 ] . Recently , a specific approach to SSL called contrastive learning has gained popularity [ Misra and van der Maaten , 2019 , Chen et al. , 2020a , b , Grill et al. , 2020 ] . Contrastive learning leverages inherent symmetries and equivariances in the data , allows to encode inductive biases and thus structure the learned representations . In this paper , we propose a novel approach to apply SSL to learn representations of the weights of NN populations . We learn representations using reconstruction , contrast , and a combination of both . To that end , we adapt a transformer architecture to NN weights . Further , we propose three novel data augmentations for the domain of NN weights . We introduce structure preserving permutations of NN weights as augmentations , which make use of the structural symmetries within the NN weights that we find necessary for learning generalizing representations . We also adapt erasing [ Zhong et al. , 2020 ] and noise [ Goodfellow et al. , 2016 ] as augmentations for NN weights . We evaluate the learned representations by linear-probing for the generating factors and characteristics of the models . An overview for our learning approach is given in Figure 1 . To validate our approach , we perform extensive numerical experiments over different populations of trained NN models . We find that hyper-representations2 can be learned and reveal the characteristics of the model zoo . We show that hyper-representations have high utility on tasks such as the prediction of hyper-parameters , test accuracy , and generalization gap . This empirically confirms our hypothesis on meaningful structures formed by NN populations . Furthermore , we demonstrate improved performance compared to the state-of-the-art for model characteristic prediction and outlay the advantages in out-of-distribution predictions . For our experiments , we use publicly available NN model zoos and introduce new model zoos . In contrast to [ Unterthiner et al. , 2020 , Eilertsen et al. , 2020 ] , our zoos contain models with different initialization points and diverse configurations , and include densely sampled model versions during training . Our ablation study confirms that the various factors for generating a population of trained NNs play a vital role in how and which properties are recoverable for trained NNs . In addition , the relation between generating factors and model zoo diversity reveals that seed variation for the trained NNs in the zoos is beneficial and adds another perspective when recovering NNs ’ properties . 2 Model Zoos and Augmentations . Model Zoo We denote as D a data set that contains data samples with their corresponding labels . We denote as λ the set of hyper-parameters used for training , ( e.g. , loss function , optimizer , learning rate , weight initialization , batch-size , epochs ) . We define as A the specific NN architecture . Training under different prescribed configurations { D , λ , A } results in a population of NNs which we refer to as model zoo . We convert the weights and biases of all NNs of each model into a vectorized form . In the resulting model zooW = { w1 , .... , wM } , wi denotes the flattened vector of dimension N , representing the weights and biases for one trained NN model . 2By hyper-representation , we refer to a learned representation from a population of NNs , i.e. , a model zoo , in analogy to HyperNetworks [ Ha et al. , 2016 ] , which are trained to generate weights for larger NN models . Augmentations . Data augmentation generally helps to learn robust models , both by increasing the number of training samples and preventing overfitting of strong features [ Shorten and Khoshgoftaar , 2019 ] . In contrastive learning , augmentations can be used to exploit domain-specific inductive biases , e.g. , know symmetries or equivariances [ Chen and Li , 2020 ] . To the best of our knowledge , augmentations for NN weights do not yet exist . In order to enable self-supervised representation learning , we propose three methods to augment individual instances of our model zoos . Neurons in dense layers can change position without changing the overall mapping of the network , if in-going and out-going connections are changed accordingly [ Bishop , 2006 ] . The relation between equivalent versions of the same network translates to permutations of incoming weights with matrix P and transposed permutation PT of the outgoing weights ( PTP = I ) . Considering the output at layer l + 1 , with weight matrices W , biases b , activations a and activation function σ , we have zl+1 = Wl+1σ ( Wlal−1 + bl ) + bl+1 = Ŵl+1σ ( Ŵlal−1 + b̂l ) + bl+1 , ( 1 ) where Ŵl+1 = Wl+1 ( Pl ) T , Ŵl = PlWl and b̂l = Plbl are the permuted weight matrices and bias vector , respectively . The equivalences hold not only for the forward pass , but also for the backward pass and weight update.3 The permutation can be extended to kernels of convolution layers . The permutation augmentation differs significantly from existing augmentation techniques . As an analogy , flips along the axis of images are similar , but specific instances from the set of possible permutations in the image domain . Each permutable layer with dimension Nl , has Nl ! different permutation matrices , and in total there are ∏ lNl ! distinct , but equivalent versions of the same NN . While new data can be created by training new models , the generation is computationally expensive . The permutation augmentation , however , allows to compute valid NN samples at almost no computational cost . Empirically , we found the permutation augmentation crucial for our learning approach . In computer vision and natural language processing , masking parts of the input has proven to be helpful for generalization [ Devlin et al. , 2019 ] . We adapt the approach of random erasing of sections in the vectorized forms of trained NN weights . As in [ Zhong et al. , 2020 ] , we apply the erasing augmentation with a probability p to an area that is randomly chosen with a lower and upper bounds blow and bup . In our experiments , we set p = 0.5 , blow = 0.03 , bup = 0.3 and erase with zeros . Adding noise augmentation is another way of altering the exact values of NN weights without overly affecting their mapping , and has long been used in other domains [ Goodfellow et al. , 2016 ] . 3 Hyper-Representation Learning . With this work , we propose to learn representations of structures in weight space formed by populations of NNs . We evaluate the representations by predicting model characteristics . A supervised learning approach has been demonstrated [ Unterthiner et al. , 2020 , Eilertsen et al. , 2020 ] . [ Unterthiner et al. , 2020 ] find that statistics of the weights ( mean , var and quintiles ) are superior to the weights to predict test accuracy , which we empirically confirm in our results . However , we intend to learn representations of the weights , that contain rich information beyond statistics . ‘ Labels ’ for NN models can be obtained relatively simply , yet they can only describe predefined characteristics of a model instance ( e.g. , accuracy ) and so supervised learning may overfit few features , as [ Unterthiner et al. , 2020 ] show . Self-supervised approaches , on the other hand , are designed to learn task-agnostic representations , that contain rich and diverse information and are exploitable for multiple downstream tasks [ LeCun and Misra , 2021 ] . Below , we present the used architectures and losses for the proposed self-supervised representation learning . Architectures and Self-Supervised Losses . We apply variations of an encoder-decoder architecture . We denote the encoder as gθ ( wi ) , its parameters as θ , and the hyper-representation with dimension L as zi = gθ ( wi ) . We denote the decoder as hψ ( zi ) , its parameters as ψ , and the reconstructed NN weights as ŵi = hψ ( zi ) = hψ ( gθ ( wi ) ) . As is common in CL , we apply a projection head pγ ( zi ) , with parameters γ , and denote the projected embeddings as z̄i = pγ ( zi ) = pγ ( gθ ( wi ) ) . In all of the architectures , we embed the hyper-representation zi in a low dimensional space , L < N . We employ two common SSL strategies : reconstruction and contrastive learning . Autoencoders ( AEs ) with a reconstruction loss are commonly used to learn low-dimensional representations [ Goodfellow 3Details , formal statements and proofs can be found in Appendix A. et al. , 2016 ] . As AEs aim to minimize the reconstruction error , the representations attempt to fully encode samples . Further , contrastive learning is an elegant method to leverage inductive biases of symmetries and inductive biases [ Bronstein et al. , 2021 ] . Reconstruction ( ED ) : For reconstruction , we minimize the MSE LMSE = 1M ∑M i=1 ‖wi − hψ ( gθ ( wi ) ) ‖22 . We denote the encoder-decoder with a reconstruction loss as ED . Contrast ( Ec ) : For contrastive learning , we use the common NT_Xent loss [ Chen et al. , 2020a ] as Lc . For a batch of MB model weights , each sample is randomly augmented twice to form the two views i and j . With the cosine similarity sim ( z̄i , z̄j ) = z̄Ti z̄j/||z̄i||||z̄j || , the loss is given as Lc = ∑ ( i , j ) − log exp ( sim ( z̄i , z̄j ) /T∑2MB k=1 Ik 6=i exp ( sim ( z̄i , z̄j ) /T , ( 2 ) where Ik 6=i is 1 if k 6= i and 0 otherwise , and T is the temperature parameter . We denote the encoder with a contrastive loss as Ec . Reconstruction + Contrast ( EcD ) : Further we combine reconstruction and contrast via L = βLMSE + ( 1 − β ) Lc in order to achieve good quality compression via reconstruction and wellstructured representations via the contrast . We denote this architecture with its loss as EcD . Reconstruction + Positive Contrast ( Ec+D ) : In contrastive learning , many methods prevented mode collapse by using negative samples . The combined loss contains a reconstruction term LMSE , which can be seen as a regularizer that prevents mode collapse . Therefore , we also experiment with replacing Lc in our loss with a modified contrastive term without negative samples : Lc+ = ∑ i − log ( exp ( sim ( z̄ji , z̄ k i ) ) /T ) = ∑ i −sim ( z̄ji , z̄ k i ) + log ( T ) . ( 3 ) [ Chen and He , 2021 , Schwarzer et al. , 2021 ] explore similar simplifications without reconstruction . We denote the encoder-decoder with the loss L = βLMSE + ( 1− β ) Lc+ as Ec+D . Attention Module . Our encoder and decoder pairs are symmetrical and of the same type . As there is no intuition on good inductive biases in the weight space , we apply fully connected feed-forward networks ( FFN ) as baselines . Further , wee use multi-head self-attention modules ( Att ) [ Vaswani et al. , 2017 ] as an architecture with very little inductive bias . In the multi-head self-attention module , we apply learned position encodings to preserve structural information [ Dosovitskiy et al. , 2020 ] . The explicit combination of value and position makes attention modules ideal candidates to resolve the permutation symmetries of NN weight spaces . We propose two methods to encode the weights into a sequence ( Figure 2 ) . In the first method , we encode each weight as a token in the input sequence . In the second method , we linearly transform the weights of one neuron or kernel and use it as a token . Further , we apply two variants to compress representations in the latent space ( Figure 3 ) . In the first variant , we aggregate the output sequence of the transformer and linearly compress it to a hyper-representation zi . In the second variant , similarly to [ Devlin et al. , 2019 , Zhong et al. , 2020 ] , we add a learned token to the input sequence that we dub compression token . After passing the input sequence trough the transformer , only the compression token from the output sequence is linearly compressed to a hyper-representation zi . Without the compression token , the information is distributed across the output sequence . In contrast , the compression token is learned as an effective query to aggregate the relevant information from the other tokens , similar to [ Jaegle et al. , 2021 ] . The capacity of the compression token can be an information bottleneck . Its dimensionality is directly tied to the dimension of the value tokens and so its capacity affects the overall memory consumption . TETRIS-SEED Rec . Eph Acc F1C0 F1C1 F1C2 F1C3 Ec - 96.7 90.8 67.7 72.0 74.4 85.8 ED 96.1 88.3 68.9 47.8 57.2 33.0 58.1 EcD 84.1 97.0 90.2 70.7 75.9 69.4 86.6 Ec+D 95.9 94.0 69.9 48.9 58.1 32.5 58.8 Table 1 : Ablation results over self-supervised learning losses . All models implemented with attention-based reference architecture . All values are given in % . TETRIS-SEED Rec . Eph Acc F1C0 F1C1 F1C2 F1C3 FF 0.0 80.0 85.3 57.3 53.9 64.5 80.1 AttW 6.8 95.3 71.1 47.9 69.0 49.9 61.3 AttW+t 74.1 95.4 88.6 65.6 69.8 69.9 86.8 AttN 89.4 97.1 88.4 71.4 80.8 69.1 82.3 AttN+t 84.1 97.0 90.2 70.7 75.9 69.4 86.6 Table 2 : Ablation results in % under EcD setup . We use feed-forward FF and attention-based variants with weight and neuron encoding AttW and AttN each with +t and without compress . token . Downstream Tasks . We use linear probing [ Grill et al. , 2020 ] as a proxy to evaluate the utility of the learned hyper-representations . As downstream tasks we use accuracy prediction ( Acc ) , generalization gap ( GGap ) , epoch prediction ( Eph ) as proxy to model versioning , F1-Score prediction ( F1C ) , learning rate ( LR ) , ` 2-regularization ( ` 2-reg ) , dropout ( Drop ) and training data fraction ( TF ) . Using such targets , we solve a regression problem and measure the R2 score [ Wright , 1921 ] . We also evaluate for hyper-parameters prediction tasks , like the activation function ( Act ) , optimizer ( Opt ) , initialization method ( Init ) . Here , we train a linear perceptron by minimizing a cross entropy loss [ Goodfellow et al. , 2016 ] and measure the prediction accuracy .

Summary: This paper proposes a self-supervised representation learning approach to uncover properties from trained neural network (NN) weights and biases. Representations are learnt through encoder-decoder architectures on the weights, using various losses. Attention-based modules are used to find interactional structure between different weights of a NN, resulting in compressed representations in different manners. To this end, data augmentation based on permutation augmentation is proposed; random erasing and noise augmentation are also explored. Empirics are performed on trained models from MNIST, Fashion-MNIST, SVHN, CIFAR-10 as well as synthetic datasets. Results explore data augmentation, architecture ablation, out-of-distribution prediction and various downstream tasks.