Generalization Guarantees for Representation Learning via Data-Dependent Gaussian Mixture Priors

TL;DR

This paper derives data-dependent generalization bounds for representation learning using Gaussian mixture priors, introduces a regularization method based on these bounds, and demonstrates improved performance over existing methods.

Contribution

It provides novel data-dependent generalization bounds for representation learning and proposes a systematic approach to learn and utilize Gaussian mixture priors as regularizers.

Findings

Bounds reflect the encoder's structure and simplicity.

The proposed regularizer outperforms VIB and CDVIB methods.

A natural weighted attention mechanism emerges in the prior learning process.

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).