pnn.py

1"""2PNN 精排模型（第12章 §12.4.1.2）34Product-based Neural Network 在 Embedding 之后插入 Product Layer，5显式构造字段之间的乘性二阶交互，再送入 DNN：67    f_i = W_0^(i) x_i8    z = Concat(f_1,...,f_N)9    l_1 = ReLU(l_z + l_p + b_1)10    ŷ = σ(DNN(l_1))1112两种乘积形式（正文 PDF 第 14–19 页）：13  - IPNN：内积 p_{i,j}=<f_i,f_j>，用秩一/低秩分解把 O(N²) 降到 O(N)14  - OPNN：外积 p_{i,j}=f_i f_jᵀ，用特征叠加 f_Σ=Σf_i 压缩为 M×M 矩阵1516本书依据：《第十二章-精排模型.pdf》§12.4.1.2，PDF 第 12–19 页。17依赖：PyTorch。运行：python3 pnn.py18"""1920import random21import torch22import torch.nn as nn23import torch.nn.functional as F24import torch.optim as optim252627def _auc(scores, labels) -> float:28    pos = [s for s, y in zip(scores, labels) if y == 1]29    neg = [s for s, y in zip(scores, labels) if y == 0]30    if not pos or not neg:31        return 0.532    return sum(1.0 if p > n else 0.5 if p == n else 0.033               for p in pos for n in neg) / (len(pos) * len(neg))343536# ── IPNN ──────────────────────────────────────────────────────────────────────3738class IPNN(nn.Module):39    """40    内积版 PNN。4142    正文关键推导（PDF 第 16–18 页）：43      l_p^(n) = Σ_i Σ_j θ_i^(n) θ_j^(n) <f_i,f_j>44               = ‖Σ_i θ_i^(n) f_i‖²         （秩一分解，O(D1·M·N)）45    低秩推广：46      l_p^(n) = Σ_{r=1}^{K} ‖Σ_i θ_{i,r}^(n) f_i‖²47    """4849    def __init__(self, n_fields: int, vocab_size: int, emb_dim: int = 8,50                 d1: int = 24, rank: int = 1, hidden: int = 16):51        super().__init__()52        self.n_fields = n_fields53        self.embs = nn.ModuleList([54            nn.Embedding(vocab_size, emb_dim) for _ in range(n_fields)])55        self.z_dim = n_fields * emb_dim5657        # 线性部分 W_z^(n)：z → D158        self.wz = nn.Linear(self.z_dim, d1)59        # 交互部分 θ^(n) ∈ ℝ^(N × K)：每个隐藏节点 n 一组字段权重60        # 用 (d1, rank, n_fields) 存储，输出 l_p = Σ_r ‖Σ_i θ_{n,r,i} f_i‖²61        self.theta = nn.Parameter(torch.randn(d1, rank, n_fields) * 0.1)62        self.b1 = nn.Parameter(torch.zeros(d1))63        self.mlp = nn.Sequential(64            nn.Linear(d1, hidden), nn.ReLU(), nn.Linear(hidden, 1))6566    def product_layer(self, fields: torch.Tensor) -> torch.Tensor:67        """fields: (B, N, M) → l1: (B, D1)"""68        B, N, M = fields.shape69        z = fields.reshape(B, N * M)7071        # l_z^(n) = <W_z^(n), z>72        l_z = self.wz(z)                                  # (B, D1)7374        # l_p^(n) = Σ_r ‖Σ_i θ_{n,r,i} f_i‖²   —— O(B·D1·K·N·M)75        # theta: (D1, K, N) → s: (B, D1, K, M)76        s = torch.einsum("dkn,bnm->bdkm", self.theta, fields)77        l_p = s.pow(2).sum(dim=(2, 3))                    # (B, D1)7879        return F.relu(l_z + l_p + self.b1)8081    def forward(self, feat_ids):82        fields = torch.stack([emb(feat_ids[:, i])83                              for i, emb in enumerate(self.embs)], dim=1)84        l1 = self.product_layer(fields)85        return self.mlp(l1).squeeze(-1)868788# ── OPNN ──────────────────────────────────────────────────────────────────────8990class OPNN(nn.Module):91    """92    外积版 PNN。9394    正文（PDF 第 18 页）用特征叠加把 N²M² 压缩为 M×M：95      f_Σ = Σ_i f_i96      p = f_Σ f_Σᵀ97      l_p^(n) = <W_p^(n), f_Σ f_Σᵀ>98    """99100    def __init__(self, n_fields: int, vocab_size: int, emb_dim: int = 8,101                 d1: int = 24, hidden: int = 16):102        super().__init__()103        self.embs = nn.ModuleList([104            nn.Embedding(vocab_size, emb_dim) for _ in range(n_fields)])105        self.z_dim = n_fields * emb_dim106        self.emb_dim = emb_dim107108        self.wz = nn.Linear(self.z_dim, d1)109        # W_p^(n) ∈ ℝ^(M×M)，但由于外积对称，只需对称部分110        self.wp = nn.Parameter(torch.randn(d1, emb_dim, emb_dim) * 0.05)111        self.b1 = nn.Parameter(torch.zeros(d1))112        self.mlp = nn.Sequential(113            nn.Linear(d1, hidden), nn.ReLU(), nn.Linear(hidden, 1))114115    def product_layer(self, fields):116        B, N, M = fields.shape117        z = fields.reshape(B, N * M)118        l_z = self.wz(z)                                   # (B, D1)119120        f_sigma = fields.sum(dim=1)                        # (B, M)121        outer = f_sigma.unsqueeze(2) * f_sigma.unsqueeze(1)  # (B, M, M)122        # l_p^(n) = <W_p^(n), outer> = Σ_{a,b} W[a,b] * outer[a,b]123        l_p = torch.einsum("bij,nij->bn", outer, self.wp)  # (B, D1)124        return F.relu(l_z + l_p + self.b1)125126    def forward(self, feat_ids):127        fields = torch.stack([emb(feat_ids[:, i])128                              for i, emb in enumerate(self.embs)], dim=1)129        return self.mlp(self.product_layer(fields)).squeeze(-1)130131132# ── 通用训练循环 ──────────────────────────────────────────────────────────────133134def train(model, ids_tr, y_tr, ids_te, y_te, epochs: int = 40, lr: float = 1e-2):135    opt = optim.Adam(model.parameters(), lr=lr)136    for ep in range(epochs):137        opt.zero_grad()138        loss = F.binary_cross_entropy_with_logits(model(ids_tr), y_tr.float())139        loss.backward()140        opt.step()141    with torch.no_grad():142        p = torch.sigmoid(model(ids_te)).tolist()143    return loss.item(), _auc(p, y_te.tolist())144145146# ── 合成数据 ──────────────────────────────────────────────────────────────────147148def make_data(n: int = 500, n_fields: int = 5, vocab: int = 8, seed: int = 42):149    """标签含字段 0×1 与 2×3 的乘性交互，突出 Product Layer 的作用。"""150    torch.manual_seed(seed)151    random.seed(seed)152    ids = torch.randint(0, vocab, (n, n_fields), dtype=torch.long)153    local = ids.float()154155    logit = (0.35 * local[:, 0] + 0.3 * local[:, 1] - 0.25 * local[:, 2]156             + 0.9 * (local[:, 0] * local[:, 1]) / vocab      # 乘性交互157             + 0.7 * (local[:, 2] * local[:, 3]) / vocab158             - 2.5)159    y = (torch.sigmoid(logit + 0.4 * torch.randn(n)) > 0.5).long()160    return ids, y161162163# ── 主演示 ────────────────────────────────────────────────────────────────────164165if __name__ == "__main__":166    random.seed(42)167    torch.manual_seed(42)168169    N_FIELDS, VOCAB = 5, 8170    ids, y = make_data(500, N_FIELDS, VOCAB)171    split = 380172    ids_tr, y_tr = ids[:split], y[:split]173    ids_te, y_te = ids[split:], y[split:]174175    print(f"数据：{N_FIELDS} 字段 × {VOCAB} 取值，训练 {split} / 测试 {500-split}")176177    # ── IPNN（秩一）──────────────────────────────────────────────────────────178    torch.manual_seed(42)179    ipnn = IPNN(N_FIELDS, VOCAB, emb_dim=8, d1=24, rank=1)180    loss_i, auc_i = train(ipnn, ids_tr, y_tr, ids_te, y_te)181    print(f"\nIPNN（K=1 秩一分解）  loss={loss_i:.4f}  AUC={auc_i:.4f}")182183    # ── IPNN（秩 K=4）────────────────────────────────────────────────────────184    torch.manual_seed(42)185    ipnn_k4 = IPNN(N_FIELDS, VOCAB, emb_dim=8, d1=24, rank=4)186    loss_k, auc_k = train(ipnn_k4, ids_tr, y_tr, ids_te, y_te)187    n_theta_k1 = ipnn.theta.numel()188    print(f"IPNN（K=4 低秩）      loss={loss_k:.4f}  AUC={auc_k:.4f}")189    print(f"  交互参数：K=1 → {n_theta_k1}，K=4 → {ipnn_k4.theta.numel()}"190          f"（{ipnn_k4.theta.numel()/n_theta_k1:.0f} 倍）")191192    # ── OPNN ─────────────────────────────────────────────────────────────────193    torch.manual_seed(42)194    opnn = OPNN(N_FIELDS, VOCAB, emb_dim=8, d1=24)195    loss_o, auc_o = train(opnn, ids_tr, y_tr, ids_te, y_te)196    print(f"OPNN（特征叠加外积）  loss={loss_o:.4f}  AUC={auc_o:.4f}")197198    # ── 精排推理演示 ──────────────────────────────────────────────────────────199    with torch.no_grad():200        scores = torch.sigmoid(ipnn(ids_te[:20])).tolist()201    top5 = sorted(enumerate(scores), key=lambda t: -t[1])[:5]202    print("\nIPNN 精排 Top-5（候选集 20 条）:")203    for rank_i, (idx, s) in enumerate(top5, 1):204        print(f"  #{rank_i}  item={idx:2d}  score={s:.4f}  label={y_te[idx].item()}")205206    # ── 复杂度对比 ────────────────────────────────────────────────────────────207    print(f"\n复杂度（正文 §12.4.1.2）：")208    print(f"  IPNN 秩一: O(D1·M·N)         本实现即此形式")209    print(f"  OPNN 叠加: O(D1·M·(M+N))     本实现即此形式")210    print("  Product Layer 是 PNN 区别于 FNN 的关键组件：")211    print("  去掉二阶交互后网络退化为 Embedding + MLP 的普通深度模型。")