Image and Video Tokenization with Binary Spherical Quantization 论文阅读

摘要

我们提出了一种基于Transformer的新型图像与视频分词器，采用了二值球面量化（Binary Spherical Quantization, BSQ）方法。BSQ 首先将高维的视觉嵌入投影到低维的超球面上，然后进行二值化量化。BSQ 具有以下优点：（1）无需显式码本，因此参数高效；（2）能够扩展至任意 token 维度，具有良好的可扩展性；（3）表示紧凑：能以极低失真的情况下将视觉数据压缩至原始数据的 1/100。

我们的分词器采用了 Transformer 编码器和解码器，并配合简单的分块式因果掩码，从而支持变长视频作为输入。最终形成的 BSQ-ViT 在图像与视频重建基准测试中，实现了当前最优的可视化重建质量，并且吞吐量比以往最佳方法高出 2.4 倍。此外，通过学习自回归先验实现自适应算术编码，BSQ-ViT 在视频压缩任务上取得了与现有最先进视频压缩标准相当的结果。BSQ-ViT 还使得掩码语言模型在图像合成质量上能够与基于 GAN 和扩散模型的方法竞争。

引言

当然，以下是你提供的英文段落的详细中文翻译：

---

1 引言

通过学习离散的图像和视频分词，可以实现最先进的视觉压缩[1, 2, 3]、识别[4, 5, 6, 7]和生成[8, 9, 10]。这些模型遵循了大型语言建模[11, 12, 13]的成熟范式：将输入和输出分词为离散单元，并通过自回归模型逐个预测这些token化的序列。当前，最广泛应用的图像编码方法是向量量化变分自编码器（VQ-VAE）[8]。该方法先将输入编码为连续的潜在嵌入，再通过最近邻查找映射到学习得到的码本。然而，VQ-VAE 风格的方法存在两个缺点：首先，大多数图像编码器建立在卷积神经网络（CNN）[9, 14]的基础上。将图像中的空间卷积扩展到视频中的时空卷积，需要复杂的结构调整[15, 16, 17]，并增加了计算成本。若将视频视为一系列图像进行处理，则会导致次优的量化效果[16]。其次，向量量化（VQ）在码本尺寸较大时扩展性较差。其运行时间与码本大小线性相关，且在小数据集上码本容易过拟合[17]。这一问题在处理视频时尤为突出，因为视频需要更大的码本来同时表示静态视觉模式和动态运动模式。

本论文提出了一种基于视觉Transformer和二值球面量化（BSQ）的统一视觉分词器。该Transformer编码器-解码器利用分块式因果掩码，仅使用当前或过去时刻的视觉token进行重建（见图3）。BSQ 首先将 Transformer 编码器生成的高维视觉嵌入投影到低维的超球面上，然后进行二值量化。Transformer 编码器、解码器和 BSQ 可无缝集成到 VQ-GAN [9] 框架中，并实现端到端训练。

我们提出的视觉分词器具有如下优势。首先，基于 Transformer 的编码器-解码器在视觉重建质量和计算效率方面，相较于标准 CNN 实现了帕累托改进。其次，分块式因果设计在训练时统一处理图像和视频输入，并在推理阶段支持变长视频。BSQ 构建了一个隐式码本，其有效词汇量随球面维度呈指数增长，而无需学习参数。码本规模的提升持续带来更优的重建效果。与近期另一种基于标量量化（SQ）构建隐式码本的方法——无查找量化（LFQ）[17]相比，BSQ 具有有界的量化误差，并且更易于训练。此外，我们展示了 BSQ 中的软量化概率可以简化为多个通道独立的伯努利分布的乘积，从而在训练中实现高效的熵正则化。具体而言，我们证明，对于 L 位的软量化，其熵的分解近似方法能将理论计算复杂度从 O(2^L × L) 降低到 O(L)，且近似误差极小，实际性能几乎无损。

我们在视觉重建和压缩基准上验证了 BSQ-ViT 的有效性。在图像重建任务上，我们的模型在像素级和语义级度量下均实现了当前最优的视觉重建质量。尤其是，表现最好的 BSQ-ViT 在 ImageNet-1k 验证集上取得了 0.41 的重建 FID，相较次优方法（SDXL-VAE [14]）降低了 43%，速度提升了 2.4 倍。在视频重建任务上，我们的最佳模型将 UCF-101 上的 FVD 指标从 8.62 降至 4.10。通过进一步学习自回归先验以实现自适应算术编码，BSQ-ViT 在视频压缩任务上的表现亦可与当前最先进的视频压缩标准（如 H.264 和 HEVC）相媲美。通过学习掩码语言模型，BSQ-ViT 能够生成与 BigGAN [18] 和 ADM [19] 等方法相当的高质量图像。代码和模型将在 https://github.com/zhaoyue-zephyrus/bsq-vit 上发布。

---

以下是你提供的第二部分“Related Work”的详细中文翻译：

---

2 相关工作

视觉分词（Visual Tokenization）
VQ-VAE [8] 首次在自编码器架构中引入了离散分词瓶颈（discrete tokenized bottlenecks）的概念。近期的改进包括更优的训练目标[20, 9]、提升 VQ 码本的利用率[4, 21]、用乘积量化（PQ）[3]或标量量化（SQ）[22]替代 VQ，以及采用更强的生成模型[9, 10]。图像分词器可以通过对每一帧单独分词简单地扩展到视频场景[23, 24]，但这种方式忽略了动态运动，会导致次优的分词结果：相同的视觉信息会在多帧中被重复压缩。

视频分词（Video Tokenization）
专门的视频分词器能更好地利用输入信号的时序相关性。VideoGPT [25] 在 VQ-VAE 框架中提出了用于视频生成的 3D 卷积与反卷积。TATS [15] 用重复填充（replicate padding）取代了零填充，以缓解视频长度变化时的时序破坏。Yu 等人提出将预训练的 2D 卷积滤波器中心扩充为 3D 卷积[16]，并进一步将其因果化[17]。Phenaki [23] 采用了分解式因果视频视觉 Transformer [26]（C-ViViT），提升了效率，但在一定程度上牺牲了对复杂时序运动的建模能力。

神经压缩（Neural Compression）
自从 Shannon 建立了基本的信息源编码定理[27]以来，基于概率模型的无损压缩（lossless compression）成为主流方法，包括 RNN[32, 33]、CNN[34, 8]、VAE[35, 36] 和 Transformer[37, 38]。L3C [39] 提出了用于无损图像压缩的快速分层概率模型。LMIC [38] 显示，主要基于文本训练的大型语言模型（如 Llama 2[13] 和 Chinchilla[40]）是文本、图像、音频的通用压缩器。然而，这些大模型过于庞大和缓慢，难以实际应用于压缩场景。我们的分词器提供了一种更轻量的替代方案：分词器首先进行局部有损压缩，然后利用约 3亿参数的轻量级、高效序列模型对全局视频结构进行压缩。

视频压缩（Video Compression）
绝大多数高性能的现代视频压缩方法依赖于混合编码器，将变换编码[41, 42]和运动补偿[43, 44]结合起来。这一思路在当前许多主流的基于学习的方法中依然主导[45, 46, 47, 48]。VCT [49] 提出了一种基于 Transformer 的时序熵模型，用于隐式学习运动信息。然而，VCT 需要高度工程化的图像压缩模型[50]，且其时序上下文窗口较短。在本研究中，我们展示了经过学习的视频分词器结合基于序列模型的算术编码器，无需显式建模运动信息，也能达到有竞争力的视频压缩性能。

---

以下是你提供的第3部分“Preliminaries”的详细中文翻译：

---

3 预备知识

基于分词的压缩算法通常包括三个基本步骤：首先，视觉分词器（如 VQ-VAE[8] 或 LFQ[17]）将原始视觉输入转化为一组离散的 token，并能实现重建。接着，序列模型对这些离散 token 建立自回归概率分布进行预测。最后，算术编码器（Arithmetic Coding）将这一分布转化为压缩后的表示。

视觉分词
VQ-VAE[8] 首次在自编码器架构中引入了基于向量量化（VQ）的离散视觉表示学习。假设有一个视频 (X \in \mathbb{R}^{T \times H \times W \times 3})，编码器 (E) 会输出一组 (d) 维的潜在嵌入 (Z = E(X) \in \mathbb{R}^{(\frac{T}{q} \times \frac{H}{p} \times \frac{W}{p}) \times d})，其中空间和时间下采样因子为 (q \times p \times p)。瓶颈模块 (q) 将实值潜在嵌入进一步转化为离散 token (\hat{z} = q(z))。

在向量量化（VQ）中，量化器 (q_{VQ}) 会把每个 (z \in Z) 映射到码本 (C = [c_1, \cdots, c_K] \in \mathbb{R}^{K \times d}) 中距离最近的一个码字：
[
\hat{z} = q_{VQ}(z) = c_k = \arg\min_{\hat{c}k \in C} |z - \hat{c}k|2
]
其中 (K) 是码本的词汇量，整数 (k) 表示 (z) 的离散化 token，可以用 (\lceil \log(K) \rceil) bit 存储。解码器 (G) 能将离散 token 映射回视觉表示 (\hat{X} = G(\hat{Z}))。整个网络（即 (E, G, q)）是端到端可训练的，通过直通估计器（straight-through estimator）[51] 让梯度能够传递过量化瓶颈，并最小化 MSE 损失 (L{MSE} = |\hat{X} - X|^2)。近年的方法还引入了感知损失 LLPIPS 和对抗损失 LGAN 以提升重建质量[9]：
[
\min{E, G, q} \mathbb{E}X[L{VQ}(E, G, q) + \eta L{LPIPS}(E, G, q) + \lambda L_{GAN}(E, G, q)]
]
其中，量化损失项 (L_{VQ}) 通过在线聚类方式学习 (c_k)。VQ-VAE 的主要问题是码本词汇量 (K) 增大时，向量量化扩展性很差[17]。解决方法包括：使用更低的码字维度[4]，引入随机性[52]，激活沉睡码字[21]，以及用承诺损失（commitment loss）正则化[8]：
[
L_{commit}(\hat{z}, z) = |sg(\hat{z}) - z|
]
其中 (sg(\cdot)) 表示停止梯度操作。

无查找量化（LFQ）
无查找量化（Lookup-Free Quantization, LFQ）[17] 使用固定的隐式码本 (C_{LFQ} = {-1, 1}^L)，即 L 维空间的超立方体顶点。对该隐式码本的最佳二值量化器为 (q_{LFQ}(z) = \operatorname{sign}(z))。Yu 等人[17] 还引入了熵目标以促使潜在码的有效性及码本利用[53]：
[
L_{entropy} = \mathbb{E}[H(q(z))] - \gamma H[\mathbb{E}[q(z)]]
]
这两个熵项都依赖于软量化[2]，以确保损失的可微性：
[
\hat{q}(c|z) = \frac{\exp(-\tau(c-z)^2)}{\sum\limits_{c \in C_{LFQ}} \exp(-\tau(c-z)^2)}
]
最终损失 (L_{LFQ}) 包含 (L_{MSE}, L_{commit}, L_{LPIPS}, L_{GAN}) 及 (L_{entropy}) 多项。LFQ 的主要计算瓶颈在于高维码本的熵优化，因为它涉及对 (2^L) 个隐式码字求和。

无论是 VQ-VAE 还是 LFQ，都可以将视觉输入 (X) 有损地压缩为 N 个离散 token ([k_1, …, k_N])，其中 (k_i \in {1,…,K})，共需 (N\lceil \log K \rceil) 位。但这两种分词策略都未能充分利用图像或视频的全局结构。而结合无损算术编码的序列模型能更好地拟合这类全局结构。

算术编码（Arithmetic Coding, AC）
算术编码[29, 30, 54]利用编码分布的统计特性，可生成长度近乎最优的比特流。给定对 token 序列的分布 (P_t: {1, \cdots, K}^n \rightarrow (0,1])，算术编码试图将 token 序列编码为 (-\lceil \log P_t(k_1, …, k_N) \rceil + 1) 位。最常见的 token 分布是假设自回归模型：
[
P_t(k_1, …, k_N) = P_t(k_1)P_t(k_2|k_1)…P_t(k_N|k_1,…,k_{N-1})
]
对该模型已有高效的增量式编码和解码算法[49]。

---

下面是你提供的文本的详细翻译：

---

4 基于Transformer的视频分词器与二值球面量化

我们的视频分词器采用编码器-解码器结构，中间有一个离散化瓶颈，如图1a所示。它结合了基于Transformer的编码器、基于Transformer的解码器和一个二值球面量化（BSQ）层。BSQ将潜在编码映射到一个低维球面空间，进行二值量化，然后再将结果投影回解码器的潜在空间。这种低维球面空间投影有若干理论优势：量化器的近似误差有界，并且熵计算在各个维度上基本独立。这些优势在实验中也带来了提升。BSQ比其他量化方案收敛更快、效果更好。

4.1 二值球面量化（BSQ）

二值球面量化（BSQ）优化于一个隐式码本 ( C_{BSQ} = { -\frac{1}{\sqrt{L}}, +\frac{1}{\sqrt{L}} }^L )，即将一个超立方体的所有顶点投影到单位球面上。每个超立方体的顶点 ( c_k \in C_{BSQ} ) 对应一个唯一的token。量化器的工作流程如下：首先将高维潜在嵌入z投影到低维单位球面u上，对每个坐标轴进行二值量化 ( \hat{u} = sign(u) )，然后再将结果投影回原始潜在空间的量化向量 ( \hat{x} )，如图1a所示。

具体来说，我们从编码后的视觉输入 ( z=E(x) \in \mathbb{R}^d ) 出发，首先用线性投影将潜在嵌入降维到L维，得到 ( v = Linear(z) \in \mathbb{R}^L )（其中L远小于d）。然后将v投影到单位球面 ( u = \frac{v}{|v|} )。接着对u的每一维独立地进行二值量化： ( \hat{u} = \frac{1}{\sqrt{L}}sign(u) ) ，其中sign(x)为符号函数，对于0则映射为1。为保证操作可微，我们采用直通估计器（STE）：
( sign_{STE}(x) = sg(sign(x) - x) + x )，其中sg(·)为停止梯度操作。
最后将量化后的 ( \hat{u} ) 线性回投影到d维空间，得到 ( \hat{z} = Linear(\hat{u}) \in \mathbb{R}^d )。

BSQ有以下优点：与LFQ类似，其隐式码本无需参数、易于计算。但与LFQ不同，BSQ的软量化有简洁的概率解释，从而可高效计算熵损失，且量化误差有界，这在实验中带来了更快、更优的收敛表现。

高效的隐式编码分配
推理时，我们只需将投影后的嵌入v按如下方式二值化成token index：
( k = \sum_{i=1}^L 1[v_i > 0] 2^{i-1} )，其中1[·]是指示函数。反向映射用位移和按位与操作即可。

软BSQ与熵
为了充分利用隐式码本 ( C_{BSQ} ) 的全部空间，我们引入熵损失 ( L_{entropy} = \mathbb{E}u[H(q(u))] - \gamma H[\mathbb{E}u[q(u)]] )。计算熵损失前，先推导一种软量化方案。由于码本和量化输入都是单位向量，软量化分布可写为
[
\hat{q}(c|u) = \frac{\exp(\tau c^\top u)}{\sum{c′\in C{BSQ}} \exp(\tau c′^\top u)} = \prod_{d=1}^L \sigma(2\tau c_d u_d)
]
（详见C.1节）。这种分解形式，使熵的第一项可以高效计算：
[
\mathbb{E}u[H(\hat{q}(c|u))] = \mathbb{E}u \left[\sum{d=1}^L H(\hat{q}(c_d|u_d))\right]
]
熵第二项无法直接假设各维独立，因为平均值 (\mathbb{E}u[\hat{q}(c|u)]) 在各维有相关性。我们找到最接近的分解分布 (\tilde{q}© = \prod{d=1}^K \tilde{q}(c_d))，并最小化其熵（KL距离最优时 (\tilde{q}(c_d) = \mathbb{E}{u_d}[\hat{q}(c_d|u_d)])），最终近似熵项为：
[
H(\mathbb{E}u[\hat{q}(c|u)]) \approx H(\tilde{q}©) = \sum{d=1}^L H(\mathbb{E}_{u_d}[\hat{q}(c_d|u_d)])
]
（详见C.3节）。这一近似是真实熵的上界，且实验证明误差很小，因此该熵项评价高效。

BSQ的量化误差
大多数量化器在训练中用直通梯度估计，假设量化前后的梯度基本一致，这成立的前提是量化误差 ( d(u, \hat{u}) = |u-\hat{u}| ) 足够小。按C.4节推导，BSQ的期望量化误差有界：
[
\mathbb{E}_u[d(u, \hat{u})] < 2 - 2/\sqrt{L} < \sqrt{2}
]

与其它量化方法的关系
BSQ和信息论、编码理论中的诸多方法密切相关。LFQ与BSQ同样用二值化技术，但未归一化输出，导致量化误差无界，也无法像BSQ那样简单地实现软量化和熵计算。球面向量量化（SVQ）也保证所有码字有固定半径，但SVQ允许半径多样，需要额外的增益编码器；而BSQ仅需单位半径，避免了增益编码。金字塔向量量化（PVQ）要求所有码字l1范数相同，但其在球面上的分布不如l2归一化均匀。

---

4.2 含因果视频Transformer的分词网络

我们建议用Vision Transformer（ViT）同时做编码器和解码器，因为其计算效率高、重建质量好。

视频Transformer
我们以ViT-VQGAN为基础，将其扩展为视频输入。输入视频 ( X \in \mathbb{R}^{T \times H \times W \times 3} ) 被分为不重叠的patch，大小为 ( 1 \times p \times p )，每个patch为 ( x_i \in \mathbb{R}^{1 \times p \times p \times 3} )。视觉token被拉平成一维序列，通过线性投影和N层Transformer编码器生成潜在表示 ((z_1, \ldots, z_N))。解码器结构相同，将潜在嵌入 ( \hat{z} ) 映射回像素空间，并重组成原始形状 ( (\hat{x}_1, \ldots, \hat{x}_N) = \text{MLP}(\text{TransformerDecoder}(\hat{z}_1, …, \hat{z}_N)) )；其中MLP为两层结构（Linear-Tanh-Linear）。

分块因果注意力
训练时输入视频帧数为T，但推理时视频长度可能变化。补零到T帧虽可行但浪费比特，尤其在压缩场景下。为此我们提出分块因果掩码注意力，类似语言建模中的因果注意力：只利用时刻t及之前的token重建t时刻的视觉token。
公式如下：

• 编码器输出：
  ( (z_{(t-1) \times H/p \times W/p+1}, …, z_{t \times H/p \times W/p}) = \text{TransformerEncoder}(x_1, …, x_{t \times H/p \times W/p}) )
• 量化输出：
  ( (\hat{z}{(t-1) \times H/p \times W/p+1}, …, \hat{z}{t \times H/p \times W/p}) = q_{LFQ}(\cdots) )
• 解码器输出：
  ( (\hat{x}{(t-1) \times H/p \times W/p+1}, …, \hat{x}{t \times H/p \times W/p}) = \text{MLP}(\text{TransformerDecoder}(\hat{z}1, …, \hat{z}{t \times H/p \times W/p})) )

这种分块因果注意力可以用下三角的block mask高效实现。T=1时，模型即为全注意力的ViT，因此可以用图像和视频混合训练。

空间-时间位置编码
我们用分解的时空位置编码。具体做法是，给每个时间帧加一个零初始化的时间位置编码 ( PE_t \in \mathbb{R}^{T \times d} )，加到原有空间位置编码 ( PE_s \in \mathbb{R}^{N \times d} ) 上，即 ( PE[i,:,:]=PE_t[i, None, :]+PE_s[None, :, :] )。

从图像分词器训练视频分词器
由于现有视频集多样性不足，我们先在图像数据上训练图像分词器，再微调为视频分词器。虽然之前有工作认为图像分词器可直接用于视频，但我们发现微调后的视频分词器在视频基准上重建质量更高（见5.1节），尤其是词表变大时这一提升更明显。我们认为BSQ带来的大词表对学习视频特有的运动和模糊很有帮助。而普通VQ方法在码本超过16K时已无法维持高码本利用率。

优化视觉分词器
我们采用VQGAN的知觉损失和对抗损失。判别器用StyleGAN而非PatchGAN（ViT-VQGAN实验发现StyleGAN更易训练）。视频微调时，我们不像MAGVIT或TATS那样将StyleGAN扩展为3D判别器，而是将所有重构帧分别输入原版StyleGAN，损失求和即可。

---

如需进一步解释原理、公式细节或代码实现，可以继续追问！