Analyzing Populations of Neural Networks via Dynamical Model Embedding

1 INTRODUCTION . A crucial feature of neural networks with fixed network architecture is that they form a manifold by virtue of their continuously tunable weights , which underlies their capability to be trained by gradient descent . However , this conception of the space of neural networks is inadequate for understanding the computational processes they perform . For example , two neural networks trained to perform the same task may have vastly different weights , and yet implement the same high-level algorithms and computational processes ( Maheswaranathan et al. , 2019b ) . In this paper , we construct an algorithm which provides alternative parametrizations of the space of RNNs and CNNs with the goal of endowing a geometric structure that is more compatible with the high-level computational processes performed by neural networks . In particular , given a set of neural networks with possibly different architectures ( and possibly trained on different tasks ) , we find a parametrization of a low-dimensional submanifold of neural networks which approximately interpolates between these chosen “ base models ” , as well as extrapolates beyond them . We can use such model embedding spaces to cluster neural networks and even compute model averages of neural networks . A key feature is that two points in model embedding space are nearby if they correspond to neural networks which implement similar high-level computational processes , in a manner to be described later . In this way , two neural networks may correspond to nearby points in model embedding space even if those neural networks have distinct weights or even architectures . The model embedding space is parametrized by a low-dimensional parameter θ ∈ Rd , and each base model is assigned a value of θ in the space . This allows us to apply traditional ideas from clustering and interpolation to the space of neural networks . Moreover , each model embedding space has an associated meta-model which , upon being given a θ , is rendered into an operable neural network . If a base model is mapped to some θ , then the meta-model , upon being given that θ , will emulate the corresponding base model . See Figure 1 for a diagrammatic depiction . An interesting application is that given two models assigned to parameters θ1 and θ2 , we can consider the averaged model corresponding to the value ( θ1 + θ2 ) /2 . We find that this averaged model performs similarly on the task for which the base models in the average were trained . We also use the model embedding space to extrapolate outside the space of base models , and find cases in which the model embedding manifold specifies models that perform better on a task than any trained base model . Later on we will explain how this can be regarded as a form of semi-supervised learning . The rest of the paper is organized as follows . We first provide the mathematical setup for our algorithmic construction of model embedding spaces . After reviewing related work , we then present results of numerical experiments which implement the algorithm and explore clustering , model averaging , and semi-supervised learning on the model embedding space . We further examine how topological features of the dynamics of RNNs in model embedding spaces are reflective of classes of high-level computation processes . Finally , we conclude with a discussion . 2 DYNAMICAL MODEL EMBEDDING . 2.1 MATHEMATICAL SETUP . In this section , we provide the mathematical setup for construction of model embedding spaces . We treat this in the RNN setting , and relegate the CNN setting to Section A of the Appendix . Algorithm 1 DYNAMO algorithm Input : Base models { ( Fn , Gn ) } Nn=1 , loss functions { ( Ln , hidden , Ln , output ) } Nn=1 , datasets { Dn } Nn=1 , output loss weight λ ≥ 0 Output : Meta-model ( F̃ , G̃ ) , state maps { Vn } Nn=1 , model embeddings { θn } Nn=1 Initialize networks F̃ , G̃ , { Vn } Nn=1 Initialize model embeddings θn = 0 for τ = 1 , . . . , τmax do Sample base model ( Fi , Gi ) uniformly at random Sample input { xt } from dataset Di Compute base model hidden states { ht } and outputs { Gi ( ht ) } Compute meta-model hidden states { hθi , t } and outputs { G̃ ( hθi , t ) } Compute mapped meta-model states { Vi ( hθi , t ) } Compute loss L̂ = Li , hidden + λLi , output Update F̃ , G̃ , Vi , θi using∇L̂ end for return ( F̃ , G̃ ) , { Vn } , { θn } Notation . Let an RNN be denoted by F ( x , h ) where x is the input and h is the hidden state . We further denote the hidden-to-output map byG ( h ) . We consider a collection ofN RNNs { ( Fn , Gn ) } Nn=1 we call the base models which may each have distinct dimensions for their hidden states , but all have the same dimension for their inputs as well as the same dimension for their outputs . A sequence of inputs is notated as { xt } Tt=0 which induces a sequence of hidden states by ht+1 = F ( xt , ht ) where the initial hidden state h0 is given . A collection of sequences of inputs , possibly each with different maximum lengths T , is denoted by D which we call an input data set . We suppose that each base model RNN has an associated input data set . 2.2 META-MODELS FOR RNNS . Given a collection base model RNNs , we would like to construct a meta-model which emulates the behavior of each of the base models . In this case , the meta-model is itself an RNN with one additional input θ ∈ Rd and a corresponding map F̃ ( θ , x , h ) whose output is the next hidden state . Given a sequence of input states { xt } Tt=0 , we have a corresponding sequence of output states hθ , t+1 = F̃ ( θ , xt , hθ , t ) starting from an initial hidden state h0 ( which we suppose does not depend on θ ) . The meta-model also includes a hidden-to-output map G̃ ( h ) that is independent of θ . For the meta-model ( F̃ , G̃ ) to emulate a particular base model ( Fn , Gn ) with respect to its corresponding data set Dn , we consider the following criteria : there is some θn for which 1 . G̃ ( hθn , t ) ≈ Gn ( ht ) for all t > 0 and all input sequences in the data set 2 . Vn ( hθ , t ) ≈ ht for all t > 0 and all input sequences in the data set , where Vn is a transformation of a meta-model ’ s hidden activity . We emphasize that θn and Vn depend on the particular base model under consideration . The first criterion means that at some particular θn , the outputs of the meta-model RNN dynamics are close to the outputs of the base model RNN dynamics . The second criterion means that at the same θn , there is a time-independent transformation Vn ( i.e. , Vn does not depend on t ) such the transformed hidden state dynamics of the meta-model are close to the hidden state dynamics of the base model . See Figure 2 for a visualization . As depicted in the figure , it is convenient to regard the meta-model RNN as having input states ( θn , x ) . As such , a sequence of inputs { xt } Tt=1 is appended by θn to become { ( θn , xt ) } Tt=1 . The desired properties of the meta-model are enforced in the loss function . Defining the functions Loutput [ F̃ , G̃ , θn ] : = 1 T T∑ t=1 d ( G̃ ( hθn , t ) , Gn ( ht ) ) ( 1 ) Lhidden [ F̃ , θn , Vn ] : = 1 T T∑ t=1 ‖Vn ( hθn , t ) − ht‖ 2 2 ( 2 ) where d is some suitable distance or divergence , we can construct the loss function E { xt } ∼Dn [ Lhidden [ F̃ , θn , Vn ] + λLoutput [ F̃ , G̃ , θn ] ] ( 3 ) where we average over the choice of sequence { xt } coming from the input data set Dn . Above , λ is a hyperparameter . Our aim is to minimize the above over a suitable class of F̃ , G̃ , Vn , as well as θn ; this can be implemented computationally via the DYNAMO algorithm ( see Algorithm 1 ) . As a side remark , it appears that a suitable alternative choice to equation 2 would be 1 T ∑T t=1 ‖hθn , t −Wn ( ht ) ‖ 2 2 where here Wn is a map from the hidden states of the base model to the hidden states of the meta-model . However , this would be problematic since minimization may pressure Wn to be the the zero map ( or otherwise have outputs which are small in norm ) and accordingly pressure the dynamics of the meta-model to be the trivial ( or have small norm ) . As such , we opt to formulate Lhidden as it is written in equation 2 . Suppose we want the meta-model to be able to emulate an entire collection of base models { ( Fn , Gn ) } Nn=1 . In particular , the meta-model will attempt to assign to the nth base model a θn and Vn so that the two criteria listed above are satisfies for that base model . These desiderata can be implemented by minimizing the loss function L [ F̃ , G̃ , { θn } Nn=1 , { Vn } Nn=1 ] : = 1 N N∑ n=1 E { xt } ∼Dn [ Ln , hidden [ F̃ , θn , Vn ] + λLn , output [ F̃ , G̃ , θn ] ] . ( 4 ) In some circumstances , we may want to consider base models with distinct dimensions for their output states . For instance , suppose half of the base models perform a task with outputs in Rd1 and the other half perform a task without outputs in Rd2 . To accommodate for this , we can have two hidden-to-output maps G̃1 , G̃2 for the meta-model , where the maps have outputs in Rd1 and Rd2 respectively . The loss function is slightly modified so that we use G̃1 where compare the metamodel to the first kind of base model , and G̃2 when we compare the meta-model to the second kind of base model . This construction generalizes to the setting where the base models can be divided up into k groups with distinct output dimensions ; this would necessitate k hidden-to-output functions G̃1 , ... , G̃k for the meta-model . 3 RELATED WORK . There is a substantial body of work on interpreting the computational processes implemented neural networks by studying their intermediate representations . Such analyses can be performed either on individual models ( Simonyan et al. , 2013 ; Zeiler & Fergus , 2014 ; Lenc & Vedaldi , 2015 ) or on collections of models ( Li et al. , 2016 ; Raghu et al. , 2017 ; Morcos et al. , 2018 ; Kornblith et al. , 2019 ) . Prior work in this latter category has focused on pairwise comparisons between models . For example , SVCCA ( Raghu et al. , 2017 ) uses canonical correlation analysis to measure the representational similarity between pairs of models . While these methods can also be used to derive model representations by embedding the pairwise distance matrix , our approach does not require the Θ ( N2 ) computational cost of comparing all pairs of base models . Moreover , DYNAMO yields an executable meta-model that can be called with model embedding vectors other than those corresponding to the base models . There is a related body of work in the field of computational neuroscience . CCA based techniques and representational geometry are standard for comparing neural networks to the neural activations of animals performing vision tasks ( Yamins & DiCarlo , 2016 ) as well as motor tasks . In an example of the latter , the authors of ( Sussillo et al. , 2015 ) used CCA techniques to compare brain recordings to those of neural networks trained to reproduce the reaching behaviors of animals , while the authors of ( Maheswaranathan et al. , 2019b ) used fixed point analyses of RNNs to consider network similarity from a topological point of view . DYNAMO can be viewed as a form of knowledge distillation ( Hinton et al. , 2015 ) , since the outputs of the base models serve as targets in the optimization of the meta-model . However , unlike typical instances of knowledge distillation involving an ensemble of teacher networks ( Hinton et al. , 2015 ; Fukuda et al. , 2017 ) where individual model predictions are averaged to provide more accurate target labels for the student , our approach instead aims to preserve the dynamics and outputs of each individual base model . FitNets ( Romero , 2015 ) employ a form a knowledge distillation using maps between hidden representations to guide learning ; this is similar to our use of hidden state maps in training the meta-model . Our treatment of the model embedding vectors { θn } as learnable parameters is similar to the approach used in Generative Latent Optimization ( Bojanowski et al. , 2018 ) , which jointly optimizes the image generator network and the latent vectors corresponding to each image in the training set . Bojanowski et al . ( 2018 ) find that the principal components of the image representation space found by GLO are semantically meaningful . We likewise find that the principal components of the model embeddings found by DYNAMO are discriminative between subsets of models . Unlike methods such as hypernetworks ( Ha et al. , 2017 ) and LEO ( Rusu et al. , 2019 ) that use a model to generate parameters for a separate network , our approach does not attempt to reproduce the parameters of the base models in the collection . Instead , the meta-model aims to reproduce only the hidden states and outputs of a base model when conditioned on the corresponding embedding vector . The core focus of this work also differs from that of the meta-learning literature ( Santoro et al. , 2016 ; Ravi & Larochelle , 2017 ; Finn et al. , 2017 ; Munkhdalai & Yu , 2017 ) , which primarily concerns the problem of few-shot adaptation when presented with data from a new task . Our empirical study centers on the post-hoc analysis of a given collection of models , which may or may not have been trained on different tasks . However , we remark that our exploration of optimization in lowdimensional model embedding space is related to LEO ( Rusu et al. , 2019 ) , where a compressed model representation is leveraged for efficient meta-learning .

Summary: The manuscript proposes to analyze the similarity of different neural networks trained to solve the same task by developing a meta-model capable of reproducing/emulating both the output and the hidden representation of each individual network.  This meta-model is parametrized by a set of parameters which is common to all the networks, and by other parameters which are network-specific (theta). In this way, each network will correspond to a different theta, and one can analyze the global features of the set of networks by unsupervised manifold learning in the  theta-space.