tdm.py

1"""2TDM 树模型召回（第10章 §10.6.2）34论文：Zhu et al., "Learning Tree-based Deep Model for Recommender Systems", KDD 201856核心思想：给全库物品建一棵平衡二叉树（叶子 = 物品），把检索建模成「最大堆概率树」7    P^(j)(n|u) = max_{n_c ∈ Child(n)} P^(j+1)(n_c|u) / α^(j)8即父节点的兴趣概率由孩子中的最大值决定。于是逐层做 Beam Search（每层只保留打分最高的 B 个9节点并展开其孩子）即可在 O(B·log|I|) 次打分内拿到近似 Top-K，无需对全库逐一打分；打分网络10还可任意复杂（这里用 node 作 Query 对行为序列做 Attention + MLP），表达能力不受内积结构11约束。树结构本身也可学习：先按类目聚类初始化，训练若干轮后用学到的叶子 Embedding 做平衡12k-means 重排叶子再继续训练，即「模型训练 ↔ 树结构学习」交替优化。1314它解决了 Hierarchical Softmax（HS）的两个问题：151) HS 的目标是加速语言模型 Softmax 的归一化，并非真正解决 Top-K 检索：它只能沿单条路径16   贪心下降，高层一旦判错就无法回退，误差逐层累积；TDM 的最大堆形式配合 Beam Search17   每层保留 B 个候选，容错能力强得多。182) HS 把每层退化成「一对兄弟节点之间的二分类」这种局部问题，而非「同层全局排序」，而相邻19   兄弟往往语义相近，强行当互斥两类会损失表达能力；TDM 改为每层独立的全局二分类——正样本20   是目标物品的祖先节点，负样本从同层随机采样。2122本实现：PyTorch 实现，树用 Python dict 表示（parent/children），节点打分 MLP 用 nn.Linear，23逐层 Beam Search 与平衡 k-means 树重建为 Python 逻辑。24"""2526import math27import random2829import torch30import torch.nn as nn31import torch.nn.functional as F32import torch.optim as optim333435class TDM(nn.Module):36    """37    简化版 TDM 树模型召回。树用堆式下标存储：根为 1，节点 n 的孩子为 2n / 2n+1，叶子为38    [2^D, 2^(D+1))，D = log2(|I|)。参数分两套（类似 word2vec 的输入/输出表）：beh_emb 是39    行为序列侧的物品 Embedding；树节点侧内部节点用 node_emb、叶子用随物品走的 item_emb，40    便于重建树时平滑继承表征。4142    打分网络：p(node|u) = sigmoid(MLP([û_n ; ê_n ; d·(û_n ⊙ ê_n)]))，其中 u_n 是43    以该节点为 Query 对行为序列做 Attention 得到的节点相关用户表征。44    """4546    def __init__(self, n_items: int, item_cate: list, emb_dim: int = 16, hid: int = 8,47                 lr: float = 0.05, reg: float = 1e-4):48        super().__init__()49        d = int(math.log2(n_items))50        assert 2 ** d == n_items, "本实现要求物品数为 2 的幂，便于构建满二叉树"51        self.n_items, self.item_cate, self.emb_dim = n_items, item_cate, emb_dim52        self.depth, self.leaf0, self.n_nodes = d, 2 ** d, 2 ** (d + 1)  # 下标 1..n_nodes-153        s = math.sqrt(1.0 / emb_dim)5455        self.beh_emb = nn.Embedding(n_items, emb_dim)        # 行为序列侧56        self.item_emb = nn.Embedding(n_items, emb_dim)       # 叶子节点侧（随物品走）57        self.node_emb = nn.Embedding(self.n_nodes, emb_dim)  # 内部节点侧58        for tab in (self.beh_emb, self.item_emb, self.node_emb):59            nn.init.uniform_(tab.weight, -s, s)6061        # 输入 [û_n ; ê_n ; d·(û_n ⊙ ê_n)]62        self.net = nn.Sequential(63            nn.Linear(3 * emb_dim, hid),64            nn.Tanh(),65            nn.Linear(hid, 1),66        )67        nn.init.xavier_uniform_(self.net[0].weight)68        nn.init.zeros_(self.net[0].bias)69        nn.init.xavier_uniform_(self.net[2].weight)70        nn.init.zeros_(self.net[2].bias)7172        self.reg = reg73        self.n_score_calls = 0                       # 用于复杂度对比74        self.leaf_item: dict = {}                    # 叶子节点 <-> 物品 的双向映射75        self.item_leaf: dict = {}76        self.build_tree_by_category()7778        self.optimizer = optim.Adam(self.parameters(), lr=lr, weight_decay=reg)7980    # ── (a) 树构建：按类目聚类初始化 ──────────────────────────────────────────8182    def build_tree_by_category(self) -> None:83        """初始树：按 (类目, id) 排序挂叶子，同类目物品相邻、共享较多祖先 —— 即论文建议的84        初始化（尚无 Embedding 时用类目充当聚类先验）。"""85        order = sorted(range(self.n_items), key=lambda i: (self.item_cate[i], i))86        self._assign_leaves(order)8788    def _assign_leaves(self, order: list) -> None:89        self.leaf_item = {self.leaf0 + p: i for p, i in enumerate(order)}90        self.item_leaf = {i: self.leaf0 + p for p, i in enumerate(order)}91        self._build_parent_children()9293    def _build_parent_children(self) -> None:94        """用普通 Python dict 显式存储树的父子结构（根为 1）。9596        注意：属性名不能叫 children / parent —— nn.Module 已有同名方法，会被覆盖。97        """98        self.child_map: dict = {}99        self.parent_map: dict = {}100        for node in range(1, self.n_nodes):101            if node < self.leaf0:102                self.child_map[node] = (2 * node, 2 * node + 1)103            else:104                self.child_map[node] = ()105            if node > 1:106                self.parent_map[node] = node // 2107108    def layer_of(self, node: int) -> int:109        return node.bit_length() - 1                # 根为第 0 层，叶子为第 depth 层110111    def ancestors(self, item: int) -> list:112        """目标物品叶子节点到根的整条路径（不含根），即逐层的正样本节点。"""113        n = self.item_leaf[item]                    # 叶子右移 k 位即它在上面第 k 层的祖先114        return [n >> k for k in reversed(range(self.depth))]115116    def _node_vec(self, node: int) -> torch.Tensor:117        # 叶子用随物品走的 item_emb，内部节点用 node_emb118        if node >= self.leaf0:119            return self.item_emb(torch.tensor(self.leaf_item[node], dtype=torch.long))120        return self.node_emb(torch.tensor(node, dtype=torch.long))121122    # ── (b) 深度打分网络：node 作为 Query 对行为序列做 Attention ───────────────123124    def _score(self, seq: list, node: int) -> torch.Tensor:125        """p(node|u) = σ(MLP([û_n ; ê_n ; d·(û_n ⊙ ê_n)]))，a_t = Softmax_t(γ·cos(e_n, h_t))126        是以树节点为 Query 的注意力，u_n = Σ_t a_t·h_t 是节点相关的用户表征 —— 同一用户在127        不同节点上表征不同，这是 TDM 比双塔更强的地方。向量先 L2 归一化、交叉项再乘 d128        （使其求和等于 d·cos），否则交叉项数量级远小于前两块，MLP 会直接忽略用户信息。"""129        self.n_score_calls, d = self.n_score_calls + 1, self.emb_dim130        e_hat = F.normalize(self._node_vec(node), dim=-1)                # [d]131        H = F.normalize(self.beh_emb(torch.tensor(seq, dtype=torch.long)), dim=-1)  # [T, d]132        a = F.softmax(4.0 * (H @ e_hat), dim=0)                          # [T] 注意力权重133        u_hat = F.normalize((a.unsqueeze(-1) * H).sum(dim=0), dim=-1)    # [d]134        x = torch.cat([u_hat, e_hat, d * u_hat * e_hat])                 # [3d]135        return self.net(x).squeeze(-1)                                   # scalar logit136137    # ── (d) 逐层正负样本构造 + 训练 ──────────────────────────────────────────138139    def _train_step(self, seq: list, target: int, n_negs: int = 3) -> float:140        """TDM 的样本构造：目标物品叶子的**所有祖先节点都是正样本**（逐层各一个），负样本则在141        **同一层内随机采样**，于是每层都是独立的二分类 —— 学的是「这层里哪些节点值得展开」。142        越深的层节点越多，故按层放大负样本数（n_negs + layer），否则叶子层信号过于稀疏。"""143        pos_path = self.ancestors(target)144        nodes, labels = [], []145        for pos in pos_path:146            layer = self.layer_of(pos)147            lo, hi = 2 ** layer, 2 ** (layer + 1)       # 该层节点下标区间148            pool = [n for n in range(lo, hi) if n != pos]149            cands = [(pos, 1.0)] + [(n, 0.0) for n in150                                    random.sample(pool, min(n_negs + layer, len(pool)))]151            for node, label in cands:152                nodes.append(node)153                labels.append(label)154155        logits = torch.stack([self._score(seq, n) for n in nodes])156        y = torch.tensor(labels, dtype=torch.float32)157        loss = F.binary_cross_entropy_with_logits(logits, y)158        return loss, pos_path159160    def _train_step_loss(self, seq: list, target: int, n_negs: int = 3) -> float:161        """训练一步并更新参数，返回 loss 值。"""162        self.optimizer.zero_grad()163        loss, _ = self._train_step(seq, target, n_negs)164        loss.backward()165        self.optimizer.step()166        return loss.item()167168    def train_model(self, samples: list, epochs: int = 6, n_negs: int = 3, tag: str = "TDM"):169        """samples: [(行为序列, 目标物品), ...]；每条样本在树的每一层各做一次二分类。"""170        self.train()171        for epoch in range(epochs):172            random.shuffle(samples)173            total = sum(self._train_step_loss(seq, tgt, n_negs) for seq, tgt in samples)174            print(f"  [{tag}] epoch {epoch}  avg_loss={total / max(len(samples), 1):.4f}")175176    # ── (c) 逐层 Beam Search 检索 ───────────────────────────────────────────177178    def recommend(self, seq: list, top_k: int = 5, beam: int = 3,179                  verbose: bool = False) -> list:180        """从根出发逐层下降：对当前层候选打分、保留 Top-B，只展开它们的孩子。共 depth 层、181        每层最多打分 2B 个节点，故复杂度 O(B·log|I|)，与 |I| 几乎无关。"""182        self.eval()183        keep = [(1, 0.0)]184        with torch.no_grad():185            for layer in range(1, self.depth + 1):186                cand = [c for n, _ in keep for c in self.child_map[n]]187                scored = sorted(((n, self._score(seq, n).item()) for n in cand),188                                key=lambda x: x[1], reverse=True)189                keep = scored[:beam] if layer < self.depth else scored[:top_k]190                if verbose:191                    detail = "  ".join(192                        f"n{n}" + (f"(i{self.leaf_item[n]})" if n >= self.leaf0 else "")193                        + f"={s:+.3f}" for n, s in keep)194                    print(f"    第 {layer} 层：候选 {len(cand)} 个 → 保留 {len(keep)} 个   {detail}")195        return [(self.leaf_item[n], s) for n, s in keep]196197    def recommend_bruteforce(self, seq: list, top_k: int = 5) -> list:198        # 全库逐个叶子打分（线性扫描），用作 Beam Search 的对照组199        self.eval()200        with torch.no_grad():201            scored = [(self.leaf_item[n], self._score(seq, n).item())202                      for n in range(self.leaf0, self.leaf0 + self.n_items)]203        return sorted(scored, key=lambda x: x[1], reverse=True)[:top_k]204205    # ── (e) 树结构学习：对叶子 Embedding 做平衡 k-means 重排 ──────────────────206207    def _leaf_vec(self, item: int) -> torch.Tensor:208        with torch.no_grad():209            return self.item_emb(torch.tensor(item, dtype=torch.long)).clone()210211    def _balanced_kmeans2(self, items: list, iters: int = 6) -> tuple:212        """手写二分 k-means，并强制两个簇等大（保证树仍然平衡）。用 Python 手动迭代。"""213        a0 = items[0]  # 用「离 a0 最远的点」做第二个簇心，避免两个簇心初始重合214        a1 = max(items, key=lambda i: float(((self._leaf_vec(a0) - self._leaf_vec(i)) ** 2).sum()))215        c = [self._leaf_vec(a0), self._leaf_vec(a1)]216        half, left, right = len(items) // 2, items[:len(items) // 2], items[len(items) // 2:]217        for _ in range(iters):218            # 分配步：按「到簇0 与到簇1 的距离差」排序，前一半归簇0（这样天然等大）219            ranked = sorted(items, key=lambda i: float(((self._leaf_vec(i) - c[0]) ** 2).sum()220                                                       - ((self._leaf_vec(i) - c[1]) ** 2).sum()))221            left, right = ranked[:half], ranked[half:]222            for ci, grp in ((0, left), (1, right)):     # 更新步：重算簇心223                c[ci] = torch.stack([self._leaf_vec(i) for i in grp]).mean(dim=0)224        return left, right225226    def rebuild_tree(self) -> list:227        """递归平衡 k-means 得到新的叶子顺序，让 Embedding 相近的物品共享更多祖先。"""228        def split(items: list) -> list:229            if len(items) <= 1:230                return items231            lf, rt = self._balanced_kmeans2(items)232            return split(lf) + split(rt)233        order = split(list(range(self.n_items)))234        self._assign_leaves(order)235        self._reseed_internal_nodes()236        return order237238    def _reseed_internal_nodes(self) -> None:239        # 树结构变了，内部节点旧表征已失效：自底向上用两个孩子的均值重新初始化240        with torch.no_grad():241            for node in range(self.leaf0 - 1, 0, -1):242                l, r = self._node_vec(2 * node), self._node_vec(2 * node + 1)243                self.node_emb.weight[node] = (l + r) / 2.0244245246def _tree_gap(md: TDM, pairs: set) -> float:247    """树结构质量：共现物品对的平均 LCA 层数 - 非共现物品对的平均 LCA 层数，越大越好 ——248    常被一起消费的物品共享越多祖先，Beam Search 就越不容易在高层把目标剪掉。"""249    acc: dict = {True: [], False: []}250    for i in range(md.n_items):251        for j in range(i + 1, md.n_items):252            a, b = md.item_leaf[i], md.item_leaf[j]       # 两叶子同层，同步右移求 LCA253            k = next(k for k in range(md.depth + 1) if a >> k == b >> k)254            acc[(i, j) in pairs].append(md.depth - k)255    return sum(acc[True]) / len(acc[True]) - sum(acc[False]) / len(acc[False])256257258def _hit_rate(model: TDM, samples: list, top_k: int = 5, beam: int = 3) -> float:259    return sum(1.0 for seq, tgt in samples260               if tgt in [i for i, _ in model.recommend(seq, top_k, beam)]) / len(samples)261262263# ── 快速验证 ──────────────────────────────────────────────────────────────────264if __name__ == "__main__":265    random.seed(42)266    torch.manual_seed(42)267268    n_items, n_users, n_cate = 16, 8, 4269    item_cate = [i % n_cate for i in range(n_items)]  # 每个类目 4 个物品270271    # 合成数据：每个用户偏好 2 个类目，用「序列前缀 → 下一个物品」构造逐层训练样本272    samples = []273    user_seq = {}274    for u in range(n_users):275        favs = random.sample(range(n_cate), 2)276        pool = [i for i in range(n_items) if item_cate[i] in favs]277        seq = user_seq[u] = random.sample(pool, k=6)278        samples += [(seq[:t], seq[t]) for t in range(2, len(seq))]279280    model = TDM(n_items, item_cate, emb_dim=16, hid=8, lr=0.05)281    print(f"物品数={n_items} 类目数={n_cate} 树深度={model.depth} "282          f"节点数={model.n_nodes - 1} 训练样本数={len(samples)}")283    print("初始树（按类目聚类初始化）叶子顺序：",284          [f"i{(j := model.leaf_item[n])}(c{item_cate[j]})" for n in range(model.leaf0, 2 * n_items)])285    cooc = {(min(a, b), max(a, b)) for sq in user_seq.values() for a in sq for b in sq if a != b}286    model._assign_leaves(list(range(n_items)))   # 对照组：无任何先验，按 id 顺序挂叶子287    gap_rand = _tree_gap(model, cooc)288    model.build_tree_by_category()               # 恢复类目初始树289    gap_cate = _tree_gap(model, cooc)290    print("\n=== 阶段一：固定树结构，训练深度打分网络（逐层二分类） ===")291    model.train_model(samples, epochs=6, n_negs=3)292    hr_before = _hit_rate(model, samples, top_k=5, beam=3)293294    print("\n=== 阶段二：树结构学习（对学到的叶子 Embedding 做递归平衡 k-means） ===")295    order = model.rebuild_tree()296    print("  重建后叶子顺序：", [f"i{i}(c{item_cate[i]})" for i in order])297    print(f"  树结构质量（共现物品对的 LCA 层数优势，越大越好）：无先验树 {gap_rand:+.2f}"298          f" / 类目初始树 {gap_cate:+.2f} / k-means 重建树 {_tree_gap(model, cooc):+.2f}")299    print("  k-means 只看学到的叶子 Embedding、完全没用类目特征，却恢复出了接近类目先验的"300          "结构（远优于无先验树）—— 这就是「树结构可学习」的意义。")301    hr_mid = _hit_rate(model, samples, top_k=5, beam=3)302    print("  树结构已变、内部节点表征需重新适配 —— 继续训练（模型 ↔ 树结构交替优化）：")303    model.train_model(samples, epochs=5, n_negs=3, tag="重训")304    hr_after = _hit_rate(model, samples, top_k=5, beam=3)305    print(f"\n  HR@5（Beam Search, B=3）：原始树 {hr_before:.3f}"306          f" → 刚重建完 {hr_mid:.3f} → 重训之后 {hr_after:.3f}")307    print("  重建瞬间 HR 会掉（节点语义被打乱），重训后反超 —— 这正是交替优化的典型表现。")308309    seq, target = samples[0]310    print(f"\n=== 逐层 Beam Search 检索（用户序列 {seq}，类目 "311          f"{[item_cate[i] for i in seq]}，真实目标 i{target}） ===")312    model.n_score_calls = 0313    recs = model.recommend(seq, top_k=5, beam=3, verbose=True)314    beam_calls, model.n_score_calls = model.n_score_calls, 0315    brute = model.recommend_bruteforce(seq, top_k=5)316    full_calls = model.n_score_calls317    fmt = lambda r: [f"i{i}(c{item_cate[i]}){s:+.2f}" for i, s in r]318    print(f"  Beam Search Top-5：{fmt(recs)}\n  全库打分    Top-5：{fmt(brute)}")319    overlap = len({i for i, _ in recs} & {i for i, _ in brute})320    print(f"\n--- 打分次数对比（B=3, |I|={n_items}） ---\n"321          f"  全库线性扫描 O(|I|)      ：{full_calls} 次\n"322          f"  Beam Search  O(B·log|I|) ：{beam_calls} 次，与全库 Top-5 重合 {overlap}/5")323    print("  |I|=16 时树的开销还摊不开，物品量越大优势越明显（|I| vs 2B·log2|I|）：")324    for n, bm in ((m, 6 * int(math.log2(m))) for m in (16, 1024, 1 << 20)):325        print(f"    |I|={n:>9}：全库 {n:>9} 次  vs  Beam {bm:>4} 次  加速 {n / bm:>9.1f}x")326    print("\n--- 最终召回示例（偏好类目内的物品标 √） ---")327    for u in range(3):328        ctx, favs = user_seq[u][:3], {item_cate[i] for i in user_seq[u]}329        print(f"  用户 {u} 序列 {ctx}（类目 {sorted(favs)}）-> "330              + str([f"i{i}(c{item_cate[i]})" + ("√" if item_cate[i] in favs else "×")331                     for i, _ in model.recommend(ctx, top_k=5, beam=3)]))