WithinVAE: Compact Visual Autoencoding with Self-Distilled Semantics

Yuzhou Ji; Jiepeng Wang; Tianyu Jin; Lizhuang Ma; Haibin Huang; Chi Zhang; Xuelong Li

cs.CVOct 27, 2026

WithinVAE: Compact Visual Autoencoding with Self-Distilled Semantics

Authors: Yuzhou Ji, Jiepeng Wang, Tianyu Jin, Lizhuang Ma, Haibin Huang, Chi Zhang, Xuelong Li

Organizations: Shanghai Jiao Tong University · Institute of Artificial Intelligence, China Telecom

Abstract

Visual autoencoders must balance reconstruction fidelity, compactness, and semantic structure. Existing semantic tokenizers typically inherit semantics from pretrained recognition encoders, whose invariances can discard the appearance detail required for inversion. We show that this inheritance is unnecessary: semantic structure can be learned inside an autoencoder's own reconstructive representation. We introduce WithinVAE, a compact variational visual autoencoder with a sequential bottleneck, DINO/iBOT-style self-distillation from an exponential-moving-average copy of the autoencoder itself, and position-dependent token masking to achieve a 4-fold compression ratio, which is free of pretrained representation encoder or external semantic target. On ImageNet-1K at $256\times256$ , WithinVAE improves all four reported reconstruction metrics over RPiAE including PSNR by 1.36 dB, while using half its latent scalar capacity. Under frozen evaluation, the same latents improve linear probing by 6.5 points over RPiAE and 5.1 points over our reconstruction-only baseline, attributing the semantic gains to self-distillation rather than architecture or external supervision. An auxiliary latent-to-latent model reproduces the ordered-compression trend, and a pixel-to-latent model transfers the objective to a pretrained generative space. These results separate the origin of semantic supervision from its use: a compact visual autoencoder can organize semantic structure without borrowing a recognition representation.

Figures & tables

Figure 1: Two diagnostics of self-distilled semantic autoencoding. A: Under matched display settings, WithinVAE attends to object extent without any pretrained recognition encoder, while RPiAE totally fails in providing reasonable attention after borrowed DINOv2 (rollout and actual latent after variational-bridge). B: Prefix decoding of the $512\times 96$ latent-to-latent variant front-loads reliability: $K=128$ retains 25.6 dB PSNR, within 0.6 dB of the full sequence, and $K=64$ (6,144 scalars, under half of RPiAE’s budget) already exceeds RPiAE’s 21.3 dB.

Figure 2: WithinVAE training. The reconstruction path maps an image $x_{s}$ through a query encoder (ViT-L patches and 256 latent queries) and a token-wise VAE to a $256\times 32$ ordered token sequence; position-dependent masking keeps a reliable prefix $z_{1:k}$ and corrupts the masked tail before the ViT-L pixel decoder produces $\hat{x}_{s}$ , supervised by a reconstruction loss $\mathcal{L}_{\rm rec}+\mathcal{L}_{\rm perc}+\mathcal{L}_{\rm adv}$ . The semantic branches below match mean-pooled compact tokens via global DINO ( $\mathcal{L}_{\rm global}$ ) and student masked patches via local iBOT ( $\mathcal{L}_{\rm local}$ ) against targets produced by an EMA copy of the encoder and VAE that observes a clean view $x_{t}$ .

Model	Ext. RM	Latent	Scalars	rFID $\downarrow$	PSNR $\uparrow$	LPIPS $\downarrow$	SSIM $\uparrow$
SD-VAE ( Rombach et al., 2022 )	No	$32\!\times\!32\!\times\!4$	4,096	0.61	26.9	0.130	0.736
VA-VAE ( Yao et al., 2025 )	Yes	$16\!\times\!16\!\times\!32$	8,192	0.27	27.7	0.097	0.779
RAE-S ( Zheng et al., 2025 )	Yes	$256\!\times\!384$	98,304	0.64	18.9	0.252	0.489
RAE-B ( Zheng et al., 2025 )	Yes	$256\!\times\!768$	196,608	0.57	18.8	0.256	0.483
FAE-d32 ( Gao et al., 2025 )	Yes	$16\!\times\!16\!\times\!32$	8,192	0.68	–	–	–
RPiAE ( Gong et al., 2026 )	Yes	$256\!\times\!64$	16,384	0.50	21.3	0.216	0.525

Table 1: ImageNet-1K reconstruction at

256\times 256

. WithinVAE improves all four metrics over RPiAE at half its scalar count, without any external representation model. Baselines are reported values from their source papers (consolidated in RPiAE Table 3); WithinVAE uses our 50K-image evaluation. Ext. RM indicates an external pretrained representation model used for semantic initialization or supervision. Scalar count is an interface-capacity proxy, not a bitrate.

	Variant	PSNR $\uparrow$	LPIPS $\downarrow$	LP $\uparrow$	20-NN $\uparrow$	View $\uparrow$
①	RPiAE	21.30	0.216	65.404	66.856	$98.339\!\pm\!1.466$
②	WithinVAE reconstruction only	21.90	0.175	66.814	–	–
③	$+$ progressive mask	21.69	0.180	67.812	69.435	$98.029\!\pm\!1.441$
④	$+$ mask $+$ global	21.47	0.186	69.794	–	–
⑤	$+$ mask $+$ local	22.13	0.162	70.312	–	–
⑥	WithinVAE full	22.66	0.159	71.886	72.394	$\mathbf{99.125\!\pm\!0.732}$

Table 2: Reconstruction and frozen-latent semantics on ImageNet-1K. The full objective improves both fidelity and linear-probe accuracy over every partial variant. Circled indices ①–⑥ label rows for reference in the text. RPiAE reconstruction values are published; its semantic columns use our frozen-latent protocol. LP and 20-NN are top-1 accuracy (%). View is same-image cross-view cosine multiplied by 100. Dashes denote evaluations not run for intermediate ablations.

Model	Guidance	gFID $\downarrow$	IS $\uparrow$	Prec. $\uparrow$	Rec. $\uparrow$
RPiAE	Off	2.25	208.7	0.81	0.60
RPiAE	On	1.51	225.9	0.79	0.65
WithinVAE	Off	9.12	124.3	0.69	0.66
WithinVAE	On	3.67	169.2	0.76	0.58

Table 3: Class-conditional ImageNet generation with a 675M LightningDiT trained for 80 epochs. RPiAE values are reported by Gong et al. (2026) ; WithinVAE uses a 50K-sample evaluation. Guidance settings follow each implementation and are not asserted to be identical beyond the reported on/off condition.

Figure 3: Qualitative results. Up: Matched validation examples showing WithinVAE reconstruction, WithinVAE head-averaged attention, and RPiAE variational-bridge attention averaged over heads and bridge queries. The displays share the input crop, per-map normalization, and color map, but the attention operators are architecture-specific. Down: Class-conditional generation samples from the 80-epoch LightningDiT. Table 3 provides the distribution-level comparison.

Interface	Bottleneck	Scalars	rFID $\downarrow$	PSNR $\uparrow$	LPIPS $\downarrow$	SSIM $\uparrow$
Pixel $\rightarrow$ pixel	$256\!\times\!32$	8,192	0.46	22.66	0.159	0.593
Pixel $\rightarrow$ pixel, prefix	$128\!\times\!32$	4,096	0.47	22.33	0.163	0.581
Pixel $\rightarrow$ Qwen latent	$512\!\times\!96$	49,152	0.43	23.91	0.144	0.614
Qwen latent $\rightarrow$ Qwen latent	$512\!\times\!96$	49,152	0.40	26.27	0.131	0.714
prefix	$128\!\times\!96$	12,288	0.42	25.63	0.136	0.698
prefix	$64\!\times\!96$	6,144	0.48	22.23	0.160	0.579

Table 4: End-to-end reconstruction across interfaces at

256\times 256

. Qwen rows use a pretrained generative autoencoder; metrics are computed in RGB after decoding. Scalars expose the unequal model-facing capacities.

Variational Autoencoders

Variational autoencoders are probabilistic encoder–decoder models for learning latent representations, generating data, and estimating uncertainty.

See all 210 papers

TC-AE: Unlocking Token Capacity for Deep Compression Autoencoders
Apr 8, 2026Teng Li, Ziyuan Huang, Cong Chen +5Autoencoders Visual Token Compression
Inclusion AI, Ant Group · HKUST · ZJU +2
Distilling Latent Manifolds: Resolution Extrapolation by Variational Autoencoders
Mar 15, 2026Jiaming Chu, Tao Wang, Lei JinPosition Extrapolation Variational Autoencoders
School of Electronic Engineering, Beijing University of Posts and Telecommunications, Beijing, China
Entropic Auto-Encoding via Implicit Free-Energy Minimization
May 15, 2026Hazhir Aliahmadi, Irina Babayan, Greg van AndersVariational Autoencoders Posterior Collapse
Department of Physics, Engineering Physics & Astronomy, Queen’s University, Kingston, ON, Canada

Apr 8, 2026J/K move · Enter open · S save

TC-AE: Unlocking Token Capacity for Deep Compression Autoencoders

Teng Li, Ziyuan Huang, Cong Chen, Yangfu Li+4

Inclusion AI, Ant Group · HKUST · ZJU+2

Autoencoders Visual Token Compression Visual Tokenization Vision Transformer

Open paper AskPDF arXiv

We propose TC-AE, a ViT-based architecture for deep compression autoencoders. Existing methods commonly increase the channel number of latent representations to maintain reconstruction quality under high compression ratios. However, this strategy often leads to latent representation collapse, which degrades generative performance. Instead of relying on increasingly complex architectures or multi-stage training schemes, TC-AE addresses this challenge from the perspective of the token space, the key bridge between pixels and image latents, through two complementary innovations: Firstly, we study token number scaling by adjusting the patch size in ViT under a fixed latent budget, and identify aggressive token-to-latent compression as the key factor that limits effective scaling. To address this issue, we decompose token-to-latent compression into two stages, reducing structural information loss and enabling effective token number scaling for generation. Secondly, to further mitigate latent representation collapse, we enhance the semantic structure of image tokens via joint self-supervised training, leading to more generative-friendly latents. With these designs, TC-AE achieves substantially improved reconstruction and generative performance under deep compression. We hope our research will advance ViT-based tokenizer for visual generation.

Abstract

Variational Autoencoders

Explore similar work

TC-AE: Unlocking Token Capacity for Deep Compression Autoencoders

Distilling Latent Manifolds: Resolution Extrapolation by Variational Autoencoders

Entropic Auto-Encoding via Implicit Free-Energy Minimization