feat(models): add WuKong (ICML 2024) via onboarding pipeline

- WuKongLayer: stackable FMB (X X^T -> MLP) + LCB (linear embedding
  recombination) with residual; registered in custom_objects.
- WuKong model: stacks N WuKong layers over field embeddings + DNN head.
- Wire registry + tests + docs; mark implemented in candidate KB.
- audit: 33/33 models fully wired.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Author: zhouyuchen.zyc

README.md

@@ -84,6 +84,7 @@ Introduction](https://zhuanlan.zhihu.com/p/53231955)) and [welcome to join us!](
 |   OneTrans                   | [industrial 2024][OneTrans: Unified Feature Interaction and Sequence Modeling with One Transformer in Industrial Recommender Systems]()   |
 |   FinalMLP                   | [AAAI 2023][FinalMLP: An Enhanced Two-Stream MLP Model for CTR Prediction](https://arxiv.org/abs/2304.00902)   |
 |   MaskNet                   | [DLP-KDD 2021][MaskNet: Introducing Feature-Wise Multiplication to CTR Ranking Models by Instance-Guided Mask](https://arxiv.org/abs/2102.07619)   |
+|   WuKong                   | [ICML 2024][Wukong: Towards a Scaling Law for Large-Scale Recommendation](https://arxiv.org/abs/2403.02545)   |
 
 ## Citation
 

benchmarks/RESULTS.md

@@ -172,3 +172,4 @@ Auto-recorded verify outcomes for models added through `benchmarks.onboard`.
 - **OneTrans**: see `benchmarks/onboard/reports/OneTrans.md`
 - **FinalMLP**: see `benchmarks/onboard/reports/FinalMLP.md`
 - **MaskNet**: see `benchmarks/onboard/reports/MaskNet.md`
+- **WuKong**: see `benchmarks/onboard/reports/WuKong.md`

benchmarks/onboard/candidates.json

@@ -102,7 +102,7 @@
       "ref_impl_url": "https://github.com/reczoo/FuxiCTR",
       "paper_metric": null,
       "difficulty": "medium",
-      "status": "candidate"
+      "status": "implemented"
     },
     {
       "name": "FCN",

benchmarks/onboard/reports/WuKong.md

@@ -0,0 +1,13 @@
+# Verify report: WuKong
+
+- category: **single**
+- correctness:
+    - audit (fully wired): PASS
+    - unit test: PASS (2 passed, 3 warnings in 18.12s)
+- effectiveness (primary metric = AUC):
+    - DeepFM: AUC=0.5771 LogLoss=0.7297 [ok] (baseline)
+    - WuKong: AUC=0.5627 LogLoss=0.5764 [ok]  <- new model
+
+## Verdict
+
+correctness PASS. WuKong AUC=0.5627; DeepFM baseline AUC=0.5771 (new model is below).

benchmarks/registry.py

@@ -10,7 +10,7 @@
 from deepctr.feature_column import DenseFeat, SparseFeat
 from deepctr.models import (AFM, BST, CCPM, DCN, DIEN, DIFM, DIN, DSIN, EDCN, FGCNN, FLEN, FNN,
                             FwFM, IFM, MLR, MMOE, NFM, ONN, PLE, PNN, WDL, AutoInt, DCNMix,
-                            DeepFEFM, DeepFM, ESMM, FiBiNET, SharedBottom, xDeepFM, OneTrans, FinalMLP, MaskNet)
+                            DeepFEFM, DeepFM, ESMM, FiBiNET, SharedBottom, xDeepFM, OneTrans, FinalMLP, MaskNet, WuKong)
 
 _TF2 = version.parse(tf.__version__) >= version.parse("2.0.0")
 _TF28 = version.parse(tf.__version__) >= version.parse("2.8.0")
@@ -68,6 +68,7 @@ def _regroup(feature_columns, n_groups=3):
     "EDCN": lambda lin, dnn, task: EDCN(_sparse_only(lin), _sparse_only(dnn), task=task),  # sparse-only,
     'FinalMLP': lambda lin, dnn, task: FinalMLP(lin, dnn, task=task),
     'MaskNet': lambda lin, dnn, task: MaskNet(lin, dnn, task=task),
+    'WuKong': lambda lin, dnn, task: WuKong(lin, dnn, task=task),
 }
 
 # --------------------------------------------------------------------------- #

deepctr/layers/__init__.py

@@ -12,6 +12,7 @@
 from .sequence import (AttentionSequencePoolingLayer, BiasEncoding, BiLSTM,
                        KMaxPooling, SequencePoolingLayer, WeightedSequenceLayer,
                        Transformer, DynamicGRU, PositionEncoding)
+from .wukong import WuKongLayer
 from .utils import NoMask, Hash, Linear, _Add, combined_dnn_input, softmax, reduce_sum, Concat
 
 custom_objects = {'tf': tf,
@@ -57,4 +58,5 @@
                   'SinusoidalPositionEncoding': SinusoidalPositionEncoding,
                   'TokenSlice': TokenSlice,
                   'InteractionAggregation': InteractionAggregation,
+                  'WuKongLayer': WuKongLayer,
                   }

deepctr/layers/wukong.py

@@ -0,0 +1,91 @@
+# -*- coding:utf-8 -*-
+"""Custom layer for the WuKong model: one stackable interaction layer made of a
+Factorization-Machine Block (FMB) and a Linear Compression Block (LCB), joined
+by a residual connection.
+
+Reference:
+    Zhang B, et al. Wukong: Towards a Scaling Law for Large-Scale Recommendation.
+    ICML 2024. (https://arxiv.org/abs/2403.02545)
+"""
+
+import tensorflow as tf
+from tensorflow.keras.layers import Layer, Dropout
+
+from .normalization import LayerNormalization
+
+
+class WuKongLayer(Layer):
+    """One WuKong interaction layer over a token matrix ``X`` of shape (B, k, d).
+
+    * FMB: explicit pairwise interactions ``X X^T`` (B, k, k) → MLP → reshaped to
+      ``(B, n_fmb, d)`` and layer-normed.
+    * LCB: a learned linear recombination of the input embeddings → ``(B, n_lcb, d)``.
+    * Output: ``concat([FMB, LCB], axis=1) + X`` (residual), shape ``(B, k, d)``,
+      so layers stack uniformly.
+
+    Input / output shape: ``(batch_size, k, d)``.
+    """
+
+    def __init__(self, num_fmb_tokens=None, fmb_hidden_dim=128, dropout_rate=0.0,
+                 seed=1024, **kwargs):
+        super(WuKongLayer, self).__init__(**kwargs)
+        self.num_fmb_tokens = num_fmb_tokens
+        self.fmb_hidden_dim = fmb_hidden_dim
+        self.dropout_rate = dropout_rate
+        self.seed = seed
+        self.dropout = Dropout(dropout_rate, seed=seed)
+        self.ln = LayerNormalization()
+
+    def build(self, input_shape):
+        self.k = int(input_shape[1])
+        self.d = int(input_shape[2])
+        # split output tokens so n_fmb + n_lcb == k (keeps the layer stackable)
+        self.n_fmb = self.num_fmb_tokens if self.num_fmb_tokens else max(1, self.k // 2)
+        self.n_fmb = min(self.n_fmb, self.k)
+        self.n_lcb = self.k - self.n_fmb
+
+        # FMB: flatten(X X^T) (k*k) -> hidden -> n_fmb*d
+        self.fmb_w1 = self.add_weight(name="fmb_w1", shape=(self.k * self.k, self.fmb_hidden_dim),
+                                      initializer="glorot_uniform")
+        self.fmb_b1 = self.add_weight(name="fmb_b1", shape=(self.fmb_hidden_dim,),
+                                      initializer="zeros")
+        self.fmb_w2 = self.add_weight(name="fmb_w2", shape=(self.fmb_hidden_dim, self.n_fmb * self.d),
+                                      initializer="glorot_uniform")
+        self.fmb_b2 = self.add_weight(name="fmb_b2", shape=(self.n_fmb * self.d,),
+                                      initializer="zeros")
+        # LCB: (n_lcb, k) linear recombination of input embeddings
+        if self.n_lcb > 0:
+            self.lcb_w = self.add_weight(name="lcb_w", shape=(self.n_lcb, self.k),
+                                         initializer="glorot_uniform")
+        super(WuKongLayer, self).build(input_shape)
+
+    def call(self, inputs, training=None, **kwargs):
+        x = inputs  # (B, k, d)
+        # ── FMB ──
+        inter = tf.matmul(x, x, transpose_b=True)            # (B, k, k)
+        flat = tf.reshape(inter, [-1, self.k * self.k])      # (B, k*k)
+        h = tf.nn.relu(tf.matmul(flat, self.fmb_w1) + self.fmb_b1)
+        h = self.dropout(h, training=training)
+        h = tf.matmul(h, self.fmb_w2) + self.fmb_b2          # (B, n_fmb*d)
+        fmb = tf.reshape(h, [-1, self.n_fmb, self.d])        # (B, n_fmb, d)
+        fmb = self.ln(fmb)
+
+        # ── LCB ──
+        if self.n_lcb > 0:
+            lcb = tf.einsum("lk,bkd->bld", self.lcb_w, x)    # (B, n_lcb, d)
+            out = tf.concat([fmb, lcb], axis=1)              # (B, k, d)
+        else:
+            out = fmb
+
+        # ── residual ──
+        return out + x
+
+    def compute_output_shape(self, input_shape):
+        return input_shape
+
+    def get_config(self):
+        config = {"num_fmb_tokens": self.num_fmb_tokens, "fmb_hidden_dim": self.fmb_hidden_dim,
+                  "dropout_rate": self.dropout_rate, "seed": self.seed}
+        base = super(WuKongLayer, self).get_config()
+        base.update(config)
+        return base

deepctr/models/__init__.py

@@ -24,6 +24,7 @@
 
 from .finalmlp import FinalMLP
 from .masknet import MaskNet
+from .wukong import WuKong
 __all__ = ["AFM", "CCPM", "DCN", "IFM", "DIFM", "DCNMix", "MLR", "DeepFM", "MLR", "NFM", "DIN", "DIEN", "FNN", "PNN",
            "WDL", "xDeepFM", "AutoInt", "ONN", "FGCNN", "DSIN", "FiBiNET", 'FLEN', "FwFM", "BST", "DeepFEFM",
-           "SharedBottom", "ESMM", "MMOE", "PLE", 'EDCN', "OneTrans", "FinalMLP", "MaskNet"]
+           "SharedBottom", "ESMM", "MMOE", "PLE", 'EDCN', "OneTrans", "FinalMLP", "MaskNet", "WuKong"]

deepctr/models/wukong.py

@@ -0,0 +1,66 @@
+# -*- coding:utf-8 -*-
+"""
+WuKong
+
+Reference:
+    [1] Zhang B, Luo L, Chen Y, et al. Wukong: Towards a Scaling Law for
+    Large-Scale Recommendation. ICML 2024. (https://arxiv.org/abs/2403.02545)
+"""
+
+from tensorflow.keras.models import Model
+from tensorflow.keras.layers import Dense, Flatten
+
+from ..feature_column import build_input_features, get_linear_logit, input_from_feature_columns
+from ..layers.core import PredictionLayer, DNN
+from ..layers.utils import concat_func, combined_dnn_input, add_func
+from ..layers.wukong import WuKongLayer
+
+
+def WuKong(linear_feature_columns, dnn_feature_columns, num_layers=3, fmb_hidden_dim=128,
+           num_fmb_tokens=None, dnn_hidden_units=(128, 64), l2_reg_linear=0.00001,
+           l2_reg_embedding=0.00001, l2_reg_dnn=0, seed=1024, dnn_dropout=0,
+           dnn_activation='relu', dnn_use_bn=False, task='binary'):
+    """Instantiates the WuKong architecture.
+
+    Stacks ``num_layers`` WuKong interaction layers (Factorization-Machine Block +
+    Linear Compression Block, with residuals) over the stacked field embeddings,
+    then a DNN head. All sparse / var-len features must share one ``embedding_dim``
+    (the field embeddings are stacked into a (B, num_fields, d) token matrix).
+
+    :param linear_feature_columns: features for the linear part.
+    :param dnn_feature_columns: features for the deep/interaction part.
+    :param num_layers: number of stacked WuKong layers.
+    :param fmb_hidden_dim: hidden size of each FMB's MLP.
+    :param num_fmb_tokens: FMB output tokens per layer (default: half the fields).
+    :param dnn_hidden_units: layer sizes of the DNN head.
+    :param task: ``"binary"`` or ``"regression"``.
+    :return: A Keras model instance.
+    """
+    features = build_input_features(linear_feature_columns + dnn_feature_columns)
+    inputs_list = list(features.values())
+
+    linear_logit = get_linear_logit(features, linear_feature_columns, seed=seed, prefix='linear',
+                                    l2_reg=l2_reg_linear)
+
+    sparse_embedding_list, dense_value_list = input_from_feature_columns(
+        features, dnn_feature_columns, l2_reg_embedding, seed)
+    if not sparse_embedding_list:
+        raise ValueError("WuKong requires at least one embedding (sparse/var-len) feature.")
+
+    # stack field embeddings into a (B, num_fields, d) token matrix
+    x = concat_func(sparse_embedding_list, axis=1)
+    for _ in range(num_layers):
+        x = WuKongLayer(num_fmb_tokens=num_fmb_tokens, fmb_hidden_dim=fmb_hidden_dim,
+                        dropout_rate=dnn_dropout, seed=seed)(x)
+
+    interaction_out = Flatten()(x)
+    dnn_input = combined_dnn_input([interaction_out], dense_value_list)
+    dnn_out = DNN(dnn_hidden_units, dnn_activation, l2_reg_dnn, dnn_dropout, dnn_use_bn,
+                  seed=seed)(dnn_input)
+    dnn_logit = Dense(1, use_bias=False)(dnn_out)
+
+    final_logit = add_func([linear_logit, dnn_logit])
+    output = PredictionLayer(task)(final_logit)
+
+    model = Model(inputs=inputs_list, outputs=output)
+    return model

docs/source/Features.md

@@ -327,6 +327,12 @@ Applies instance-guided multiplicative masks on feature embeddings / hidden laye
 
 [MaskNet: Introducing Feature-Wise Multiplication to CTR Ranking Models by Instance-Guided Mask](https://arxiv.org/abs/2102.07619)
 
+### WuKong
+
+Stacks factorization-machine blocks (LCB/FMB) with MLPs into a dense architecture whose quality scales predictably with model size (a recommendation scaling law).
+
+[Wukong: Towards a Scaling Law for Large-Scale Recommendation](https://arxiv.org/abs/2403.02545)
+
 ## Sequence Models
 
 ### DIN (Deep Interest Network)