two_tower_debias.py

1"""2双塔召回负采样偏差处理（第10章 §10.3.2）34核心思想：5- In-Batch 负采样会让热门物品以更高概率成为别人的负样本，导致热门物品被"过度惩罚"，6  Batch Softmax 与全局 Softmax 之间产生采样偏差（Sampling Bias）。7- LogQ 矫正：s_c(x_i, y_j) = s(x_i, y_j) - log p_j，用采样概率的对数抵消高频采样偏差。8- 流式频率估计：不直接计数，而是用哈希数组记录物品两次被采样之间的 Step Gap，用 EMA 在线9  更新平均间隔 δ，则 p_j ≈ 1/δ；借鉴 Count-Min Sketch 的多重哈希思想缓解哈希冲突。10- 训练技巧：Embedding L2 归一化（学余弦相似度）+ 温度系数 τ + 混合负采样（Batch 内负样本11  + 全局均匀负样本）。1213论文出处：Yi et al., Sampling-Bias-Corrected Neural Modeling for Large Corpus Item14Recommendations, RecSys 2019.（PyTorch autograd 实现）15"""1617import math18import random19from collections import Counter2021import torch22import torch.nn as nn23import torch.nn.functional as F24import torch.optim as optim252627class StreamingFrequencyEstimator:28    """29    流式频率估计器：在线估计物品进入 Batch 的采样概率 p_j ≈ 1/δ。3031    多重哈希时取各哈希桶 δ 的最大值（等价于取 p 的最小值），以抑制哈希冲突导致的32    频率高估，这与 Count-Min Sketch 取 min 的思想一致。33    """3435    def __init__(self, n_hash: int = 2, n_bucket: int = 64, alpha: float = 0.05, init_delta: float = 100.0):36        self.n_bucket, self.alpha = n_bucket, alpha37        self.A = [[0] * n_bucket for _ in range(n_hash)]38        self.B = [[init_delta] * n_bucket for _ in range(n_hash)]39        self.primes = [(73, 19), (151, 37), (257, 53)][:n_hash]4041    def _buckets(self, item: int) -> list:42        return [(item * p + q) % self.n_bucket for p, q in self.primes]4344    def update(self, item: int, step: int) -> None:45        for h, idx in enumerate(self._buckets(item)):46            gap = float(step - self.A[h][idx]) if self.A[h][idx] > 0 else self.B[h][idx]47            self.B[h][idx] = (1.0 - self.alpha) * self.B[h][idx] + self.alpha * max(gap, 1.0)48            self.A[h][idx] = step4950    def prob(self, item: int) -> float:51        delta = max(self.B[h][idx] for h, idx in enumerate(self._buckets(item)))52        return min(1.0, 1.0 / max(delta, 1.0))535455class Tower(nn.Module):56    """单塔 MLP：emb_dim → hidden → output_dim（ReLU 激活，最后一层线性输出）。"""5758    def __init__(self, in_dim: int, hidden: int, out_dim: int):59        super().__init__()60        self.net = nn.Sequential(61            nn.Linear(in_dim, hidden),62            nn.ReLU(),63            nn.Linear(hidden, out_dim),64        )6566    def forward(self, x: torch.Tensor) -> torch.Tensor:67        return self.net(x)686970class DebiasedTwoTower(nn.Module):71    """72    带采样偏差修正的双塔召回模型。7374    logits[c] = <u_hat, v_hat_c> / τ - log p_c      （LogQ 矫正）75    L_B = -1/B Σ_i log softmax(logits)[i]            （矫正后的 Batch Softmax）76    """7778    def __init__(self, n_users: int, n_items: int, emb_dim: int = 16, hidden: int = 16,79                 out_dim: int = 16, temperature: float = 0.1,80                 logq_correction: bool = True, n_global_neg: int = 4):81        super().__init__()82        self.n_items = n_items83        self.tau = temperature84        self.logq = logq_correction85        self.n_global_neg = n_global_neg8687        self.user_emb = nn.Embedding(n_users, emb_dim)88        self.item_emb = nn.Embedding(n_items, emb_dim)89        self.user_tower = Tower(emb_dim, hidden, out_dim)90        self.item_tower = Tower(emb_dim, hidden, out_dim)91        self.freq = StreamingFrequencyEstimator()92        self.step = 09394    def _encode_user(self, uid_tensor: torch.Tensor) -> torch.Tensor:95        x = self.user_emb(uid_tensor)96        out = self.user_tower(x)97        return F.normalize(out, dim=-1)9899    def _encode_item(self, iid_tensor: torch.Tensor) -> torch.Tensor:100        x = self.item_emb(iid_tensor)101        out = self.item_tower(x)102        return F.normalize(out, dim=-1)103104    def forward(self, uids: torch.Tensor, iids: torch.Tensor,105                log_p: torch.Tensor) -> torch.Tensor:106        """107        uids: (B,)  iids: (C,) 候选物品（正样本 + 全局负样本）108        log_p: (C,) 各候选物品的 log 采样概率109        正样本 label = 对角线：label[i] = i110        """111        B = uids.shape[0]112        u_hat = self._encode_user(uids)    # (B, d)113        v_hat = self._encode_item(iids)    # (C, d)114115        # 相似度矩阵 (B, C)116        sim = u_hat @ v_hat.T / self.tau117        if self.logq:118            sim = sim - log_p.unsqueeze(0)  # (B, C) - (1, C) = (B, C)119120        # 正样本为对角线（前 B 个候选）121        labels = torch.arange(B)122        return F.cross_entropy(sim, labels)123124    def train_model(self, interactions: list, epochs: int = 8,125                    batch_size: int = 8, lr: float = 0.08) -> None:126        optimizer = optim.Adam(self.parameters(), lr=lr)127        data = list(interactions)128        for ep in range(epochs):129            random.shuffle(data)130            tot, nb = 0.0, 0131            self.train()132            for st in range(0, len(data), batch_size):133                bt = data[st: st + batch_size]134                if len(bt) < 2:135                    continue136                self.step += 1137                B = len(bt)138139                # 混合负采样：Batch 内正样本物品 + 全局均匀负样本140                batch_iids = [iid for _, iid in bt]141                global_neg = [random.randrange(self.n_items) for _ in range(self.n_global_neg)]142                cand_iids = batch_iids + global_neg143144                for iid in cand_iids:145                    self.freq.update(iid, self.step)146147                # log 采样概率：Batch 内用流式估计，全局负样本用均匀概率148                log_p_list = [math.log(max(self.freq.prob(iid), 1e-6)) for iid in batch_iids] + \149                             [math.log(1.0 / self.n_items)] * len(global_neg)150151                uids_t = torch.tensor([u for u, _ in bt], dtype=torch.long)152                iids_t = torch.tensor(cand_iids, dtype=torch.long)153                log_p_t = torch.tensor(log_p_list, dtype=torch.float)154155                optimizer.zero_grad()156                loss = self.forward(uids_t, iids_t, log_p_t)157                loss.backward()158                optimizer.step()159                tot += loss.item()160                nb += 1161            print(f"    epoch {ep:2d}  avg_loss={tot / max(nb, 1):.4f}")162163    def warmup_frequency(self, interactions: list, batch_size: int = 8) -> None:164        data = list(interactions)165        for _ in range(6):166            random.shuffle(data)167            for st in range(0, len(data), batch_size):168                self.step += 1169                for _, i in data[st: st + batch_size]:170                    self.freq.update(i, self.step)171172    def recommend(self, uid: int, top_k: int = 10) -> list:173        self.eval()174        with torch.no_grad():175            u_hat = self._encode_user(torch.tensor([uid]))   # (1, d)176            all_ids = torch.arange(self.n_items)177            v_hat = self._encode_item(all_ids)               # (n, d)178            scores = (u_hat @ v_hat.T).squeeze(0).tolist()179            pairs = list(zip(range(self.n_items), scores))180            return sorted(pairs, key=lambda x: -x[1])[:top_k]181182183def _hot_item_ratio(m: DebiasedTwoTower, n_users: int, hot: set, top_k: int = 5) -> float:184    """统计召回结果 Top-K 中热门物品的占比。"""185    return sum(1 for u in range(n_users) for i, _ in m.recommend(u, top_k) if i in hot) / (n_users * top_k)186187188# ── 快速验证 ──────────────────────────────────────────────────────────────────189if __name__ == "__main__":190    random.seed(42)191    torch.manual_seed(42)192    n_users, n_items, batch_size = 12, 30, 8193194    # 物品流行度服从长尾分布（Zipf）：物品 0、1、2 ... 依次变冷195    pop = [1.0 / (rank + 1) ** 1.3 for rank in range(n_items)]196    cum = [sum(pop[:i + 1]) for i in range(n_items)]197    interactions = []198    for uid in range(n_users):199        for _ in range(24):200            r = random.random() * cum[-1]201            interactions.append((uid, next(i for i in range(n_items) if cum[i] >= r)))202203    emp, total = Counter(iid for _, iid in interactions), len(interactions)204    rank_by_pop = sorted(range(n_items), key=lambda k: -emp[k])205    hot_items = set(rank_by_pop[:5])206    print(f"合成数据：{n_users} 用户 / {n_items} 物品 / {total} 条交互（长尾 Zipf 分布）")207    print(f"真实曝光最高的 5 个热门物品：{sorted(hot_items)}（理想 Top-5 占比 = 100%）")208209    models = {}210    for name, use_logq in [("A 未矫正（原始 In-Batch 负采样）", False),211                           ("B LogQ 矫正 + 流式频率估计 + 混合负采样", True)]:212        print(f"\n【模型 {name}】")213        random.seed(42)214        torch.manual_seed(42)215        m = DebiasedTwoTower(n_users, n_items, logq_correction=use_logq)216        m.warmup_frequency(interactions, batch_size)217        m.train_model(interactions, epochs=8, batch_size=batch_size, lr=0.08)218        models[use_logq] = m219    print("  （注：A 与 B 优化的是不同的 logits，loss 数值不可直接横向比较）")220221    print("\n流式频率估计效果（p_j ≈ 1/δ，对照「物品出现在一个 Batch 中的真实概率」）：")222    seen = [i for i in rank_by_pop if emp[i] > 0]223    for iid in seen[:3] + seen[-2:]:224        f = emp[iid] / total225        p_true = 1.0 - (1.0 - f) ** batch_size226        tag = "热门" if iid in hot_items else "长尾"227        print(f"  item {iid:2d} [{tag}]  估计 p={models[True].freq.prob(iid):.4f}  "228              f"真实 p={p_true:.4f}  样本频率={f:.4f}")229    print("  （长尾物品估计值偏高源于哈希冲突，增大桶数或哈希个数可进一步降低误差）")230231    ra = _hot_item_ratio(models[False], n_users, hot_items, top_k=5)232    rb = _hot_item_ratio(models[True], n_users, hot_items, top_k=5)233    print("\n召回结果 Top-5 中热门物品的占比（理想值 100%）：")234    print(f"  模型 A（未矫正）  ：{ra:.2%}\n  模型 B（LogQ 矫正）：{rb:.2%}")235    print("  结论：未矫正时热门物品被反复当作负样本而遭过度惩罚、排序被压低；LogQ 矫正把 logit")236    print(f"        减去 log p_j 抵消该偏差，占比变化 {rb - ra:+.2%}，Batch Softmax 更接近全局 Softmax。")