target_attention.py

1"""2Target Attention（目标注意力）精排示例（第12章 §12.4.3.2）34本书依据：《第十二章-精排模型.pdf》§12.4.3.2，代码 12.3。56推荐精排中的 Target Attention 用「当前候选物品」对用户行为序列做加权汇聚，7从而提取与候选相关的兴趣表示。它与缩放点积 Cross-Attention 的主要区别在于：8打分函数是一个可学习的非线性 MLP，而非简单的向量内积。910    缩放点积 Cross-Attention：α_i = softmax( (e_a · e_i) / √d )11    Target Attention（DIN 风格）：α_i = MLP([e_a; e_i; e_a−e_i; e_a⊙e_i])1213后者能将候选、历史行为及两者之间的差异与交互作为打分的输入，14因此当「差分」或「元素乘积」携带有效判别信息时，MLP 打分函数的表达能力更强。1516正文指出：完整的 DIN 模型和 BST 模型将在第20章《用户兴趣与序列建模》详细介绍；17本节只展示 Target Attention 模块在精排特征交互框架中的作用。1819本文件的对照实验：20  - 基线：平均池化（忽略候选物品，静态汇聚序列兴趣）21  - 缩放点积 Cross-Attention（内积打分，不使用差分/交互项）22  - Target Attention（MLP 打分，使用 [e_a; e_i; e_a−e_i; e_a⊙e_i]）2324依赖：PyTorch。运行：python3 target_attention.py25"""2627import random28import torch29import torch.nn as nn30import torch.nn.functional as F31import torch.optim as optim32import math333435def _auc(scores, labels) -> float:36    pos = [s for s, y in zip(scores, labels) if y == 1]37    neg = [s for s, y in zip(scores, labels) if y == 0]38    if not pos or not neg:39        return 0.540    return sum(1.0 if p > n else 0.5 if p == n else 0.041               for p in pos for n in neg) / (len(pos) * len(neg))424344# ── Target Attention（代码 12.3，正文 §12.4.3.2）────────────────────────────4546class TargetAttention(nn.Module):47    """48    DIN 风格目标注意力（正文代码 12.3）。4950    打分网络输入 [target; history; target−history; target⊙history]，51    包含四类关系：原始向量、差分（方向）、逐元素乘积（交互）。52    与仅使用内积相比，能捕获非线性的相关性模式。5354    与缩放点积 Cross-Attention 的关键区别：55      - 缩放内积：α_i = softmax(e_a · e_i / √d)  — 线性相似度56      - 本模块：  α_i = MLP([e_a; e_i; e_a−e_i; e_a⊙e_i]) — 可学习非线性打分57    正文强调，推荐系统不必局限于向量内积，可以使用更丰富的匹配特征。58    """5960    def __init__(self, d_model: int, hidden_dim: int = 16):61        super().__init__()62        self.score_net = nn.Sequential(63            nn.Linear(4 * d_model, hidden_dim),64            nn.ReLU(),65            nn.Linear(hidden_dim, 1),66        )6768    def forward(self, target: torch.Tensor, history: torch.Tensor,69                valid_mask: torch.Tensor | None = None):70        """71        target     : (B, D)   当前候选物品嵌入72        history    : (B, L, D) 用户历史行为嵌入序列73        valid_mask : (B, L) bool，True 表示有效位置（可选）74        返回: pooled (B, D)，weights (B, L)75        """76        t = target.unsqueeze(1).expand_as(history)             # (B, L, D)77        feat = torch.cat([t, history, t - history, t * history], dim=-1)78        scores = self.score_net(feat).squeeze(-1)               # (B, L)7980        if valid_mask is not None:81            scores = scores.masked_fill(~valid_mask,82                                        torch.finfo(scores.dtype).min)83        weights = torch.softmax(scores, dim=-1)                 # (B, L)84        if valid_mask is not None:85            weights = weights * valid_mask.float()86            weights = weights / weights.sum(dim=-1, keepdim=True).clamp_min(1e-8)8788        pooled = (weights.unsqueeze(-1) * history).sum(dim=1)  # (B, D)89        return pooled, weights909192class TargetAttentionUnnormalized(TargetAttention):93    """94    DIN 原文的「局部激活单元」：直接用原始相关性打分做加权求和，不做 softmax 归一化。9596    正文代码 12.3 为演示简洁采用了 softmax；而 DIN 论文强调不做归一化，97    以便在候选与某条历史行为高度相关时，让该行为的表示被**放大**，98    而不是仅仅在若干行为之间重新分配 1 的总权重。99    当序列中只有少数行为与候选真正相关时，归一化会把注意力「摊薄」，100    未归一化则保留了兴趣强度信息。101    """102103    def forward(self, target: torch.Tensor, history: torch.Tensor,104                valid_mask: torch.Tensor | None = None):105        t = target.unsqueeze(1).expand_as(history)106        feat = torch.cat([t, history, t - history, t * history], dim=-1)107        weights = self.score_net(feat).squeeze(-1)             # (B, L) 原始打分108        if valid_mask is not None:109            weights = weights.masked_fill(~valid_mask, 0.0)110        pooled = (weights.unsqueeze(-1) * history).sum(dim=1)111        return pooled, weights112113114# ── 缩放点积 Cross-Attention（对照组）────────────────────────────────────────115116class ScaledDotCrossAttention(nn.Module):117    """缩放点积 Cross-Attention：α_i = softmax((W_q e_a)·(W_k e_i) / √d)。"""118119    def __init__(self, d_model: int):120        super().__init__()121        self.q = nn.Linear(d_model, d_model, bias=False)122        self.k = nn.Linear(d_model, d_model, bias=False)123        self.scale = math.sqrt(d_model)124125    def forward(self, target: torch.Tensor, history: torch.Tensor,126                valid_mask: torch.Tensor | None = None):127        q = self.q(target).unsqueeze(1)                        # (B,1,D)128        k = self.k(history)                                    # (B,L,D)129        scores = (q * k).sum(-1) / self.scale                  # (B,L)130        if valid_mask is not None:131            scores = scores.masked_fill(~valid_mask,132                                        torch.finfo(scores.dtype).min)133        weights = torch.softmax(scores, dim=-1)134        pooled = (weights.unsqueeze(-1) * history).sum(dim=1)135        return pooled, weights136137138# ── CTR 预测器（三种汇聚方式共用相同 MLP 头）────────────────────────────────139140class CTRPredictor(nn.Module):141    """142    嵌入 → 行为序列汇聚 → MLP → CTR。143144    pool_mode 取值：145      "avg"         平均池化（基线）146      "dot"         缩放点积 Cross-Attention147      "target_attn" Target Attention（MLP 打分 + softmax，正文代码 12.3）148      "target_unnorm" Target Attention（MLP 打分 + 不归一化，DIN 原文做法）149    """150151    def __init__(self, vocab: int, n_fields: int, emb_dim: int = 16,152                 hidden: int = 32, pool_mode: str = "target_attn"):153        super().__init__()154        self.emb = nn.Embedding(vocab, emb_dim)155        self.pool_mode = pool_mode156        self.n_fields = n_fields157158        if pool_mode == "dot":159            self.attn = ScaledDotCrossAttention(emb_dim)160        elif pool_mode == "target_attn":161            self.attn = TargetAttention(emb_dim, hidden_dim=16)162        elif pool_mode == "target_unnorm":163            self.attn = TargetAttentionUnnormalized(emb_dim, hidden_dim=16)164        # avg 不需要额外模块165166        # 输入：汇聚兴趣 + 候选物品 + 场景字段167        mlp_in = emb_dim * (n_fields + 2)168        self.mlp = nn.Sequential(169            nn.Linear(mlp_in, hidden), nn.ReLU(),170            nn.Linear(hidden, hidden // 2), nn.ReLU(),171            nn.Linear(hidden // 2, 1),172        )173174    def forward(self, seq_ids: torch.Tensor, target_id: torch.Tensor,175                feat_ids: torch.Tensor) -> torch.Tensor:176        seq_emb = self.emb(seq_ids)                            # (B, T, D)177        target_emb = self.emb(target_id)                       # (B, D)178179        if self.pool_mode == "avg":180            v = seq_emb.mean(dim=1)181        else:182            v, _ = self.attn(target_emb, seq_emb)183184        fields = [self.emb(feat_ids[:, i]) for i in range(self.n_fields)]185        z = torch.cat([v, target_emb] + fields, dim=1)186        return self.mlp(z).squeeze(-1)187188    def fit(self, seq_ids, target_id, feat_ids, y,189            epochs: int = 40, lr: float = 5e-3):190        opt = optim.Adam(self.parameters(), lr=lr, weight_decay=1e-4)191        for ep in range(epochs):192            opt.zero_grad()193            loss = F.binary_cross_entropy_with_logits(194                self.forward(seq_ids, target_id, feat_ids), y.float())195            loss.backward()196            opt.step()197            if (ep + 1) % 20 == 0:198                print(f"  epoch {ep+1:3d}  loss={loss.item():.4f}")199        return loss.item()200201    def predict_proba(self, seq_ids, target_id, feat_ids) -> torch.Tensor:202        with torch.no_grad():203            return torch.sigmoid(self.forward(seq_ids, target_id, feat_ids))204205206# ── 合成数据：差分信息比内积更能刻画「相关」────────────────────────────────────207208def make_data(n: int = 3000, vocab: int = 16, n_cat: int = 4,209              seq_len: int = 6, n_fields: int = 3, seed: int = 42):210    """211    用户偏好与物品类目绑定：候选物品与偏好同类目时点击率高；212    行为序列中约 60% 是偏好类目物品，其余为其他类目的随机浏览。213214    这种构造要求模型从序列里挑出「与候选同类目」的那部分行为。215      - 平均池化与候选无关，只能给出静态摘要，无法区分候选；216      - 内积打分只能利用线性相似度；217      - MLP 打分函数可以利用 [e_a; e_i; e_a−e_i; e_a⊙e_i] 中的差分与交互项，218        表达更复杂的相关性判断。219    """220    torch.manual_seed(seed)221    random.seed(seed)222223    item_cat = torch.randint(0, n_cat, (vocab,))224    user_pref = torch.randint(0, n_cat, (n,))225    target = torch.randint(0, vocab, (n,))226    feat = torch.randint(0, vocab, (n, n_fields))227228    match = (item_cat[target] == user_pref).float()229    logit = 2.4 * match - 1.0230    y = (torch.sigmoid(logit + 0.5 * torch.randn(n)) > 0.5).long()231232    seqs = []233    for p in user_pref.tolist():234        s = []235        for _ in range(seq_len):236            if random.random() < 0.6:237                pool = (item_cat == p).nonzero(as_tuple=True)[0].tolist()238            else:239                pool = (item_cat != p).nonzero(as_tuple=True)[0].tolist()240            s.append(random.choice(pool) if pool else random.randrange(vocab))241        seqs.append(s)242    seq = torch.tensor(seqs)243244    return dict(seq=seq, target=target, feat=feat, y=y)245246247# ── 主演示 ────────────────────────────────────────────────────────────────────248249if __name__ == "__main__":250    random.seed(42)251    torch.manual_seed(42)252253    VOCAB, N_FIELDS, SEQ_LEN = 16, 3, 6254    data = make_data(3000, vocab=VOCAB, seq_len=SEQ_LEN, n_fields=N_FIELDS)255    split = 2400256    tr = {k: v[:split] for k, v in data.items()}257    te = {k: v[split:] for k, v in data.items()}258259    def build(mode):260        torch.manual_seed(42)261        return CTRPredictor(VOCAB, N_FIELDS, emb_dim=8, hidden=32,262                            pool_mode=mode)263264    def evaluate(model):265        p = model.predict_proba(te["seq"], te["target"], te["feat"])266        return _auc(p.tolist(), te["y"].tolist())267268    results, models = {}, {}269    for mode, label in [("avg", "平均池化（基线）"),270                        ("dot", "缩放点积 Cross-Attention"),271                        ("target_attn", "Target Attention（MLP + softmax，代码 12.3）"),272                        ("target_unnorm", "Target Attention（MLP，不归一化，DIN 原文）")]:273        print(f"\n=== {label} ===")274        m = build(mode)275        m.fit(tr["seq"], tr["target"], tr["feat"], tr["y"], epochs=40)276        models[mode] = m277        results[mode] = evaluate(m)278        print(f"  AUC = {results[mode]:.4f}")279280    print(f"\n对比（正文 §12.4.3.2）：")281    print(f"  平均池化                    AUC={results['avg']:.4f}")282    print(f"  缩放点积 CrossAttn          AUC={results['dot']:.4f}"283          f"  (Δ={results['dot'] - results['avg']:+.4f})")284    print(f"  Target Attn + softmax       AUC={results['target_attn']:.4f}"285          f"  (Δ={results['target_attn'] - results['avg']:+.4f})")286    print(f"  Target Attn（不归一化）     AUC={results['target_unnorm']:.4f}"287          f"  (Δ={results['target_unnorm'] - results['avg']:+.4f})")288    print("  正文要点：Target Attention 的打分函数可以是可学习的非线性网络，")289    print("  将目标、历史行为及二者之间的差分与逐元素乘积作为输入，")290    print("  从而表达比向量内积更丰富的关系；它按候选动态汇聚兴趣，")291    print("  而平均池化只能给出与候选无关的静态兴趣摘要。")292    print("  另外，归一化与不归一化的差别很关键：softmax 把总权重固定为 1，")293    print("  序列中只有少数行为真正相关时会摊薄注意力；DIN 原文不做归一化，")294    print("  让高度相关的行为被放大，保留兴趣强度。")295296    # ── 关键性质：同一序列、不同候选，汇聚表示应变化 ────────────────────────297    print("\n同一序列在不同候选下的汇聚表示变化（L2 距离，16 个候选的平均）：")298    with torch.no_grad():299        for mode in ("avg", "target_attn"):300            m = models[mode]301            se = m.emb(te["seq"][:1])302            if mode == "avg":303                delta = 0.0304            else:305                base = se.mean(dim=1)306                vs = torch.cat([m.attn(m.emb(torch.tensor([c])), se)[0]307                                for c in range(VOCAB)])308                delta = (vs - base).norm(dim=1).mean().item()309            print(f"  {mode:<12} 汇聚表示随候选的变化 = {delta:.4f}")310311    # ── 注意力权重分析：同一序列、不同候选，权重应不同 ──────────────────────312    print("\nTarget Attention 权重（前 3 个测试样本，对比同一序列对不同候选的关注）：")313    m_ta = models["target_attn"]314    with torch.no_grad():315        for i in range(3):316            seq_i = te["seq"][i:i + 1]317            seq_emb = m_ta.emb(seq_i)318            cand_ids = [te["target"][i].item(), (te["target"][i].item() + 5) % VOCAB]319            for c in cand_ids:320                c_emb = m_ta.emb(torch.tensor([c]))321                _, w = m_ta.attn(c_emb, seq_emb)322                print(f"  样本{i} 候选={c:2d} 权重="323                      f"{[f'{v:.2f}' for v in w[0].tolist()]}")324325    # ── 精排推理演示 ─────────────────────────────────────────────────────────326    p = m_ta.predict_proba(te["seq"][:20], te["target"][:20],327                           te["feat"][:20]).tolist()328    top5 = sorted(enumerate(p), key=lambda t: -t[1])[:5]329    print("\n精排 Top-5（候选集 20 条）:")330    for rank_i, (idx, s) in enumerate(top5, 1):331        print(f"  #{rank_i}  item={idx:2d}  p_click={s:.4f}  "332              f"label={te['y'][idx].item()}")