Generalization Guarantees for Representation Learning via Data-Dependent Gaussian Mixture Priors
Milad Sefidgaran, Abdellatif Zaidi, Piotr Krasnowski
TL;DR
This work derives in-expectation and tail generalization bounds for representation-learning algorithms that depend on the minimum description length (MDL) of latent variables relative to a data-dependent symmetric prior, improving over MI-based bounds by capturing encoder structure. It then optimizes a data-driven Gaussian mixture prior to regularize the encoder, presenting both lossless and lossy variants; the lossy version yields a weighted-attention-like mechanism that dynamically concentrates latent representations around learned centroids. The proposed GM-MDL regularizers translate MDL-based guarantees into practical training objectives, with EM-like updates for mixture components and an approximate KL estimator for Gaussian mixtures. Empirically, GM-MDL regularization, particularly in the lossy form, yields superior generalization performance compared to Variational Information Bottleneck and Category-Dependent VIB across multiple datasets and encoder architectures, illustrating the practical impact of MDL-centric regularization in representation learning.
Abstract
We establish in-expectation and tail bounds on the generalization error of representation learning type algorithms. The bounds are in terms of the relative entropy between the distribution of the representations extracted from the training and "test'' datasets and a data-dependent symmetric prior, i.e., the Minimum Description Length (MDL) of the latent variables for the training and test datasets. Our bounds are shown to reflect the "structure" and "simplicity'' of the encoder and significantly improve upon the few existing ones for the studied model. We then use our in-expectation bound to devise a suitable data-dependent regularizer; and we investigate thoroughly the important question of the selection of the prior. We propose a systematic approach to simultaneously learning a data-dependent Gaussian mixture prior and using it as a regularizer. Interestingly, we show that a weighted attention mechanism emerges naturally in this procedure. Our experiments show that our approach outperforms the now popular Variational Information Bottleneck (VIB) method as well as the recent Category-Dependent VIB (CDVIB).
