AI Research 技术调研

MAGVIT-v2 — Language Model Beats Diffusion: Tokenizer is Key to Visual Generation

15分钟阅读

一句话定位

MAGVIT-v2 是一个统一图像/视频的离散视觉 tokenizer，核心创新是 Lookup-Free Quantization（LFQ，查表无关量化）——把 VQ-VAE 码本嵌入维度降到 0、用符号函数二值化，从而把词表扩到 2^18 ≈ 262K 而生成质量仍随词表单调提升。配上它的 token，掩码语言模型（MLM）在 ImageNet 与 Kinetics 上首次反超扩散模型：ImageNet 512×512 达 有引导 FID 1.91（超扩散最好 VDM++ 的 2.65，↓28%），且用更小模型（307M）、更少采样步（64 vs 250）。论文中心命题：「语言模型生成不如扩散，问题不在模型而在 tokenizer」。ICLR 2024。

背景与定位

2023 年的共识是：用 LM 做视觉生成（把图像 token 化后当「词」自回归/掩码预测）质量明显逊于扩散模型。论文给出的标尺差距是——ImageNet 256×256 上最好的 LM（Contextual RQ-Transformer，FID 3.41）比当时最好的扩散（MDT，FID 1.79）落后约 48%。

作者的诊断不是「LM 架构弱」，而是「缺少一个像自然语言那样好的离散视觉表示」。这把矛头指向 visual tokenizer。技术脉络上，本工作直接承接 vq-vae → taming-transformers-vqgan → MAGVIT(v1, CVPR 2023, 同组) 这条离散 token 化主线，并与 maskgit（掩码图像生成）、Muse、Phenaki 同源。相对扩散侧（ddpm→latent-diffusion-ldm→dit-scalable-diffusion-transformers），论文刻意只在「同数据、同量级模型、同训练预算」的受控条件下，用 ImageNet/Kinetics 这类公开 benchmark 做科学对比，回避了 t2i 大模型因私有数据无法公平比较的问题。

核心痛点与前作 MAGVIT(v1) 的两个缺陷：(1) v1 用普通 3D CNN，因时间感受野无法 token 化单张图像，做不到图像/视频共享词表；(2) 长视频有明显闪烁。MAGVIT-v2 用 LFQ + 因果架构改造同时解决这两点。

模型架构

整体框架仍是 VQ-VAE（CNN 编码器 + 量化瓶颈 + CNN 解码器 + GAN/感知损失，即 VQGAN 范式），但量化器与卷积时序结构被重新设计。

1) Lookup-Free Quantizer（LFQ）— 最核心创新

  • 把码本嵌入 C ∈ R^{K×d} 的维度 d 降为 0，码本退化成一个整数集合 C（|C|=K），不再需要在 K 个 d 维向量里查最近邻。
  • 论文采用最简变体：独立维度 + 二值码本。潜空间分解为 log₂K 个一维变量的笛卡尔积；每维码本取 {-1, +1}，量化就是符号函数 q(z_i)=sign(z_i)。token 索引由各维符号按二进制位组合得到（Index = Σ 2^{i-1}·1{z_i>0}）。
  • 直觉：限制单个 token 的表示容量，反而更利于在大词表分布上学习。结果（Fig.1）是 LFQ 的重建 FID 与生成 FID 都随词表增大单调下降，而传统 VQ 的生成 FID 在词表变大后反而恶化——这是 LFQ 让大词表「可用」的关键证据。

2) 因果 3D CNN（统一图像/视频的架构）

  • 对比了三种因果架构：C-ViViT（空间 transformer + 因果时序 transformer）、C-ViViT+MAGVIT（3D CNN 块 + 因果 transformer）、以及本文采用的 Causal 3D CNN。
  • 因果 3D 卷积只在时间维前向 padding（前补 kt-1 帧、后不补），使每帧输出只依赖历史帧；这样第一帧独立于其它帧 → 可单独 token 化为图像，实现图像/视频共享词表。时间下采样按 1+s·t → 1+t 映射，上采样后丢弃前 s-1 帧。消融（Tab.5a）证明 Causal 3D CNN 最好（UCF FVD 96.33，远优于 C-ViViT 的 437.54、C-ViViT+MAGVIT 的 316.70）。

3) 其它架构改造

  • 编码器下采样：average pooling → strided conv（用可学习核）；解码器上采样：nearest+conv → depth-to-space。
  • 延后时间下采样：从编码器前几个 block 推迟到后几个 block。
  • 判别器下采样改用 3D blur pooling 增强平移不变性。
  • 解码器每个分辨率残差块前加 Adaptive GroupNorm，把量化潜变量当控制信号注入（借鉴 StyleGAN）。

4) Token 因式分解（便于小 transformer 在大词表上预测）

  • 2^18 词表对 ~300M 的 transformer 太大，于是把 LFQ 潜空间切成等大子空间：用两个各 2^9=512 的码本拼接来代替单个 2^18 码本，每个子空间 token 单独 embedding 后求和作为输入 embedding；输出端用两个预测头分别出 logits。配合 weight tying（embedding 与 softmax 权重共享）。

用 tokenizer 的下游生成模型：主要用 MLM（掩码语言模型，BERT 式，迭代非自回归解码），复用 MAGVIT(v1) 的 transformer（约 300M）；附录也验证了 AR-LM（自回归）同样能借此 tokenizer 超过 MAGVIT。注意：本文不含 text encoder（无 t2i/t2v 文本条件），条件仅为 class label / 帧前缀。

tokenizer 规格（附录）：视频输入 17 帧 / 128×128；base channels 128；VQVAE channel multipliers 1,2,2,4；判别器 multipliers 2,4,4,4,4；残差块 4；潜形状 5×16×16；词表 2^18。图像版另设 16× 与 32× 下采样两档（分别用于 256² 与 512² 生成，均压成 16×16 token）。

数据

本文是 tokenizer/方法论文，刻意只用公开学术数据集做受控对比，不涉及网络规模 t2i/t2v 数据。

  • 图像生成 / 图像 tokenizer：ImageNet（train set，270 epoch 训 tokenizer，1080 epoch 训 MLM）。预处理：随机裁 80–100% 保持长宽比 + 随机水平翻转；10% batch 丢弃 class label 以支持 classifier-free guidance。
  • 视频生成：Kinetics-600（K600，帧预测，190 epoch 训 tokenizer / 360 epoch 训 MLM）、UCF-101（类条件生成）。
  • 视频压缩评测：MCL-JCV（30 段视频，缩到 640×360）。
  • 视频理解评测：Kinetics-400、Something-Something-v2（SSv2）。
  • 定性重建对比里另用了「web images」训练的一版 tokenizer（来自 PaLI 数据，Chen et al. 2022），但生成数字均来自 ImageNet 训练版。
  • 未涉及：re-captioning、合成数据、美学/安全过滤、图文对——因为这是离散重建型 tokenizer + class-conditional 生成，不需要文本标注。

训练方法

  • Tokenizer 训练目标：标准 VQGAN 损失组合——重建 + GAN（非饱和生成器损失）+ 感知（perceptual）+ commitment，因 LFQ 无嵌入故去掉 codebook loss；外加 熵惩罚 L_entropy = E[H(q(z))] − H[E(q(z))]（鼓励码本利用率，借鉴 MaskGIT/VQGAN）。用 LeCAM 正则稳住 GAN 训练。
  • 关键超参（视频 tokenizer，附录 A.2）：熵损失权重 0.1（退火：起点 ×3、2000 步内线性衰减到 0.1）；重建损失权重 5.0；生成器对抗损失权重 0.1；判别器梯度惩罚 r1 cost 10；感知损失 0.1；commitment 0.25；LeCAM 0.001；peak LR 1e-4（linear warmup + cosine decay）；Adam β1=0、β2=0.99；EMA 衰减 0.999；batch size 256。
  • 初始化（inflation）：视频 tokenizer 由 128×128 上训好的 2D 图像 tokenizer 中心 inflation 得到（与 v1 不同处：填充时间最后一片以匹配因果 padding）；判别器不 inflate、从头训以求稳定。
  • 下游 MLM 训练：复用 MAGVIT(v1) 的掩码 token 目标 + 因式分解 + weight tying。图像 MLM 训 1080 epoch（对齐 MDT），batch 1024。
  • 解码（推理）：MLM 非自回归迭代解码（cosine masking schedule）。无引导用 temperature 30（512²）/15（256²）；有引导借用 MDT 的 guidance schedule + temperature scaling（guidance scale 25、temperature 15）。512² 仅 64 步、256² 也 64 步；视频用 temperature 32。
  • 未使用：diffusion / flow matching / 步数蒸馏 / RLHF / DPO / reward model——这些都不适用于「离散 tokenizer + 掩码 LM」路线。

Infra（训练 / 推理工程）

  • 训练/推理在 TPU 上（论文压缩章节明确提到「在 TPU 上有不错结果，但要像标准编解码器那样在 CPU 上高效运行尚需后续研究」）。
  • 未披露：具体 TPU 数量、芯片型号、总 GPU·时/TPU·时、并行策略、吞吐量、训练墙钟时间。
  • 推理效率优势体现在采样步数：MLM 512² 用 64 步 vs 扩散 250–2000 步；视频帧预测 24 步 vs 扩散 256–1000 步——这是 token+掩码 LM 相对扩散的工程卖点。
  • 论文强调 token 表示「天然兼容 LLM 生态」（FlashAttention、MoE、缩放配方、TPU/GPU 优化等都可直接复用），并提出视觉 token 可作为新视频压缩格式直接喂生成模型、省去解压步骤（利好边缘计算）。

评测 benchmark（把效果讲清楚）

数字均来自已落盘的 ICLR 2024 论文（正文 Tab.1–4 + 附录 Tab.7–8）。

图像生成 · ImageNet 512×512（Tab.2，ADM 协议）

  • MAGVIT-v2（MLM+LFQ，307M，64 步）：无引导 FID 3.07 / IS 213.1；有引导 FID 1.91 / IS 324.3。
  • 对比扩散最好 VDM++（2B，512 步）：无引导 FID 2.99、有引导 FID 2.65。有引导下 1.91 vs 2.65，↓28%（论文正文明确给出此对比），且参数量约为其 1/6、步数 1/8。
  • 也优于 simple diffusion（无引导 3.54 → 有引导 3.02）、DPC+Upsample（MLM+VQ，有引导 2.65 / IS 278.1，619M/72 步）。

图像生成 · ImageNet 256×256（附录 Tab.7）

  • MAGVIT-v2（307M，64 步）：无引导 FID 3.65 / IS 200.5；有引导 FID 1.78 / IS 319.4。
  • 此前最好 LM 是 Contextual RQ-Transformer（无引导 FID 3.41，1.4B/72 步）；扩散侧无引导最好 VDM++（2.40，2B/512 步），有引导最好 MDT（FID 1.79，676M/250 步）。MAGVIT-v2 有引导 FID 1.78 险胜 MDT 的 1.79，成为该 benchmark 上首个超扩散的 LM（256² 这档优势比 512² 小，但模型 307M、仅 64 步）。

视频生成（Tab.1，MAGVIT 评测协议）

  • K600 帧预测 FVD 4.3、UCF-101 类条件 FVD 58（均 307M，24 步）。
  • 大幅超前作 MAGVIT(v1)（K600 FVD 9.9 / UCF FVD 76，306M，12 步）；并显著优于自家 non-causal baseline（K600 FVD 11.6，307M）——4.3 vs 11.6 凸显因果设计对帧预测的贡献。
  • 超 Phenaki（K600 36.4）、TATS（UCF 332）、Video Diffusion（K600 16.2）、RIN（K600 10.8）等。

视频生成 · AR-LM 版（附录 Tab.8，UCF-101）

  • AR-LM + MAGVIT-v2 tokenizer：FVD 109（840M，1280 步），远好于 AR-LM + MAGVIT v1 tokenizer（FVD 265，306M）。证明 tokenizer 增益对 AR 与 MLM 两种 LM 都成立。

视频压缩（Tab.3 + Fig.6，MCL-JCV，0.0384 bpp）

  • LPIPS 0.104（优于 HEVC 0.199、VVC 0.153、MAGVIT 0.144，全场最佳）；PSNR 26.18、MS-SSIM 0.894（这两项不及标准编解码器，因 neural 模型在等码率下牺牲局部细节换感知质量）。
  • 16 名评审、人均约 800 对偏好题、Elo 评分：在多个码率下 raters 更偏好 MAGVIT-v2，质量超 HEVC(H.265)、与下一代 VVC(H.266) 持平——「视觉生成用 tokenizer 首次达到标准编解码器水平」。

视频理解 · 动作识别（Tab.4，分类准确率）

  • 作为掩码建模目标：SSv2 62.40 / K400 75.34 / K600 77.93，均超 MAGVIT v1（57.34 / 72.29 / 74.65），接近 raw pixel（63.08 / 76.13 / 78.92）。

关键消融

  • LFQ vs VQ + 词表扩展（Fig.1）：LFQ 让生成 FID 随词表单调降，VQ 不能——这是全文最核心结论。
  • ImageNet 128² 逐项消融（Tab.5b）：MAGVIT 基线 FID 2.65 → +LFQ 2.48 → +大词表 1.34 → +up/downsampler 1.21 → +更深 1.20 → +adaptive norm 1.15。「LFQ+大词表」一步带来最大跳变（2.48→1.34）。
  • UCF-101 视频 tokenizer 消融（Tab.5c）：24.55 →(LFQ+大词表) 16.12 → 15.37 →(延后时间下采样) 11.11 →(更深) 8.90 →(3D blur pooling) 8.62。

创新点与影响

核心贡献

  1. LFQ：首个让大词表（262K）对 LM 生成「有益而非有害」的视觉量化方法，把「重建好≠生成好」的错配关系打通。
  2. 统一图像/视频 tokenizer：因果 3D CNN 让单图与视频共享一套词表，解决 v1 无法 token 化图像、长视频闪烁的问题。
  3. 首个证据：在同数据/同量级/同预算下，LM（掩码或自回归）在 ImageNet/Kinetics 上超过 SOTA 扩散——把「LM 视觉生成弱」从「模型问题」重新定义为「tokenizer 问题」。
  4. 神经视频压缩：人评质量首次追平 VVC 这类标准编解码器。

影响

  • 直接成为 VideoPoet（Google 的零样本视频 LLM）的核心 tokenizer（项目页明确点名）。
  • LFQ 成为离散视觉生成的标配组件之一，催生开源复现 Open-MAGVIT2（LFQ tokenizer + AR 生成）等后续工作，并被大量统一多模态/视觉自回归研究引用为「大词表离散 token 化」的基线方案。
  • 为「视觉与语言共享同一 token 空间 → 真正多模态 LLM」提供了关键的表示侧支撑。

已知局限

  • 本文无 t2i/t2v：作者明确说 text-to-image「可能超出本文范围」、t2v 模型因算力/时间未能放出结果（附录自陈），文本条件能力留给 VideoPoet 等后续。
  • 压缩在 PSNR/MS-SSIM（保真度指标）上不及标准编解码器，且 TPU 上跑、未解决 CPU 高效运行，离实用编解码器尚远。
  • Infra 细节（算力、并行、吞吐）几乎未披露。
  • LFQ 只验证了「独立维度 + 二值」最简形式，多值码本等变体留作 future work。

原始链接

一手源存档（sources/）

关系图谱

反向链接