TransVAE: Combining CNN Local Bias with Transformer Global Modeling for Scalable Visual Tokenizers

1. Limitations of CNN‑VAE and ViT‑VAE as visual tokenizers

CNN‑VAE recovers local detail reliably because convolutional kernels encode strong priors for color, texture and edges, but its limited receptive field makes modeling long‑range relationships inefficient, requiring deeper stacks.

ViT‑VAE captures global context via self‑attention, yet learns pixel‑level color and texture slowly, can produce high‑frequency artifacts, and its absolute positional embeddings must be interpolated for higher‑resolution inference, degrading spatial priors. Scaling the model size does not guarantee stable reconstruction gains.

2. TransVAE architecture

2.1 RoPE‑only positional encoding for resolution‑agnostic scaling

Instead of absolute position embeddings, TransVAE inserts a 2‑D Rotary Position Embedding (RoPE) inside the attention mechanism. RoPE rotates the query and key vectors to encode relative spatial relations, so a model trained at a single resolution can be applied to higher resolutions without interpolation.

2.2 Multi‑stage CNN‑Transformer backbone

The single ViT trunk is split into several stages. The first two stages use ResBlocks for down‑sampling or up‑sampling, providing efficient local feature extraction. Deeper stages replace the ResBlock with a TransVAE block that models global relationships.

This division lets shallow layers preserve color, texture and edges, while deeper layers capture shape, object and scene structure. Early training shows TransVAE recovers fine details faster than a pure ViT and forms coherent structures earlier than a pure CNN.

2.3 ConvFFN and training stabilisation

Standard Transformer feed‑forward networks (FFN) treat each token independently, breaking spatial connections. TransVAE adds a convolutional branch with a residual connection after the input projection, forming a ConvFFN that mixes local spatial information in deep layers.

Additional stabilisation includes RMSNorm before the Q/K/V projections and parallel shortcuts in high‑compression down‑sampling/up‑sampling paths to retain fine details.

2.4 Model configurations

TransVAE is instantiated in a family ranging from 44 M to 2.3 B parameters, covering various compression ratios to demonstrate effectiveness across deployment scenarios.

3. Experimental results

3.1 Faster high‑fidelity reconstruction

On ImageNet‑1k with a 16×32 token grid, the 44 M‑parameter TransVAE‑T reaches 28.61 PSNR after 25 epochs, surpassing the 70 M‑parameter VA‑VAE’s 28.57 PSNR after 100 epochs. After 100 epochs, TransVAE‑T attains 29.08 PSNR.

At a higher compression of 8×16, the 719 M‑parameter TransVAE‑L achieves 33.30 PSNR on the ImageNet‑1k validation set and 33.27 PSNR on MS‑COCO 2017.

3.2 Robust resolution extrapolation

All 16×32 models are trained only at that resolution. At inference, the 545 M‑parameter TransVAE‑L reaches 38.25 PSNR on ImageNet‑1k and 36.96 PSNR on MS‑COCO 2017 when evaluated at 1024 × 1024 resolution, demonstrating stable up‑scaling.

3.3 Scaling trends from 44 M to 2.3 B

Comparisons among CNN‑VAE, ViT‑VAE, Swin‑VAE and TransVAE across parameter scales show that pure ViT and Swin gains are unstable, while CNN‑VAE’s reconstruction plateaus. TransVAE consistently improves both reconstruction quality and latent alignment as size grows.

Early‑stage (3‑epoch) training shows monotonic improvement from Tiny to Giant variants; after 50 epochs, the 2.3 B TransVAE‑G maintains this trend. Linear probing on aligned latents yields 65.18 % top‑1 accuracy for TransVAE‑VF‑G, compared with 42.05 % for VA‑VAE‑VF and 44.13 % for TransVAE‑VF‑T.

4. Downstream generation and feature visualisation

Using LightningDiT as the downstream generator, the 16×32 TransVAE‑VF‑L (LightningDiT‑XL/1) after 80 epochs achieves an FID‑50k of 3.5, lower than the same‑budget VA‑VAE‑VF (4.51). At 8×16, after 100 k steps, TransVAE‑VF‑L’s FID approaches that of FLUX‑VAE after 200 k steps.

Zero‑shot segmentation analysis freezes the encoder, extracts stage‑2/3/4 features on ADE20K, upsamples them to a common resolution and applies K‑Means clustering without training a segmentation head. Deeper stages produce more coherent object‑level clusters and clearer boundaries, illustrating qualitative latent quality.

5. Conclusion

TransVAE integrates CNN’s local inductive bias with Transformer’s global modelling within a single visual tokenizer. RoPE enables resolution‑agnostic inference, ConvFFN restores local spatial mixing in deep Transformer layers, and RMSNorm plus parallel shortcuts stabilise large‑scale training. Experiments from 44 M to 2.3 B parameters confirm consistent gains in reconstruction fidelity, representation alignment and downstream generation, demonstrating that VAE representation capacity can continue to scale with appropriate architecture.

Code example

[2] Junyu Chen et al. Deep Compression Autoencoder for Efficient High-Resolution Diffusion Models. ICLR 2025.
[3] Philippe Hansen-Estruch et al. Learnings from Scaling Visual Tokenizers for Reconstruction and Generation. ICML 2025.
[4] Jiawei Yang et al. Latent Denoising Makes Good Visual Tokenizers. ICLR 2026.