inttower.py

1"""2IntTower 粗排模型（第11章 §11.4.4）34IntTower = Interaction enhanced Two-Tower5论文出处：Li et al., "IntTower: the Next Generation of Two-Tower Model for Pre-Ranking6          System", CIKM 2022（华为诺亚方舟实验室）。78精度-时延权衡（对应书中 §11.4.4）：朴素双塔是 Late Interaction——两侧各自编码到最后一层9才做一次内积，细粒度交互信息被压缩殆尽，精度受限；精排式交互 DNN 精度高却无法预计算物品10向量，时延扛不住。IntTower 在「物品侧表示仍可离线预计算」的效率前提下，把交互提前到塔内部11（Early Interaction），用几乎无参数的轻量算子换取接近交互式模型的精度。三个子模块：12    1) LightSE：轻量 SENET。squeeze 用 Mean Pooling 压成标量 z_i = (1/d)·Σ_t e_{i,t}；13       excitation 只用一层 FC + Softmax，k = Softmax(W·z + b)；再重加权 ẽ_i = k_i·e_i。14    2) FEBlock：用户塔第 i 层投影到 H 个子空间 m_u^{i,h} = W_u^{i,h}·h_u^i + b，物品塔只投影15       最后一层 m_v^{L,h}；S^i = Σ_{hu} max_{hv} ⟨m_u^{i,hu}, m_v^{L,hv}⟩/√d_h，ŷ = Σ_i S^i。16       交互算子本身无参数，物品侧只需缓存最后一层多头表示 → 可预计算。17    3) CIR：对比交互正则（Contrastive Interaction Regularization）。在最后一层多头表示18       M_u^L / M_v^L 上做 InfoNCE，同批次其他物品为负样本，温度 τ。19训练目标：L = L_ctr + λ1·L_cir + λ2·‖Θ‖²（λ2 以逐参数权重衰减实现）。2021PyTorch 实现（autograd 自动微分），可直接 `python3 inttower.py` 运行。22"""23import math24import random25import collections26from typing import List, Tuple2728import torch29import torch.nn as nn30import torch.nn.functional as F31import torch.optim as optim323334# ── 工具函数 ──────────────────────────────────────────────────────────────────3536def _dot_py(a, b):37    return sum(x * y for x, y in zip(a, b))3839def _sigmoid_py(x):40    return 1.0 / (1.0 + math.exp(-max(-30.0, min(30.0, x))))4142def _auc(pairs: List[Tuple[float, int]]) -> float:43    pos, neg = [s for s, y in pairs if y == 1], [s for s, y in pairs if y == 0]44    if not pos or not neg:45        return 0.546    hit = sum(1.0 if p > n else (0.5 if p == n else 0.0) for p in pos for n in neg)47    return hit / (len(pos) * len(neg))484950class LightSE(nn.Module):51    """Light-SE 模块：Mean-Pooling squeeze → 单层 FC → Softmax excitation → 逐字段重加权。"""5253    def __init__(self, m: int):54        super().__init__()55        self.m = m56        self.fc = nn.Linear(m, m)5758    def forward(self, fields: List[torch.Tensor]):59        # fields: list of m tensors, each shape (field_dim,)60        z = torch.stack([f.mean() for f in fields])        # (m,)61        lg = self.fc(z)                                    # (m,)62        k = torch.softmax(lg, dim=0)                       # (m,)63        out = [k[i] * fields[i] for i in range(self.m)]64        return out, k656667class FEBlock(nn.Module):68    """FE-Block：用户塔「每一层」表示 × 物品塔「最后一层」表示的多头 Max-Similarity 早期交互。"""6970    def __init__(self, user_dims: List[int], item_dim: int, n_heads: int = 4, d_head: int = 8):71        super().__init__()72        self.L = len(user_dims)73        self.H = n_heads74        self.d = d_head75        self.scale = math.sqrt(d_head)76        # user_proj[i][h]: project user layer-i rep to head h77        self.user_proj = nn.ModuleList([78            nn.ModuleList([nn.Linear(du, d_head) for _ in range(n_heads)])79            for du in user_dims80        ])81        self.item_proj = nn.ModuleList([nn.Linear(item_dim, d_head) for _ in range(n_heads)])8283    def forward(self, user_hs: List[torch.Tensor], item_h: torch.Tensor):84        """85        user_hs: list of L tensors, item_h: shape (item_dim,)86        Returns scalar score (sum of MaxSim over all layers and heads).87        Also returns last-layer multi-head reps for CIR.88        """89        # Item projections (only last layer needed, but computed once)90        m_v = [self.item_proj[h](item_h) for h in range(self.H)]  # list of H tensors (d_head,)9192        total = torch.tensor(0.0)93        for i in range(self.L):94            for hu in range(self.H):95                m_u = self.user_proj[i][hu](user_hs[i])            # (d_head,)96                # MaxSim over item heads97                sims = torch.stack([torch.dot(m_u, m_v[hv]) for hv in range(self.H)])98                total = total + sims.max() / self.scale99        return total, m_v100101    def last_layer_reps(self, user_hs: List[torch.Tensor], item_h: torch.Tensor):102        """M_u^L / M_v^L 拼接，作为 CIR 模块输入。"""103        m_u_L = torch.cat([self.user_proj[self.L - 1][h](user_hs[-1]) for h in range(self.H)])104        m_v_L = torch.cat([self.item_proj[h](item_h) for h in range(self.H)])105        return m_u_L, m_v_L106107108class IntTower(nn.Module):109    """IntTower 粗排模型：Light-SE + 双塔 DNN + FE-Block 早期交互 + CIR 对比正则。110111    梯度：PyTorch autograd 自动微分，无手写梯度。112    """113114    def __init__(self, n_users: int, n_items: int, n_fields: int = 4, field_dim: int = 4,115                 hidden: int = 24, out_dim: int = 16, n_heads: int = 4, d_head: int = 8,116                 lr: float = 0.05, reg: float = 1e-5, lam_cir: float = 0.1, tau: float = 0.2,117                 emb_scale: float = 2.0):118        super().__init__()119        self.n_items = n_items120        self.m = n_fields121        self.fd = field_dim122        self.lam_cir = lam_cir123        self.tau = tau124125        in_dim = n_fields * field_dim126127        # Embeddings: each user/item has n_fields × field_dim parameters128        self.user_emb = nn.Parameter(129            torch.empty(n_users, n_fields, field_dim).uniform_(-emb_scale, emb_scale))130        self.item_emb = nn.Parameter(131            torch.empty(n_items, n_fields, field_dim).uniform_(-emb_scale, emb_scale))132133        self.user_se = LightSE(n_fields)134        self.item_se = LightSE(n_fields)135136        self.user_l1 = nn.Linear(in_dim, hidden)137        self.user_l2 = nn.Linear(hidden, out_dim)138        self.item_l1 = nn.Linear(in_dim, hidden)139        self.item_l2 = nn.Linear(hidden, out_dim)140141        self.user_dims = [hidden, out_dim]142        self.fe = FEBlock(self.user_dims, out_dim, n_heads, d_head)143144        self.optimizer = optim.Adam(self.parameters(), lr=lr, weight_decay=reg)145146    def _tower(self, fields_param: torch.Tensor, se: LightSE,147               l1: nn.Linear, l2: nn.Linear):148        """fields_param: (n_fields, field_dim); returns (list of layer outputs, LightSE k)"""149        fields = [fields_param[i] for i in range(self.m)]150        gated, k = se(fields)151        h = torch.cat(gated)                    # (n_fields * field_dim,)152        h1 = F.relu(l1(h))153        h2 = F.relu(l2(h1))154        return [h1, h2], k155156    def forward(self, uid: int, iid: int):157        u_hs, u_k = self._tower(self.user_emb[uid], self.user_se, self.user_l1, self.user_l2)158        i_hs, i_k = self._tower(self.item_emb[iid], self.item_se, self.item_l1, self.item_l2)159        y, m_v = self.fe(u_hs, i_hs[-1])160        return y, u_hs, i_hs, m_v, u_k, i_k161162    def score(self, uid: int, iid: int) -> float:163        with torch.no_grad():164            y, *_ = self.forward(uid, iid)165            return float(torch.sigmoid(y).item())166167    def train_model(self, pos_pairs, epochs: int = 6, batch_size: int = 8) -> None:168        data, pos_set = list(pos_pairs), set(pos_pairs)169        for epoch in range(epochs):170            random.shuffle(data)171            l_ctr = l_cir = 0.0172            nb = 0173            for st in range(0, len(data), batch_size):174                batch = data[st:st + batch_size]175                if len(batch) < 2:176                    continue177                self.optimizer.zero_grad()178                loss_ctr = torch.tensor(0.0)179                mu_list, mv_list = [], []180                for uid, iid in batch:181                    y, u_hs, i_hs, m_v, _, _ = self.forward(uid, iid)182                    loss_ctr = loss_ctr + F.binary_cross_entropy_with_logits(183                        y, torch.tensor(1.0))184                    # Negative sample185                    nid = random.randrange(self.n_items)186                    while (uid, nid) in pos_set:187                        nid = random.randrange(self.n_items)188                    yn, u_hs_n, i_hs_n, _, _, _ = self.forward(uid, nid)189                    loss_ctr = loss_ctr + F.binary_cross_entropy_with_logits(190                        yn, torch.tensor(0.0))191                    # Collect last-layer multi-head reps for CIR192                    mu, mv = self.fe.last_layer_reps(u_hs, i_hs[-1])193                    mu_list.append(mu)194                    mv_list.append(mv)195196                # CIR loss: InfoNCE over batch197                B = len(mu_list)198                mu_stack = torch.stack(mu_list)   # (B, H*d_head)199                mv_stack = torch.stack(mv_list)   # (B, H*d_head)200                mu_norm = F.normalize(mu_stack, dim=1)201                mv_norm = F.normalize(mv_stack, dim=1)202                sim_matrix = torch.mm(mu_norm, mv_norm.t()) / self.tau  # (B, B)203                labels = torch.arange(B)204                loss_cir = F.cross_entropy(sim_matrix, labels)205206                loss = loss_ctr / (len(batch) * 2) + self.lam_cir * loss_cir207                loss.backward()208                self.optimizer.step()209                l_ctr += (loss_ctr / (len(batch) * 2)).item()210                l_cir += loss_cir.item()211                nb += 1212213            l_ctr /= max(1, nb)214            l_cir /= max(1, nb)215            print(f"  epoch {epoch}  L_ctr={l_ctr:.4f}  L_cir={l_cir:.4f}  "216                  f"L_total={l_ctr + self.lam_cir * l_cir:.4f}")217218    def pre_rank(self, uid: int, candidates: List[int], top_k: int = 10) -> List[Tuple[int, float]]:219        """粗排：用户侧实时前向各层表示，物品侧多头表示可预计算，FE-Block 聚合打分取 Top-K。"""220        scored = [(iid, self.score(uid, iid)) for iid in candidates]221        scored.sort(key=lambda x: x[1], reverse=True)222        return scored[:top_k]223224225if __name__ == "__main__":226    random.seed(42)227    torch.manual_seed(42)228    n_users, n_items, n_fields, fd, epochs = 8, 24, 4, 4, 8229    lat_u = [[random.gauss(0, 1) for _ in range(4)] for _ in range(n_users)]230    lat_i = [[random.gauss(0, 1) for _ in range(4)] for _ in range(n_items)]231    pos_pairs = [(u, i) for u in range(n_users) for i in range(n_items)232                 if _dot_py(lat_u[u], lat_i[i]) / 2.0 > 0.5]233    pos_set = set(pos_pairs)234    model = IntTower(n_users, n_items, n_fields, fd, lam_cir=0.1, tau=0.2)235    print(f"IntTower 粗排：n_users={n_users}  n_items={n_items}  正样本对={len(pos_pairs)}")236237    print("\n[1] 各模块输出维度流转（结构自洽性检查）")238    with torch.no_grad():239        y, u_hs, i_hs, m_v, u_k, i_k = model.forward(0, 0)240        mu_L, mv_L = model.fe.last_layer_reps(u_hs, i_hs[-1])241    print(f"  输入字段 Embedding ：{n_fields} 字段 × {fd} 维 = {n_fields * fd}")242    k_sum = float(u_k.sum().item())243    k_list = [round(float(v), 3) for v in u_k.tolist()]244    print(f"  Light-SE 重加权    ：{n_fields} × {fd}（门控 k 和={k_sum:.3f}，k={k_list}）")245    print(f"  用户塔 / 物品塔    ：{n_fields * fd} → {model.user_dims[0]} → {model.user_dims[1]}")246    print(f"  FE-Block 多头投影  ：用户 {model.fe.L} 层 × {model.fe.H} 头 × d_h={model.fe.d}，物品仅最后一层")247    print(f"  M_u^L / M_v^L      ：{len(mu_L)} / {len(mv_L)} 维（CIR 输入）")248    print(f"  输出 ŷ = Σ_i S^i   ：标量 = {float(y):.4f}")249    print(f"\n[2] 联合训练 L = L_ctr + λ1·L_cir + λ2·‖Θ‖²（λ1={model.lam_cir}，τ={model.tau}）")250    model.train_model(pos_pairs, epochs=epochs, batch_size=8)251    auc = _auc([(model.score(u, i), 1 if (u, i) in pos_set else 0)252                for u in range(n_users) for i in range(n_items)])253    print(f"  训练集 AUC = {auc:.4f}（早期交互使 ŷ 能区分正负样本，而物品侧仍可预计算）")254    uid = 0255    cands = [i for i in range(n_items) if (uid, i) not in pos_set]256    print(f"\n[3] 用户 {uid} 粗排 Top-5（候选 {len(cands)} 个，排除已交互物品）")257    for iid, s in model.pre_rank(uid, cands, top_k=5):258        print(f"  item {iid:2d}: {s:.4f}")