eges.py

1"""2EGES 图嵌入召回（第10章 §10.2.1 「EGES召回」）34核心思想：5    DeepWalk / Node2Vec 只用图结构学习 Item ID Embedding，新物品缺少游走路径就学不到好的6    表示。EGES（Enhanced Graph Embedding with Side Information）为每个物品额外引入 n 种7    Side Information（品类、品牌、店铺……），每种 Side Info 维护独立的 Embedding 矩阵 W_s，8    并用「每个物品私有且可学习」的权重矩阵 A ∈ R^{|V|x(n+1)} 做加权融合（s=0 表示 Item ID）：9        H_v = sum_{s=0..n} exp(a_v^s) * W_s(v) / sum_{s=0..n} exp(a_v^s),  a_v^s = A_{vs}10    exp 保证权重为正、分母做归一化，本质是 Softmax 加权平均（Weighted Average Layer，11    思想类似 Attention）。融合后的 H_v 代替原始 Item Embedding 输入 Skip-Gram，用负采样12    联合优化上下文 Embedding Z_u、各类 Side Info Embedding W_s 与权重参数 a_v^s。13    由于 Side Info Embedding 在物品间共享，冷启动物品（无交互但有属性）也能得到合理表示。14论文出处：Wang J. et al. "Billion-scale Commodity Embedding for E-commerce Recommendation15    in Alibaba"（EGES）, KDD 2018. PyTorch 实现，自动求导。16"""1718import math19import random20from collections import Counter, defaultdict21from typing import Dict, List, Optional, Sequence, Tuple2223import torch24import torch.nn as nn25import torch.optim as optim2627Graph = Dict[str, Dict[str, float]]282930# ── 1. 构图与采样工具 ────────────────────────────────────────────────────────31def build_item_graph(sessions: List[List[str]]) -> Graph:32    """由 Session 行为序列构建带权有向物品图，边权 w_ij = Count(i -> j)。"""33    graph: Graph = defaultdict(dict)34    for seq in sessions:35        for a, b in zip(seq, seq[1:]):36            if a != b:37                graph[a][b] = graph[a].get(b, 0.0) + 1.038    return dict(graph)394041class AliasTable:42    """别名采样表：O(N) 预处理、O(1) 采样，工业界负采样与随机游走的常用实现。"""43    def __init__(self, weights: List[float]) -> None:44        n = len(weights)45        total = sum(weights) or 1.046        scaled = [w * n / total for w in weights]47        self.prob, self.alias = [0.0] * n, [0] * n48        small, large = [], []49        for i, s in enumerate(scaled):50            (small if s < 1.0 else large).append(i)51        while small and large:52            si, li = small.pop(), large.pop()53            self.prob[si], self.alias[si] = scaled[si], li54            scaled[li] -= 1.0 - scaled[si]55            (small if scaled[li] < 1.0 else large).append(li)56        for i in large + small:57            self.prob[i], self.alias[i] = 1.0, i5859    def sample(self) -> int:60        i = random.randrange(len(self.prob))61        return i if random.random() < self.prob[i] else self.alias[i]626364def random_walk(graph: Graph, start: str, walk_len: int) -> List[str]:65    """按边权在带权有向图上做随机游走（轮盘赌采样），生成节点访问序列。"""66    path = [start]67    while len(path) < walk_len:68        nbrs = graph.get(path[-1])69        if not nbrs:70            break71        r, acc = random.random() * sum(nbrs.values()), 0.072        for it, w in nbrs.items():73            acc += w74            if acc >= r:75                path.append(it)76                break77    return path787980def cosine_sim(a: Sequence[float], b: Sequence[float]) -> float:81    na, nb = math.sqrt(sum(x * x for x in a)), math.sqrt(sum(x * x for x in b))82    return sum(x * y for x, y in zip(a, b)) / (na * nb + 1e-9)838485# ── 2. EGES 模型 ────────────────────────────────────────────────────────────86class EGES(nn.Module):87    """EGES：Side Info Embedding + 物品私有可学习权重的加权融合 + Weighted Skip-Gram。88    使用 PyTorch autograd 自动求导。"""89    def __init__(self, side_info: Dict[str, Dict[str, str]], dim: int = 16,90                 fields: Sequence[str] = ("item_id", "cate", "brand", "shop")) -> None:91        super().__init__()92        self.side_info = side_info93        self.fields = list(fields)  # fields[0] = "item_id" 对应公式中的 s = 094        self.dim = dim95        self.items = sorted(side_info)96        self.item2id = {v: i for i, v in enumerate(self.items)}9798        # 每种 Side Info 一张独立的 Embedding 矩阵 W_s99        self.value2id: Dict[str, Dict[str, int]] = {}100        self.W = nn.ModuleDict()101        for f in self.fields:102            vals = sorted({self._value(v, f) for v in self.items})103            self.value2id[f] = {val: i for i, val in enumerate(vals)}104            emb = nn.Embedding(len(vals), dim)105            nn.init.uniform_(emb.weight, -1.0 / math.sqrt(dim), 1.0 / math.sqrt(dim))106            self.W[f] = emb107108        # A ∈ R^{|V|x(n+1)}：每个物品「私有」的可学习融合权重，初始化为 0 等价于均匀平均109        self.A = nn.Embedding(len(self.items), len(self.fields))110        nn.init.zeros_(self.A.weight)111        self.Z = nn.Embedding(len(self.items), dim)  # 上下文 Embedding Z_u112        nn.init.zeros_(self.Z.weight)113        self._neg: Optional[AliasTable] = None114115    def _value(self, item: str, field: str) -> str:116        """取物品在某个字段上的取值；item_id 字段的取值就是物品自身。"""117        return item if field == "item_id" else self.side_info[item][field]118119    # ---- 2.1 Weighted Average Layer ----120    def fuse(self, idx: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:121        """Softmax 加权融合 H_v = Σ_s β_s W_s(v)，β_s = exp(a_v^s)/Σ_j exp(a_v^j)；122        返回 (H_v, β, 各 Side Info Embedding)。idx 为物品下标 tensor。"""123        a = self.A(idx)                       # (B, n_fields)124        betas = torch.softmax(a, dim=-1)      # (B, n_fields)125        vecs = torch.stack(126            [self.W[f](torch.tensor([self.value2id[f][self._value(it, f)] for it in _items_of(idx, self.items)],127                                    dtype=torch.long)) for f in self.fields], dim=1128        )                                     # (B, n_fields, dim)129        emb = (betas.unsqueeze(-1) * vecs).sum(dim=1)   # (B, dim)130        return emb, betas, vecs131132    def fuse_item(self, item: str) -> Tuple[List[float], List[float], List[List[float]]]:133        """单个物品的融合结果（纯 Python 返回，便于观察与检索）。"""134        idx = torch.tensor([self.item2id[item]], dtype=torch.long)135        with torch.no_grad():136            emb, betas, vecs = self.fuse(idx)137        return (emb[0].tolist(), betas[0].tolist(),138                [[vecs[0, s, d].item() for d in range(self.dim)] for s in range(len(self.fields))])139140    def attention(self, item: str) -> Dict[str, float]:141        """返回该物品各类 Side Info 的归一化融合权重，便于观察学到的重要性。"""142        _, betas, _ = self.fuse_item(item)143        return {f: b for f, b in zip(self.fields, betas)}144145    # ---- 2.2 Weighted Skip-Gram + 负采样 ----146    def score(self, center_idx: torch.Tensor, ctx_idx: torch.Tensor) -> torch.Tensor:147        """σ(H_v · Z_u)，Skip-Gram 的正样本/负样本得分。"""148        emb, _, _ = self.fuse(center_idx)   # (B, dim)149        z = self.Z(ctx_idx)                 # (B, dim)150        return (emb * z).sum(dim=-1)151152    def _build_negative_table(self, walks: List[List[str]]) -> None:153        """按节点出现频率的 3/4 次幂构建负采样 alias 表。"""154        cnt = Counter(v for path in walks for v in path)155        self._neg = AliasTable([math.pow(cnt.get(v, 0) + 1e-3, 0.75) for v in self.items])156157    def train_model(self, walks: List[List[str]], epochs: int = 3, window: int = 2,158                    neg_k: int = 3, lr: float = 0.05, verbose: bool = True) -> None:159        """Weighted Skip-Gram 训练：联合优化 Item/Side Info Embedding 与融合权重。160        梯度由 PyTorch autograd 自动完成。"""161        self._build_negative_table(walks)162        optimizer = optim.Adam(self.parameters(), lr=lr)163164        for ep in range(1, epochs + 1):165            total, num = 0.0, 0166            optimizer.zero_grad()167            loss_sum = None168169            for path in walks:170                for i, center in enumerate(path):171                    if center not in self.item2id:172                        continue173                    ctxs = path[max(0, i - window):i] + path[i + 1:i + window + 1]174                    for ctx in ctxs:175                        c_idx = torch.tensor([self.item2id[center]], dtype=torch.long)176                        p_idx = torch.tensor([self.item2id[ctx]], dtype=torch.long)177                        logit = self.score(c_idx, p_idx)178                        l = -torch.log(torch.sigmoid(logit) + 1e-10)179                        loss_sum = l if loss_sum is None else loss_sum + l180                        total += l.item()181                        num += 1182                        for _ in range(neg_k):  # 负采样：频率 3/4 次幂分布183                            neg = self._neg.sample() if self._neg else 0184                            if self.items[neg] == ctx:185                                continue186                            n_idx = torch.tensor([neg], dtype=torch.long)187                            logit_n = self.score(c_idx, n_idx)188                            l_n = -torch.log(1.0 - torch.sigmoid(logit_n) + 1e-10)189                            loss_sum = loss_sum + l_n190                            total += l_n.item()191                            num += 1192193            if loss_sum is not None:194                loss_sum.backward()195                optimizer.step()196            if verbose:197                print(f"  epoch {ep}/{epochs}  平均损失 = {total / max(num, 1):.4f}")198199    # ---- 2.3 召回 ----200    def embedding(self, item: str) -> List[float]:201        """物品最终 Embedding，即融合后的 H_v。"""202        return self.fuse_item(item)[0]203204    def cold_start_embedding(self, values: Dict[str, str]) -> List[float]:205        """冷启动物品 Embedding：无交互时 Item ID Embedding 不可靠，退化为只用206        Side Info Embedding，并采用全量物品的平均融合权重。"""207        side = [f for f in self.fields if f != "item_id"]208        avg = {f: sum(self.attention(v)[f] for v in self.items) for f in side}209        tot = sum(avg.values()) or 1.0210        emb = [0.0] * self.dim211        with torch.no_grad():212            for f in side:213                idx = torch.tensor([self.value2id[f][values[f]]], dtype=torch.long)214                vec = self.W[f](idx)[0]215                for k in range(self.dim):216                    emb[k] += avg[f] / tot * vec[k].item()217        return emb218219    def recommend(self, query, top_k: int = 3,220                  exclude: Sequence[str] = ()) -> List[Tuple[str, float]]:221        """Trigger 可以是物品 ID，也可以是一个 Embedding 向量；返回 Top-K 近邻物品。"""222        src = self.embedding(query) if isinstance(query, str) else list(query)223        skip = set(exclude) | ({query} if isinstance(query, str) else set())224        scores = [(v, cosine_sim(src, self.embedding(v)))225                  for v in self.items if v not in skip]226        scores.sort(key=lambda x: x[1], reverse=True)227        return scores[:top_k]228229230def _items_of(idx: torch.Tensor, items: List[str]) -> List[str]:231    """把物品下标 tensor 转回物品名列表（供 fuse 内部查 side info 取值）。"""232    return [items[i] for i in idx.tolist()]233234235# ── 3. 演示 ─────────────────────────────────────────────────────────────────236if __name__ == "__main__":237    random.seed(42)238    torch.manual_seed(42)239    # 24 个热门物品分属 3 个品类，每个品类 2 个品牌；另有 1 个冷启动物品（无任何交互）240    cates = ["数码", "服饰", "食品"]241    side_info: Dict[str, Dict[str, str]] = {}242    groups: Dict[str, List[str]] = {c: [] for c in cates}243    for idx in range(24):244        item, c = f"i{idx + 1}", cates[idx // 8]245        groups[c].append(item)246        side_info[item] = {"cate": c, "brand": f"{c}-b{idx % 2}", "shop": f"shop{idx % 4}"}247    cold_item = "i_new"248    side_info[cold_item] = {"cate": "数码", "brand": "数码-b0", "shop": "shop0"}249    sessions = []250    for _ in range(60):  # 同品类内部的连续点击行为（冷启动物品不出现在任何 Session 中）251        pool = groups[random.choice(cates)]252        sessions.append([random.choice(pool) for _ in range(4)])253    graph = build_item_graph(sessions)254    walks = [random_walk(graph, v, 10) for _ in range(5) for v in sorted(graph)]255    print(f"物品图 {len(graph)} 个节点；随机游走 {len(walks)} 条；{cold_item} 不在图中\n")256    model = EGES(side_info, dim=16)257    print("训练前 i1 的融合权重（A 初始化为 0，等价于平均池化）：",258          {k: round(v, 2) for k, v in model.attention("i1").items()})259    model.train_model(walks, epochs=3, window=2, neg_k=3, lr=0.05)260261    print("\n训练后各物品私有的融合权重（每个物品一组，彼此不同，说明权重是可学习的）：")262    for item in ["i1", "i2", "i12", "i20"]:263        w = model.attention(item)264        print(f"  {item}（{side_info[item]['cate']}）: "265              + "  ".join(f"{f}={w[f]:.3f}" for f in model.fields))266267    print("\n热门物品 i3（数码）的 Top-3 召回：")268    for it, s in model.recommend("i3", top_k=3, exclude=[cold_item]):269        print(f"  {it}（{side_info[it]['cate']}）: {s:.4f}")270271    print(f"\n冷启动物品 {cold_item}（无交互，仅有 {side_info[cold_item]}）的 Top-3 召回：")272    cold_emb = model.cold_start_embedding(side_info[cold_item])273    cold_recall = model.recommend(cold_emb, top_k=3, exclude=[cold_item])274    for it, s in cold_recall:275        print(f"  {it}（{side_info[it]['cate']}，{side_info[it]['brand']}）: {s:.4f}")276    hit = sum(1 for it, _ in cold_recall if side_info[it]["cate"] == "数码")277    print(f"  命中同品类物品 {hit}/3 → Side Info Embedding 在物品间共享，无交互物品也能召回")