TLDR: AIMMerging is a novel continual learning framework for large language models that uses learning and forgetting signals from training trajectories to dynamically schedule iterative model merges. This adaptive approach effectively balances acquiring new knowledge and preventing catastrophic forgetting, demonstrating significant performance improvements over state-of-the-art methods across various benchmarks and model sizes.
Large language models (LLMs) are constantly evolving, and keeping them updated with new information without forgetting old knowledge is a significant challenge. This process, known as continual learning (CL), is crucial for deploying LLMs in real-world scenarios where new data emerges continuously. Traditional methods often struggle to balance learning new information with retaining previously acquired knowledge, leading to what is called “catastrophic forgetting.”
A new research paper introduces a novel framework called Adaptive Iterative Model Merging, or AIMMerging, designed to tackle these challenges. The core idea behind AIMMerging is to dynamically monitor an LLM’s training status by observing its “training trajectory” – essentially, how its parameters change and how well it remembers old data during the learning process. This dynamic monitoring allows the model to decide the optimal times to merge new knowledge without compromising old knowledge.
AIMMerging operates through two main components: a Training Trajectory-guided Merge Controller and a Rehearsal-based Knowledge Fusion Module. The merge controller acts like a smart scheduler, using two key signals: a “learning signal” and a “forgetting signal.” The learning signal tracks how quickly the model is acquiring new knowledge by observing changes in its parameters. If the model is learning rapidly, the controller might increase the merging frequency to integrate this new knowledge efficiently. Conversely, if learning slows down, it might reduce the merging frequency to avoid unnecessary merges.
The forgetting signal is equally important. It monitors the model’s performance on historical data. If the model starts to forget previously learned information (indicated by an increase in loss on old data), the forgetting signal triggers an immediate merge. This proactive approach helps prevent catastrophic forgetting before it becomes a major issue. By combining both signals, the controller adaptively adjusts when and how often merges occur, striking a delicate balance between learning new tasks and preserving old ones.
Once the merge controller decides it’s time to merge, the Rehearsal-based Knowledge Fusion Module takes over. This module calculates specific “merging weights” based on the strength of the learning and forgetting signals. These weights determine how much emphasis to place on new knowledge versus historical knowledge during the fusion process. It then integrates the new and old knowledge into a unified model, effectively updating the LLM without retraining it from scratch.
The researchers conducted extensive experiments across three continual learning benchmarks and with various LLM sizes, ranging from 770 million to 13 billion parameters. The results were impressive: AIMMerging showed significant performance improvements over existing state-of-the-art methods. On average, it achieved an 80% relative improvement in Forward Transfer (FWT), which measures the model’s ability to learn new tasks, and a 59% relative improvement in Backward Transfer (BWT), which indicates how well it retains old knowledge. This demonstrates AIMMerging’s effectiveness in both acquiring new skills and mitigating forgetting.
The study also included ablation studies, which confirmed the importance of both the learning and forgetting signals for the merge controller’s effectiveness. Removing either signal led to a decline in performance, highlighting their synergistic role. Furthermore, the method proved robust across different model architectures and scales, showcasing its generalizability. For more technical details, you can refer to the full paper available here.
Also Read:
- Designing AI Models That Can Forget On Demand
- Unlocking Efficiency in Language Models: A New Bias-Selection Method for Fine-Tuning
In essence, AIMMerging offers a smarter, more adaptive way for large language models to continuously learn and evolve. By intelligently managing when and how new information is integrated, it helps LLMs stay relevant and performant in dynamic environments, paving the way for more robust and adaptable AI systems.
