Online Adaptation of Language Models with a Memory of Amortized Contexts

Summary

This paper focuses on how to adapt static language models (LMs) with streaming documents during inference time. There are two high-level challenges here: 1) how to store new domain/task relevant information, 2) how to utilize the stored information for downstream task-solving, i.e., doing question answering (QA). In this paper, the authors propose, MAC, a parameter-efficient adaptor approach to 1) encode new documents into a sizable vector memory bank, and 2) utilize those encoded knowledge via extra attention mechanism. The proposed method is compared with other fine-tuning baselines on three QA tasks tailored for evaluating online adaption scenarios, showing that the method is more effective. Additional analyses and ablative studies are also provided to drive more insights into the model designs and behaviors.

Strengths

1) The paper is well-written with enough background and details for readers to follow. 2) The authors apply their proposed method to different LMs with various architectures and training protocols, which provide support for the generalization ability of the approach. 3) The design of experiments are mostly reasonable (baselines are OK, multiple datasets are good), and the results suggest the proposed method is effective. There are additional analyses and ablative studies, e.g., i) the proposed method can do better knowledge retention and help RAG, ii) most design choices are validated.

Weaknesses

1) The proposed method is a natural extension of memory-augmented LLM (as cited in paper [76,79]) with a token compression module. Rather than comparing with only closed-book QA models, it is good to compare with context compression methods, e..g, text token compression w/ RAG or long-context models, e.g., [1]J.-H. Kim, J. Yeom, S. Yun, and H. O. Song. Compressed context memory for online language [2]H. Jiang, Q. Wu, X. Luo, D. Li, C.-Y. Lin, Y. Yang, and L. Qiu. Longllmlingua: Accelerating and enhancing llms in long context scenarios via prompt compression [3]T. Ge, J. Hu, L. Wang, X. Wang, S.-Q. Chen, and F. Wei. In-context autoencoder for context compression in a large language model model interaction, 2024 2) The experiment settings can be problematic. As the goal is to adapt the model to new knowledge, it is not clear whether those evaluated datasets manifest that. Both news articles and wikipedia pages are highly used in pretraining LMs, e.g., LLaMA-2. It is necessary to report zero-shot and few-shot results with the base model. Without that, it is hard to judge the benefit of online adaption. As all reported models have very low EM or F1 scores, it is good to report the base model for sanity check. It is noticeable that the proposed method has less improvement on more capable models (larger sizes, e.g., LLaMA-2). It is good to dig a little bit into this, e.g., experiment with similar sized models ***with or without instruction tuning*** (LLAMA and Vicuna/Alpaca). Specifically, the instruction tuning might be relevant for the model to memorize certain information.

Questions

Are those methods reported in Table 1 sensitive to the input order? For example, SQuAD questions do not have any temporal dependency on documents, it is good to see the performance on different streams, e.g., recency bias (putting the relevant documents in the beginning and irrelevant in the end). What is the setup for decoding? Beam search or greedy? Using sampling techniques? If so, what is the temperature? What is the training cost? Vs baselines? What are the inference prompts? How data efficient is the proposed method? Is it possible to achieve similar performance with less training?

Rating

6

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

NA

Summary

This paper proposes Memory of Amortized Contexts (MAC) which can encode the documents into compact modulations stored in a memory bank, which can later be retrieved to answer questions.

Strengths

1. The proposed method is efficient compared to the baselines. 2. The paper is well-written and easy to follow.

Weaknesses

1. **Unfair Comparison**: This paper only compares with [1], which may be somewhat unfair. In the paper [1], they propose to distill the information into a parameter vector $\phi$. However, in the current paper, the proposed method MAC stores all the modulations constructed from each context, which naturally contains all the knowledge from the contexts, leading to better knowledge retention. The setting might be unfair. 2. **Missing Baselines**: What I think could be the fair comparisons might be comparing MAC with most recent retrieval methods, such as DPR[2], BM25, RAPTOR[3], IRCoT[4] etc. Although the paper proposes to combine MAC and BM25, there is no direct comparison between them as MAC is also essentially doing retrieval. Having two retrievers is better than only having one retriever may not show the effectiveness of MAC. 3. **Doubt on the Performances**: Even MAC shows improvements over [1] (in the paper's setting), the best performance of MAC on Llama2-7B, ArchivalQa-Seq is 20.12 (EM), while BM25 can easily achieve 52.81 (EM) as shown in Table 2. It seems that MAC performs much worse than BM25. Can I expect the more advanced RAG methods can easily beat MAC even equipped with BM25? Can I really expect there is going to be improvements over other strong RAG methods when equipped with MAC? These are unanswered questions. 4. **Missing Related Works**: The paper talked about Retrieval Augmentation for LMs and Memory augmented LMs. However, for RAG, [2][3][4] are not mentioned in related work (there are more RAG method out there); for Memory augmented LMs, some methods such as MemoryLLM [5], Memoria [6], MemLLM [8], MemoryBank[7], Camelot[8] are not mentioned which can also perform online learning. Maybe the memory-based methods should also be used for comparison in the experiments. [1] Meta-Learning Online Adaptation of Language Models. [2] Dense Passage Retrieval for Open-Domain Question Answering. [3] RAPTOR: Recursive Abstractive Processing for Tree-Organized Retrieval. [4] Interleaving Retrieval with Chain-of-Thought Reasoning for Knowledge-Intensive Multi-Step Questions. [5] MemoryLLM: Towards Self-Updatable Large Language Models. [6] Memoria: Resolving Fateful Forgetting Problem through Human-Inspired Memory Architecture. [7] MemoryBank: Enhancing Large Language Models with Long-Term Memory. [8] MemLLM: Finetuning LLMs to Use An Explicit Read-Write Memory. [9] CAMELoT: Towards Large Language Models with Training-Free Consolidated Associative Memory.

Questions

See the weaknesses part

Rating

5

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

yes the authors have addressed the limitations

Summary

The paper proposes an online learning framework called MAC (Memory of Amortized Contexts) designed to efficiently adapt large language models (LLMs) online. By using feature extraction and memory augmentation methods, MAC compresses and stores new document information in a memory bank, retrieving relevant knowledge when answering questions. This method utilizes amortization-based meta-learning, achieving efficient modulation learning through a single forward pass, avoiding the need for gradient updates during testing. Experiments demonstrate that MAC outperforms existing methods in online adaptation performance, time, and memory efficiency, and can be combined with popular alternatives like Retrieval-Augmented Generation (RAG) to further enhance performance.

Strengths

1. The paper presents a novel memory-augmented online adaptation framework based on amortization, which can efficiently adapt to new information without requiring gradient updates. 2. Experimental results show that MAC performs exceptionally well across multiple datasets and architectures, significantly improving online adaptation performance, time, and memory efficiency. 3. MAC can be combined with existing methods like RAG, enhancing the quality of retrieved documents, demonstrating good scalability and compatibility. 4. The use of two memory-efficient techniques during training and inference stages reduces memory requirements, ensuring the method's scalability.

Weaknesses

1. The method involves several complex steps and model components, such as the amortization network and aggregation network, which may increase implementation difficulty. 2. Although the method has been experimentally evaluated on multiple datasets, the tasks are primarily focused on question-answering. Its adaptability to other task types remains to be verified. 3. The paper mainly demonstrates the method's effectiveness through experimental results but lacks in-depth theoretical analysis of certain key design choices, particularly in amortization and aggregation strategies.

Questions

1. How is the storage overhead of MAC controlled when facing large-scale data streams? Are there any further optimization measures? 2. How does the method perform on other task types (e.g., text classification, generation tasks)? Have any related experimental evaluations been conducted? 3. How were the architectures of the amortization and aggregation networks determined? Have other architectures been tried, and what were the outcomes? 4. Is there a concern of outdated or redundant modulation parameters in the memory bank? How are these issues addressed to maintain model efficiency?

Rating

5

Confidence

4

Soundness

2

Presentation

3

Contribution

3

Limitations

yes

Summary

This paper presents a novel online adaptation framework (Memory of Amortised Contexts, MAC), which effectively solves the problem of rapid updating of large language models (LLMs). MAC successfully preserves the knowledge learned by the model during the original training phase and the new data streams through memory augmentation and efficient fine-tuning of parameters. Experimental results show that MAC outperforms existing online fine-tuning methods regarding online adaptive performance (Table 1 and Figure 3), time, and memory efficiency (Figures 2, 4, and 5). MAC can be combined with popular methods such as Retrieval-Augmented Generation (RAG) to improve performance further (Table 2).

Strengths

1. In a time of rapidly updating information, MAC provides a practical approach to model updating that helps to keep language models current. With the proposed memory-efficient technique and forward propagation optimization, MAC significantly reduces the memory footprint and time consumption during online adaptation. 2. MAC effectively avoids the problem of catastrophic forgetting by constructing memory banks and ensures that the model integrates and utilizes old and new knowledge. 3. This paper validated the effectiveness of MAC through experiments with multiple datasets and different models, and the results are convincing.

Weaknesses

1. Memory bank growth issues. As online adaptation proceeds, the size of the memory bank is likely to grow, which may pose a challenge for memory management. It is recommended that the authors explore more effective memory bank reduction techniques in future work. 2. The ability to generalize adaptations from different domains remains unknown. Often, online learning may not be the same type of task (e.g., knowledge answering and coding), and it remains unclear whether the current method can cope with such scenarios.

Questions

1. The current experiments are mainly validation on small models, but according to the author's description llama was also trained for 50 epochs, which seems to be problematic? 2. The current fine-tuning is based on P-Tuning. Has there been any consideration of comparing more parameter-efficient fine-tuning methods, such as Lora fine-tuning? 3. Is it possible to add scenario experiments for online learning related to different tasks?

Rating

6

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

Yes

Please respond to the responses from the authors

Dear reviewer QtU4 Could you please take a look at the responses of the authors and let us know your thoughts on them? Are you satisfied with the responses and do you have some updates on your comments? AC

Thank you very much for the response

Dear Reviewer kPwd Thank you for letting us know! We are delighted to hear that our rebuttal addressed your questions well.\ If you have any further questions or suggestions, please do not hesitate to let us know. Thank you very much,\ Authors