Vector Quantization for Recommendation 笔记

嵌入表的内存爆炸随着物品和用户规模激增完整的嵌入表Embedding Table会占用数百 GB 甚至 TB 级显存远超单机容量迫使系统采用复杂的参数服务器带来高昂的通信开销语义空间利用不足庞大的嵌入表参数中大量长尾物品的嵌入向量训练不充分而表示能力却受限于表的大小难以在有限的计算资源下进一步提升模型的表达力传统解法如特征哈希虽然能压缩内存但会带来严重的哈希冲突损失语义信息。这催生了对更高效、更灵活的向量表示方法的需求Vector Quantization正是其中一类重要方法组合式编码与码本向量量化在推荐中的核心思想是不再为每个物品学习一个独立的嵌入向量而是学习多组共享的码本codebook每个物品的最终表示由其所属码本中基向量的组合来构成。这套框架和 NLP 中的 BPE 分词或 VQ-VAE 思路相似Product Quantization想象你有一个 1024 维的向量PQ 的处理方式是切分成子空间把 1024 维向量横向切割成 8 段每段 128 维。这 8 个 128 维的小向量构成 8 个独立的子空间子空间码本为每个子空间独立训练一个码本。假如码本大小 K256那么每个子空间码本包含 256 个 128 维的基底向量8 个子空间共有 8 个独立码本独立量化原向量的第1段前128维只在第1号码本内寻找最近邻得到一个索引比如 15第2段下一个128维只在第2号码本内寻找得到一个索引比如 203……依此类推最终表示一个完整的 1024 维物品向量被表示为一个 8 元组的索引序列 比如Residual Quantization同样用 1024 维向量RQ 的处理方式是不分段而是用多级码本进行迭代式逼近第一级粗量化用第一级码本全局唯一1024 维假设K256对整个向量做最近邻查找找到第一个码字 c1计算残差求出第一级量化后的残差余量 r1 原始向量 - c1。这个残差代表了第一级没能捕捉的细节信息第二级残差量化用第二级码本专门学习如何量化残差也是 1024 维对残差 r1 进行最近邻查找得到第二个码字 c2迭代计算新的残差 再用第三级码本去量化它……如此重复 L 级最终表示与重建物品被表示为一个 L 元组的索引序列 。重建时直接把对应索引的各个码字加起来即可最终向量 ≈VQ 的学习向量量化VQ的学习通常不是一个独立的任务而是和推荐模型的最终目标如点击率预估联合训练的。其核心挑战在于“选码”操作在码本中找最近邻是不可导的梯度无法通过离散的索引回传到编码器。围绕这一点主流的学习方式分成了两大类直通估计器STE这是VQ-VAE中最经典的方法也是推荐系统中联合训练最常用的方式前向传播对于一个输入的稠密向量 比如物品经过某层网络后的表示在大小为 K 的码本中找到距离最近的码字。假设第 k 个码字 距离向量 最近那么量化后的向量就是 但是在实际上只需要存储索引 k 就行了反向传播由于 argmin 这个操作是不可导的因此 STE 会假装寻找最近邻这一步不存在直接把 的梯度作为 的梯度原封不动地拷贝给编码器的输出这相当于在反向传播时把整个量化操作近似为一个恒等映射 。只要量化误差不太大这个梯度近似就能引导编码器往正确的方向更新。为了让这个近似有效训练时会在推荐主损失如交叉熵上额外增加两个 VQ 专用的辅助损失减小码本和输入的稠密向量 之间的距离所以码本内容也是在不断更新的码本损失Commitment Loss只更新码本不更新编码器。让选中的码字 靠近编码器输出 达到训练码本的目的。形式为 其中 表示 stop gradient即 在这里被截断梯度编码器约束损失Codebook Loss 的对偶只更新编码器不更新码本。约束编码器的输出不要离码字太远避免量化误差过大。形式为 其中 对 截断梯度两者的权重通常由一个超参数 β 控制典型值如 β 0.25软量化 Gumbel-Softmax软量化就是把离散选择变成可导的连续近似代表性方法是Gumbel-Softmax不再做硬最近邻查找而是计算输入 与所有码字 的相似度如内积得到 logits然后通过 Gumbel-Softmax 产生一个近似 one-hot 的软分配向量最终的量化向量就是所有码字的软分配加权和此时前向和反向都是连续可导的在推荐中的应用嵌入层压缩直接用 VQ 技术替换传统的巨型嵌入表作为内存高效的嵌入层。这是最直接的应用适用于召回和精排阶段的特征表达semantic ID 生成将物品的完整嵌入经过 RQ 量化后得到的多级码本索引序列如可视作该物品的 semantic ID。这种离散 ID 能直接接入大语言模型成为连接推荐系统和生成式 LLM 的桥梁用于端到端的生成式推荐