dat.py

1"""2DAT 双塔增强粗排模型（第11章 §11.1.4 DAT模型）34论文出处：Yu et al., "A Dual Augmented Two-tower Model for Online Large-scale5Recommendation", DLP-KDD / CIKM 2021（美团）。67核心思想：8传统双塔模型中用户塔看不到物品信息、物品塔看不到用户信息，两塔之间缺乏交互。9DAT 为每个用户和每个物品各配一个可学习的「增强向量」（Augmented Vector）10a_u / a_v，拼接进本侧塔的输入，让单塔能隐式地「看到」对侧的历史交互语义：11    z_u = [e_u, a_u],  z_v = [e_v, a_v]12    p_u = L2Norm(MLP_u(z_u)),  p_v = L2Norm(MLP_v(z_v))13    s(u, v) = <p_u, p_v>1415三层训练目标（§11.1.4 的 L = L_p + λ1·L_u + λ2·L_v + λ3·L_CA）：161) L_p  主损失：二分类交叉熵，监督匹配准确性；172) L_u / L_v  自适应模仿机制（Adaptive-Mimic Mechanism, AMM）：18       L_u = (1/T)·Σ y·‖a_u − sg(p_v)‖²,   L_v = (1/T)·Σ y·‖a_v − sg(p_u)‖²19   只在正样本（y=1）上拟合，sg(·) 为 Stop Gradient，PyTorch autograd 用 .detach() 实现；203) L_CA 类别对齐损失（Category Alignment Loss）：对齐不同类别 item embedding 的21   二阶统计量（均值差 + 协方差差），缓解类别数据分布不均衡：22       L_CA = Σ_i ‖C(S_major) − C(S_i)‖_F²2324本实现：PyTorch autograd，三项损失联合训练（超参 α, β 固定为合理默认值）。25"""2627import math28import random29import collections30from typing import List, Tuple, Dict, Optional3132import torch33import torch.nn as nn34import torch.nn.functional as F35import torch.optim as optim363738class TowerMLP(nn.Module):39    """双塔 MLP：多层线性 + ReLU，最后一层无激活。"""4041    def __init__(self, dims: List[int]):42        super().__init__()43        layers = []44        for in_d, out_d in zip(dims[:-1], dims[1:]):45            layers.append(nn.Linear(in_d, out_d))46        self.layers = nn.ModuleList(layers)4748    def forward(self, x: torch.Tensor) -> torch.Tensor:49        h = x50        for i, layer in enumerate(self.layers):51            h = layer(h)52            if i < len(self.layers) - 1:53                h = F.relu(h)54        return h555657class DAT(nn.Module):58    """59    DAT（Dual Augmented Two-tower）粗排模型。6061    结构：user_emb + a_u → 用户塔 → L2Norm → p_u62         item_emb + a_v → 物品塔 → L2Norm → p_v63         s(u, v) = <p_u, p_v>64    """6566    def __init__(67        self,68        n_users: int,69        n_items: int,70        item_category: List[int],71        emb_dim: int = 16,72        aug_dim: int = 16,73        hidden_dims: Optional[List[int]] = None,74        output_dim: int = 16,75        lambda_u: float = 0.3,76        lambda_v: float = 0.3,77        lambda_ca: float = 0.05,78    ):79        super().__init__()80        hidden_dims = hidden_dims or [32, 16]81        # 增强向量维度必须等于对侧塔输出维度，模仿损失才能逐元素做差82        assert aug_dim == output_dim, "增强向量维度需与塔输出维度一致（模仿损失要求）"83        self.n_items = n_items84        self.output_dim = output_dim85        self.lambda_u = lambda_u86        self.lambda_v = lambda_v87        self.lambda_ca = lambda_ca8889        # 注册 item_category 为 buffer，方便随模型一起保存/移动设备90        self.register_buffer("item_category", torch.tensor(item_category, dtype=torch.long))9192        self.user_emb = nn.Embedding(n_users, emb_dim)93        self.item_emb = nn.Embedding(n_items, emb_dim)9495        # (a) 双塔增强结构：每个用户/物品各一个可学习增强向量96        self.aug_u = nn.Embedding(n_users, aug_dim)97        self.aug_v = nn.Embedding(n_items, aug_dim)9899        dims = [emb_dim + aug_dim] + hidden_dims + [output_dim]100        self.user_tower = TowerMLP(dims)101        self.item_tower = TowerMLP(dims)102103        # 初始化 Embedding104        scale = math.sqrt(1.0 / emb_dim)105        nn.init.uniform_(self.user_emb.weight, -scale, scale)106        nn.init.uniform_(self.item_emb.weight, -scale, scale)107        nn.init.zeros_(self.aug_u.weight)108        nn.init.zeros_(self.aug_v.weight)109110    def encode_user(self, uid: torch.Tensor) -> torch.Tensor:111        """z_u = [e_u, a_u] → MLP → L2Norm  →  p_u  [B, D]"""112        z = torch.cat([self.user_emb(uid), self.aug_u(uid)], dim=-1)113        h = self.user_tower(z)114        return F.normalize(h, p=2, dim=-1)115116    def encode_item(self, iid: torch.Tensor) -> torch.Tensor:117        """z_v = [e_v, a_v] → MLP → L2Norm  →  p_v  [B, D]"""118        z = torch.cat([self.item_emb(iid), self.aug_v(iid)], dim=-1)119        h = self.item_tower(z)120        return F.normalize(h, p=2, dim=-1)121122    # ── (b) 自适应模仿损失 ────────────────────────────────────────────────────123    def _mimic_loss(124        self,125        uid: torch.Tensor,126        iid: torch.Tensor,127        y: torch.Tensor,128        p_u: torch.Tensor,129        p_v: torch.Tensor,130    ) -> torch.Tensor:131        """132        L_u = Σ y·‖a_u − sg(p_v)‖²,  L_v = Σ y·‖a_v − sg(p_u)‖²133        只在正样本上拟合；监督目标走 stop_gradient（.detach()），134        梯度只更新增强向量 aug_u / aug_v 本身。135        """136        mask = (y > 0.5).float()  # [B]137        z_v = p_v.detach()        # stop_gradient138        z_u = p_u.detach()        # stop_gradient139        a_u = self.aug_u(uid)140        a_v = self.aug_v(iid)141        loss_u = ((a_u - z_v).pow(2).sum(dim=-1) * mask).sum()142        loss_v = ((a_v - z_u).pow(2).sum(dim=-1) * mask).sum()143        n = mask.sum().clamp(min=1.0)144        return (loss_u + loss_v) / n145146    # ── (c) 类别对齐损失 ──────────────────────────────────────────────────────147    def _category_alignment_loss(148        self,149        pos_iids: torch.Tensor,150        p_v_pos: torch.Tensor,151        cov_dim: int = 6,152    ) -> torch.Tensor:153        """154        L_CA = Σ_{i≥2} ‖C(S_major) − C(S_i)‖_F²155        C(S) = (1/|S|)·Σ (p − μ)(p − μ)^T 为类别内表示的协方差矩阵。156        对齐不同类别 item embedding 的二阶统计量（均值差 + 协方差差）。157        """158        cats = self.item_category[pos_iids]  # [M]159        unique_cats = cats.unique()160        reprs: Dict[int, torch.Tensor] = {}161        for c in unique_cats.tolist():162            mask = cats == c163            if mask.sum() >= 2:164                reprs[c] = p_v_pos[mask, :cov_dim]  # [k, cov_dim]165166        if len(reprs) < 2:167            return torch.tensor(0.0)168169        major = max(reprs, key=lambda c: reprs[c].size(0))170171        def cov(r: torch.Tensor) -> torch.Tensor:172            mu = r.mean(dim=0, keepdim=True)173            diff = r - mu174            return diff.t().mm(diff) / r.size(0)  # [cov_dim, cov_dim]175176        C_major = cov(reprs[major])177        total = torch.tensor(0.0)178        for c, r in reprs.items():179            if c == major:180                continue181            total = total + (C_major - cov(r)).pow(2).sum()182        return total183184    # ── (d) 总训练目标 ────────────────────────────────────────────────────────185    def forward(186        self,187        uid: torch.Tensor,188        iid: torch.Tensor,189        y: torch.Tensor,190    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:191        """192        L = L_p + λ1·L_u + λ2·L_v + λ3·L_CA193        Returns: (total_loss, L_p, L_mimic, L_CA)194        """195        p_u = self.encode_user(uid)   # [B, D]196        p_v = self.encode_item(iid)   # [B, D]197198        # L_p：二分类交叉熵199        logits = (p_u * p_v).sum(dim=-1)  # [B]200        loss_p = F.binary_cross_entropy_with_logits(logits, y)201202        # L_mimic：AMM 损失（b）203        loss_mimic = self._mimic_loss(uid, iid, y, p_u, p_v)204205        # L_CA：类别对齐损失（c）—— 只在正样本物品上计算206        pos_mask = y > 0.5207        if pos_mask.any():208            loss_ca = self._category_alignment_loss(iid[pos_mask], p_v[pos_mask])209        else:210            loss_ca = torch.tensor(0.0)211212        total = loss_p + (self.lambda_u + self.lambda_v) * loss_mimic + self.lambda_ca * loss_ca213        return total, loss_p, loss_mimic, loss_ca214215    def pre_rank(216        self, uid: int, candidates: List[int], top_k: int = 20217    ) -> List[Tuple[int, float]]:218        """粗排：p_v 可离线预计算，线上仅算 p_u 与内积。"""219        with torch.no_grad():220            u = torch.tensor([uid])221            p_u = self.encode_user(u)  # [1, D]222            cand_t = torch.tensor(candidates)223            p_v = self.encode_item(cand_t)  # [N, D]224            scores_t = (p_u * p_v).sum(dim=-1)  # [N]225            scores = [(candidates[j], scores_t[j].item()) for j in range(len(candidates))]226        scores.sort(key=lambda x: x[1], reverse=True)227        return scores[:top_k]228229230def train(231    model: DAT,232    samples: List[Tuple[int, int, float]],233    epochs: int = 6,234    batch_size: int = 16,235    lr: float = 0.05,236) -> None:237    """L = L_p + λ1·L_u + λ2·L_v + λ3·L_CA"""238    optimizer = optim.Adam(model.parameters(), lr=lr)239    data = list(samples)240    for epoch in range(epochs):241        random.shuffle(data)242        acc_p = acc_m = acc_ca = acc_total = 0.0243        n_batch = 0244        for start in range(0, len(data), batch_size):245            batch = data[start: start + batch_size]246            if not batch:247                continue248            uids = torch.tensor([s[0] for s in batch])249            iids = torch.tensor([s[1] for s in batch])250            ys = torch.tensor([s[2] for s in batch], dtype=torch.float)251252            optimizer.zero_grad()253            total, lp, lm, lca = model(uids, iids, ys)254            total.backward()255            optimizer.step()256257            acc_total += total.item()258            acc_p += lp.item()259            acc_m += lm.item()260            acc_ca += lca.item() if isinstance(lca, torch.Tensor) else lca261            n_batch += 1262263        nb = max(n_batch, 1)264        print(265            f"  epoch {epoch:2d}  L_total={acc_total / nb:.4f}  L_p={acc_p / nb:.4f}  "266            f"L_mimic={acc_m / nb:.4f}  L_CA={acc_ca / nb:.4f}"267        )268269270# ── 快速验证 ──────────────────────────────────────────────────────────────────271if __name__ == "__main__":272    random.seed(42)273    torch.manual_seed(42)274275    n_users, n_items = 8, 24276    # 类别分布故意不均衡：类别 0 为主类别（长尾类别样本很少）→ 触发 CAL277    item_category = [0 if i < 14 else (1 if i < 20 else 2) for i in range(n_items)]278279    samples: List[Tuple[int, int, float]] = []280    for uid in range(n_users):281        pos = {(uid * 3 + k) % n_items for k in range(4)}282        for iid in pos:283            samples.append((uid, iid, 1.0))284        for _ in range(4):285            neg = random.randint(0, n_items - 1)286            if neg not in pos:287                samples.append((uid, neg, 0.0))288289    model = DAT(290        n_users, n_items, item_category,291        emb_dim=16, aug_dim=16, hidden_dims=[32, 16], output_dim=16,292        lambda_u=0.3, lambda_v=0.3, lambda_ca=0.05,293    )294    print("DAT 模型训练（主损失 + 模仿损失 + 类别对齐损失）：")295    train(model, samples, epochs=6, batch_size=16, lr=0.05)296297    uid = 0298    seen = {iid for u, iid, y in samples if u == uid and y > 0}299    candidates = [i for i in range(n_items) if i not in seen]300    print(f"\n用户 {uid} 粗排 Top-5：")301    for iid, s in model.pre_rank(uid, candidates, top_k=5):302        print(f"  item {iid:2d} (类别 {item_category[iid]}): {s:.4f}")303304    # 演示 L_mimic 效果：增强向量 a_u 的模长反映了 AMM 拟合程度305    with torch.no_grad():306        a_u_norm = model.aug_u(torch.tensor([uid])).norm().item()307    print(308        f"\n用户 {uid} 增强向量 a_u 模长 = {a_u_norm:.4f}"309        f"（由 AMM 拟合正样本物品表示 sg(p_v) 得到）"310    )