基于重要性采样的期望估计——sampled softmax推导

一、背景

在推荐召回系统中，通常会采用tow-tower模型并利用log softmax作为损失进行优化，设$[B]$为mini-batch，$[C]$为全局语料库，$s(x,y)$为query x和item y的相似度分数，则有如下的损失函数：
L=−1B∑i∈[B]log(es(xi,yi)∑j∈[C]es(xi,yj))=−1B∑i∈[B]{ s(xi,yi)−log∑j∈[C]es(xi,yj)} \begin{align} \mathcal{L} &= - \frac {1}{B} \sum_{i \in [B]}log(\frac {e^{s(x_i,y_i)}}{\sum_{j\in [C]} e^{s(x_i,y_j)}}) \\ &= - \frac{1}{B} \sum_{i \in [B]}\{s(x_i,y_i) - log\sum_{j\in [C]} e^{s(x_i,y_j)}\} \end{align} L​=−B1​i∈[B]∑​log(∑j∈[C]​es(xi​,yj​)es(xi​,yi​)​)=−B1​i∈[B]∑​{ s(xi​,yi​)−logj∈[C]∑​es(xi​,yj​)}​​

对损失求导
∇θL=−1B∑i∈[B]{ ∇θs(xi,yi)−∑j∈[C]es(xi,yj)∑k∈[C]es(xi,yk)∇θs(xi,yj)}=−1B∑i∈[B]{ ∇θs(xi,yi)−∑j∈[C]P(yj∣xi)∇θs(xi,yj)}=−1B∑i∈[B]{ ∇θs(xi,yi)⏟part one−EP[∇θs(xi,yj)]⏟part two} \begin{align} \mathcal{\nabla_\theta L} &=- \frac{1}{B} \sum_{i \in [B]} \{ \nabla_{\theta} s(x_i, y_i) - \sum_{j \in [C]} \frac{e^{s(x_i, y_j)}}{\sum_{k\in [C]} e^{s(x_i, y_k)}} \nabla_ \theta s(x_i, y_j)\} \\ &= - \frac{1}{B} \sum_{i \in [B]} \{ \nabla_{\theta} s(x_i, y_i) - \sum_{j \in [C]} P(y_j|x_i) \nabla_ \theta s(x_i, y_j)\} \\ &= - \frac{1}{B} \sum_{i \in [B]} \{ \underbrace{\nabla_{\theta} s(x_i, y_i)}_{part\ one} - \underbrace{E_{P}[\nabla_\theta s(x_i, y_j)]}_{part \ two}\} \end{align} ∇θ​L​=−B1​i∈[B]∑​{ ∇θ​s(xi​,yi​)−j∈[C]∑​∑k∈[C]​es(xi​,yk​)es(xi​,yj​)​∇θ​s(xi​,yj​)}=−B1​i∈[B]∑​{ ∇θ​s(xi​,yi​)−j∈[C]∑​