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: This paper proposes an approach to combine together multiple (pretrained) neural networks into a single model that can simulate each of the considered neural networks, both in terms of output and hidden states. The setup here relies on all the considered networks to operate on the same task, and in particular the same input data type (but perhaps with different sampled training sets). The combined meta-model then considers as inputs not only the input data, but also an encoding (which is learned during its training) of each model. This allows model-dependent latent features to be computed by the meta-model, which can then be decoded to hidden states of each network via specialized decoders that can be thought of as auxiliary computational pathways that operate in parallel to the one producing the main output of the meta-model. The model encoding are then used to visualize and explore the relations between different networks (or instantiations), and furthermore, to further tune the combined model towards an interpolation or extrapolation that goes beyond (and seems to outperform) each of the individual networks used to train the meta-model.