The Low-Dimensional Linear Geometry of Contextualized Word Representations (2021)

Abstract: Black-box probing models can reliably extract linguistic features like tense, number, and syntactic role from pretrained word representations. However, the manner in which these features are encoded in representations remains poorly understood. We present a systematic study of the linear geometry of contextualized word representations in ELMO and BERT. We show that a variety of linguistic features (including structured dependency relationships) are encoded in low-dimensional subspaces. We then refine this geometric picture, showing that there are hierarchical relations between the subspaces encoding general linguistic categories and more specific ones, and that low-dimensional feature encodings are distributed rather than aligned to individual neurons. Finally, we demonstrate that these linear subspaces are causally related to model behavior, and can be used to perform fine-grained manipulation of BERT’s output distribution.

## introduction contextual word representations @xcite encode general linguistic features (e.g. semantic class; @xcite and sentence-specific relations (e.g. syntactic role; @xcite . used as input features, pretrained representations enable efficient training of models for a variety of nlp tasks @xcite . an enormous body of recent work in nlp has attempted to enumerate what features are encoded by pretrained word representations (e.g. @xcite @xcite , and more recent approaches have studied how accessible this information is, characterizing the tradeoff between generic notions of model complexity and accuracy needed to recover word features from learned representations @xcite . but the manner in which these features are encoded in representations remains poorly understood. are subspaces encoding linguistic features… hierarchical? or unstructured? axis-aligned? an additional set of experiments (d) uses these findings to identify linear operators that predictably affect outputs of a masked language model. best viewed in color. in this paper, we investigate the shape and structure of learned representation spaces (figure 1 ). we build on previous studies that have identified specific linear subspaces responsible for gender bias in word embeddings @xcite and word sense disambiguation @xcite . we show that these linear subspaces are the rule, not the exception: a variety of other lexical and syntactic phenomena are encoded in low-dimensional linear subspaces, and these subspaces exhibit relational structure that has not yet been fully characterized by previous work. our approach involves training a sequence of expressive but rank constrained probing models, building on similar exploratory experiments by @xcite and @xcite that study the effects of rank constraints on probe accuracy and complexity. we generalize the use of rank constraints to identify the lowest-dimensional subspace that encodes a task, finding that (1) linguistic features specifically reside in low-dimensional subspaces (section 4) and (2) these subspaces exhibit some degree of hierarchical structure (section 5.1), but (3) are mostly not aligned with individual neurons (section 5.2). additional findings include that linguistic features tend to be encoded in lower-dimensional subspaces in early layers of both elmo and bert, and that relational features (like dependency relations between pairs of words) are encoded less compactly than categorical features like part of speech. as an example of how this kind of information can inform ongoing nlp research, we conclude with a demonstration that discovered subspaces can be used to exert fine-grained control over masked language models: specific linear transformations of the last layer of bert narrowly ablate its ability to model agreement in nouns and verbs (section 6). ## related work the predominant paradigm for interpreting learned word representations is based on training auxiliary probing models to predict linguistic attributes from fixed representations. for standard pretraining schemes, simple probes can successfully recover parts of speech, dependency relations, semantic role labels, coreference links, and named entities @xcite . these features are better encoded by language models than encoders trained for machine translation @xcite . more complex probes target structured labels such as dependency trees @xcite and parser state @xcite . other probes localize linguistic information further: in time, as a function of training dynamics @xcite , and in space, e.g. as a function of individual hidden units @xcite . see @xcite for a detailed survey. our work falls into the latter category: we aim to identify global, linear structure in representation spaces. some previous work has identified surprising non-linear geometry in word embeddings @xcite , but other work suggests that simple linear subspaces encode meaningful attributes like gender information @xcite @xcite and word sense information @xcite . we extend this account to a variety of other linguistic features, identify relations among feature subspaces themselves, and show that these subspaces are causally linked to model behavior. more recent work has focused on addressing shortcomings of the probing paradigm itself. hewitt and liang (2019) design non-linguistic control tasks to benchmark how selective probes are for linguistic information, and show that high probe accuracy is sometimes attributable to task simplicity rather than explicit representation of linguistic content. voita et al. (2019), voita and titov (2020) and pimentel et al. (2020b ) describe an informationtheoretic approach to probing. @xcite in particular argue that measurements of probe quality based on description length characterize encodings more precisely than raw selectivity measurements, showing that part-of-speech tags and dependency edges can be extracted with a smaller code length than controls. @xcite and @xcite further question the use of probe accuracy, finding that language models encode linguistic information even when it is not causally linked to task performance. in contrast, this paper aims to characterize what information is encoded in any form in lowdimensional representation subspaces, regardless of the complexity of the encoding. ## method consider a corpus of words @xmath0, . . . , w n ) with d-dimensional word representations (r 1 , . . . , r n ) written as a matrix @xmath1. a probing task on w is defined by a mapping f (w ) from words to discrete labels (part of speech tags, head-dependent relations, semantic roles, etc.) we say the representations encode the task if there exists a probe g(r) that predicts f (w ) according to a score s(g(r), f (w )), which we choose to be held-out accuracy. there are other choices for s, e.g. selectivity @xcite and mdl @xcite , but these measure different aspects of probe complexity, whereas we are interested in the existence of any information in subspaces of r that is predictive of f (w ). we aim to find the lowest-dimensional subspace r d of r d within which the representations r encode the linguistic task f (w ). if π is a projection from r d to the lower-dimensional r d , we say the subspace encodes the task if the optimal probe predicting task labels from πr is approximately as accurate as the optimal probe predicting them from r. formally, given a tolerance α, we seek the smallest positive integer @xmath2: in practice, we cannot optimize directly for d via gradient descent because matrix rank is neither convex nor differentiable. instead, we simply enumerate values of d and, for each rank, learn a d-dimensional projection foot_0 jointly with a probe g implemented as a multilayer perceptron. we use the same architecture for all tasks, beginning with an investigation of the value of d in a set of standard probing experiments. the use of an expressive model downstream of the projected representation disentangles questions about representation geometry from the questions about probe capacity considered in previous work. 2 4 are linguistic variables encoded in low-dimensional subspaces? our first set of experiments aims to characterize the subspaces that encode linguistic tasks in contextualized representations. following @xcite , we experiment with two contextualized representations alongside their non-contextual word embeddings, and we probe them on three linguistic tasks and three non-linguistic control tasks. ## conclusions we have described a procedure for probing the linear geometry of word representations. our method identifies low-dimensional subspaces of the representations that encode predefined sets of attributes. we find that these subspaces are smallest when they encode linguistic attributes. the subspaces also exhibit hierarchical structure present in the variables they encode and appear to be distributed across neurons. ablation experiments reveal that bert relies on subspaces with as few as 3 dimensions to make fine-grained part of speech distinctions when enforcing subject-verb agreement. future work might explore richer geometric structure in word representations and its effect on model behavior. ric cistac, tim rault, rémi louf, morgan funtowicz, 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: state-of-the-art natural language processing. in proceedings of the 2020 conference on empirical methods in natural language processing: system demonstrations, pages 38-45, online. association for computational linguistics. kelly zhang and samuel bowman. 2018. language modeling teaches you more than translation does: lessons learned through auxiliary syntactic task analysis. in proceedings of the 2018 emnlp workshop blackboxnlp: analyzing and interpreting neural networks for nlp, pages 359-361, brussels, belgium. association for computational linguistics.