COBRA 精读笔记

生成式召回 —— COBRA 精读笔记

文章链接：https://arxiv.org/pdf/2503.02453
机构：百度
发布时间：2025.03
参考文献：
https://zhuanlan.zhihu.com/p/1921680262242206140
Copyright (c) Wang-Luning. All Rights Reserved.

TIGER等通过预测下一个物品的ID（SID）的生成式召回方法相比于稠密检索方法的计算效率更高，但其将量化生成SID的过程和序列建模过程进行了分离，导致信息丢失，难以建模精细的相似性关系。

因此，COBRA（Cascaded Organized Bi-Represented generAtive retrieval）希望结合生成式方法和稠密检索方法的优势，通过一个级联的过程来结合稀疏SID和稠密向量的优势。其交替地依次生成下一个物品的SID和稠密向量表征，分别用来捕捉粗粒度的语义信息和细粒度的细节信息。先生成的SID作为粗粒度类别语义信息可以作为条件来辅助该物品稠密向量的生成，降低了稠密表征的学习难度，并促进了两种表征协同学习。在推理时结合Beam Search的思路来高效地生成一组多个候选ID，然后再通过BeamFusion分数得到兼顾准确性和多样性的召回集。

物品表征的构建：

稀疏SID的构建：COBRA将物品的属性转化成一段文字描述，压缩成一个稠密向量后通过RQ-VAE来生成它的稀疏SID表示。稀疏SID可以捕捉物品的粗粒度类别信息，为接下来的细粒度处理提供基础。这里简洁起见假设每个SID里只有1个code，而实际上也可以扩展成多个code。
稠密向量表征的构建：仍然将物品的属性（如类别文字、标题文字、品牌文字等）转化为一段文字描述，加上一个[CLS][CLS] $v_t$ 。

$t$ $ID_t$ $v_t$ $(ID_t,v_t)$ ，它既可以提供稳定的粗粒度类别基础信息，也可以提供连续稠密的细粒度特征信息，使得模型既能捕捉到物品的高层次语义信息又能捕捉到细粒度细节信息。

序列建模与训练：

$i$ $ID_i$ $e_t=\text{Embed}(ID_i)$ ，则：

\begin{aligned} S_{1 : t} & = [h_{1}, \dots, h_{t}] \\ = [e_{1}, v_{1}, e_{2}, v_{2}, \dots, e_{t}, v_{t}] \end{aligned}

$S_{1:t}$ $ID_{t+1}$ $v_{t+1}$ ：

P (I D_{t + 1}, v_{t + 1} | S_{1 : t}) = P (I D_{t + 1} | S_{1 : t}) P (v_{t + 1} | I D_{t + 1}, S_{1 : t})

$\hat v_{t+1}$ 。

$S_{1:t}$ $y_t$ $z_{t+1}$ $e_{t+1}$ $\bar S_{1:t}$ $\hat v_{t+1}$ ：

\begin{matrix} y_{t} = TransformerDecoder (S_{1 : t}) \\ z_{t + 1} = SparseHead (y_{t}) \\ {\bar{S}}_{1 : t} = [S_{1 : t}, e_{t + 1}] = [e_{1}, v_{1}, \dots, e_{t}, v_{t}, e_{t + 1}] \\ {\hat{v}}_{t + 1} = TransformerDecoder ({\bar{S}}_{1 : t}) \end{matrix}

训练时，端到端的训练过程可以联合优化稀疏ID和稠密表征的预测。其中，稀疏ID的预测loss就是一个多分类loss，希望最大化每个位置的ground truth的ID的预测概率：

L_{s p a r s e} = - \sum_{t = 1}^{T - 1} \log \frac{\exp (z_{t + 1}^{I D_{t + 1}})}{\sum_{j = 1}^{C} \exp (z_{t + 1}^{j})}

$T$ $C$ 为码本大小。

稠密表征的loss是希望预测生成的稠密向量能够尽可能接近下一个物品的真实向量，同时尽可能远离和其他负样本的向量的距离：

L_{d e n s e} = - \sum_{t = 1}^{T - 1} \log \frac{\exp (\cos ({\hat{v}}_{t + 1}, v_{t + 1}))}{\sum_{i t e m_{j} \in B a t c h} \exp (\cos ({\hat{v}}_{t + 1}, v_{i t e m_{j}}))}

生成每个物品的ground truth稠密向量表征的Encoder在该训练过程中也是可训练的，可以不断地迭代来为物品生成更好的表征。

总的loss为稀疏和稠密两部分之和：

L = L_{s p a r s e} + L_{d e n s e}

推理过程：

在推理中也使用了由粗到精到生成思路，首先生成稀疏的IDs，然后再进一步通过生成稠密向量来修正。具体而言，首先生成一组候选ID和一组对应的稠密向量：

稀疏ID的生成：
$S_{1:T}$ $M$ 个概率值最高的下一个ID，也即一次生成了一组多个ID，用来描绘下一个可能的物品：
${{\hat{I D}}_{T + 1}^{k}}_{k = 1}^{M} = BeamSearch (TransformerDecoder (S_{1 : T}), M)$
$\phi_{\hat{ID}_{T+1}^k}$ ，代表了该ID的置信度
稠密向量的生成：
对于每个ID，将它们分别接到历史序列后，再分别输入模型得到各自预测的稠密向量：
${\hat{v}}_{T + 1}^{k} = TransformerDecoder ([S_{1 : T}, Embed ({\hat{I D}}_{T + 1}^{k})])$

$\hat ID_{T+1}^k$ $\mathcal A_{k}$ ：

A_{k} = ANN ({\hat{I D}}_{T + 1}^{k}, C ({\hat{I D}}_{T + 1}^{k}), N)

$\mathcal C(\hat {ID}_{T+1}^k)$ $N$ 为召回数量。

$a$ $\phi$ ）和它们与对应ID的预测稠密向量之间的相似度（也即同一个ID的召回集合内部比较），为其生成一个全局可比的BeamFusion分数：

Φ^{({\hat{v}}_{T + 1}^{k}, {\hat{I D}}_{T + 1}^{k}, a)} = Softmax (τ ϕ_{{\hat{I D}}_{T + 1}^{k}}) \times Softmax (ψ \cos ({\hat{v}}_{T + 1}^{k}, a))

$\tau,\psi$ 为系数。

$\mathcal R$ ：

R = TopK (⋃_{k = 1}^{M} A_{k}, Φ, K)