Wukong: 一种推荐系统的 Scaling Law 方案

原文: 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$$