Title: Branching Memory: Task-Specific Expansion for Continual Learning in Large Language Models
Primary area: transfer learning, meta learning, and lifelong learning
Keywords: Continual Learning, Catastrophic Forgetting, LLMs, MLP, Feed-Forward Layers

Abstract: Large Language Models (LLMs) face the challenge of catastrophic forgetting in continual learning scenarios, where learning new tasks often overwrites previously acquired knowledge, leading to performance degradation and limiting their applicability in dynamic task environments. Existing approaches can be categorized into rehearsal-based, regularization-based, and architecture-based methods. Among these, architecture-based methods are more suitable for LLMs as they dynamically adjust model structures to handle large-scale parameters and task interference while avoiding the storage overhead of replay-based methods and the adaptability issues of regularization-based methods. However, existing architecture-based methods face challenges in efficiently training parameters within the large parameter space of LLMs. Additionally, these methods do not fully leverage the characteristics of the Transformer architecture, resulting in limited effectiveness in mitigating catastrophic forgetting.
In this work, we propose Branching Memory, a novel method that leverages the organization of knowledge within transformer models. By modeling knowledge as key-value (KV) representations within the FFN layers, our approach dynamically allocates dedicated capacity for new tasks, allowing the model to store and integrate task-specific knowledge without overwriting existing information. To further improve knowledge retention and reduce task interference, we employ an orthogonality-based regularization strategy and a task-aware initialization method  to stabilize training and minimize parameter conflicts.
Experimental results on standard continual learning benchmarks show that Branching Memory outperforms most existing methods in LLM architectures. For example, in short-sequence tasks with T5-Large, Branching Memory with the regularization method achieves an improvement of over 11\% in accuracy compared to other baselines. This improvement highlights the effective mitigation of forgetting, task adaptability, and practicality in continual learning.