Paper: A Framework for Inference Inspired by Human Memory Mechanisms (ICLR-2024)

Reviewer comment (Theory):
The primary limitation concerns the framing of the architecture as instantiating both working memory and longterm memory. It is not clear to me that the architecture actually involves longterm memory in any meaningful sense. I think the approach would be better described as a form of relational working memory (utilizing a tensor product to capture relational information). This of course doesn't concern the method itself, which seems to perform well across multiple tasks. But I think the contribution would be much more clearly framed as a kind of working memory that exploits *relational* information. The role of relations in working memory is very well-studied in cognitive science (see references below), and I think this would make an interesting topic for discussion.

Author response:
We are enthused that the reviewers gained a profound understanding of this paper. We would like to provide the following clarifications to address the concerns of the reviewer. In cognitive neuroscience, working memory and long-term memory differ in the following aspects: (i) purpose: working memory is primarily employed for processing and manipulating current information, while long-term memory is more geared towards storing and retrieving extensive information, including acquired knowledge and experiences. (ii) capacity: working memory is constrained by a limited capacity, whereas long-term memory boasts a significantly larger capacity. In our research, we utilize a content-based sparse cross-attention mechanism for updating working memory, where the current perception serves as both the key and value, while working memory acts as the query. In contrast, long-term memory is primarily updated through tensor outer product operations, which are not directly tied to the current perception but rather involve extracting higher-order relationships from the cumulative content of working memory. With a larger three-dimensional tensor, long-term memory can accommodate a greater volume of information. This can be interpreted as working memory paying more attention to current perceptual information, while long-term memory involves extracting higher-order relationships from the accumulated content of all previous working memories. During our research, we have indeed reviewed numerous relevant papers the role of relations in working memory. We would like to make appropriate modifications in the next version. Thank you once more for your insightful suggestions.

Reviewer comment (Experiments):
Is it possible to study an ablation model that includes the 'longterm' memory component but not the 'working' memory? It seems likely that the tensor product in the longterm memory component is primarily driving the gain in performance, and it would be nice if this could be isolated.

Author response:
We prefer to believe that the presence of the working memory component is justified, both from a cognitive science perspective and in the context of Transformers' attention mechanism. This is because it utilizes an attention mechanism to encode crucial information from the current perception into the working memory, which is consistent with the human cognitive process. The reviewers also acknowledge that long-term memory involves higher-order extraction of relations in working memory. If the structure of working memory were absent, it could require substantial changes to the research.
However, to investigate whether the tensor product in long-term memory is the primary factor driving the gain in performance, we conducted an ablation experiment during the inference phase by removing the working memory (i.e., not utilizing the understanding of the current perception from the working memory, omitting formulas 8 and 11 in the paper, and directly extracting the understanding from long-term memory). The results are reported in the table below. Across various tasks, the absence of working memory led to a slight performance decline compared to the original best cases. However, this performance decline is far less than the impact of removing long-term memory, which indicates the importance of tensor outer product in long-term memory.

| Tasks                              |Layers |MITR                  | MITR without WM      |
| --------------------------------- | ----------------------- | ---------------------- |---------------------- |
| SORT-OF-CLEVR                     |8| 99.34, 87.61, 62.45     | 99.26 , 84.72 , 61.86   |
| bAbI                              |8| 2.55                    | 2.64                   |
| DETECTING EQUILATERAL TRIANGLES   |8| 97.9                    | 97.2                   |
| Enwik8  [1]                         |12| 0.96 bpc                | 1.01 bpc               |
| WikiText-103 [2]                  |16   | 16.5 ppl                | 17.2 ppl               |

Here, 99.34, 87.61 and 62.45 represent the accuracy of the unary, binary and ternary problems, respectively.

Reviewer comment (Novelty):
It would be good to cite work from cognitive science on the role of tensor product representations in working memory [1,2] as this is highly related to the outer product mechanism in the 'longterm' memory module.

Author response:
We appreciate your valuable insights, and we plan to include these references in the next revision.

Reviewer comment (Theory):
In what sense does the 'longterm' memory module involve long term memory more than the 'working' memory module? They both seem to operate over the same timescale, the only difference being the presence of the tensor product to capture relational interactions (which is not related to longterm vs. working memory).

Author response:
In this study, the difference between working memory and long-term memory is not solely based on distinct capacities but also on their distinct processing capabilities for information. Working memory is primarily involved in storing and processing temporary information, while long-term memory denotes a more advanced stage involving the extraction and processing of information, which, as we previously hypothesized, may include the extraction of relations within working memory.
The relationship between working memory and the GWT theory is closely intertwined, supported by numerous studies [1,2]. Consequently, we believe that the depiction of the term "working memory" might be relatively clear and reasonable. Regarding the term "long-term memory", after carefully considering your suggestions and reviewing relevant literature, we propose describing it as "relation memory" for enhanced clarity. This term captures and stores certain outer product relations within working memory, a concept extensively explored in various studies [3,4,5]. 

However, due to time constraints and the substantial textual adjustments required, we find it challenging to achieve the desired version within the limited timeframe. Nevertheless, we commit to revise the description of the term "long-term memory" in the next version, attributing the main contributions of this work to the extraction of high-order relations within working memory.

Reviewer comment (Experiments):
Is it possible to ablate the 'working' memory module?

Author response:
We prefer to believe that the presence of the working memory component is justified, both from a cognitive science perspective and in the context of Transformers' attention mechanism. This is because it utilizes an attention mechanism to encode crucial information from the current perception into the working memory, which is consistent with the human cognitive process. The reviewers also acknowledge that long-term memory involves higher-order extraction of relations in working memory. If the structure of working memory were absent, it could require substantial changes to the research.
However, to investigate whether the tensor product in long-term memory is the primary factor driving the gain in performance, we conducted an ablation experiment during the inference phase by removing the working memory (i.e., not utilizing the understanding of the current perception from the working memory, omitting formulas 8 and 11 in the paper, and directly extracting the understanding from long-term memory). The results are reported in the table below. Across various tasks, the absence of working memory led to a slight performance decline compared to the original best cases. However, this performance decline is far less than the impact of removing long-term memory, which indicates the importance of tensor outer product in long-term memory.

| Tasks                              |Layers |MITR                  | MITR without WM      |
| --------------------------------- | ----------------------- | ---------------------- |---------------------- |
| SORT-OF-CLEVR                     |8| 99.34, 87.61, 62.45     | 99.26 , 84.72 , 61.86   |
| bAbI                              |8| 2.55                    | 2.64                   |
| DETECTING EQUILATERAL TRIANGLES   |8| 97.9                    | 97.2                   |
| Enwik8  [1]                         |12| 0.96 bpc                | 1.01 bpc               |
| WikiText-103 [2]                  |16   | 16.5 ppl                | 17.2 ppl               |

Here, 99.34, 87.61 and 62.45 represent the accuracy of the unary, binary and ternary problems, respectively.

Reviewer comment (Writing):
The text appears to be excessively embellished. I would like to encourage the author to employ conventional terminology, as exemplified by the authors referencing "relation calculation" in the abstract.

Author response:
We sincerely appreciate the reviewer for giving us thorough feedback regarding the writing style of this paper, we will incorporate the feedback in the next revision of the paper to avoid being artificial. We will also make appropriate adjustments to the presentation structure, particularly in the visualization section, to enhance clarity.

Reviewer comment (Evaluation):
The paper includes visualizations of attention patterns between perceptions and memories, but it could benefit from providing more detailed explanations and interpretations of these visualizations.

Author response:
We would like to emphasize that we do provide a more detailed explanation in Appendix B.2 (given the limited length in the main text), illustrating the meaning of visualizations with specific examples. Detailed explanations for each case can be found in the legends of Figures 6a and 6b. Taking your valuable suggestions into account, we will diligently revise this presentation style in the next revision. We plan to appropriately integrate textual explanations into both the main text and the body of the appendix, rather than presenting them solely in the form of legends. If you find the explanations insufficient, we would like to provide additional explanations and interpretations in the next revision.

Reviewer comment (Evaluation):
Examining the qualitative impact of your modules on various types of tasks would provide valuable insights, rather than solely relying on quantitative results. This approach would enhance the paper's overall credibility. You can achieve this by employing various visualization techniques and similar methods.

Author response:
For this suggestion, we conduct a qualitative analysis of the proposed modules from a priori perspective on various types of tasks. We posit that as the computational steps increase, the content in working memory and long-term memory can serve as a prior for reasoning. Moreover, we use the bertviz library [14] to add visualizations of attention distributions for perception and memory on the bAbI dataset in the appendix. In addition, we add attention histograms for working memory and long-term memory on the WikiText-103 [2] dataset (new experiment) in the appendix. When predicting the next token, the attention histograms of these two different types are notably distinct. The attention histogram for long-term memory indicates its ability to focus on more long-range and relevant information, which is significantly longer than working memory.

Reviewer comment (Writing):
Enhance the clarity of the text to facilitate a deeper comprehension of the paper. Despite grammatical accuracy, the writing occasionally comes across as artificial.

Author response:
We sincerely appreciate the reviewer for giving us thorough feedback regarding the writing style of this paper, we will incorporate the feedback in the next revision of the paper to avoid being artificial. We will also make appropriate adjustments to the presentation structure, particularly in the visualization section, to enhance clarity.

Reviewer comment (Experiments):
Consider expanding the scope of this research to include more memory-based and cognition-inspired tasks. You may find the paper titled "Decoding the Enigma: Benchmarking Humans and AIs on the Many Facets of Working Memory" by Sikarwar et al. to be a relevant reference in this context.

Author response:
We conduct additional experiments by adding the PMI framework into the decoder of Transformers for language modeling, covering enwik8 [1], WikiText-103 [2] and PG-19 [3], to broaden the scope of this research. The results are as follows. 

In addition, we thoroughly reviewed the paper "Decoding the Enigma: Benchmarking Humans and AIs on the Many Facets of Working Memory", which you kindly shared. We consider it an outstanding benchmark. In the next revision, we plan to conduct experiments on this dataset to further evaluate the performance of our PMI architecture.

**Table 1**: Test set bits-per-character on enwik8.

| Models| Layers|Params |BPC|
|--|--|-|--|
| Transformer-XL  [8]| 12 | 41M |1.06|
| Transformer-XL  [8] | 24 | 277M | 0.99|
| Compressive Transformer  [3]  | 24| - | 0.97|
| Sparse Transformer [12]  | 30| 95M | 0.99|
| Adaptive Transformer [13] |12| 39M| 1.02 |
| RMT [9] | 12     | - | 1.222|
| TIMS+HSW [5]*  | 12| 43M| 1.36|
| MITR (ours) | 12| 45M| **0.96**|

**Table 2**: Comparison with other models on WikiText-103.

| Models | Layers | Params |Valid PPL |Test PPL |
|--|-|-|--|--|
| LSTM [6]|  - | - | -  | 48.7   |
| RMC [7]|-| 30.8   | - | 31.6 |
| Standard Transformer-XL [8] | - | 151M   | - | 24 |
| RMT [9]  | 16|- |-  | 24.85|
| Compressive Transformer [3] | 18| - | 16 | 17.1|
| Transformer-XL Large [8]  | 18| 257M   | - | 18.3|
| TIMS+HSW [5]* | 8| 112M   | 35.9 | 36.7|
| MITR (ours)   | 8 | 116M | 24.9 | 23.8 |
| MITR (ours)   | 16| 233M| 15.3| **16.5** |

**Table 3**: Results on language modeling on PG-19 dataset.

| Models | Layers|Valid PPL|Test PPL| 
|--|--|-|--|
| Transformer-XL [8] | 36| 45.5| 36.3|
| Compressive Transformer [3]   | 36 | 43.4| 33.6|
| ∞-former [10]  |12| - | 32.48|
| Routing Transformer [11] | 12| - | 33.2|
| TR+HSW [5]* | 12 | 39.46 |32.46|
| MITR (ours) | 12 |37.12| **31.04**|
Here, * signifies the results obtained from our experiments.

Reviewer comment (Experiments):
The experiments is hard to understand, and i think it is not necessary to conduct experiments on image classification. And there is little work on the memory augment image model due to the image's too long context.

Author response:
**3.Image Classification**

For image classification, our aim is to demonstrate that the architecture embedded with our PMI, whether it belongs to the transformers or convolutional series, could outperform the original models. This is to establish the generality of the PMI, showcasing its applicability beyond Transformers. Additionally, the original intent of the paper is to prove that the introduction of PMI could enhance the model’s inference ability, which we think maybe include the ability for image classification, while not approaching it from the perspective of long sequences. However, considering your constructive suggestion, we will place this experiment in the appendix, while adding language modeling experiments into the main text. What are your thoughts on this arrangement?

Reviewer comment (Presentation):
I don't see any connection between your work and the title, the author maybe need change a title due to this model hard to help us underanding AI.

Author response:
**4.Title**

After thoughtful consideration of your valuable feedback, we are planning to revise the title to "A PMI Framework for Inference Inspired by Human Memory Mechanisms", considering that the focal point of this paper lies in leveraging theoretical insights from cognitive science to construct a cognitive architecture named PMI, and embed it into mainstream architecture such as Transformers and convolutional networks to help us build AI systems for inference that align more closely with human cognition.

Reviewer comment (Experiments):
I hope the author takes more experiments on the language model to solve the longer context challenge.

Author response:
**1.More Experiments on Language Modeling**

To further validate the effectiveness of the proposed approach in long-sequence language modeling, we applied MITR—a decoder-only transformer embedded with our PMI module, to a variety of datasets on both character-level and word-level language modeling, including enwik8 [1], WikiText-103 [2] and PG-19 [3]. Enwiki8 is utilized for character-level language modeling and comprises 100M bytes of unprocessed Wikipedia text. Both WikiText-103 and PG-19 are benchmarks for word-level language modeling with long-term dependency, with the former containing 103M tokens from 28K English Wikipedia articles and the latter from English books published before 1919.
Part of the implementation is based on Transformer-XL repository [4]. The experiment on Enwik8 employs a 12-layer MITR (8 heads, 512 hidden size, 2048 intermediate FF), WikiText-103 uses a 16-layer MITR (10 heads, 410 hidden size, 2100 intermediate FF), while PG19 uses a 12-layer MITR (8 heads, 1024 embedding size, 4096 intermediate FF). In addition, we use Adam optimizer with linear schedule learning rate starting from 0.00025 for 500,000 steps and set the working memory size M=8, the long-term memory size N=5 and top-k=5 on all the datasets. The complete hyperparameter settings are available in our repository as well as in the Appendix. 

The test bits per character (BPC) on the Enwiki8 dataset and the perplexity (PPL) on WikiText-103 and PG-19 are reported in the table below. Notably, we improve the results of bpc/perplexity to 0.96 on enwiki8, 16.5 on WikiText-103 and 31.04 on PG19, which demonstrates the superiority of the PMI architecture. 

**Table 1**: Test set bits-per-character on enwik8.

| Models| Layers|Params |BPC|
|-|-|-|-|
| Transformer-XL  [8]|12| 41M |1.06|
| Transformer-XL  [8] |24| 277M|0.99|
| Compressive Transformer  [3]|24|-|0.97|
| Sparse Transformer [12]| 30| 95M |0.99|
| Adaptive Transformer [13] |12| 39M|1.02 |
| RMT [9] |12| - |1.222|
| TIMS+HSW [5]*|12| 43M|1.36|
| MITR (ours) |12|45M| **0.96**|

**Table 2**: Comparison with other models on WikiText-103.

|Models|Layers|Params|Valid PPL|Test PPL|
|-|-|-|-|-|
| LSTM [6]|-|-|-|48.7|
| RMC [7]|-| 30.8| -|31.6 |
| Standard Transformer-XL [8] |-|151M|-| 24 |
| RMT [9]|16|-|-| 24.85|
| Compressive Transformer [3] |18|-|16|17.1|
| Transformer-XL Large [8]| 18| 257M| - |18.3|
| TIMS+HSW [5]* | 8| 112M| 35.9 |36.7|
| MITR (ours) |8|116M |24.9|23.8|
| MITR (ours) |16| 233M|15.3|**16.5**|

**Table 3**: Results on language modeling on PG-19 dataset.

| Models | Layers|Valid PPL|Test PPL| 
|-|-|-|-|
| Transformer-XL [8]|36|45.5| 36.3|
| Compressive Transformer [3]| 36 |43.4| 33.6|
| ∞-former [10]|12|-|32.48|
| Routing Transformer [11] |12|-| 33.2|
| TR+HSW [5]*|12 |39.46 |32.46|
| MITR (ours) |12|37.12|**31.04**|
Here, * signifies the results obtained from our experiments.