Unifying CNNs and ViTs for Learning-Efficient and Scalable Variational AutoEncoder
Abstract
Continuous variational autoencoders (VAEs) are widely usedas visual tokenizers for latent diffusion models. While CNN-based VAEsreliably capture local details, they model long-range dependencies ineffi-ciently. Vision Transformers (ViTs) provide global context, but existingViT-VAEs exhibit three practical limitations in this setting: (i) weakresolution extrapolation when trained at a single resolution, (ii) slowconvergence with suboptimal fine-grained color and texture fidelity, and(iii) inconsistent scaling gains when increasing model size. To addressthese challenges, we propose TransVAE, a hybrid CNN-ViT VAE thatcombines a shallow CNN front-end for local feature extraction with adeep Transformer backbone for global modeling. We apply minimal yetprincipled architectural modifications and systematically analyze their ef-fectiveness, including a pure Rotary Position Embedding (RoPE) strategy,a multi-stage design, and a convolutional feed-forward network. Acrossmodel sizes from 44M to 2.3B parameters, TransVAE shows learning-efficient training and predictable scaling improvements, and it supportssingle-resolution training (256×256) with robust inference at higher resolu-tions (512 × 512/1024 × 1024). Moreover, TransVAE achieves competitivereconstruction and downstream generation performance compared tostrong tokenizers, such as FLUX-VAE and VA-VAE, offering a favorabletrade-off between reconstruction fidelity and generation-friendly latents.