Denoising Vision Transformer Autoencoder with Spectral Self-Regularization

TL;DR

This paper introduces a spectral self-regularization method for Denoising-VAE that reduces high-frequency noise in latent spaces, leading to faster convergence and improved image reconstruction and generation quality.

Contribution

It proposes a novel spectral self-regularization strategy for ViT-based autoencoders that enhances generative performance without relying on external foundation models.

Findings

Denoising-VAE produces cleaner, lower-noise latents.

Generative models converge approximately 2× faster.

Achieves state-of-the-art reconstruction quality on ImageNet.

Abstract

Variational autoencoders (VAEs) typically encode images into a compact latent space, reducing computational cost but introducing an optimization dilemma: a higher-dimensional latent space improves reconstruction fidelity but often hampers generative performance. Recent methods attempt to address this dilemma by regularizing high-dimensional latent spaces using external vision foundation models (VFMs). However, it remains unclear how high-dimensional VAE latents affect the optimization of generative models. To our knowledge, our analysis is the first to reveal that redundant high-frequency components in high-dimensional latent spaces hinder the training convergence of diffusion models and, consequently, degrade generation quality. To alleviate this problem, we propose a spectral self-regularization strategy to suppress redundant high-frequency noise while simultaneously preserving reconstruction quality. The resulting Denoising-VAE, a ViT-based autoencoder that does not rely on VFMs, produces cleaner, lower-noise latents, leading to improved generative quality and faster optimization convergence. We further introduce a spectral alignment strategy to facilitate the optimization of Denoising-VAE-based generative models. Our complete method enables diffusion models to converge approximately 2$\times$ faster than with SD-VAE, while achieving state-of-the-art reconstruction quality (rFID = 0.28, PSNR = 27.26) and competitive generation performance (gFID = 1.82) on the ImageNet 256$\times$256 benchmark.