瓶颈才是信息杀手：TC-AE 如何让 64 个 token 生成出 2.57 gFID 的图像
一个反直觉的数字
在 ImageNet 256×256 图像生成任务上，当前主流的深度压缩 tokenizer（如 DC-AE）需要 607 GFLOPs 才能达到 gFID 26.44。而蚂蚁集团 inclusionAI 和港科大的合作论文 TC-AE（arXiv:2604.07340）用仅 164 GFLOPs（不到前者的三分之一）达到了 gFID 2.57——性能提升 10 倍，计算量降低 4 倍。
这个结果不是靠更大的模型、更多的数据、或更复杂的架构。它靠的是一个反直觉的洞察：问题不在你看到了多少，而在你如何压缩你看到的。
背景：深度压缩的"通道膨胀"陷阱
潜在扩散模型（Latent Diffusion Models, LDMs）的工作流程很简单：tokenizer 把图像压缩成潜在表示，扩散模型在潜在空间里生成。tokenizer 压缩得越狠，扩散模型越高效——因为潜在空间越小，扩散计算越便宜。
"深度压缩"指的是空间压缩比达到 32 倍甚至更高（比如把 256×256 的图像压成 8×8 的潜在表示）。为了在这种极端压缩下保持重建质量，主流方法（DC-AE、ViTok 等）采用了一个朴素策略：空间压得狠，通道补回来——把潜在表示的通道数从 4 扩到 16、32 甚至 128。
但这个策略有一个副作用：潜在表示崩溃（latent representation collapse）。通道数膨胀带来的重建质量提升，并不能转化为生成质量提升。模型能完美重建图像，但在潜在空间里做扩散时，生成的图像质量很差。
为什么？因为膨胀的通道数捕获的是像素级冗余，不是语义级结构。扩散模型需要的是后者。
关键诊断：瓶颈才是信息杀手
TC-AE 的第一个贡献是一个精确的诊断。论文用 ViT-based tokenizer 做了一组实验：固定像素到潜在的总压缩比，改变 patch size（32→16→8），观察 token 数量增加时会发生什么。
他们用线性探测准确率（Linear Probing Accuracy）作为语义结构的度量：
- $A_1$：瓶颈前的 token 语义结构
- $A_2$：瓶颈后的潜在表示语义结构
- 结构损失：$1 - A_2/A_1$
| Patch Size | gFID ↓ | $A_1$ ↑ | $A_2$ ↑ | 结构损失 |
|---|---|---|---|---|
| 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 |
patch size 越小，token 数量越多，瓶颈前的语义结构越好（$A_1$ 从 29.5 涨到 62.9）。但瓶颈后的语义结构反而崩了（$A_2$ 从 6.3 跌到 5.33）。patch size 8 时，92% 的语义信息在瓶颈处被摧毁了。
这就是"token 数量扩展失败"的根源：问题不在 token，在瓶颈。你给模型更多的 token，它确实能捕获更多语义，但最后那一下暴力压缩把这些语义全毁了。
这就像你写了一封 10000 字的深度报告，最后被要求用一句话总结——信息损失是必然的。但如果你先把它分成 10 段，每段写一个 100 字的摘要，最后再把 10 个摘要合并成一句话，损失会小得多。
解决方案一：分阶段 token 压缩（Staged Token Compression）
TC-AE 的第一个解法就是这个"先分段再合并"的思路。
原始方法：ViT encoder 处理所有 token → 一次性在瓶颈处压缩到潜在表示
TC-AE：ViT encoder 分两阶段 → 第一阶段处理大量 token 学习语义 → 中间压缩层聚合 token → 第二阶段处理压缩后的 token → 最终瓶颈压缩
具体来说，总压缩比 $f_{\text{pix} \to \text{lat}} = 32$ 被分解为：
- patch embedding：$f_{\text{pix} \to \text{tok}} = p^2$（patch size $p$ 决定）
- ViT 内部分阶段压缩：$f_{\text{stage1}} = 2$
- 瓶颈压缩：$f_{\text{bottleneck}} = f_{\text{pix} \to \text{lat}} / (f_{\text{pix} \to \text{tok}} \cdot f_{\text{stage1}})$
tcae/models/tcae.py）：self.compress_rate_bottleneck = int(
compress_rate_total // compress_rate_stage1 // patch_size
)
# 第一阶段：ViT blocks 处理大量 token
moments, output2 = self.encoder(recon_input)
# 中间压缩：pixel-shuffle + MLP
x = rearrange(x, 'b (h p1) (w p2) c -> b h w (p1 p2 c)',
p1=sub_pz, p2=sub_pz)
latent = self.enc_merger(latent_before_compress)
关键在于：压缩不再是一次性的暴力操作，而是分两步完成。第一步在 ViT 内部用注意力机制学习如何聚合语义，第二步才做最终压缩。这样瓶颈处的压缩比大幅降低，语义损失显著减少。
效果：结构保留率 $A_2/A_1$ 从 0.08 提升到 0.31——4 倍提升。
解决方案二：联合自监督学习（Joint SSL）
第二个创新是引入 iBOT（Zhou et al., 2021）作为辅助训练目标。
iBOT 是一个学生-教师蒸馏框架：教师是学生的指数移动平均（EMA），学生输入带 mask 的图像，教师输入完整图像，两者在 token 级别做对齐。这个自监督目标强迫 token 空间学习语义结构——不只是能重建像素，还要能预测被遮蔽的 token。
为什么这很重要？因为传统的 tokenizer 训练目标（像素重建 + 感知损失 + GAN 损失）只关心像素级保真，不关心语义级结构。但扩散模型需要的是后者——它要在潜在空间里做生成，潜在表示必须有足够的语义结构才能支撑生成。
iBOT 的巧妙之处在于：它不需要外部预训练模型（如 DINO、CLIP），而是用自身做教师。这避免了预训练配置（图像分辨率、压缩比）与 tokenizer 设计耦合的问题。
论文的消融实验证明两个创新是互补的：
| 配置 | rFID ↓ | gFID ↓ | IS ↑ |
|---|---|---|---|
| 无 SSL + 无 STC | 1.33 | 44.72 | 33.62 |
| 无 SSL + 有 STC | 0.75 | 32.92 | 43.58 |
| 有 SSL + 无 STC | 0.90 | 25.36 | 52.38 |
| 有 SSL + 有 STC | 0.90 | 16.39 | 71.33 |
两者一起的效果远大于单独效果之和——这说明它们解决的是不同层面的问题：STC 解决"压缩过程中的信息损失"，SSL 解决"压缩前的语义结构不足"。
反直觉发现：token 扩展比参数扩展更有效
论文还有一个值得注意的发现：在相同计算量下，增加 token 数量比增加模型参数更有效。
论文比较了两种扩展策略：
- 参数扩展：固定 patch size，增大 decoder（S→B→L）
- token 扩展：固定 decoder，减小 patch size（32→16→8）
这和 LLM 的 scaling law 有趣地呼应：在语言模型里，token（上下文长度）和参数是两个正交的扩展维度，两者协同效果最好。TC-AE 在视觉 tokenizer 上发现了类似规律。
系统级表现：64 token 打败 256 token
最终系统级比较（Table 4）是论文的杀手锏：
| 方法 | token 数 | GFLOPs | gFID ↓ |
|---|---|---|---|
| DC-AE | 64 | 607 | 26.44 |
| DC-AE-1.5 | 64 | - | 17.31 |
| TC-AE (无 CFG) | 64 | 164 | 7.16 |
| TC-AE (CFG=1.4) | 64 | 164 | 2.57 |
| ViTok | 256 | - | 2.45 |
| VAR | 256 | - | 2.95 |
这意味着：深度压缩不再是生成质量的妥协选项。通过解决瓶颈信息损失，TC-AE 让 64 token 的深度压缩达到了此前只有低压缩比才能达到的生成质量。
代码与可复现性
论文配套代码完全开源在 github.com/inclusionAI/TC-AE，预训练权重在 HuggingFace。仓库结构清晰：
TC-AE/
├── tcae/
│ ├── models/
│ │ ├── tcae.py # 主模型
│ │ ├── encoder.py # ViT encoder + 分阶段压缩
│ │ └── decoder.py # ViT decoder
│ └── utils.py
├── configs/
│ ├── TC-AE-SL.yaml # tokenizer 配置
│ └── DiT-XL.yaml # 下游 DiT 配置
├── diffusion/ # 扩散模型训练
└── requirements.txt
配置文件设计很贴心——关键超参数（compress_rate_total、compress_rate_stage1、patch_size）都在 YAML 里，改起来很方便。模型支持 conv 和 pixel_shuffle 两种压缩类型，SSL 通过 use_ssl: True 开关控制。
对 AI 研究的启示
1. 诊断先于治疗
TC-AE 的故事起点不是"我们提出了一个新架构"，而是"我们做了一个精确的诊断"——用 $A_1/A_2$ 这个度量定位了信息损失发生在哪里。没有这个诊断，就不会有分阶段压缩这个解法。
这和 AI 研究的很多突破类似：真正的创新往往不是"新方法"，而是"新视角"——看到别人没看到的问题。
2. "通道膨胀"是视觉 tokenizer 的"参数膨胀"
当前视觉 tokenizer 领域有一个和 LLM 类似的迷思：越大越好。通道数越膨胀，重建质量越好，但生成质量不一定跟上。TC-AE 证明：问题不在容量，在结构。这和 LLM 里"参数规模不等于能力"的发现同构。
3. 自监督是结构化的通用工具
iBOT 作为辅助目标在 TC-AE 里的成功，再次证明自监督学习的通用性。它不需要外部预训练数据，不需要额外标注，只需要在训练时加一个 mask-predict 目标。这种"内生"的结构化方法对资源有限的研究者特别友好。
4. 深度压缩的"黄金比例"
TC-AE 的消融实验暗示了一个"黄金比例"：瓶颈处的压缩比不宜过大。当 $f_{\text{bottleneck}} > 4$ 时，语义损失急剧上升。分阶段压缩的本质是把大压缩比分解为多个小压缩比——这和信号处理里的多级降采样原理一致。
局限与未来方向
论文承认的局限：
- 只在 ImageNet 上验证：更复杂的数据集（如 COCO、Open Images）表现如何未知
- SSL 增加训练成本：iBOT 的学生-教师框架需要额外的 forward pass
- 两阶段是否最优：论文只验证了两阶段，三阶段或更多阶段是否有进一步收益？
- 把 TC-AE 扩展到视频 tokenizer（时间维度压缩）
- 探索更高效的自监督目标（比 iBOT 更轻量）
- 和最新的扩散模型架构（如 LightningDiT）协同优化
结语：瓶颈是信息论的永恒主题
TC-AE 的核心故事可以用一句话概括：信息损失不在你看到了什么，而在你如何压缩你看到的。
这让人想起 Shannon 的信息论：信道容量决定信息传输速率，信源编码再好也受限于信道。在视觉 tokenizer 里，瓶颈就是信道——无论 token 捕获了多少语义，瓶颈处的压缩比决定了最终能保留多少。
TC-AE 的解法——分阶段压缩 + 自监督结构化——本质上是把信道扩容。不是靠增加通道数（那是假扩容），而是靠减少每次压缩的比率（真扩容）。
当 64 个 token 能生成 gFID 2.57 的图像，当 164 GFLOPs 能打败 607 GFLOPs，我们看到的不仅是一个新 SOTA，而是一个方法论胜利：精确诊断 + 针对性解法 > 暴力扩展。
在 AI 研究越来越崇尚"scale is all you need"的当下，TC-AE 提醒我们：有时候，瓶颈才是问题。
---
论文链接：arXiv:2604.07340 代码：GitHub 模型权重：HuggingFace