WisPaper
WisPaper
学术搜索
学术问答
价格
TrueCite
[ICLR 2025/Archived] CaTok:驯服平均流(Mean Flows),构建真正的 1D 因果视觉分词器
总结
问题
方法
结果
要点
摘要

本文提出了 CaTok,一种新型一维因果图像分词器(1D Causal Image Tokenizer)。该方法结合了提供因果约束的 MeanFlow 解码器与基于寄存器的 ViT 编码器,在 ImageNet 上实现了 SOTA 的重建性能(0.75 rFID),并支持极速单步生成与高质量多步采样。

TL;DR

在视觉生成领域,如何让图像像文本一样具备“因果逻辑”一直是自回归(AR)模型的痛点。复旦大学等机构的研究者提出了 CaTok,这是一种基于 MeanFlow 的 1D 图像分词器。它不仅解决了传统方法在训练过程中的不平衡问题,还实现了高效的单步采样。在 ImageNet 重建任务中,CaTok 刷写了多项指标,并在 AR 生成测试中展现了极高的训练效率。

背景定位:视觉进入“1D 因果”时代

自回归模型(如 Llama)的核心在于因果分词(Causal Tokenization)。文本有天然的先后顺序,但图像是一个 2D 网格。目前的做法要么是生硬地展平(Flatten),要么是像 VAR 模型那样进行多尺度预测,这都与 LLM 的“下一个 Token 预测”范式存在偏差。

CaTok 的出现,旨在利用最近兴起的扩散自动编码器(Diffusion Autoencoders),将图像压缩为一串具备严格时序逻辑和语义因果的 1D 向量。


痛点深挖:不平衡的代价

作者指出,目前的 1D 分词器(如 Semanticist, FlexTok)为了引入因果性,通常采用嵌套 Dropout(Nested Dropout):训练时随机给解码器输入前 个 Token。

  • 问题:早期的 Token(编号靠前的)被采样的频率远高于后期 Token。这种训练不平衡会导致 AR 模型在生成长序列时,后期的预测质量迅速坍塌。
  • 代价:为了弥补这种不平衡,往往需要复杂的重新加权机制或超长的训练周期。

核心方法论:MeanFlow 与时间间隔绑定

1. 架构解析

CaTok 由一个带**因果掩码(Causal Mask)**和寄存器的 ViT 编码器以及一个基于 MeanFlow 的 DiT 解码器组成。

  • Encoder:图像通过 ViT 提取特征,通过寄存器(Registers)压缩成 1D Token。关键在于 Causal Mask 确保后方的 Token 无法看到前方的信息。
  • Decoder:不同于传统的单点流匹配,CaTok 使用了 MeanFlow 目标函数。

模型架构图

2. 数学直觉:为什么选 MeanFlow?

传统扩散模型拟合的是瞬时速度场 。CaTok 通过 MeanFlow 拟合时间区间 内的平均速度场 通过将 Token 序列的区间映射到扩散步的时间区间,CaTok 实现了:

  • 因果性:特定段落的 Token 负责生成扩散路径中对应的特定阶段。
  • 单步生成:MeanFlow 天然支持通过一次推理估算出整个区间的变化,从而实现单步从噪声到图像的重建()。

3. REPA-A:来自基础模型的“援军”

为了让训练更稳、收敛更快,作者提出了 REPA-A。它将编码器提取的特征直接与 DINOv2 等预训练视觉大模型(VFM)的特征进行余弦相似度对齐。这一步确保了分词器在训练初期就能捕捉到高质量的语义信息。


实验结果:效率与质量的双赢分析

Reconstruction 战绩

在 ImageNet 256x256 重建测试中,CaTok-L 在仅使用 160 个 epoch 的情况下,FID 达到了 0.75。

| Method | Tokens | rFID ↓ | PSNR ↑ | SSIM ↑ | | :--- | :--- | :--- | :--- | :--- | | TiTok-L-32 | 32 | 2.21 | 15.60 | 0.359 | | CaTok-L-256 (Ours) | 256 | 0.75 | 22.53 | 0.674 |

实验结果对比 图 1:展示了从 256 个 Token 到 16 个 Token 的重建效果。可以看到随着 Token 减少,图像呈现出“由精细到粗糙”的退化,这有力证明了其分词的因果层级结构。

自回归生成表现

CaTok 产生的 Token 非常契合 Llama 架构。在与 LlamaGen 结合后,由于其 Token 具备更强的因果连贯性,生成效果优于传统非因果分词器。


深度洞察与总结

为什么 CaTok 更好?

  1. 解决了“头重脚轻”问题:通过时间间隔绑定,避开了 Nested Dropout 带来的步数不平衡,让所有 Token 在训练中各司其职。
  2. 多尺度重建的优雅实现:不需要像 VAR 那样设计复杂的多尺度网格,仅仅通过截断 1D 序列就能实现从语义轮廓到细节纹理的平滑控制。
  3. 训练稳定性:REPA-A 有效抑制了 MeanFlow 训练中常见的损失尖峰。

局限性与未来

虽然 CaTok 在重建指标上表现惊人,但在纯单步生成(无 AR 辅助)的情况下,FID 仍与 GAN 系列方法有微小差距。未来的研究方向可能在于如何进一步将这种 1D 因果表征拓展到更高分辨率(如 1024px)以及视频理解领域。

Takeaway: CaTok 证明了,视觉分词不需要为了因果性而牺牲训练效率。通过数学上的流对齐,我们可以获得兼具 LLM 兼容性与扩散模型强大生成能力的完美表征。

发现相似论文

试试这些示例

  • 查找最近其他利用 MeanFlow 或 Rectified Flow 改进扩散自动编码器图像重建能力的论文。
  • 哪篇论文最早在 ViT 中引入了 Register(寄存器)概念用于特征提取,本文是如何将其扩展到因果分词任务中的?
  • 探究 CaTok 提出的 REPA-A 特征对齐方法是否可以应用到视频分词器或多模态理解模型的预训练中?
目录
[ICLR 2025/Archived] CaTok:驯服平均流(Mean Flows),构建真正的 1D 因果视觉分词器
1. TL;DR
2. 背景定位:视觉进入“1D 因果”时代
3. 痛点深挖:不平衡的代价
4. 核心方法论:MeanFlow 与时间间隔绑定
4.1. 1. 架构解析
4.2. 2. 数学直觉:为什么选 MeanFlow?
4.3. 3. REPA-A:来自基础模型的“援军”
5. 实验结果:效率与质量的双赢分析
5.1. Reconstruction 战绩
5.2. 自回归生成表现
6. 深度洞察与总结
6.1. 为什么 CaTok 更好?
6.2. 局限性与未来