图像与视频分词

图像与视频分词

图像与视频分词把连续的视觉数据转换成离散的词元序列，让 transformer 能像处理文本一样处理它们。本文件涵盖 VQ-VAE、VQ-GAN、码本学习、DALL-E 的 dVAE、视频分词以及无查表量化。

为什么要给图像分词

• 把语言想象成一个有限的字母表：英语大约有 26 个字母，而现代语言模型把文本切成 30,000 到 100,000 个子词词元。每个句子都变成一串离散符号序列，transformer 可以一个一个地预测。而图像则生活在一个连续、高维的空间里：一张 256x256 的 RGB 图像是 \mathbb{R}^{256 \times 256 \times 3} \approx \mathbb{R}^{196{,}608} 中的一个点。如果你希望一个语言模型用与「说英语」相同的机制来「说图像」，就需要把这些连续的像素数组转换成一串可管理的、来自有限词表的离散词元序列。这种转换就是图像分词（image tokenisation）。

• 想象你是一位马赛克艺术家。你没有无穷多种色度的瓷砖，你只有一个固定的调色板，比如 8192 种不同的颜色。要把一张照片复刻成马赛克，你必须（1）决定每块瓷砖代表照片的哪个区域，（2）为每个区域挑选最接近的瓷砖颜色，（3）接受一些细节会丢失，但整体画面仍然能辨认出来。图像分词做的正是这件事：编码器把空间图块压缩成潜在向量，码本把每个向量映射到最接近的码本项，结果就是一个整数索引组成的网格，每个图块一个索引，离散模型可以处理它。

• 分词的好处有三方面。第一，它极大地压缩了图像：一张 256x256 的图像可能变成 16x16 的词元网格，序列长度从 65,536 个像素降到 256 个词元，这对开销随序列长度二次增长的注意力模型来说很友好。第二，它统一了表征：文本词元和图像词元生活在同一个离散词表里，从而一个自回归 transformer 就可以生成交错的文本与图像。第三，它施加了一个有用的瓶颈，迫使模型学习语义上有意义的编码，而不是记住像素噪声。

• 回想第 8 章中卷积网络如何从图像中抽取层次化特征图，以及第 7 章中文本分词器如何把字符串转成整数序列。图像分词正处在两者的交汇处：它用 CNN 或视觉 transformer 编码器（第 8 章）产生空间特征，再借用离散词表的思想（第 7 章）把这些特征转换成词元索引。

VQ-VAE：向量量化

• 正如我们在第 6 章所见，标准的**变分自编码器（variational autoencoder，VAE）**把输入编码成一个连续的潜在分布，再从该分布中采样并解码回重建结果。潜在空间是连续的，这让它难以喂给离散序列模型。**向量量化变分自编码器（Vector Quantised Variational Autoencoder，VQ-VAE）**由 van den Oord 等（2017）提出，它通过引入一个可学习的嵌入向量码本、把每个编码器输出「吸附」到最近的码本项上，从而用离散潜在表示替代了连续潜在表示。

• 把它想象成一个正好有 K 个贴了标签的书架的图书馆。当一本新书（编码器输出）到来时，图书管理员会把它放到与它最相似（与已有书籍，即码本向量，最相近）的那个书架上，并记下书架编号。之后要取回这本书，你只需要书架编号：那个书架上的码本项就是一个足够好的替身。这就是向量量化。

• 形式上，VQ-VAE 有三个组件：

• 一个编码器 E，把输入图像 \mathbf{x} \in \mathbb{R}^{H \times W \times 3} 映射为一个连续潜在向量组成的网格 \mathbf{z}_e = E(\mathbf{x}) \in \mathbb{R}^{h \times w \times d}，其中 h \times w 是下采样后的空间分辨率，d 是嵌入维度。

• 一个码本 \mathcal{C} = \{\mathbf{e}_1, \mathbf{e}_2, \ldots, \mathbf{e}_K\} \subset \mathbb{R}^d，包含 K 个可学习的嵌入向量。典型的码本大小从 512 到 16,384 不等。

• 一个解码器 D，从量化后的潜在表示重建图像。

• 量化步骤把空间位置 (i, j) 上的每个编码器输出 \mathbf{z}_e(\mathbf{x}) 替换为最近的码本项：

• 这是嵌入空间中的最近邻查找，与 k-means 的分配步骤完全相同（第 6 章）。索引 k^\ast 就是空间位置 (i,j) 的离散词元，整张图像被表示为一个 h \times w 的整数网格，取值来自 \{1, \ldots, K\}。

• 挑战在于 \arg\min 不可微：你无法通过一个离散选择进行反向传播。VQ-VAE 用**直通估计器（straight-through estimator）**解决这个问题：在前向传播中，解码器收到 \mathbf{z}_q（量化后的向量）；在反向传播中，重建损失对 \mathbf{z}_q 的梯度被直接复制给 \mathbf{z}_e，就好像量化步骤是恒等函数一样。可以紧凑地写成：

• 其中 \text{sg}(\cdot) 是停止梯度算子。在前向传播中它求值为 \mathbf{z}_q；在反向传播中，梯度只流过 \mathbf{z}_e 这一项。

• 完整的 VQ-VAE 损失有三项：

• **重建损失（reconstruction loss）**训练编码器和解码器忠实地复现输入。码本损失（codebook loss）（也叫 VQ 损失）把码本向量拉向编码器输出；注意 \text{sg}(\mathbf{z}_e) 意味着编码器不会从这一项收到梯度，所以它只更新码本。**承诺损失（commitment loss）**则相反：它鼓励编码器输出靠近码本向量，防止编码器从码本「跑掉」。超参数 \beta（通常为 0.25）控制码本项与承诺项之间的平衡。

• 在实践中，码本通常用**指数移动平均（exponential moving average，EMA）**而不是梯度下降来更新，这样更稳定。设 \mathbf{n}_k 为分配给码本项 k 的编码器输出计数，\mathbf{s}_k 为它们的和。EMA 更新为：

• 其中 \gamma 是衰减率（通常为 0.99）。这等价于在编码器输出上跑一个在线 k-means 算法。

码本坍缩

• VQ-VAE 一个臭名昭著的失败模式是码本坍缩（codebook collapse）（也叫索引坍缩）：模型学会只使用 K 个码本项中的一小部分，让大多数码本项「死掉」。想象一个图书馆，90% 的书架是空的，因为图书管理员总是把书送到同样的几个热门书架上。这浪费了表征能力。

• 码本坍缩之所以发生，是因为编码器、码本和解码器在训练中协同适应。如果某个项在几个批次中都没有被选中，它就会漂离编码器流形，使得被选中的可能性更低，从而形成正反馈循环。

• 有几种技术可以缓解码本坍缩：

  • 码本重置（codebook reset）：定期把死掉的项重新初始化为随机采样的编码器输出。这给死项在潜在空间的活跃区域附近一个重新开始的机会。
  • 带拉普拉斯平滑的 EMA 更新：给 \mathbf{n}_k 加一个小常数，防止任何项的计数为零，确保所有项都能收到梯度信号。
  • 承诺损失调参：增大 \beta 迫使编码器输出更紧密地聚集在码本项附近，使分配更均匀。
  • 因式化编码（factorised codes）：把码本查找分解成若干更小查找的乘积（例如两个各 \sqrt{K} 大小的码本），通过减小每次查找的有效码本规模来提升利用率。
  • 熵正则化（entropy regularisation）：加一个惩罚项，鼓励码本使用分布均匀，最大化熵 H = -\sum_k p_k \log p_k，其中 p_k 是经验分配概率。

VQ-GAN：用对抗训练获得更高保真度

• VQ-VAE 能产生不错的重建结果，但像素级 \ell_2 损失倾向于产生模糊的输出，因为它对每个像素偏差一视同仁地惩罚，于是对所有看似合理的细节取平均，而不是选择锐利的细节。想象一下，让一个人画一张使「与所有可能人脸的平均差异最小」的脸——他会画出一张模糊的「平均脸」，而不是一张清晰具体的脸。

• VQ-GAN（Esser 等，2021）通过把 VQ-VAE 框架与生成对抗网络（第 6 章）中的**判别器（discriminator）**结合起来，解决了这个问题。判别器是一个基于图块的卷积网络，判断一个局部图像图块是真实的（来自训练数据）还是假的（来自解码器）。这种对抗损失鼓励解码器产生感知上锐利、真实的纹理，而不是逐像素的平均。

• VQ-GAN 的目标在 VQ-VAE 损失上加了两个项：

• 对抗损失（adversarial loss） \mathcal{L}_\text{adv} 是施加在解码器输出上的标准 GAN 目标。判别器 \mathcal{D} 试图区分真实图块和解码图块，解码器（生成器）则试图骗过它。非饱和形式为：

• 感知损失（perceptual loss） \mathcal{L}_\text{perc} 用一个预训练网络（通常是 VGG 或 LPIPS）比较原始图像与重建图像之间的特征激活：

• 其中 \phi_l 表示预训练网络第 l 层的特征图。这个损失捕捉的是高层结构相似性，而非像素级准确度。

• 权重 \lambda_\text{adv} 被自适应地设置，使对抗梯度和重建梯度保持平衡，防止训练早期（重建还很差时）对抗损失占据主导。

• 结果是一个在相同码本大小下能产生比 VQ-VAE 锐利得多的重建的分词器。VQ-GAN 是许多主流图像生成系统背后的骨干分词器，包括最初的 DALL-E、Parti 以及众多文生图模型。它把一张 256x256 的图像转换成 16x16 或 32x32 的离散词元网格，码本大小从 1024 到 16384，每个空间维度上达到 16 倍到 64 倍的压缩比。

残差量化与多尺度码本

• 单个码本给重建质量设了一个硬上限：每个空间位置由恰好一个码本向量表示，任何比码本所能表达的更细的细节都会丢失。想象用一个固定调色板里的单个词来描述一种颜色：「青色」很接近但不够精确。如果你能加一句修饰——「青色，但更偏蓝一点、再亮一点」——你就会接近得多。

• **残差量化（residual quantisation，RQ）**迭代地应用这个想法。在第一步量化产生 \mathbf{z}_q^{(1)} 之后，计算残差 \mathbf{r}^{(1)} = \mathbf{z}_e - \mathbf{z}_q^{(1)}，再用第二个码本对残差量化得到 \mathbf{z}_q^{(2)}，如此进行 T 个层级：

• 最终的量化表示是 \hat{\mathbf{z}} = \sum_{t=1}^{T} \mathbf{z}_q^{(t)}。T 个层级每个用大小为 K 的码本，有效词表大小是 K^T，而你只需存储 T \times K 个向量，而不是 K^T 个。例如，8 个层级、K = 1024，就能给出有效 1024^8 \approx 10^{24} 项，而只存储 8192 个向量。

• 每个后续层级捕捉更细的细节：第一个码本捕捉粗略结构，第二个捕捉中频修正，依此类推。这类似于 JPEG 中的逐次逼近，或网页图片中的渐进式渲染——先出现一个粗糙版本，然后细节逐步填入。

• **多尺度码本（multi-scale codebook）**通过在不同空间分辨率上操作来扩展这个想法。你不再反复量化同一个空间网格，而是在多个尺度上量化：粗网格捕捉全局结构，细网格捕捉局部细节。这与第 8 章目标检测部分的特征金字塔思想相关，那里不同尺度的特征捕捉不同层次的细节。

• **乘积量化（product quantisation）**是一种相关技术，把 d 维潜在向量拆成 M 个维度为 d/M 的子向量，每个子向量用自己的码本独立量化。这给出 K^M 的有效词表，而只存储 M \times K 个向量。乘积量化被广泛用于近似最近邻搜索（第 13 章），也被改造用于图像分词。

• **有限标量量化（finite scalar quantisation，FSQ）**由 Mentzer 等（2023）提出，采取了完全不同的思路：它不学习码本，而是简单地把潜在向量的每个维度舍入到一组固定的整数等级之一（例如 \{-2, -1, 0, 1, 2\}）。每维 L 个等级、共 d 维，隐式码本大小为 L^d。FSQ 完全避免了码本坍缩，因为根本没有可学习的码本向量，只有可学习的编码器输出，再被确定性舍入。直通估计器处理舍入的不可微性。

实践中的图像分词器

• 从 VQ-VAE 到 VQ-GAN 再到残差量化的演进，催生了一系列用于最先进生成模型的实用图像分词器。

DALL-E 分词器（dVAE）

• 最初的 DALL-E（Ramesh 等，2021）用一个离散 VAE（dVAE）把 256x256 的图像分词成 32x32 的词元网格，码本大小为 8192。dVAE 用 Gumbel-Softmax 松弛替代了硬 \arg\min 量化，使训练时的前向传播可微。在推理时，用 \arg\max 产生硬词元分配。dVAE 用重建损失、对均匀先验的 KL 散度以及 Gumbel-Softmax 的可学习温度调度组合训练。DALL-E 随后训练了一个 120 亿参数的自回归 transformer，建模 256 个文本词元和 1024 个图像词元（32x32）的联合分布。

LlamaGen

• LlamaGen（Sun 等，2024）表明，只要有一个好的图像分词器，你就可以把标准的 Llama 式语言模型架构（第 7 章）改用于自回归图像生成。LlamaGen 用一个改进的 VQ-GAN 分词器，码本较大（16,384 项），并训练一个朴素的自回归 transformer（除了分词器之外没有任何针对图像的特殊改动），按光栅扫描顺序从左到右预测图像词元。关键洞见是：一旦图像被分词成离散序列，对语言有效的下一个词元预测范式对图像同样有效，这验证了「分词真正桥接了模态鸿沟」这一想法。

Cosmos 分词器

• Cosmos 分词器（NVIDIA，2024）为图像和视频设计了一个统一框架。它用一个因果 3D 架构把图像当作单帧视频处理，使同一个分词器能同时处理两种模态。Cosmos 同时支持连续和离散两种分词模式：连续模式输出实值潜在向量（用于扩散模型后端），离散模式则应用有限标量量化产生整数词元（用于自回归模型后端）。编码器使用因果 3D 卷积，使每一帧的词元只依赖当前帧和过去帧，从而支持流式视频分词。

视频分词

• 视频在图像的空间维度之上加了第三个轴——时间。视频是帧的序列，通常是每秒 24-30 帧，相邻帧高度冗余，因为视觉世界在 33 毫秒内不会剧烈变化。视频分词利用这种时间冗余，达到比逐帧独立分词高得多的压缩率。

• 把视频压缩想象成一本翻页书。如果你每一页都从头画起，你需要成千上万张精细的画。但大多数页与邻居几乎相同，所以你可以每 10 页画一张完整的「关键帧」，只在中间的页上记下小改动。视频分词器会自动学会这一招。

3D VQ-VAE

• 把 VQ-VAE 扩展到视频最直接的方式是 3D VQ-VAE，它把编码器和解码器中的 2D 卷积替换为同时作用于空间和时间维度的 3D 卷积。如果编码器在空间上下采样 f_s 倍、在时间上下采样 f_t 倍，那么一段 T \times H \times W 的视频片段就变成 (T/f_t) \times (H/f_s) \times (W/f_s) 的词元网格。

• 例如，f_s = 16、f_t = 4 时，一段 16 帧、256x256 的视频片段变成 4 \times 16 \times 16 = 1024 的词元序列。这紧凑到足以让 transformer 自回归地建模，而原始像素数则是 16 \times 256 \times 256 \times 3 \approx 310 万个值。

• 3D 卷积联合学习空间和时间特征。浅层捕捉局部运动（帧之间移动的边缘），更深的层捕捉更高层的动态（物体出现、消失或形变）。这正是第 8 章卷积网络的层次化特征提取原理，沿时间轴的延伸。

因果视频分词器

• 标准 3D 卷积会同时看过去、当前和未来帧，这意味着你需要整段视频片段才能开始分词。**因果视频分词器（causal video tokeniser）**约束时间卷积，使每个输出只依赖当前帧和过去帧，绝不依赖未来帧。这类似于自回归 transformer 中的因果掩码（第 7 章）：信息只能沿时间向前流动，不能向后。

• 因果分词对两种用例至关重要。第一，流式处理（streaming）：你可以在帧到达时实时分词视频，无需缓冲未来帧。第二，自回归生成（autoregressive generation）：当 transformer 逐帧生成视频时，第 t 帧的词元必须能在不知道第 t+1 帧的情况下计算，因为第 t+1 帧还没生成。

• 因果约束通过对时间卷积进行非对称填充来实现：时间大小为 k 的核在过去的方向填 k-1 个零，在未来的方向填零个零，保证时刻 t 的输出只依赖时刻 t-k+1, \ldots, t 的输入。

• 因果视频分词器有一个优雅的性质：它能不经过任何特殊处理就分词单张图像（一段「单帧视频」）。第一帧没有过去上下文，所以它的词元只从这一帧本身计算。这种**图像–视频统一（image-video unification）**意味着一个分词器就能同时服务于两种模态，简化了架构，并使模型能用同一个解码器生成图像和视频。

时间压缩策略

• 不同的应用需要不同的时间压缩比。对于动作识别（细微运动很重要），温和的压缩（f_t = 2）保留时间细节。对于长视频生成（存储数千帧不可行），激进的压缩（f_t = 8 或更高）是必要的。

• 有些分词器采用因式化压缩（factorised compression）：空间压缩和时间压缩在分开的阶段进行。首先，一个 2D 编码器独立地压缩每一帧，产生逐帧的潜在网格。然后，一个 1D 时间编码器沿时间维度压缩。这种因式化比完整 3D 卷积计算更便宜，并允许空间和时间使用不同的压缩比。代价是它无法像联合 3D 编码那样高效地捕捉时空模式（比如斜向移动的球）。

• **时间插值词元（temporal interpolation token）**是一项较新的创新：分词器只对关键帧做完整编码，把中间帧表示为轻量的插值码，描述如何在关键帧之间变形。这镜像了经典视频压缩（H.264/HEVC 中的 I 帧和 P 帧），只不过是在学到的潜在空间中。

连续词元 vs 离散词元

• 并非每个下游模型都需要离散词元。扩散模型（diffusion model）（第 10 章第 04 节）原生支持连续值——它们迭代地去噪一个高斯样本，损失函数（去噪分数匹配）定义在连续空间上。对于扩散后端，分词器的编码器产生连续的潜在向量，从不被量化。潜在扩散模型（latent diffusion model）（Stable Diffusion、DALL-E 3、Flux）使用类似 VQ-GAN 的编码器–解码器，但完全跳过码本，在连续潜在空间中运作。

• 自回归模型（autoregressive model）（GPT 式）则相反，它通过对 K 个类做 softmax 来从有限词表预测下一个词元。它们本质上需要离散词元。每一个使用自回归 transformer 的图像生成系统（DALL-E、Parti、LlamaGen、Chameleon）都依赖离散分词器。

• 因此，在连续和离散词元之间的选择由生成后端决定：

• 当满足以下条件时使用离散词元：模型是自回归的（带交叉熵损失的下一个词元预测）、你希望与文本词元共享词表以构建统一的多模态模型，或者你需要精确的词元级控制（例如通过替换词元来检索或编辑）。

• 当满足以下条件时使用连续词元：模型是扩散模型或流匹配（flow-matching）模型、任务需要非常高保真的重建（连续潜在表示完全避免量化误差），或者你想用作用于实值向量的回归损失。

• 一些最近的架构同时支持两种模式。例如 Cosmos 分词器，既可以从同一个编码器输出连续潜在表示（用于它的扩散模式），也可以输出 FSQ 离散化的词元（用于它的自回归模式），通过一个可开关的轻量量化头实现。

• **软量化（soft quantisation）**是一种折中：不做硬 \arg\min 分配，而是计算最接近的 top-k 个码本项的加权平均，权重由负距离上的 softmax 给出。这比硬量化保留更多信息，同时又近似离散。一些系统在训练时用软量化，在推理时用硬量化。

应用

自回归图像生成

• 一旦图像成为离散词元序列，你就可以训练一个标准的自回归 transformer 来建模它。图像词元被展平为一维序列（通常按光栅扫描顺序：从左到右、从上到下），transformer 用标准交叉熵损失学习 p(\text{token}_i | \text{token}_1, \ldots, \text{token}_{i-1})。在生成时，词元被一个一个地采样，完整的网格再通过分词器的解码器转换为像素。

• 以文本为条件很直接：把文本词元前置到图像词元序列之前，这样模型学习的就是 p(\text{image tokens} | \text{text tokens})。这正是 DALL-E、Parti 和 LlamaGen 做文生图生成的方式。文本和图像词元共享同一个 transformer、同一个注意力机制，往往还共享同一个嵌入表（文本和图像词元占据不同的索引区间）。

• 光栅扫描顺序引入了一种人为的不对称性：图像左上角最先被生成，没有任何关于右下角的上下文。有几项工作针对这一点。掩蔽图像建模（masked image modelling）（MaskGIT）训练一个双向 transformer，同时生成所有词元但置信度不同，迭代地解掩蔽最自信的词元。**多尺度生成（multi-scale generation）**先生成粗粒度词元（捕捉全局构图），再用残差词元精修。这些方法以牺牲纯粹从左到右生成的简洁性，换取更好的全局连贯性。

统一的视觉–语言词元

• 图像分词最深层的动机是统一（unification）：把视觉和语言放进同一种表征格式，让单一模型架构同时处理两者。正如我们在第 7 章讨论的，语言模型是极为强大的序列到序列机器。通过把图像表示为词元序列，我们免费继承了语言建模的全部基础设施——预训练配方、缩放律、RLHF、上下文长度扩展。

• Chameleon（Meta，2024）是一个突出的例子：它用一个码本为 8192 项的 VQ-GAN 分词器把图像转成词元，再与文本词元交错，放进一个约 65,000 项（文本+图像）的统一词表。一个标准 transformer 在混合的文本–图像序列上训练，使它能在同一次前向传播中，根据图像生成文本、根据文本生成图像，或生成交错的图文内容。

• Gemini（Google，2024）在大规模上采取类似方法，在一个 transformer 内原生理解和生成图像、音频与文本，由模态专属的分词器送入一个共享序列。

• 统一模型的关键工程挑战是词表平衡（vocabulary balance）：如果 65,000 个词表项中有 8192 个是图像词元，模型可能给视觉分配不足的容量。解决方案包括为每种模态使用独立的嵌入层（仅在注意力层面共享）、模态专属的损失加权，以及预训练时精心调配的数据混合比例。

编程练习（使用 CoLab 或 notebook）

1. 在 JAX 中实现一个最小化的 VQ 层：给定一批编码器输出向量，执行最近邻码本查找，并计算 VQ-VAE 损失（重建 + 码本 + 承诺）。把码本利用率可视化为直方图。

import jax
import jax.numpy as jnp
import matplotlib.pyplot as plt

# --- 最小化 VQ 层 ---
key = jax.random.PRNGKey(42)
d = 8          # 嵌入维度
K = 64         # 码本大小
n_vectors = 256  # 一批编码器输出

# 随机的编码器输出和码本
k1, k2 = jax.random.split(key)
z_e = jax.random.normal(k1, (n_vectors, d))       # 编码器输出
codebook = jax.random.normal(k2, (K, d)) * 0.1     # 码本（小初始化）

# 最近邻查找：为每个 z_e 找最近的码本项
# distances[i, k] = ||z_e[i] - codebook[k]||^2
distances = (
    jnp.sum(z_e ** 2, axis=1, keepdims=True)
    - 2 * z_e @ codebook.T
    + jnp.sum(codebook ** 2, axis=1, keepdims=True).T
)
indices = jnp.argmin(distances, axis=1)       # 词元索引
z_q = codebook[indices]                        # 量化后的向量

# VQ-VAE 损失项
beta = 0.25
loss_codebook = jnp.mean((jax.lax.stop_gradient(z_e) - z_q) ** 2)
loss_commit   = jnp.mean((z_e - jax.lax.stop_gradient(z_q)) ** 2)
loss_total    = loss_codebook + beta * loss_commit
print(f"Codebook loss: {loss_codebook:.4f}, Commitment loss: {loss_commit:.4f}")

# 码本利用率
unique, counts = jnp.unique(indices, return_counts=True, size=K, fill_value=-1)
plt.figure(figsize=(10, 4))
plt.bar(range(K), counts, color='#3498db', alpha=0.8)
plt.xlabel('Codebook Index'); plt.ylabel('Assignment Count')
plt.title(f'Codebook Utilisation ({jnp.sum(counts > 0)}/{K} entries used)')
plt.grid(True, alpha=0.3); plt.tight_layout(); plt.show()
# 试试：把 K 增大到 512 观察坍缩。然后加入码本重置逻辑。
​

2. 构建一个玩具级的二维向量量化器，学习铺贴一个二维分布。生成随机二维点，通过 EMA 更新学习一个码本，并可视化 Voronoi 区域。

import jax
import jax.numpy as jnp
import matplotlib.pyplot as plt

# 从高斯混合生成二维数据
key = jax.random.PRNGKey(0)
n_points = 2000
K = 16  # 码本项数
gamma = 0.99  # EMA 衰减

# 四个簇
keys = jax.random.split(key, 5)
centres = jnp.array([[2, 2], [-2, 2], [-2, -2], [2, -2]], dtype=jnp.float32)
data = jnp.concatenate([
    jax.random.normal(keys[i], (n_points // 4, 2)) * 0.5 + centres[i]
    for i in range(4)
])

# 从随机数据点初始化码本
idx = jax.random.choice(keys[4], n_points, (K,), replace=False)
codebook = data[idx]
ema_count = jnp.ones(K)
ema_sum = codebook.copy()

# 跑若干轮基于 EMA 的码本学习
for epoch in range(30):
    # 把每个点分配到最近的码本项
    dists = jnp.sum((data[:, None, :] - codebook[None, :, :]) ** 2, axis=2)
    assignments = jnp.argmin(dists, axis=1)
    # EMA 更新
    for k in range(K):
        mask = (assignments == k)
        count_k = jnp.sum(mask)
        ema_count = ema_count.at[k].set(gamma * ema_count[k] + (1 - gamma) * count_k)
        if count_k > 0:
            sum_k = jnp.sum(data[mask], axis=0)
            ema_sum = ema_sum.at[k].set(gamma * ema_sum[k] + (1 - gamma) * sum_k)
    codebook = ema_sum / ema_count[:, None]

# 可视化分配和码本
fig, ax = plt.subplots(1, 1, figsize=(8, 8))
colors = plt.cm.tab20(jnp.linspace(0, 1, K))
for k in range(K):
    mask = assignments == k
    ax.scatter(data[mask, 0], data[mask, 1], c=[colors[k]], s=5, alpha=0.3)
ax.scatter(codebook[:, 0], codebook[:, 1], c='black', s=120, marker='X',
           edgecolors='white', linewidths=1.5, zorder=10, label='Codebook')
ax.set_title(f'Learned VQ Codebook ({K} entries) on 2D Data')
ax.legend(); ax.set_aspect('equal'); ax.grid(True, alpha=0.3)
plt.tight_layout(); plt.show()
# 试试：把 K 增大到 64 观察更细的铺贴。减小 gamma 看不稳定性。
​

3. 演示残差量化：用 T 个连续的量化阶段编码一批向量，衡量重建误差如何随每个层级下降。

import jax
import jax.numpy as jnp
import matplotlib.pyplot as plt

key = jax.random.PRNGKey(7)
d = 16         # 嵌入维度
K = 32         # 每个层级的码本大小
T = 8          # 残差层级数
n_vectors = 512

# 要量化的随机数据
k1, *cb_keys = jax.random.split(key, T + 1)
z = jax.random.normal(k1, (n_vectors, d))

# 每个层级独立随机码本
codebooks = [jax.random.normal(cb_keys[t], (K, d)) * (0.5 ** t)
             for t in range(T)]

# 残差量化循环
residual = z.copy()
z_hat = jnp.zeros_like(z)
errors = []

for t in range(T):
    cb = codebooks[t]
    dists = (jnp.sum(residual ** 2, axis=1, keepdims=True)
             - 2 * residual @ cb.T
             + jnp.sum(cb ** 2, axis=1, keepdims=True).T)
    indices = jnp.argmin(dists, axis=1)
    z_q_t = cb[indices]
    z_hat = z_hat + z_q_t
    residual = residual - z_q_t
    mse = jnp.mean(jnp.sum((z - z_hat) ** 2, axis=1))
    errors.append(float(mse))
    print(f"Level {t+1}: MSE = {mse:.4f}")

plt.figure(figsize=(8, 5))
plt.plot(range(1, T + 1), errors, 'o-', color='#e74c3c', linewidth=2, markersize=8)
plt.xlabel('Residual Quantisation Level')
plt.ylabel('Reconstruction MSE')
plt.title('Error Reduction with Residual Quantisation')
plt.xticks(range(1, T + 1)); plt.grid(True, alpha=0.3)
plt.tight_layout(); plt.show()
# 试试：用单个大小为 K*T 的码本与 RQ 比较。哪个更好？
​

4. 模拟一个简单的一维「视频分词器」：生成一串一维信号（模拟视频帧），应用因果时间压缩，并与非因果压缩比较重建质量。

import jax
import jax.numpy as jnp
import matplotlib.pyplot as plt

key = jax.random.PRNGKey(99)
n_frames = 16
frame_len = 64

# 生成一段「视频」：一个高斯凸起在帧之间缓慢移动
x_axis = jnp.linspace(-3, 3, frame_len)
frames = jnp.stack([
    jnp.exp(-0.5 * (x_axis - (-2 + 4 * t / n_frames)) ** 2)
    for t in range(n_frames)
])  # 形状: (n_frames, frame_len)

# 因果时间压缩：每一帧的编码只依赖过去的帧
# 简单做法：当前帧与过去的指数衰减加权和
alpha_causal = 0.6
causal_codes = jnp.zeros_like(frames)
causal_codes = causal_codes.at[0].set(frames[0])
for t in range(1, n_frames):
    causal_codes = causal_codes.at[t].set(
        alpha_causal * frames[t] + (1 - alpha_causal) * causal_codes[t - 1]
    )

# 非因果：与过去和未来都取平均（双边平滑）
kernel = jnp.array([0.2, 0.6, 0.2])  # 过去、当前、未来
padded = jnp.concatenate([frames[:1], frames, frames[-1:]], axis=0)
noncausal_codes = jnp.stack([
    kernel[0] * padded[t] + kernel[1] * padded[t+1] + kernel[2] * padded[t+2]
    for t in range(n_frames)
])

# 重建误差
mse_causal = jnp.mean((frames - causal_codes) ** 2)
mse_noncausal = jnp.mean((frames - noncausal_codes) ** 2)
print(f"Causal MSE: {mse_causal:.6f}, Non-causal MSE: {mse_noncausal:.6f}")

fig, axes = plt.subplots(1, 3, figsize=(15, 5))
for ax, data, title in zip(axes,
    [frames, causal_codes, noncausal_codes],
    ['Original Frames', f'Causal (MSE={mse_causal:.5f})',
     f'Non-causal (MSE={mse_noncausal:.5f})']):
    ax.imshow(data, aspect='auto', cmap='viridis', origin='lower')
    ax.set_xlabel('Spatial Position'); ax.set_ylabel('Frame Index')
    ax.set_title(title)
plt.tight_layout(); plt.show()
# 试试：改变 alpha_causal 和核权重。alpha=1.0 时会发生什么？
​