Adaptive Compression of the Latent Space in Variational Autoencoders

TL;DR

This paper tackles the problem of selecting the latent-space dimensionality $n_z$ in Variational Autoencoders, a major hyperparameter influencing reconstruction quality and clustering. It proposes ALD-VAE, an automated method that gradually prunes neurons to adapt $n_z$ during training, using a multi-metric stopping rule that combines reconstruction loss, Fréchet Inception Distance (FID), and Silhouette score. Across four image datasets, ALD-VAE achieves performance close to or matching a model trained with the optimal fixed $n_z$, but with significantly reduced computation compared to grid search. The approach is simple to implement, end-to-end trainable, and can alleviate manual hyperparameter tuning in practical VAE deployments, with potential applicability to other data domains beyond images.

Abstract

Variational Autoencoders (VAEs) are powerful generative models that have been widely used in various fields, including image and text generation. However, one of the known challenges in using VAEs is the model's sensitivity to its hyperparameters, such as the latent space size. This paper presents a simple extension of VAEs for automatically determining the optimal latent space size during the training process by gradually decreasing the latent size through neuron removal and observing the model performance. The proposed method is compared to traditional hyperparameter grid search and is shown to be significantly faster while still achieving the best optimal dimensionality on four image datasets. Furthermore, we show that the final performance of our method is comparable to training on the optimal latent size from scratch, and might thus serve as a convenient substitute.

Figures (4)

• Figure 1: Overview of our training procedure with shrinking latent space size. After each epoch, we encode a batch of 500 random samples from the validation dataset and run K-means clustering. We then calculate the Silhouette score based on the detected clusters. The Decision rule observes the approximate slope of the Silhouette score, FID and the reconstruction loss over epochs. We prune 1-5 neurons after each 5 epochs from the encoder's $\boldsymbol{\mu}$ and $\boldsymbol{\sigma}$ layers and the decoder's first layer until the slopes meet our criteria (see Section \ref{['sec:reduction']}), then the latent space size remains fixed.
• Figure 2: The values of the four observed metrics (FID for reconstructions, FID for generations, Silhouette score and Reconstruction loss) during training when the latent space is gradually reduced until Dim=2. The upper axis shows the dimensionality at the given epoch, the vertical red line shows when would our ALD algorithm stop the compression. We show the results for four datasets.
• Figure 3: Comparison of the four metrics when using a fixed latent dimensionality for the whole training (blue curves) or our ALD-VAE latent space compression that converged to the same final dimensionality (orange curves). We show FID for reconstructed and generated images (lower is better), Silhouette score (higher is better) and Reconstruction loss (lower is better). Shown on four different datasets.
• Figure 4: Training our ALD-VAE with different initial conditions. $ALD_1$ was trained with initial latent dimensionality $n_z=100$, $ALD_2$ was trained with initial $n_z=80$ and $ALD_3$ was trained with initial $n_z$=64. We show the results for the MNIST (left) and FashionMNIST (right) datasets. The $Dim$ values in brackets for each ALD-VAE show the final converged $n_z$.