TC-AE：解锁深度压缩自编码器的 token 容量

首次从 token 空间视角解决深度压缩自编码器的表征坍缩问题

前置知识

潜在扩散模型 (LDM)

一种两阶段图像生成框架：先用自编码器把图像压缩到紧凑的 latent 表示，再用扩散模型在 latent 空间训练生成。生成效率高、效果好，是当前主流方案。

TC-AE 是 LDM 中的 tokenizer，必须理解两阶段流水线才能明白深度压缩 tokenizer 的定位与价值。

视觉 tokenizer (图像自编码器)

把高分辨率图像压缩到少量 latent token 的网络，由 encoder 和 decoder 组成。重建质量 (rFID/PSNR) 和 latent 的生成友好性 (gFID/IS) 是两个独立指标。

TC-AE 本身是一种 ViT 图像自编码器，需要理解其输入输出和评估指标才能看懂实验。

Vision Transformer (ViT)

把图像切成固定大小 $p \times p$ 的 patch，线性投影为 token 序列后用 Transformer 处理。patch 越小则 token 数量越多，模型表达能力越强。

TC-AE 的 encoder 是 ViT 架构，patch size $p$ 决定 token 数，是论文核心的缩放维度之一。

iBOT 自监督学习

结合 masked image modeling 和 class-token 蒸馏的自监督方法，用 EMA teacher-student 框架同时学习局部 patch 语义和全局语义结构。

TC-AE 用 iBOT 作为辅助损失来正则化 token 空间，是其与同类工作的关键区别点。

研究动机

在潜在扩散模型中，研究者追求更激进的压缩比 (例如从 256 个 token 压到 64 个) 以提升效率，方法是在 ViT 编码器末端用大通道数来弥补空间分辨率损失。但这条路在 $32\times$ 压缩比下会遭遇 latent 表征坍缩 (representation collapse)，导致生成质量严重退化。在 ImageNet 256×256 上，DC-AE 仅 64 tokens 时 gFID 高达 26.44，DC-AE-1.5 也只到 17.31，远差于 256 token 的标准方案 (gFID 2.27)。直观上似乎可以靠增加 token 数 (减小 ViT patch size) 来缓解空间信息损失，但论文通过控制实验 (patch size $\in \{32,16,8\}$，$f_{\text{pix}\to\text{lat}}=32$ 固定) 发现一个反直觉的现象：重建质量 rFID 从 14.97 提升到 1.33，但生成质量 gFID 却始终在 30-44 之间徘徊，token 数翻倍也无济于事。

本文的目标是本文目标是在保持 64 token 深度压缩的前提下，让 ViT tokenizer 的生成质量真正受益于 token 数扩展。论文从 token 空间这一被忽视的中间视角切入，目标是同时解决 (1) 为什么 token 数扩展对生成无效 (诊断) 以及 (2) 如何让生成质量随 token 数增长而增长 (方法)。最终在 ImageNet 256×256、64 token 设置下，TC-AE 把 gFID 推到 7.16 (w/o CFG) 和 2.57 (w/ CFG)，与 256 token 的低压缩方案持平甚至更优。

与已有工作不同的是，现有工作主要从两个方向解决表征坍缩：堆参数 (Hansen-Estruch et al. 2025, Xiong et al. 2025) 或对齐大模型特征 (VA-VAE, MAETok, RAE 等)。前者带来 FLOPs 飙升 (DC-AE 需 607 GFLOPs)，后者依赖外部大规模预训练、隐式引入了强数据先验，且对预训练配置 (分辨率、压缩比) 高度敏感。本文首次系统地把视角放在 token 空间上，识别出 token-to-latent 瓶颈才是限制 token 数扩展的根因，并提出轻量的分段压缩 + 内部自监督方案，无需外部大模型、算力开销小 (164 GFLOPs)。

核心方法

TC-AE 的整体思路是「从单点压缩改为两段压缩 + 给 token 注入语义结构」。在 ViT 编码器中，作者把 patch size 减小到 $p=8$ 以获得高分辨率 token (例如 $32\times32=1024$ 个 token)，先经过前 6 个 Transformer block 学习丰富语义，再用一个中间瓶颈 (intermediate bottleneck) 把 token 序列长度压缩到 $1/4$，最后由剩余 6 个 block 加最终瓶颈输出 64 个 compact latent。在训练时，除了标准 $\ell_1 + $ 感知 + 对抗重建损失外，还附加 iBOT 自监督损失 (masked image modeling + class-token 蒸馏) 联合优化。这样设计的核心思想是：让 token 在被压到 latent 之前先学到良好结构，并通过两段压缩避免一次性激进压缩导致的语义破坏。

TC-AE 的核心创新有两点。**第一，分段 token 压缩 (STC)**：传统 ViT tokenizer 只有末端瓶颈 $\mathcal{B}(\cdot)$ 把全部 $N$ 个 token 一次压成 $h\cdot w$ 个 latent。$p$ 减小时 $N$ 激增使 $f_{\text{tok}\to\text{lat}}$ 急剧放大，linear probing 显示 semantic retention $A_2/A_1$ 从 0.31 暴跌到 0.08。STC 把单步压缩拆成 encoder 内 + 末端两步，先在中间把 token 减到 $N/4$，让剩余 block 继续优化结构再压到 latent，$A_2/A_1$ 恢复到 0.31。**第二，联合 iBOT 自监督训练**：无需 DINOv2 等外部教师，直接在 tokenizer 训练中结构化 token 空间 (Table 6: iBOT gFID 17.22 vs DINOv2 28.22)。两者结合 gFID 从 44.72 降到 16.39，IS 从 33.62 涨到 71.33。

方法步骤详情

方法分四步。**(1) Patch embedding**：$X \in \mathbb{R}^{H\times W\times 3}$ 经 $\phi_p(\cdot)$ 切成 $N=HW/p^2$ 个 token 序列，$p=8$ 时 $N=1024$，比 $p=32$ 保留更多细节。**(2) 第一段 ViT**：$M=6$ 个 block 处理高分辨率 token 学局部细节与语义。**(3) 中间瓶颈**：两层卷积把 token 减到 $N/4$，避免一次性激进压缩。**(4) 第二段 ViT + 最终瓶颈**：剩余 6 个 block 继续优化，pixel-shuffle + MLP 输出 $z \in \mathbb{R}^{h\times w\times c}$，实验中 $h\cdot w=64$、$c=128$。训练损失为 $\mathcal{L}_{\text{TC-AE}} = \alpha\mathcal{L}_{\text{rec}} + \mathcal{L}_{\text{iBOT}}$，无需外部大模型即可结构化 token 空间。

技术新颖性

技术新颖性体现在三个层面。诊断层面：首次定量指出 token-to-latent 压缩是 token 数扩展失效的根因，并用 linear probing + 结构损失率 $(1 - A_2/A_1)$ 给出可量化的解释 (Table 1 显示 p=8 时结构损失高达 92%)。架构层面：STC 是个简洁但有效的结构性创新，对比单步压缩 gFID 改善 11.8 个绝对点 (44.72→32.92)，且与 CNN-based 渐缩下采样 (DC-AE) 异曲同工但实现更简单。训练策略层面：联合自监督训练 (iBOT) 让 ViT 编码器自身学到良好 token 语义，避免对 DINOv2 等大型预训练模型的依赖，相比 VTP 显著降低算力需求 (Table 6 显示 iBOT 在 8 卡 batch=512 设置下表现优于 DINOv2，gFID 17.22 vs 28.22)。三者结合构成一个「轻量、纯 ViT、无外部依赖」的深度压缩方案。

**Figure 4**: Token number scaling with staged token compression (STC)

实验结果

核心发现可归纳为四条。**第一，STC 显著缓解瓶颈信息损失**：Table 1 显示 p=8 时单步压缩的 $A_2/A_1$ 仅 0.08，加入 STC 后恢复到 0.31，gFID 从 25.36 降到 16.39。**第二，SSL 在所有 patch size 下都提升生成质量** (Table 3)：p=32 时 gFID 37.09→22.93，p=16 时 31.37→17.22，p=8 时 44.72→25.36，IS 普遍提升 50% 以上。**第三，token 数扩展与模型参数扩展互补** (Fig. 8)：相同 GFLOPs 下 token 数扩展效果优于参数扩展，二者叠加 gFID < 16。**第四，系统级对比** (Table 4) 中 TC-AE 以 64 tokens、164 GFLOPs 的轻量配置，gFID 2.57 (w/ CFG) 优于 DC-AE-1.5 的 17.31，FLOPs 仅为其约 1/4。扩散收敛方面 (Fig. 6) STC 单独使用使 gFID 收敛快 2.7×，完整 TC-AE 达到 4.7× 加速。

**Table 1**: Latent structure loss under token-to-latent compression

**Table 2**: Staged token compression improves generation performance

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

**Table 4**: System-level comparison on ImageNet 256×256

**Table 5**: Layer depth for staged token compression

**Table 6**: Training with other self-supervised learning methods

**Figure 6**: Diffusion model convergence comparison

**Figure 7**: Qualitative comparison between TC-AE and the baseline at 256×256 resolution

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

查看结构化数据

任务	指标	本文	基线	提升
ImageNet 256x256 图像生成 (64 token 深度压缩)	gFID (w/ CFG, ↓)	2.57	DC-AE 26.44 / DC-AE-1.5 17.31	相比 DC-AE-1.5 提升 14.74 个绝对点 (约 85% 相对下降)
ImageNet 256x256 图像生成 (64 token, 无 CFG)	gFID (w/o CFG, ↓)	7.16	DC-AE 26.44 / DC-AE-1.5 17.31	相比最强基线 DC-AE-1.5 提升 10.15 个绝对点
深度压缩 tokenizer 重建质量 (rFID)	rFID (↓)	0.35	DC-AE / DC-AE-1.5 = 0.26	略低于最优基线 (差 0.09)，重建略损换取生成大幅提升
Inception Score (生成质量辅助指标)	IS (↑)	71.33 (Table 2, 完整 TC-AE)	无 STC 无 SSL = 33.62	提升超过 1 倍 (约 112%)
扩散模型收敛速度 (相同 gFID 目标)	加速比	4.7× (gFID 收敛) / 3.5× (IS 收敛)	Vanilla ViT tokenizer (1×)	训练迭代数减少 78%-79%

局限与改进

论文的局限性主要有三点。其一，重建-生成权衡：SSL 加入后重建质量会下降，Table 3 中 p=64 时 rFID 从 14.97 恶化到 30.96，虽然大 patch size 受影响更大，但说明 SSL 在小 token 空间下与重建目标存在一定冲突。其二，对 SSL 方法的依赖性：Table 6 显示 DINOv2 在小 batch 设置下效果较差 (gFID 28.22)，表明 iBOT 的成功不能简单推广到所有自监督方案，TC-AE 的实现需要仔细调参。第三，作者坦诚未在更大规模数据 (如 LAION) 或更高分辨率上验证，64 token 方案是否在文生图场景中同样有效仍未知。此外论文没有分析推理时延迟 (latency)，只比较了 GFLOPs，在硬件上的实际加速未量化。

独立分析的弱点

独立分析可识别三个值得改进的弱点。第一，STC 的中间压缩层数 (M) 选择敏感：Table 5 显示 $M=6$ 时生成最好，但 $M=9$ 或 $M=12$ 时 IS 大幅下降 (从 71.33 降到 52.38)，说明架构对超参敏感，迁移到不同 ViT 规模需要重新搜索 M，建议探索动态或自适应调度方案。第二，SSL 联合训练增加显存和训练时间：iBOT 需要 student-teacher 双网络 + 多 crop 增强 + masking，相比纯重建训练开销翻倍，对算力受限团队是门槛，可考虑用更轻的代理任务 (如简单 masked token prediction) 替代。第三，论文把「token 空间」作为新视角，但 Table 1 中 p=8 + STC 的 $A_2/A_1$ (0.31) 仍低于 p=16 (0.31) 和 p=32 (0.33)，说明即使加入 STC，结构损失并未完全消除，可考虑在 latent 端再做一次 refinement 或加 post-hoc 正则。

未来方向

作者在结论中提到「token 空间优化是一个新且有效的方向」，可延伸的研究至少包括四个层面。其一，扩展到视频/3D tokenizer：视频 token 空间更复杂，token-to-latent 瓶颈更明显，STC 思想有望迁移。其二，结合 DiT 之外的生成范式 (如自回归、MAR)，验证 latent 结构提升是否具有范式无关性。其三，把 STC 与参数扩展做联合 scaling law 研究：Fig. 8 已显示二者互补，但缺乏明确的算力-性能曲线。其四，文生图与多模态场景下，64 token 极度压缩的语义承载能力是潜在瓶颈，可探索条件生成任务中 token 结构化对 prompt-following 的影响。

复现评估

复现可行性较好。代码和模型已在 GitHub 开源 (github.com/inclusionAI/TC-AE)，主实验 ImageNet 256×256 训练细节在 Sec 4.1 和附录中给出：tokenizer 50 epoch (batch=512) + decoder finetune 16 epoch；下游 DiT 800 epoch (batch=1024) 用于系统级对比，80 epoch 用于消融。算力需求上，论文提到 8 张 NVIDIA GPU + batch=512 即可复现 SSL 消融 (Table 6)，完整 scale-up 实验未明示但参照 DiT 量级。难点在于 (a) iBOT 联合训练对超参敏感 (crop 数量、mask 比例、EMA momentum) 容易调崩；(b) 中间瓶颈和最终瓶颈的卷积/pixel-shuffle 实现需要严格对齐；(c) 完整的 800 epoch 训练对个人研究者算力门槛较高。整体可复现性 7/10。

TC-AE：解锁深度压缩自编码器的 token 容量 TC-AE: Unlocking Token Capacity for Deep Compression Autoencoders