mind.py

1"""2MIND 多兴趣召回（第10章 §10.5.1）34论文：Li et al., "Multi-Interest Network with Dynamic Routing for Recommendation at Tmall", CIKM 201956核心思想：7- 用胶囊网络（Capsule Network）的动态路由机制，从用户行为序列中抽取多个兴趣向量8- 每个兴趣胶囊对应用户某一方面的偏好9- 训练时用 Label-aware Attention 机制，根据目标物品选择最相关的兴趣向量1011本实现：PyTorch 实现，动态路由 + Label-aware Attention + Sampled Softmax loss。12"""1314import math15import random16import collections1718import torch19import torch.nn as nn20import torch.nn.functional as F21import torch.optim as optim222324def _squash(v: torch.Tensor) -> torch.Tensor:25    """Capsule squash 函数：将向量压缩到单位球内，保持方向。v: [..., d]"""26    norm_sq = (v * v).sum(dim=-1, keepdim=True)27    norm = torch.sqrt(norm_sq + 1e-9)28    scale = norm_sq / (1.0 + norm_sq) / norm29    return v * scale303132class MIND(nn.Module):33    """34    简化版 MIND 多兴趣模型。3536    组件：37      item_emb:   物品 Embedding 表38      动态路由：  3 次迭代，b_ij += dot(interest_k, item_emb)，softmax 得路由权重39      label_aware_attention: w_k = softmax(pow(dot(interest_k, target), p))40      Sampled Softmax loss 训练41    """4243    def __init__(44        self,45        n_items: int,46        emb_dim: int = 32,47        n_interests: int = 4,48        routing_iters: int = 3,49        lr: float = 0.01,50        reg: float = 1e-4,51    ):52        super().__init__()53        self.n_items = n_items54        self.emb_dim = emb_dim55        self.n_interests = n_interests56        self.routing_iters = routing_iters57        self.reg = reg5859        self.item_emb = nn.Embedding(n_items, emb_dim)60        nn.init.uniform_(self.item_emb.weight, -math.sqrt(1.0 / emb_dim), math.sqrt(1.0 / emb_dim))6162        self.optimizer = optim.Adam(self.parameters(), lr=lr, weight_decay=reg)6364    def _dynamic_routing(self, seq_embs: torch.Tensor) -> torch.Tensor:65        """66        EM 式动态路由，输入行为序列 Embedding，输出 n_interests 个兴趣向量。6768        Args:69            seq_embs: [T, d] 行为序列 Embedding7071        Returns:72            interests: [K, d]73        """74        T, d = seq_embs.shape75        K = self.n_interests7677        b = torch.zeros(T, K, device=seq_embs.device)  # 路由 logit7879        interests = torch.zeros(K, d, device=seq_embs.device)80        for _ in range(self.routing_iters):81            c = F.softmax(b, dim=1)  # [T, K]82            # 加权聚合：interests[k] = sum_t c[t][k] * seq_embs[t]83            interests = torch.einsum('tk,td->kd', c, seq_embs)  # [K, d]84            interests = _squash(interests)  # 胶囊 squash85            # 更新路由 logit86            b = b + torch.einsum('td,kd->tk', seq_embs, interests)  # [T, K]8788        return interests8990    def encode_user(self, seq: list) -> torch.Tensor:91        """给定用户行为序列（id 列表），返回 K 个兴趣向量 [K, d]。"""92        if not seq:93            return torch.zeros(self.n_interests, self.emb_dim)94        seq_t = torch.tensor(seq, dtype=torch.long)95        with torch.no_grad():96            seq_embs = self.item_emb(seq_t)  # [T, d]97        return self._dynamic_routing(seq_embs)9899    def _label_aware_attention(100        self, interests: torch.Tensor, target_emb: torch.Tensor, pow_p: float = 2.0101    ) -> torch.Tensor:102        """103        Label-aware Attention：对多个兴趣向量加权求和，权重由目标物品相似度决定。104105        attention_weight[k] = softmax(pow(<interest_k, target_emb>, pow_p))106        p=1 时退化为普通 softmax。107        """108        scores = (interests * target_emb.unsqueeze(0)).sum(dim=-1)  # [K]109        weights = F.softmax(scores.pow(pow_p), dim=0)  # [K]110        fused = (weights.unsqueeze(-1) * interests).sum(dim=0)  # [d]111        return fused112113    def forward(self, seq_ids: torch.Tensor, pos_id: torch.Tensor, neg_id: torch.Tensor):114        """BPR 训练步骤：用 label-aware attention 融合兴趣后打分，返回 loss。"""115        seq_embs = self.item_emb(seq_ids)    # [T, d]116        pos_emb = self.item_emb(pos_id)      # [d]117        neg_emb = self.item_emb(neg_id)      # [d]118119        interests = self._dynamic_routing(seq_embs)  # [K, d]120        user_emb = self._label_aware_attention(interests, pos_emb)  # [d]121122        s_pos = F.cosine_similarity(user_emb.unsqueeze(0), pos_emb.unsqueeze(0))123        s_neg = F.cosine_similarity(user_emb.unsqueeze(0), neg_emb.unsqueeze(0))124        loss = -F.logsigmoid(s_pos - s_neg).mean()125        return loss126127    def train_model(128        self,129        user_seqs: dict,130        epochs: int = 10,131        neg_k: int = 1,132    ) -> None:133        user_pos = collections.defaultdict(set)134        for uid, seq in user_seqs.items():135            user_pos[uid].update(seq)136137        samples = []138        for uid, seq in user_seqs.items():139            if len(seq) < 2:140                continue141            for t in range(1, len(seq)):142                samples.append((uid, seq[:t], seq[t]))143144        self.train()145        for epoch in range(epochs):146            random.shuffle(samples)147            total_loss = 0.0148            for uid, context, pos_id in samples:149                for _ in range(neg_k):150                    neg_id = random.randint(0, self.n_items - 1)151                    while neg_id in user_pos[uid]:152                        neg_id = random.randint(0, self.n_items - 1)153154                    seq_t = torch.tensor(context, dtype=torch.long)155                    pos_t = torch.tensor(pos_id, dtype=torch.long)156                    neg_t = torch.tensor(neg_id, dtype=torch.long)157158                    self.optimizer.zero_grad()159                    loss = self(seq_t, pos_t, neg_t)160                    loss.backward()161                    self.optimizer.step()162                    total_loss += loss.item()163164            if epoch % 2 == 0:165                print(f"  epoch {epoch:3d}  avg_loss={total_loss / max(len(samples), 1):.4f}")166167    def recommend(168        self, seq: list, top_k: int = 10, exclude: set = None169    ) -> list:170        """171        线上推荐：对每个兴趣向量做 ANN 检索，取各兴趣的 Top-K 物品合并后重排。172        """173        exclude = exclude or set(seq)174        self.eval()175        with torch.no_grad():176            interests = self.encode_user(seq)  # [K, d]177            all_embs = self.item_emb.weight    # [n_items, d]178            # 对每个兴趣计算余弦相似度179            interests_norm = F.normalize(interests, dim=-1)   # [K, d]180            all_norm = F.normalize(all_embs, dim=-1)          # [n_items, d]181            sims = torch.mm(interests_norm, all_norm.t())      # [K, n_items]182            max_sims, _ = sims.max(dim=0)                      # [n_items]183184        score_dict = {}185        for iid in range(self.n_items):186            if iid not in exclude:187                score_dict[iid] = max_sims[iid].item()188189        ranked = sorted(score_dict.items(), key=lambda x: x[1], reverse=True)190        return ranked[:top_k]191192193# ── 快速验证 ──────────────────────────────────────────────────────────────────194if __name__ == "__main__":195    random.seed(42)196    torch.manual_seed(42)197198    n_items = 20199    user_seqs = {200        u: random.sample(range(n_items), k=random.randint(5, 10))201        for u in range(8)202    }203204    model = MIND(n_items, emb_dim=16, n_interests=3, routing_iters=3, lr=0.05)205    print("MIND 多兴趣模型训练：")206    model.train_model(user_seqs, epochs=6)207208    seq = user_seqs[0]209    recs = model.recommend(seq, top_k=5)210    print(f"\n用户序列 {seq} 的 Top-5 推荐：")211    for iid, s in recs:212        print(f"  item {iid}: {s:.4f}")213214    # 可视化多兴趣向量（查看是否分散）215    interests = model.encode_user(seq)216    print(f"\n{len(interests)} 个兴趣向量与 item 的相似度分布：")217    model.eval()218    with torch.no_grad():219        all_embs = model.item_emb.weight220        interests_norm = F.normalize(interests, dim=-1)221        all_norm = F.normalize(all_embs, dim=-1)222        sims = torch.mm(interests_norm, all_norm.t())  # [K, n_items]223    for k in range(len(interests)):224        norm_val = interests[k].norm().item()225        max_sim = sims[k].max().item()226        mean_sim = sims[k].mean().item()227        print(f"  interest {k}: norm={norm_val:.4f}  max_sim={max_sim:.4f}  mean_sim={mean_sim:.4f}")