README.md

metadata

license: mit
pipeline_tag: feature-extraction
tags:
  - visual-tokenizer
  - image-reconstruction

TC-AE: Unlocking Token Capacity for Deep Compression Autoencoders

 

Teng Li^1,2*, Ziyuan Huang^1,*,✉, Cong Chen^1,3,*, Yangfu Li^1,4, Yuanhuiyi Lyu^1,5,
Dandan Zheng¹, Chunhua Shen³, Jun Zhang^2✉
1Inclusion AI, Ant Group, ²HKUST, ³ZJU, ⁴ECNU, ⁵HKUST (GZ)
^*Equal contribution, ✉ Corresponding authors

Introduction

TC-AE is a novel Vision Transformer (ViT)-based tokenizer for deep image compression and visual generation. Traditional deep compression methods typically increase channel dimensions to maintain reconstruction quality at high compression ratios, but this often leads to representation collapse that degrades generative performance. TC-AE addresses this fundamental challenge from a new perspective: optimizing the token space — the critical bridge between pixels and latent representations. By scaling token numbers and enhancing their semantic structure, TC-AE achieves superior reconstruction and generation quality. Key Innovations:

• Token Space Optimization: First to address representation collapse through token sapce optimization
• Staged Token Compression: Decomposes token-to-latent mapping into two stages, reducing structural information loss in the bottleneck
• Semantic Enhancement: Incorporates self-supervised learning to produce more generative-friendly latents

Usage

Environment Setup

To set up the environment for TC-AE, follow these steps:

conda create -n tcae python=3.9
conda activate tcae
pip install -r requirements.txt

Image Reconstruction Demo

python tcae/script/demo_recon.py \
    --img_folder /path/to/your/images \
    --output_folder /path/to/output \
    --ckpt_path results/tcae.pt \
    --config configs/TC-AE-SL.yaml \
    --rank 0

ImageNet Reconstruction Evaluation

Evaluate reconstruction quality on ImageNet validation set:

python tcae/script/eval_recon.py \
    --ckpt_path results/tcae.pt \
    --dataset_root /path/to/imagenet_val \
    --config configs/TC-AE-SL.yaml \
    --rank 0

Citation

@article{li2026tcae,
  title={TC-AE: Unlocking Token Capacity for Deep Compression Autoencoders},
  author={Li, Teng and Huang, Ziyuan and Chen, Cong and Li, Yangfu and Lyu, Yuanhuiyi and Zheng, Dandan and Shen, Chunhua and Zhang, Jun},
  journal={arXiv preprint arXiv:2604.07340},
  year={2026}
}