Sparse Interventions in Language Models with Differentiable Masking (2022)

Abstract: There has been a lot of interest in understanding what information is captured by hidden representations of language models (LMs). Typically, interpretation methods i) do not guarantee that the model actually uses the information found to be encoded, and ii) do not discover small subsets of neurons responsible for a considered phenomenon. Inspired by causal mediation analysis, we propose a method that discovers a small subset of neurons within a neural LM responsible for a particular linguistic phenomenon, i.e., subsets causing a change in the corresponding token emission probabilities. We use a differentiable relaxation to approximately search through the combinatorial space. An L_0 regularization term ensures that the search converges to discrete and sparse solutions. We apply our method to analyze subject-verb number agreement and gender bias detection in LSTMs. We observe that it is fast and finds better solutions than alternatives such as REINFORCE and Integrated Gradients. Our experiments confirm that each of these phenomena is mediated through a small subset of neurons that do not play any other discernible role.

## introduction the success of language models (lms) in many natural language processing tasks is accompanied by an increasing interest in interpreting and analyzing such models. one goal in this direction is to identify how a model employs its hidden representations to arrive at a prediction @xcite . a popular line of research studies lms with "diagnostic classifiers" or "probes" that are trained to predict linguistics properties from hidden units, with the purpose of analyzing what information is encoded by the network and where @xcite @xcite . ## related work the l 0 regularization was proposed by @xcite in the context of pruning neural network weights and biases. it has been used in a variety of works in nlp as a tool for generating rationales and attribution @xcite @xcite . masking weights and groups of weights was also used by @xcite to investigate the functional modularity of neural networks. studies suggested that some of the linguistic phenomena are encoded, at least to a large degree, in a disentangled and sparse fashion. for example, @xcite detected a neuron encoding sentiment polarity and @xcite showed that individual facts learned by an lm can be manipulated by modifying a small number of neurons. in a similar spirit, @xcite observed that many transformer attention heads in a neural machine translation model are specialized; interestingly, they also used l 0 regularization but only to prune less important heads; the roles played by the heads were identified by studying their attention patterns. our technique can facilitate such studies by effectively identifying sets of neural network's subcomponents playing a given function. @xcite use different kinds of correlations between neurons from different models to measure their importance. the authors find that many individual neurons capture common linguistic phenomena, also showing how to control translations in predictable ways by modifying their activations. similarly to @xcite , the work of @xcite instead focuses on models' preferences for grammatical inflections, as well as whether neurons process subject-verb agreement. the authors include causal mediation analysis in their methodology. conversely, @xcite criticize recently proposed methodologies for analyzing individual neurons in lms. in particular, they discuss methods that rely on an external probe to rank neurons according to their relevance to some linguistic attribute. they indicate two main pitfalls: 1) these methodologies confound probe quality and ranking quality, and 2) they focus on encoded information rather than information that the model uses. their analysis does not apply to ours since we do not use probes explicit. finally, we refer the reader to @xcite for a recent survey of neuron-level interpretation of nlp models, which includes methods to discover neurons, evaluation methods, significant findings and future research directions. ## method we investigate if we can find groups of neurons for which a modification of their value -which we call an intervention -systematically leads to a change of probability for the single token emission related to a specific phenomenon. because there is no direct supervision for interventions, we need to learn them with a proxy objective. let's assume we have an autoregressive model (e.g., an lstm; hochreiter and schmidhuber 1997) that assigns a probability to sequences. for a set of input tokens @xmath0, . . . , x n , we obtain the model's probability of the token of interest p(x n |@xmath1, . . . , h n where h @xmath2. we then intervene in the model's computation by modifying a group of neurons from one or multiple hidden states. the intervention at a certain time step @xmath3, 1} k indicating which hidden units need intervention and which can be left unchanged. the intervention is then made substituting the ith hidden state with the altered state in addition, as the main objective of this work, we are looking for sparse interventions, which we define as finding a defined small percentage (e.g., 1-5%) of neurons where to apply an intervention to while keeping all the rest untouched. ## experimental setting we study the pre-trained lstm language model made available by @xcite foot_1 , which has been studied extensively in previous works and therefore provides a good testing ground for our method. the studied model is a standard twolayered lstm with a hidden dimension of 650. the embedding layer also has dimensionality 650, and it is not tied with the output layer. the vocabulary size is 50, 000 and the model was trained on english wikipedia data (with around 80m tokens training tokens and 10m for validation). we used this model to compare to previous findings of @xcite . we also pre-train this lm several times with different weights initializations to make sure our results generalize. we study the original model, as well as newly trained models with the same architecture, on two tasks described below: subject-verb number agreement and gender bias. the evaluation for tasks naturally follows the defined objective l ratio ( ĥi , x) (see § 3.1). without intervention, the ratio is @xmath4. thus, we define a successful intervention when we find a mask and baseline values such that the ratio @xmath5. then, we define the accuracy of interventions as the average number of times that the ratio @xmath6. the accuracy thus reflects how often we can flip the model's decision. subject-verb number agreement here, we seek the neurons responsible for predicting the number of verb forms: for a given sentence, we wish the intervention to change the number of the verb from singular to plural or vice versa. for this task, we employ data made available by @xcite foot_2 . the data are synthetic and generated with a modified version from @xcite and @xcite . each synthetic number-agreement instance has a fixed syntax and varied lexical material. sentences were randomly sampled by choosing words from pools of 20 subject/object nouns, 15 verbs, 10 adverbs, 5 prepositions, 10 proper nouns and 10 location nouns. we used a total of 11,000 training sentences and 1,000 evaluation sentences. we apply the single-step intervention to the subject of the (only) verb. we apply two intervention here (i.e., two sets of mask and baseline values): one where we train the model to turn the verb into the singular form and one into the plural one. ## results for the single-step intervention (with regularization), our method achieves 91.5 and 93.9 accuracies for the number agreement and gender bias tasks, respectively. on average, our method finds 5.7 and 5.3 units for the two tasks, respectively. considering that the lm has 1,300 hidden units, this intervention is relatively sparse as desired (we @xmath7.41% of the total units). in figure 1 and 2 , we show examples of hidden state activations with original intervention the ranger yelled because 1 0 unit 563 figure 2: activations of four units we intervene on (single step intervention at "ranger") for changing the pronoun (after "because"). and without interventions for both tasks (see appendix a for additional examples). from these figures, we can see that only one time-step is heavily affected (the one of the intervention) while the others are minimally corrupted after that time step. we hypothesize that the model stores the information of number or gender in other units (or in cell states), but the discovered units are the ones responsible for the initialization of such memory. in for the every-step intervention, our method achieved an almost perfect accuracy of 95.8 and 99.9 for the number agreement and gender bias tasks, respectively, while using 3 units or less on average for both tasks. this type of intervention is much more effective and more intrusive-the number of changes is larger as it happens at every step). in efficiency to demonstrate the efficiency and efficacy of our estimation employing the hard concrete distribution, we compare to the standard score function estimation (aka rein-force; williams 1992) with a moving average baseline for variance reduction @xcite and trying different values of α to achieve a good trade-off between accuracy and number of units used. we also compare to integrated gradients @xcite where we intervene on the top-k influential neurons by setting them to zero. in robustness to demonstrate that our method is robust, we tested it on 5 language models initialized with different seeds and trained with the original script by @xcite . we run our method for the single-step intervention 3 times for each language model. the average accuracy at convergence is 88.7 ± 2.6, and we discover 4.7 ± 0.5 units on average. ## conclusions in this work, we present a new method that employs constraint optimization to efficiently find hidden units that are responsible for particular language phenomena of language models. we use an l 0 regularization to find a sparse solution-., our methodology discovers few units in the order of 2-6 that @xmath8.41% of all units in the studied lm. we show such sparse solutions can be found for multiple phenomena (number and gender agreement) and is an useful tool for analysis of what a lm has learned and how units influence its token emissions. although this work focuses on lstm models, the proposed technique is not architecture-dependent and thus easily applicable to transformers, convolutionbased models and many others.