pfd.py

1"""2PFD 粗排蒸馏模型（第11章 §11.3.3）34PFD = Privileged Feature Distillation（特权特征蒸馏）5论文出处：Xu et al., "Privileged Features Distillation at Taobao Recommendations", KDD 2020。67核心思想（对应书中 §11.3.3）：8    训练阶段可以拿到一批「在线推理阶段不可得 / 代价过高」的特权特征 x_privileged9    （例如详情页停留时长、是否查看评论、是否咨询客服等用户后验行为，10      以及用户-商品的复杂实时交叉统计特征）。11    PFD 让教师模型同时吃「在线特征 + 特权特征」，学生模型只吃在线特征，12    通过蒸馏把特权信息迁移给学生，在不增加任何在线推理开销的前提下提升粗排效果。13    目标函数（书中公式）：14        L = (1-λ)·L_S(y, f_S(x)) + λ·L_KD(f_T(x, x_priv), f_S(x)) + L_T(y, f_T(x, x_priv))1516本实现同时覆盖书中强调的两项工程机制：17    (a) 参数共享（Parameter Sharing）：教师与学生共享 Embedding 层，仅高层预测分支独立；18    (b) 渐进式 Warm-up（Synchronous Training + Warm-up）：训练初期 λ=0 只学各自监督损失，19        待教师收敛后 λ 随 epoch 线性上升到 λ_max，避免早期教师给出不稳定监督信号。2021实现方式：PyTorch autograd，nn.Embedding 由教师与学生两个分支共享，22梯度阻断用 .detach()（教师只被真实标签 L_T 优化）。23可直接 `python3 pfd.py` 运行。24"""25import math26import random27from typing import List, Tuple, Dict2829import torch30import torch.nn as nn31import torch.nn.functional as F32import torch.optim as optim333435# ── 工具：AUC ─────────────────────────────────────────────────────────────────3637def _auc(pairs: List[Tuple[float, int]]) -> float:38    """AUC：正负样本对中被正确排序的比例。"""39    pos = [s for s, y in pairs if y == 1]40    neg = [s for s, y in pairs if y == 0]41    if not pos or not neg:42        return 0.543    hit = sum(1.0 if p > n else (0.5 if p == n else 0.0) for p in pos for n in neg)44    return hit / (len(pos) * len(neg))454647class Branch(nn.Module):48    """教师 / 学生的高层预测分支（共享 Embedding 之外的独立参数）。"""4950    def __init__(self, in_dim: int, hidden_dim: int = 32):51        super().__init__()52        self.net = nn.Sequential(53            nn.Linear(in_dim, hidden_dim),54            nn.ReLU(),55            nn.Linear(hidden_dim, 16),56            nn.ReLU(),57            nn.Linear(16, 1),58        )5960    def forward(self, x: torch.Tensor) -> torch.Tensor:61        return self.net(x).squeeze(-1)  # [B]626364class PFD(nn.Module):65    """66    PFD 特权特征蒸馏粗排模型。6768    教师 f_T(x, x_priv)、学生 f_S(x) 共享同一对 nn.Embedding(user_emb / item_emb)，69    仅高层 Branch 独立。推理只用学生分支 f_S(x)。70    """7172    def __init__(73        self,74        n_users: int,75        n_items: int,76        emb_dim: int = 16,77        priv_dim: int = 3,78        lr: float = 0.05,79        lam_max: float = 0.8,80        warmup_epochs: int = 2,81        temperature: float = 1.0,82    ):83        super().__init__()84        self.n_users, self.n_items = n_users, n_items85        self.lr = lr86        self.lam_max = lam_max87        self.warmup_epochs = warmup_epochs88        self.temperature = temperature8990        # (a) 参数共享：教师与学生共用这层 Embedding91        self.user_emb = nn.Embedding(n_users, emb_dim)92        self.item_emb = nn.Embedding(n_items, emb_dim)93        scale = math.sqrt(1.0 / emb_dim)94        nn.init.uniform_(self.user_emb.weight, -scale, scale)95        nn.init.uniform_(self.item_emb.weight, -scale, scale)9697        # 教师吃 [在线 emb, 特权特征]，学生只吃 [在线 emb]98        self.teacher_branch = Branch(emb_dim * 2 + priv_dim)99        self.student_branch = Branch(emb_dim * 2)100101    # ── 在线特征（用户 / 物品 Embedding 拼接）──────────────────────────────102    def online_emb(self, uid: torch.Tensor, iid: torch.Tensor) -> torch.Tensor:103        return torch.cat([self.user_emb(uid), self.item_emb(iid)], dim=-1)  # [B, 2D]104105    def teacher_logit(self, uid: torch.Tensor, iid: torch.Tensor, priv: torch.Tensor) -> torch.Tensor:106        x = torch.cat([self.online_emb(uid, iid), priv], dim=-1)107        return self.teacher_branch(x)108109    def student_logit(self, uid: torch.Tensor, iid: torch.Tensor) -> torch.Tensor:110        return self.student_branch(self.online_emb(uid, iid))111112    # ── (b) warm-up：λ 随 epoch 线性上升 ────────────────────────────────────113    def lam_at(self, epoch: int, epochs: int) -> float:114        if epoch < self.warmup_epochs:115            return 0.0116        span = max(epochs - self.warmup_epochs, 1)117        return self.lam_max * min(1.0, (epoch - self.warmup_epochs + 1) / span)118119    # ── KD 损失：温度缩放后对教师软标签做 KL ────────────────────────────────120    def _kd_loss(self, s_T: torch.Tensor, s_S: torch.Tensor) -> torch.Tensor:121        # 把标量 logit 扩展成二分类 logits [B, 2] = [0, s]，得到合法概率分布122        T = self.temperature123        logits_T = torch.stack([torch.zeros_like(s_T), s_T / T], dim=-1)  # [B, 2]124        logits_S = torch.stack([torch.zeros_like(s_S), s_S / T], dim=-1)125        p_T = F.softmax(logits_T, dim=-1)126        log_p_S = F.log_softmax(logits_S, dim=-1)127        # .detach()：教师不被 KD 梯度更新，只被真实标签 L_T 优化（梯度阻断）128        return F.kl_div(log_p_S, p_T.detach(), reduction='batchmean') * (T * T)129130    def forward(131        self,132        uid: torch.Tensor,133        iid: torch.Tensor,134        y: torch.Tensor,135        priv: torch.Tensor,136        lam: float,137    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:138        """139        L = (1-λ)·L_S + λ·L_KD + L_T140        Returns: (total, L_S, L_KD, L_T)141        """142        s_T = self.teacher_logit(uid, iid, priv)143        s_S = self.student_logit(uid, iid)144145        loss_T = F.binary_cross_entropy_with_logits(s_T, y)146        loss_S = F.binary_cross_entropy_with_logits(s_S, y)147        loss_KD = self._kd_loss(s_T, s_S)148149        total = (1 - lam) * loss_S + lam * loss_KD + loss_T150        return total, loss_S, loss_KD, loss_T151152    # ── 训练：Synchronous Training ──────────────────────────────────────────153    def train_pfd(154        self,155        samples: List[Tuple[int, int, int, List[float]]],156        epochs: int = 6,157        batch_size: int = 32,158        verbose: bool = True,159    ) -> None:160        optimizer = optim.Adam(self.parameters(), lr=self.lr)161        data = list(samples)162        for epoch in range(epochs):163            lam = self.lam_at(epoch, epochs)164            random.shuffle(data)165            acc = [0.0, 0.0, 0.0, 0.0]166            nb = 0167            for start in range(0, len(data), batch_size):168                batch = data[start: start + batch_size]169                uids = torch.tensor([s[0] for s in batch])170                iids = torch.tensor([s[1] for s in batch])171                ys = torch.tensor([s[2] for s in batch], dtype=torch.float)172                priv = torch.tensor([s[3] for s in batch], dtype=torch.float)173174                optimizer.zero_grad()175                total, l_s, l_kd, l_t = self(uids, iids, ys, priv, lam)176                total.backward()177                optimizer.step()178                for k, v in enumerate((total, l_s, l_kd, l_t)):179                    acc[k] += v.item()180                nb += 1181            nb = max(nb, 1)182            if verbose and epoch % 2 == 0:183                print(f"  epoch {epoch:2d}  λ={lam:.3f}  "184                      f"L={acc[0] / nb:.4f} (L_S={acc[1] / nb:.4f}  "185                      f"L_KD={acc[2] / nb:.4f}  L_T={acc[3] / nb:.4f})")186187    # ── 推理接口（只用学生分支，零特权特征、零额外在线开销）────────────────188    def score(self, uid: int, iid: int) -> float:189        with torch.no_grad():190            s = self.student_logit(torch.tensor([uid]), torch.tensor([iid]))191            return torch.sigmoid(s).item()192193    def teacher_score(self, uid: int, iid: int, priv: List[float]) -> float:194        with torch.no_grad():195            s = self.teacher_logit(196                torch.tensor([uid]), torch.tensor([iid]),197                torch.tensor([priv], dtype=torch.float),198            )199            return torch.sigmoid(s).item()200201    def pre_rank(202        self, uid: int, candidates: List[int], top_k: int = 10203    ) -> List[Tuple[int, float]]:204        with torch.no_grad():205            u = torch.tensor([uid]).expand(len(candidates))206            c = torch.tensor(candidates)207            s = torch.sigmoid(self.student_logit(u, c))208            scores = [(candidates[j], s[j].item()) for j in range(len(candidates))]209        scores.sort(key=lambda x: x[1], reverse=True)210        return scores[:top_k]211212213def make_samples(n_users: int, n_items: int, lat_dim: int = 4):214    """215    构造带特权特征的合成样本：在线特征即用户/物品 Embedding；216    特权特征 = [详情页停留时长, 是否查看评论, 是否咨询客服]，均为点击之后才可得的后验行为。217    """218    def _dot(a, b):219        return sum(x * y for x, y in zip(a, b))220221    def _sigmoid(x):222        return 1.0 / (1.0 + math.exp(-max(-30.0, min(30.0, x))))223224    lat_u = [[random.gauss(0, 1) for _ in range(lat_dim)] for _ in range(n_users)]225    lat_i = [[random.gauss(0, 1) for _ in range(lat_dim)] for _ in range(n_items)]226    samples = []227    for u in range(n_users):228        for i in range(n_items):229            s = _dot(lat_u[u], lat_i[i]) / 2.0230            y = 1 if random.random() < _sigmoid(s) else 0231            dwell = max(0.0, s + random.gauss(0, 0.3)) if y == 1 else max(0.0, s - 1.0 + random.gauss(0, 0.3))232            comment = 1.0 if (y == 1 and random.random() < 0.8) or random.random() < 0.15 else 0.0233            consult = 1.0 if (y == 1 and random.random() < 0.6) else 0.0234            samples.append((u, i, y, [dwell, comment, consult]))235    return samples236237238if __name__ == "__main__":239    random.seed(42)240    torch.manual_seed(42)241    n_users, n_items, emb_dim, epochs = 8, 24, 16, 6242    all_samples = make_samples(n_users, n_items)243244    # 划分训练 / 测试集：蒸馏的收益体现在「泛化」而非训练集拟合上245    random.shuffle(all_samples)246    n_test = len(all_samples) // 4247    test_samples = all_samples[:n_test]248    train_samples = all_samples[n_test:]249    print(f"PFD 特权特征蒸馏：n_users={n_users}  n_items={n_items}  "250          f"训练 {len(train_samples)} 条 / 测试 {len(test_samples)} 条")251252    print("\n[1] PFD 同步训练（Warm-up 2 轮后渐进引入蒸馏损失）")253    pfd = PFD(n_users, n_items, emb_dim=emb_dim, priv_dim=3, lr=0.05,254              lam_max=0.8, warmup_epochs=2)255    pfd.train_pfd(train_samples, epochs=epochs)256257    print("\n[2] 消融：关闭蒸馏（λ≡0，学生只用在线特征+硬标签）")258    torch.manual_seed(42)259    base = PFD(n_users, n_items, emb_dim=emb_dim, priv_dim=3, lr=0.05,260               lam_max=0.0, warmup_epochs=0)261    base.train_pfd(train_samples, epochs=epochs, verbose=False)262263    def auc_of(scorer):264        return _auc([(scorer(u, i, p), y) for u, i, y, p in test_samples])265266    t_auc = auc_of(lambda u, i, p: pfd.teacher_score(u, i, p))267    s_auc = auc_of(lambda u, i, p: pfd.score(u, i))268    b_auc = auc_of(lambda u, i, p: base.score(u, i))269    b_t_auc = auc_of(lambda u, i, p: base.teacher_score(u, i, p))270    print("\n[3] 测试集 AUC 对比")271    print(f"  教师 f_T(x, x_priv)       ：{t_auc:.4f}   ← 吃特权特征，线上不可用")272    print(f"  PFD 学生 f_S(x)           ：{s_auc:.4f}   ← 线上实际部署的粗排模型")273    print(f"  无蒸馏学生 f_S(x), λ=0    ：{b_auc:.4f}")274275    # 与教师的排序一致性：学生 Top-5 与教师 Top-5 的平均重合度276    priv_of = {(u, i): p for u, i, _, p in all_samples}277278    def agree_with_teacher(student: "PFD", k: int = 5) -> float:279        tot = 0.0280        for u in range(n_users):281            t_top = sorted(range(n_items), key=lambda i: -pfd.teacher_score(u, i, priv_of[(u, i)]))[:k]282            s_top = sorted(range(n_items), key=lambda i: -student.score(u, i))[:k]283            tot += len(set(t_top) & set(s_top)) / k284        return tot / n_users285286    print("\n[4] 学生与教师的排序一致性（Top-5 平均重合度）")287    print(f"  PFD 学生：{agree_with_teacher(pfd):.3f}    无蒸馏学生：{agree_with_teacher(base):.3f}")288    uid = 0289    cands = [i for i in range(n_items) if i not in {i2 for u, i2, y, _ in train_samples if u == uid and y == 1}]290    print(f"\n[5] 用户 {uid} 粗排 Top-5（学生模型，零额外在线开销）")291    for iid, s in pfd.pre_rank(uid, cands, top_k=5):292        print(f"  item {iid:2d}: {s:.4f}")