Paper: Memoria: Hebbian Memory Architecture for Human-Like Sequential Processing (ICLR-2024)

Reviewer comment (Evaluation):
My biggest concern regarding the proposed architecture is its computational cost. Unless I've missed something obvious, the authors do not quantify the cost of their method (relative to the previously-proposed extensions to transformer models with which they compare their model's performance), which is a key determinant in whether it can prove useful in machine learning.

Author response:
We have revised the paper by incorporating discussions on the computational complexity of Memoria. Theoretically, we have computed and presented the computational complexity of Memoria operations by denoting Big-O notation, while empirically measuring and presenting the corresponding execution times and memory usage. For detailed information, please refer to Appendix E, Algorithm & Computational Complexity section of the revised paper.

Reviewer comment (Experiments):
The computational cost will in turn depend on the values chosen for Memoria's many hyperparameters, which the authors do not seem to systematically sweep.

Author response:
We have revised the paper by incorporating discussions on the computational complexity of Memoria. Theoretically, we have computed and presented the computational complexity of Memoria operations by denoting Big-O notation, while empirically measuring and presenting the corresponding execution times and memory usage. For detailed information, please refer to Appendix E, Algorithm & Computational Complexity section of the revised paper.

Reviewer comment (Evaluation):
Can the authors state precisely how the computational and memory cost of Memoria scales with its various hyperparameters, and quantify the cost incurred in the performance tests shown?

Author response:
We have revised the paper by incorporating discussions on the computational complexity of Memoria. Theoretically, we have computed and presented the computational complexity of Memoria operations by denoting Big-O notation, while empirically measuring and presenting the corresponding execution times and memory usage. For detailed information, please refer to Appendix E, Algorithm & Computational Complexity section of the revised paper.

Reviewer comment (Experiments):
The authors do not justify the hyperparameter choices used in their experiments, nor do they probe how sensitive performance is to those choices. Some exploration of these effects is required.

Author response:
We conducted additional experiments by modifying various hyperparameters to investigate the sensitivity of hyperparameters. The experimental results show that there is generally no significant difference in performance compared to what was initially reported in the paper. There were even cases where the performance improved compared to the original scores reported in the paper.

|  | Wikitext PPL |
| --- | --- |
| Original | 23.471 |
| initial lifespan 9 → 5 | 23.485 |
| lifespan extend scale 8 → 4 | 23.518 |
| ltm search depth 10 → 5 | 23.491 |
| memoria reset period 500 → 100 | 23.407 |
| $N_{wm}$ 100; $N_{stm}^{rem}$ 25; $N_{ltm}^{rem}$ 25 | 23.376 |
| $N_{wm}$ 25; $N_{stm}^{rem}$ 100; $N_{ltm}^{rem}$ 25 | 23.831 |
| $N_{wm}$ 25; $N_{stm}^{rem}$ 25; $N_{ltm}^{rem}$ 100 | 23.670 |

Reviewer comment (Experiments):
In Appendix B.2, the authors write "Furthermore, to prevent potential interference with the learning process, we periodically reset all memory in Memoria every 500 steps during training (1500 steps for enwik8 dataset). This was done to avoid referencing memory generated at stages where learning was insufficient, as it could impede the training progress." This procedure is a substantial alteration of the memory architecture described in Section 3, and is not mentioned in the main text. Can the authors quantify how much performance is degraded if this step is not included?

Author response:
Incorporating periodic memory resets during the training process is not due to inherent deficiencies in our forgetting capabilities. Memoria maintains the exact past representations unchanged regardless of the passage of time unless they expire. However, during training, the unavoidable fluctuation of model weights at each gradient descent step poses a challenge. A model trained for an additional 500 steps will inevitably produce different representations for the same input, introducing noise to Memoria, which operates based on the distance of representations. As it is impractical to infer all past information and rebuild memories based on a new model, we have opted for a simple solution—periodically resetting the memory to eliminate embeddings derived from past models.

Reviewer comment (Theory):
The periodic memory reset procedure suggests that the proposed architecture is unable to gracefully forget low-quality or corrupted memories, which is a substantial limitation.

Author response:
Incorporating periodic memory resets during the training process is not due to inherent deficiencies in our forgetting capabilities. Memoria maintains the exact past representations unchanged regardless of the passage of time unless they expire. However, during training, the unavoidable fluctuation of model weights at each gradient descent step poses a challenge. A model trained for an additional 500 steps will inevitably produce different representations for the same input, introducing noise to Memoria, which operates based on the distance of representations. As it is impractical to infer all past information and rebuild memories based on a new model, we have opted for a simple solution—periodically resetting the memory to eliminate embeddings derived from past models.

Reviewer comment (Writing):
There are many typos and grammatical errors, which can make parts of the paper harder to parse. I will not give line-by-line comments, but, for instance, the opening sentences of Section 3.3 do not read cleanly. Moreover, the citation to Atkinson and Shiffrin is duplicated, with one instance being corrupted with acknowledgements in the title field.

Author response:
We have reviewed the grammatical aspects of the paper, but during the modification process, it appears that some overlooked issues persisted. We have conducted a comprehensive grammar check on all sections of the paper and addressed the citation issues you mentioned, incorporating them into the revised version of the paper.

Reviewer comment (Novelty):
The authors motivate their memory architecture in terms of interest in Hebbian learning rules, but they do not mention what is perhaps the most impactful memory model based on Hebbian learning: the Hopfield network. For completeness, Hopfield's original work should be cited, along with its recent extensions (Krotov & Hopfield 2016) and relationship to transformer attention (Ramsauer et al., 2020).

Author response:
Hopfield network is based on Hebbian mechanisms for modeling associative memory and became integrable into deep learning as Ramsauer proposed a differentiable structure. The Hopfield network, however, differs slightly from our goal of improving long-term memory for extended sequences. Nonetheless, it shares some theoretical background with our research. In the revised version of our paper, we have incorporated the research mentioned, along with additional research on Hebbian learning and associative learning, into the related work section.

Reviewer comment (Presentation):
It would be useful to include smaller versions of Figures 6 and 7 in the main text.

Author response:
The figures were initially intended for inclusion in the main text. However, due to space constraints and the fact that the architecture primarily illustrates a method of interfacing with Memoria rather than detailing the structure and functions of Memoria itself, we opted to provide a textual description in the main text and included the figures in the appendix, assigning them a relatively lower priority. Upon further review, if it is possible to reduce unnecessary elements or conserve space, we will certainly consider including a condensed version of the figure in the main text.

Reviewer comment (Reproducibility):
The main problem when reading this submission is insufficient details on the proposed memory module and intuition on why this specific instantiation of memory is chosen. Even after reading the paper several times (including the appendices), there is still no clear understanding of what is actually done in this paper. The main request to the authors is to expand section 3 (possibly in appendix C) on the mathematical implementation of their architecture so that each operation performed by their memory module is clearly defined in the text.

Author response:
In the revised version of the paper, we have presented the operations of each memory operation stage in algorithmic form in Appendix E, Algorithm & Computational Complexity section. If you review Section 3, Memoria, along with this portion, it should facilitate a better understanding.

Reviewer comment (Reproducibility):
Q1. Could the authors please expand section 3 or appendix C with step-by-step operations that are performed in the memory module both during training and inference?

Author response:
In the revised version of the paper, we have presented the operations of each memory operation stage in algorithmic form in Appendix E, Algorithm & Computational Complexity section. If you review Section 3, Memoria, along with this portion, it should facilitate a better understanding.

Reviewer comment (Presentation):
Q2. It is unclear what is illustrated in Figure 2 (remind process). Could the authors please explain in the text how this is done in their architecture?

Author response:
Figure 2 illustrates the remind stage of Memoria. This stage involves selecting memory information from long-term memory to be utilized in model inference. The key to this selection is working memory. Initially, using the embedding distance between short-term memory and working memory, a few of the closest engrams are selected. Subsequently, in long-term memory, it is necessary to find associated memories. However, as LTM has an unspecified capacity and is much larger than short-term memory, exploring based on the embedding distance with all engrams is inefficient and biologically implausible. Since the engrams discovered in short-term memory are connected to the engrams in LTM, the goal is to consecutively activate the most strongly connected engrams based on this connection. The engrams thus explored are utilized to aid in the model's inference. Additionally, as mentioned earlier, you can refer to Appendix E for algorithmic operations, and in the source code we provided, you can see all the logic we implemented for Memoria.

Reviewer comment (Presentation):
Q3. The only reference to Appendix D seems to be in section 3: “We provided the visualizations of changes of connection in Appendix D to help understand these processes.” It remains unclear what these visualizations are, given that there are no explanations in Appendix D.

Author response:
This figure visualizes the evolving connections of long-term memory of Memoria while continuously inferring data. It was conducted to structurally depict the long-term memory formed in Memoria. Figure 9 provides visual information at timesteps 20, 60, and 130. At $t=20$, overall, there are still relatively few engrams (nodes), and the connections are dense. As time passes, the connections gradually spread throughout the entire network. At $t=130$, it can be observed that old memories have mostly faded, but some memories remain valid even after a long time, indicating Memoria's ability to preserve crucial information over extended periods.

Figure 10 demonstrates that Memoria autonomously establishes better connections for similar information. In the experiment, the same data was inferred twice. The lower half of the dots represent memories from the first inference, while the upper half represents memories from the second inference. An important point here is that dots on the same vertical line encode the same information. The figure shows strong connections along vertical lines, indicating that Memoria can form connections between similar information even if they are temporally distant.

In the revised version of our paper, we have supplemented explanations in the captions of each figure.

Reviewer comment (Experiments):
Q4. It seems that the proposed memory module is very complicated with many operations and sophisticated plasticity rules. Is this complexity all necessary for the performance? Ablation studies, or at least qualitative discussion, would be appreciated to clarify which aspects of the proposed architecture are important and which are not.

Author response:
We conducted an ablation study, and the revised version of the paper includes this in Appendix C. We compared the effectiveness of each memory module in a sorting task across various sequence lengths. In summary, when the sequence length is short, the contribution of working memory to performance is significant. However, as the sequence length increases, the function and role of working memory gradually diminish. The deficient performance is compensated as the roles of short-term memory and long-term memory become more pronounced. This observation highlights the complementary functionality of the three modules of Memoria based on the characteristics of the dataset.

Reviewer comment (Reproducibility):
Q5. Given the high complexity of the proposed architecture, could the authors comment on its computational complexity (both theoretical and empirical)? It seems it would be quite resource-demanding to run.

Author response:
We have revised the paper by incorporating discussions on the computational complexity of Memoria. Theoretically, we have computed and presented the computational complexity of Memoria operations, while empirically measuring and presenting the corresponding execution times and memory usage. In terms of execution time, Memoria demonstrated a slightly faster performance compared to the other techniques, while utilizing a larger GPU memory footprint. However, a significant portion of this memory usage is largely independent of the model's embedding or gradient update processes. Therefore, there seems to be considerable room for improvement when enhancing the algorithmic or programming efficiency. For detailed information, please refer to Appendix E, Algorithm & Computational Complexity section of the revised paper.

Reviewer comment (Evaluation):
The method is over-complicated with many memory retrieval steps. The authors need to consider the running times and computing resource requirements when comparing their methods and other baselines.

Author response:
We have revised the paper by incorporating discussions on the computational complexity of Memoria. Theoretically, we have computed and presented the computational complexity of Memoria operations, while empirically measuring and presenting the corresponding execution times and memory usage. For detailed information, please refer to Appendix E, Algorithm & Computational Complexity section of the revised paper.

Reviewer comment (Experiments):
The baseline set should include more recent memory-augmented Transformers such as Recurrent Memory Transformer (Bulatov et al., 2022) and Memorizing Transformers (Wu et al., 2022) or long-range modelling techniques (Mehta et al., 2022).

Author response:
It is hard to compare Memorizing Transformers with our research or the baseline we dealt with in our study. The reason is that the techniques covered in this study retrieve and utilize memory for each segment of inference, whereas Memorizing Transformers retrieve memory at the token level. In Memorizing Transformers, 32 pieces of memory information were fetched for every token inference. When applied to a single segment, such as in the case of a sequence length of 512, it results in utilizing $512 \times 32 = 16384$ pieces of memory information. This substantial difference sets it apart from other baselines. Since other techniques operate at the block level, it is not possible to increase their memory size as big as in their method. If we reduce the memory of Memorizing Transformer significantly, it would not be a fair comparison and there is a risk of denigrating the study. Thus, it was not included in the comparison.  

In the case of RMT, the highest performance achieved in that paper is the result of applying the RMT technique simultaneously with Transformer-XL. However, when RMT is applied alone, its highest performance is lower than Transformer-XL. Therefore, we decided to include Transformer-XL in our baseline. Additionally, the authors of RMT reported training instability when the memory size exceeds 10, making it difficult to make a fair comparison as in Memorizing Transformers.

As for Gated State Spaces, it is a research aiming to optimize the expensive attention operations of Transformer to enhance performance. This is in the same context as attention operation optimization studies. This focuses on a slightly different problem compared to Memoria or other baselines, aiming to maintain a long context while inferring with a short segment length using external memory.

Reviewer comment (Novelty):
Reference on Hebbian learning for deep learning is a bit out of date. Please consider more recent works on attention/Transformer (Le et al., 2020, Ramsauer et al., 2020, Limbacher et al., 2020).

Author response:
Limbacher's research proposes a new type of memory neural network, applying Hebbian synaptic plasticity. Ramsauer's research proposes a differentiable structure for the Hopfield network, suggesting it as a replacement for traditional neural network layers. This research, despite differences in the goals and functions of Memoria, partially shares theoretical background. In the revised version of our paper, we have incorporated the research mentioned, along with additional research on Hebbian learning and associative learning, into the related work section.

Reviewer comment (Presentation):
Fig. 1: Are the connections arbitrarily drawn? Should there be connections between any pair of engrams?

Author response:
Figure 1 is drawn arbitrarily. Not all connections exist for every pair. During the memorize and forget stage, connections are formed among all reminded engrams, representing their interrelations.

Reviewer comment (Presentation):
Fig. 2 caption provides little information. It is very hard to understand the figure.

Author response:
In the revised version of the paper, we have added a caption for Figure 2.

Reviewer comment (Reproducibility):
Please include an algorithm summarizing all steps presented in the method.

Author response:
In the revised version of the paper, we have presented the operations of each memory operation stage in algorithmic form in Appendix E, Algorithm & Computational Complexity section.

Reviewer comment (Experiments):
Experiments: can you compare the size of baselines? Do you ensure that the sizes of these models are similar and everything is fair?

Author response:
Yes, in each task, experiments were conducted by applying the same model hyperparameters to all models compared together. The relevant information is specified in Appendix B, Training Details And Additional Results, and in Section 4, Experiments, at the beginning of each subsection of the task.

Reviewer comment (Theory):
Several claims this paper makes regarding Hebbian theory and biological plausibility are inaccurate, or not sufficiently qualified. Hebbian theory is not a theory of human memorization, as the authors state, but rather a theory for explaining synaptic plasticity, i.e. changes in synaptic strength between neurons. While we can use Hebbian learning as a learning algorithm to train computation models for associative memory (such as Hopfield Nets), there is insufficient evidence to suggest that such models are an accurate reflection of how the brain store and retrieve memories (in fact, they typically require symmetric weights, which make them biologically-implausible).

Author response:
> Several claims this paper makes regarding Hebbian theory and biological plausibility are inaccurate, or not sufficiently qualified.

In the revised paper, especially in the introduction and related works sections, we have deliberately chosen expressions related to Hebbian learning or aspects associated with human function more attentively, aiming to facilitate the interpretation that some processes of human memory are abstractly reflected in Memoria.

Reviewer comment (Evaluation):
The authors compare the performance of Memoria against several previous methods. However, there appears to be a large discrepancy between the performance of those previous methods in the authors’ experiments, versus the experiments of their original papers. This makes it difficult to judge the performance of Memoria in comparison to current methods.

Author response:
> The authors compare the performance of Memoria against several previous methods. However, there appears to be a large discrepancy between the performance of those previous methods in the authors’ experiments, versus the experiments of their original papers. This makes it difficult to judge the performance of Memoria in comparison to current methods. I’ll elaborate in the questions section.

> The perplexity and BPC achieved by Transformer-XL, Compressive Transformer, and ∞-former in this paper all differ quite significantly from the results their respective papers claim. For example, in their original papers, on WikiText-103, Transformer-XL claims to achieve 18.3 perplexity, Compressive Transformer claims 17.1, and ∞-former claims 16.61. But in the authors experiments, their perplexities are 24.543, 24.794 and 24.685 respectively, which are much higher. Is this discrepancy mostly attributed to differences in network hyperparameters? If so, are the hyperparameters you select in your experiments biased to favor Memoria over these other methods in any way? How would Memoria compare with these methods under the conditions they achieve their published results— would its lead still hold?

The mentioned perplexity scores of 18.3, 17.1, and 16.61 may appear quite similar, but the backgrounds from which they were obtained are actually quite different. Referring to the original papers, we presented the key variables in a table for the settings where the scores were obtained. For the Compressive Transformer, the sequence length and memory length are significantly different from Transformer-XL, while for the $\infty$-former, the model is smaller, with a different memory length, but the most crucial difference is the use of a pre-trained model.

The conditions of all previous studies cannot be tested, but we conducted additional experiments using the settings of the $\infty$-former (last two rows of the table), which showed the best performance as mentioned. In the case of the $\infty$-former, we achieved performance similar to that reported in the original paper, and Memoria exhibited even better performance. Thus, Memoria maintained its lead even in the setting of the previous paper. The hyperparameters for Memoria were set as follows: $N_{wm}$ = 256, $N_{ltm}^{rem}$ and $N_{stm}^{rem}$ = 128, and a short-term memory capacity of 1024.

|  | PPL | layers | dim | Sequence Length | Memory Length | Pretraining |
| --- | --- | --- | --- | --- | --- | --- |
| Transformer XL | 18.3 | 18 | 1024 | 384 | 384 | X |
| Compressive Transformer | 17.1 | 18 | 1024 | 512 | 512+1536 | X |
| $\infty$-former | 16.61 | 12 | 768 | 512 | 512 | O |
| $\infty$-former (Ours) | 16.411 | 12 | 768 | 512 | 512 | O |
| Memoria Transformer | 15.834 | 12 | 768 | 512 | 512 | O |

Reviewer comment (Experiments):
An important novel contribution of Memoria to me is the introduction of a 3-level memory model inspired by the Multi-Store Memory Model; this is significantly different from previous approaches. But I’m not sure how each of these stages each contribute to the model’s performance (in other words, is the complexity justified? Can the same results be achieved with a simpler design?). Have you conducted any ablation studies? I think demonstrating how each part of your memory model improves performance would better convince the audience of the necessity of your design, and the importance of your contributions.

Author response:
> An important novel contribution of Memoria to me is the introduction of a 3-level memory model inspired by the Multi-Store Memory Model; this is significantly different from previous approaches. But I’m not sure how each of these stages each contribute to the model’s performance (in other words, is the complexity justified? Can the same results be achieved with a simpler design?). Have you conducted any ablation studies? I think demonstrating how each part of your memory model improves performance would better convince the audience of the necessity of your design, and the importance of your contributions.

We conducted an ablation study, and the revised version of the paper includes this in Appendix C. We compared the effectiveness of each memory module in a sorting task across various sequence lengths. In summary, when the sequence length is short, the contribution of working memory to performance is significant. However, as the sequence length increases, the function and role of working memory gradually diminish. The deficient performance is compensated as the roles of short-term memory and long-term memory become more pronounced. This observation highlights the complementary functionality of the three modules of Memoria based on the characteristics of the dataset.

Reviewer comment (Novelty):
There’s been a lot of research on auto and heteroassociative memory models that employ Hebbian learning, dating back to Hopfield Networks and their heteroassociative counterpart Sparse Distributed Memories. Some more recent papers include Krotov et al. (2016), Demircigil et al. (2017), Rae et al. (2018), and Ramsauer et al. (2020), to name a few. This line of work seem intimately relevant to Memoria, as both aim to use Hebbian learning to train memory models for deep learning. How does Memoria compare to these approaches, in terms of design, performance, and potential applications?

Author response:
> There’s been a lot of research on auto and heteroassociative memory models that employ Hebbian learning, dating back to Hopfield Networks and their heteroassociative counterpart Sparse Distributed Memories. Some more recent papers include Krotov et al. (2016), Demircigil et al. (2017), Rae et al. (2018), and Ramsauer et al. (2020), to name a few. This line of work seem intimately relevant to Memoria, as both aim to use Hebbian learning to train memory models for deep learning. How does Memoria compare to these approaches, in terms of design, performance, and potential applications?

The studies you mentioned appear to utilize Hebbian learning in the context of deep learning or investigate the modeling of associative memory. However, these studies differ in their objectives from Memoria, despite sharing some theoretical background. The research by Krotov and Demircigil explores the functionality of association memory for pattern recognition. Both studies focus on the Hopfield Network, where the extraction of associated information is based on pre-compressed weights that are based on the content. In contrast, Memoria aims to utilize external memory for maintaining and utilizing memories to process long sequence segments. Hence, direct comparisons are not easy due to the disparity in goals. Rae's study also applies Hebbian learning but for the purpose of aiding learning for low-frequency classes in classification problems. Ramsauer's research proposes a differentiable structure for the Hopfield network, suggesting it as a replacement for traditional neural network layers. Thus, it appears difficult to directly compare these studies with Memoria but we included some of the works in related works.