SC-VAE: Sparse Coding-based Variational Autoencoder with Learned ISTA

TL;DR

SC-VAE introduces a sparse coding-based VAE that represents latent variables as sparse combinations of fixed orthogonal atoms, learned end-to-end via a Learnable ISTA. By using a fixed DCT dictionary and a differentiable ISTA unrolled network, SC-VAE avoids posterior and codebook collapse while enabling high-quality reconstruction, controllable generation, and unsupervised image segmentation through patch-level sparse codes. Experiments on FFHQ and ImageNet show superior image reconstruction compared to strong baselines, and the learned sparse codes support meaningful image generation, smooth interpolation, and effective clustering of image patches for segmentation. Supplementary materials provide detailed architectural choices, dictionary visualizations, and expanded qualitative and quantitative results across datasets and tasks.

Abstract

Learning rich data representations from unlabeled data is a key challenge towards applying deep learning algorithms in downstream tasks. Several variants of variational autoencoders (VAEs) have been proposed to learn compact data representations by encoding high-dimensional data in a lower dimensional space. Two main classes of VAEs methods may be distinguished depending on the characteristics of the meta-priors that are enforced in the representation learning step. The first class of methods derives a continuous encoding by assuming a static prior distribution in the latent space. The second class of methods learns instead a discrete latent representation using vector quantization (VQ) along with a codebook. However, both classes of methods suffer from certain challenges, which may lead to suboptimal image reconstruction results. The first class suffers from posterior collapse, whereas the second class suffers from codebook collapse. To address these challenges, we introduce a new VAE variant, termed sparse coding-based VAE with learned ISTA (SC-VAE), which integrates sparse coding within variational autoencoder framework. The proposed method learns sparse data representations that consist of a linear combination of a small number of predetermined orthogonal atoms. The sparse coding problem is solved using a learnable version of the iterative shrinkage thresholding algorithm (ISTA). Experiments on two image datasets demonstrate that our model achieves improved image reconstruction results compared to state-of-the-art methods. Moreover, we demonstrate that the use of learned sparse code vectors allows us to perform downstream tasks like image generation and unsupervised image segmentation through clustering image patches.

Figures (23)

• Figure 1: An illustration of using sparse coding to model the latent representations of VAEs. The majority of VAEs can be categorized into two classes based on whether the latent representations are continuous, using a static prior, or discrete, utilizing vector quantization (VQ) with a codebook. Combining the VAE framework with sparse coding can be conceptualized as representing the middle ground between continuous and discrete VAEs.
• Figure 2: (a) The diagram of the ISTA algorithm for sparse coding. (b) The diagram of the Learnable ISTA, which is a time-unfolded version of the ISTA algorithm.
• Figure 3: A schematic representation of the proposed Sparse Coding-VAE with Learned ISTA (SC-VAE). SC-VAE integrates a Learnable ISTA network within VAE framework to learn sparse code vectors in the latent space for the input image. Each image can be represented as one or several sparse code vectors, depending on the number of downsampling blocks in the encoding process.
• Figure 4: Image reconstructions from different models trained on ImageNet dataset. Original images in the top two rows are from the validation set of ImageNet dataset. Two external images are shown in the last two rows to demonstrate the generalizability of different methods. The numbers denote the shape of latent codes and learned codebook (dictionary) size, respectively. SC-VAE achieved improved image reconstruction compared to the baselines. Zoom in to see the details of the red square area.
• Figure 5: Manipulating sparse code vectors on FFHQ. Each row represents a different seed image used to infer the latent sparse code vector in the SC-VAE$^\dag$ model. The disentangled attributes associated with the $i$-th component of a sparse code vector $z$ and a traversal range are shown in the first column.