CaTok：用 MeanFlow 驯服一维因果图像分词 CaTok: Taming Mean Flows for One-Dimensional Causal Image Tokenization

完整训练流水线分为四个阶段。第一阶段是编码器前向：将图像 patch $x$ 与 K 个 register 拼接输入因果 ViT 编码器 $E_\delta$，得到图像特征 $H_e$ 与 1D 令牌 $V_K = [V_1,...,V_K]$。因果掩码规则是：图像 patch 之间全连接，patch 可以被后续 register 看到，但 register 之间严格只看到之前的 register 与全部 patch。第二阶段是解码条件构造：在 $[0,1]$ 上均匀采样 $r<t$，令 $z_t=(1-t)x+t\epsilon$，并把区间 $[r\cdot K, t\cdot K]$ 内的 register 切片 $V_{r:t}$ 作为 DiT 的 cross-attention 条件；同时按比例 $q=75\%$ 把 $r$ 重置为 $t$，此时区间退化为单点 $V_K$，DiT 拟合瞬时速度 $v_\theta$ 用于稳定训练。第三阶段是四个损失联合优化：(a) MeanFlow 目标 $\mathcal{L}_{MF}=\mathbb{E}\|u_\theta - (\epsilon-x) - \text{sg}[(t-r)((\epsilon-x)\partial_z u_\theta+\partial_t u_\theta)]\|^2$，其中 $u_\theta = D_\theta(z_t,r,t,V_{r:t})$，sg 表示停止梯度；(b) Rectified Flow 目标 $\mathcal{L}_{RF}=\mathbb{E}\|v_\theta - (\epsilon-x)\|^2$，只在 $r=t$ 时启用；(c) REPA 把 DiT 中间层 patch 特征投影后与 DINOv2 patch 表征做余弦相似度 $\mathcal{L}_{REPA}$；(d) 本文新加的 REPA-A 把编码器输出 $H_e$ 直接与 DINOv2 patch 表征做对齐 $\mathcal{L}_{REPA-A}=-\mathbb{E}[\frac{1}{N}\sum_n \text{sim}(H^n_{vfm}, H^n_e)]$。总损失为四项加权和，$c=10^{-3}$、$w=1.0$ 用于自适应 L2。第四阶段是 AR 训练与采样：冻结 CaTok 编码器，用 ϵLlamaGen-L (343M 参数) 配合 MAR [29] 的 diffusion loss 训练 400 epoch；采样时以类别 token 为起点自回归预测连续令牌，再用 MeanFlow 一步解码器生成图像。

技术新颖性体现在三个层面。第一是范式层面：首次把扩散自编码器从'无条件 + 多步'推向'有条件区间 + 一步'，用 $[r,t]$ 这个与去噪过程同构的物理量替代人为的'前 k 个'选择，使因果性、平衡性、采样效率三个目标首次在统一框架下达成。第二是数学层面：把 MeanFlow 平均速度 $u(z_t,r,t)$ 的两个时间参数与令牌索引 $r\cdot K$、$t\cdot K$ 做精确绑定，这种'时间-令牌'同构映射是新颖设计；论文也展示了在 $r=t$ 退化情形下自然退化为普通整流流，从而保留了训练稳定性。第三是表征对齐层面：与 REPA [71] 对齐 DiT 中间层、REPA-E [26] 对齐 VAE 不同，REPA-A 直接对齐编码器 patch 输出，这要求作者从'对齐潜在变量'转向'对齐条件信号源'，对应了一类新的正则化设计。整体上，CaTok 没有引入全新的网络结构 (仍然是 ViT + DiT + registers)，其新颖性主要来自于把多个已有组件重新组装为更自洽的因果分词框架。

独立分析几个具体弱点。第一，连续令牌 + diffusion loss 的组合让 AR 训练本身比 VQ 方案慢且更难调，作者在 Tab. 5 中承认 CaTok-VQ 的 rFID 3.81 远不如 LlamaGen 2.19，说明如果不专门为 VQ 调训练超参、加入感知损失或后训练技巧，离散化版本质量损失很大。改进方向：可以借鉴 MAR-VAE 的微调策略或加入 perceptual loss + GAN loss 来缩小这一差距。第二，论文主要实验在 256×256 完成，512 分辨率要么依赖 SD-VAE + 训练 (rFID 1.07) 要么用训练免费的 patchwise 策略 (rFID 0.60)，但前者性能下降、后者依赖额外的 blend 步骤，不够干净。改进方向：在更高分辨率上重新训练 DiT 与 MeanFlow 联合目标，并验证不依赖 patchwise 技巧的高分辨率能力。第三，MeanFlow 的 Jacobian-vector 积实现限制了 batch size 与分辨率，论文未充分讨论显存与吞吐。改进方向：可以借鉴 SiT 中的 gradient checkpointing 或分段 JVP 进一步降低开销。第四，论文的 [r,t] 采样采用均匀策略，缺乏对不同采样分布的消融；当 AR 模型仅用前 32 个令牌时 (gFID 4.91 vs. CaTok-L-32 3.40)，性能反而比用前 128 更好，暗示令牌排序与训练分布并不完全对齐——这值得进一步探索自适应采样或课程式调度。第五，CaTok 的 AR 生成 gFID 2.95 (128 令牌) 仍略输 Semanticist 2.57 (256 令牌)，论文没有给出用 256 令牌做 AR 的结果，可能存在长序列退化。

作者在结论与附录中点明了几个未来方向：在更大规模数据集 (如 LAION、ImageNet-21K) 上做 tokenization 预训练；将 CaTok 用于文本到图像、图像编辑、视频等更广泛的任务；探索 REPA-A 与不同 VFM 教师 (DINOv3、SigLIP2 等已在 Tab. 中展示 1.16/1.12 的 rFID) 的组合。基于论文成果还可以延伸：第一，'时间-令牌'同构思想可以推广到音频、视频等多模态分词，把时间维度天然作为令牌顺序；第二，REPA-A 的思路可以泛化为'对齐编码器而非解码器'的正则化族，用于其他条件生成模型；第三，论文附录显示 CaTok 与非自回归生成器 (MaskGiT-L, 8 步) 也兼容 (gFID 2.69 vs. TiTok 2.77)，说明该分词器是生成范式无关的，可以与任何 token-wise 模型组合；第四，MeanFlow 一步采样与一步 AR 结合后，端到端推理只需 AR 自回归 + 一步解码，对部署极其友好，可在边缘设备上做进一步优化；第五，结合 LLM 的 instruction tuning 范式，CaTok 的语义可解释性 (论文 Fig. 1 第 7-10 列的'概念解码') 有望支持可控图像编辑——给定一段语义描述，仅替换对应区间的令牌即可。

复现评估整体较好但有一定门槛。论文提供了项目主页 https://sharelab-sii.github.io/catok-web，附录 Tab. 4 给出了完整的训练超参：CaTok-B 与 CaTok-L 的优化器 (AdamW)、峰值学习率 (1e-4 tokenization、5e-5 AR)、batch size (1024/2048)、weight decay (0.05)、EMA (0.999)、gradient clipping (3.0)、warmup (96 epoch for AR) 全部公开。训练 epoch 80 (CaTok-B) / 160 (CaTok-L) tokenization 与 400 epoch AR，论文也明确指出 MeanFlow 目标在 epoch 10/20 引入、选 [r,t] 机制在 epoch 40/80 引入。开源方面，论文提到 Semanticist 用官方代码库复现，但 CaTok 代码未在论文中给出明确 GitHub 链接 (仅项目主页)，需要读者自行访问网站确认。算力上，CaTok-L-256 用 552M 参数 + 1024 batch + 160 epoch，单卡 H100/A100 (80G) 大约需要数天到一周，AR 模型 343M × 400 epoch × 2048 batch 显著更贵。复现难度属于中等偏上：(a) MeanFlow 的 JVP 实现需要自定义 autograd 函数；(b) 因果 ViT 掩码与 register 索引映射容易出错；(c) REPA-A 要求把 DINOv2-B/16 的 patch 表征与编码器 patch 对齐，patch 网格尺寸需对齐 (ViT-B/8 输出与 DINOv2-B/16 输出均为 32×32，所以 OK)，但若换不同 teacher 需要重新检查；(d) AR 模型使用 MAR [29] 的 diffusion loss，需理解其连续令牌的高斯噪声建模。整体而言，读者有 DiT/ViT/LlamaGen 经验的话可以基于公开 Semanticist/MAR 代码做迁移。