Variational Rank Reduction Autoencoders

TL;DR

VRRAEs address the tension between deterministic, rank-regularized latent spaces and probabilistic generative modeling by introducing a variational treatment of the rank-reduced bottleneck from RRAEs. They sample SVD-based coefficients from a learned Gaussian while retaining a fixed SVD basis, enabling both sharper generation and reduced posterior collapse. Empirical results on synthetic data and real-world datasets (MNIST, CelebA, CIFAR-10) show VRRAEs outperform VAEs and RRAEs in interpolation and random generation as measured by FID, and exhibit lower reconstruction error with improved robustness to posterior collapse. The work provides an open-source JAX implementation and demonstrates the practical value of integrating rank-based regularization into variational generative modeling for robust and high-quality samples.

Abstract

Deterministic Rank Reduction Autoencoders (RRAEs) enforce by construction a regularization on the latent space by applying a truncated SVD. While this regularization makes Autoencoders more powerful, using them for generative purposes is counter-intuitive due to their deterministic nature. On the other hand, Variational Autoencoders (VAEs) are well known for their generative abilities by learning a probabilistic latent space. In this paper, we present Variational Rank Reduction Autoencoders (VRRAEs), a model that leverages the advantages of both RRAEs and VAEs. Our claims and results show that when carefully sampling the latent space of RRAEs and further regularizing with the Kullback-Leibler (KL) divergence (similarly to VAEs), VRRAEs outperform RRAEs and VAEs. Additionally, we show that the regularization induced by the SVD not only makes VRRAEs better generators than VAEs, but also reduces the possibility of posterior collapse. Our results include a synthetic dataset of a small size that showcases the robustness of VRRAEs against collapse, and three real-world datasets; the MNIST, CelebA, and CIFAR-10, over which VRRAEs are shown to outperform both VAEs and RRAEs on many random generation and interpolation tasks based on the FID score. We developed an open-source implementation of VRRAEs in JAX (Equinox), available at https://github.com/JadM133/RRAEs.git.

Figures (13)

• Figure 1: Schematic illustrating the architecture of Variational Rank Reduction Autoencoders (VRRAEs). Both $E$ and $D$ are trainable Neural Networks representing an encoding and a decoding map. $SVD_{k^*}$ is a truncated SVD of rank $k^*$ as detailed in Equation \ref{['eq:RRAE']}. Note that $f=I$ (identity) for VRRAEs.
• Figure 2: Visualization of the second latent mean plotted against the first one for a certain seed in the latent space for VAEs (left) and VRRAEs (right). The collapse of the second latent dimension in VAEs is evidenced by the small variations along the $y$-axis (i.e., $0.02\leq\zeta_i\leq0.09$ in Equation \ref{['eqn:zeta']}). In contrast, VRRAEs do not suffer from collapse, and preserve the structure of a diagonal motion (highlighted in red), maintaining interpretability.
• Figure 3: Randomly generated samples for different architectures on the 2D gaussian problem
• Figure 4: Example of an interpolation (linear, in the latent space) between two CelebA samples. Note that the VRRAE is the only model to capture the right skin color, even though all models have the same bottleneck size.
• Figure 5: Randomly generated MNIST samples for each one of the selected models.

Theorems & Definitions (1)

• Remark 1