The Effect of Architecture During Continual Learning
Allyson Hahn, Krishnan Raghavan
TL;DR
The paper addresses the instability of continual learning under distribution shifts when using a fixed neural-architecture. It introduces a Sobolev-space framework that jointly models architecture and weights, and proves that weight-only updates cannot prevent forgetting; a bilevel optimization with derivative-free architecture search and a low-rank transfer mechanism is proposed to adapt both architecture and parameters across tasks. The approach, termed AWB (architecture search with low-rank transfer and weight training), is theoretically grounded in continuity/absolute continuity of the forgetting cost and demonstrated empirically across regression, image, and graph learning, achieving substantial improvements in accuracy and reduced forgetting (often by large margins) while maintaining robustness to noise. The practical impact is a scalable, architecture-adaptive continual learning method that preserves and transfers knowledge across evolving task distributions, applicable to FFN, CNN, and GNN architectures with strong empirical performance gains. Overall, the work provides both a rigorous mathematical foundation and a comprehensive empirical validation for jointly learning architecture and weights in continual learning, highlighting the critical role of capacity expansion via architecture changes when distributions drift.
Abstract
Continual learning is a challenge for models with static architecture, as they fail to adapt to when data distributions evolve across tasks. We introduce a mathematical framework that jointly models architecture and weights in a Sobolev space, enabling a rigorous investigation into the role of neural network architecture in continual learning and its effect on the forgetting loss. We derive necessary conditions for the continual learning solution and prove that learning only model weights is insufficient to mitigate catastrophic forgetting under distribution shifts. Consequently, we prove that by learning the architecture and weights simultaneously at each task, we can reduce catastrophic forgetting. To learn weights and architecture simultaneously, we formulate continual learning as a bilevel optimization problem: the upper level selects an optimal architecture for a given task, while the lower level computes optimal weights via dynamic programming over all tasks. To solve the upper level problem, we introduce a derivative-free direct search algorithm to determine the optimal architecture. Once found, we must transfer knowledge from the current architecture to the optimal one. However, the optimal architecture will result in a weights parameter space different from the current architecture (i.e., dimensions of weights matrices will not match). To bridge the dimensionality gap, we develop a low-rank transfer mechanism to map knowledge across architectures of mismatched dimensions. Empirical studies across regression and classification problems, including feedforward, convolutional, and graph neural networks, demonstrate that learning the optimal architecture and weights simultaneously yields substantially improved performance (up to two orders of magnitude), reduced forgetting, and enhanced robustness to noise compared with static architecture approaches.