一句话定位
OmniTokenizer 是复旦大学与字节跳动 2024 年 6 月提出的 Transformer 式联合图像-视频分词器（tokenizer）：用时空解耦架构（空间维窗口注意力 + 时间维因果注意力）和”先图后视频”的两阶段渐进式训练，一套模型、一套权重同时服务图像与视频的离散化分词（VQVAE）与连续隐空间分词（VAE），在图像和视频重建上都刷新 SOTA，并证明同一个分词器可以底层通用地喂给自回归 Transformer 生成模型和扩散生成模型两条路线。
背景与定位
视觉生成的两大范式——语言模型式的自回归生成（vit-vqgan、maskgit）与扩散模型（DiT、Latent Diffusion）——都依赖一个”分词器”把高维像素压缩到紧凑隐空间：前者用 VQVAE 离散化（如 taming-transformers-vqgan、vit-vqgan），后者用 VAE 建模连续分布。论文指出现有分词器普遍是”单模态专用”的：图像分词器（taming-transformers-vqgan、vit-vqgan）不管视频，视频分词器（videogpt、TATS）不管图像；即便 magvit-v2（MAGVITv2）用因果 3D 卷积尝试同时处理两种模态，也仍需为图像和视频分别训练模型，没有真正实现跨模态协同——这正是 OmniTokenizer 要填的空白：一个共享框架+共享权重的联合分词器，既能让视频数据稀缺的问题借图像数据缓解，又能让分词器学到更通用的视觉表征。
方法上，论文把这类需求归纳为”时空解耦架构”（借鉴 TimeSformer 的空间/时间注意力分离思路），并提出”渐进式训练”作为连接图像与视频这两种异质数据的桥梁。与 magvit-masked-generative-video-transformer/magvit-v2 的非自回归、dit-scalable-diffusion-transformers 的扩散路线不同，OmniTokenizer 本身只是分词器层，刻意保持与下游生成范式无关，因此论文分别在其 VQVAE 分词之上接自回归 Transformer、在 VAE 分词之上接 DiT/Latte 扩散模型做验证。范式命名：joint image-video tokenization / spatial-temporal decoupled tokenizer / progressive image-then-video training。作为世界模型评测脉络里的方法学条目，它的定位是”通用视觉分词器”这一底座能力，可类比 ivideogpt 中为世界模型专门定制的条件 VQGAN——OmniTokenizer 提供的是不针对任何具体下游（无动作条件）、开箱即用的联合分词替代方案。
模型架构
Patchify（分块嵌入）：输入 x ∈ ℝ^((1+T)×H×W×3)，首帧 x₀ 与后续帧 x_{1:T} 分开处理（延续 MAGVITv2 的”首帧独立”思路，服务图像/视频统一编码）。x₀ 按 p×p 分块，x_{1:T} 按 t×p×p 的时空块分块，各自过一层独立线性投影得到 patch embedding e₀ ∈ ℝ^(L₁×c)（L₁=H/p×W/p）与 e_{1:T} ∈ ℝ^(L₂×c)（L₂=T/t×H/p×W/p），拼接为时空联合序列 e；总体把 (1+T)×H×W 压缩到 (1+T/t)×H/p×W/p。
编解码器（时空解耦）：空间维用窗口注意力（借鉴 Swin，局部聚合、效率高），时间维用因果注意力（对齐后续自回归生成、保证时序一致）；解码器结构与编码器对称，最终用两层线性投影（无激活函数）把时空 token 映回像素空间。官方参考训练脚本（scripts/recons/train.sh）给出的默认配置与论文 Implementation Details 完全对应：4 层窗口注意力空间块 + 4 层因果注意力时间块（spatial_depth=4/temporal_depth=4）、窗口大小 8（twod_window_size=8）、隐藏维 512（embedding_dim=512）、注意力头数 8、每头维 64、空间 patch size 8、时间 patch size 图像阶段 2 / 图像-视频联合阶段 4（对应论文”视频序列长度 17 帧”的设置）。
量化 / 隐变量：VQVAE 分支通过码本查表得到离散码 z_q，码本大小 8192，隐（码本）维度 8（跟随 vit-vqgan 用 factorized code + l2-normalized code 提升码本利用率）；VAE 分支从高斯分布采样连续隐变量。两个分支共享同一套编码器骨干,区别只在量化层（VQ 训练目标 vs. KL 微调目标）。
下游生成头：
- 自回归语言模型：把 VQVAE 码索引按光栅序展平成序列 y，训练 Transformer 做下一 token 预测，配置沿用 VQGAN，扩大版把隐藏维提到 1535（沿用 ViT-VQGAN 的放大方式）；论文报告了 227M 与 650M 两档参数量。
- 隐扩散模型：在 OmniTokenizer-VAE 的连续隐空间上，直接套用 DiT（图像）与 Latte（视频）的架构与训练流程，仅把它们原本的 Stable-Diffusion VAE 换成 OmniTokenizer-VAE。
分词器自身总参数量论文未单独披露（只给出隐藏维/隐变量维/层数），标记为未披露。
数据
分词器训练/评测所用数据集：图像——ImageNet、CelebA-HQ、FFHQ；视频——UCF-101、Kinetics(-600)、Moments-in-Time (MiT)、Something-Something v2 (SSv2)。均为公开真实世界数据集，无仿真、无动作标注（OmniTokenizer 本身不做动作条件，是纯视觉重建/生成分词器）。
渐进两阶段训练的数据配方：
- 阶段一（图像预训练）：单一固定分辨率 256×256，仅图像数据，训练 50 万迭代。
- 阶段二（图像-视频联合训练）：图像与视频数据交替前向，视频序列长度固定 17 帧，空间分辨率从 {128, 192, 256, 320, 384} 中随机采样（多分辨率训练），同样训练 50 万迭代。
- 数据增强：论文正文只提到随机水平翻转；官方训练脚本另外开启了
--apply_noise（各阶段）与--apply_blur（仅阶段一），两者口径略有出入，如实并陈。 - 论文明确说明：“除非另有说明，本文汇报结果均为 ImageNet + UCF-101 联合训练”，即默认数据组合是一图一视频。
发布的 checkpoint 矩阵（来自 GitHub/HF 模型仓库，13 组）：单域（ImageNet-only / CelebAHQ-only / FFHQ-only）与联合域两大类；联合域覆盖 3 个图像源（ImageNet / CelebAHQ / FFHQ）×最多 4 个视频源（UCF / K600 / MiT / SSv2）的组合，每组单独训练、各自报告 rFID/rFVD（VQVAE 8 组 + VAE 2 组，见下方评测表）。论文与开源代码均未披露各数据集参与训练的具体视频小时数/帧数/样本条数，只给出数据集名称与训练时用的分辨率、序列长度。
训练方法
目标函数：
- VQ 训练阶段：ℒ_VQ = λ₁‖sg[E(e)]−z_q‖² + λ₂‖E(e)−sg[z_q]‖²（λ₁=λ₂=1，sg 为 stop-gradient），叠加 L2 重建损失、GAN 损失（各阶段图像/视频判别器权重不同，见训练脚本）与感知损失（权重 4）。
- KL 微调阶段：把 ℒ_VQ 换成 ℒ_KL = λ₃·D_KL(Q(z|x)‖P(z))（λ₃=1e-6，P(z) 为高斯先验），复用同样的重建/GAN/感知损失，从阶段二训练好的 VQVAE checkpoint 初始化继续微调，得到 OmniTokenizer-VAE。
渐进式三阶段流程：图像单模态预训练（阶段一，500K 步，LR warmup 到 1e-3 后 cosine 衰减到 0，Adam β1=0.9/β2=0.99）→ 图像-视频联合训练（阶段二，500K 步，LR 显著调低到 5e-5，从阶段一 checkpoint 继续）→ KL 微调（阶段三，同样从阶段二 checkpoint 继续，LR 同为 5e-5）。消融实验（下节 Table 7）证明这一”先图后视频、渐进式”顺序显著优于任何单阶段替代方案（纯视频训练、固定分辨率联合训练等）。
下游生成训练：AR Transformer 用标准交叉熵下一 token 预测（对 VQVAE 码索引序列，条件 c 为类别标签）；扩散模型直接复用 DiT / Latte 各自的训练流程，无需为接入新分词器改动训练本身。
关键消融：
- 架构消融（Table 3，ImageNet）：对比联合时空注意力（JointAttn）、解耦但用 plain attention（DePlainAttn）、本文的解耦窗口+因果注意力（Ours）——Ours 推理开销最低（51 iGFLOPs / 262 vGFLOPs，VQGAN 基线为 167/167×17）且 rFID 最优（1.28，JointAttn 1.89，DePlainAttn 1.33）。
- 训练范式消融（Table 7，跨 ImageNet + 4 个视频数据集）：Joint(Multi)（本文最终方案）在几乎所有列上最优；Joint(Fix)（联合训练但固定分辨率）虽然视频指标优于纯视频训练，但 ImageNet rFID 反而从 1.28 恶化到 1.35——说明”多分辨率”是联合训练不损害图像质量的关键。
- 隐变量维度/压缩率消融（Figure 3，定性）：压缩率越高 rFID 越差；隐维度=8 是 rFID 与码本利用率的最佳平衡点，据此定为默认配置。
Infra（训练 / 推理工程）
- 训练硬件（论文正文）：8× NVIDIA A100 GPU，训练约 2 周（针对论文汇报的分词器训练）。
- 官方参考脚本默认配置：
scripts/recons/train.sh三个阶段均设--num_nodes 4 --gpus 8，即默认多机 32 卡分布式训练——与正文”8×A100、2 周”的表述不完全一致（可能正文数字针对某个具体训练配置，脚本给的是可扩展的通用模板），如实并陈、不做主观调和。 - 训练精度：论文正文与开源脚本参数均未披露 fp16/bf16/fp32 等精度设置。
- 推理成本：论文只用硬件无关的 GFLOPs 衡量效率（Table 3：Ours 51 iGFLOPs/262 vGFLOPs vs. VQGAN 167/167×17，JointAttn 72/358，DePlainAttn 72/299），未报告实际推理 FPS、控制频率或端侧延迟，未披露。
评测 benchmark
图像重建 rFID（同压缩率、同码本大小 8192）：
| 方法 | 数据集 | rFID |
|---|---|---|
| ViT-VQGAN | CelebAHQ | 4.66 |
| Ours-VQVAE | CelebAHQ | 1.93 |
| ViT-VQGAN | FFHQ | 3.13 |
| Ours-VQVAE | FFHQ | 1.91 |
| DALL-E | ImageNet | 32.01 |
| VQGAN* | ImageNet | 1.49 |
| ViT-VQGAN | ImageNet | 1.28 |
| Ours-VQVAE | ImageNet | 1.11（-13% vs ViT-VQGAN） |
| Ours-VAE | ImageNet | 0.69 |
视频重建 rFVD：
| 方法 | 类型 | UCF-101 | MiT |
|---|---|---|---|
| MaskGIT | 图像式 | 240 | - |
| VQGAN | 图像式 | 299 | 306 |
| ViT-VQGAN | 图像式 | - | 167 |
| CViViT | 视频式 | - | 66 |
| TATS | 视频式 | 162 | - |
| MAGVIT | 视频式 | 58 | - |
| Ours-VQVAE | 联合 | 42（-26% vs MAGVIT） | 20 |
| Ours-VAE | 联合 | 23 | 13 |
自回归图像生成（ImageNet 256×256，class-conditional）：227M 档，VQGAN* 18.65 FID/80.4 IS、RQ-Transformer 15.72/86.8、Ours 10.13/94.5；650M 档，VQVAE-2 31.11/~45、VQGAN 15.78/74.3、RQ-Transformer 13.11/104.3、ViT-VQGAN 8.81/110.8、Ours 7.45/146.7。
自回归视频生成 FVD：UCF-101（class-conditional）——TATS(321M) 332，Ours(227M) 314；K600（frame prediction）——Phenaki(227M,NAR) 36.4、MAGVIT(306M,NAR) 9.9、MAGVITv2(307M,NAR) 4.3、LVT(50M) 224.7、ViTrans(373M) 170.0、CogVideo(9.4B) 109.2、ViVQVAE 64.3、Ours(227M) 34.2、Ours(650M) 32.9。注：正文另有一句”FVD(283 v.s. 314)“描述 Ours 相对 TATS 的优势，但与表格中 TATS=332 一行对不上，无法进一步核实这处文本与表格的差异，如实标注。
扩散模型 · 图像（ImageNet 256×256，DiT-XL/2 架构）：DiT-XL/2 原版（用 SD-VAE）9.62 FID/121.50 IS/0.67 Prec/0.67 Rec；DiT-XL/2-CFG 原版 2.27/278.24/0.83/0.57；Ours-DiT-XL/2 12.25/109.94/0.73/0.64；Ours-DiT-XL/2-CFG 3.48/244.23/0.89/0.52。如实说明：换成 OmniTokenizer-VAE 后，FID/IS 相比原版 SD-VAE 基线均略逊（3.48 vs 2.27，244.23 vs 278.24），论文的表述”achieves a better inception score”指的是加 CFG 相对不加 CFG 的提升（109.94→244.23），并非超越原版基线；Precision 略优（0.89 vs 0.83）而 Recall 略低（0.52 vs 0.57）。
扩散模型 · 视频（UCF-101 无条件生成，Latte 架构）：MoCoGAN 2886.9、VideoGPT 2880.6、MoCoGAN-HD 1729.6、DIGAN 1630.2、StyleGAN-V 1431.0、PVDM(隐压缩 1×4×4) 1141.9、MoStGAN-V 1380.3、Latte 原版(隐压缩 1×8×8) 478.0、Ours-Latte(隐压缩 4×8×8) 209.2——在更高压缩率（时间维额外压 4×）下仍大幅超越原版 Latte 及所有 GAN 基线，是本文扩散侧最干净的胜利。
创新点与影响
- 贡献：(1) 首个用共享框架、共享权重同时做图像与视频分词的 tokenizer，填补了此前”图像分词器/视频分词器各自为战”（连 magvit-v2 都仍需分模态训练）的空白；(2) 时空解耦（窗口注意力+因果注意力）架构在效率（GFLOPs）与重建质量上同时优于联合注意力/解耦纯注意力两个替代设计；(3) 渐进式（先图后视频、多分辨率）训练策略被消融证明是”联合训练不损害单模态质量”的关键，且反直觉地证明视频数据的加入还能提升图像重建质量（ImageNet rFID 1.28→1.11）；(4) 证明同一套 VQVAE/VAE 分词器可以分别驱动自回归 Transformer 与隐扩散模型两条主流生成范式，并开源了覆盖 13 种数据组合的 checkpoint 矩阵。
- 改变了什么：把”视觉分词器”从生成范式（AR vs. 扩散）和输入模态（图像 vs. 视频）的双重绑定中解耦出来，为后续世界模型/视频生成工作提供了一个可以直接拿来用、无需为图像或视频单独重训的通用分词底座（相对于 ivideogpt 那种为具体下游任务定制的条件 VQGAN 是另一种更通用的选择）。
- 作者自陈局限：结论部分仅提出”未来将探索扩大 OmniTokenizer 模型容量以获得更强分词性能”——即论文本身未做规模化（scaling law）实验，尽管其动机部分援引了 LM/ViT 的 scaling law 文献；fetch 到的 v1 版本正文中，题为”Conclusion and Discussion of Broader Impact”的章节实际只谈了这一条扩展方向,未见展开的社会影响/滥用风险讨论。DiT 扩散实验里换用 OmniTokenizer-VAE 后 FID/IS 相对原版 SD-VAE 基线略有下降，是论文自身数据里唯一”未反超基线”的实测结果。
原始链接
- 论文（arXiv abs，v1，2024-06）：https://arxiv.org/abs/2406.09399
- PDF：https://arxiv.org/pdf/2406.09399
- 代码：https://github.com/FoundationVision/OmniTokenizer
- 预训练权重（HuggingFace，第三方账号 Daniel0724 托管）：https://huggingface.co/Daniel0724/OmniTokenizer
- 项目主页（GitHub README 中链接，抓取时连接失败，未能核实内容）：https://www.wangjunke.info/OmniTokenizer/
一手源存档（sources/）
- omnitokenizer—github-readme — GitHub README（模型库表、tokenizer/LM/扩散训练与评测说明）
- omnitokenizer—hf-card — HuggingFace 模型仓库 README（与 GitHub README 内容一致，确认 checkpoint 命名与托管账号）
- omnitokenizer—train-script — 官方参考训练脚本
scripts/recons/train.sh全文（三阶段超参数，印证正文 Implementation Details 中的架构/训练配置数字） - arXiv 2406.09399 原文（HTML v1 通读，PDF 不入 git）：https://arxiv.org/abs/2406.09399