Dynamic Narrowing of VAE Bottlenecks Using GECO and L0 Regularization

TL;DR

This paper introduces a method to automatically and dynamically reduce the latent space dimensionality of VAEs during training, ensuring optimal size without extensive cross-validation.

Contribution

We develop a novel technique combining GECO and L0 regularization to prune VAE latent spaces on-the-fly, improving efficiency and adaptability.

Findings

Effective latent space pruning without violating GECO constraints

Stable training procedure demonstrated across five datasets

Automatic dimensionality adjustment reduces need for cross-validation

Abstract

When designing variational autoencoders (VAEs) or other types of latent space models, the dimensionality of the latent space is typically defined upfront. In this process, it is possible that the number of dimensions is under- or overprovisioned for the application at hand. In case the dimensionality is not predefined, this parameter is usually determined using time- and resource-consuming cross-validation. For these reasons we have developed a technique to shrink the latent space dimensionality of VAEs automatically and on-the-fly during training using Generalized ELBO with Constrained Optimization (GECO) and the $L_0$-Augment-REINFORCE-Merge ($L_0$-ARM) gradient estimator. The GECO optimizer ensures that we are not violating a predefined upper bound on the reconstruction error. This paper presents the algorithmic details of our method along with experimental results on five different datasets. We find that our training procedure is stable and that the latent space can be pruned effectively without violating the GECO constraints.