Efficient Bilevel Optimization for CKA-Guided MoE Upcycling
Abstract
Upcycling, a strategy that initializes Mixture-of-Experts (MoE) by replicating pre-trained feed-forward or MoE networks to expand model capacity, has become a popular method in continual learning due to its effectiveness in mitigating catastrophic forgetting. However, existing paradigms indiscriminately expand capacity to prioritize performance at the cost of severe inefficiency, introducing severe parameter redundancy and failing to exploit structural heterogeneity. To address this, we investigate the determinants of forgetting in training dynamics using Centered Kernel Alignment (CKA) and loss landscape flatness to analyze the behavior of pre- and post-expansion MoE layers, uncovering instability in deep-layer representations and heterogeneous expert sensitivity to new tasks, thereby demonstrating the potential of selective upcycling to eliminate redundancy. Consequently, we propose a dynamic bilevel optimization framework to guide adaptive upcycling, featuring an outer loop employing a Gumbel-Softmax differentiable mask to perform Neural Architecture Search (NAS) for adaptive growth, while an inner loop optimizes weight updates via task objectives and CKA-regularized replay. Experiments on the TRACE benchmark demonstrate that our proposed method achieves better average accuracy with 80\% forgetting reduction, while effectively eliminating 60\% of redundant parameter expansion that standard upcycling would introduce.
Lay Summary
When artificial intelligence models learn new skills, they often suffer from catastrophic forgetting, where they overwrite and lose their previously acquired knowledge. To prevent this, developers expand the model's capacity by copying its internal components, a brute-force strategy that wastes massive amounts of memory and computing power. To solve this inefficiency, we created a smart, targeted expansion method for AI models. Instead of growing the whole network uniformly, our system acts like a careful architect. It mathematically analyzes the model's internal layers to identify exactly which parts are struggling with new information and which parts are stable, adding new capacity only where it is absolutely needed. By dynamically allocating this brainpower, our method allows the AI to seamlessly learn a sequence of new tasks without forgetting old ones. Most importantly, it matches the high accuracy of standard expansion methods while eliminating 60% of the unnecessary bulk. This makes the deployment of advanced, adaptable AI significantly cheaper, faster, and more sustainable.