Recurrent Neural Networks Learn to Store and Generate Sequences using Non-Linear Representations (2024)

Abstract: The Linear Representation Hypothesis (LRH) states that neural networks learn to encode concepts as directions in activation space, and a strong version of the LRH states that models learn only such encodings. In this paper, we present a counterexample to this strong LRH: when trained to repeat an input token sequence, gated recurrent neural networks (RNNs) learn to represent the token at each position with a particular order of magnitude, rather than a direction. These representations have layered features that are impossible to locate in distinct linear subspaces. To show this, we train interventions to predict and manipulate tokens by learning the scaling factor corresponding to each sequence position. These interventions indicate that the smallest RNNs find only this magnitude-based solution, while larger RNNs have linear representations. These findings strongly indicate that interpretability research should not be confined by the LRH.

## introduction it has long been observed that neural networks encode concepts as linear directions in their representations @xcite , and much recent work has articulated and explored this insight as the linear representation hypothesis (lrh; @xcite @xcite @xcite . a strong interpretation of the lrh says that such linear encodings are entirely sufficient for a mechanistic analysis of a deep learning model @xcite . in this paper, we present a counterexample to the strong lrh by showing that recurrent neural networks with gated recurrent units (grus; @xcite ) learn to represent the token at each position using magnitude rather than direction when solving a simple repeat task (memorizing and generating a sequence of tokens). this leads to a set of figure 1 : we find that grus solve a repeat task by learning a scaling factor corresponding to each sequence position, leading to layered onion-like representations. in this simplified illustration, the learned token embeddings (a) are rescaled to have magnitudes proportional to their sequence positions (b). to change an element of the sequence, remove (c) and replace (d) the token embedding at the given positional magnitude. the layered nature of the representations makes them non-linear; any direction will cross-cut multiple layers of the onion. we made this surprising finding in a hypothesisdriven fashion. our hypothesis 1 was that grus would store each token in a linear subspace. to test this hypothesis, we employed a variant of distributed alignment search (das; @xcite ) that uses a gumbel softmax to select dimensions for intervention. this revealed that the larger grus do in fact have linear subspaces for each position, but we found no evidence of this for the smaller ones (section 5). this led to hypothesis 2: grus learn to represent input bigrams in linear subspaces. a das-based analysis supports this for the medium-sized models but not for the smallest ones (section 6). this left the task success of the smallest models to be explained. for the smallest models, we observed that the update gates of the grus got gradually lower as the sequence progressed. this led to hypothesis 3: onion representations. to evaluate this hypothesis, we learned interventions on the hidden vector encoding a sequence of tokens that replaces token a with token b at position j. the intervention adds the scaled difference of learned embeddings for a and b, where the scaling factor is determined by the position j with learned linear and exponential terms. across positions, this intervention works with ≈90% accuracy, demonstrating the existence of layered features stored at different scales. the existence of non-linear representations is a well-formed theoretical possibility. for example, under the framework of @xcite and @xcite , any bijective function can be used to featurize a hidden vector, and interventions can be performed on these potentially non-linear features. however, the typical causal analysis of a neural networks involves only interventions on linear representations (see section 2 for a brief review of such methods). we hope that our counterexample to the strong version of the lrh spurs researchers to consider methods that fall outside of this class, so that we do not overlook concepts and mechanisms that our models have learned. ## related work the linear representation hypothesis much early work on 'word vectors' was guided by the idea that linear operations on vectors could identify meaningful structure @xcite @xcite . more recently, @xcite articulated the linear representation hypothesis (lrh), which says that (1) features are represented as directions in vector space and (2) features are one-dimensional (see also @xcite @xcite . @xcite by showing some features are irreducibly multi-dimensional. @xcite subsequently argued that (1) is the more significant aspect of the hypothesis, and it is the one that we focus on here. @xcite adds important nuance to the lrh by distinguishing a weak version (some concepts are linearly encoded) from a strong one (all concepts are linearly encoded). our concern is with the strong form; there is ample evidence that linear encoding is possible, but our example shows that other encodings are possible. in onion representations, multiple concepts can be represented in a linear subspace by storing each concept at a different order of magnitude, i.e., a 'layer' of the onion, and any direction will cross-cut multiple layers of the onion. intervention-based methods recent years have seen an outpouring of new methods in which interventions are performed on linear representations, e.g., entire vectors @xcite @xcite , individual dimensions of weights @xcite and hidden vectors @xcite @xcite , linear subspaces @xcite @xcite , or linear features from a sparse dictionary @xcite . these methods have provided deep insights into how neural networks operate. however, the vast and varied space of non-linear representations is woefully underexplored in a causal setting. rnns recurrent neural networks (rnns) were among the first neural architectures used to process sequential data @xcite . many variants arose to help networks successfully store and manage information across long sequences, including lstms @xcite and grus @xcite . bidirectional lstms provided the basis for one of the first large-scale pretraining efforts (elmo; @xcite . with the rise of transformer-based models @xcite , rnns fell out of favor somewhat, but the arrival of structured state-space models @xcite @xcite has brought rnns back into the spotlight, since such models seek to replace the transformer's potentially costly attention mechanisms with recurrent connections. we chose grus for our studies, with an eye towards better understanding structured state space models as well. ## models in this paper, we focus on how rnns solve the repeat task. as noted in section 2, this question has taken on renewed importance with the development of structured state-space models that depend on recurrent computations and are meant to provide efficient alternatives to transformers. define an rnn as h @xmath0, x t ), h 0 = 0, where f (•, •) is the state update function, @xmath1, . . . , t } is the current timestep, x @xmath2, and h @xmath3. the output of the model is y @xmath4. vectorized inputs x t are obtained with a learned embedding @xmath5, using the indexing operator x @xmath6, where i @xmath7, . . . , n s } is the index of the token at timestep t. in our experiments, we use gru cells over the more widely-used lstm cells because they have a single state to intervene on, as opposed to the two states of the lstm. gru-based rnns defined as: for output generation, we use g(h t ) = softmax(h t w o + b o ). the learned parameters are weights w * , u * ∈ r n ×n , and biases b * ∈ r n . we will investigate how the final hidden state h l of a gru represents an input token sequence @xmath8, i 2 , . . . i l . the final state is a bottle-neck between the input token sequence and the output. ## repeat task experiments our over-arching research question is how different models learn to represent abstract concepts. the repeat task is an appealingly simple setting in which to explore this question. in this task, the network is presented with a sequence of random tokens @xmath9, i 2 , . . . , i l , where each i j is chosen with replacement from a set of symbols n s and the length l is chosen at random from {1 . . . l max }. this is followed by a special token, i l+1 = 's', that indicates the start of the repeat phase. the task is to repeat the input sequence: y l+1+@xmath10. the variables in this task will represent positions in the sequence and take on token values. as a preliminary step, we evaluate rnn models on the repeat task. the core finding is that all of the models solve the task. this sets us up to explore our core interpretability hypotheses in sections 5-7. ## hypothesis 1: unigram variables intuitively, to solve the repeat task, the token at each position will have a different feature in the state vector h l (the boundary between the input and output phrases). in line with the lrh, we hypothesize these features will be linear subspaces. ## hypothesis 2: bigram variables our second hypothesis is a minor variant of hypothesis 1. here, we posit that, instead of representing variables for unigrams, the model instead stores tuples of inputs (i t , i t+1 ) we call bigram variables. ## hypothesis 3: onion representations in an effort to better understand how the smallest grus solve the repeat task, we inspected the gate values z t as defined in equation 1 from the gru definition (section 3). figure 2a visualizes the first 64 input gates for the @xmath11. the x-axis is the sequence (temporal dimension) and the y-axis depicts the gate for each dimension. one can see that this model uses gates to store inputs by closing position-dependent channels sharply, creating a position-dependent subspace for each input. (this gating pattern is consistent across all inputs.) a b 1 c a d f 2 s a b 1 c a d f i t z t [j] model. here, the picture looks substantially different. this model gradually closes its gates simultaneously, suggesting that the network might be using this gate to encode token positions. this led us to hypothesis 3: rnns learn to encode each position in a sequence as a magnitude. this hypothesis relies heavily on the autoregressive nature of the gru, the discriminative capacity of the output classifier g(h t ), and the sequential nature of the problem. multiple features can be stored in the same subspace, at different scales. when the gru begins to generate tokens at timestep @xmath12, if the scales s t ′ associated with position t ′ > t are sufficiently small (s t ′ ≪ s t ), the output classifier y @xmath13. in the following step, i 1 is fed back to the model as an input, and the model is able to remove the scaled representation corresponding to i 1 from h t , obtaining h t+1 . in this new representation, the input with the next largest scale, i 2 , will be dominant and will be decoded in the next step. this can be repeated to store a potentially long sequence in the same subspace, limited by the numerical precision. we call these 'onion representations' to invoke peeling back layers corresponding to sequence positions. hypothesis 3 falls outside of the lrh. in lin-ear representations, tokens are directions and each position has its own subspace. all positions are independently accessible; tokens can be read-out and manipulated given the right target subspace. onion representations have very different characteristics. first, tokens have the same direction regardless of which position they are stored in; the magnitude of the token embedding determines the position rather than its direction. as a result, if multiple positions contain the same token, the same direction will be added twice with different scaling factors (see figure 1d where the token c occurs in positions 2 and 3). second, because the memory is the sum of the scaled token embeddings, it is impossible to isolate the position associated with a given scale. only the token with the most dominant scale can be extracted at a given time, by matching it to a dictionary of possible token directions. this is done by the final classifier for our grus. the autoregressive feedback for grus in effect peels off each layer, clearing access to the next variable. appendix f provides a toy implementation of the onion solution to elucidate the underlying concepts. ## discussion and conclusion the preceding experiments show that grus learn highly structured and systematic solutions to the repeat task. it should not be overlooked that two of these solutions (those based in unigram and bigram subspaces) are consistent with the general guiding intuitions behind the lrh and so help to illustrate the value of testing hypotheses in that space. however, our primary goal is to highlight the onion solution, as it falls outside the lrh. our hope is that this spurs researchers working on mechanistic interpretability to consider a wider range of techniques. the field is rapidly converging around methods that can only find solutions consistent with the lrh, as we briefly reviewed in section 2. in this context, counterexamples to the lrh have significant empirical and theoretical value, as olah (2024) makes clear: but if representations are not mathematically linear in the sense described above [in a definition of the lrh], it's back to the drawing board -a huge number of questions like "how should we think about weights?" are reopened. our counterexample is on a small network, but our task is also very simple. very large networks solving very complex tasks may also find solutions that fall outside of the lrh. there is also a methodological lesson behind our counterexample to the lrh. much interpretability work is guided by concerns related to ai safety. the reasoning here is that we need to deeply understand models if we are going to be able to certify them as safe and robust, and detect unsafe mechanisms and behaviors before they cause real harm. given such goals, it is essential that we analyze these models in an unbiased and open-minded way.