WithinVAE: Compact Visual Autoencoding with Self-Distilled Semantics
Abstract Overview
This paper investigates whether semantic structure can be learned directly inside a visual autoencoder without relying on external pretrained recognition encoders. The authors introduce WithinVAE, a compact variational autoencoder featuring a sequential latent bottleneck, position-dependent token masking, and EMA-based global and local self-distillation inspired by DINO and iBOT. Evaluated on ImageNet-1K at 256×256, the method is tested for reconstruction fidelity, frozen semantic quality, prefix decoding, downstream latent diffusion generation, and transfer across alternative reconstruction interfaces. The findings show that self-distilled semantics improve both reconstruction quality and frozen feature representations while maintaining a compact latent interface.
Novelty
WithinVAE eliminates the need for an external pretrained representation model by generating semantic supervision online via an exponential moving average (EMA) copy of the autoencoder itself. Combined with position-dependent token masking, this self-distillation framework organizes semantic and reconstructive structure directly within a compact 1D latent bottleneck.
Results
On ImageNet-1K, WithinVAE improves all four reported reconstruction metrics over RPiAE using half the latent scalar capacity (8,192 vs. 16,384), achieving a 1.36 dB PSNR improvement. Under frozen evaluation, its latents boost linear probing by 6.482 points over RPiAE and by 5.072 points over the reconstruction-only baseline, while a 128-token prefix outperforms RPiAE in PSNR and SSIM at one-quarter of RPiAE's scalar capacity.
Key Points
- WithinVAE arranges latents into an ordered 1D variational sequence with position-dependent corruption, establishing a reliability gradient where early tokens carry the most critical reconstructive information for prefix decoding without retraining.
- Ablation experiments confirm that semantic performance stems directly from self-distillation rather than architecture alone, with the combination of global and local EMA-based distillation delivering improvements in both linear probing and reconstruction metrics.
- The self-distillation framework successfully transfers to auxiliary pixel-to-latent and latent-to-latent interfaces, and the representation supports downstream diffusion generation where it surpasses RPiAE in unguided recall despite trailing in overall generation metrics.