sdm.py

1"""2SDM 长短期兴趣序列召回（第10章 §10.4.4「SDM召回」）34论文：Lv et al., "SDM: Sequential Deep Matching Model for Online Large-scale5      Recommender System", CIKM 2019（阿里巴巴）6核心思想（长短期兴趣分离 + 门控融合）：7- 短期兴趣：最近 Session 行为序列 → LSTM → 多头自注意力 → 以用户画像 Embedding8  e_u 为 query 的 User-Attention Pooling，得到短期兴趣向量 s；9- 长期兴趣：长期行为按多个特征维度（item / 类目 / 店铺 / 品牌）分别 Attention10  Pooling 得到 z_f，拼接后经全连接融合得到长期兴趣向量 p；11- 门控融合（Gated Fusion，SDM 的核心）：12      G = sigmoid(W1·e_u + W2·s + W3·p + b);  o = (1 - G) ⊙ p + G ⊙ s13- 训练用 Sampled Softmax（下一次点击为正样本 + 负采样），线上用 o 做 ANN 检索。1415本实现：PyTorch autograd，用 nn.LSTM / nn.MultiheadAttention / nn.Linear 实现各模块；16完整反向传播由 autograd 自动完成。17"""1819import math20import random2122import torch23import torch.nn as nn24import torch.nn.functional as F25import torch.optim as optim262728# ── SDM 模型 ──────────────────────────────────────────────────────────────────2930class SDM(nn.Module):31    """32    简化版 SDM：LSTM + 多头自注意力 + User-Attention（短期兴趣），多特征维度33    Attention Pooling + FC（长期兴趣），门控融合输出最终用户 Embedding。3435    参数：item/cate/shop/brand_emb（物品与 side info）、user/age/gender_emb36    （用户画像）、LSTM、MultiheadAttention、W_p、W_g1/g2/g3/b_g。37    """3839    LONG_FIELDS = ("item", "cate", "shop", "brand")4041    def __init__(self, n_items: int, n_cates: int, n_shops: int, n_brands: int,42                 n_users: int, n_ages: int = 3, n_genders: int = 2,43                 dim: int = 16, n_heads: int = 2, lr: float = 0.05, reg: float = 1e-4):44        super().__init__()45        assert dim % n_heads == 0, "dim 必须能被 n_heads 整除"46        self.n_items = n_items47        self.dim = dim48        self.n_heads = n_heads49        self.lr = lr50        self.reg = reg5152        # 物品与 side info Embedding（长期兴趣的多个特征维度）53        self.item_emb = nn.Embedding(n_items, dim)54        self.cate_emb = nn.Embedding(n_cates, dim)55        self.shop_emb = nn.Embedding(n_shops, dim)56        self.brand_emb = nn.Embedding(n_brands, dim)5758        # 用户画像 Embedding：uid + 年龄段 + 性别，加和后经 W_u 投影为 e_u59        self.user_emb = nn.Embedding(n_users, dim)60        self.age_emb = nn.Embedding(n_ages, dim)61        self.gender_emb = nn.Embedding(n_genders, dim)62        self.W_u = nn.Linear(dim, dim)   # 投影（原文拼接三者，简化为加和后投影）6364        # 短期兴趣：LSTM + 多头自注意力65        self.lstm = nn.LSTM(input_size=dim, hidden_size=dim, batch_first=True)66        self.self_attn = nn.MultiheadAttention(embed_dim=dim, num_heads=n_heads,67                                               batch_first=True)6869        # 长期兴趣融合：p = tanh(W_p · Concat(z_item, z_cate, z_shop, z_brand) + b_p)70        self.W_p = nn.Linear(len(self.LONG_FIELDS) * dim, dim)7172        # 门控融合参数：W_g1 作用在 e_u，W_g2 作用在短期 s，W_g3 作用在长期 p73        self.W_g1 = nn.Linear(dim, dim, bias=False)74        self.W_g2 = nn.Linear(dim, dim, bias=False)75        self.W_g3 = nn.Linear(dim, dim)   # 带 bias（b_g）7677        scale = math.sqrt(1.0 / dim)78        for emb in (self.item_emb, self.cate_emb, self.shop_emb, self.brand_emb,79                    self.user_emb, self.age_emb, self.gender_emb):80            nn.init.uniform_(emb.weight, -scale, scale)8182    def _emb_table(self, field: str) -> nn.Embedding:83        return {"item": self.item_emb, "cate": self.cate_emb,84                "shop": self.shop_emb, "brand": self.brand_emb}[field]8586    def user_profile(self, uid: torch.Tensor, age: torch.Tensor,87                     gender: torch.Tensor) -> torch.Tensor:88        """e_u = LayerNorm(W_u(uid_emb + age_emb + gender_emb))，shape [B, dim]。"""89        cat = self.user_emb(uid) + self.age_emb(age) + self.gender_emb(gender)90        return F.layer_norm(self.W_u(cat), [self.dim])9192    # ── 短期兴趣：LSTM → 多头自注意力 → User-Attention ────────────────────93    def short_term(self, session_ids: torch.Tensor,94                   e_u: torch.Tensor) -> torch.Tensor:95        """96        session_ids: [B, T]97        e_u:         [B, dim]98        返回短期兴趣向量 s，shape [B, dim]99        """100        embs = self.item_emb(session_ids)              # [B, T, dim]101        lstm_out, _ = self.lstm(embs)                   # [B, T, dim]102103        # 多头自注意力（不加因果 mask，SDM 短期序列较短，双向看也合理）104        H_hat, _ = self.self_attn(lstm_out, lstm_out, lstm_out)  # [B, T, dim]105        H_hat = F.layer_norm(H_hat + lstm_out, [self.dim])       # 残差 + LN106107        # User-Attention Pooling：以 e_u 为 query，对 H_hat 做软注意力加权求和108        # scores: [B, T] = bmm([B,1,dim], [B,dim,T])109        scores = torch.bmm(e_u.unsqueeze(1), H_hat.transpose(1, 2)).squeeze(1)110        attn_weights = scores.softmax(dim=-1)           # [B, T]111        s = (attn_weights.unsqueeze(-1) * H_hat).sum(dim=1)  # [B, dim]112        return s113114    # ── 长期兴趣：多特征维度 Attention Pooling + FC ───────────────────────115    def long_term(self, long_feats: dict[str, torch.Tensor],116                  e_u: torch.Tensor) -> torch.Tensor:117        """118        long_feats: {field: [B, K_f]}  各特征维度的历史 id 序列119        e_u:        [B, dim]120        返回长期兴趣向量 p，shape [B, dim]121        """122        z_all = []123        for f in self.LONG_FIELDS:124            ids = long_feats[f]                              # [B, K_f]125            g = self._emb_table(f)(ids)                      # [B, K_f, dim]126            # Attention Pooling：scores = g · e_u，α = softmax(scores)127            scores = torch.bmm(g, e_u.unsqueeze(-1)).squeeze(-1)  # [B, K_f]128            alpha = scores.softmax(dim=-1)                         # [B, K_f]129            z_f = (alpha.unsqueeze(-1) * g).sum(dim=1)            # [B, dim]130            z_all.append(z_f)131        concat = torch.cat(z_all, dim=-1)                    # [B, 4*dim]132        p = torch.tanh(self.W_p(concat))                     # [B, dim]133        return p134135    # ── 门控融合（SDM 核心）───────────────────────────────────────────────136    def gated_fusion(self, e_u: torch.Tensor, s: torch.Tensor,137                     p: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:138        """139        G = σ(W1·e_u + W2·s + W3·p + b)；o = (1-G)⊙p + G⊙s140        返回 (o, G)，shape 均为 [B, dim]141        """142        G = torch.sigmoid(self.W_g1(e_u) + self.W_g2(s) + self.W_g3(p))143        o = (1.0 - G) * p + G * s144        return o, G145146    def user_vector(self, sample: dict) -> tuple[torch.Tensor, dict]:147        """前向：画像 → 短期 / 长期 → 门控融合，返回最终用户 Embedding o 和缓存。"""148        p = sample["profile"]149        uid = torch.tensor([p["uid"]], dtype=torch.long)150        age = torch.tensor([p["age"]], dtype=torch.long)151        gender = torch.tensor([p["gender"]], dtype=torch.long)152        e_u = self.user_profile(uid, age, gender)            # [1, dim]153154        session = torch.tensor([sample["session"]], dtype=torch.long)155        s = self.short_term(session, e_u)                    # [1, dim]156157        long = {f: torch.tensor([sample["long"][f] or [0]], dtype=torch.long)158                for f in self.LONG_FIELDS}159        p_vec = self.long_term(long, e_u)                    # [1, dim]160161        o, G = self.gated_fusion(e_u, s, p_vec)             # [1, dim]162        return o.squeeze(0), {"e_u": e_u, "s": s, "p": p_vec, "G": G}163164    # ── 训练（Sampled Softmax）─────────────────────────────────────────────165    def train_model(self, samples: list[dict], epochs: int = 6, neg_k: int = 4) -> None:166        """samples: 每条包含 profile / session / long / target 四个字段。"""167        optimizer = optim.Adam(self.parameters(), lr=self.lr, weight_decay=self.reg)168        for epoch in range(epochs):169            random.shuffle(samples)170            total = 0.0171            for smp in samples:172                pos = smp["target"]173                neg_candidates = []174                while len(neg_candidates) < neg_k:175                    c = random.randrange(self.n_items)176                    if c != pos:177                        neg_candidates.append(c)178                cands = [pos] + neg_candidates                  # [1 + neg_k]179                cand_t = torch.tensor(cands, dtype=torch.long)180181                optimizer.zero_grad()182                o, _ = self.user_vector(smp)                    # [dim]183                logits = self.item_emb(cand_t) @ o              # [1+neg_k]184                label = torch.zeros(1, dtype=torch.long)        # pos is index 0185                loss = F.cross_entropy(logits.unsqueeze(0), label)186                loss.backward()187                optimizer.step()188                total += loss.item()189            print(f"  epoch {epoch:2d}  avg_loss={total / max(len(samples), 1):.4f}")190191    def recommend(self, sample: dict, top_k: int = 5) -> tuple[list[tuple[int, float]], dict]:192        """在线推理：实时算出用户 Embedding o，再对全量物品打分（等价 ANN 检索）。"""193        with torch.no_grad():194            o, cache = self.user_vector(sample)195            scores = self.item_emb.weight @ o                   # [n_items]196            exclude = set(sample["session"]) | set(sample["long"]["item"])197            for iid in exclude:198                scores[iid] = -1e9199            top_vals, top_idx = scores.topk(top_k)200        return ([(int(i), float(s)) for i, s in zip(top_idx.tolist(), top_vals.tolist())],201                cache)202203204# ── 快速验证 ──────────────────────────────────────────────────────────────────205if __name__ == "__main__":206    random.seed(42)207    torch.manual_seed(42)208    n_users, n_items, n_cates, n_shops, n_brands = 8, 20, 5, 4, 6209    item_cate = {i: i % n_cates for i in range(n_items)}210    item_shop = {i: i % n_shops for i in range(n_items)}211    item_brand = {i: i % n_brands for i in range(n_items)}212    samples = []213    for uid in range(n_users):214        profile = {"uid": uid, "age": uid % 3, "gender": uid % 2}215        hist = random.sample(range(n_items), k=12)216        long_items, session = hist[:6], hist[6:]217        long_feats = {218            "item": long_items,219            "cate": sorted({item_cate[i] for i in long_items}),220            "shop": sorted({item_shop[i] for i in long_items}),221            "brand": sorted({item_brand[i] for i in long_items}),222        }223        for t in range(3, len(session)):224            samples.append({"profile": profile, "long": long_feats,225                            "session": session[:t], "target": session[t]})226227    model = SDM(n_items, n_cates, n_shops, n_brands, n_users, dim=16, n_heads=2, lr=0.05)228    print(f"SDM 训练（{len(samples)} 条样本，Sampled Softmax，负样本数=4）：")229    model.train_model(samples, epochs=6, neg_k=4)230231    demo = samples[0]232    recs, cache = model.recommend(demo, top_k=5)233    print(f"\n用户 {demo['profile']['uid']} 最近 Session={demo['session']}"234          f"  长期行为 item={demo['long']['item']}  cate={demo['long']['cate']}")235    print(f"门控向量 G 均值={cache['G'].mean().item():.4f}（越接近 1 越偏向短期兴趣）")236    print("Top-5 召回结果：")237    for iid, sc in recs:238        print(f"  item {iid:2d}  score={sc:.4f}  cate={item_cate[iid]}")239240    with torch.no_grad():241        g_avg = sum(242            model.user_vector(smp)[1]["G"].mean().item() for smp in samples243        ) / len(samples)244    print(f"\n全样本平均门控权重 G={g_avg:.4f}")