item2vec.py

1"""2Item2Vec 物品 Embedding 学习（第10章 §10.4.1「List2Vec召回 - Item2Vec模型」）34论文：Barkan & Koenigstein, "Item2Vec: Neural Item Embedding for Collaborative5      Filtering", MLSP 201667核心思想：8- 把一个用户的一次行为列表当作一个「句子」，把物品当作「词」，9  直接用 Word2Vec 的 Skip-Gram with Negative Sampling（SGNS）学习物品向量。10- 共同出现在同一用户行为集合中的物品，向量距离更近（协同过滤信号）。1112与 word2vec 的关键区别（本实现的重点）：13- word2vec：中心词只与固定大小滑动窗口 [t-c, t+c] 内的词构成正样本，14  因此「空间/时间距离」会影响共现关系。15- Item2Vec：丢弃了窗口概念，把一次行为集合 S={i_1,...,i_K} 当作「动态窗口」，16  集合内任意两个物品（i_a != i_b）都构成正样本，与先后顺序无关。17  因此 A->B->C 与 C->B->A 产生完全相同的训练样本，Item2Vec 学到的是18  共现关系（Co-occurrence）而非时序依赖。1920优化目标（对应书中公式）：21    L = -(1/K) * sum_{i=1..K} sum_{j != i} log p(i_j | i_i)22其中 p(i_j|i_i) 用 SGNS 近似：23    log p(j|i) ≈ log σ(u_i · v_j) + sum_{n in Neg} log σ(-u_i · v_n)2425本实现：PyTorch autograd，SGNS 损失用 F.binary_cross_entropy_with_logits 计算。26"""2728import math29import random30from collections import defaultdict3132import torch33import torch.nn as nn34import torch.nn.functional as F35import torch.optim as optim363738# ── Item2Vec 模型 ─────────────────────────────────────────────────────────────3940class Item2Vec(nn.Module):41    """42    Item2Vec（Skip-Gram + 负采样）。4344    参数矩阵：45      item_emb: [n_items × dim]  中心物品向量 u_i（最终用于召回的 Item Embedding）46      ctx_emb:  [n_items × dim]  上下文物品向量 v_j（训练辅助参数）47    """4849    def __init__(50        self,51        n_items: int,52        dim: int = 16,53        lr: float = 0.05,54        n_neg: int = 5,55        reg: float = 1e-5,56    ):57        super().__init__()58        self.n_items = n_items59        self.dim = dim60        self.lr = lr61        self.n_neg = n_neg62        self.reg = reg6364        self.item_emb = nn.Embedding(n_items, dim)65        self.ctx_emb = nn.Embedding(n_items, dim)6667        nn.init.uniform_(self.item_emb.weight, -math.sqrt(1.0 / dim), math.sqrt(1.0 / dim))68        nn.init.zeros_(self.ctx_emb.weight)6970        self._neg_table: list[int] = []7172    # ── 样本构造：整个 session 两两共现（Item2Vec 的核心差异点） ──────────────73    def _build_pairs(self, user_seqs: dict[int, list[int]]) -> list[tuple[int, int]]:74        """75        把每个 session 看作无序集合，集合内任意有序对 (i_a, i_b), a != b 都是正样本。76        注意：这里没有任何窗口大小参数，这正是 Item2Vec 与 word2vec 的区别。77        """78        pairs: list[tuple[int, int]] = []79        for seq in user_seqs.values():80            items = list(dict.fromkeys(seq))  # 集合语义：同一 session 内去重81            for a in range(len(items)):82                for b in range(len(items)):83                    if a != b:84                        pairs.append((items[a], items[b]))85        return pairs8687    def _build_neg_table(self, user_seqs: dict[int, list[int]], table_size: int = 2000) -> None:88        """按 freq^0.75 构造负采样表（与 word2vec 的 unigram 分布一致）。"""89        freq: dict[int, int] = defaultdict(int)90        for seq in user_seqs.values():91            for iid in seq:92                freq[iid] += 193        weights = {iid: (c ** 0.75) for iid, c in freq.items()}94        total = sum(weights.values()) or 1.095        table: list[int] = []96        for iid, w in weights.items():97            cnt = max(1, int(round(table_size * w / total)))98            table.extend([iid] * cnt)99        self._neg_table = table or list(range(self.n_items))100101    def _sample_negs(self, exclude: int, k: int) -> list[int]:102        negs = []103        for _ in range(k * 3):104            cand = random.choice(self._neg_table)105            if cand != exclude:106                negs.append(cand)107            if len(negs) == k:108                break109        while len(negs) < k:110            negs.append((exclude + 1) % self.n_items)111        return negs112113    def forward(self, center_ids: torch.Tensor, pos_ids: torch.Tensor,114                neg_ids: torch.Tensor) -> torch.Tensor:115        """116        center_ids: [B]117        pos_ids:    [B]118        neg_ids:    [B, n_neg]119120        loss = BCE(u·v_pos, 1) + sum_neg BCE(u·v_neg, 0)121             = F.binary_cross_entropy_with_logits on stacked logits122        """123        u = self.item_emb(center_ids)           # [B, dim]124        v_pos = self.ctx_emb(pos_ids)           # [B, dim]125        v_neg = self.ctx_emb(neg_ids)           # [B, n_neg, dim]126127        pos_logits = (u * v_pos).sum(-1)        # [B]128        neg_logits = torch.bmm(v_neg, u.unsqueeze(-1)).squeeze(-1)  # [B, n_neg]129130        pos_loss = F.binary_cross_entropy_with_logits(131            pos_logits, torch.ones_like(pos_logits))132        neg_loss = F.binary_cross_entropy_with_logits(133            neg_logits, torch.zeros_like(neg_logits))134        return pos_loss + neg_loss135136    def train_model(self, user_seqs: dict[int, list[int]], epochs: int = 5,137                    batch_size: int = 256) -> None:138        """用所有用户行为集合训练 Item Embedding。"""139        self._build_neg_table(user_seqs)140        pairs = self._build_pairs(user_seqs)141        print(f"  共现物品对数量：{len(pairs)}（整 session 两两组合，无滑动窗口）")142143        optimizer = optim.Adam(self.parameters(), lr=self.lr, weight_decay=self.reg)144        for epoch in range(epochs):145            random.shuffle(pairs)146            total_loss = 0.0147            n_batches = 0148            for start in range(0, len(pairs), batch_size):149                batch = pairs[start:start + batch_size]150                center_ids = torch.tensor([p[0] for p in batch], dtype=torch.long)151                pos_ids = torch.tensor([p[1] for p in batch], dtype=torch.long)152                neg_ids = torch.tensor(153                    [self._sample_negs(p[1], self.n_neg) for p in batch],154                    dtype=torch.long)155156                optimizer.zero_grad()157                loss = self(center_ids, pos_ids, neg_ids)158                loss.backward()159                optimizer.step()160161                total_loss += loss.item()162                n_batches += 1163            print(f"  epoch {epoch:2d}  avg_loss={total_loss / max(n_batches, 1):.4f}")164165    # ── 召回接口 ──────────────────────────────────────────────────────────────166    def item_vector(self, iid: int) -> torch.Tensor:167        """返回物品的最终 Embedding（线上写入 ANN 索引的向量）。"""168        with torch.no_grad():169            return self.item_emb.weight[iid]170171    def most_similar(self, iid: int, top_k: int = 5) -> list[tuple[int, float]]:172        """I2I 检索：给定 Trigger 物品，返回余弦最相似的 Top-K 物品。"""173        with torch.no_grad():174            W = F.normalize(self.item_emb.weight, dim=-1)  # [n_items, dim]175            base = W[iid]176            sims = (W @ base)  # [n_items]177            sims[iid] = -1e9178            top_vals, top_idx = sims.topk(top_k)179        return [(int(i), float(s)) for i, s in zip(top_idx.tolist(), top_vals.tolist())]180181    def recommend(self, seq: list[int], top_k: int = 5) -> list[tuple[int, float]]:182        """183        多 Trigger 融合召回：用行为列表中的每个物品做 I2I 检索，分数取最大值融合。184        """185        exclude = set(seq)186        fused: dict[int, float] = {}187        for trigger in seq:188            for cand, s in self.most_similar(trigger, top_k=self.n_items):189                if cand in exclude:190                    continue191                fused[cand] = max(fused.get(cand, -2.0), s)192        ranked = sorted(fused.items(), key=lambda x: x[1], reverse=True)193        return ranked[:top_k]194195196# ── 快速验证 ──────────────────────────────────────────────────────────────────197if __name__ == "__main__":198    random.seed(42)199    torch.manual_seed(42)200201    n_users, n_items = 8, 20202    # 构造合成行为序列：把物品分成 4 个「兴趣簇」，用户主要在某个簇内行为203    clusters = [list(range(0, 5)), list(range(5, 10)), list(range(10, 15)), list(range(15, 20))]204    user_seqs: dict[int, list[int]] = {}205    for u in range(n_users):206        main = clusters[u % len(clusters)]207        seq = random.sample(main, k=4)                       # 簇内 4 个物品208        seq += random.sample(range(n_items), k=4)            # 混入 4 个随机噪声物品209        random.shuffle(seq)210        user_seqs[u] = seq211212    print("合成用户行为序列（session 视为无序集合）：")213    for u in range(n_users):214        print(f"  user {u}: {user_seqs[u]}")215216    model = Item2Vec(n_items, dim=16, lr=0.05, n_neg=5)217    print("\nItem2Vec 训练：")218    model.train_model(user_seqs, epochs=6)219220    print("\nI2I 相似物品（Top-5，簇内物品应彼此靠近）：")221    for trigger in (0, 6, 12):222        sims = model.most_similar(trigger, top_k=5)223        pretty = ", ".join(f"{i}({s:.3f})" for i, s in sims)224        print(f"  trigger item {trigger} -> {pretty}")225226    seq = user_seqs[0]227    print(f"\n用户 0 的行为列表 {seq}")228    print("多 Trigger 融合召回 Top-5：")229    for iid, s in model.recommend(seq, top_k=5):230        print(f"  item {iid}: sim={s:.4f}")231232    # 验证「无序性」：顺序颠倒后训练样本完全相同233    fwd = sorted(model._build_pairs({0: [1, 2, 3]}))234    bwd = sorted(model._build_pairs({0: [3, 2, 1]}))235    print(f"\n顺序颠倒后共现样本是否完全一致：{fwd == bwd}（Item2Vec 不建模时序）")