Understanding the Inner-workings of Language Models Through Representation Dissimilarity (2023)

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.

## introduction the defining feature of deep neural networks is their capability of learning useful feature representations from data in an end-to-end fashion. perhaps ironically, one of the most pressing scientific challenges in deep learning is understanding the features that these models learn. this challenge is not merely philosophical: learned features are known to impact model interpretability/explainability, generalization, and transferability to downstream tasks, and it can be the case that none of these effects are visible from the point of view of model performance on even carefully selected validation sets. when studying hidden representations in deep learning, a fundamental question is whether the internal representations of a given pair of models are similar or not. dissimilarity measures @xcite are a class of functions that seek to † work done at pacific northwest national laboratory. address this by measuring the difference between (potentially high-dimensional) representations. in this paper, we focus on two such functions that, while popular in computer vision, have seen limited application to language models: model stitching @xcite and centered kernel alignment (cka) @xcite . 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 a simple transformation to such an 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 models/layers from a fixed set of input datapoints, ignoring any relationship to performance on the task for which the models were trained. in this paper we make the case that dissimilarity measures are a 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 to better understand the changes to representations that result from using different nonlinear activations in a model. in particular, we find evidence that feeding gaussian error linear unit (gelu) model representations into a softmax linear unit (solu) model incurs a smaller penalty in loss compared to feeding solu activations into a 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 dissimilarity measures can localize differences in models that are invisible via test set performance. following the experimental set-up of @xcite , 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 @xcite , a sequence of generative transformers of increasing width and depth, finding a high degree of similarity between early layer features even as scale varies from 70 million to 1 billion parameters, an emergent "block structure" previously observed in cka analysis of image classifiers such as resnets and vision transformers, and showing that cka identifies a pythia model (pythia-2.8b-deduped) exhibiting remarkably low levels of feature similarity when compared with the remaining models in the family (as well as inconsistent architectural characteristics). this last finding is perhaps surprising considering the consistent trend towards lower perplixity and higher performance on benchmarks with increasing model scale seen in the original pythia evaluations @xcite -one might have expected to see an underlying consistent evolution of hidden features from the perspective of cka. ## background in this section we review the two model dissimilarity measures appearing in this paper. model stitching @xcite : informally, model stitching asks how well the representation extracted by the early layers of one model can be used by the later layers of another model to solve a specific task. let f be a neural network and for a layer l of f let @xmath0. given another network g the model obtained by stitching layer l of f to layer m of g with stitching layer ϕ is @xmath1. the performance of this stitched network measures the similarity of representations of f at layer l and representations of g at layer m. centered kernel alignment (cka) @xcite : let @xmath2, . . . , x d } be a set of model inputs. for models f and g with layers l and m respectively, let a f,l,g,m be the covariance matrix of @xmath3. then the cka score for models f and g at layers l and m respec-tively and evaluated at d is ## model stitching reveals an asymmetry between gelu and interpretable-by-design solu the design of more interpretable models is an area of active research. interpretable-by-design models often achieve comparable performance on downstream tasks to their non-interpretable counterparts @xcite . however, these models (almost by definition) have differences in their hidden layer representations that can impact downstream performance. many contemporary transformers use the gaussian error linear unit activation function, approximately gelu(x) = x * sigmoid(1.7x). the softmax linear unit solu(x) = x•softmax(x) is an activation function introduced in @xcite in an attempt to reduce neuron polysemanticity: the softmax has the effect of shrinking small and amplifying large neuron outputs. one of the findings of @xcite was that solu transformers yield comparable performance to their gelu counterparts on a range of downstream tasks. however, those tests involved zero-shot evaluation or finetuning of the full transformer, and as such they do not shed much light on intermediate hidden feature representations. we use stitching to do this. using the notation from section 2, we set our stitching layer ϕ to be a learnable linear layer (all other parameters of the stitched model are frozen) between the residual streams following a layer l. we stitch small 3-layer (9.4m parameter) and 4layer (13m parameter) solu and gelu models trained by @xcite , using the pile validation set @xcite to optimize ϕ. figure 1 displays the resulting stitching losses calculated on the pile validation set. for both the 3-layer and 4-layer models, and at every layer considered, we see that when f is a solu model for the stitched model @xmath4, i.e. the "gelu-into-solu" stitching for "solu-→gelu," i.e. stitching with a 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-3l-→gelu-3l") are given to account for potential inherent differences in learning the stitching layer ϕ between the activation functions. we also measure the stitching performance between pairs of identical models, the solu-into-solu and gelu-into-gelu stitching baselines. this is meant to delineate between two factors influencing the stitching penalties being recorded: the ease of optimizing stitching layers and the interplay between hidden features of possibly different architectures. we seek to measure differences between hidden features, while 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 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-solu stitches better than gelu-into-gelu gives 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. as pointed out in @xcite ) such an analysis is not possible using cka, which only detects distance between distributions of hidden features (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 expressiveness between gelu and solu models are not easily elicited by evaluation on downstream tasks, but can be seen through linear stitching. ## 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 stochasticity of batching in fine-tuning optimization) on the multi-genre natural language inference (mnli) dataset @xcite , the authors of (mc-coy et al., 2020) observed a striking phenomenon: despite the fine-tuned berts' near identical performance on mnli, their performance on heuristic analysis for nli systems (hans), an out of distribution variant of mnli, @xcite was highly variable. on a specific 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. 1 recent work of @xcite 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 @xcite for further details. but at what layer(s) do the hidden features of the generalizing and heuristic models diverge? for example, are these two subpopulations of models distinct because their earlier layers are different or because their later layers are different? or both? neither analysis of ood performance nor mode connectivity analysis can provide an answer, but both stitching and cka reveal that the difference 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 consuming. that identity stitching and cka using indistribution mnli data alone differentiate between heuristic and generalizing behavior on hans-lo suggests that they are useful tools for model error analysis and debugging. figure 2a displays performance of pairs of bert models stitched with the identity function ϕ = id, i.e. models of the form @xmath5. 2 we see that at early layers of the models identity stitching incurs almost no penalty, but that stitching at later layers (when the fraction of layers occupied by at the bottom model exceeds 75 %) stitching between generalizing and heuristic models incurs significant accuracy drop (high stitching penalty), whereas stitching within the generalizing and heuristic groups incurs almost no penalty. a similar picture emerges in fig. 2b , where we plot cka values between features of fine-tuned berts on the mnli dataset. at early layers, cka is insensitive to the generalizing or heuristic nature of model a 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. 2 paint the following picture: the nli models of @xcite seem to all 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. ## 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 neural network architectures with similar components, but with variable width and/or depth resulting in a sequence of models with increasing size (as measured by number of parameters). pioneering work in computer vision used cka to discover interesting properties of hidden features that vary with network width and depth @xcite , but to the best of our knowledge no similar studies have appeared in the natural language domain we take a first step in this direction by computing cka measures of features within and between the models of the pythia family @xcite , up to the second-largest model with 6.9 billion parameters. in all experiments we use the "deduped" models (trained on the deduplicated version of the pile @xcite ). in the case of inter-model cka measurements on pythia-6.9b we see in fig. 3 (bottom right) that similarity between layers l and m gradually decreases as |m-l| decreases. we also see a pattern reminiscent of the "block structure" analyzed in @xcite , where cka values are relatively high when comparing 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 remaining intermediate layers, while cka values are relatively low when comparing two layers between these three groups. using cka to compare features between pythia-{1b,1.4b,2.8b} and pythia-6.9b we observe high similarity between features at the beginning of the model. a plausible reason for this early layer similarity is the common task of "detokenization," @xcite where early neurons in models may change unnatural tokenizations to more useful representations, for example responding strongly to compound words (e.g., 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.9b) and (pythia-1.4b, pythia-6.9b) substantial inter-model feature similarity in intermediate layers. the low cka values obtained when comparing intermediate layers of (pythia-2.8b, pythia-6.9b) break this trend -in fact, as illus-trated in fig. 9 , pythia-2.8b exhibits such trendbreaking intermediate layer feature dissimilarity with every pythia model with 1 billion or more parameters. 3 upon close inspection, we see that while the pythia-{1b,1.4b,6.9b} models all have a early layer "block" consisting of 3 layers with relatively high cka similarity, in pythia-2.8b this block consists of only 2 layers, suggesting the features of pythia-2.8b diverge from the rest of the family very early in the model. in