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

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.

• The paper provides a counterexample to the strong LRH by revealing that RNNs can store sequences using non-linear, magnitude-based representations.
• The paper shows that smaller RNNs predominantly use non-linear magnitude solutions, while larger networks exhibit a mix of linear and non-linear encoding.
• The paper validates its claims with targeted interventions achieving about 90% accuracy, emphasizing the need to broaden current interpretability frameworks.

Recurrent Neural Networks Learn to Store and Generate Sequences Using Non-Linear Representations

The paper "Recurrent Neural Networks Learn to Store and Generate Sequences Using Non-Linear Representations" challenges the strong version of the Linear Representation Hypothesis (LRH) through empirical findings in gated recurrent neural networks (RNNs). The authors demonstrate that RNNs are capable of encoding sequence information through non-linear, magnitude-based representations, which they refer to as `onion representations'. This is in contrast to the strong LRH, which posits that neural networks encode all features as linear directions in their representation spaces.

Key Findings and Contributions

1. Counterexample to Strong LRH:
   • The paper provides a detailed counterexample to the strong LRH, showing that when RNNs are tasked to repeat a sequence of input tokens, they frequently resort to encoding these sequences using magnitudes rather than directions in the activation space. Specifically, in smaller RNNs, the sequence position of tokens is stored with varying orders of magnitude that are not linearly separable.
2. Layered Non-Linear Representations:
   • The RNNs learned layered representations, where the hidden state at each time step embodies the previous layers, making it impossible to isolate distinct features in simple linear subspaces. This is manifested in how the smallest RNNs (48, 64 units) almost exclusively utilize magnitude-based solutions, while larger RNNs (128, 512, 1024 units) tend to fit within the LRH framework but still maintain compatibility with these non-linear mechanisms.
3. Experimental Validation:
   • Through a series of carefully designed interventions, the authors validate their hypotheses. They learn the scaling factor associated with each sequence position and demonstrate interventions with approximately 90% accuracy, showcasing the presence of magnitude-based features in the model's hidden states. These results suggest that the scope of current interpretability research needs to extend beyond the confines of the LRH.

Implications and Speculative Outlook

The findings pose significant implications for the field of AI interpretability:

• Broadened Scope of Interpretability Research:
  • By demonstrating that non-linear, magnitude-based representations can be fundamental in certain RNN models, the paper urges the research community to explore beyond the linear paradigms. This could pave the way to novel interpretability methods that account for complex mechanisms underlying neural models.
• Impact on Model Design and Analysis:
  • The insights gained from this paper may inform the design of future neural architectures, particularly in the context of sequence modeling. Understanding that small models with limited parameters might develop fundamentally different encoding strategies compared to larger models could be crucial for applications requiring precise and interpretable behavior.
• Potential for New Mechanisms in Complex Tasks:
  • While the paper focuses on relatively simple sequence tasks, the mechanisms identified, especially in smaller networks, could manifest in more intricate scenarios such as large-scale LLMs or structured state-space models. This underscores the need for continuous re-examination of our assumptions about neural representation as model complexities evolve.

Conclusion

In conclusion, the paper robustly challenges the strong LRH by empirically demonstrating that RNNs, when trained on sequence tasks, can employ non-linear, magnitude-based encoding strategies—a significant departure from purely linear representations. These findings not only contest existing interpretability paradigms but also open avenues for novel analyses and methods that can further our understanding of neural network behavior in complex settings. The consistent observation of non-linear representations in small RNNs also hints at a broader underlying complexity in neural mechanisms that warrants deeper exploration and may influence future AI system designs.