原文: Wukong: Towards a Scaling Law for Large-Scale Recommendation

网络结构

参考 LLM 领域 Transformer 的结构设计, 对推荐场景进行了适配。 网络结构如下:

网络主要有多个悟空层堆叠而成。 每个悟空层包括两个模块: FMB和LCB。其中 FMB (Factorization Machine Block) 通过 DeepFM 建模二阶特征交叉, LCB (Linear Compress Block) 则对特征维度进行压缩。

$$
X_{i+1} = \text{LN}(\text{concat}(FMB_i(X_i), LCB_i(X_i)) + X_i)
$$

FMB

$$
\operatorname{FMB}(X_i) = \operatorname{reshape}(\operatorname{MLP}(\operatorname{LN}(\operatorname{flatten}(\operatorname{FM}(X_i))))))
$$

其实就是简单的FM。 在实现时, 由于FM的复杂度是$O(n^2)$, 在此做了一点优化。 即计算 $XX^T Y$ 来代替 $XX^T$ , 其中, $Y$是一个 $n \times k$的矩阵。 由于 $k << n$ , 因此 , 可以先计算 $X^T Y$ , 再计算 $X \times X^TY$。

LCB

$$
LCB(X_i) = WLX_i
$$

在 GitHub 讨论 ↗