Sparse Compression VAE (SC-VAE)

• SC-VAE is a framework that embeds sparsity into variational autoencoders to overcome issues like posterior and codebook collapse while offering a structured latent space.
• It employs learned ISTA for efficient sparse coding in images and sparse-convolutional networks for point cloud attribute compression, optimizing rate-distortion tradeoffs.
• Empirical results show SC-VAE achieves superior reconstruction metrics and competitive compression performance compared to traditional and state-of-the-art methods.

Sparse Compression Variational Autoencoder (SC-VAE) refers to two distinct, state-of-the-art frameworks that integrate sparsity into the variational autoencoder paradigm: one for image modeling via sparse coding with learned ISTA, and one for point cloud attribute compression via sparse convolutions. Both models exploit the representational advantages of sparsity within the VAE framework to overcome limitations of conventional continuous and discrete latent variable approaches, achieving superior or competitive results in data compression, generative modeling, and unsupervised structuring of high-dimensional data (Xiao et al., 2023, Wang et al., 2022).

1. Motivation and Conceptual Foundations

Traditional VAEs leverage either (i) continuous, static Gaussian priors or (ii) discrete latent representations via vector quantization (VQ). Continuous-prior VAEs (e.g., vanilla VAE, β-VAE, InfoVAE) are susceptible to posterior collapse—strong decoders ignore the latent variables, driving the approximate posterior $q(z|x)$ towards the prior $p(z)$ and degrading the information-carrying capacity of $z$. Additionally, simple Gaussian priors poorly fit multi-modal or highly structured data, limiting reconstruction and generative capabilities.

Discrete (VQ-VAE family) sidestep posterior collapse via codebook lookup but suffer from codebook collapse (under-utilized embeddings), quantization artifacts, and the need for surrogate gradient estimators. SC-VAE for images introduces latent variables as sparse linear combinations of atoms from a fixed, orthonormal dictionary, minimizing both forms of collapse while supporting a smooth, interpretable latent space (Xiao et al., 2023).

In point cloud attribute compression, the classical impediment is the representation of large, irregular data with efficient encoding and rate-distortion tradeoffs. Here, SC-VAE exploits sparse-convolutional networks to encode color attributes, using adaptive entropy models and context-aware priors to minimize bitrate and error (Wang et al., 2022).

2. Mathematical Formulation

2.1 SC-VAE for Images (Sparse Coding-based VAE with Learned ISTA)

Let $x\in\mathbb{R}^{H\times W\times C}$ be an image, $z_{ij}\in\mathbb{R}^K$ a sparse latent code at spatial location $(i,j)$, and $D\in\mathbb{R}^{n\times K}$ a fixed orthonormal dictionary ($D^T D = I$).

Decoder and Generation:

• Each patch latent is decoded via $\tilde u_{ij}=D z_{ij}$.
• The reconstructed image is $\hat{x}=G(\{ \tilde u_{ij}\})$, with $G$ a deterministic deep decoder.

Probabilistic model:

• $$p_{\theta}(x|\{z_{ij}\}) = \mathcal{N}(x; G(Dz), \sigma^2 I)$$
• Prior: $p(z_{ij}) \propto \exp(-\alpha \|z_{ij}\|_1)$ (Laplace/L1 prior).

Inference via LISTA:

• $$z_{ij}^{*} = \operatorname*{argmin}_z \frac{1}{2}\|u_{ij} - D z\|_2^2 + \alpha \|z\|_1$$
• Solved by learned ISTA (LISTA): for layer $k$, $v^{(k+1)} = h_{\theta}(W_e u + S v^{(k)})$, where $$h_{\theta}(r)_i = \operatorname{sign}(r_i)\cdot\max(|r_i|-\theta_i,0)$$.
• $W_e = (1/L)D^T$, $S = I - (1/L)D^T D$, $\theta_i = \alpha / L$; in LISTA, these are learnable.

Loss (no KL—explicit sparsity):

• $\mathcal{L}_{rec} = \|G(D Z) - x\|_2^2$
• $$\mathcal{L}_{latent} = \sum_{i,j} \left[ \frac{1}{2} \|u_{ij} - D z_{ij}\|_2^2 + \alpha \|z_{ij}\|_1 \right]$$
• Total: $$\mathcal{L} = \mathcal{L}_{rec} + \frac{1}{h w} \mathcal{L}_{latent}$$ (Xiao et al., 2023).

2.2 SC-VAE for Point Cloud Attribute Compression

Let $N$ points $(x_i,y_i,z_i)$ with color $(R_i,G_i,B_i)$ as input, represented as a sparse tensor.

Encoder/decoder: 6 sparse-convolutional layers (stride $2$ at layers $2,4,6$), outputting quantized latents $\widehat{Y}$. Decoder mirrors the encoder.

Hyperprior and Context:

• Hyper encoder yields hyperlatents $\widehat{Z}$; hyperdecoder reconstructs side-information $\psi$.
• Autoregressive context interrogates masked causal neighborhoods, with context and hyper-decoded features producing locally adaptive mean/scale for latents.

VAE objective:

$$\mathcal{L}_{ELBO}(\theta,\phi; X) = \mathbb{E}_{q_{\phi}(z|X)}[-\log p_{\theta}(X|z)] + \mathrm{KL}[q_{\phi}(z|X)\|p(z)]$$

Rate-distortion implemented as:

$\mathcal{L}(X) = R + \lambda D$

with

$$R = \sum_i [-\log_2 p_{\widehat{y}_i|\widehat{z}}(\widehat{y}_i|\widehat{z},\widehat{y}_{<i})] + \sum_j [-\log_2 p_{\widehat{z}_j}(\widehat{z}_j)]$$

$D = \sum_{n=1}^{N} \|x_n - \widehat{x}_n\|_2^2$

Entropy model: Each latent $\widehat{y}_i$ modeled by a Laplacian, convolved with uniform: $$p_{\widehat{y}_i|\widehat{z},\widehat{y}_{<i}}(\widehat{y}_i) = \left( \mathcal{L}(x; \mu_i, \sigma_i) * \mathcal{U}(-1/2,1/2) \right)(\widehat{y}_i)$$ (Wang et al., 2022).

3. Model Architecture and Training

3.1 Image SC-VAE

• Encoder: Deep ResNet blocks (VQGAN architecture), downsampling to latent map $u\in\mathbb{R}^{h\times w \times n}$.
• LISTA module: $s=16$ unfolded ISTA layers, producing $z_{ij}\in\mathbb{R}^K$ for each patch. Parameters of unfolding are trained end-to-end.
• Decoder: Symmetric to encoder with upsampling.
• Dictionary $D$: Fixed orthonormal DCT, typically $n=256, K=512$.
• Optimization: Adam, $10$ epochs, batch size $16$, learning rate $1e^{-4}$, $\alpha$ initialized to $1$ and learned.

Table: Key Model Hyperparameters for Image SC-VAE

Component	Choice/Size	Notes
Encoder	VQGAN + ResNet	Downsampling
LISTA depth	$s=16$	Best at $s=5$–$15$
Dictionary	DCT, $n\times K$	e.g., $256\times 512$
$\alpha$	Learnable, init $1$	Promotes sparsity

3.2 Point Cloud Attribute SC-VAE

• Encoder/Decoder: 6 SConv layers, $3^3$ kernels (MinkowskiEngine). Output feature width $128$.
• Hyperpath: Additional down/up SConvs for hyperlatents, context modeling via masked $5^3$ SConv.
• Optimization: Learning rate $10^{-4} \rightarrow 2 \times 10^{-5}$, $50$ epochs, Adam optimizer typical.
• Batch size: Not specified; small batches common due to data size.

4. Empirical Results and Comparative Evaluation

4.1 Image SC-VAE

Experiments on FFHQ and Imagenet, at $256 \times 256$ resolution, compare SC-VAE to VQGAN, Mo-VQGAN, RQ-VAE, and continuous/sparse-prior VAEs across metrics (PSNR, SSIM, LPIPS, recon-FID). Notably, at $16\times 16$ latent grid, SC-VAE achieves (FFHQ):

Metric	VQGAN	RQ-VAE	Mo-VQGAN	SC-VAE
PSNR	22.24	24.53	26.72	29.70
SSIM	0.6641	0.7602	0.8212	0.8347
LPIPS	0.1175	0.0895	0.0585	0.1956

At $32\times 32$, SC-VAE yields PSNR=$34.92$, SSIM=$0.9497$, LPIPS=$0.0080$, rFID=$4.21$, outperforming all baselines. Analogous gains are observed on ImageNet.

LISTA ablation: $s=5$–$15$ yields highest PSNR ($\approx31.4$) and best sparsity ($72$–$75\%$). Lower or higher $s$ degrades performance.

4.2 Point Cloud Attribute SC-VAE

Evaluation on 8i Full Bodies (“longdress,” etc.), against TMC13 v6, v14, RAHT, and prior learned methods.

• Bjøntegaard gains vs TMC13 v6: BD-BR reduction $24\%$, BD-PSNR $+0.97$ dB.
• vs RAHT: $34\%$ reduction, $+1.38$ dB PSNR.
• vs TMC13 v14: SC-VAE lags by $+55\%$ BD-BR ($-1.25$ dB PSNR).
• Qualitative: Smoother color, reduced artifacts; rivals TMC13 v14 visually.
• Entropy model ablation: Joint hyperprior+AR gives $38.9\%$ BD-BR saving over factorized baseline.

5. Downstream Applications

5.1 Image Generation and Disentanglement

Image SC-VAE supports interpretable latent traversals and interpolations:

• Latent traversal: Vary $z_k$ within $[-3\sigma, +3\sigma]$ can effect semantic edits (smile, pose, lighting).
• Interpolation: Linear interpolation between two latent codes generates smooth image morphs.

5.2 Patch Clustering and Unsupervised Segmentation

• Each $z_{ij}$ serves as a patch descriptor; $K$-means clusters ($K=1000$) segregate regions of similar texture and semantics (e.g., sky, foliage).
• Unsupervised segmentation by clustering all $z_{ij}$ into classes ($c\in\{2,3,5\}$), upsampling to produce segmentation masks. Achieves IoU=$81.2\%$, DICE=$88.5\%$ on Flowers dataset, and competitive results on additional datasets without fine-tuning.
• SC-VAE-based segmenters are more robust to Gaussian noise than GAN-driven approaches.

6. Practical Considerations and Limitations

• Orthogonal dictionary (e.g., DCT): Facilitates disentanglement, avoids scale ambiguity.
• Loss balancing by $1/(h w)$: Ensures latent penalty does not overshadow image reconstruction.
• Optimal LISTA rollout ($s\approx10$–$20$): Trade-off between sparsity ($\approx70$–$80\%$ by Hoyer metric) and reconstruction fidelity.
• Latent Map Size: Finer grids ($32\times 32$) yield best fidelity but incur higher computational cost; coarser ($1\times1$) degrade results.
• Current limits: In point cloud compression, SC-VAE is on par with (or slightly lags) the latest TMC13 v14 in rate-distortion, but outperforms earlier learning-based and standardized systems. Enhancements such as cross-scale prediction are anticipated to close this gap in future work.

7. Significance and Future Directions

SC-VAE frameworks establish that sparse priors and sparse coding principles, when combined with deep architectures and modern learning paradigms (e.g., LISTA, sparse convolutions), yield latent representations that are compact, highly informative, and conducive to both high-fidelity reconstruction and structured downstream analysis. This hybridization addresses foundational limitations of both continuous and discrete VAE approaches—specifically, posterior and codebook collapse—while supporting broader generative and unsupervised learning. A plausible implication is that further integration with cross-scale prediction, transform coefficient prediction, and advanced context models will enable even higher compression ratios and richer semantic decompositions for both images and irregular 3D data (Xiao et al., 2023, Wang et al., 2022).