Understanding the Inner-workings of Language Models Through Representation Dissimilarity
Authors: Davis Brown
Charles Godfrey
Nicholas Konz
Jonathan Tu
Henry Kvinge
Venue: EMNLP 2023 Main (2023)
Keywords: language models
interpretability
stitching
TL;DR: Using representation dissimilarity methods to understand generalization and scaling behavior in language models.

Abstract: As language models are applied to an increasing number of real-world applications, understanding their inner workings has become an important  issue in model trust, interpretability, and transparency. In this work we show that representation dissimilarity measures, which are functions that measure the extent to which two model's internal representations differ, can be a valuable tool for gaining insight into the mechanics of language models. Among our insights are: (i) an apparent asymmetry in the internal representations of model using SoLU and GeLU activation functions, (ii) evidence that dissimilarity measures can identify and locate generalization properties of models that are invisible via in-distribution test set performance, and (iii) new evaluations of how language model features vary as width and depth are increased. Our results suggest that dissimilarity measures are a promising set of tools for shedding light on the inner workings of language models.

Understanding the Inner Workings of Language Models Through Representation Dissimilarity
Davis Brown'
Charles Godfrey2.1 Nicholas Konz
Jonathan Tu Henry Kvinge
Pacifiic Northwest National Laboratory
Thhomon Reuters Labs
Duke University UUniversity of Washington
{first. Jastn
Abstract
address this by measuring the difference between (potentially high-dimensional) representations. this paper, we focus on two such functions that, while popular in computer vision, have seen lim- ited application to language models: model stitch- ing (Lenc and Vedaldi, 2014; Bansal et al., 2021) and centered kernel alignment (CKA) (Kornblith et al., 2019). Model stitching extracts features from earlier layers of model f and plugs them into the later layers of model g (possibly mediated by a small, learnable, connecting layer), and evaluates downstream performance of the resulting "stitched" model. Stitching takes a task-centric view towards representations, operating under the assumption that if two models have similar representations then these representations should be reconcilable by simple transformation to such extent that the downstream task can still be solved. On the other hand, CKA compares the statistical structure of the representations obtained in two different model- s/layers from fixed set of input datapoints, ignor- ing any relationship to performance on the task for which the models were trained.
As language models are applied to an increas- ing number of real-world applications, under- standing their inner workings has become an important issue in model trust, interpretability, and transparency. In this work we show that representation dissimilarity measures, which are functions that measure the extent to which two model's internal representations differ, can be a valuable tool for gaining insight into the mechanics of language models. Among our insights are: (i) an apparent asymmetry internal representations of model using SoLU and GeLU activation functions, (ii) evi- dence that dissimilarity measures can identify and locate generalization properties of models that are invisible via in-distribution test set per- formance, and (iii) new evaluations of how lan- guage model features vary width and depth are increased. Our results suggest that dissim- ilarity measures are a promising set of tools for shedding light on the inner workings of lan- guage models.
1 Introduction
The defining feature of deep neural networks is their capability of learning useful feature represen- tations from data in an end-to-end fashion. Perhaps ironically, one of the most pressing scientific chal- lenges in deep learning is understanding the fea- tures that these models learn. This challenge is not merely philosophical: learned features are known to impact model iiteppecbbliiyyeepllaabiii gen- cralization, and transfcrability to downstream tasks, and it can be the case that none of these effects are visible from the point of view of model perfor- mance on even carefully selected validation sets.
In this paper we make the case that dissimilar- ity measures are tool that has been underutilized in the study of language models. We support this claim through experiments that shed light on the inner workings of language models: (i) We show that stitching can be used better understand the changes to representations that result from using different nonlinear activations a model. In partic- ular, we find evidence that feeding Gaussian error linear unit (GeLU) model representations into a softmax linear unit (SoLU) model incurs smaller penalty in loss compared to feeding SoLU activa- tions into GeLU model, suggesting that models with GeLU activations may form representations that contain strictly more useful information for the training task than the representations of models using SoLU activations. (ii) We show that dissimi- larity measures can localize differences in models
When studying hidden representations in deep learning, fundamental question is whether the in- ternal representations of given pair of models are similar or not. Dissimilarity measures (Klabunde et al., 2023) are a class of functions that seek to
Work done at Pacific Northwest National Laboratory.

---

that are invisible via test set performance. Fol- lowing the experimental set-up of (Juneja al., 2022), we show that both stitching and CKA detect the difference between models which generalize to an out-of-distribution test set and models that do not generalize. (iii) Finally, we apply CKA to the Pythia networks (Biderman et al., 2023), se- quence of generative transformers of increasing width and depth, finding high degree of simi- larity between early layer features even as scale varies from 70 million emergent "block structure" previously observed in CKA unalysis of image classifiers such as ResNets and Vision Transformers, and showing that CKA identifies Pythia model (pythia-2 Seriiaa2..b-deuuupe 8b-deduped) exhibiting remarkably low levels of feature simi- larity when compared with the remaining models in the family (as well inconsistent architectural characteristics). This last finding is perhaps surpris- ing considering the consistent trend towards lower perplixity and higher performance on benchmarks with increasing model scale seen in the original Pythia evaluations (Biderman et al., 2023) - one might have expected to see underlying consistent evolution of hidden features from the perspective of CKA.
tively and evaluated at D is
f.A.A.gmll
(1)
|| AA Fl Al..m.....mm
where || |||| is the Frobenious norm. Higher CKA scores indicate more structural similarity between representations. In our experiments we use un- biased estimator of eq. (1) to calculate CKA in batches (see appendix C for details).
billion parameters, an
Model stitching reveals an asymmetry between GeLU and
interpretable-by-deess SoLU
The design of more interpretable models an area of active research, ineerpretablebyydeeee models often achieve comparable performance on down- stream tasks to their non-interpretable counterparts (Rudin, 2019). However, these models (almost by definition) have differences in their hidden layer representations that can impact downstream perfor- mance.
Many contemporary transformers use the Gaus- sian error linear unit activation function, approxi- mately GeLU(z) = x * sigmoid((17x)) The soft- max linear unit SoLU((x) = s-sofmmax(( (1r) is an ac- tivation function introduced in (Elhage et al., 2022) in an attempt to reduce neuron polysemanticity: the softmax has the effect of shrinking small and ampli- fying large neuron outputs. One of the findings of (Elhage et al., 2022) was that SoLU transformers yield comparable performance their GeLU coun- terparts on a range of downstream tasks. However, those tests involved zero-shot evaluation or fine- tuning of the full transformer, and as such they do not shed much light on intermediate hidden feature representations. We use stitching to do this.
2 BBckkround
In this section we review the two model dissimilar- ity measures appearing in this paper.
Model stitching (Bansal al., 2021): Infor- mally, model stitching asks how well the represen- tation extracted by the early layers of one model can be used by the later layers of another model to solve specific task. Let f be a neural network and for a layer of f let fgr (respectively f2() be the composition of the first / layers of (respectively the layers m of f with > l). Given another network g the model obtained by stitching layer of to layer m of g with stitching layer yp is I>m o up fei- The performance of this stitched network measures the similarity of representations of ff at layer and representations of g at layer m.
Using the notation from Section 2, we set our stitching layer yp to be learnable linear layer (all other parameters of the stitched model are frozen) between the residual streams following a layer We stitch small 3-layer (9.4M parameter) and 4- layer (13M parameter) SoLU and GeLU models trained by (Nanda, 2022), using the Pile validation set (Gao et al., 2020) to optimize s.
Figure displays the resulting stitching losses calculated on the Pile validation set. For both the 3-layer and 4-layer models, and every layer con- sidered, we see that when f is a SoLU model for the stitched model 9>m o oypo f&i - the "Sol.U-into- GeLU" cases larger penalties are incurred than when f uses GeLUs, i.e. the "GeLUUinto-SoLUU"
Centered kernel alignment (CKA) (Korn- blith et al., 2019): Let D [x1,..., be set of model inputs. For models f and g with layers I and m respectively, let A f.f.g.m m be the covariance matrix of DD) and (gm(D), Then the CKA score for models of and at layers and m respec-

---

1.70
1.60
solu - gelu
1.85
gelu solu
1.58
Loss of stitched model
1.65
1.56
1.75
gelu-41 gelu-41
1.62
1.54
50l0-31 solu-31
gelu-41 gelu-4]
1.70
1.60
1.52
1.65
gelu 3
1.50
1.58
1.60
solu-41
- solu-4]
1.48
1.55
gelu-41 geelu--)
1.55
gelu- 31 - gelu-31
1.46
1.53
sollu44-- 50lu-41
Solu-30 solu-31 31
solu-41 soluu-41
1.50
1.44
1.50
layers
layers
3 layers
4 layers
layers
(a) Stitching at layer
(b) Stitching at layer
(c) Stitching at layer 3
Figure 1: Comparing stitching loss between 3-layer and 4-layer GeLU and SoLU models on the Pile validation set. Stitching for "solu-ogelu." iee stitching with SoLU 'head' and GeLU 'tail,' incurs systematic penalties (larger loss is worse) compared to the "gelu->solu" stitching. Baselines for stitching identical models (e.g.. "gelu- 31--gelu-31") are given account for potential inherent differences in learning the stitching layer ş between the activation functions.
cases. We conjecture that this outcome results from the fact that SoLU activations effectively reduce capacity of hidden feature layers; there may be an analogy between the experiments fig. I and those of (Bansal et al., 2021, Fig. 3c) stitching vision models of different widths.
As pointed out in (Bansal al., 2021) such an analysis is not possible using CKA, which only de- tects distance between distributions of hidden fea- tures (displayed for our GeLU and SoLU models in figs. 4 and 5), not their usefulness for a machine learning task. Further, the differences in layer ex- pressiveness between GeLU and SoLU models are not easily clicited by evaluation downstream tasks, but can be seen through linear stitching.
We also measure the stitching performance be- tween pairs of identical models, the SoLU-into- SoLU and GeLU-into-GeLU stitching baselines. This is meant to delineate between two factors in- fluencing the stitching penaltics being recorded: the ease of optimizing stitching layers and the inter- play between hidden features of possibly different architectures. We seek to measure differences be- tween hidden features, while the former is an addi- tional unavoidable factor inherent in model stitch- ing experiments. The identical SoLU-into-Sou and GeLU-into-GeLU stitching baselines serve as a proxy measure of stitching optimization success, since in-principle the stitching layer should be able to learn the identity matrix and incur a stitching penalty of 0. So, that SoLU-into-Sod SSLU stitches better than Bel.U-into-GeLA. gives evidence for SoLU models being casier to stitch with than GeLU models, from an optimization perspective. That both SoLU-into-GeLU and GeLU-into-SoLU incur higher stitching penalties than Gel.U-into-GelA. suggests that the penalties cannot only be a result stitching layer optimization issues.
4 Locating generalization failures
Starting with a single pretrained BERT model and fine-tuning 100 models (differing only in random initialization of their classifier heads and stochastic- ity of batching in fine-tuning optimization) the Multi-Gena Natural Language Inference (MNLI) dataset (Williams al., 2018), the authors of (Mc- Coy et al., 2020) observed a striking phenomenon: despite the fine-tuned BERTs' near identical perfor- mance on MNLI, their performance on Heuristic Analysis for NLI Systems (HANS), an out of dis- tribution variant of MNLI, (MeCoy et al., 2019) was highly variable. On specifie subset of HANS called "lexical overlap" (HANS-LO) one subset of fine-tuned BERTs (the generalizing) models) were relatively successful and used syntactic features; models in its complement (the heuristic models) failed catastrophically and used a strategy akin to

---

bag-of-words.' ! Recent work of (Juneja al., 2022) found that partitioning the 100 fine-tuned BERTs on the basis of their HANS-LO performance is essentially equivalent to partitioning them on the basis of mode connectivity: that is, the generalizing and heuristic models lie in separate loss-landscape basins. We refer the interested reader to (Juneja et al., 2022) for further details.
BERTs on the MNLI dataset. At early layers, CKA insensitive the generalizing or heuristic nature of model and B. At later layers (in the same range where we saw identity stitching penalties appear), the CKA measure between generalizing and heuristic models significantly exceeds its value within the generalizing and heuristic groups.
Together, the results of fig. paint the following picture: the NLT models of (McCoy et al., 2020) seem to all build up generally useful features in their early layers, and only decide to memorize (i.e., use lexical overlap heuristic) or generalize (use syntactic features) in their final layers.
But at what layer(s) do the hidden features of the generalizing and heuristic models diverge? For example, are these two subpopulations of mod- els distinct because their earlier layers are differ- ent because their later layers different? Or both? Neither analysis of OOD performance nor mode connectivity analysis can provide an answer, but both stitching and CKA reveal that the differ- ence between the features of the generalizing and heuristic models is concentrated in later layers of the BERT encoder. It is worth highlighting that conceptualizing and building relevant datasets for distribution shifts is often expensive and time con- suming. That identity stitching and CKA using in- distribution MNLI data alone differentiate between heuristic and generalizing behavior on HANS-LO suggests that they are useful tools for model error analysis and debugging.
5 Representation dissimilarity in scaling families
Enormous quantities of research and engineering energy have been devoted to design, training and evaluation of scaling families of language models. By a "scaling family" we mean a collection of neu- ral network architectures with similar components, but with variable width and/or depth resulting in a sequence of models with increasing size (as mea- sured by number of parameters). Pioneering work in computer vision used CKA to discover inter- esting properties of hidden features that vary with network width and depth (Nguyen et al., 2021), but to the best of our knowledge no similar studies have appeared in the natural language domain
Figure 2a displays performance of pairs of BERT models stitched with the identity function yp = id, i.e. models of the form J>m o f's on the MNLI finetuning task.? 2 We see that at carly layers of the models identity stitching incurs almost no penalty, but that stitching at later layers (when the frac- tion of layers occupied by at the bottom model exceeds 75 %) stitching benween generalizing and heuristic models incurs significant accuracy drop (high stitching penalty), whereas stitching within the generalizing and heuristic groups incurs almost no penalty.
We take first step in this direction by comput- ing CKA measures of features within and between the models of the Pythia family (Biderman et al., 2023), up to the second-largest model with 6.9 billion parameters. In all experiments we use the "deduped" models (trained on the deduplicated ver- sion of the Pile (Gao et al., 2020)). In the case of inter-model CKA measurements on py pythia-6.. 9b we see in fig. 3 (bottom right) that similarity be- tween layers 1 and m gradually decreases as |m I| decreases. We also see a pattern reminiscent of the "block structure" analyzed in (Nguyen et al., 2021), where CKA values are relatively high when com- paring two layers within one of the following three groups: (i) early layers (the first 3), (ii) late layers (in this case only the final layer) and (iii) the re- maining intermediate layers, while CKA values are relatively low when comparing two layers between these three groups.
A similar picture emerges in fig. 2b, where plot CKA values between features of fine-tuned
'By "relatively successful" mean "Achieving up 50% accuracy on an adversariallyydesignee test set." While variable OOD performance with respect fine-tuning seed was also seen on other subscts of HANS, follow (McCoy et al., 2020; Juneja al., 2022) in focusing on HANS-L.O where such variance was most pronounced.
"The motivation for stitching with a constant identity func- tion comes from evidence that stitching type algorithms per- form symmeerr correction (accounting for the fact that features of model A could differ from those of model B by an architec- ture symmetry e.g. permuting neurons, see e.g. (Ainsworth et al., 2023)). but that networks obtuined from multiple fine- tuning runs from a fixed pretrained model seem to not differ by such symmetries (sce for example (Wortsman et al., 2022)).
Using CKA to compare features between 1 1.bb, 2bb and pythia-6. 9b we observe high similarity between features at the be-

---

1.00
heur-heur gen-heur gen-gen
stitching accuracy
0.8
CKA
gen-to-heur
gen-to-gen
heur-to-gen
0.85
0.6
heur-to-heur
0
2
6 layer
8
10
2
10
layer
Figure 2: Left: Identity stitching on the MNLI dataset between pairs of the top 10 ("generalizing") and bottom ("heuristic") performing models of (McCoy al., 2020) on lexical overlap subset HANS (an out-of- distribution NLI dataset). Right: Corresponding CKA values. Both: Confidence intervals obtained evaluat- ing on all distinct pairs of generalizing and heuristic models, for a total of 2- (2) ) + -102 190 model comparisons.
pythia- 1b-deduped
0
pythia-1 db-deduped
0
2
0.8
3
0.75
6
9
8
12
0.50
10
15
12
18
0.25
21
24
30
12
18
24
30
gychia-6.9b-dedupae
pythia-61 6.9hia-..bbedduee
pythia-2 80-deduped
pythia-6 wwrhia-6.990110 9b-deddped
1.00
0.75
0.75
15
0.50
15
18
0.50
21
21
24
0.25
24
0.25
27
30
30
12
24
30
12
18
24
30
pythia 6.9b-deduped
pythia 6.9b-deduped
Figure 3: CKA between 1pythia-1b, 1.eiia--.44..... 1..44b,ypiiaa--p 2iiiaa....bb..ttttt bythaaatha-6.90 and pythia-6... evaluated
the Pile dataset (Gao et al., 2020) (higher more similar).
ginning of the model. A plausible reason for this early layer similarity is the common task of "de- tokenization," (Elhage et al., 2022) where early neurons in models may change unnatural tokeniza- tions more useful representations, for exam- ple responding strongly to compound words (e.g.. bir birthday party). We also continue to see "block structure" even when comparing between two models, and in the case of the pairs (pythia-1b, pythia-6. 6. 9b) and (pythia-1..) pythia-6. (.chia-6.9b)) 9b) substantial inter-model feature similarity in inter- mediate layers.
trated fig. 9, pythia-2. 8b exhibits such trend- breaking intermediate layer feature dissimilarity with every Pythia model with 1 billion more pa- rameters." 3 Upon close inspection, the pythia-(15, 1.b,, 6.b669b) models all have early layer "block" consisting of layers with rel- atively high CKA similarity, in pythia-2.. this block consists only 2 layers, suggesting the fea- tures of pythia-2. 2. 8b diverge from the rest of the family very early in the model. In table 1 we point out that some aspects of pythia 2b''s architec- ture are inconsistent with the general scaling trend of the Pythia family whole.
that while
The low CKA values obtained when com- paring intermediate layers of (pythia-2.0 8b, pythia-6.8 9b) break this trend in fact, as illus-
Excepp itself, of course.

---

Limitations
Samuel K. Ainsworth, Jonathan Hayase, and Sid- dhartha Srinivasa. 2023. Git Re-Basin: Merging Models modulo Permutation Symmetries.
In the current work we examine relatively small language models. Larger models have qualitatively different features, and it is not obvious if our exper- iments on the differences in layer expressiveness between GeLU and SoLU models will scale to significantly larger models. While we show that identity stitching and CKA distinguish the heuristic and generalizing BERT models, we did not attempt use stitching/CKA to automatically cluster the models (as was done with mode connectivity in (Juneja et al., 2022)).
Yamini Bansal, Preetum Nakkiran, and Boaz Barak. 2021. Revisiting model stitching to compare neu- ral representations. In Advances in Neural Informa- fion Processing Systems, volume 34, pages 225-236. Curran Associates, Inc.
Serguei Barannikov, Ilya Trofimov, Nikita Balabin, and Evgeny Burnaev. 2021. Representation topology di- vergence: A method for comparing neural network representations. arXiv preprint arXiv:2201.000 00058..
Yonatan Belinkov, Nadir Durrani. Fahim Dalvi, Has- san Sajjad, and James Glass. 2017. What do neu- ral machine translation models learn about morphol- ogy? Proceedings of the 55th Annual Meeting of the Association for Computational Linguistics ( (Vol- ume I: Long Papers), pages 861-872, Vancouver, Canada. Association for Computational Linguistics.
Ethics Statement
In this paper we present new applications of rep- resentation analysis to language model hidden fea- tures. Large language models have the potential impact human society in ways that we are only beginning to glimpse. deeper understanding of the features they learn could be exploited for both positive and negative effect. Positive use cases in- clude methods for enhancing language models to be more safe, generalizable and robust (one such approach hinted at in the experiments of section 4), methods for explaining the decisions of language models and identifying model components respon- sible for unwanted behavior. Unfortunately, these tools can at times be repurposed to do harm, for example extracting information from model train- ing data and inducing specific undesirable model predictions.
Nora Belrose, Zach Furman, Logan Smith, Danny Ha- lawi, Igor V. Ostrovsky, Lev McKinney, Stella Rose Biderman, and J. Steinhardt. 2023. Eliciting latent predictions from transformers with the tuned lens. ARXIVOOG.
Stella Biderman, Hailey Schoelkopf, Quentin An- thony. Herbie Bradley, Kyle O'Brien, Eric Hal- lahan, Mohammad Aflah Khan, Shivanshu Puro- hit, USVSN Sai Prashanth, Edward Raff, Aviya Skowron, Lintang Sutawika, and Oskar van der Wal. 2023. Pythia: A Suite for Analyzing Large Lan- guage Models Across Training and Scaling.
Frances Ding. Jean-Stanislas Demain, and Jacob Stein- hardt. 2021. Grounding representation similarity through statistical testing. Advances in Neural formation Processing Systems, 34:1556-1566
Acknowledgements
Nelson Elhage, Tristan Hume, Catherine Olsson, Neel Nanda, Tom Henighan, Scott Johnston, Sheer ElShowk, Nicholas Joseph, Nova Das- Sarma, Ben Mann, Danny Hernandez, Amanda Askell, Kamal Ndousse, Andy Jones, Dawn Drain, Anna Chen, Yuntao Bai, Deep Gan- guli, Liane Lovitt, Zac Hatfield-Dodds, Jackson Kernion, Tom Conerly, Shauna Kravec, Stanislav Fort, Saurav Kadavath, Josh Jacobson, Tran- Johnson, Jared Kaplan, Jack Clark, Tom Brown, Sam McCandlish, Dario Amodei, and Christo- pher Olah. 2022. Softmax linear units. Trans- former Circuits Thread. Https//transformer: hitiii/ppu//0222ss//nneeee
This research was supported by the Mathemat- ies for Artificial Reasoning Science (MARS) initiative via the Laboratory Directed Research and Development (LDRD) investments at Pacific Northwest National Laboratory (PNNL). PNNL is a multi-program national laboratory operated for the U.S. Department of Energy (DOE) by Battelle Memorial Institute under Contract No. AC-5-05- 76RL0-1830.
References
Rahim Entezari, Hanie Sedghi, Olga Saukh, and Behnam Neyshabur. 2022. The role of permutation invariance in linear mode connectivity of neural net- works. In The Tenth International Conference Learning Representations, ICLR 2022, Virtual Event, April 25-29, 2022. OpenReview.net.
Yossi Adi, Einat Kermany, Yonatan Belinkov, Ofer Lavi, and Yoav Goldberg. 2016. Fine-grained anal- ysis of sentence embeddings using auxiliary predic- tion tasks. International Conference Learning Representations.

---

Jonathan Frankle, G. Dziugaite, Daniel M. Roy, and Michael Carbin. 2019. Linear mode connectivity and the lottery ticket hypothesis. JCMI..
Tom McCoy, Ellie Pavlick, and Tal Linzen. 2019. Right for the wrong reasons: Diagnosing syntactic heuristics in natural language inference. In Proceed- ings of the 57th Annual Meeting of the Association for Computational Linguisties, pages 3428-3448, Florence, Italy. Association for Computational Lin- guisties.
Leo Gao, Stella Biderman, Sid Black, Laurence Gold- ing. Travis Hoppe, Charles Foster, Jason Phang. Horace He, Anish Thite, Noa Naheshima, Shawn Presser, and Connor Leahy. 2020. The pile: An 800gb dataset of diverse text for language modeling.
Neel Nanda. 2022. Transformerlens.
arXiv preprint arXiv: 2101.00027.
Thao Nguyen, Maithra Raghu, and Simon Kornblith. 2021. Do Wide and Deep Networks Learn the Same Things? Uncovering How Neural Network Repre- sentations Vary with Width and Depth.
Charles Godfrey, Davis Brown, Tegan Emerson, and Henry Kvinge, 2022. On the symmetriee of deep learning models and their internal representations. arXiv preprint 1205:22205.12558.
Maithra Raghu, Thomas Unterthiner, Simon Kornblith, Chiyuan Zhang, and Alexey Dosovitskiy. 2022. Do Vision Transformers See Like Convolutional Neural Networks?
John Hewitt and Percy Liang. 2019. Designing and in- terpreting probes with control tasks. Conference On Empirieal Methods In Natural Language Process-
ing.
Yuxin Ren, Qipeng Guo, Zhijing Jin, Shauli Ravfogel, Mrinmaya Sachan, Bernhard Schölkopf, and Ryan Cotterell. 2023. All roads lead to rome? explor- ing the invariance of transformers' representations. arXiv preprint arXiv: 2305.14555.
Jeevesh Juneja, Rachit Bansal, Kyunghyun Cho, João Sedoe, and Naomi Saphra. 2022. Linear connectiv- ity reveals generalization strategies. arXiv preprint arXiv: 2idi-22205.121111
Max Klabunde, Tobias Schumacher, Markus Strohmaier, and Florian Lemmerich. 2023. Sim- ilarity of neural network models: survey of functional and representational measures. arXiv preprint 2305.0699.
Cynthia Rudin. 2019. Stop explaining black box ma- chine learning models for high stakes decisions and use interpretable models instead. Nature machine intelligence, 1(5):206-215.
Naomi Saphra and Adam Lopez. 2019. Understand- ing learning dynamics of language models with SVCCA. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies. Volume 1 (Long and Short Papers). pages 3257-3267, Minneapolis, Minnesota. Associ- ation for Computational Linguistics.
Simon Korbblihh, Mohammad Norouzi, Honglak Lee, and Geoffrey E. Hinton. 2019. Similarity of Neural Network Representations Revisited. ICML.
Karel Lenc and A. Vedaldi. 2014. Understanding image representations by measuring their equivari- ance and equivalence. Computer Vision And Pattern Recognition.
Elena Voita, Rico Sennrich, and Ivan Titov. 2019. The bottom-up evolution of representations in the trans- former: A study with machine translation and lan- guage modeling objectives. In Proceedings of the 2019 Conference on Empirical Methods in Natu- ral Language Processing and the 9th International Joint Conference Natural Language Processing (EMNLP-ICCNLP,, pages 4396-4406. Hong Kong, China. Association for Computational Linguistics.
Nelson F. Liu, Matt Gardner, Yonatan Belinkov, Matthew E. Peters, and Noah A. Smith. 2019. Lin- guistic knowledge and transferability of contextual representations. In Proceedings of the 2019 Confer- ence of the North American Chapter of the Associ- ation for Computational Linguisties: Human Lan- guage Technologies, Volume I (Long and Short Pa- pers). pages 1073-1094, Minneapolis, Minnesota. Association for Computational Linguistics.
Adina Williams, Nikita Nangia, and Samuel Bowman. 2018. A broad-coverage challenge corpus for sen- tence understanding through inference. In Proceed- ings of the 2018 Conference of the North American Chapter of the Association for Computational Lin- guistics: Human Language Technologies. Volume (Long Papers), pages 12-11122, New Orleans, Louisiana. Association for Computational Linguis- tics.
Ekdeep Singh Lubana, Eric J. Bigelow, Robert P. Dick, David Krueger, and Hidenori Tanaka. 2022. Mech- anistic mode connectivity. arXiv preprint arXiv: 2214.0.
Thomas McCoy, Junghyun Min, and Tal Linzen. 2020. BERTs of a feather do not generalize to- gether: Large variability in generalization across models with similar test set performance. In Pro- ceedings of the Third BlackboxNI.P Workshop on An- alyzing and Interpreting Neural Networks for NLP, pages 217-227, Online. Association for Computa- tional Linguistics.
Alex H Williams, Erin Kunz, Simon Kornblith, and Scott Linderman. 2021. Generalized shape metrics on neural representations. Advances in Neural Infor- mation Processing Systems, 34:4738-477

---

Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pier- ric Cistac, Tim Rault, Rémi Louf, Morgan Funtow- icz, Joe Davison, Sam Shleifer, Patrick von Platen, Clara Ma, Yacine Jernite, Julien Plu, Canwen Xu, Teven Le Scao, Sylvain Gugger, Mariama Drame, Quentin Lhoest, and Alexander M. Rush. 2020. Transformers: Sfate-of-the-ars natural language pro- cessing. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing: System Demonstrations, pages 38 38-45, Online. Asso- ciation for Computational Linguistics.
Mitchell Wortsman, Gabriel Ilharco, Samir Yitzhak Gadre, Rebecca Roelofs, Raphael Gonejj-Lopes, Ari S. Morcos, Hongseok Namkoong, Ali Farhadi, Yair Carmon, Simon Kornblith, and Ludwig Schmidt. 2022. Model soups: Averaging weights of multiple fine-tumed models improves accuracy with- out increasing inference time.

---

A Related Work
The problem of understanding how to compare the internal representations of deep learning models in way that captures differences meaningful to the machine learning task has inspired rich line of research (Klabunde al., 2023). Popular approaches quantifying dissimilarity range from shape based measures (Ding et al., 2021; Williams et al., 2021; Godfrey et al., 2022) to topological approaches (Barannikov et al., 2021).
The vast majority of research into these approaches has concentrated on computer vision applications rather than language models. Exceptions have looked at the evolution of model hidden layers with respect training task (Voita et al., 2019) or downstream tasks (Saphra and Lopez, 2019), statistical testing of sensitivity and specificity of dissimilarity metrics applied to BERT models (Ding et al., 2021), and more flexible representation dissimilarity measures aligning hidden features with invertible neural networks (Ren et al., 2023). While representation dissimilarity measures have not been extensively explored language models, there a rich line of work examining the structure of and relationship between hidden layers of NLP transformer models via probing for specific target tasks (Adi et al., 2016; Liu et al., 2019; Belinkov et al., 2017; Hewitt and Liang, 2019).
Our Section 4 borrows an experiment set-up from (Juneja al., 2022), which found that partitioning BERT models using geometry of the fine-tuning loss surface (namely, linear mode connectivity (Frankle et al., 2019; Entezari et al., 2022)) can differentiate between models that learn generalizing features and heuristic features for a natural language inference task. Similar mode connectivity experiments for computer vision models were performed in (Lubana et al., 2022).
B Additional results
figs. and 5 we display CKA representation dissimilarity measurements for the GeLU and SoLU models used in the experiments of section 3. The purpose of these plots is simply to illustrate that in contrast to model stitching, CKA only detects distance between distributions of hidden features, not their usefulness for machine learning task.
1.0
IE-n105
0.8
0.6
0.4
gelu-31
solu-31
gelu:31
(a)
(b)
(c)
Figure 4: CKA for three layer GeL.U and SoLU models (higher means more similar) evaluated on Pile dataset (Gao et al., 2020). (a, b) Inter-model CKA for GeLU (respectively SoLU) models with 3 transformer blocks. (c) Between-model CKA for 3-black GeLUU,SoLU models (rows correspond to SoLU layers, columns to GeLU layers
The plot fig. (left) shows identity stitching penalties for the MNLI fine-tuned BERT models on the MNLI dataset. In practice data distribution shifts are generally not known the time of model develop- ment, and (as noted) the fact that identity stitching using MNLI data alone differentiates between heuristic and generalizing behavior on HANS-LO is potentially quite useful. Nevertheless, for comprehensiveness the performance of the identity stitched models on HANS-LO, displayed in fig. 7, also shows distinct differences between generalizing and heuristic models in the final stitching layer.
In fig. 2 (right) we display CKA between MNLI fine-tuned BERT models each layer. This has the benefit of showing that the features of top (generalizing) and bottom (heuristic) models diverge in late BERT encoder layers (as we also saw with identity stitching fig. 2 (left)). We aggregate the per-layer CKA measurements into a single distance measurement for each of the three cases (top-top, top-bottom, bottom-bottom) bottom) in fig. 8. This illustrates that the top (generalizing) and bottom (heuristic) models form

---

1.0
1.0
0.8
5004
0.7
300
0.4
gelu-41
solu-4]
gelu-4)
(a)
(b)
(c)
Figure 5: CKA for four layer GeLU and SoLU models (higher means more similar) evaluated on the Pile dataset (Gao et al., 2020). (a, b) Inter-model CKA for GeLU (respectively SoLU) models with transformer blocks. (c) Between-model CKA for 3-block GeLU,SoLU models (rows correspond to SoLU layers, columns to GeLU layers.
Stlccing after layer
Stitching second layer
Stitching third layer
gdu
petu-ll
z0lu-40
solu
gelu solu-d
getu stluctl
l0s3
ento ge
golu solu
gel 20
t0tu ge
soli--gelo-t)
getu
2
1500
2000
500
1000 batches
1500
2000
1000 batches
1500
2000
batches
(a)
(b)
(c)
Figure 6: Learning linear stitching layers for the stitching GeLU and SoLU models in Figure 1.
two distinct clusters in natural metric space of hidden features of MNLI datapoints; for the technical definition of this metric space refer appendix C. In fig. we saw high CKA values between intermediate features of pythia- pythia-(1b,11 (1b,1.4b) and pythia-6. 6, 9b, but comparatively low CKA values between intermediate features of the 2.8 and 6.9 billion parameter models. Figure 9 shows that more generally there is a high level of intermediate feature similarity for models in the set bythia-(1b, 1.1,,... 1b4b, 6. 9b} but low level of intermediate feature similarity between pythia-2. 8b and any of the models in pythia-(1b) 1. 6.)..
hypothesis for why pythia-2.5 8b possesses unique hidden features from the perspective of CKA, namely that its attention head dimension is quite bit smaller than those of the pythia-(1b, (1b, 1.4b, 6. 9b} models.
Table suggests
model_name
head rotary_ dim
pythia deddu-e
64
16
16
pythia-410m 1i0h-4410mddeeee dodddu-e deduped
64
16
pythia- Ib-dedued
256
64
pythia- 1.4iia-1.4bdeddupee
128
32
pythia-2.8b-dedee
80
20
pythia-6.9b-de -rbhia--.bb-ddduee
128
32
pythia- 1 12b-deduped
128
32
Table 1: Select HuggingFace (Wolf et al., 2020) model configuration dictionary keys for the Pythia models, show- ing that the 2.8b model has an unexpectedly small attention head dimension and hence rotary embedding demen- sion (the latter appears be tied 1/4 of the former).
We make a couple final notes on our analysis of Pythia models as it relates to prior work on image models. First, restricting attention to intra-model CKA heatmaps in fig. 10 we see an increase in intra-

---

0.7
label
heur-to-gen.
gen--to-gen.
0.6
heur-to-heur.
gen-ttohheur.
accs
0.4
0.3
0.2
0.1
4
6
10
12
layer
Figure 7: Identity stitching on the HANS-LO subset. The "heur-tt-heur:" and "gen.-to-gen." models still exhibit distinct behavior.
model feature similarity as scale increases to billion parameters, 4 analogous to an observation of (Nguyen et al., 2021). Note however that our experimental set up is not as clean as theirs: in our case both depth and width are scaled simultaneously, and perplexity of the Pythia models decreases parameter count increases. In particular, our results would also be consistent with hypothesis that intra-model feature similarity increases validation loss decreases. Second, speculate that the large blocks of intermediate features with high CKA similarity figs. and related to the finding of (Raghu et al., 2022) that (compared to ResNets) Vision Transformers exhibit a high degree of feature uniformity across model layers. In other words, we hypothesize that this phenomenon is not specific to Vision Transformers, but rather transformer models more generally. However, as we only experiment with transformer models we are unable to make such a comparative statement and leave a broader study including e.g. RNNs and/or state-space sequence models to future work.
"See also relevant subplots of figs. 3 and 9 for larger model scales.

---

CKA distances
1.5
distance w.r.t. product of arccos(CKA) metrics
1.4
1.3
1.2
1.1
1.0
heur-heur gen-heur
gen-gen
Figure 8: CKA-derived distances between the models of (McCoy et al., 2020) (higher means less similar, see appendix C for details on this distance metric). Top-top: pairwise CKAs between the top 10 models on HANS-LO. Top-bottom: pairwise CKAs between the top 10 and bottom 10 models on HANS-LO Bot-bot: pairwise CKAs between the bottom models HANS-LO.

---

0
1.00
0
bythia-1b-dedupee
2
2
0.8
4
0.75
4
6
6
0.6
8
0.50
8
0.4
10
10
12
0.25
12
0.2
14
14
6
10
12
12
115
18
21
pythia-lb-deduper
pythia-1.4b-dedupeee
pythia-1, 4b--dddupppd
0
1.00
0
3
Atthia-bb-edduer
2
0.8
6
0.75
4
0.6
9
6
12
0.50
8
pythia-
0.4
15
10
18
0.25
12
0.2
21
14
6
9
12
15
18
21
12
18
24
30
pythia-1.4b-dedupee
2ythia-2.8b-dedupel
pythia-1.49h 41.4-deduuped
0
pythia-2 2a..-2...p 1..-8-dee w...-bddddpee
1.00
3
0.75
6
0.75
9
15
0.50
12
0.50
18
15
21
18
0.25
24
0.25
21
27
30
12
18
24
30
12
18
24
30
pythia-2.8bb 2ete--2288bdddeeee deduped
pythia 2....-ddbddueee
Figure 9: CKA for all pairs of models in the set pythia-(1b, (1,, 1,, 8b), evaluated on the Pile dataset (Gao et al.,
2020) (higher means more similar).

---

1.00
0
1.00
0
2
pythia- 70m
1
0.75
pythia-160m
0.75
4
2
6
0.50
3
0.50
8
4
0.25
5
0.25
10
0
2
4
10
pythia-70m
pythia-160m
1.00
0
1.00
3
aythia-1p wa------dddppp
2
2ythia-410m
6
0.75
4
0.75
9
6
12
8
0.50
0.50
15
10
18
12
0.25
0.25
21
14
12
15
18
21
2
6
8
10
12
14
pythia-410m
bythia-1b-dedupe
Figure 10: Intra-model CKA for models in the set pythia-(70m,, 140m, 41 10m, 1b), evaluated on the Pile dataset
(Gao et al., 2020) (higher means more similar).

---

C Experimental details
C.1 SoLU/GeLUU and Pythia models
For the SoLUUGeLLU and Pythia experiments we use the Transformerl.ens library (Nanda, 2022), whose developers pre-trained the SoLU/GeLU models on the Pile (Gao et al., 2020) (architecture details can be found here and some training details available here). We use heavily adapted version of the tuned-lens library to performing stitching (Beirose et al., 2023). In particular, for stitching we learn an affine transformation (along with an additional LayerNorm) yp with the PILE validation set. Figure 6 plots learning curves for GeLU-SoLU stitching layer optimization. We borrow many of the Tuned Lens hyperparameters, and train for 2000 steps (200 warm-up steps) with SGD with Nesterov momentum with a learning rate of 1.0 and cross-entropy loss.
When computing CKA, we extract features after each transformer block (after addition with the residual, technically the hook_resid post hook point of TransformerLenn) Letting f and g be the two models under consideration and letting X be an input batch of token sequences, this yields feature tensors say
H; = fss(X) and Hj- = (X)
for i - 1,...,,.,444 and j- 1........ L' where L, L' are the respective depths of f and g. In our experiments we use batches consisting of 1 (!) sequence of 1024 tokens (the default context length for the Pythia models); these sequences are constructed by "packing" datapoints from the Pile.'. When one or more of S. is large (e.g. has more than billion parameters), we run the necessary forward passes of f and g on separate GPUs
used a mix of NVIDIA A100 GPUs with either 40 80 GB of RAM.
We then flatten the tensors of shape (batch, sequence, feature) (in our experiments (1, 1024, hidden_dim) to a matrix of features with shape (batch sequence, feature) (so in our experiments (1024, hidden_dim), 6 and then compute CKA values in batches using an implementation of the unbiased estimator described in (Nguyen et al., 2021). Explicitly, if H, and H' are (already flattened) hidden feature matrices of shape (batch * sequence, feature) (with possibly distinct feature dimensions), we first mean center both matrices (HI; + H; 11 Hi, similarly for H²)) and compute the kernels
K; H.H. and Lg H
Next, make use of the unbiased Hilbert-Schmice independence criterion HSIC (A, B), defined for symmetric matrices A and B of equal shape (n, n) as follows: first, let À and B be obtained by replacing the diagonals of and respectively with Os.? Then,
2
HSTC (4, := n(n-3) -3) (Ir(AB") + (n 1)(n 2) tr ((A11))tr(BB1)) n
A)
This unweildy-looking expression can be computed efficiently by noting for example that tr(Aß²) is simply the dot-product of the vectorizations of and B. tr(A11²) is just the sum of entries of Å, and A11"' 7 B is just the dot product of the vectors obtained by summing columns of Á and B. Finally, we compute a matrix S of shape (L, L') with entries
HSIC LT) Sij HSIC K2) FISICI(L,) L4)
and average these batch-estimated matrices S over many randomly sampled batches, in our experiments 1000.
"Fxplicitly,
the chunk_and__r tone_to function from Transformerl.com Lens
"In other words, we are comparing activations at token level, An alternative which have not yet evaluated would he flattening the sequence and feature indices to single dimension and comparing activations at sample level. "In PyTorch or NumPy, accomplished with = A - diag(diag(A))

---

C.2 BERT models
For the fine-tuned BERTs experiments we use fork of the codebase made available by (Juneja et al., 2022) (as well as their models made available on HuggingFace (Wolf et al., 2020)). When computing CKA values, we exiract features after cach transformer block (in this case before addition with the residual, as it was not immediately clear how to obtain features after the residual addition). The rest of our CKA calculation is very similar to that described for SoLU, GeLU and Pythia models above, with the following modifications: first, since MNLI is classification task input sequences are padded the maximum context length of the model. in this case 512, rather than "packed." Second, we collected features H = (X) and H (X) for many batches iin our case a total N 1000 datapoints and stack them to obtain a matrix of shape (N, max length, feature) for each layer of each model. Then we flatten the first two indices to get matrices of shape (N (M+max length, feature), which are "chunked" along the first index- - in our implementation to matrices of shape (1024, feature) - and passed to the batched, unbiased CKA estimator using HSIC; as above.
We do not expect these differences in implementation make much difference, what is described in the preceding paragraph is essentially equivalent to the procedure used for SoLU, GeLU and Pythia models bit with batch size of 2 instead of 1.
To collapse layer-by-layer CKA to a single scalar value representing distance between the hidden features of two models with the same depth (as in fig. 8). letting Hi,... Hr, and Hi,.... Hy, be the matrices of hidden features in the respective models and compute
L arccos(CKA(HH, H
(2)
The explanation of this expression is as follows: pointed out in (Williams et al., 2021), while CKA is not a metric (H H' implies CKA(H, H') 1, but for metric it would have to be 0), the arccos if CKA is metric in the mathematical sense. Equation (2) is then metric product of hidden feature spaces obtained from the arccos-CKi CKA metric each of the factors.


Reviews:
[{"review_id": "73z3rb6Ikm", "reviewer": "Reviewer_rr9F", "summary": "", "questions": "- Rework figure 1 and clean up abbreviations and rethink terminology.\n- Make the text in your subfigs larger (I couldn't read  Fig 1 until I zoomed in to 180%)", "limitations": "", "rating": 4, "confidence": 5, "soundness": 4, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you very much for your in-depth analysis of our results and valuable feedback. \n\n**Regarding interpretation of stitching results**: first, here is an outline of our thought process in slightly more detail. There are at least two factors influencing the stitching penalties being recorded – ease of optimizing stitching layers and the interplay between hidden features of possibly different architectures (the latter is what we seek to measure, the former is an additional unavoidable factor inherent in model stitching experiments). The identical SoLU-into-SoLU and GeLU-into-GeLU stitching baselines serve as a measure of stitching optimization success, since in principle the stitching layer should be able to learn the identity matrix and incur a stitching penalty of 0. So, that SoLU-into-SoLU stitches better than GeLU-into-GeLU gives weak evidence for SoLU models being easier to stitch with than GeLU models, from an optimization perspective. That both SoLU-into-GeLU and GeLU-into-SoLU incur higher stitching penalties than GeLU-into-GeLU suggests that the penalties cannot only be a result of stitching layer optimization issues. \n\nIn this context, one way of unfolding your question is: it is unclear to us whether higher-capacity layers result in easier or harder stitching optimization. This is an interesting issue for us to follow up on; one could imagine that simpler features are easier to align, or alternatively since we are learning a linear layer that richer features yield a better-conditioned stitching optimization problem. While we might speculate that the stitching loss penalty is dominated by the downstream usefulness of the features rather than the optimization difficulty, we have not seen this question addressed in previous stitching work [1, 2] (or the intra-model stitching of [3]). Such optimization considerations, as well as other potential contributions to differences in stitching performance between SoLU-into-GeLU and GeLU-into-SoLU models, may change or modify our capacity interpretation. We will add language reflecting this experimental uncertainty and hope that our results may inform future approaches for specifying capacity or generalization differences (e.g., by guiding evaluation suites designed to identify performance differences in SoLU vs GeLU models).\n\n**Regarding presentation**: \n\n- Apologies, “resp.” was meant to denote “respectively.” We will update it and other ambiguous abbreviations where appropriate.\n- Regarding heads and tails, we will revise to be explicit here in the text (e.g. “stitching early layers of a SoLU model to later layers of a GeLU model”) and are considering using notation of the form “SoLU -> GeLU” in the plots to indicate the direction in which features are passed. \n- Regarding Figure 1, on rereading we empathize with your comments. To make the experiment more readable and remove unnecessary complications, we have 1) changed the stitching metric in the figure to be the absolute loss over the validation set (rather than the relative loss wrt the GeLU-into-GeLU stitched model), 2) will adopt “more positive is worse” phrasing, and 3) will rework the color scheme (and increase font size). To make the comparisons more obvious, we are also considering changing the bar plots to dumbbell plots to show the loss delta between SoLU -> GeLU vs GeLU -> SoLU (but maintaining the plot-per-stitching-layer scheme). \n\n[1] Bansal, Y., Nakkiran, P., & Barak, B. (2021). Revisiting model stitching to compare neural representations. Advances in neural information processing systems, 34, 225-236.\n\n[2] Csiszárik, A., Kőrösi-Szabó, P., Matszangosz, A., Papp, G., & Varga, D. (2021). Similarity and matching of neural network representations. Advances in neural information processing systems, 34, 225-236.\n\n[3] Belrose, N., Furman, Z., Smith, L., Halawi, D., Ostrovsky, I., McKinney, L., ... & Steinhardt, J. (2023). Eliciting latent predictions from transformers with the tuned lens. arXiv preprint arXiv:2303.08112.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "Timely demonstration of the utility of the 2 model analysis tools on a diverse set of problems\n\nOverall this is a great example of an exploratory short paper that makes substantiated claims and will inspire future work. \n\nDespite covering 3 very different applications of the two representation dissimilarities measures in only 4 pages, the authors do a good job of justifying their appropriateness and motivating them for unfamiliar readers. Overall I am impressed by how cohesive the story is in spite of its breadth and brevity\n\n**I would like to see this paper accepted.**", "reasons_to_reject": "~Some of the discussion of claim (1) feels a little shaky to me (see my summary).~ **The authors have addressed this concern in their rebuttal and I now agree. Hoping to see the camera ready changed in light of these suggestions.**\n\nThe rapid discussion of 3 wildly different application areas may feel hasty or too thin for some. This isn't a problem in my opinion.\n\n**Overall I would not like to see this paper rejected**", "questions_to_authors": "- Rework figure 1 and clean up abbreviations and rethink terminology.\n- Make the text in your subfigs larger (I couldn't read  Fig 1 until I zoomed in to 180%)", "paper_topic_and_main_contributions": "This paper is an exploration of applying representation dissimilarity measures to a small selection of LMs and natural language tasks.\n\nSpecifically they use them to investigate the role that activation functions play, and the differences of representations captured in different layers of Pythia models of varying sizes. The dissimilarity functions they use are model stitching and centered kernel alignment (CKA).\n\nThe main 3 claims are:\n1. model stiching works better for GeLU into SoLU vs SoLU into GeLU for LM\n2. dissimilarity measures show differences not surfaced via test set performance\n3. wider and deeper Pythia LMs have highly similar early layers across sizes, while intermediate layers grow more distant in representations\n\n(1) matters because SoLU is designed to be interpretable, and assessing how this impacts learned reps as well as downstream performance are important. SoLU is intended to reduce polysemanticity using softmax rather than sigmoid as the nonlinear component in the activation function. While the authors who introduced it showed comparable SoLU performance to GeLU, they didn't analyze the representations.\n\nThe takeaway is that solu stitches together to itself better than gelu does, and *also* solu stitches on top of gelu better than gelu does on top of solu, when levels are held equal. The authors think this is because SoLU activations make lower-capacity layers. ~*I'm not sure if I buy this interpretation considering solu-solu stitches better than gelu-gelu?*~ **After the authors' rebuttal I am more convinced of this claim. The key point they make is that solu-solu vs gelu-gelu stitching baseline is intended to separate the inherent \"stitchability\" of the model from the utility of the features learned in the model for downstream adaptation. With this explanatory text added I think the camera ready will be more convincing to readers.**\n\n(2) Is very important for a variety of questions, such as understanding the difference between generalizating and heuristic models for NLI by identifying divergent layers.\n\nThey show that MNLI test accuracy under model stitching alone can actually surface the difference between heuristic and generalizing models without deferring to HANS accuracy, just based on the marked dip in test accuracy performance when stitching in the high layers. This suggests that representations in the lower layers of a BERT model trained on MNLI contain consistent features, while the \"decision\" to memorize heuristics or learn the generalization features happens in the final few layers. Similarly, separation in CKA distance wrt layer between the classes of models also starts to occur at higher layers, further confirming this hypothesis.\n\n(3) Is a bit more of an afterthought as it is introduced on the last page, but is a nice demonstration of the utility of CKA for analyzing LLM feature representations. In future work this could be expanded, and their narrow claim is demonstrated sufficiently by Figure 3.", "excitement": 4, "reproducibility": 5, "ethical_concerns": "No", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"Missing_References": "N/a"}, "extra_text": {"Typos_Grammar_Style_and_Presentation_Improvements": "*I would not like to see this rejected, but I think there is some confusing writing in the early sections that needs to be addressed:*\n\nrelated to claim (1): \n- Lots of abbreviations like \"resp.\" (eg, line 103) Is this wrt? respective? Idk.\n- **Figure 1 really does not stand alone, and requires you to track through a bit of confusing double-negative logic to get the takeaways.** \"larger is worse\" phrasing for a value that can go negative (the performance penalty) is a little confusing because that can mean magnitude, I would frame it as \"more positive is worse\" or something. Additionally, the phrase \"solu-gelu stitching (stitching with a SoLU 'head' and GeLu 'tail')\" is very counterintuitive to me, as in LM we usually refer to upper layers as being the \"head\" of the model, but in this case the first activation function is the $f$ function which is on the bottom. It took me multiple rereads to parse this correctly, please reconsider phrasing.\nAlso, please explain what the red bars are in the caption for Figure 1 because you are overriding the de-facto standard \"red is bad\" connotation, which having a red bar in Figure 1 (a) be the most negative just made things really confusing."}}, {"review_id": "j88N3QMZye", "reviewer": "Reviewer_EMMQ", "summary": "", "questions": "1. I am not an expert in evaluating the activation functions so the following question is a clarification question: on p.1 in the last paragraph, there is an assumption that less penalty for loss leads to more useful information for the training task. Could you provide citations or footnotes explaining the rationale?\n\n2. What is the theoretical contribution of your final step? p.2, apply CKA to the Pythia networks.\n\n3. What is the theoretical rationale for choosing GeLU and SoLU as the particular model representations? You mentioned that most contemporary transformers are using them but are there more reasons?\n\nAfter the rebuttal: All these questions are properly solved.", "limitations": "", "rating": 4, "confidence": 1, "soundness": 3, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks very much for your thoughtful comments and questions.\n - Re: “… there is an assumption that less penalty for loss leads to more useful information for the training task. Could you provide citations or footnotes explaining the rationale?” Our main reference here is [1] which found in a variety of experiments that stitching earlier layers of a model with more parameters/more training data/trained for more epochs (all things most practitioners generally feel lead to better internal representations) into later layers of a model with fewer parameters/fewer training data/trained for fewer epochs resulted in better performance than stitching in the opposite direction. We will add more discussion of this prior work as context.\n - Re: contribution of CKA, using this tool to understand how architecture choices (CNNs vs. Vision Transformers) and model scale (in terms of width and depth) impact features of computer vision models. To our knowledge no similar studies have been conducted on language models. One can view our results as showing some phenomena previously observed in image classifiers can also be seen in LMs:\n      - An emergent block structure previously observed in CKA analysis of ResNets and ViTs [2], attributed to residual connections (see in particular Fig. 8). We will add a reference to [2] and this qualitative observation.\n      - Some evidence that as model width increases intra-model layer similarity increases (compare Figs. 3a and 8), parallel to another finding from [2]. Since submitting, we reworked our CKA evaluation to run on larger models (up to the 6.9 billion parameter Pythia model) and have some evidence that scaling width also increases inter-model layer similarity.\n      - Finally, the significant amount of similarity between the layers of a given Pythia model (seen in Fig. 8) is reminiscent of a finding on feature uniformity in ViTs in [3], although in our case we do not compare to other architectures. Using CKA to compare the representations of e.g., transformer models/self-attention layers to RNNs or state space models may be an interesting avenue for future work. \n - Re: rationale for choosing the GeLU vs SoLU comparison. The SoLU activation function was designed to improve the neuron-level interpretability of transformer language models by reducing neuron superposition in MLP layers [1]. An interesting finding from [1] was that language models with the SoLU activation function achieved comparable performance to their GeLU counterparts in certain downstream evaluations. We hypothesized that the hidden feature representations (and maybe even generalization abilities) of small SoLU models are different (plausibly weaker) compared to their GeLU counterparts. Our stitching experiments support this hypothesis. However, we agree that this is only preliminary evidence, in contrast to the more straightforward conclusion from the BERT NLI study.\n\n[1] Bansal, Y., Nakkiran, P., & Barak, B. (2021). Revisiting model stitching to compare neural representations. Advances in neural information processing systems, 34, 225-236. \n\n[2] Nguyen, T., Raghu, M., & Kornblith, S. (2020). Do wide and deep networks learn the same things? uncovering how neural network representations vary with width and depth. arXiv preprint arXiv:2010.15327.\n\n[3] Raghu, M., Unterthiner, T., Kornblith, S., Zhang, C., & Dosovitskiy, A. (2021). Do vision transformers see like convolutional neural networks?. Advances in Neural Information Processing Systems, 34, 12116-12128.\n\n[4] Elhage et al., “Softmax linear units,” Transformer Circuits Thread, 2022.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "Novel techniques", "reasons_to_reject": "None", "questions_to_authors": "1. I am not an expert in evaluating the activation functions so the following question is a clarification question: on p.1 in the last paragraph, there is an assumption that less penalty for loss leads to more useful information for the training task. Could you provide citations or footnotes explaining the rationale?\n\n2. What is the theoretical contribution of your final step? p.2, apply CKA to the Pythia networks.\n\n3. What is the theoretical rationale for choosing GeLU and SoLU as the particular model representations? You mentioned that most contemporary transformers are using them but are there more reasons?\n\nAfter the rebuttal: All these questions are properly solved.", "paper_topic_and_main_contributions": "This paper brings up a very important issue regarding the interpretability of deep neural networks. I can see their novelty is utilizing underrepresented techniques in NLP/NLU, e.g., model stitching and CKA, for such evaluation. Had I known more about this specific area, I would love to see these kinds of work published in EMNLP. However, I have to say I am not familiar with topics like this. I will defer the acceptance opinion to the other two reviewers. Wish the authors good luck!", "excitement": 4, "reproducibility": 3, "ethical_concerns": "No", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"Typos_Grammar_Style_and_Presentation_Improvements": "### LaTeX style\n\n1. Double check the in-text citation format. E.g., the experimental set up of (Juneja et al. 2022) -> the experimental set up of Juneja et al. (2022)"}, "extra_text": {}}, {"review_id": "5s1DyoT0bk", "reviewer": "Reviewer_Mj3F", "summary": "", "questions": "A. If SoLU activation reduces the capacity of hidden feature layers compared to GeLU activation functions, why the solu-4l models perform better than gelu counterpart?", "limitations": "", "rating": 4, "confidence": 3, "soundness": 4, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your thorough read of our submission and suggestions for improvement.\n\n- “Why only do model stitching on models with different activation functions?” This is a valid question, and we think that a set of experiments more directly inspired by our references in the computer vision literature (i.e. using models of varying parameter count/training time/training data size) would be worth carrying out as well (to our knowledge this would still be novel for LMs). Our interest in SoLU layers stemmed from their interest from the point of view of neuron interpretability, their relevance to the phenomenon of superposition, and the not-fully-understood effects on hidden feature representations [1]. In a longer paper we would expand this experiment suite considerably.\n- “For such a GeLU baseline, do you still put a linear stitching layer?” Yes, we do.\n- “What is HANS-“LO” different from HANS? Why the Figure 2’s experiment only test on the lexical overlap subset of HANS?” HANS-LO is a particular subset of HANS testing whether an MNLI-trained model uses the (generally incorrect) heuristic “assume that a premise entails all hypotheses constructed from words in the premise” (see [2] and references therein). It would be interesting to expand our experiments to include the other 2 subsets, however we follow [3] in focusing on LO where the most striking generalization differences were observed in [2].\n- “For experiments in Figure 2, are these figures do a ten times ten combination stitching /CKA score calculation for each line?” Yes, exactly. \n- “… I suggest the author considering extend the paper to long paper and perform more thorough studies to support each conjecture.” We agree that a longer paper with more comprehensive experiments would provide deeper insights, and our decision to submit as a short paper was motivated by considerations very similar to what you outline here. Moreover, the conclusions of the BERT generalization experiments are more clear-cut than the SoLU-GeLU experiments you mention, identifying specific layers where feature-distortion is correlated with generalization performance. While we can continue to expand the paper and hope to build on these results in follow-up work, we felt that the results on hand were worth sharing.\n\n- R.e.: question on the stitching performance of SoLU-into-SoLU models, we include the stitching performance of identical SoLU-into-SoLU and GeLU-into-GeLU as a control for optimization differences for the linear stitching layer.  SoLU-into-SoLU stitches better than GeLU-into-GeLU, and this gives weak evidence for SoLU models being easier to stitch with than GeLU models, from an optimization perspective. That both SoLU-into-GeLU and GeLU-into-SoLU incur higher stitching penalties than GeLU-into-GeLU suggests that the penalties cannot only be a result of stitching layer optimization issues. We provide more context for our interpretation for this experiment in our response to reviewer rr9F above.\n\n[1] Elhage et al., “Softmax linear units,” Transformer Circuits Thread, 2022.\n\n[2] McCoy, R. T., Min, J., & Linzen, T. (2019). BERTs of a feather do not generalize together: Large variability in generalization across models with similar test set performance. arXiv preprint arXiv:1911.02969.\n\n[3] Juneja, J., Bansal, R., Cho, K., Sedoc, J., & Saphra, N. (2022). Linear connectivity reveals generalization strategies. arXiv preprint arXiv:2205.12411.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "Two interesting methods are introduced and tried to analyze language models, despite using small models.", "reasons_to_reject": "1. The clarity of the paper and the motivation of several design choices can be further improved. For example, \n- Why only do model stitching on models with different activation functions?\n- What's the Nanda, 2022 's model pre-trained on?\n- (Line 158~160) For such a GeLU baseline, do you still put a linear stitching layer?\n- What is HANS-\"LO\" different from HANS?\n- Why the Figure 2's experiment only test on the lexical overlap subset of HANS?\n- For experiments in Figure 2, are these figures do a ten times ten combination stitching /CKA score calculation for each line?\n\n2. Overall, although I like the proposed methods, the conclusions we can draw from the paper is relatively sparse and not providing enough support. For example, if the author conjecture that SoLU activations effectively reduce the capacity of hidden feature layers, what is the corresponding influence? Is this asymmetry property correlate to final task performances? I suggest the author considering extend the paper to long paper and perform more thorough studies to support each conjecture.", "questions_to_authors": "A. If SoLU activation reduces the capacity of hidden feature layers compared to GeLU activation functions, why the solu-4l models perform better than gelu counterpart?", "paper_topic_and_main_contributions": "This paper presents two ways to understand the language model's behavior by considering the models' dissimilarity. Specifically, they consider using model stitching to combine different but similar models, and CKA score analysis is also utilized to measure the layer-wise similarity between models.", "excitement": 4, "reproducibility": 3, "ethical_concerns": "No", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"Typos_Grammar_Style_and_Presentation_Improvements": "Appendix B, Figure 4 & 5, \"dblcheck -CG\" is not removed."}, "extra_text": {}}, {"review_id": "JbSVkWDDJ4", "reviewer": "Reviewer_Dvu9", "summary": "", "questions": "", "limitations": "", "rating": 4, "confidence": 3, "soundness": 4, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for your consideration of our submission and valuable comments. \n\nRegarding stitching models with different activation functions, our interest in SoLU layers stemmed from their interest from the point of view of neuron interpretability, superposition, and not-fully-understood effects of activation functions on hidden feature representations [1]. We realize that a set of experiments more directly inspired by related work in the computer vision literature (i.e. using models of varying parameter count/training time/training data size) would be worth carrying out as well (to our knowledge this would still be novel for LMs). In a longer paper we could of course expand this experiment suite considerably.\n\n\n[1] Elhage et al., “Softmax linear units,” Transformer Circuits Thread, 2022.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "The paper discusses a very relevant concern about language models: how different are language models and why do they (or not) generalize on out-of-distribution data. ", "reasons_to_reject": "Three insights are discussed in the paper. The insights on generalizability and scaling are understandable but the authors did not motivate the need for an in-depth study on activation functions.  ", "questions_to_authors": "", "paper_topic_and_main_contributions": "The paper introduces representation dissimilarity as a tool for investigating the inner workings of language models. The paper posits that understanding the inner workings of the models is necessary for interpreting their behavior. Three insights are provided in the paper to support this premise: 1) an apparent asymmetry in the internal representations of model using SoLU and GeLU activation functions; 2) evidence that dissimilarity measures can identify generalization properties of models that are invisible via in-distribution test set performance, and 3) new evaluations of how language model features vary as width and depth are increased. Two methods for representation dissimilarity are discussed: model stitching and centered kernel alignment.", "excitement": 4, "reproducibility": 3, "ethical_concerns": "No", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"Typos_Grammar_Style_and_Presentation_Improvements": "OOD is not defined in text (out of distribution)?"}, "extra_text": {}}, {"review_id": "1nraFdFQpI", "reviewer": "Reviewer_ur3B", "summary": "", "questions": "1. For each of the study, could you elaborate on the potential impact of the phenomenon observed on the way models encode and process language? \n2. Do you consider some of your observed phenomena to be unique to language models, or in general NLP models, or could they also apply for deep learning models in general?", "limitations": "", "rating": 4, "confidence": 3, "soundness": 4, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your feedback and comments.\n\n - “For each of the study, could you elaborate on the potential impact of the phenomenon observed on the way models encode and process language?” \n   - Our stitching/dissimilarity study for NLI models is perhaps worth highlighting. It was found that there are striking generalization differences in BERT models trained with different random seeds on NLI tasks, where some models learn a bag-of-words heuristic while others learn better generalizing syntactic heuristics [1]. Follow-up work demonstrated that such generalization differences can be identified without a dataset designed to exhibit domain shift (like the HANS-LO dataset) by examining linear mode connectivity [2]. We build on this result by demonstrating that the NLI models seem to build up generally useful features in their early layers, and only decide to memorize (i.e., use a lexical overlap heuristic) or generalize (use syntactic features) in their final layers. This is perhaps surprising as these modelling strategies appear to be quite difference and are in different loss basins [2]. We plan to add a description of the bag-of-words vs syntactic strategies when we cite [1, 2].\n   - We expect the differences between the SoLU and GeLU activation functions (if such generalization differences exist) to manifest as differences in how the transformer models encode and process language. Our experiments only demonstrate that such a difference may exist, and we leave further development to future work. However, we hope that our experiments indicate that model stitching is a useful preliminary tool for understanding how models encode language. Likewise, while our CKA experiments for the Pythia suite give evidence for the early layers of language models being similar across scales, we remain uncertain of the mechanistic behavior of the models. Nonetheless, we speculate that this is related to de-tokenization [3], where early layer MLPs may convert unnational token representations (e.g., separate tokens for a compound word “social” and “security\") to representations more natural for downstream use (a single discrete representation for “social security”).\n\n - “Do you consider some of your observed phenomena to be unique to language models, or in general NLP models, or could they also apply for deep learning models in general?” TL;DR our experiments are at least specific to sequence models, and we argue some are specific to NLP. \n   - The BERT NLI experiments test a generalization/robustness phenomenon that is specific to NLP. \n   - One technical aspect that we would further highlight in a camera-ready revision is that both stitching and CKA are conducted at the token level, i.e. we apply the same stitching matrix at every point in the sequence. Similarly, the features we apply CKA to are of size the hidden dimension of the transformer model in question, so we are comparing the similarity of representations at the token level, not just the sample level as is often done in other domains (like computer vision) – we outline the latter point in Appendix C. In this way, our experiments are at least specific to sequence models, and of course all models occurring in the paper are text models. \n   - Regarding our Pythia CKA evaluations, we hope that there is value in demonstrating that some phenomena observed in computer vision models can also be seen in LMs. Without running such experiments, the question of whether these are phenomena general to a class of neural networks regardless insensitive to the data they are trained on, or if they are domain specific, would remain open.\n\n[1] McCoy, R. T., Min, J., & Linzen, T. (2019). BERTs of a feather do not generalize together: Large variability in generalization across models with similar test set performance. arXiv preprint arXiv:1911.02969.\n\n[2] Juneja, J., Bansal, R., Cho, K., Sedoc, J., & Saphra, N. (2022). Linear connectivity reveals generalization strategies. arXiv preprint arXiv:2205.12411.\n\n[3] Elhage et al., “Softmax linear units,” Transformer Circuits Thread, 2022.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "The author(s) show how techniques used more commonly in e.g., computer vision, can be used to study NLP models. \n\nFor each of the three techniques presented, the methods and experiments are well described, with insights of different nature for each study.", "reasons_to_reject": "The insights derived remain at an abstract machine learning level and the discussion remains rather disconnected from the fact that these are models specifically encoding and processing language. That is, the author(s) do not discuss the potential implications of these results on the way a LM encode and process a linguistic sequence, or the linguistic feature that it learns.", "questions_to_authors": "1. For each of the study, could you elaborate on the potential impact of the phenomenon observed on the way models encode and process language? \n2. Do you consider some of your observed phenomena to be unique to language models, or in general NLP models, or could they also apply for deep learning models in general?", "paper_topic_and_main_contributions": "The paper proposes to use dissimilarity metrics applied to hidden representation of different language models to derive insights into their inner workings.", "excitement": 4, "reproducibility": 4, "ethical_concerns": "No", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: This work explores using dissimilarity measures for model interpretability, including findings that these measures can give insights about model generalization. All reviewers appreciated the execution and motivation of this work, and it would be a good short paper at EMNLP.
Decision: Accept-Main