Table of Contents
Fetching ...

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

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

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

Paper Structure

This paper contains 38 sections, 5 theorems, 134 equations, 4 figures, 2 tables.

Key Result

Theorem 1

Consider a $C$-class classification problem and a learning algorithm $\mathcal{A}\colon \mathcal{Z}^n\to \mathcal{W}$ that induces the joint distribution $(S',S,W,\mathbf{U},\mathbf{U'},\hat{\mathbf{Y}},\hat{\mathbf{Y}}') \sim P_{S'} P_{S,W} P_{\mathbf{U}, \mathbf{U}'|\mathbf{X},\mathbf{X}',W_e} where $\hat{p}_{\mathbf{Y}}$ and $\hat{p}_{\mathbf{Y}'}$ are empirical distributions of $\mathbf{Y}

Figures (4)

  • Figure 1: Studied representation learning setup.
  • Figure 2: Values of $h_C\left( \mathcal{\hat{L}}(\mathbf{y},\mathbf{ \hat{y}}),\mathcal{\hat{L}}(\mathbf{y}',\mathbf{ \hat{y}}') ;\epsilon\right)$ for various values of the generalization error for the CIFAR10 dataset.
  • Figure 3: Comparison of the generalization bounds of Theorem \ref{['th:generalizationExp_hd']} (for various values of $\hat{\mathcal{L}}(S,W)$) and sefidgaran2023minimum for the CIFAR10 dataset.
  • Figure 4: Test performance of the CNN-based encoder trained on CIFAR10 using standard VIB alemi2016deep regularization, Category-dependent VIB (CDVIB) sefidgaran2023minimum regularization, and our proposed Gaussian Mixture MDL (GM-MDL) regularization.

Theorems & Definitions (10)

  • Definition 1: Symmetric prior
  • Theorem 1
  • Theorem 2
  • Definition 2: Strongly symmetric prior
  • Definition 3: Differentially private prior
  • Proposition 1
  • Definition 4: Partially symmetric prior
  • Proposition 2
  • Lemma 1
  • Definition 5