Episodic Memory Theory for the Mechanistic Interpretation of Recurrent Neural Networks

We thank the reviewer for the comments.

"The current analysis is overly limited. ..."

Our analysis is restricted to linear RNNs, but linearization approaches allow researchers to study complex non-linear dynamical systems like RNNs. What we described in A.4 is only one approach to linearization typically used in the literature. Alternate approaches, like the Koopman theory, allow a more complicated linearization while maintaining applicability in regions farther from fixed points. Our theory extends to these spaces, although we find in the paper that the variable binding tasks are encoded in a simple inner product (the finite-dimensional real vector inner product) space.

“, the analysis pertains to a task, the VARIABLE BINDING TASK, which is notably distinct from a translation task”

The class of VARIABLE BINDING TASKS is analogous to a deterministic version of s-gram tasks used in NLP. We only restrict the composition function to be linear. Future work will explore more complex non-linear composition functions.

1. The current analysis lacks novelty. The article primarily employs linear projection techniques to examine the components of RNN weights, a methodology that has been in use for a considerable period.

In our restriction of the VARIABLE BINDING TASKS, we find that a linear basis transformation reveals the operator. To our knowledge, an explicit operator that enables memory storage mechanisms in RNNs has not yet been proposed and experimentally validated. We further showed that the operator is not just a theoretical construct but emerges naturally after training. We also want to note that Fig 4,5 has not been extracted from prior RNN analyses.

Simple tasks enable researchers to probe the mechanisms of complex dynamical systems like RNNs. For instance, simple tasks like 3-bit flip flop, Frequency cued sine wave, Context-dependent integration are used to probe various dynamical behaviors of RNNs [3]. Our tasks are more complex than these as the dynamical behavior for VARIABLE BINDING is very high dimensional and inscrutable (Appendix Figure 5 shows how high dimensional and diverse the Jacobian spectrum is after linearizing around the origin).

Among the VARIABLE BINDING tasks, repeat copy ($T_1$) has been used frequently to analyze the memory storage behavior of RNNs [1]. Recently, the operator we proposed for repeat copy was also found in connection to traveling waves in RNNs [2], suggesting the theory's applicability is broader than the VARIABLE BINDING tasks.

[1] Alex Graves, Greg Wayne, and Ivo Danihelka. Neural turing machines. ArXiv, abs/1410.5401, 2014 [2] Thomas Anderson Keller, Lyle E. Muller, Terrence J. Sejnowski, and Max Welling. Traveling waves encode the recent past and enhance sequence learning. ArXiv, abs/2309.08045, 2023. URL https://api.semanticscholar.org/CorpusID:262013982 [3] Maheswaranathan, Niru et al. “Universality and individuality in neural dynamics across large populations of recurrent networks.” Advances in neural information processing systems 2019 (2019): 15629-15641 .

“Additionally, mentioning GSEMM seems unnecessary since the current work solely involves a simple linear RNN. There is a lack of a concise summary of the core mechanisms of the RNN.”

GSEMM is not unnecessary, as this connection to neurocomputational memory models enabled interpreting the weights of the RNN as stored memories and their interactions. To our knowledge, this is the first connection between Hopfield-like memory models studied in theoretical neuroscience and Recurrent Neural Networks.

"There is a lack of a concise summary of the core mechanisms of the RNN."

Section 5 and Figure 2 summarize the main parts of the theory. Please let us know if any particular aspects of the figures and exposition are unclear.

We thank the reviewer for the comments and for pointing out the typo.

• Only a Repeat Copy task is shown to reveal the stored information in hidden neurons. As the authors mentioned in section 7.3, there are some cases in which the computed basis is converged, but it cannot give interpretable representations. In other words, under which tasks do we expect this framework to fail?

We have added a section in the Appendix A.6 for a more formal analysis of the error in approximating variable memories. In short, the algorithm assumes specific interactions of the variable memory are negligible. This assumption is valid for the cases of repeat copy and compose copy we showed in the paper, but for some $f$ where this assumption is broken, the algorithm fails. We want to note that the failure of the algorithm does not mean that the theoretical mechanisms are not present (the convergence of the theoretical $\Phi$ in Table 1 still stands). This, however, means that it is not trivial to formulate an algorithm to find such a human interpretable basis for a general class of tasks.

• This deviation from the theory is a result of the sensitivity of the basis definition to minor errors in the pseudo-inverse required to compute the dual.

We were incorrect on this point in the paper. More recent analysis has revealed that the deviation results from the errors accumulated when approximating the variable memories by power iteration. A section has been added to Appendix A.6 that formally explains the approximation error of the variable memory algorithm.

• What is the meaning of "dual"? is this related to the conjugate transpose computation ��∗ in Algorithm 1?

"dual" is a misnomer. We meant the projection operator that extracts the components in the variable memory space into the standard basis. This is updated in the manuscript.

• Could you please discuss in detail the difficulties of applying the proposed theory to nonlinear RNNs? In Appendix A.4, you have shown that a nonlinear RNN has a similar form of the linear system as linear RNNs instead of a different �ℎℎ.

The proposed algorithm (not the theory) does not work for non-linear RNNs when any of the two assumptions fail – (1) The geometry of the representation is not easily captured by a linear basis, in which case a more complex set of basis vectors needs to be taken, (2) If the RNN dynamics are far from the fixed point, which means that the Taylor series expansion in Appendix A.4 needs to account for the higher order terms, even if the correction from the Jacobian is taken into account. (1) and (2) are not mutually exclusive; sometimes failure of (2) implies the failure of (1).

“The notion of storing different items in different linear subspaces has also been explored”

We thank the reviewer for the relevant literature. Empirical works have investigated the linear subspaces we presented, mainly from the neuroscience community. To our knowledge, there is no work formalizing these linear spaces, how they interact, and how the two (space and interactions) solve the task at hand. Using the new formalism, we made theoretical predictions on what class of operators could solve the tasks and validated them with experimental results. Note that our theory goes beyond the task T_1 that is analogous to the delayed sequence reproduction task considered in the linked paper on sequence working memory.

“This is obviously not the first time people have studied how neurons process information in RNNs (see the work of Omri Barak and David Sussillo, e.g.).”

Our work builds on the Jacobian spectrum analysis approach of Omri Barak and David Sussilo. This connection is discussed in the Related Works section. To our knowledge, Jacobian spectrum research has not explicitly delved into the phenomenon of variable binding and defined a notion of storing variables in subspaces of the hidden state.

Further, the tasks used in the previous studies using spectral interpretations (3-bit flip flop, Frequency-cued sine wave, Context-dependent integration) are intentionally low-dimensional, making it easy to plot and draw conclusions. The variable binding tasks in our paper exhibit very high-dimensional behaviors with complicated spectrum distribution (Figure 5 in the appendix), which means that spectral analysis will not yield easily understandable results. We showed in the paper that all these complicated spectral distributions have a straightforward interpretation with variable memories and their interactions albeit restricted to the variable binding tasks in Section 4. We have made changes to clarify these points in the manuscript.

"What are the tasks? "

We have amended the document with the matrix representation of the composition function for these tasks with a figure above Table 1. They all differ only by the composition function $f$ acting on all the variables.

"I thought u(t) was the input vector, which is 0 for t>s, yet Eqn 2 shows the evolution of u(t) for t>s."

Thank you for pointing this out. Corrected: Eqn 2 is supposed to be y(t) – the network output rather than u(t).

"Why are they focusing on mechanistic interpretability for such a limited behavior? As I understand it, MI can be used to explain any behavior of a neural network."

There are different levels of explainability that MI is used for. Sussilo proposed a very general method, which can be applied to any behavior of neural networks. However, this approach makes interpreting the variable binding tasks in our work challenging. Specifically, the high dimensional nature of the variable binding tasks meant no straightforward interpretation of the Jacobian spectrum.

Our method is specific to the class of variable problems described in Section 4. It provides a better understanding of the RNN behavior beyond what the Jacobian spectrum allows in these tasks. With this specificity, we better understand RNN behavior on these tasks but lose some generality to other tasks. There is a clear gap between what the general method enables in these tasks that the current approach fills. Although the tasks we consider are still simple, repeat copy has been used to test the memory capabilities of recurrent neural networks [1] and more recently found in connection with traveling waves in RNNs [2].

We have clarified these points in the updated paper.

“On the whole I also don't see the specific value in claiming that this analysis is related to episodic memory. Sequential memory, yes. But there is nothing specifically episodic about the motivation for the analyses and the tasks represent serial working memory.”

This analysis is episodic in the sense that it uses an episodic memory model from computational neuroscience and the formalisms (stored memories, inter-memory interactions) it provides to perform the analysis. The definition is in line with GSEMM, and episodic memory as described in neuroscience [3] - which describes the ability of neural networks to store and process temporal and contiguous memory sequences.

[1] Alex Graves, Greg Wayne, and Ivo Danihelka. Neural turing machines. ArXiv, abs/1410.5401, 2014

[2] Thomas Anderson Keller, Lyle E. Muller, Terrence J. Sejnowski, and Max Welling. Traveling waves encode the recent past and enhance sequence learning. ArXiv, abs/2309.08045, 2023. URL https://api.semanticscholar.org/CorpusID:262013982

[3] Umbach, Gray et al. “Time cells in the human hippocampus and entorhinal cortex support episodic memory.” Proceedings of the National Academy of Sciences of the United States of America 117 (2020): 28463 - 28474.

We thank the reviewer for the insightful comments and for pointing us to the relevant literature.

"The tasks (insofar as they are described) seem like weak, or at least very specific, tests of variable binding. ..."

Without the linear assumption, the variable binding tasks can represent any dynamic system with a historical context of s states or, in NLP terms, an s-gram model. Probablistic versions of s-gram models have been used in the literature before; we take a restricted $f$ for better analytic tractability.

Variable binding is a very diverse phenomenon, and tackling all of its aspects in a single paper is intractable. Our class of tasks is a simple restriction that can be used to probe some mechanisms of variable binding. Although simple, these tasks reveal the mechanisms of RNNs and contribute to developing a better, more complete theory of variable binding.

We also note that simple tasks are used liberally in the literature to improve our understanding of RNN computations. Starting with small-scale tasks on standard architectures yields alternative benefits, including interpretable results and potential generality of conclusions.

Regarding the impact of the theory beyond what we showed, we want to bring to attention a recent work studying traveling waves in RNNs that found the operator for $T_1$ independently (https://arxiv.org/pdf/2309.08045.pdf). Our work accommodates more general cases that $T_1$.

“a novel perspective on Recurrent Neural Networks (RNNs), framing them as dy- namical systems performing sequence memory retrieval”

Previously, RNNs were used to study memory, but to our knowledge, memory models were not used to explore the behavior of RNNs, and we showed in the paper that this view has benefits.

We showed that RNNs are discrete-time analogs of a sequence memory model GSEMM (a newer class of Hopfield Networks with precisely defined memories and inter-memory interactions). To our knowledge, this is the first mathematically rigorous connection between RNNs and neurocomputational memory models. In this paper, this connection enabled reinterpreting the learned weight matrix of RNNs as stored memories and their interactions. Future work can flesh out the relationship between RNNs and the energy function in memory models.

We thank the reviewer for deeper insights and for constructive feedback to improve the clarity of the work.

(1) The definition of Episodic Memory is stipulative: that is, we took an episodic memory model (GSEMM) and showed that RNNs are equivalent to its discretization. Here, episodic memory is defined as the ability of a neural system to store episodes (or distinct sequences of memories). In the context of the paper, we showed it was useful to interpret RNNs using this notion of episodic memory from the memory modeling literature. There are, of course, other definitions of episodic memory. For instance, episodic memory is defined originally as the subjective recall of personal experiences, focusing on the human experience rather than the underlying neural network. Given the variety of definitions for episodic memory, we don’t think it is fruitful to formalize a single definition encapsulating all aspects of episodic memory. Further, such a precise definition is not necessary for the technical contributions of the paper. We take the point that the addition example in Fig 1 is misleading. We have replaced that example with the generation of the Fibonacci series, which necessitates an explicit variable memory store.

(2) Thank you for the comments to improve the clarity, we have modified our document clarifying the points of confusion and removing the notation. The intention was that the model of variable binding is defined in the more general vector space – the Hilbert space. We don’t assume any geometry (the definition of the inner product) on the neural representations to formulate the theory. The theory will work irrespective of the geometry of the neural networks representing the solution (Appendix A.3). It so happens that the tasks and RNNs we considered in the paper all represent the solution in a space where the inner product is the vector inner product, and hence, linear algebra can be used. When formulating a theory of neural computation, the rich representational diversity of neural networks needs to be considered, and this notation enables that. We accept that the notation is not standard in the general literature. We take the point that the new notation needs to be better motivated. Given that we found only a simple geometry in these tasks, we have deferred introducing it to future work (and removed the notation from the main document in the updated version). Still, given the utility of the notation to form principled theories for how RNNs work, we don’t see why it cannot be used to inform future developments.

(3) Figure 3, 4A are results of running the approximation algorithm which can have errors (added a section in the Appendix with formal error analysis). The effect described in the paper is present across all the 4 tasks we considered in the paper (Table 1). Since the variable memory algorithm is approximate, sometimes correct bases cannot be obtained by power iteration, which is why only some tasks can be supported by the empirical algorithm. Figures 3, 4 show the result of running the algorithm to approximate the variable memories, which may have errors when applied in specific tasks. We only show 3 for repeatCopy as the variables there are easier to comprehend than the other composition functions.

"The intended narrative is that the authors are using a class of tasks to elegantly highlight the feature of the proposed method, but my impression is that the method will not be producing anything meaningful beyond this class of tasks."

We want to bring into attention a recent work, that found the operator to solve repeat copy (Fig 4A) in connection to the presence of traveling waves in RNNs (https://arxiv.org/pdf/2309.08045.pdf), suggesting a broader applicability of the theory. Starting with small-scale experiments on standard architectures yields alternative benefits, including interpretable results and potential generality of conclusions. Future work will extend it to more general tasks of interest to the machine learning community.

There is always the added caveat in analyzing RNNs that behaviors like chaos and quasi-periodicity cannot be captured analytically. This is a disadvantage for any analysis of RNNs and not unique to our approach. That said, it should still be interesting that RNNs capture the inherent periodicity of the problems they solve in their dynamical behavior without introducing any additional implicit biases.

We added the 4 task functions as a figure above the Table 1. All these tasks differ only in the definition of $f$, the function computing the final variable memory.