TC-AE: Unlocking Token Capacity for Deep Compression Autoencoders

TL;DR

TC-AE is a ViT-based deep-compression image tokenizer from Inclusion AI / Ant Group that attacks “representation collapse under high spatial compression” from the token-space side instead of the usual latent-channel side. Two ingredients: (1) staged token compression that splits the single aggressive token-to-latent bottleneck into an intermediate ViT-internal compression plus a final bottleneck, preventing semantic structure from being destroyed; and (2) joint iBOT self-supervised training to give the token space stronger semantic structure without needing an external pretrained foundation model. On ImageNet 256², a 64-token TC-AE paired with a DiT reaches gFID 7.16 (no CFG) / 2.57 (CFG) at ~164 GFLOPs versus DC-AE’s 26.44 / 607 GFLOPs at the same token count.

Key claims

• Naively scaling token count by reducing patch size improves reconstruction monotonically but does not improve generation under a fixed latent budget — a published failure of the “more tokens = better generator” intuition [§3.2, Fig. 3, Observation 1].
• The root cause is the token-to-latent compression ratio: under a fixed pixel-to-latent ratio, shrinking patch size forces a higher token-to-latent compression at the bottleneck, which destroys semantic structure even as image tokens become more semantically structured pre-bottleneck [§3.3, Eq. 3, Table 1].
• Linear-probing accuracy demonstrates this: at patch size 8 the pre-bottleneck probe hits 62.9% but the post-bottleneck probe collapses to 5.33% (structure loss 0.92); staged token compression at the same patch size recovers post-bottleneck probe to 12.1% (structure loss 0.69) [Table 1].
• Staged token compression alone (no SSL) cuts rFID 1.33 → 0.75 and gFID 44.72 → 32.92 vs. a single-bottleneck baseline at matched compute [§4.2, Table 2].
• Joint iBOT self-supervised training as an auxiliary objective gives the token space stronger semantic structure without external pretraining, replacing the external-DINO-distillation step used by VA-VAE / MAETok / GigaTok [§3.4, Eq. 4-5].
• SSL + staged token compression are complementary: combined they reach gFID 16.39 with IS 71.33 vs. baseline 44.72 / 33.62 [§4.2, Table 2].
• Token-count scaling and decoder-parameter scaling are complementary, not substitutive, and at matched GFLOPs token scaling yields stronger generative gains than parameter scaling [§4.3, Fig. 8].
• At 64 latent tokens, TC-AE + DiT reaches gFID 7.16 (no CFG) / 2.57 (CFG) at ~164 GFLOPs, vs. DC-AE 26.44 and DC-AE-1.5 17.31 at 607 GFLOPs [§4.4, Table 4].
• A 64-token TC-AE matches or beats 256-token low-compression tokenizers (ViTok 2.45, VAR 2.95) — i.e. it claims to close the deep-compression vs. low-compression generative-quality gap [§4.4, Table 4].

Method

A standard ViT autoencoder factors the total spatial compression into pixel-to-token (patch embedding) × token-to-latent (bottleneck). TC-AE keeps the same overall compression ratio but redistributes work. The encoder begins with a small patch size to produce a long, high-resolution token sequence, runs N₁ ViT blocks to let semantic structure emerge across tokens, then applies an intermediate compression (a 2-layer conv) that reduces the sequence length to 1/4, runs N₂ further ViT blocks, and finally a second bottleneck (pixel-shuffle + MLP) produces the compact latent. The decoder is structurally symmetric. Training combines the usual reconstruction stack (pixel L1 + LPIPS perceptual + adversarial) with an iBOT auxiliary objective: a student–teacher (EMA) distillation where the student sees masked global crops and additional local crops, predicting teacher patch-token outputs on the masked global crops (masked-image-modeling loss) and aligning class-token outputs across local/global views (cross-view distillation loss). Crucially, iBOT supervision is internal — no external DINOv2 / MAE features are required, in contrast to VA-VAE, MAETok, GigaTok, or RAE which all rely on a frozen pretrained vision backbone.

Results

System-level on ImageNet 256² with the same DiT backbone:

• TC-AE (64 tokens, ~164 GFLOPs): gFID 7.16 (no CFG), 2.57 (CFG).
• DC-AE baseline (64 tokens, 607 GFLOPs): gFID 26.44.
• DC-AE-1.5 (64 tokens): gFID 17.31.
• Low-compression reference: ViTok S-B (256 tokens) 2.45; VAR-Tok (680 tokens) 2.95; SD-VAE + DiT (256 tokens) 2.27.

Ablations (Table 2, no-CFG ImageNet @ 80 epochs DiT training):

• baseline (no STC, no SSL): rFID 1.33, gFID 44.72, IS 33.62.
• • STC only: rFID 0.75, gFID 32.92, IS 43.58.
• • SSL only: rFID 0.90, gFID 25.36, IS 52.38.
• • STC + SSL (full TC-AE): rFID 0.90, gFID 16.39, IS 71.33.

Patch-size sweep (Table 3): SSL consistently improves gFID at every patch size (8 → 64); reconstruction is hurt more at large patch size (rFID 14.97 → 30.96 at p=64) and roughly flat at small patch size (rFID 1.33 → 0.90 at p=8).

Scaling figure (Fig. 8): at matched GFLOPs, scaling token count beats scaling decoder parameters, and the two combine additively.

Why it’s interesting

TC-AE sits at the same junction as DINO-SAE: DINO Spherical Autoencoder for High-Fidelity Image Reconstruction and Generation and the RAE/RAEv2 lineage (Improved Baselines with Representation Autoencoders, Scaling Text-to-Image Diffusion Transformers with Representation Autoencoders) — all are recipes that fix “the VFM tokenizer trades reconstruction for semantics” — but proposes the most architecturally radical fix: instead of changing what feature the tokenizer aligns to (RAE: use DINO; DINO-SAE: cosine-only on DINO; RAEv2: multi-layer DINO), TC-AE drops the external pretrained encoder entirely and structures the token space via iBOT self-supervision during tokenizer training. This puts it in the same “no external foundation model needed” camp as No Other Representation Component Is Needed: Diffusion Transformers Can Provide Representation Guidance by Themselves (SRA) (SRA) and Self-Flow: Self-Supervised Flow Matching for Scalable Multi-Modal Synthesis (Self-Flow), but operates at the tokenizer layer rather than the generator layer — which is a strictly upstream lever. The staged-token-compression idea is also independently interesting: it identifies a specific architectural pathology (single-bottleneck token-to-latent compression destroys structure) that almost every deep-compression VAE / VAE-VAE / DC-AE inherits silently. If the diagnosis holds, this should drop in cleanly on top of UL (Unified Latents (UL): How to train your latents), DC-AE-1.5, and the LTX-2 / Wan tokenizer stacks Luma cares about.

See also

• DINO-SAE: DINO Spherical Autoencoder for High-Fidelity Image Reconstruction and Generation — same problem (VFM-tokenizer reconstruction-vs-semantics tradeoff), different fix (drop magnitude matching + Riemannian FM on sphere) versus TC-AE’s staged compression + iBOT.
• Improved Baselines with Representation Autoencoders — RAEv2; concurrent fix via multi-layer DINO aggregation + REPA, still relying on external DINO supervision that TC-AE eliminates.
• Scaling Text-to-Image Diffusion Transformers with Representation Autoencoders — RAE for T2I; the “use DINO as the encoder” lineage TC-AE positions against.
• No Other Representation Component Is Needed: Diffusion Transformers Can Provide Representation Guidance by Themselves (SRA) — SRA: “no external representation component needed” thesis at the generator level; TC-AE makes the analogous argument at the tokenizer level.
• Self-Flow: Self-Supervised Flow Matching for Scalable Multi-Modal Synthesis — Self-Flow; internal-self-distillation in the generator. Pairs naturally with TC-AE’s internal-iBOT in the tokenizer to give a fully external-encoder-free stack.
• Unified Latents (UL): How to train your latents — UL; latent-design-as-efficiency-lever from the encoder noise floor side. Complementary to TC-AE’s token-count side.
• Diffusion training efficiency — concept page; TC-AE is the tokenizer-side variant of “internal self-distillation beats external DINO” that the concept page already tracks at the generator side.