ple.py

1"""2PLE 多任务精排模型（第12章 §12.4.2.2）34Progressive Layered Extraction（腾讯 RecSys 2020）：56CGC（Customized Gate Control）模块：7    在共享专家之外，每个任务有专属专家；8    每个任务的门控混合「本任务专属专家 + 共享专家」：910    f^(k) = Σ g_i^(k) · e_i^(k 或 shared)1112    共享专家另有独立的共享门控，聚合出跨任务的共享表示，13    作为下一层 CGC 共享专家的输入。1415PLE：将多个 CGC 模块逐层堆叠，逐步分离不同任务的特征表示：1617    第 1 层 CGC → 第 2 层 CGC → … → 各任务塔网络1819与 MMoE 的区别：20  - MMoE  : 所有任务仅使用共享专家，通过独立门控选择组合；21  - CGC   : 引入任务专属专家，共享专家补充跨任务信号；22  - PLE   : 逐层堆叠 CGC，使任务特异性随层深递增，减少负迁移。2324本书依据：《第十二章-精排模型.pdf》§12.4.2.2，PDF 第 35–40 页（多任务部分）。25依赖：PyTorch。运行：python3 ple.py26"""2728import random29import torch30import torch.nn as nn31import torch.nn.functional as F32import torch.optim as optim333435def _auc(scores, labels) -> float:36    pos = [s for s, y in zip(scores, labels) if y == 1]37    neg = [s for s, y in zip(scores, labels) if y == 0]38    if not pos or not neg:39        return 0.540    return sum(1.0 if p > n else 0.5 if p == n else 0.041               for p in pos for n in neg) / (len(pos) * len(neg))424344def _mlp(dims: list[int]) -> nn.Sequential:45    layers = []46    for i in range(len(dims) - 1):47        layers.append(nn.Linear(dims[i], dims[i + 1]))48        if i < len(dims) - 2:49            layers.append(nn.ReLU())50    return nn.Sequential(*layers)515253# ── CGC 抽取层 ────────────────────────────────────────────────────────────────5455class CGCLayer(nn.Module):56    """57    一个 CGC 抽取层。5859    正文说明（PDF §12.4.2.2）：60      「CGC 模块在共享专家之外，为每个任务设置专属专家；61        每个任务的门控混合本任务专属专家和共享专家的输出。」6263    共享专家的输入是上一层聚合的共享表示（首层为原始输入），64    而非某个任务的输出——这是 PLE 与朴素 MMoE 堆叠的关键差异。6566    参数：67      in_dim        : 输入维度（上一层输出或原始特征）68      hidden        : 专家内部隐藏/输出维度69      n_task_experts: 每个任务的专属专家数70      n_shared      : 共享专家数71      n_tasks       : 任务数72    """7374    def __init__(self, in_dim: int, hidden: int,75                 n_task_experts: int, n_shared: int, n_tasks: int):76        super().__init__()77        self.n_tasks = n_tasks78        self.n_task_exp = n_task_experts79        self.n_shared = n_shared8081        # 任务专属专家: n_tasks 组，每组 n_task_experts 个82        self.task_experts = nn.ModuleList([83            nn.ModuleList([84                nn.Sequential(nn.Linear(in_dim, hidden), nn.ReLU(),85                              nn.Linear(hidden, hidden), nn.ReLU())86                for _ in range(n_task_experts)87            ])88            for _ in range(n_tasks)89        ])9091        # 共享专家（接收共享输入）92        self.shared_experts = nn.ModuleList([93            nn.Sequential(nn.Linear(in_dim, hidden), nn.ReLU(),94                          nn.Linear(hidden, hidden), nn.ReLU())95            for _ in range(n_shared)96        ])9798        # 每个任务的门控：任务输入 → 任务专属(n_task_experts) + 共享(n_shared) → softmax99        gate_out = n_task_experts + n_shared100        self.task_gates = nn.ModuleList([101            nn.Linear(in_dim, gate_out) for _ in range(n_tasks)])102103        # 共享门控：共享输入 → n_shared → softmax104        self.shared_gate = nn.Linear(in_dim, n_shared)105106    def forward(self, task_inputs: list[torch.Tensor],107                shared_input: torch.Tensor108                ) -> tuple[list[torch.Tensor], torch.Tensor]:109        """110        task_inputs  : list of T tensors, 每个 (B, in_dim)111                       第 k 个元素是任务 k 的上一层输出（首层均为原始输入）112        shared_input : (B, in_dim) 上一层聚合的共享表示（首层为原始输入）113        返回:114          new_task_outputs : list of T tensors, 每个 (B, hidden)115          new_shared       : (B, hidden)116        """117        # 共享专家：用共享输入计算一次，所有任务共享118        shared_outs = [e(shared_input) for e in self.shared_experts]119        shared_stack = torch.stack(shared_outs, dim=1)  # (B, n_shared, H)120121        # 共享门控 → 聚合共享专家为下一层的共享输入122        sg = F.softmax(self.shared_gate(shared_input), dim=1)  # (B, n_shared)123        new_shared = (sg.unsqueeze(-1) * shared_stack).sum(dim=1)  # (B, H)124125        new_task_outputs = []126        for k in range(self.n_tasks):127            # 任务 k 的专属专家（用任务 k 自己的上一层输入）128            task_outs = [e(task_inputs[k]) for e in self.task_experts[k]]129130            # 任务专属专家 + 共享专家拼在一起供门控选择131            all_outs = torch.stack(task_outs + shared_outs, dim=1)  # (B, E_k+S, H)132133            g = F.softmax(self.task_gates[k](task_inputs[k]), dim=1)  # (B, E_k+S)134            f_k = (g.unsqueeze(-1) * all_outs).sum(dim=1)            # (B, H)135            new_task_outputs.append(f_k)136137        return new_task_outputs, new_shared138139140# ── PLE 模型 ─────────────────────────────────────────────────────────────────141142class PLE(nn.Module):143    """144    PLE：逐层堆叠 CGC，再接各任务塔网络。145146    正文（PDF §12.4.2.2）：147      「PLE 通过多层网络逐步完成信息提取与任务分离，148        随层深递增，减少任务间的负迁移（Negative Transfer）。」149    """150151    def __init__(self, in_dim: int, n_tasks: int,152                 hidden: int = 16, n_layers: int = 2,153                 n_task_experts: int = 2, n_shared: int = 2,154                 tower_hidden: int = 16):155        super().__init__()156        self.n_tasks = n_tasks157        self.n_layers = n_layers158159        # 逐层 CGC（各层维度相同，方便堆叠）160        self.cgc_layers = nn.ModuleList()161        cur_dim = in_dim162        for _ in range(n_layers):163            self.cgc_layers.append(164                CGCLayer(cur_dim, hidden, n_task_experts, n_shared, n_tasks))165            cur_dim = hidden166167        # 各任务塔网络168        self.towers = nn.ModuleList([169            _mlp([cur_dim, tower_hidden, 1]) for _ in range(n_tasks)])170171    def forward(self, x: torch.Tensor) -> torch.Tensor:172        # 初始时各任务输入与共享输入均为原始特征173        task_inputs = [x] * self.n_tasks174        shared_input = x175176        # 逐层 CGC177        for layer in self.cgc_layers:178            task_inputs, shared_input = layer(task_inputs, shared_input)179180        # 各任务塔181        return torch.cat([self.towers[k](task_inputs[k])182                          for k in range(self.n_tasks)], dim=1)  # (B, T)183184    def fit(self, x, ys, epochs: int = 100, lr: float = 5e-3,185            task_weights=None):186        w = task_weights or [1.0] * self.n_tasks187        opt = optim.Adam(self.parameters(), lr=lr, weight_decay=1e-5)188        for ep in range(epochs):189            opt.zero_grad()190            logits = self.forward(x)191            loss = sum(192                w[k] * F.binary_cross_entropy_with_logits(193                    logits[:, k], ys[:, k].float())194                for k in range(self.n_tasks))195            loss.backward()196            opt.step()197            if (ep + 1) % 20 == 0:198                print(f"  epoch {ep+1:3d}  total_loss={loss.item():.4f}")199200    def predict_proba(self, x) -> torch.Tensor:201        with torch.no_grad():202            return torch.sigmoid(self.forward(x))203204205# ── 合成数据 ──────────────────────────────────────────────────────────────────206207def make_data(n: int = 2400, in_dim: int = 16, seed: int = 42):208    """209    两任务的输入特征互不重叠：210      任务 0（点击）依赖 x0–x3 的线性与 x1*x2 交互；211      任务 1（完播）依赖 x6–x9 的线性与 x7*x8 交互。212    每个任务都有自己的专属信号，正是 CGC/PLE 的任务专属专家213    要捕捉的对象；共享专家则负责跨任务的共性表示。214    与 mmoe.py 共用同一份数据生成逻辑，方便公平对比。215    """216    torch.manual_seed(seed)217    random.seed(seed)218    x = torch.randn(n, in_dim)219220    l0 = 2.2 * x[:, 0] + 1.8 * (x[:, 1] * x[:, 2]) + 1.0 * x[:, 3] - 0.3221    l1 = 2.2 * x[:, 6] + 1.8 * (x[:, 7] * x[:, 8]) + 1.0 * x[:, 9] - 0.3222223    y0 = (torch.sigmoid(l0 + 0.35 * torch.randn(n)) > 0.5).long()224    y1 = (torch.sigmoid(l1 + 0.35 * torch.randn(n)) > 0.5).long()225226    return x, torch.stack([y0, y1], dim=1)227228229# ── 主演示 ────────────────────────────────────────────────────────────────────230231if __name__ == "__main__":232    random.seed(42)233    torch.manual_seed(42)234235    IN_DIM, N_TASKS = 16, 2236    x, y = make_data(2400, IN_DIM)237    split = 1920238    x_tr, y_tr = x[:split], y[:split]239    x_te, y_te = x[split:], y[split:]240241    # ── 对比基线 1：Shared-Bottom（与 mmoe.py 配置一致）───────────────────────242    from mmoe import MMoE, SharedBottom243    torch.manual_seed(42)244    sb = SharedBottom(IN_DIM, N_TASKS, bottom_hidden=8, tower_hidden=16)245    opt_sb = optim.Adam(sb.parameters(), lr=5e-3, weight_decay=1e-5)246    for _ in range(120):247        opt_sb.zero_grad()248        logits = sb(x_tr)249        loss = sum(F.binary_cross_entropy_with_logits(250            logits[:, k], y_tr[:, k].float()) for k in range(N_TASKS))251        loss.backward()252        opt_sb.step()253    with torch.no_grad():254        p_sb = torch.sigmoid(sb(x_te))255    auc_sb = [_auc(p_sb[:, k].tolist(), y_te[:, k].tolist())256              for k in range(N_TASKS)]257    print(f"\nShared-Bottom  Task0 AUC={auc_sb[0]:.4f}  Task1 AUC={auc_sb[1]:.4f}")258259    # ── 对比基线 2：MMoE（来自正文 §12.4.2.1 的对比） ──────────────────────────260    torch.manual_seed(42)261    mmoe = MMoE(IN_DIM, n_experts=4, n_tasks=N_TASKS,262                expert_hidden=8, tower_hidden=16)263    mmoe.fit(x_tr, y_tr, epochs=120, lr=5e-3)264    p_mmoe = mmoe.predict_proba(x_te)265    auc_mmoe = [_auc(p_mmoe[:, k].tolist(), y_te[:, k].tolist())266                for k in range(N_TASKS)]267    print(f"\nMMoE    Task0 AUC={auc_mmoe[0]:.4f}  Task1 AUC={auc_mmoe[1]:.4f}")268269    # ── PLE（2 层 CGC，每任务 2 专属专家 + 2 共享专家）─────────────────────────270    torch.manual_seed(42)271    ple = PLE(IN_DIM, N_TASKS, hidden=8, n_layers=2,272              n_task_experts=2, n_shared=2, tower_hidden=16)273    print(f"\nPLE 训练（2 层 CGC，各任务 2 专属 + 2 共享专家，专家宽度 8）")274    ple.fit(x_tr, y_tr, epochs=100, lr=5e-3)275276    p_ple = ple.predict_proba(x_te)277    auc_ple = [_auc(p_ple[:, k].tolist(), y_te[:, k].tolist())278               for k in range(N_TASKS)]279    print(f"\nPLE     Task0 AUC={auc_ple[0]:.4f}  Task1 AUC={auc_ple[1]:.4f}")280281    # ── 三模型对比 ────────────────────────────────────────────────────────────282    print(f"\n对比（正文 §12.4.2.2）：")283    print(f"  Shared-Bottom avg AUC={sum(auc_sb)/2:.4f}")284    print(f"  MMoE         avg AUC={sum(auc_mmoe)/2:.4f}")285    print(f"  PLE          avg AUC={sum(auc_ple)/2:.4f}")286    print("  PLE 通过任务专属专家 + 共享专家逐层分离，")287    print("  在任务特征支持不重叠时进一步减少负迁移，超越 MMoE。")288289    # ── 门控分析（PLE 最后一层 CGC 的任务专属专家权重）──────────────────────────290    print(f"\n最后一层 CGC 门控权重分析（前 20 样本平均）：")291    with torch.no_grad():292        # 取得倒数第二层（即最后 CGC 的输入）293        task_in = [x_te[:20]] * N_TASKS294        shared_in = x_te[:20]295        for layer in ple.cgc_layers[:-1]:296            task_in, shared_in = layer(task_in, shared_in)297        last_cgc = ple.cgc_layers[-1]298        n_te = last_cgc.n_task_exp299        n_sh = last_cgc.n_shared300        for k in range(N_TASKS):301            g = F.softmax(last_cgc.task_gates[k](task_in[k]), dim=1)302            g_mean = g.mean(dim=0)303            task_w = g_mean[:n_te].sum().item()304            shared_w = g_mean[n_te:].sum().item()305            print(f"  Task{k}: 专属专家权重合计={task_w:.3f}  "306                  f"共享专家权重合计={shared_w:.3f}")307    print("  正文（PDF §12.4.2.2）：任务差异越大，专属专家权重占比应越高。")308309    # ── 精排推理演示（任务 0）─────────────────────────────────────────────────310    scores_t0 = p_ple[:20, 0].tolist()311    top5 = sorted(enumerate(scores_t0), key=lambda t: -t[1])[:5]312    print("\n按任务0（点击）精排 Top-5（候选集 20 条）:")313    for rank_i, (idx, s) in enumerate(top5, 1):314        print(f"  #{rank_i}  item={idx:2d}  p_click={s:.4f}  "315              f"label_click={y_te[idx, 0].item()}")