Paper: Enhancing Length Extrapolation in Sequential Models with Pointer-Augmented Neural Memory (ICLR-2024)

Reviewer comment (Experiments):
The baselines used for comparison are outdated, and large language models (LLMs) with over 10B parameters could outperform the proposed method on reasoning tasks. The paper should either compare the method with these LLMs or show how PANM enhances LLMs' reasoning abilities.

Author response:
As explained in Appendix D's Baseline choice section, our baselines are competitive with SOTA in the considered tasks. Regarding your mention of LLM, could you clarify the setting when LLM with 10B could outperform our method?  If the LLM is not finetuned on the task, we have shown that Chat-GPT with 175B still fails in our task, and can be considered to underperform our method (Appendix D6). In case the setting is to finetune LLM to the task, we need to conduct experiments to confirm that the finetuned LLM could perform our PANM+LLM, and this task is not trivial because finetuning LLM of 10B requires huge computing resources.

Reviewer comment (Writing):
The paper downplays the similarities between PANM and the Transformer, particularly in terms of positional encodings and attention computations. The current notation makes these connections unclear, and the paper should work to make them more explicit. By presenting PANM as a modification of the Transformer, rather than something entirely different, the distinctions and innovations of PANM could be more easily understood. The choice of notation, such as using h and g instead of query and key, further obscures the connection, and should be addressed to enhance clarity.

Author response:
We spent a paragraph in Appendix A to compare positional encoding and our physical pointers. Similarly, we highlighted the difference between our attention and softmax attention used by Transformer in the introduction and Sec. 2.3.1. In this revision, we have added another table (Appendix A Table 3) to summarize the difference between Transformer and our method (see below).
| Difference | Transformer |	PANM (Our) |
| ----------- | ----------- |-----------|
| Key Generation  | Keys are computed based on input data. Hence, when meeting novel  data during testing, Transformer will observe novel keys, and cannot work properly. | The keys in our approach are  generated as fixed numbers,  specifically physical memory addresses. These keys are entirely separate from the data.|
|Extendable to Longer Sequences|The dimension of attention weights varies with input length, making arithmetic transformations on these attention weights infeasible as the sequence length increases.|The fixed nature of our physical addresses allows our pointers to be easily manipulated and extendable to longer sequences.|
|Symbol Processing Advantages| The use of attention weights as implicit pointers may lack the explicitness needed for effective symbol processing.| Using physical addresses as keys in our approach is crucial for symbol processing as it explicitly allows pointer assignment, dereference, and arithmetic operations.|
|Physical Address vs Positional Encoding| Positional encoding can be generated independently from data. However, they are not separated from the input data as our physical addresses. There is no explicit mechanism in Transformer to attend only to these | Our physical addresses are detached from the data, supporting the transformation of pointers through timesteps and isolating pointer manipulation from the input.|

We respectfully disagree with the reviewer that our paper should "presenting PANM as a modification of the Transformer". PANM is based on memory-augmented neural network framework (MANN) [Graves, Alex, Greg Wayne, and Ivo Danihelka. "Neural turing machines." arXiv preprint arXiv:1410.5401 (2014).], which was introduced even before Transformer. The notation we used follows the classical notation used in MANN, which was later rebranded into those in Transformer paper. In particular, h means the hidden state the GRU (kind of a controller in MANN), g is our function to map the memory content to the same space with h. They are not required to be written as K and V in Transformer.
That said, we agree that there are similarities between PANM and Transformer, and we have added detailed comparisons above to clarify the difference, showing the innovation of our approach in regards to not only Transformer but also other approaches relying on softmax attention. The choice of notations does not hinder the readability of our paper.

Reviewer comment (Experiments):
Why not conduct an experiment on LLM? I do think some billion-level network could be used in normal research lab, especially with advanced packages such as DeepSpeed, etc.

Author response:
It is not trivial for our resources to finetune LLM (billions of parameters) and we aim to improve fundamental and general deep neural networks, not to scale with LLMs, which can be pursued as a separate work. It is crucial to realize that our paper is not an NLP paper. Rather, it contributes a fundamental method for the general machine learning community. Please contribute to the discussion if you still insist that fine-tuning LLM is required to accept our paper.

Reviewer comment (Reproducibility):
How many parameters does PANM add? The authors state that DNC and NTM have ~1.2 million parameters, but I think I am missing something, I cannot see the parameter count for PANM.

Author response:
It depends on the task and choice of the Controller. Usually, PANM's module's size is similar to the size of an attention-layer in Transformer.  We report PANM model size in Appendix D1, the final sentence of the fifth paragraph. In the considering task, the whole architecture with PANM has ≈1.1 million parameters.

Reviewer comment (Novelty):
Say I am a practitioner and I want to decide if I want to choose PANM or U. TRM— looking at the performance of both the models, it is not obvious to me why I should prefer PANM over U. TRM since U. TRM performs very similarly to PANM on the tasks considered— could you highlight its advantages a bit more strongly over U. TRM, not necessarily in an experimental sense, but intuitively speaking what the benefits/non-benefits of an explicit pointer based memory module system are vis-a-vis a generalised transformer model design wise?

Author response:
As clarified earlier, PANM outperforms U. TRM (with and without RPE) in the toughest generalization testbed. We will clarify the benefit of our approach vs Transformer:
| Difference | Transformer |	PANM (Our) |
| ----------- | ----------- |-----------|
| Key Generation  | Keys are computed based on input data. Hence, when meeting novel  data during testing, Transformer will observe novel keys, and cannot work properly. | The keys in our approach are  generated as fixed numbers,  specifically physical memory addresses. These keys are entirely separate from the data.|
|Extendable to Longer Sequences|The dimension of attention weights varies with input length, making arithmetic transformations on these attention weights infeasible as the sequence length increases.|The fixed nature of our physical addresses allows our pointers to be easily manipulated and extendable to longer sequences.|
|Symbol Processing Advantages| The use of attention weights as implicit pointers may lack the explicitness needed for effective symbol processing.| Using physical addresses as keys in our approach is crucial for symbol processing as it explicitly allows pointer assignment, dereference, and arithmetic operations.|
|Physical Address vs Positional Encoding| Positional encoding can be generated independently from data. However, they are not separated from the input data as our physical addresses. There is no explicit mechanism in Transformer to attend only to these | Our physical addresses are detached from the data, supporting the transformation of pointers through timesteps and isolating pointer manipulation from the input.|

Reviewer comment (Experiments):
The comparison with Chat-GPT on algorithmic tasks is interesting— in Table 3, what would the failure examples of PANM look like, as are shown for Chat-GPT? Also, I maybe missing something— but are there any comparison tables on the other tasks apart from the couple of examples— mathematics, and SCAN? You can just point me to them, or maybe state why looking at them would be or not be relevant.

Author response:
PANM failure will focus on the end of the sequence. When sequence length increases, more characters at the end of the output sequence are likely wrong. For example, in the reverse task, input 123456789, true output 987654321, PANM output can be 987654322. The reason for this failure pattern is that PANM uses a Pointer Unit to recurrently transform the pointer across timestep. As the transformation is not always perfect, there could be error accumulation as the number of timesteps increases. Looking at Fig 2c as an example, attention errors (red) become clearer at the end of the decoding. Although in this example, the majority of attention is still correct, pointing to the right address. As the test length increases, these errors will amplify and possibly make the final prediction incorrect. We explained the reason for not carefully testing ChatGPT on other tasks in Appendix D6's first paragraph. There are two main reasons: (1) the task input is not suitable for text representation. (2) the test data is public and can be seen by ChatGPT. Another reason is the cost of systematically evaluating ChatGPT on these tasks: we have to manually extract the relevant answer from ChatGPT's output to have a reliable evaluation.

Reviewer comment (Presentation):
Just a small suggestion, but could be helpful to see what p_t^a, *p_t^a, and z_t correspond to in the input sequence in Figure 1 as an example.

Author response:
Thank you for your suggestion. We have revised the figure to connect the p_t^a, *p_t^a to the input sequence as you suggested. z_t is the decoding input, not part of the input sequence.

Reviewer comment (Writing):
Also, is their name for the combined module consisting of the address bank and the data memory?

Author response:
We have not named the module consisting of the address bank and the data memory. We can call {the address bank, data memory and pointer unit} PANM memory (Fig 1a,b,c and Appendix Fig. 3).