Source-linked AI summary
Uncertainty-guided Continual Learning with Bayesian Neural Networks
Sayna Ebrahimi, Mohamed Elhoseiny, Trevor Darrell, Marcus Rohrbach
TL;DR
Continual learning must acquire sequential tasks without forgetting when prior data are unavailable and model capacity is limited. The paper proposes UCB, which adapts Bayesian weight learning rates using uncertainty, and UCB-P, which prunes and masks important weights. Across diverse short and long task sequences, the methods perform superiorly or on par with existing approaches and can operate without task information at test time.
Problem
Continual-learning systems often forget earlier tasks when data arrive sequentially and previous training data are unavailable, especially with limited model capacity.
Method
UCB uses Bayesian per-weight uncertainty to adapt learning rates, while UCB-P uses uncertainty to prune important weights and preserve them with task-specific binary masks.
Results
UCB and UCB-P achieve superior or on-par performance with state-of-the-art approaches across short and long task sequences, including generalized accuracy without task information at test time.
Takeaways & Limitations
Bayesian weight uncertainty can support continual learning across diverse object-classification task sequences and single-head evaluation without task labels.
Takeaways & Limitations
Hyperparameters are tuned using validation data from only the first two tasks rather than all classes in the dataset.
Abstract
from arXiv · showhide
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.
1 INTRODUCTION
Continual learning addresses catastrophic forgetting when models learn sequential tasks without access to prior data. UCB uses Bayesian weight uncertainty to adapt learning rates, while UCB-P freezes selected weights through pruning and task-specific masks.
- Catastrophic forgetting is the significant performance drop that occurs when parameters learned for an initial task change to favor a new objective.
- Limited-capacity continual-learning methods traditionally estimate parameter importance explicitly to penalize or freeze changes, whereas UCB represents importance implicitly through uncertainty.
- Bayesian neural networks represent each parameter with a distribution over possible values, naturally providing uncertainty estimates for parameter adaptation.
- UCB adapts each parameter’s learning rate according to uncertainty, taking larger gradient steps for more uncertain and therefore more learnable weights.
- UCB-P is a hard-threshold variant that prunes and freezes important parameters while preserving task-specific binary masks for subsequent learning.
- The approach is evaluated on task sequences and supports inference without task-boundary information, using a single-head generalized-accuracy setting.
2 RELATED WORK
Continual-learning research includes dynamic architectures, memory-based rehearsal, and regularization methods that protect prior representations. UCB extends Bayesian approaches by using uncertainty for parameter adaptation and evaluates convolutional Bayesian networks beyond the usual MLP setting.
- Continual-learning methods are categorized as dynamic architectural, memory-based, or regularization approaches.
- Dynamic architectural methods: Dynamic architectural methods preserve past knowledge by expanding model capacity with task-specific layers, nodes, or modules.
- Memory-based methods: Memory-based methods store portions of previous data for rehearsal, with GEM constraining later gradient updates using stored examples.
- Regularization methods: Regularization methods prevent major changes to prior representations by penalizing or directly restricting updates to parameters deemed important.
- Bayesian approaches: Bayesian continual-learning approaches use probabilistic parameter estimates, while UCB uses classic gradient descent rather than natural-gradient optimization.
- The paper evaluates convolutional Bayesian neural networks for continual learning rather than limiting evaluation to fully connected MLP networks.
3 BACKGROUND: VARIATIONAL BAYES-BY-BACKPROP
Bayes-by-Backprop learns a variational posterior over neural-network weights by optimizing a tractable approximation to the true posterior. The framework samples Gaussian-distributed weights through reparameterization and supplies per-weight uncertainty for continual learning.
- Bayes-by-Backprop learns a probability distribution over network parameters using a back-propagation-compatible variational framework.
- Variational inference approximates the posterior over latent weights by finding the closest member of a parameterized density family under KL divergence.
- The objective function expresses the variational-learning criterion used to fit the approximate posterior.
- The objective can be estimated with N Monte Carlo samples drawn from the variational posterior.
- The variational posterior uses a diagonal-covariance Gaussian parameterized by θ=(µ,ρ), with weights sampled by reparameterization as w=µ+σ ◦ ϵ.
- Per-parameter uncertainty from this framework is used in UCB to learn sequential tasks without forgetting.
4 UNCERTAINTY-GUIDED CONTINUAL LEARNING IN BAYESIAN NEURAL NETWORKS
UCB uses Bayesian weight uncertainty as an implicit importance measure, adapting parameter learning rates to reduce changes in important parameters while allowing uncertain ones to learn new tasks. UCB-P hardens this idea by pruning low-importance weights and preserving task-specific binary masks.
- UCB with learning rate regularization: UCB regulates each parameter’s gradient update through a learning rate determined by its importance.The method performs regularization through learning rates rather than directly modifying the objective or weights.
- UCB with learning rate regularization: Parameter importance is set inversely proportional to posterior standard deviation, so more uncertain parameters receive larger learning updates.The posterior standard deviation σ represents parameter uncertainty, and larger uncertainty permits larger gradient steps for learning the current task.
- UCB with learning rate regularization: The ablation study found Ωµ = 1/σ with Ωρ = 1 achieved the highest accuracy and least forgetting.This setting adapts the learning rate for the posterior mean while leaving the ρ learning-rate importance unchanged.
- UCB with learning rate regularization: UCB does not require additional memory or tracking parameter changes relative to previously learned tasks.This distinguishes its learning-rate regularization from pruning-based memory and common weight-regularization methods.
- UCB using weight pruning (UCB-P): UCB-P freezes parameters in the top p% of layer-wise SNR values by setting their learning rates to zero.The hard-threshold update uses Ω = |µ|/σ as the SNR-based importance measure.
- UCB using weight pruning (UCB-P): UCB-P ranks parameters by SNR within each layer, prunes the lowest-ranked weights, and saves binary masks for later task learning.Important parameters remain fixed during future-task training, while pruned parameters remain available for reuse.
5 RESULTS
UCB is evaluated across class-incremental, alternating-dataset, and multiple-dataset continual-learning scenarios using Bayesian and non-Bayesian baselines. Experiments report average accuracy and backward transfer, including pruning, hyperparameter, and task-label-free settings.
- Experimental setup: Experiments cover class-incremental learning, alternating CIFAR10/CIFAR100 tasks, and an eight-dataset sequence with varied distributions.The evaluation includes single-dataset and multiple-dataset task sequences, with no data augmentation used in the eight-dataset analysis.
- Experimental setup: UCB experiments compare against Bayesian fine-tuning, feature extraction, and joint training, plus established continual-learning baselines including EWC, IMM, LWF, LFL, PathNet, PNNs, and HAT.Joint training is treated as an upper bound for average accuracy because it does not follow the continual-learning setup.
- Evaluation metrics: The evaluation measures ACC as average test accuracy across tasks and BWT as the effect of later-task learning on earlier-task performance, with BWT < 0 indicating catastrophic forgetting.For UCB-P, pre-pruning accuracy is used for Ri,i and post-pruning accuracy for Ri,n.
- MNIST results: On 2-Split MNIST, UCB and UCB-P achieve zero forgetting, while UCB has marginally higher accuracy than PackNet, HAT, and UCB-P.The reported comparison includes PackNet, HAT, LWF, UCB-P, and UCB.
- Permuted MNIST results: 91.44 ± 0.04% ACC with BWT=(0.03% ± 0.00%) lets small-network UCB outperform SI, EWC, VCL, VCL-Vadam, and VCL-GNG in reported comparisons, although HAT reaches 91.6%.For the 1.9M-parameter network, UCB reaches ACC=97.42±0.01% and performs better than HAT at ACC=97.34 ± 0.05%, with both showing zero forgetting.
- Multi-dataset results: Across alternating CIFAR10/CIFAR100 and eight-dataset sequences, UCB variants outperform or closely compete with strong baselines, including 3.6% higher ACC than PNN for UCB-P.In the eight-task experiment, UCB achieves 2.4% higher ACC than HAT, while UCB-P outperforms PNN by 3.6% in ACC.
6 SINGLE HEAD AND GENERALIZED ACCURACY OF UCB
UCB can operate with a single classification head when task identity is unavailable at inference, with only small accuracy reductions compared with multi-head evaluation. Generalized accuracy also remains close to standard accuracy across the tested task sequences.
- UCB uses one head covering all task classes instead of separate task-specific heads during training and inference.In the 8-dataset experiment, the single head has 293 output classes.
- 0.3%, 2.6%, 5.1%, and 4.1% ACC reductions occur for 2-Split MNIST, Permuted MNIST, Alternating CIFAR10/100, and the 8-task sequence, respectively.
- 0.2%, 2.6%, 3.1%, and 3.1% performance degradation occurs from ACC to Generalized ACC for the same four experiments, respectively.Generalized accuracy evaluates predictions over classes from all tasks without task information at test time.
- The small ACC-to-Generalized-ACC reductions indicate limited confusion from presenting more classes at test time.
7 CONCLUSION
The conclusion presents UCB as a Bayesian continual-learning formulation that uses weight uncertainty either to preserve important parameters or regulate their adaptation. Across short and long benchmark sequences, UCB is superior or on par with prior methods and can operate without task identity at inference.
- UCB uses predicted weight uncertainty to preserve important parameters with binary masks in UCB-P or regulate their changes in UCB.
- UCB-P prevents forgetting after initial pruning by saving a small binary mask per task, whereas UCB uses no additional memory and permits more flexible learning with some forgetting.
- UCB performs superior or on par with state-of-the-art models such as HAT across the reported experiments.
- UCB supports continual learning on both short and long sequences of benchmark datasets using probabilistic uncertainty distributions for each weight.
- A single-head UCB model remains competitive when the task-specific class subset is unknown during inference.
A.1 DATASETS
The experiments use continual-learning benchmarks spanning class-incremental tasks within datasets and sequences in which each task is a different dataset. Images are standardized to three-channel 32 × 32 inputs when necessary.
- Images are resized to 32 × 32 × 3 when necessary, and monochromatic images are replicated across all RGB channels.
A.2 IMPLEMENTATION DETAILS
Implementation studies examine UCB components, Bayesian-network tuning, architecture choices, and Monte Carlo sampling while comparing continual-learning models under matched capacity where described.
- A focused MNIST study evaluates parameter-regularization variants, importance measurements, and the number of posterior samples.
- BBB hyperparameters are tuned on the validation set of the first two tasks using a specified search space.
- Split and Permuted MNIST use a two-layer 1200-unit perceptron, while multi-dataset and alternating CIFAR10/100 experiments use Bayesian ResNet18 networks.The reported Bayesian-network parameter counts are 1.9M for the MNIST settings and 7.1–11.3M for the larger scenarios, depending on the experiment.
- The CIFAR10/100 comparison evaluates continual learning with AlexNet and ResNet18 for UCB and HAT using BWT and ACC percentages.
- UCB uses multiple posterior samples for robustness to random noise and selects 10 samples for all experiments.
A.3 ADDITIONAL RESULTS
Additional results cover continual learning on 2-split MNIST, 5-split MNIST, Permuted MNIST, and CIFAR10/100, reporting BWT and ACC in percentages. The tables also clarify setup exceptions, upper bounds, reporting sources, and standard deviations.
- Standard deviations for the main-text Table 2a, 2b, and 2c results are provided in Tables 9, 10, and 11, respectively.
- Tables 8 and 9 report BWT and ACC for continual learning on 2-split and 5-split MNIST, respectively.
- For 2-split and 5-split MNIST, BBB-JT and ORD-JT serve as upper bounds for ACC for BBB and ORD networks, respectively.
- Tables 10 and 11 report BWT and ACC for Permuted MNIST and CIFAR10/100, respectively.
- In the Permuted MNIST table, BBB-JT is identified as the BBB network's ACC upper bound, while some methods do not follow the continual learning setup.
- Results are reproduced by the authors except where Tables 10 and 11 identify results from prior reports or original work.