Paper: Episodic Memory Theory for the Mechanistic Interpretation of Recurrent Neural Networks (ICLR-2024)

Reviewer comment (Theory):
The current analysis is overly limited. It applies solely to a single-layer RNN with linear dynamics. The assertion that any general RNN can be treated as a linear RNN in A.4 is fundamentally incorrect. In reality, a general RNN may not behave near a fixed point, resulting in significant higher-order terms. Furthermore, the analysis pertains to a task, the VARIABLE BINDING TASK, which is notably distinct from a translation task.

Author response:
"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.

Reviewer comment (Novelty):
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.

Author response:
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 .

Reviewer comment (Presentation):
The presentation lacks clarity. Figures 1 and 2 are difficult to comprehend, and their intended message is unclear. Additionally, the mention of 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.

Author response:
“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.

Reviewer comment (Experiments):
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?

Author response:
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.

Reviewer comment (Theory):
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.

Author response:
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.

Reviewer comment (Writing):
What is the meaning of "dual"? is this related to the conjugate transpose computation $EE^{*}$ in Algorithm 1?

Author response:
"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.

Reviewer comment (Theory):
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 $W_{hh}$.

Author response:
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).

Reviewer comment (Theory):
While the mathematical formulation is meaningful, I believe this is not episodic memory (and this proposed method is not a theory of episodic memory). The model and mathematical framework may be inspired by episodic memory models, but the task, objective and entire narrative is largely unrelated to episodic memory. RNNs receiving variable inputs, performing computations based on those variable inputs, and subsequently producing an output is straightforward decision-making or information-processing in many cases, including this work. This is further supported by Figure 1, where performing an addition operation is merely a simple computation that does not even require any memory storage or retrieval. Knowing how to add is not episodic memory. I am open to discussion if the authors still feel it is correctly defined.

Author response:
(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.

Reviewer comment (Writing):
Right now a reader requires knowledge in RNNs and their applications in neuroscience, as well as some experience in mathematical tools commonly found in modern physics to fully understand the work, when in reality this paper could be written to a pure RNN audience without the Dirac and Einstein notations (at least in the main text), by simply stating that information about the latent is being carefully tracked by a basis transformation and formulating the equations in that context.

Author response:
(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.

Reviewer comment (Experiments):
The effects described in Figure 3 and 4A are specific to the repeat-copy task. More complex tasks (especially those with non-periodic solutions) may result in weight matrices that are not interpretable or offer any additional insight. Similarly, the compose-copy task, which is surprisingly never defined in this paper (but can be inferred from the Figure 4B to be generating the sequence one unit at a time), is the only task that will give rise to such interpretable values. In general, the authors summarize this class of tasks as $f$ in equation (2). 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.

Author response:
(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.