Uncertainty-guided Continual Learning with Bayesian Neural Networks

TL;DR

The paper tackles catastrophic forgetting in continual learning under fixed capacity by leveraging uncertainty in Bayesian neural networks. It introduces Uncertainty-guided Continual Bayesian Neural Networks (UCB), which adapt learning rates per parameter using posterior uncertainty $\sigma$, and a hard-pruning variant UCB-P that freezes important weights via task-specific masks. Across MNIST variants and multi-dataset sequences, UCB and UCB-P achieve state-of-the-art or competitive results, often with minimal forgetting and without requiring task labels at test time (single-head/generalized accuracy). The work highlights uncertainty as a natural criterion to decide what to remember and what to change, offering practical deployment without memory replay or explicit task boundaries.

Abstract

Continual learning aims to learn new tasks without forgetting previously learned ones. This is especially challenging when one cannot access data from previous tasks and when the model has a fixed capacity. Current regularization-based continual learning algorithms need an external representation and extra computation to measure the parameters' \textit{importance}. In contrast, we propose Uncertainty-guided Continual Bayesian Neural Networks (UCB), where the learning rate adapts according to the uncertainty defined in the probability distribution of the weights in networks. Uncertainty is a natural way to identify \textit{what to remember} and \textit{what to change} as we continually learn, and thus mitigate catastrophic forgetting. We also show a variant of our model, which uses uncertainty for weight pruning and retains task performance after pruning by saving binary masks per tasks. We evaluate our UCB approach extensively on diverse object classification datasets with short and long sequences of tasks and report superior or on-par performance compared to existing approaches. Additionally, we show that our model does not necessarily need task information at test time, i.e. it does not presume knowledge of which task a sample belongs to.

Figures (1)

• Figure 1: Illustration of the evolution of weight distributions -- uncertain weights adapt more quickly -- when learning two tasks using $\text{UCB}$. (a) weight parameter initialized by distributions initialized with mean and variance values randomly sampled from $\mathcal{N}(0,0.1)$. (b) posterior distribution after learning task one; while $\theta_1$ and $\theta_2$ exhibit lower uncertainties after learning the first task, $\theta_3$, $\theta_4$, and $\theta_5$ have larger uncertainties, making them available to learn more tasks. (c) a second task is learned using higher learning rates for previously uncertain parameters ($\theta_1$, $\theta_2$, $\theta_3$, and $\theta_4$) while learning rates for $\theta_1$ and $\theta_2$ are reduced. Size of the arrows indicate the magnitude of the change of the distribution mean upon gradient update.