Revisiting Model Stitching to Compare Neural Representations

1 Introduction . The success of deep neural networks can , arguably , be attributed to the intermediate features or representations learnt by them [ Rumelhart et al. , 1985 ] . While neural networks are trained in an end-to-end fashion with no explicit constraints on their intermediate representations , there is a body of evidence that suggests that they learn rich a representation of the data along the way [ Goh et al. , 2021 , Olah et al. , 2017 ] . However , theoretically we understand very little about how to formally characterize these representations , let alone why representation learning occurs . For instance , there are various ad-hoc pretraining methods [ Chen et al. , 2020a , b ] , that are purported to perform well by learning good representations , but it is unclear which aspects of these methods ( objective , training algorithm , architecture ) are crucial for representation learning and if these methods learn qualitatively different representations at all . Moreover , we have an incomplete understanding of the relations between different representations . Are all “ good representations ” essentially the same , or is each representation “ good ” in its own unique way ? That is , even when we train good end-to-end models ( i.e. , small test loss ) , the internals of these models could potentially be very different from one another . A priori , the training process could evolve in either one of the following extreme scenarios ( see Figure 1 ) : ( 1 ) In the “ snowflakes ” scenario , training with different initialization , architectures , and objectives ( e.g. , supervised vs self-supervised ) will result in networks with very different internals , which are completely incompatible with one another . For example , even if we train two models with identical data , architecture , and task , but starting from two different initializations , we may end up at local 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . minima with very different properties ( e.g. , Liu et al . [ 2020 ] ) . If models are trained with different data ( e.g. , different samples ) , different architecture ( e.g. , different width ) , or different task ( e.g. , self-supervised vs supervised ) then they could end up being even more different from one another . ( 2 ) In the “ Anna Karenina ” scenario1 , all successful models end up learning roughly the same internal representations . For example , all models for vision tasks will have internal representation corresponding to curve detectors , and models that are better ( for example , trained on more data , are bigger , or trained for more time ) will have better curve detectors . The “ Anna Karenina ” scenario implies the following predictions : “ All roads lead to Rome : ” Successful models learned with different initializations , architectures , and tasks , should learn similar internal representations , and so if A and B are two such models , it should be possible to “ plug in ” the internals from A into B without a significant loss in performance . See Figure 2A-B . “ More is better : ” Better models trained using more data , bigger size , or more compute , should learn better versions of the same internal representations . Hence if A is a more successful model than B , it should be possible to “ plug in ” A ’ s internals to B and obtain improved performance . See Figure 2C . In this work , we revisit the empirical methodology of “ model stitching ” to test the above predictions . Initially proposed by Lenc and Vedaldi [ 2015 ] , model stitching is natural way of “ plugging in ” the bottom layers of one network into the top layers of another network , thus forming a stitched network ( however care must be taken in the way it is performed , see Section 2 ) . We show that model stitching has some unique advantages that make it more suitable for studying representations than representational similarity measures such as CKA [ Kornblith et al. , 2019 ] and SVCCA [ Raghu et al. , 2017 ] . Our work provides quantitative evidence for the intuition , shared by many practitioners , that the internals of neural networks often end up being very similar in a certain sense , even when they are trained under different settings . 1.1 Summary of Results . Model stitching as an experimental tool . We establish model stitching as a way of studying the representations of neural networks . A version of model stitching was proposed in Lenc and Vedaldi [ 2015 ] to study the equivalence between representations . In this work , we argue that the idea behind model stitching is more powerful than has been appreciated : we analyze the benefits of modelstitching over other methods to study representations , and we then use model-stitching to establish an number of intuitive properties , including new results on the properties of SGD . 1 “ All happy families are alike ; each unhappy family is unhappy in its own way. ” Leo Tolstoy , 1877. https : //en.wikipedia.org/wiki/Anna_Karenina_principle In this paper , we use model stitching in the following way . Suppose we have a neural network A ( which we ’ ll think of as the “ top model ” ) for some task with loss function L ( e.g . the CIFAR-10 or ImageNet test error ) . Let r : X → Rd be a candidate “ representation ” function , which can come from the first ( bottom-most ) layers of some “ bottom model ” B . Our intuition is that r has better quality than the first ` layers of network A if “ swapping out ” these layers with r will improve performance . “ Swapping out ” is performed by introducing an additional trainable stitching layer to r withA ( defined more formally in Section 2 ) . The stitching layers have very low capacity , and are only meant to “ align ” representations , rather than improving the model . Comparison to Representational Similarity Metrics . Much of the current work studying representations focuses on similarity metrics such as CKA . However , we argue that model stitching can be a better suited tool to study representations in various scenarios , and can give qualitatively different conclusions about the behavior of neural representations compared to these metrics . We analyze the differences between model stitching and prior similarity metrics in Section 3 . Quantitative evidence for intuitions . Using model stitching , we are able to provide formal and quantitative evidence to the intuitions mentioned above . In particular we give evidence for the “ all roads lead to rome ” intuition by showing compatibility of networks with representation that are trained using ( 1 ) different initializations , ( 2 ) different subsets of the dataset , ( 3 ) different tasks ( e.g. , self-supervised or coarse labels ) . See Figure 2A-B for results . We also show that network with different random initialization enjoy a property which we call stitching connectivity , wherein almost all minima reachable via SGD can be “ stitched ” to each other with minimal loss of accuracy . We also give evidence for the “ more is better ” intuition by showing that we can improve the performance of a network A by plugging in a representation r that was trained with ( 1 ) more data , ( 2 ) larger width , or ( 3 ) more training epochs . See example with more samples in Figure 2C . The results above are not surprising , in the sense that they confirm intuitions that practitioners might already have . However model stitching allows us to obtain quantitative and formal measures of these in a way that is not achievable by prior representation measures . Organization . We first define model stitching formally in Section 2 . We compare it with prior work on representational similarity measures in Section 3 . Then , we formally define stitching connectivity and provide experimental evidence for it in Section 4 . Finally , in Section 5 and 6 , we provide quantitative evidence for the `` all roads lead to Rome '' and `` more is better '' intuitions respectively . Related works . As mentioned above , stitching was introduced by Lenc and Vedaldi [ 2015 ] who studied equivalence of representations . They showed that certain early layers in a network trained on the Places dataset [ Zhou et al. , 2014 ] are compatible with AlexNet ( see Table 4 in Lenc and Vedaldi [ 2015 ] ) . After completing this work , we were made aware of concurrent work Csiszárik et al . [ 2021 ] that also proposes model stitching to compare neural representations . The results from this work complement ours by studying the effect of changing the stitching layer in various ways ( for instance , imposing a sparsity penalty on the stitching layer ) , and further clarifying the relationship of stitching with other representational similarity measures . In contrast , our work demonstrates stitching compatibility under different scenarios , like the comparison between supervised and self-supervised methods and experiments that show `` more is better '' . Our work on stitching connectivity is related to the work on mode connectivity [ Garipov et al. , 2018 , Freeman and Bruna , 2017 , Draxler et al. , 2018 ] . These works show that the local minima found by SGD are often connected through low-loss paths . These paths are generally non-linear , though it was shown that there are linear paths between these minima if they are identically initialized but then use different SGD noise ( order of samples ) after a certain point in training [ Frankle et al. , 2020 ] . Stitching connectivity is complementary to mode connectivity . Stitching connectivity corresponds to a discrete path ( with as many steps as layers ) but one where the intermediate steps are interpretable . We also show stitching connectivity of networks that are trained on different tasks . Finally , most of the prior work on concrete metrics for relating representations was in the context of representation similarity measures . We describe this work and compare it to ours in Section 3 . 2 Model Stitching . LetA be a neural network of some architectureA , and let r : X → Rd be a candidate “ representation ” function . We consider a family S of stitching layers which are simple ( e.g . linear 1× 1 convolutional layers for a convolutional network A ) functions mapping Rd to Rd ` where d ` is the width of A ’ s ` -th layer . Given some loss function L ( e.g. , CIFAR-10 or ImageNet test accuracy ) we define L ` ( r ; A ) = inf s∈S L ( A > ` ◦ s ◦ r ) ( 1 ) where A > ` denotes the function mapping the activations of A ’ s ` th layer to the final output , and ◦ denotes function composition . That is , L ` ( r ; A ) is the smallest loss obtained by stitching r into all but the first ` layers of A using a stitching layer from S. We define the stitching penalty of the representation r with respect to A ( as well as ` and the loss L ) as L ` ( r ; A ) − L ( A ) . If the penalty is non-negative , then we say that the representation r is at least as good as the first ` layers of A . Consider the simple case of a fully-connected network and a stitching family of linear functions . In this case , model stitching tells us if there is a way to linearly transform the representation r into that of the first ` layers of A , but only in the subspace that is relevant to the achieving low loss . In practice , we approximate the infimum in Equation 1 by using gradient methods— concretely , by optimizing the stitching-layer using the train set of the task L. The stitching penality itself is then estimated on the test set . Choosing the stitching family S : The stitching family can be chosen flexibly depending on the desired invariances between representations , as long as it is simple ( see below ) . For instance , if we choose S to be all permutations or orthogonal matrices , the stitching penalty will be invariant up to permutations or orthogonal transformations respectively . In this work , we mainly consider cases where r = B≤l for B with architecture similar to A . We restrict the stitching family per architecture A such that the composed model A > ` ◦ s ◦ r lies in A . This way the stitched model consists of layers identical to the layers of either A or B , with the exception of just one layer . For example , for convolutional networks we consider a 1 × 1 convolution and for transformers we consider a token-wise linear function between transformer blocks . We perform an ablation with kernels of different sizes in Appendix B.1 . Stitching is not learning : One concern that may arise is that if the stitching family S is sufficiently powerful , it can learn to transform any representation into any other representation . This would defeat the aim of faithfully studying the original representations . To avoid this , the family of stitchers should be chosen to be simple ( e.g. , linear ) . Nevertheless , to verify that our experiments are not in such a regime , we stitch an untrained , randomly initialized network with a fully trained network . Figure 2A-B shows that the penalty of such a stitched network is high , specially for high l. CKA has the same trend for the same networks ( See Appendix B.2 ) Model stitching as a “ happy middle ” . Model stitching can be considered as a compromise between the following extremes : ( 1 ) Direct plugging in : the most naive interpretation of “ plugging in ” r into A would be to use no stitching at all . However , even in cases where networks are identical up to a permutation of neurons , plugging in r into A would not work . ( 2 ) Full fine tuning : the other extreme interpretation is to perform full fine tuning . That is , start with the initialized network A > ` ◦ r and optimize over all choices of A . The problem with this approach is that there are so many degrees of freedom in the choices for A > ` that the resulting network could achieve strong performance regardless of the quality of r. For example , in B.3 we show that fine tuning can fail to distinguish between a trained network and a random network . ( 3 ) Linear probe : a popular way to define quality of representations is to use linear probes [ Alain and Bengio , 2016 ] . However , linear probes are not as well suited for studying the representation of early layers , which have low linear separability . In particular , the linear probe accuracy of an early layer of network A would generally always be much worse than a later layer of network B , even if A was of “ higher quality ” ( e.g. , trained with more data ) . Linear probes also don ’ t have the operational interpretation of compatibility . By choosing a trainable but low-capacity layer , model stitching “ threads the needle ” between simply plugging in , and full fine tuning , and unlike linear probes , enables the study of early layers using powerful nonlinear decoders . Experimental setup . Unless specified otherwise , the CIFAR-10 experiments are conducted on the ResNet-18 architecture ( with first layer width 64 ) and the ImageNet experiments are conducted on the ResNet-50 architecture [ He et al. , 2015 ] . The ResNets are trained with the standard hyperparameters ( See Appendix A.1 for training parameters of all base models ) . The stitching layer in the the convolutional networks consist of a 1× 1 convolutional layer with input features equal to the number of channels in r , and output features equal to the output channels of A≤l ( x ) . We add a BatchNorm ( BN ) layer before and after this convolutional layer . Note that the BN layer does not change the representation capacity of the stitching layer and only aides with optimization . We perform stitching only between ResNet blocks ( and not inside a block ) , but note that it is possible to stitch within the block as well . We use the Adam optimizer with the cosine learning rate decay and an initial learning rate of 0.001 . Full experimental details for each experiment are described in Appendix A .

Summary: The work proposes to utilize model stitching as a tool to compare the quality of learned neural representations. Authors show that model stitching based comparison can avoid many problems often observed in other representation comparison methods, such as CCA/CKA. Using this tool, under image classification setting, authors empirically show that (1) networks with different random initializations would converge to "similar" representations, (2) networks trained via different methods also have "similar" representations given these networks have close performances, and (3) networks trained with more data consistently have better quality.