Unifying CNNs and ViTs for Learning-Efficient and Scalable Variational AutoEncoder

Conference: ECCV 2026
Paper: ECCV 2026
Area: Image Generation
Keywords: Variational AutoEncoder, Visual Tokenizer, Vision Transformer, Resolution Extrapolation, Hybrid Architecture

TL;DR

TransVAE unifies a shallow CNN front-end for local feature extraction with a deep Transformer backbone for global context modeling, leveraging a pure RoPE strategy, multi-stage overlapping convolutions, and a convolutional feed-forward network to achieve 4× faster convergence, robust zero-shot resolution extrapolation, and predictable scaling from 44M to 2.3B parameters.

Background & Motivation

Continuous-value variational autoencoders (VAEs) serve as the cornerstone visual tokenizers for modern latent diffusion models (LDMs), tasked with compressing high-resolution pixel space into compact, information-rich latent representations. For years, visual tokenizers have been overwhelmingly dominated by convolutional neural networks (CNNs), which harness local connectivity and translation equivariance priors to learn fine-grained textures and pixel-level color fidelity with high sample efficiency. Nevertheless, the strictly localized receptive fields of convolutions render CNNs intrinsically inefficient at capturing long-range dependencies, and stacking ever-deeper convolutional layers leads to severe optimization bottlenecks.

In contrast, Vision Transformers (ViTs) provide native all-to-all global context modeling via self-attention mechanisms. However, when deployed as continuous visual tokenizers, pure ViT-VAEs suffer from three critical shortcomings: first, fragile resolution extrapolation, where models trained at a single fixed resolution fail during higher-resolution inference because interpolating absolute position embeddings (APE) disrupts learned spatial priors and induces grid artifacts; second, sluggish convergence and poor local fidelity, where the lack of local inductive biases leads to slow optimization of high-frequency colors and textures; and third, inconsistent parameter scaling gains, where increasing the capacity of pure ViT-VAEs fails to yield reliable reconstruction improvements.

Recognizing that pure CNNs lack global field-of-view while pure ViTs lack local inductive bias, this paper seeks an architecturally principled synthesis of their complementary strengths. Core idea: construct a hybrid multi-stage architecture combining a shallow CNN front-end for local feature extraction with a deep Transformer backbone for global modeling, augmented with pure RoPE positional encoding, sequential overlapping convolutions, and a convolutional feed-forward network (ConvFFN) to achieve rapid convergence, artifact-free resolution extrapolation, and predictable scaling up to 2.3B parameters.

Method

Overall Architecture

TransVAE adheres to an architecturally symmetric encoder-decoder topology designed to produce complete visual representations that balance fine-grained local fidelity with expressive global semantics. Given an input image \(x \in \mathbb{R}^{H \times W \times 3}\), the encoder initially processes features through a shallow CNN front-end composed of ResBlocks and sequential downsampling convolutions to capture fine-grained pixel details, subsequently routes intermediate features through multi-stage deep TransVAE blocks for long-range global context modeling, and maps the latent code via Gaussian modeling with KL regularization into \(z \in \mathbb{R}^{\frac{H}{f} \times \frac{W}{f} \times d}\). The decoder mirrors this layout in exact reverse: deep TransVAE blocks invert the long-range latent dependencies before symmetric CNN upsampling stages restore fine pixel structures into the reconstructed image \(\hat{x}\).

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["Input Image x<br/>H × W × 3"] --> B["Shallow CNN Front-end<br/>Sequential Overlapping Convolutions & ResBlocks"]
    B --> C["Deep TransVAE Backbone<br/>Multi-Stage Attention for Global Semantics"]
    C --> D["Latent Projection & Sampling<br/>Gaussian Modeling & KL Regularization to z"]
    D --> E["Deep Inverse TransVAE Backbone<br/>Multi-Stage Long-Range Context Inversion"]
    E --> F["Shallow CNN Decoding Back-end<br/>Symmetric ResBlock Upsampling & Detail Recovery"]
    F --> G["Reconstruction Output x_hat<br/>High-Fidelity Reconstructed Image"]

Key Designs

1. Pure RoPE Positional Encoding: Eliminating Spatial Prior Distortion and Grid Artifacts in Resolution Extrapolation

Standard ViT tokenizers rely on additive learnable absolute positional embeddings (APE) \(X = X + P_{APE}\), which require bilinear or bicubic interpolation whenever the test sequence length deviates from the training resolution (e.g., \(256 \times 256 \to 512 \times 512 / 1024 \times 1024\)), severely distorting the learned spatial topology and triggering visible grid artifacts. TransVAE discards APE entirely in favor of a 2D Rotary Position Embedding (RoPE) applied directly within the multi-head self-attention module, binding positional awareness to query \(q_m\) and key \(k_n\) representations through orthogonal rotation operators:

\[({q'}_m)^\top {k'}_n = g(q_m, m)^\top g(k_n, n) = h(q_m, k_n, m - n)\]

Because the attention score is conditioned exclusively on the relative spatial displacement \(m - n\) rather than absolute coordinate indices, the formulation naturally generalizes to arbitrary sequence lengths without parameter interpolation, enabling flawless zero-shot resolution extrapolation while accelerating early convergence on pixel colors.

2. Hierarchical Multi-Stage Hybrid Architecture: Structuring Complementary Local and Global Representations

Conventional ViT-VAEs employ single-stage, non-overlapping patch embeddings (such as a direct \(16 \times 16\) convolution) paired with pixel-shuffle decoders, treating neighboring patches as isolated islands and neglecting local continuity. TransVAE adopts a macro-level multi-stage hybrid design: the initial two stages of the encoder are constructed using standard convolutional Residual Blocks (ResBlocks) interconnected by sequential \(3 \times 3\) stride-2 convolutions (SeqConv), ensuring overlapping receptive fields across downsampling steps; subsequent deep stages transition to Transformer blocks to model long-range context across high-level feature maps, mirrored symmetrically by the decoder. This hierarchical allocation lets convolutions handle low-level textures, relieving the Transformer from expending capacity on local translation equivariance.

3. Convolutional Feed-Forward Network (ConvFFN): Restoring Local Inductive Bias in Deep Transformer Layers

Standard Transformer FFN blocks operate strictly point-wise across individual tokens, severing local spatial relationships at deep computational stages. TransVAE introduces a convolutional bypass with residual connections directly into the FFN architecture: following the initial linear expansion, an intermediate \(3 \times 3\) full convolution (FullConv) performs spatial feature mixing across neighboring tokens, which is then added to the linear projection before activation and down-projection. Compared to depth-wise convolutions (DWConv), FullConv delivers significantly richer local spatial modeling under comparable parameter budgets, allowing the deep network to maintain local coherence without compromising global attention.

4. Stable Scaling Normalization and Information-Preserving Shortcuts: Enabling Predictable Scaling to 2.3B

To ensure robust optimization when scaling parameters from 44M to the 2.3B Giant configuration, TransVAE integrates RMSNorm immediately prior to the linear \(Q, K, V\) projections in each self-attention block, preventing gradient explosions and numerical divergence (NaN loss) during long training schedules. Additionally, across high spatial compression ratios (\(f16\) and \(f32\)), all downsampling and upsampling transitions incorporate parallel information-preserving shortcut connections, ensuring that high-frequency geometric cues propagate unhindered between shallow and deep stages.

Loss & Training

During reconstruction training, TransVAE optimizes a compound loss function combining pixel-level \(L_1\) reconstruction loss, perceptual loss \(L_{\text{LPIPS}}\), and Kullback-Leibler divergence \(L_{\text{KL}}\) to enforce Gaussian priors on the latent space:

\[\mathcal{L}_{\text{total}} = \mathcal{L}_1(x, \hat{x}) + \lambda_{\text{lpips}} \mathcal{L}_{\text{LPIPS}}(x, \hat{x}) + \lambda_{\text{kl}} \mathcal{L}_{\text{KL}}(q_\phi(z|x) \parallel \mathcal{N}(0, I))\]

For downstream generation-oriented variants (denoted with the -VF suffix), a Vector Field alignment loss (VF loss) is incorporated to align latent representations with semantic feature spaces. All models are trained with BFloat16 mixed precision using AdamW on \(256 \times 256\) resolution with a batch size of 32.

Key Experimental Results

Main Results

Evaluation of zero-shot resolution extrapolation on ImageNet-1k validation set (all models trained exclusively at \(256 \times 256\) resolution):

Tokenizer	Variant	# Params	Eval Resolution	rFID ↓	PSNR ↑	LPIPS ↓	SSIM ↑
DeTok-BB-FT	f16d16	171M	256 × 256	0.54	25.24	0.137	0.71
ViTok S-B/16	f16d16	129M	256 × 256	0.50	24.36	-	0.75
VA-VAE	f16d32	70M	256 × 256	0.27	28.57	0.090	0.80
TransVAE-T (Ours)	f16d32	44M	256 × 256	0.38	29.08	0.083	0.83
TransVAE-L (Ours)	f16d32	545M	256 × 256	0.31	29.78	0.075	0.84
DeTok-BB-FT	f16d16	171M	512 × 512	3.43	25.22	0.262	0.70
VA-VAE	f16d32	70M	512 × 512	0.12	31.05	0.095	0.84
TransVAE-T (Ours)	f16d32	44M	512 × 512	0.17	31.10	0.090	0.85
TransVAE-L (Ours)	f16d32	545M	512 × 512	0.13	32.44	0.077	0.88
DeTok-BB-FT	f16d16	171M	1024 × 1024	8.86	25.06	0.371	0.73
VA-VAE	f16d32	70M	1024 × 1024	0.06	36.15	0.071	0.93
TransVAE-T (Ours)	f16d32	44M	1024 × 1024	0.08	36.38	0.063	0.93
TransVAE-L (Ours)	f16d32	545M	1024 × 1024	0.03	38.25	0.051	0.95

Ablation Study

Ablation of key architectural components on early convergence (5 epochs, averaged across 3 random seeds) and resolution extrapolation behavior:

Component	Variant / Setup	5-Epoch Val Loss ↓	512 Extrapolation Artifacts	Note
Position Embedding (PE)	w/o PE	~0.530	Severe color shift	Lacks coordinate cues; fails to converge on fine details
Position Embedding (PE)	APE (Absolute PE)	~0.490	Severe grid artifacts	Parameter interpolation disrupts learned spatial priors
Position Embedding (PE)	RoPE-only (Ours)	~0.370	Clean & sharp	Relative distance invariance supports robust extrapolation
Downsampling & Staging	Single-stage ViT + PixelShuffle	~0.480	Blurry edges	Lacks localized inductive bias
Downsampling & Staging	Dual SeqConv (Enc & Dec)	~0.340	Smoother boundaries	Overlapping receptive fields accelerate local feature learning
Downsampling & Staging	CNN Stem + Multi-Stage (Ours)	~0.250	High-fidelity textures	Initial ResBlocks match CNN convergence efficiency
Feed-Forward Network (FFN)	Standard MLP-FFN (37M)	~0.245	Slightly smoothed textures	Point-wise operations sever spatial continuity
Feed-Forward Network (FFN)	DWConv-FFN (37M)	~0.210	Well-preserved textures	Depth-wise convolution reintroduces local inductive bias
Feed-Forward Network (FFN)	FullConv-FFN (42M, Ours)	~0.185	Best detail preservation	Spatial feature mixing surpasses pure CNN baseline

Key Findings

• 4× Acceleration in Learning Efficiency: TransVAE-T requires only 25 epochs to reach 28.61 dB PSNR, surpassing the 100-epoch performance of VA-VAE (28.57 dB) and demonstrating substantial training compute savings.
• Predictable and Consistent Scaling Gains: Scaling TransVAE across 44M (Tiny), 140M (Base), 545M (Large), 1.3B (Huge), to 2.3B (Giant) yields strictly monotonic improvements in validation loss, reconstruction PSNR (from 27.2 dB to 28.7 dB at 3 epochs), and linear probing accuracy (from 10.91% to 65.18% at 50 epochs), breaking the scaling stagnation seen in pure ViT-VAEs.
• Fidelity vs. Downstream Friendliness Trade-off: Applying VF representation alignment introduces a modest reconstruction trade-off (~0.5–1.0 dB PSNR decrease) but markedly boosts downstream DiT generative quality, lowering 80-epoch generation FID with LightningDiT-XL from 4.51 / 1.96 (VA-VAE) to 3.58 / 1.85.

Highlights & Insights

• Minimalist Yet Principled Hybridization: Rather than designing convoluted dual-stream cross-attention architectures, TransVAE systematically targets the empirical root causes of ViT-VAE failures by placing convolutions in the outer stages and Transformer blocks in the deep core.
• Relative Distance is Essential for Spatial Extrapolation: The paper definitively proves that absolute coordinate parameterization is the root culprit behind extrapolation failure in visual tokenizers, whereas pure relative RoPE unlocks zero-shot extrapolation across arbitrary image resolutions.
• Broad Transferability to Generative Representations: The strategy of inserting spatial convolutional bypasses within deep Transformer blocks offers an effective architectural template for visual tokenizers across autoregressive generation, flow matching, and video latent spaces.

Limitations & Future Work

• Unoptimized Kernel-Level Inference Throughput: The initial implementation focuses on architectural validation rather than fused hardware kernels; while Flash Attention reduces peak memory usage, overall throughput lags behind highly-engineered pure CNNs (e.g., 44M TransVAE-T achieves 193 img/s encode throughput vs. 371 img/s for 70M VA-VAE).
• Extrapolation Vulnerability in Single-Resolution Alignment: Single-resolution VF semantic alignment tends to overfit semantic features to the training scale, causing semantic degradation during high-resolution extrapolation. Exploring multi-scale or invariant representation alignment objectives remains a promising future direction.

Related Work & Insights

• vs VA-VAE: VA-VAE relies primarily on conventional CNN-with-attention configurations that encounter capacity saturation when scaled; TransVAE scales smoothly to 2.3B parameters and achieves superior reconstruction bounds and extrapolation fidelity.
• vs DeTok / ViTok: DeTok collapses under resolution extrapolation due to APE; ViTok adopts RoPE but lacks sufficient local inductive priors, yielding suboptimal reconstruction quality; TransVAE reconciles RoPE extrapolation with CNN-grade local fidelity.
• vs DC-AE: DC-AE relies heavily on convolutions in earlier stages and uses only sparse Transformer blocks near the bottleneck; TransVAE is an authentic deep Transformer tokenizer that yields a more semantically organized latent space for downstream generative modeling.

Rating

• Novelty: ⭐⭐⭐⭐ [Precise diagnosis of ViT-VAE failure modes coupled with an elegant, highly effective hybrid architecture]
• Experimental Thoroughness: ⭐⭐⭐⭐⭐ [Exhaustive scaling from 44M to 2.3B parameters, covering reconstruction, extrapolation, linear probing, and downstream DiT synthesis]
• Writing Quality: ⭐⭐⭐⭐⭐ [Clear motivation, methodical progression from macro to micro design, and insightful empirical analysis]
• Value: ⭐⭐⭐⭐⭐ [Provides a practical, highly scalable visual tokenizer blueprint for the generative modeling community]

Related Papers

• [ICLR 2026] Pareto Variational Autoencoder
• [ECCV 2026] Learning to Stylize by Learning to Destylize: A Scalable Paradigm for Supervised Style Transfer
• [CVPR 2026] SRA 2: Variational Autoencoder Self-Representation Alignment for Efficient Diffusion Training
• [ICLR 2026] Latent Diffusion Model without Variational Autoencoder
• [ECCV 2026] FlowInOne: Unifying Multimodal Generation as Image-in, Image-out Flow Matching