Online Adaptation of Language Models with a Memory of Amortized Contexts (2403.04317v2)

Abstract: Due to the rapid generation and dissemination of information, LLMs quickly run out of date despite enormous development costs. To address the crucial need to keep models updated, online learning has emerged as a critical tool when utilizing LLMs for real-world applications. However, given the ever-expanding corpus of unseen documents and the large parameter space of modern LLMs, efficient adaptation is essential. To address these challenges, we propose Memory of Amortized Contexts (MAC), an efficient and effective online adaptation framework for LLMs with strong knowledge retention. We propose a feature extraction and memory-augmentation approach to compress and extract information from new documents into compact modulations stored in a memory bank. When answering questions, our model attends to and extracts relevant knowledge from this memory bank. To learn informative modulations in an efficient manner, we utilize amortization-based meta-learning, which substitutes an otherwise required optimization process with a single forward pass of the encoder. Subsequently, we learn to choose from and aggregate selected documents into a single modulation by conditioning on the question, allowing us to adapt a frozen LLM during test time without requiring further gradient updates. Our experiment demonstrates the superiority of MAC in multiple aspects, including online adaptation performance, time, and memory efficiency. In addition, we show how MAC can be combined with and improve the performance of popular alternatives such as retrieval augmented generations (RAGs). Code is available at: https://github.com/jihoontack/MAC.

• The paper introduces MAC, a framework that efficiently adapts LLMs online by compressing new information into memory modulations.
• It leverages amortization-based meta-learning with backpropagation dropout to minimize training costs and manage memory constraints.
• Empirical evaluations demonstrate MAC’s superior capability in preserving knowledge and reducing adaptation time and GPU memory usage compared to existing methods.

Online Adaptation of LLMs with a Memory of Amortized Contexts

Overview

LLMs have rapidly become a cornerstone of contemporary NLP, driving improvements across a myriad of tasks and applications. However, the static nature of these models poses significant challenges in keeping their knowledge up-to-date, given the dynamic and evolving landscape of information. In light of this challenge, this paper introduces a Memory of Amortized Contexts (MAC), an innovative online adaptation framework designed to efficiently and effectively update LLMs to incorporate new information without the need for extensive retraining.

Addressing the Online Adaptation Challenge

Online adaptation of LLMs is a critical problem, especially for applications necessitating up-to-the-minute information. Traditional approaches, including retrieval-augmented models and gradient-based online fine-tuning, each carry limitations such as computational inefficiency, potential loss of previously acquired knowledge (catastrophic forgetting), or limited applicability in memory-restrained settings. In contrast, MAC proposes an efficacious solution, leveraging amortized feature extraction and memory augmentation to compress new information into compact modulations, which are then efficiently utilized to update a static LLM.

Methodology

At MAC's core is the use of amortization-based meta-learning, notably substituting traditional optimization processes with a more computationally efficient forward pass of an encoder. This is instrumental in generating a compact modulation for any new document, encapsulating relevant knowledge without necessitating direct adjustments to the LLM's parameters. Subsequently, relevancy-driven selection and aggregation of context modulations, conditioned on incoming queries, enable the model to dynamically adapt its responses based on preserved knowledge.

The framework introduces two memory-efficient techniques for both its training and inference phases:

• Backpropagation Dropout: This technique mitigates memory demands during training by processing only a subset of documents for gradient computation, ensuring both manageability and efficiency.
• Hierarchical Modulation Aggregation: This approach addresses memory constraints during inference by a divide-and-conquer strategy, iteratively aggregating information in manageable groups to derive a final, relevant modulation, thus significantly reducing GPU memory usage.

Empirical Validation

MAC's efficacy is comprehensively validated across multiple datasets and model architectures, showcasing superior online adaptation performance, notable for both its accuracy and efficiency compared to existing methods. Experiments highlight MAC's ability to retain knowledge effectively through its adaptation process, underlining its practical utility for real-world applications.

Furthermore, efficiency evaluations elucidate MAC's advantage in memory and computational resource utilization, an essential consideration given the prevailing constraints associated with deploying large-scale models. These findings underscore MAC’s potential in significantly reducing adaptation time and memory usage without compromising on performance.

Conclusion and Future Directions

This paper's exploration and subsequent introduction of MAC reiterates the importance of efficient and effective online adaptation for LLMs. By addressing the limitations of existing approaches and presenting a robust framework, MAC paves the way for more dynamic, up-to-date, and efficient utilization of LLMs in various applications. Future research avenues might include exploring MAC's applicability in federated learning contexts or implementing privacy-preserving mechanisms for sensitive information in the documented memory bank, further broadening MAC's utility and applicability.