[论文速览] Language Model Beats Diffusion - Tokenizer is Key to Visual Generation

Pre

title: Language Model Beats Diffusion - Tokenizer is Key to Visual Generation
accepted: ICLR 2024
paper: https://arxiv.org/abs/2310.05737
code: non-official
ref: https://mp.weixin.qq.com/s/oYlV93VHFfTjzYKUIld8JA

关键词： tokenizer, language model, diffusion
阅读理由： tokenizer

Idea

将 VQ-VAE codebook 中每个向量的维度降为0，即码表变成了一个整数集，通过将向量z每一个元素二值化[-1,1]来量化，对应整数集里一个整数

Motivation&Solution

1. LM生成文本不错，但在图像和视频上不如Diffusion —— 换个tokenizer

Background

使用离散视觉token有如下好处：

1. 与LLM兼容
2. 压缩表征，减少存储和带宽
3. 有利视觉理解。离散token在自监督学习里是有价值的预训练目标，同时用token作输入还能改进鲁棒性和泛化性

Method（Model）

Overview

该tokenizer(MAGVIT-v2)基于MAGVIT改进，而后者又是 VQ-VAE 框架的视频 tokenizer。

LOOKUP-FREE QUANTIZER

VQ-VAE模型取得巨大进展，但改进它的重建质量的与它后续生成质量之间的关系仍不明确。很多人误以为改进重建就等于改进语言模型的生成，例如扩大vocabulary可以提高重建质量。然而vocabulary并非越大越好，太大反而损害性能。

训练更大的codebook的一个简单技巧是在增加词汇量时减少编码的嵌入维度。这基于一个直觉：限制单个token的表示能力，这会促进大词汇表分布的学习。受此启发，将 VQ-VAE codebook的嵌入维数减少到零。形式上，将codebook \(\mathcal{C} \in \mathbb{R}^{K\times d}\) 替换为整数集\(\mathbb{C}\)，其中\(|\mathbb{C}|=K\)。

在VQ-VAE模型中，当计算与编码器输出最接近的codebook编码时，必须查找codebook的所有K个d维嵌入，d通常为256。而新的设计完全消除了对这种嵌入查找的需要，因此称之为无查找量化（LFQ）。

用于视觉生成的LM其词汇表约1-8k，而自然语言的词汇量则超过200k，图1也能看出随词汇量增大，VQ（VQ-VAE）的生成质量明显下降，而LFQ能继续提升

图1 ImageNet的重建和生成质量曲线 比较 Vector Quantization(VQ) 和 Lookup-Free Quantization (LFQ) 比较基于128分辨率，306-372M参数的MLM

目前有很多 LFQ 方法，但本文讨论了一种简单的变体：假设独立的codebook维度和二元隐变量，LFQ的隐空间被分解为单维度变量的笛卡尔乘积，即 \(\mathbb{C} = \times_{i=1}^{\log_{2}^K}C_{i}\)。给定一个特征向量 \(z \in \mathbb{R}^{\log_2^K}\)，其量化表示\(q(z)\)的每一维通过如下公式取得：

\[q(z_{i})=C_{i,j},\; where\ j = \arg \min_{k}\left\|z_{i}-C_{i,k}\right\|, \tag{2} \]

其中 \(C_{i,j}\) 是 \(C_i\) 里的第j个值，若 \(C_i = \{-1,1\}\) 则 \(\arg \min\) 可以由 \(sign\) 函数进行计算：

\[q(z_{i}) = sign(z_{i}) = -\mathbb{1}\left\{z_{i} \le 0\right\} + \mathbb{1}\left\{z_{i} \gt 0\right\}. \tag{3} \]

\(q(z)\)的token索引则由以下公式给出：

\[In d e x(z) = \sum_{i=1}^{\log_2 K} \arg \min_{k}\left\|z_{i}-C_{i,k}\right\| \prod_{b=0}^{i-1}\left|C_{b}\right| = \sum_{i=1}^{\log_2 K} 2^{i-1}\mathbb{1} \{z_{i} > 0\}, \tag{4} \]

其中\(|C_{b}|=1\)应该是方便公式4计算设置的，对应\(2^0 = 1\)的情况（where \(|C_{b}|=1\) sets the virtual basis）。公式4中间那坨先对i从1到\(\log_2 K\)遍历，因为每个\(C_i\)都是\(\{-1,1\}\)两种取值，要表示\(K\)种情况只需要\(\log_2 K\)个\(C_i\)。然后对于\(z\)的每一维找出最接近的\(C_i,k\)，即下标k，也只有两种情况：0或1。然后k与前面i-1个\(C_b\)的“基”相乘，然后所有i的情况累加。实际上求index过程相当于计算第三坨的二进制序列的十进制值，例如\(q(z) = \{1, -1, -1, 1, -1\}\)的index即\(0b01001 = 9\)，本质与VQ系列方法差不多，为每种可能的量化结果按顺序编号罢了。而这个索引index就对应了\(\mathbb{C}\)里的一个整数。

额外添加熵惩罚项用以鼓励codebook编码利用率：

\[\mathcal{L}_{e n t r o p y}=\mathbb{E}[H(q(z))]-H[\mathbb{E}(q(z))]. \tag{5} \]

该惩罚项受VQGAN里相似的loss启发，也被用于基于熵的聚类里。在LFQ中，考虑到维度间的独立性，重写 \(H(q(z)) = \sum_{i=1}^{\log_2 K}H(q(z_i))\)，而对于 \(K > 2^{18}\) 的情况， \(H[\mathbb{E}(q(z))]\) 可以用维度分组（sub-groups of dimensions）来近似。没看懂，为啥熵的期望要等于期望的熵？？

作者说作为第一个提出概念的，只用了二元编码本这最简单的LFQ形式，其他方式留给后面探索。

VISUAL TOKENIZER MODEL IMPROVEMENT

期望视觉tokenization能够用同一个codebook编码图片和视频，然而 MAGVIT tokenizer 使用3D CNN，由于时序感受野的限制不能处理图片。

图2 因果tokenizer架构比较。图中忽略解码器，其架构与编码器对称。

为了建立图片-视频联合tokenizer，重新思考C-ViViT，如图2a所示，它使用完全空间Transformer块联合因果时序Transformer块，效果不错但有两个缺点：

1. 不像CNN，位置嵌入使其难以处理训练时没见过的空间分辨率。
2. 实验发现3D CNN比空间Transformer更好，能生成对应patch的具有更优空间因果性的token

本文探索了两种可行的设计方案，如图 2b 将 C-ViViT 与 MAGVIT 进行结合；图 2c 使用时间因果 3D 卷积来代替常规 3D CNN，大概是讲常规3D卷积会在输入帧的前后进行pad，而因果3D卷积只在输入前进行pad，因此每帧输出只依赖于前面的帧，而第一帧总是独立于其他帧，因此可以处理单张图片。

除了使用因果 3D CNN 层之外，本文还进行了其他架构的修改，以提高 MAGVIT 性能，比如本文将编码器下采样器从平均池化更改为跨步卷积；又比如在解码器中每个分辨率的残差块之前添加一个自适应组归一化层等。

此外考虑到大词汇表的预测效率，提出将LFQ的token潜空间分解为相等的子空间，例如\(2^{18}\)大小的codebook可以用两个\(2^9\)的替代，将二者嵌入结果的和作为Transformer输入即可。

Experiment

图3 不同tokenizer的图片重建样本。这里的VQGAN是MaskGIT用的那个

Dataset

video generation: Kinetics-600 (K600), UCF-101
image generation: ImageNet
video compression: MCL-JCV
video understanding: Kinetics-400 (K400), SSv2

VISUAL GENERATION

使用MAGVIT的MLM Transformer，大码表（\(2^{18}\)），以两个子码表凑成。

表1 视频生成结果。Kinetics-600上进行帧预测和UCF-101上的类别条件生成。

图4 Kinetics-600的帧预测样本

Video generation. 本文方法大幅超越了之前的MAGVIT，二者backbone一样，说明好的tokenizer的重要性。

表2 图片生成结果。ImageNet上512分辨率的类别条件生成。guidance指的是 “classifier-free diffusion guidance” *指额外训练数据的使用。

图5 ImageNet 512分辨率的类别条件生成样本。跟其他每个工作都用一个随机样本来比较（下面一行都是本文的结果？最后一列不是比烂？）

Image generation on ImageNet. 本文模型在采样质量（FID，IS）和推理时效率（采样步数）都超越了扩散方法。MAGVIT的采样步数难道是序列长度？但为啥也能上guidance。

VIDEO COMPRESSION

视频分辨率：640ˆ360

图6 视频压缩评价？（Video compression rater）研究

表3 视频压缩指标

跟 HEVC (H.265) 和 VVC (H.266) 比较，都更优秀。也使用 LPIPS, PSNR, MS-SSIM 在 0.0384 bpp 条件下比较，是所有模型里LPIPS最低的，作者说这个指标比另两个更加客观（但也不能就此忽略吧？）。

VIDEO UNDERSTANDING

做视频理解，选了动作识别任务，有两条路子：

1. token充当Transformer的预测目标，即输出
2. token用作Transformer的输入

表4 视频动作识别表现（类准确度↑ x100）

仍然比MAGVIT好，虽然做输入的情况不如像素直接输入，但就是比MAGVIT好

ABLATION STUDY

表5 验证关键设计的消融。

图1对比了VQ，表5a对应图2的因果 3D CNN 选择，表5b,c则展示了他模型怎么一步步堆trick来涨点的。后面deeper model这种细碎的、众所周知的、非创新点的没必要特意拿出来消融吧。而且表5a和bc应该要两两搭配来测试更合理吧。

Conclusion

本文结果展示了LM在视觉生成上的巨大潜力，倡导对先进视觉tokenization方法的进一步探索。

Critique

MAGVIT-v2, LFQ

没有代码？，好像没有官方实现，没有放出预训练权重比较可惜。
这个标题好熟悉：Diffusion Models Beat GANs on Image Synthesis，真是天道有轮回