Wake-Sleep Energy Based Models for Continual Learning

Vaibhav Singh, Anna Choromanska, Shuang Li, Yilun Du; Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR) Workshops, 2024, pp. 4118-4127

This paper introduces a novel approach for continually training Energy-Based Models (EBMs) on the classification problems in the challenging setting of class incremental learning. Despite the fact that EBMs offer longer retention of knowledge on prior tasks training EBMs contrastively remains a challenge. Driven by biological plausibility we leverage the observation that sleep in humans supports active system consolidation and propose a new approach for training EBMs which we call Wake-Sleep Energy Based Models (WS-EBMs) which rely on wake-sleep cycles. Our training approach consists of short wake phases followed by long sleep phases. During the short wake phase the free energy associated with ground truth labels is minimized which conditions the model towards the correct solutions. This is followed by a long sleep phase where the free energy of the whole system is minimized contrastively which allows the model to push the energy of incorrect solutions further from the correct response. We provide a theoretical analysis of WS-EBM showing that it satisfies the sufficient condition for designing proper EBM loss. Our empirical evaluation confirms the plausibility of our approach and demonstrates favorable performance of WS-EBM compared to traditional EBM training as well as state-of-the-art class-incremental continual learning techniques. Furthermore our proposed two-phase training strategy can be easily integrated with existing techniques resulting in substantial boosts in their performance. Finally we also provide interesting insights justifying our approach by analyzing the orthogonality between the sequential task vectors and flatness of the optimized energy surfaces which may guide the design of class incremental continual learning strategies.