Augmenting Language Models with Long-Term Memory

Summary

This paper proposes an architectural extension to a frozen LLM that allows for incorporation of external memory caches (also constructed with the frozen LLM), allow the model to handle extremely long contexts.

Strengths

The approach presented is sensible (although I am not sufficiently familiar with the latest long-context modeling literature to determine if it is novel), and the experiments are convincing. In particular, I think the few-shot learning results are very strong.

Weaknesses

- The explanation of the model architecture is fairly unclear to me, and I found it difficult to match the description of the architecture and the presented diagram. I recommend the authors rewrite the explanation of the model and obtaining feedback from a fresh reader for suggestions on improving the clarity. Alternatively, presenting the architecture in terms of pseudo-code may also be clearer. - The evaluations presented are largely newly proposed (besides ChapterBreak). They appear to be sensibly constructed, but it does mean the results are slightly difficult to compare to or contextualize to other works. I would be interested in seeing some experiments on more established tasks (e.g. long document summarization).

Questions

- I believe GPT-2* refers to the reproduced GPT-2, but I don't think this is mentioned in the text. - Perplexity should be a positive number. Are the authors referring to negative log loss instead? - I would like to see experiments on with inference DEMA but with no memory cache (if I understand the model architecture correctly, this would still involve the SideNet parameters, which would differ from GPT-2*). This should apply to all current experiments.

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

2 fair

Contribution

3 good

Limitations

N/A

Summary

The paper proposes DeMA as a learning-based adapter for existing Language Models (LMs) to increase long-term memory and performance on long inputs.

Strengths

The method proposed is clear and sound and it could potentially be useful to the community since using transformers with a long context is an important open problem.

Weaknesses

Experiments that prove the advantage of DeMA are limited: - Most of the comparisons are limited to just a single memory-augmented model (MemTRM) and GPT-2 (that does not have any memory mechanism). Table 3 shows comparisons with several memory-augmented models but just on a single task and with old memory-augmented models. - The experiments do not show comparisons with several recent memory-augmented competitors: - “Memformer: A Memory-Augmented Transformer for Sequence Modeling” (Wu et al.). Here a comparison on WikiText-103 might have provided additional evidence of the superiority of DeMA. - “Unlimiformer: Long-Range Transformers with Unlimited Length Input” (Bertsch et al.) - “Training Language Models with Memory Augmentation” (Zhong et al.) - “RWKV: Reinventing RNNs for the Transformer Era” (Peng et al.) - I’m not convinced by the memory-augmented In-context learning experiments (Table 5). What is the purpose of this experiment? Is it really a fair “few-shot” comparison since DeMA is using additional 2000 demonstration examples loaded into cached memory? I think a comparison with other memory-augmented models will leave any doubt.

Questions

- L205-206 “Then, we concatenate shuffled documents within one group and truncate them into ordered segments.” Is it not clear to me the procedure to batch the input. I also read the appendix, but it is still not clear. Could you please elaborate a little bit on this aspect? - Why did you test your method just on the AO3 split of Chapterbreak and not on the PJ19 split?

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

2 fair

Contribution

3 good

Summary

This paper proposes a side network as a memory retriever and reader to help the LLMs model long context input. Specifically, the key-value pairs of the history inputs are saved in the cached memory bank. The side network takes in the current hidden states and the retrieved past key-value pairs to compute the memory-augmented representations. Their approach achieves better performance in perplexity and accuracy on multiple datasets (i.e., Project Gutenberg 2020-2022 Language Modeling Dataset, Arxiv Dataset, ChapterBreak Benchmark).

Strengths

1. The side network can retrieve and fuse memory-augmented long-context, which enables the LLM to model context with ultimate length. 2. Experiments demonstrate the effectiveness of the proposed approach on long context datasets like Arxiv. Besides, this paper shows that LLM’s ability to model long context better could help both in context learning, and zero-shot performance.

Weaknesses

1. The evaluation mainly focuses on perplexity score, and accuracy of suffix identification. To evaluate whether LLM models the long context well and thus benefit its’ generation ability, tasks like summarization/title generation should be included. 2. Missing comparison with Efficient Long-Text Understanding with Short-Text Models, which could take maximum 16K tokens as input.

Questions

na

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

na

Summary

To address the limitation of existing large language models (LLMs), the authors introduce a framework called Decoupled-Memory-Augmented LLMs (DEMA) that utilizes long-context information. DEMA utilizes a decoupled network architecture with a frozen LLM as a memory encoder and an adaptive residual side-network as a memory retriever and reader. Experimental results show that DEMA outperforms strong long-context models and achieves significant improvements in memory-augmented in-context learning over LLMs.

Strengths

1.The idea of this work is novel and the motivation is strong. It tries to store and utilize long-form previous context or knowledge to address the input length limit of LLMs. 2.The main experiments seem solid as this paper verifies their method on multiple datasets and the experimental results validate the effectiveness in helping LLMs to memorize and utilize long-form information. 3.This paper is well-written and easy to follow. Their main claim is clear and easy to understand.

Weaknesses

1.The paper introduces a novel lightweight residual SideNet, but there is a lack of discussion on the contribution of the SideNet module to the DEMA model in the ablation experiments. 2.In Section 3.4, the ablation experiments show that Figure 2(b) does not provide a clear visualization of the impact of different memory sizes on perplexity, especially when msz=32k (the orange line) and msz=16k (the green line) are very close to each other. This lack of distinction makes it less evident to draw a conclusive result that suggests better performance when msz=16k, as mentioned in the paper. 3.The baselines compared in the paper lack recent advancements, such as GPT-2 in 2019, GPT-3 and Bigbird in 2020, LocalTRM and RoutTRM in 2021, and MemTRM in March 2022. 4.The experiments include only main results (comparing with baselines) and the ablation studies. It would be better to conduct more detailed analyses or qualitative analysis to verify the proposed methods (e.g. whether the issues proposed in the motivation are fixed). 5.Related work section should discuss other memory augmented methods, especially the memory augment text generation methods. 6.The inputs section in Figure 1 is ambiguous and confusing. It is unclear whether the letters a, b, ..., h are meant to represent individual letters or tokens. If they represent individual letters, then combining them does not form a word. However, if they refer to tokens, it is advisable to use more precise symbols that represent tokens to avoid any ambiguity. 7.Figure2(a): sst-5 should be changed to SST-5, consistent with Table 5.

Questions

In line 111, you mention ‘total # layers’. Can you explain the meaning of the "#" symbol in this context?

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

None

Summary

This paper proposes a novel approach in handling long input context using memory framework. To be specific, any inputs that are longer than the maximum input size are cached in memory network (with the outcome of forward pass of LLM), and those cached memories are retrieved when "current" input is processed. The LLM remains fixed throughout training, hence being efficient; the trainable parameters are the fusion network and SideNet layers. The authors illustrate the effectiveness of the proposed approach on long-context benchmarks.

Strengths

- This paper proposes to a novel approach in handling long-context problem that transformer-based language models suffer from. - The paper evaluates the proposed idea on several benchmarks, and illustrate significant gain, especially on NLU tasks. - The idea is sounding, and the proposed approach conducts thorough ablation studies on modules introduced.

Weaknesses

- some parts of the paper need proper citations. i.e. Line 63-65: "directly adapting the entire ~~~suffers from catastrophic forgetting." needs a proper citation for the claim. - There are baselines that could be compared against in long-context language modeling tasks. For example, in Table 2, the authors compared against MemTRM and GPT2. LocalTRM is one of the baselines that the authors can compare to in language modeling task as well, as LocalTRM has similar number of parameters and has the context length of 8K. - It would be interested to see the inference speed of the proposed method, compared to the other baselines. It would be insightful to check how much of benefit is gained at how much expense of inference speed and computation cost. - Much performance gains are seen over the NLU tasks, yet the gains are somewhat trivial on language modeling benchmarks (PPL).

Questions

Please refer to the Weakness

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

Please refer to the Weakness