mmoe.py

1"""2MMoE 多任务精排模型（第12章 §12.4.2.1）34Multi-gate Mixture-of-Experts（Google KDD 2018）：56    e_i = f_i(x),            i=1…n   # n 个共享专家（MLP）7    g^(k)(x) = softmax(W_g^(k) x),   # 第 k 个任务的门控向量 ∈ ℝ^n8    f^(k)(x) = Σ_i g_i^(k)(x) e_i   # 任务 k 的加权专家输出9    ŷ_k = h_k(f^(k)(x))              # 任务 k 的独立塔网络1011与 Shared-Bottom / OMoE 对比：12  - Shared-Bottom : f^(k)(x) = f(x)              所有任务固定共享同一底层13  - OMoE          : f^(k)(x) = Σ g_i(x) e_i      共享门控，各任务使用相同权重14  - MMoE          : f^(k)(x) = Σ g_i^(k)(x) e_i  每任务独立门控，自适应专家选择1516本书依据：《第十二章-精排模型.pdf》§12.4.2.1，PDF 第 30–35 页（多任务部分）。17依赖：PyTorch。运行：python3 mmoe.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# ── Shared-Bottom（对比基线）────────────────────────────────────────────────3738class SharedBottom(nn.Module):39    """所有任务固定共享同一底层 MLP，再各自过塔网络。"""4041    def __init__(self, in_dim: int, n_tasks: int,42                 bottom_hidden: int = 32, tower_hidden: int = 16):43        super().__init__()44        self.bottom = nn.Sequential(45            nn.Linear(in_dim, bottom_hidden), nn.ReLU(),46            nn.Linear(bottom_hidden, bottom_hidden), nn.ReLU(),47        )48        self.towers = nn.ModuleList([49            nn.Sequential(nn.Linear(bottom_hidden, tower_hidden), nn.ReLU(),50                          nn.Linear(tower_hidden, 1))51            for _ in range(n_tasks)52        ])5354    def forward(self, x):55        h = self.bottom(x)56        return torch.cat([t(h) for t in self.towers], dim=1)  # (B, T)575859# ── MMoE 模型 ─────────────────────────────────────────────────────────────────6061class MMoE(nn.Module):62    """63    MMoE：n 个共享专家 + T 个任务独立门控 + T 个塔网络。6465    正文（PDF §12.4.2.1）核心：66      「不同任务可以根据输入特征，自适应地组合不同专家的输出……67        即使输入特征相同，各任务可以生成不同的专家权重，68        从而形成不同的特征组合。」69    """7071    def __init__(self, in_dim: int, n_experts: int, n_tasks: int,72                 expert_hidden: int = 32, tower_hidden: int = 16):73        super().__init__()74        self.n_experts = n_experts75        self.n_tasks = n_tasks7677        # 共享专家（结构相同，参数独立）78        self.experts = nn.ModuleList([79            nn.Sequential(80                nn.Linear(in_dim, expert_hidden), nn.ReLU(),81                nn.Linear(expert_hidden, expert_hidden), nn.ReLU(),82            )83            for _ in range(n_experts)84        ])8586        # 每个任务独立的门控线性层：x → ℝ^n_experts → softmax87        self.gates = nn.ModuleList([88            nn.Linear(in_dim, n_experts) for _ in range(n_tasks)])8990        # 每个任务独立的塔网络91        self.towers = nn.ModuleList([92            nn.Sequential(93                nn.Linear(expert_hidden, tower_hidden), nn.ReLU(),94                nn.Linear(tower_hidden, 1),95            )96            for _ in range(n_tasks)97        ])9899    def forward(self, x: torch.Tensor) -> torch.Tensor:100        """101        x : (B, in_dim)102        返回 logits : (B, T)103        """104        # 所有专家并行计算105        expert_outs = torch.stack(106            [e(x) for e in self.experts], dim=1)          # (B, n, expert_hidden)107108        logits = []109        for k in range(self.n_tasks):110            g = F.softmax(self.gates[k](x), dim=1)        # (B, n)111            # 加权融合专家输出112            f_k = (g.unsqueeze(-1) * expert_outs).sum(dim=1)  # (B, expert_hidden)113            logits.append(self.towers[k](f_k))            # (B, 1)114115        return torch.cat(logits, dim=1)                    # (B, T)116117    def fit(self, x, ys, epochs: int = 40, lr: float = 1e-2,118            task_weights=None):119        """120        ys : (B, T) 多任务标签121        task_weights : 各任务损失权重（默认均等）122        """123        w = task_weights or [1.0] * self.n_tasks124        opt = optim.Adam(self.parameters(), lr=lr, weight_decay=1e-5)125        for ep in range(epochs):126            opt.zero_grad()127            logits = self.forward(x)                      # (B, T)128            loss = sum(129                w[k] * F.binary_cross_entropy_with_logits(130                    logits[:, k], ys[:, k].float())131                for k in range(self.n_tasks))132            loss.backward()133            opt.step()134            if (ep + 1) % 10 == 0:135                print(f"  epoch {ep+1:3d}  total_loss={loss.item():.4f}")136137    def predict_proba(self, x) -> torch.Tensor:138        with torch.no_grad():139            return torch.sigmoid(self.forward(x))         # (B, T)140141142# ── 合成数据：两任务存在显式冲突 ─────────────────────────────────────────────143144def make_data(n: int = 2400, in_dim: int = 16, seed: int = 42):145    """146    两任务的输入特征互不重叠：147      任务 0（点击）依赖 x0–x3 的线性与 x1*x2 交互；148      任务 1（完播）依赖 x6–x9 的线性与 x7*x8 交互。149    模拟正文所述「不同任务关注重点存在明显差异」的场景——150    每个任务都有自己的专属信号，Shared-Bottom 被迫用同一底层表示151    同时拟合两组不相关的特征，容易出现负迁移。152    """153    torch.manual_seed(seed)154    random.seed(seed)155    x = torch.randn(n, in_dim)156157    l0 = 2.2 * x[:, 0] + 1.8 * (x[:, 1] * x[:, 2]) + 1.0 * x[:, 3] - 0.3158    l1 = 2.2 * x[:, 6] + 1.8 * (x[:, 7] * x[:, 8]) + 1.0 * x[:, 9] - 0.3159160    y0 = (torch.sigmoid(l0 + 0.35 * torch.randn(n)) > 0.5).long()161    y1 = (torch.sigmoid(l1 + 0.35 * torch.randn(n)) > 0.5).long()162163    return x, torch.stack([y0, y1], dim=1)164165166# ── 主演示 ────────────────────────────────────────────────────────────────────167168if __name__ == "__main__":169    random.seed(42)170    torch.manual_seed(42)171172    IN_DIM, N_TASKS = 16, 2173    x, y = make_data(2400, IN_DIM)174    split = 1920175    x_tr, y_tr = x[:split], y[:split]176    x_te, y_te = x[split:], y[split:]177178    # ── Shared-Bottom 基线（容量受限：bottom_hidden=8）───────────────────────179    # 正文（PDF §12.4.2.1）指出 Shared-Bottom 的局限：180    # 「所有任务必须使用相同的底层表示。当不同任务的关注重点存在明显差异时，181    #   共享表示可能难以同时满足所有任务的需求。」182    # 用较小的底层容量来凸显这一局限。183    torch.manual_seed(42)184    sb = SharedBottom(IN_DIM, N_TASKS, bottom_hidden=8, tower_hidden=16)185    opt_sb = optim.Adam(sb.parameters(), lr=5e-3, weight_decay=1e-5)186    for ep in range(120):187        opt_sb.zero_grad()188        logits = sb(x_tr)189        loss = sum(F.binary_cross_entropy_with_logits(190            logits[:, k], y_tr[:, k].float()) for k in range(N_TASKS))191        loss.backward()192        opt_sb.step()193    with torch.no_grad():194        p_sb = torch.sigmoid(sb(x_te))195    auc_sb = [_auc(p_sb[:, k].tolist(), y_te[:, k].tolist()) for k in range(N_TASKS)]196    print(f"Shared-Bottom（h=8）  Task0 AUC={auc_sb[0]:.4f}  Task1 AUC={auc_sb[1]:.4f}")197198    # ── MMoE ─────────────────────────────────────────────────────────────────199    # 4 个共享专家，每个专家容量与 Shared-Bottom 底层一致（h=8），200    # 但通过独立门控让两个任务选择不同的专家组合。201    torch.manual_seed(42)202    mmoe = MMoE(IN_DIM, n_experts=4, n_tasks=N_TASKS,203                expert_hidden=8, tower_hidden=16)204    print(f"\nMMoE 训练（4 个共享专家 h=8，2 个独立门控）")205    mmoe.fit(x_tr, y_tr, epochs=120, lr=5e-3)206207    p_mmoe = mmoe.predict_proba(x_te)208    auc_mmoe = [_auc(p_mmoe[:, k].tolist(), y_te[:, k].tolist())209                for k in range(N_TASKS)]210    print(f"\nMMoE（4 专家 h=8）    Task0 AUC={auc_mmoe[0]:.4f}  Task1 AUC={auc_mmoe[1]:.4f}")211212    # ── 门控权重分析 ──────────────────────────────────────────────────────────213    print(f"\n门控权重（前 20 样本平均，正文：不同任务应学到不同专家偏好）：")214    with torch.no_grad():215        for k in range(N_TASKS):216            g = F.softmax(mmoe.gates[k](x_te[:20]), dim=1)217            g_mean = g.mean(dim=0)218            print(f"  Task{k} 专家权重: " +219                  " ".join(f"E{i}={g_mean[i].item():.3f}"220                            for i in range(4)))221222    # ── 与 Shared-Bottom 对比 ────────────────────────────────────────────────223    print(f"\n对比（正文 §12.4.2.1）：")224    print(f"  Shared-Bottom avg AUC={sum(auc_sb)/2:.4f}  "225          f"MMoE avg AUC={sum(auc_mmoe)/2:.4f}")226    print("  MMoE 通过任务独立门控，让冲突任务选择不同专家组合，")227    print("  缓解了 Shared-Bottom 中「所有任务被迫使用相同底层表示」的局限。")228229    # ── 精排推理演示（任务 0：点击）────────────────────────────────────────────230    scores_t0 = p_mmoe[:20, 0].tolist()231    top5 = sorted(enumerate(scores_t0), key=lambda t: -t[1])[:5]232    print("\n按任务0（点击）精排 Top-5（候选集 20 条）:")233    for rank_i, (idx, s) in enumerate(top5, 1):234        print(f"  #{rank_i}  item={idx:2d}  p_click={s:.4f}  "235              f"label_click={y_te[idx, 0].item()}")