Source-linked AI summary
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
TL;DR
Deep compression autoencoders need compact latents without losing reconstruction or generative quality, yet naive token scaling can intensify token-to-latent compression and latent structure collapse. TC-AE addresses this through staged token compression and joint self-supervised training, achieving improved reconstruction and generative performance under deep compression.
Problem
Whether scaling the intermediate token space improves latent generatability in latent diffusion models remains an open problem.
Method
TC-AE redistributes token-to-latent compression across encoder stages and adds a joint self-supervised objective to structure image tokens.
Results
TC-AE significantly improves both reconstruction and generative performance under deep compression, while accelerating diffusion-model convergence and improving the efficiency–quality trade-off.
Takeaways & Limitations
Token-space optimization is an effective and complementary scaling direction for ViT-based tokenizers, alongside model-parameter scaling.
Abstract
from arXiv · showhide
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.
1 Introduction
TC-AE studies deep-compression ViT tokenizers through the intermediate token space, identifying token-to-latent compression as a bottleneck and combining staged compression with joint self-supervised training. The resulting tokenizer improves reconstruction and generation while making token-number scaling effective.
- Method: Joint self-supervised training improves the semantic structure of image tokens and produces more generative-friendly latents.This objective is integrated directly into tokenizer training rather than relying on external large-scale pretraining.
- Implication: Token-number scaling and model-parameter scaling are complementary dimensions that jointly improve generative performance.The result motivates token-space optimization as an effective direction for ViT-based tokenizers.
- Motivation: Token-to-latent compression is identified as a key factor limiting effective token-number scaling and contributing to representation collapse.Naively increasing image tokens improves reconstruction but can amplify compression at the fixed latent bottleneck, causing structural information loss and degraded generation.
- Method: TC-AE introduces staged token compression to alleviate structural information loss during token-to-latent mapping.The approach redistributes compression across encoder stages rather than applying it in a single aggressive bottleneck.
- Results: TC-AE achieves consistently better reconstruction and generative performance than existing tokenizers.The paper presents the architecture as balancing generative performance and efficiency under deep compression.
2 Related Work
Related work pursues deep compression and richer ViT token spaces, but increasing token granularity alone does not ensure better generation under a fixed latent budget. TC-AE addresses this reconstruction–generation conflict with staged token compression.
- Deep compression: Deep-compression tokenizers reduce latent spatial resolution to improve generation efficiency, but this can cause severe reconstruction degradation.The related-work discussion frames extremely high compression ratios as an active efficiency goal with associated quality challenges.
- Token scaling: Increasing the number of ViT image tokens can provide richer semantic representations and more comprehensive visual content.This benefit has been widely observed in visual representation learning, including self-supervised ViTs.
- Token scaling: Under a fixed latent resolution, increasing token numbers forces more aggressive token-to-latent compression and can cause latent structure collapse.Figure 2 identifies small patch sizes, such as p = 8, as an example of this failure mode.
- TC-AE: TC-AE proposes staged token compression to resolve the conflict between reconstruction and generation in naive token-scaling strategies.The strategy redistributes token-to-latent compression across stages to reduce structural information loss.
- Semantic structure: Representation-learning objectives have been integrated into generative frameworks, while tokenizer methods have used semantic alignment or related representation-structuring approaches.The cited related methods include REPA, iREPA, VA-VAE, MAETok, and GigaTok.
3 Method
TC-AE improves ViT-based deep-compression autoencoders by optimizing the token space through staged token compression and semantic structuring. Under a fixed latent budget, staged compression preserves more token semantics and allows generative quality to scale with token number.
- Preliminary: Compression in ViT-based Image Autoencoder: ViT autoencoders convert image patches into tokens, process them with Transformer layers, and compress the tokens into latent representations for generation.Image tokens bridge pixels and latents, so their capacity and semantic structure affect reconstruction and generative quality.
- Naive Scaling Token Numbers Fails to Improve Generation: Under a fixed latent budget, increasing token numbers improves reconstruction quality but does not improve generative performance.The controlled study fixes fpix→lat = 32 and c = 128 while varying patch size p ∈ {32, 16, 8}.
- Scaling Token Numbers with Staged Token Compression: Staged token compression redistributes compression across Transformer stages instead of applying one aggressive bottleneck compression step.A first stage learns rich structure with many tokens, an intermediate layer aggregates semantics, and a later bottleneck produces the final latent.
- Scaling Token Numbers with Staged Token Compression: 0.31 versus 0.08: staged token compression increases the semantic preservation ratio A2/A1 and enables generative quality to improve with more image tokens.The strategy mitigates structural information loss and achieves the strongest generative performance in the reported comparison.
- Naive Scaling Token Numbers Fails to Improve Generation: Token-to-latent compression is the primary bottleneck limiting token-number scaling because smaller patches increase bottleneck compression and destroy semantic structure.Although token semantics improve as patch size decreases, latent semantic preservation can collapse; at p = 8, the reported information loss is 1 − A2/A1 = 0.92.
- TC-AE: TC-AE combines staged token compression with joint self-supervised training to structure the token space and produce more generative-friendly latents.The architecture is based on a large, semantically structured token space and incorporates self-supervision in addition to staged compression.
4 Experiments
Experiments show that staged token compression and self-supervised token structuring improve reconstruction, generation, and diffusion efficiency under deep compression. Token scaling complements decoder scaling, while layer allocation balances reconstruction against generative performance.
- Diffusion convergence: 2.7× and 1.9× convergence acceleration is achieved in gFID and IS, respectively, with staged token compression over the baseline.Full TC-AE reaches the same gFID with 4.7× fewer iterations and comparable IS with 3.5× fewer steps.
- Staged token compression: Staged token compression reduces rFID from 1.33 to 0.75, gFID from 44.72 to 32.92, and increases IS from 33.62 to 43.58 without SSL.The ablation reports improvements in both reconstruction and generation quality.
- Self-supervision: SSL consistently improves generative performance across patch sizes, substantially reducing gFID and increasing IS, while generally degrading reconstruction quality.At smaller patch sizes, reconstruction becomes comparable to the baseline while generative performance remains significantly improved.
- Scaling effects: For comparable GFLOPs, token number scaling improves generation more strongly than parameter scaling, and combining both scaling dimensions yields further gains.The experiments identify token capacity and decoder capacity as complementary sources of generative improvement.
- System-level comparison: With 64 tokens, TC-AE reduces gFID to 7.16 without CFG and 2.57 with CFG, compared with 26.44 for DC-AE and 17.31 for DC-AE-1.5.TC-AE requires about 164 GFLOPs versus 607 GFLOPs for DC-AE.
- Layer-depth design: M = 6 provides the best generative performance, whereas reconstruction SSIM consistently improves as more layers process high-resolution tokens.The selected split reflects a trade-off between semantic structure learning and token-to-latent compression capacity.
- Self-supervised objective: iBOT achieves stronger generative performance than DINO with comparable reconstruction quality, while DINOv2 provides weaker gains in the moderate-batch setting.The comparison supports iBOT as a practical auxiliary objective for joint tokenizer training.
5 Conclusion
The conclusion attributes failed generation scaling to structural information loss at the compression bottleneck. TC-AE redistributes compression across encoder stages and structures tokens with self-supervision, improving quality and efficiency under deep compression.
- Naively increasing image-token count fails to improve generation because severe structural information loss occurs at the compression bottleneck.
- TC-AE redistributes token-to-latent compression across encoder stages and adds a self-supervised objective to structure the token space.
- ImageNet experiments show improved reconstruction and generation, faster diffusion convergence, and a favorable efficiency–quality trade-off under deep compression.
A TC-AE Implementation Details
The implementation details specify the training setup for self-supervised and adversarially enhanced TC-AE experiments. These details include augmentation, optimization stabilization, and discriminator design.
- TC-AE training details are reported in Table 7.
- Self-supervision training: Self-supervised training uses the corresponding methods’ augmentation pipelines and a reduced learning rate for stable joint optimization with reconstruction.
- Adversarial training: Adversarial training uses a discriminator built on a frozen DINO-S/8 backbone, with inputs resized to 224 × 224.
B More Ablations and Results
Additional ablations clarify how compression modules and objective weighting affect the reconstruction–generation trade-off. They also show that staged compression and SSL remain beneficial, while larger latent channels can worsen generative quality.
- Compression module: Two-layer convolution yields better gFID and IS than pixel-shuffle followed by an MLP, likely through local-neighborhood aggregation.The convolutional module is adopted as TC-AE’s default compression module.
- Objective weighting: Increasing α improves reconstruction at the expense of generative quality, reflecting a trade-off between high-frequency detail and semantic latent structure.The default setting is α = 0.1.
- Latent channel dimension: With c = 256, staged token compression and SSL still improve both metrics relative to their baselines, but generative quality is inferior to c = 128.The larger channel dimension has higher gFID and lower IS, suggesting aggravated representation collapse.
- Additional DC-AE comparison: TC-AE achieves better generative performance than DC-AE under matched latent shapes and training configurations.The comparison is reported in Table 11.
- Qualitative results: Additional generation results using TC-AE at 256 × 256 resolution are provided in Fig. 9.