神经网络Transformer架构中的向量量化 基本概念与动机 在Transformer架构中，激活值（即每层神经元的输出）和嵌入（如词嵌入）通常是高维的连续向量。向量量化是一种将连续向量空间映射到离散码本空间的技术。其核心动机是 压缩 和 高效表示 ：用一个有限大小的“码本”中的“码字”来近似表示大量高精度的向量，从而大幅减少存储和传输开销，并可能加速某些计算。 核心组件与过程 码本 ：一个可学习的、固定大小的向量集合。例如，一个包含 K 个 D 维向量的集合 C = {e_1, e_2, ..., e_K} 。 K 是码本大小， D 与待量化向量的维度相同。 量化过程 ：对于一个输入向量 z （例如Transformer中间层的某个特征向量），量化器在码本 C 中查找与之最相似的码字（通常使用欧几里得距离或余弦相似度作为度量）。然后，用该码字的索引 k （一个离散的整数，范围从1到 K ）来代表原始向量 z 。在后续计算中，通常使用这个码字 e_k 来替代 z 。 反量化/查找过程 ：当需要使用量化后的表示时，只需根据索引 k 从码本 C 中取出对应的码字 e_k 即可。这类似于一个“查表”操作，非常高效。 在Transformer中的典型应用点 激活值量化 ：对Transformer各层前馈网络或注意力机制的输出激活进行量化。这可以极大压缩模型运行时的中间状态，对于模型部署、尤其是边缘设备上的推理至关重要。 嵌入量化 ：对庞大的词嵌入矩阵进行量化。原始的嵌入矩阵参数数量庞大，通过向量量化，可以用一个较小的码本和一系列索引来替代，显著减少模型文件大小。 构建离散潜在表示 ：在某些生成式模型（如VQ-VAE）的变体中，编码器输出的连续特征被量化为离散的码字索引序列，这个离散序列随后被解码器使用。这为在离散空间中进行生成建模或推理提供了可能。 训练与优化挑战 梯度传递问题 ：量化操作（即根据最近邻选择码字索引）本身是不可导的，这阻碍了梯度通过量化器反向传播，使得码本无法通过标准的反向传播算法进行端到端训练。 直通估计器 ：解决上述问题的关键技术。在反向传播时，绕过不可导的“硬”量化操作（索引选择），直接将输出码字 e_k 的梯度完整地复制给输入向量 z 。这样，码本 C 可以通过接收到来自 z 的梯度来更新，而编码器（产生 z 的网络）则收到一个近似梯度，使其输出的向量趋向于与某个码字对齐。 码本崩塌 ：在训练过程中，优化可能只使用码本中少数几个码字，而大部分码字被闲置，导致表示能力下降。通常通过一些技术来缓解，如对未使用的码字进行重启或引入额外的损失项来鼓励码本使用均衡。 优势与影响 高效压缩 ：大幅减少模型存储和内存占用，是模型压缩和轻量化部署的关键技术之一。 计算加速 ：离散的索引表示可以启用高效的查找和基于整数的运算，在某些硬件上能提升推理速度。 结构化先验 ：离散的码本空间可以自然地引入符号化或概念化的结构，有时能提高表示的鲁棒性和可解释性。 高级变体与扩展 乘积量化 ：将高维向量分割成多个子向量，分别为每个子空间学习一个子码本。一个向量由多个子码本的索引组合表示，能在总码字数量（组合数）巨大的情况下，保持总参数量（各子码本之和）相对较小，实现极高的压缩率。 残差向量量化 ：分级进行量化。先对原始向量进行粗量化，然后对量化残差（原始向量与粗量化码字的差）再进行一次量化，如此迭代。这能更精细地逼近原始向量。 与Transformer特定模块结合 ：研究如何将VQ更无缝地整合进Transformer架构，例如设计专用的量化注意力头或量化前馈层，以在压缩的同时最小化性能损失。