Spherical VAE with Cluster-Aware Feasible Regions: Guaranteed Prevention of Posterior Collapse

TL;DR

This paper introduces a spherical VAE framework with cluster-aware constraints that guarantees prevention of posterior collapse, ensuring informative latent representations without restricting decoder capacity.

Contribution

It provides a novel theoretical approach leveraging spherical geometry and cluster constraints to strictly prevent posterior collapse in VAEs.

Findings

Achieves 100% collapse prevention on synthetic and real datasets.

Matches or exceeds state-of-the-art reconstruction quality.

Requires no explicit stability conditions or architecture restrictions.

Abstract

Variational autoencoders (VAEs) frequently suffer from posterior collapse, where the latent variables become uninformative as the approximate posterior degenerates to the prior. While recent work has characterized collapse as a phase transition determined by data covariance properties, existing approaches primarily aim to avoid rather than eliminate collapse. We introduce a novel framework that theoretically guarantees non-collapsed solutions by leveraging spherical shell geometry and cluster-aware constraints. Our method transforms data to a spherical shell, computes optimal cluster assignments via K-means, and defines a feasible region between the within-cluster variance $W$ and collapse loss $\delta_{\text{collapse}}$. We prove that when the reconstruction loss is constrained to this region, the collapsed solution is mathematically excluded from the feasible parameter space. \textbf{Critically, we introduce norm constraint mechanisms that ensure decoder outputs remain compatible with the spherical shell geometry without restricting representational capacity.} Unlike prior approaches, our method provides a strict theoretical guarantee with minimal computational overhead without imposing constraints on decoder outputs. Experiments on synthetic and real-world datasets demonstrate 100\% collapse prevention under conditions where conventional VAEs completely fail, with reconstruction quality matching or exceeding state-of-the-art methods. Our approach requires no explicit stability conditions (e.g., $\sigma^2 < \lambda_{\max}$) and works with arbitrary neural architectures. The code is available at https://github.com/tsegoochang/spherical-vae-with-Cluster.