TC-AE: Unlocking Token Capacity for Deep Compression Autoencoders

Teng Li; Ziyuan Huang; Cong Chen; Yangfu Li; Yuanhuiyi Lyu; Dandan Zheng; Chunhua Shen; Jun Zhang

cs.CVApr 8, 2026

TC-AE: Unlocking Token Capacity for Deep Compression Autoencoders

Authors: Teng Li, Ziyuan Huang, Cong Chen, Yangfu Li, Yuanhuiyi Lyu, Dandan Zheng, Chunhua Shen, Jun Zhang

Organizations: Inclusion AI, Ant Group · HKUST · ZJU · ECNU · HKUST (GZ)

Abstract

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.

Figures & tables

Figure 1: Architectural Comparasion of existing methods. (a) CNN-based tokenizers achieve strong generation with high FLOPs. (b) Vanilla ViT-based tokenizers are computationally efficient but suffer from severe information loss in the patchify layer, as spatial resolution is not gradually reduced as in CNNs. (c) TC-AE introduces staged token compression in the ViT encoder, balancing generative performance and efficiency.

Figure 2: Latent structure collapse caused by token-to-latent compression when scaling token numbers. Although increasing token numbers improves the token-level structure, the fixed latent resolution forces more aggressive token-to-latent compression, causing severe information loss and latent structure collapse at small patch sizes ( e.g. , $p=8$ ).

Figure 3: Naive scaling of token numbers for reconstruction and generation. We increase the number of image tokens by reducing the ViT patch size. While reconstruction quality consistently improves with more tokens (shown in the left two figures), generative performance does not exhibit corresponding gains (shown in the right two figures).

Figure 4: Token number scaling with staged token compression (STC). This strategy mitigates structural information loss at the bottleneck under large token spaces, enabling generative performance to scale with increasing token numbers.

Patch Size	gFID $\downarrow$	Linear Probing Accuracy
Patch Size	gFID $\downarrow$	Before Bottl. ( $A_{1}$ ) $\uparrow$	After Bottl. ( $A_{2}$ ) $\uparrow$	Str. Loss ( $1-\frac{A_{2}}{A_{1}}$ ) $\uparrow$
64	64.96	29.5	6.3	0.79
32	22.93	41.9	13.8	0.67
16	17.22	58.4	18.1	0.69
8	25.36	62.9	5.33	0.92
8 (w/ STC)	16.39	39.1	12.1	0.69

Table 1: Latent structure loss under token-to-latent compression. Increasing token numbers improves semantics of image tokens (

A_{1}

), but aggressive bottleneck compression causes severe semantic loss in the latent space. Bottl.: bottleneck. Str.: structure.

Figure 5: Architecture of TC-AE. TC-AE introduces staged token compression to alleviate structural information loss at the bottleneck, enabling generative performance to scale with the number of image tokens. In addition, self-supervised training is incorporated to further structure the token space.

Figure 6: Diffusion model convergence comparison. Both staged token compression strategy and joint self-supervised training accelerate the diffusion convergence.

Figure 7: Qualitative comparison between TC-AE and the baseline at $256\times 256$ resolution. Images are sampled from DiT checkpoints at different training steps using the same initial noise. TC-AE produces more visually appealing results compared to the baseline.

Staged token compression	rFID $\downarrow$	PSNR $\uparrow$	LPIPS $\uparrow$	SSIM $\uparrow$	gFID $\downarrow$	IS $\uparrow$
w/o SSL
✗	1.33	25.13	0.068	0.794	44.72	33.62
✓	0.75	25.72	0.072	0.807	32.92	43.58
w/ SSL
✗	0.90	24.84	0.085	0.772	25.36	52.38
✓	0.90	24.21	0.098	0.749	16.39	71.33

Table 2: Staged token compression improves generation performance.

ViT patch size	SSL	rFID $\downarrow$	PSNR $\uparrow$	LPIPS $\uparrow$	SSIM $\uparrow$	gFID $\downarrow$	IS $\uparrow$
64	✗	14.97	21.93	0.302	0.602	65.23	20.84
64	✓	30.96	20.25	0.383	0.501	64.96	24.89
32	✗	5.31	23.72	0.177	0.721	37.09	38.51
32	✓	4.63	22.15	0.195	0.632	22.93	57.22
16	✗	1.12	25.03	0.080	0.779	31.37	45.49
16	✓	1.27	23.76	0.104	0.723	17.22	69.00

Table 3: Self-supervision improves generative performance across token numbers.

Figure 8: Synergy of token number scaling and parameter scaling.

Tokenizer	Latent Tokens	rFID $\downarrow$	Generative Models	Param	gFID (w/o CFG) $\downarrow$	gFID (w/ CFG) $\downarrow$
low compression tokenizer
MaskGIT ( Chang et al., 2022 )	256	2.28	MaskGIT ( Chang et al., 2022 )	227M	6.18	-
VQGAN ( Esser et al., 2021 )	256	0.59	LlamaGen ( Sun et al., 2024 )	3B	9.38	2.18
VAR-Tok. ( Tian et al., 2024 )	680	1.00	VAR-d20 ( Tian et al., 2024 )	600M	-	2.95
LlamaGen-Tok. ( Sun et al., 2024 )	256	2.19	LlamaGen ( Sun et al., 2024 )	1.4B	-	3.09
VAE ( Rombach et al., 2022 )	256	0.53	MAR ( Li et al., 2024 )	943M	2.35	1.55

Table 4: System-level comparison of tokenizers and downstream generative models on ImageNet 256

\times

256. The CFG scale is set to 1.4 for TC-AE.

\dagger

: Results are taken from the original paper; additional comparisons are provided in the appendix.

Table 13

Appendix figures & tables6 assets

Supplementary material from the paper’s appendix.

Appendix

Component	Joing Training	Decoder Finetuning
Training data	ImageNet	ImageNet
Training epochs	50	16
Batch size	512	256
Learning Rate schedule	Cosine Decay	Cosine Decay
Optimizer betas	(0.9, 0.95)	(0.9, 0.95)
Weight decay	0.04	0

Appendix

Table 7: TC-AE training details.

Compression module	rFID $\downarrow$	PSNR $\uparrow$	LPIPS $\uparrow$	SSIM $\uparrow$	gFID $\downarrow$	IS $\uparrow$
Pixel shuffle + MLP	0.89	25.72	0.070	0.808	37.07	38.91
Two-layer convolution	0.75	25.72	0.072	0.807	32.92	43.58

Appendix

Table 8: Compression module design in ViT.

$\alpha$	rFID $\downarrow$	PSNR $\uparrow$	LPIPS $\uparrow$	SSIM $\uparrow$	gFID $\downarrow$	IS $\uparrow$
0.1	1.27	23.76	0.104	0.723	17.22	69.00
1.0	1.31	24.79	0.092	0.768	27.74	48.92
10	1.26	25.22	0.086	0.789	34.10	40.62

Appendix

Table 9: Trainging weights between reconstruction and self-supervision.

Channel	SSL	STC	rFID $\downarrow$	PSNR $\uparrow$	LPIPS $\uparrow$	SSIM $\uparrow$	gFID $\downarrow$	IS $\uparrow$
128			1.33	25.13	0.068	0.794	44.72	33.62
		✓	0.75	25.72	0.072	0.807	32.92	43.58
	✓		0.57	27.98	0.037	0.875	39.93	36.49
256			0.73	28.06	0.040	0.879	47.33	30.73
		✓	0.57	27.98	0.037	0.875	39.93	36.49
	✓		1.23	24.99	0.091	0.780	24.27	55.20

Appendix

Table 10: Expanding latent channel number to 256.

Tokenizer	Latent shape	GFLOPS	rFID $\downarrow$	gFID (w/o CFG) $\downarrow$	gFID (w/ CFG) $\downarrow$
DC-AE	$8\times 8\times 32$	607	0.71	15.90	5.59
TC-AE	$8\times 8\times 32$	164	0.81	10.72	4.77

Appendix

Table 11: Additional comparison with DC-AE. The CFG scale is set to 1.4.

Figure 9: Image generated by diffusion models using TC-AE at $256\times 256$ resolution.

Explore similar work

DC-SAE: Deep Compression Semantic Autoencoder for Faster Diffusion Convergence
Sep 30, 2026Xu Huang, Ye Huang, Zijun Liao +6Autoencoders Diffusion-Based Image Generation
Peking University · Singapore University of Technology and Design
ViTok-v2: Scaling Native Resolution Auto-Encoders to 5 Billion Parameters
May 6, 2026Philippe Hansen-Estruch, Jiahui Chen, Vivek Ramanujan +9Image Tokenization Autoencoders
University of Texas, Austin · University of Washington · 3Stanford University +3
Balancing Image Compression and Generation with Bootstrapped Tokenization
Jun 4, 2026Haozhe Chi, Jinghan Li, Hao Jiang +4Learned Image Compression Image Tokenization
1Peking University · 2Central Media Technology Institute, Huawei

Sep 30, 2026J/K move · Enter open · S save

DC-SAE: Deep Compression Semantic Autoencoder for Faster Diffusion Convergence

Xu Huang, Ye Huang, Zijun Liao, Yuwei Niu+5

Peking University · Singapore University of Technology and Design

Autoencoders Diffusion-Based Image Generation Diffusion Model Acceleration T2I Diffusion Models Visual Token Compression

Open paper AskPDF arXiv

High-compression tokenizers are essential for scaling latent image generative models. However, aggressive compression creates a fundamental tradeoff between reconstruction fidelity and generation efficiency: high compression image encoder always increases the learning difficulty of diffusion training, resulting in slow model convergence. Recent representation autoencoders speed up the diffusion training by improving the latent feature's expressive capability by replacing VAE encoders with pretrained semantic encoders, yet they are typically limited to moderate compression and lose pixel-level details necessary for faithful reconstruction. To achieve both high compression and fast diffusion training, we propose DC-SAE, a Decoupled Compact Semantic Autoencoder designed for high-compression image generation with accelerated diffusion model convergence. DC-SAE consists of two key components: (1) a macro-level architecture design that leverages semantic encoders to enable higher compression ratios, and (2) a pixel-level encoder that preserves low-level details, ensuring high-fidelity image reconstruction. We empirically demonstrate that DC-SAE performs strongly on image generation tasks, achieving both compact latent representations and efficient training dynamics. Specifically, on the ImageNet dataset with

512 \times 512

resolution, DC-SAE achieves

32\times

spatial compression, with 29.79 PSNR and 3.37 gFID, substantially outperforming the previous state-of-the-art high-compression tokenizer baselines DC-AE by 13.5% and 54.9% on PSNR and gFID, respectively, maintaining comparable throughput and faster diffusion model training convergence. Beyond class-conditional generation, a

1.6

B-parameter DiT using DC-SAE achieves 0.84 on GenEval and 86.007 on DPG-Bench for text-to-image generation at

1024\times1024

resolution.

Abstract

Explore similar work

DC-SAE: Deep Compression Semantic Autoencoder for Faster Diffusion Convergence

ViTok-v2: Scaling Native Resolution Auto-Encoders to 5 Billion Parameters

Balancing Image Compression and Generation with Bootstrapped Tokenization