xdeepfm.py

1"""2xDeepFM 精排模型（第12章 §12.4.1.4）34xDeepFM = 线性组件 + CIN 组件 + DNN 组件。56CIN（Compressed Interaction Network）逐层把上一层的交互表示与原始字段嵌入7再做逐元素乘积，显式构造任意阶交互：89    X^(k)_{h,:} = Σ_i^{H_{k-1}} Σ_j^{m} W^(k,h)_{i,j} (X^(k-1)_{i,:} ∘ X^(0)_{j,:})10    p^(k) = sum-pooling(X^(k))11    p⁺ = [p^(1); p^(2); ...; p^(T)]1213    ŷ = σ(w_linearᵀx + w_CINᵀp⁺ + z_DNN)1415本书依据：《第十二章-精排模型.pdf》§12.4.1.4，PDF 第 24–30 页。16依赖：PyTorch。运行：python3 xdeepfm.py17"""1819import random20import torch21import torch.nn as nn22import torch.nn.functional as F23import torch.optim as optim242526def _auc(scores, labels) -> float:27    pos = [s for s, y in zip(scores, labels) if y == 1]28    neg = [s for s, y in zip(scores, labels) if y == 0]29    if not pos or not neg:30        return 0.531    return sum(1.0 if p > n else 0.5 if p == n else 0.032               for p in pos for n in neg) / (len(pos) * len(neg))333435# ── CIN 组件 ──────────────────────────────────────────────────────────────────3637class CIN(nn.Module):38    """39    压缩交互网络。4041    实现约定：42      X^(0) ∈ ℝ^(m × D)，m = 字段数，D = 嵌入维度43      第 k 层 H_k 个交互向量，H_k 由参数决定（默认各层相同）44      每层的 W^(k,h) ∈ ℝ^(H_{k-1} × m)45    """4647    def __init__(self, n_fields: int, emb_dim: int, layer_sizes=(8, 8)):48        super().__init__()49        self.n_fields = n_fields50        self.emb_dim = emb_dim51        self.layer_sizes = list(layer_sizes)5253        # 每层的联合权重 (H_k, H_{k-1}, m)54        self.weights = nn.ParameterList()55        prev_h = n_fields                    # H_0 = m56        for h in self.layer_sizes:57            self.weights.append(58                nn.Parameter(torch.randn(h, prev_h, n_fields) * 0.05))59            prev_h = h6061    def forward(self, x0: torch.Tensor) -> torch.Tensor:62        """63        x0: (B, m, D)64        返回 p⁺: (B, ΣH_k)65        """66        B = x0.size(0)67        x_prev = x0                                    # (B, H_0=m, D)68        pooled = []69        for layer_idx, w in enumerate(self.weights):70            H_k, H_prev, m = w.shape71            # z[i,j] = x_prev[:, i, :] ∘ x0[:, j, :] → (B, H_prev, m, D)72            z = x_prev.unsqueeze(2) * x0.unsqueeze(1)73            # X^(k)_{h} = Σ_{i,j} W[h,i,j] z[i,j]74            # einsum: (h,i,j) × (b,i,j,d) → (b,h,d)75            x_cur = torch.einsum("hij,bijd->bhd", w, z)76            # Sum pooling over embedding dim → p^(k) ∈ ℝ^(B, H_k)77            pooled.append(x_cur.sum(dim=2))78            x_prev = x_cur79        return torch.cat(pooled, dim=1)                # (B, ΣH_k)808182# ── xDeepFM 完整模型 ─────────────────────────────────────────────────────────8384class xDeepFM(nn.Module):85    def __init__(self, n_fields: int, vocab_size: int, emb_dim: int = 6,86                 cin_layers=(6, 6), hidden: int = 24):87        super().__init__()88        self.n_fields = n_fields89        self.emb = nn.Embedding(vocab_size * n_fields, emb_dim)90        self.linear = nn.Embedding(vocab_size * n_fields, 1)91        self.bias = nn.Parameter(torch.zeros(1))9293        self.cin = CIN(n_fields, emb_dim, cin_layers)94        cin_out = sum(cin_layers)95        self.cin_proj = nn.Linear(cin_out, 1)9697        self.dnn = nn.Sequential(98            nn.Linear(n_fields * emb_dim, hidden), nn.ReLU(),99            nn.Linear(hidden, 1),100        )101        nn.init.zeros_(self.linear.weight)102        nn.init.normal_(self.emb.weight, std=0.01)103104    def _gids(self, feat_ids):105        offsets = (torch.arange(self.n_fields, device=feat_ids.device)106                   * (self.emb.num_embeddings // self.n_fields))107        return feat_ids + offsets108109    def forward(self, feat_ids: torch.Tensor) -> torch.Tensor:110        gids = self._gids(feat_ids)111        x0 = self.emb(gids)                            # (B, m, D)112113        z_linear = self.linear(gids).sum(dim=(1, 2))114        p_plus = self.cin(x0)115        z_cin = self.cin_proj(p_plus).squeeze(-1)116        z_dnn = self.dnn(x0.reshape(x0.size(0), -1)).squeeze(-1)117        return self.bias + z_linear + z_cin + z_dnn118119    def predict_proba(self, feat_ids):120        with torch.no_grad():121            return torch.sigmoid(self.forward(feat_ids))122123    def fit(self, feat_ids, y, epochs: int = 40, lr: float = 1e-2):124        opt = optim.Adam(self.parameters(), lr=lr, weight_decay=1e-5)125        for ep in range(epochs):126            opt.zero_grad()127            loss = F.binary_cross_entropy_with_logits(128                self.forward(feat_ids), y.float())129            loss.backward()130            opt.step()131            if (ep + 1) % 10 == 0:132                print(f"  epoch {ep+1:3d}  loss={loss.item():.4f}")133        return loss.item()134135    def rank(self, feat_ids, top_k):136        scores = self.predict_proba(feat_ids).tolist()137        return sorted(enumerate(scores), key=lambda t: -t[1])[:top_k]138139140# ── 合成数据 ──────────────────────────────────────────────────────────────────141142def make_data(n: int = 500, n_fields: int = 5, vocab: int = 8, seed: int = 42):143    """144    含三阶交互的数据：字段 0×1×2 的联合模式对标签影响较大，145    用于观察 CIN 显式建模高阶交互的能力（DNN 只能隐式学习）。146    """147    torch.manual_seed(seed)148    random.seed(seed)149    ids = torch.randint(0, vocab, (n, n_fields), dtype=torch.long)150    local = ids.float()151    triple = (local[:, 0] > 3).float() * (local[:, 1] < 4).float() * \152             (local[:, 2] > 2).float()153    logit = (0.3 * local[:, 0] + 0.2 * local[:, 3]154             + 1.6 * triple                                # 三阶交互155             + 0.7 * (local[:, 0] == local[:, 1]).float()156             - 2.3)157    y = (torch.sigmoid(logit + 0.35 * torch.randn(n)) > 0.5).long()158    return ids, y159160161# ── 主演示 ────────────────────────────────────────────────────────────────────162163if __name__ == "__main__":164    random.seed(42)165    torch.manual_seed(42)166167    N_FIELDS, VOCAB = 5, 8168    ids, y = make_data(600, N_FIELDS, VOCAB)169    split = 480170    ids_tr, y_tr = ids[:split], y[:split]171    ids_te, y_te = ids[split:], y[split:]172173    print("=== xDeepFM 训练（线性 + CIN + DNN）===")174    model = xDeepFM(N_FIELDS, VOCAB, emb_dim=6, cin_layers=(6, 6))175    model.fit(ids_tr, y_tr, epochs=30)176177    p_te = model.predict_proba(ids_te).tolist()178    auc_te = _auc(p_te, y_te.tolist())179    print(f"\n测试集 AUC={auc_te:.4f}")180181    # ── 三组件贡献 ────────────────────────────────────────────────────────────182    with torch.no_grad():183        gids = model._gids(ids_te[:20])184        x0 = model.emb(gids)185        z_lin = model.linear(gids).sum(dim=(1, 2))186        z_cin = model.cin_proj(model.cin(x0)).squeeze(-1)187        z_dnn = model.dnn(x0.reshape(x0.size(0), -1)).squeeze(-1)188    print(f"\n三组件得分量级（前 20 样本平均绝对值）：")189    print(f"  |z_linear|={z_lin.abs().mean().item():.4f}   "190          f"|z_CIN|={z_cin.abs().mean().item():.4f}   "191          f"|z_DNN|={z_dnn.abs().mean().item():.4f}")192193    # ── CIN 维度流转 ──────────────────────────────────────────────────────────194    with torch.no_grad():195        p_plus = model.cin(model.emb(model._gids(ids_te[:4])))196    print(f"\nCIN 池化输出维度：{tuple(p_plus.shape)} "197          f"(B=4, ΣH_k={sum(model.cin.layer_sizes)})")198    print(f"  第 1 层构造二阶交互（H_0=m={N_FIELDS} → H_1={model.cin.layer_sizes[0]}）")199    print(f"  第 2 层构造三阶交互（H_1 → H_2={model.cin.layer_sizes[1]}）")200    print("  正文 §12.4.1.4：交互阶数随网络深度递增，这是 CIN 与 DNN 的本质区别。")201202    # ── 精排推理演示 ──────────────────────────────────────────────────────────203    top5 = model.rank(ids_te[:20], top_k=5)204    print("\n精排 Top-5（候选集 20 条）:")205    for rank_i, (idx, s) in enumerate(top5, 1):206        print(f"  #{rank_i}  item={idx:2d}  score={s:.4f}  label={y_te[idx].item()}")