mv_dssm.py

1"""2MV-DSSM 多视角多塔粗排模型（第11章 §11.2.1 MV-DSSM模型）34论文出处：Elkahky et al., "A Multi-View Deep Learning Approach for Cross Domain5User Modeling in Recommendation Systems", WWW 2015（微软）。67核心思想：8原始 DSSM 只能建模「一类用户 vs 一类物品」的匹配关系。MV-DSSM 提出9「单用户塔、多物品塔」的多视角（Multi-View）结构：10    e_u          = f_u(x_u)                （所有业务域共享的用户塔）11    e_i^(m)      = f_i^(m)(x_i^(m))        （第 m 个业务域独立的物品塔）12    s_m(u, i)    = cos(e_u, e_i^(m))13训练时每条样本只激活其所属域的物品塔，但所有域的监督信号都会共同优化共享14用户塔，从而实现跨域知识迁移（Knowledge Transfer）：数据量大的域帮助数据量15小的新域学到更鲁棒的用户兴趣表示。1617本实现：PyTorch autograd，三个业务域（新闻 / App / 视频），每域各自的物品18Embedding 空间与物品塔维度可以不同（支持异构输入特征）。19一个共享用户塔 + 多个域独立 item 塔（nn.ModuleDict 存各域 item tower）。20演示同一 u_emb 在三个域各出 Top-3。21"""2223import math24import random25import collections26from typing import List, Tuple, Dict, Optional2728import torch29import torch.nn as nn30import torch.nn.functional as F31import torch.optim as optim323334class TowerMLP(nn.Module):35    """通用 MLP 塔：多层线性 + ReLU，最后一层无激活。"""3637    def __init__(self, dims: List[int]):38        super().__init__()39        layers = []40        for in_d, out_d in zip(dims[:-1], dims[1:]):41            layers.append(nn.Linear(in_d, out_d))42        self.layers = nn.ModuleList(layers)4344    def forward(self, x: torch.Tensor) -> torch.Tensor:45        h = x46        for i, layer in enumerate(self.layers):47            h = layer(h)48            if i < len(self.layers) - 1:49                h = F.relu(h)50        return h515253class ItemView(nn.Module):54    """单个业务域（视角）：独立的物品 Embedding 表 + 独立的物品塔。"""5556    def __init__(57        self,58        name: str,59        n_items: int,60        emb_dim: int,61        hidden_dims: List[int],62        output_dim: int,63    ):64        super().__init__()65        self.name = name66        self.n_items = n_items67        # 各域物品塔参数完全独立（输入维度可不同 → 支持异构特征）68        self.item_emb = nn.Embedding(n_items, emb_dim)69        self.tower = TowerMLP([emb_dim] + hidden_dims + [output_dim])70        scale = math.sqrt(1.0 / emb_dim)71        nn.init.uniform_(self.item_emb.weight, -scale, scale)7273    def encode(self, iid: torch.Tensor) -> torch.Tensor:74        return self.tower(self.item_emb(iid))757677class MVDSSM(nn.Module):78    """79    MV-DSSM 模型：一个共享用户塔 + 多个业务域独立物品塔（nn.ModuleDict）。8081    所有域的物品塔输出维度必须与用户塔输出维度一致，才能做余弦匹配。82    """8384    def __init__(85        self,86        n_users: int,87        view_specs: List[Tuple[str, int, int]],88        emb_dim: int = 16,89        hidden_dims: Optional[List[int]] = None,90        output_dim: int = 16,91        temperature: float = 0.1,92    ):93        """view_specs: [(域名, 物品数, 该域物品 Embedding 维度), ...]"""94        super().__init__()95        hidden_dims = hidden_dims or [32, 16]96        self.output_dim = output_dim97        self.temperature = temperature9899        self.user_emb = nn.Embedding(n_users, emb_dim)100        # 共享用户塔：所有域共用一套参数101        self.user_tower = TowerMLP([emb_dim] + hidden_dims + [output_dim])102        # nn.ModuleDict 存各域 item tower103        self.views = nn.ModuleDict({104            name: ItemView(name, n_items, v_emb, hidden_dims, output_dim)105            for name, n_items, v_emb in view_specs106        })107108        scale = math.sqrt(1.0 / emb_dim)109        nn.init.uniform_(self.user_emb.weight, -scale, scale)110111    def encode_user(self, uid: torch.Tensor) -> torch.Tensor:112        """e_u = f_u(x_u)，跨域共享，一次计算可服务所有域。"""113        return self.user_tower(self.user_emb(uid))114115    def forward(116        self,117        uid: torch.Tensor,118        view_name: str,119        pos_iid: torch.Tensor,120        neg_iids: torch.Tensor,121    ) -> torch.Tensor:122        """123        Sampled Softmax：只激活当前域的物品塔，124        但共享用户塔会被所有域的样本共同更新 → 跨域知识迁移。125        pos_iid: [B], neg_iids: [B, n_neg]126        """127        view: ItemView = self.views[view_name]128        e_u = self.encode_user(uid)                         # [B, D]129        e_pos = view.encode(pos_iid)                        # [B, D]130        e_neg = view.encode(neg_iids.view(-1)).view(         # [B, n_neg, D]131            neg_iids.size(0), neg_iids.size(1), -1132        )133        # logits: pos + negs → [B, 1+n_neg]134        s_pos = F.cosine_similarity(e_u, e_pos, dim=1) / self.temperature   # [B]135        s_neg = (136            F.cosine_similarity(137                e_u.unsqueeze(1).expand_as(e_neg), e_neg, dim=2138            ) / self.temperature139        )  # [B, n_neg]140        logits = torch.cat([s_pos.unsqueeze(1), s_neg], dim=1)  # [B, 1+n_neg]141        labels = torch.zeros(logits.size(0), dtype=torch.long)  # 第 0 列为正样本142        return F.cross_entropy(logits, labels)143144    def pre_rank(145        self,146        uid: int,147        candidates: List[int],148        top_k: int = 20,149        view: Optional[str] = None,150    ) -> List[Tuple[int, float]]:151        """粗排：指定业务域对候选集打分并返回 Top-K。默认取第一个域。"""152        view = view or next(iter(self.views))153        v: ItemView = self.views[view]154        with torch.no_grad():155            u = torch.tensor([uid])156            e_u = self.encode_user(u)                      # [1, D]157            cand_t = torch.tensor(candidates)158            e_i = v.encode(cand_t)                         # [N, D]159            sims = F.cosine_similarity(e_u, e_i, dim=1)   # [N]160            scores = [(candidates[j], sims[j].item()) for j in range(len(candidates))]161        scores.sort(key=lambda x: x[1], reverse=True)162        return scores[:top_k]163164165def train(166    model: MVDSSM,167    samples: List[Tuple[int, str, int]],168    epochs: int = 6,169    n_neg: int = 4,170    lr: float = 0.05,171) -> None:172    """samples: [(uid, 域名, 正样本 item_id), ...]；多域样本交替进行 SGD。"""173    optimizer = optim.Adam(model.parameters(), lr=lr)174    data = list(samples)175176    # 按域分组，方便快速采样177    view_items: Dict[str, int] = {178        name: v.n_items for name, v in model.views.items()179    }180181    for epoch in range(epochs):182        random.shuffle(data)183        per_view_loss: Dict[str, List[float]] = collections.defaultdict(list)184185        # 按域 mini-batch 处理（每个样本单独一个 step，保持与原逻辑一致）186        for uid, view_name, pos_iid in data:187            v_size = view_items[view_name]188            # 采 n_neg 个随机负样本（不严格去重，保持效率）189            neg_ids = [random.randrange(v_size) for _ in range(n_neg)]190191            uid_t = torch.tensor([uid])192            pos_t = torch.tensor([pos_iid])193            neg_t = torch.tensor([neg_ids])  # [1, n_neg]194195            optimizer.zero_grad()196            loss = model(uid_t, view_name, pos_t, neg_t)197            loss.backward()198            optimizer.step()199            per_view_loss[view_name].append(loss.item())200201        total = sum(sum(v) for v in per_view_loss.values()) / max(len(data), 1)202        detail = "  ".join(203            f"{name}={sum(ls) / len(ls):.4f}(n={len(ls)})"204            for name, ls in sorted(per_view_loss.items())205        )206        print(f"  epoch {epoch:2d}  avg_loss={total:.4f}  |  {detail}")207208209# ── 快速验证 ──────────────────────────────────────────────────────────────────210if __name__ == "__main__":211    random.seed(42)212    torch.manual_seed(42)213214    n_users = 8215    # 三个业务域：新闻样本多，视频域为数据稀疏的新业务 → 依赖跨域迁移216    view_specs = [("news", 24, 16), ("app", 18, 12), ("video", 20, 20)]217    sample_cnt = {"news": 48, "app": 24, "video": 10}218219    samples: List[Tuple[int, str, int]] = []220    for name, n_items, _ in view_specs:221        for _ in range(sample_cnt[name]):222            samples.append((random.randrange(n_users), name, random.randrange(n_items)))223224    model = MVDSSM(n_users, view_specs, emb_dim=16, hidden_dims=[32, 16],225                   output_dim=16, temperature=0.1)226    print("MV-DSSM 多视角训练（共享用户塔 + 3 个独立物品塔）：")227    train(model, samples, epochs=6, n_neg=4, lr=0.05)228229    uid = 0230    print(f"\n同一用户表示 e_u 在多个业务域下分别粗排（用户 {uid}）：")231    for name, n_items, _ in view_specs:232        res = model.pre_rank(uid, list(range(n_items)), top_k=3, view=name)233        pretty = ", ".join(f"item {i}:{s:.4f}" for i, s in res)234        print(f"  [{name:5s}] Top-3 → {pretty}")235236    with torch.no_grad():237        e_u = model.encode_user(torch.tensor([uid]))238        norm = e_u.norm().item()239    print(f"\n共享用户 Embedding 模长 = {norm:.4f}"240          f"（一次计算，复用于全部 {len(view_specs)} 个域的粗排）")