bert4rec.py

1r"""2BERT4Rec 双向 Transformer 序列召回（第10章 §10.4.3「Transformer序列召回 - BERT4Rec模型」）34论文：Sun et al., "BERT4Rec: Sequential Recommendation with Bidirectional Encoder5      Representations from Transformer", CIKM 201967核心思想：8- 用双向 Transformer Encoder 编码用户行为序列，训练目标为 Cloze（完形填空 / MLM）：9  随机把序列中的若干物品替换成 [MASK]，再用「前后双向上下文」还原被遮蔽的物品。10      [i1,i2,i3,i4,i5] => [i1,[MASK],i3,[MASK],i5]11      L = - Σ_{t in M} log P(i_t | S_\t)12- 推理时在序列末尾追加一个 [MASK] token，用该位置的隐状态预测下一个物品，13  再用它作为 User Embedding 去 ANN 索引中检索 Top-K。1415与 SASRec 的关键区别（见同目录 sasrec.py）：16- 注意力掩码：SASRec 用因果 Mask（位置 t 只能看 ≤t）；BERT4Rec 不加因果 Mask，双向看全序列。17- 训练目标：SASRec 自回归预测下一个物品；BERT4Rec 随机 Mask 若干位置后用双向上下文还原。18- 监督信号：SASRec 每个位置 1 个；BERT4Rec 一条序列可采样多组 Mask，数据利用率更高。19- 推理方式：SASRec 直接取最后位置隐状态；BERT4Rec 需在末尾追加 [MASK] 再取该位置隐状态。2021本实现：PyTorch autograd，使用 nn.TransformerEncoder（无因果 Mask）+ Tied Embedding 输出层。22完整反向传播由 autograd 自动完成。23"""2425import math26import random2728import torch29import torch.nn as nn30import torch.nn.functional as F31import torch.optim as optim323334# ── BERT4Rec 模型 ─────────────────────────────────────────────────────────────3536class BERT4Rec(nn.Module):37    """38    精简版 BERT4Rec：双向多头 Transformer Encoder + Cloze 训练目标。3940    参数：41      item_emb: nn.Embedding [(n_items+1) × dim]  物品 Embedding，下标 n_items 为 [MASK] token42      pos_emb:  nn.Embedding [max_len × dim]       位置 Embedding43      transformer: nn.TransformerEncoder（无因果 Mask，双向注意力）44      out_bias: [n_items]  输出层 bias（输出矩阵与 item_emb 共享，Tied Embedding）45    """4647    def __init__(self, n_items: int, dim: int = 16, max_len: int = 10, n_layers: int = 2,48                 n_heads: int = 2, d_ff: int = 32, lr: float = 0.05, mask_prob: float = 0.3,49                 dropout: float = 0.1):50        super().__init__()51        assert dim % n_heads == 0, "dim 必须能被 n_heads 整除"52        self.n_items = n_items53        self.mask_id = n_items      # 下标 n_items 即 [MASK] token54        self.dim = dim55        self.max_len = max_len56        self.lr = lr57        self.mask_prob = mask_prob5859        scale = math.sqrt(1.0 / dim)60        self.item_emb = nn.Embedding(n_items + 1, dim)   # +1 for [MASK]61        self.pos_emb = nn.Embedding(max_len, dim)62        nn.init.uniform_(self.item_emb.weight, -scale, scale)63        nn.init.uniform_(self.pos_emb.weight, -scale, scale)6465        self.out_bias = nn.Parameter(torch.zeros(n_items))6667        # 双向 Transformer Encoder（关键点 1：无因果 Mask，batch_first=True）68        encoder_layer = nn.TransformerEncoderLayer(69            d_model=dim,70            nhead=n_heads,71            dim_feedforward=d_ff,72            dropout=dropout,73            activation="relu",74            batch_first=True,75            norm_first=False,76        )77        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=n_layers)7879    def _forward(self, tokens: torch.Tensor) -> torch.Tensor:80        """81        完整前向：Embedding → 双向 Transformer Encoder（无因果 Mask）82        tokens: [B, T]  包含 [MASK] id 的 token 序列83        返回: [B, T, dim]  每个位置的隐状态84        """85        B, T = tokens.shape86        positions = torch.arange(T, device=tokens.device).unsqueeze(0)  # [1, T]87        x = self.item_emb(tokens) + self.pos_emb(positions)             # [B, T, dim]88        # 注意：这里不传 src_mask，双向看全序列（与 SASRec 的因果 mask 的关键区别）89        return self.transformer(x)  # [B, T, dim]9091    # ── Cloze 训练（关键点 2：随机 Mask + 双向上下文还原） ───────────────────92    def _make_cloze_sample(self, seq: list[int]) -> tuple[list[int], list[tuple[int, int]]]:93        """以 mask_prob 随机遮蔽序列中的物品；至少遮蔽一个位置。"""94        tokens, targets = list(seq), []95        for t, iid in enumerate(seq):96            if random.random() < self.mask_prob:97                tokens[t] = self.mask_id98                targets.append((t, iid))99        if not targets:                           # 保证至少有一个监督信号100            t = random.randrange(len(seq))101            tokens[t] = self.mask_id102            targets.append((t, seq[t]))103        return tokens, targets104105    def _cloze_loss(self, tokens: torch.Tensor,106                    mask_positions: list[list[int]],107                    mask_labels: list[list[int]]) -> torch.Tensor:108        """109        批量 Cloze 交叉熵损失。110        tokens:         [B, T]111        mask_positions: list of per-sample masked position indices112        mask_labels:    list of per-sample true item ids at masked positions113        Tied Embedding: logits = h @ item_emb.T + out_bias114        """115        hs = self._forward(tokens)   # [B, T, dim]116        total_loss = torch.zeros((), device=tokens.device)117        n = 0118        for b, (positions, labels) in enumerate(zip(mask_positions, mask_labels)):119            if not positions:120                continue121            pos_t = torch.tensor(positions, dtype=torch.long, device=tokens.device)122            lbl_t = torch.tensor(labels, dtype=torch.long, device=tokens.device)123            h = hs[b, pos_t, :]                          # [M, dim]124            # Tied Embedding 输出层（共享 item_emb，不含 [MASK] 行）125            logits = h @ self.item_emb.weight[:self.n_items].T + self.out_bias  # [M, n_items]126            total_loss = total_loss + F.cross_entropy(logits, lbl_t)127            n += 1128        return total_loss / max(n, 1)129130    def train_model(self, user_seqs: dict[int, list[int]], epochs: int = 6,131                    batch_size: int = 32) -> None:132        """133        Cloze 训练主循环。每条序列产出两类样本：134          (1) 随机 Mask 样本：学习双向上下文还原能力；135          (2) 仅 Mask 末尾位置的样本：论文中用于缓解「训练目标与线上 next-item 预测」136              之间的不一致（与推理时末尾追加 [MASK] 的做法对齐）。137        """138        seqs = [s[-self.max_len:] for s in user_seqs.values() if len(s) >= 2]139        optimizer = optim.Adam(self.parameters(), lr=self.lr)140141        for epoch in range(epochs):142            random.shuffle(seqs)143            total, n_batches = 0.0, 0144145            # Build all samples: (tokens, positions, labels) per sequence × 2 variants146            all_samples: list[tuple[list[int], list[int], list[int]]] = []147            for seq in seqs:148                toks, tgts = self._make_cloze_sample(seq)149                all_samples.append((toks, [t for t, _ in tgts], [v for _, v in tgts]))150                # Last-position variant151                last_toks = list(seq[:-1]) + [self.mask_id]152                all_samples.append((last_toks, [len(seq) - 1], [seq[-1]]))153154            for start in range(0, len(all_samples), batch_size):155                chunk = all_samples[start:start + batch_size]156                T = max(len(s[0]) for s in chunk)157                # Pad tokens to uniform length (left-pad with 0)158                padded = torch.zeros(len(chunk), T, dtype=torch.long)159                for i, (toks, _, _) in enumerate(chunk):160                    offset = T - len(toks)161                    padded[i, offset:] = torch.tensor(toks, dtype=torch.long)162                    # Adjust positions to account for left-padding163                mask_positions = []164                mask_labels = []165                for i, (toks, positions, labels) in enumerate(chunk):166                    offset = T - len(toks)167                    mask_positions.append([p + offset for p in positions])168                    mask_labels.append(labels)169170                optimizer.zero_grad()171                loss = self._cloze_loss(padded, mask_positions, mask_labels)172                loss.backward()173                optimizer.step()174175                total += loss.item()176                n_batches += 1177178            print(f"  epoch {epoch:2d}  cloze_loss={total / max(n_batches, 1):.4f}")179180    # ── 推理（关键点 3：末尾追加 [MASK]） ───────────────────────────────────181    def user_embedding(self, seq: list[int]) -> torch.Tensor:182        """在序列末尾追加 [MASK]，取该位置隐状态作为 User Embedding，shape [dim]。"""183        with torch.no_grad():184            tokens = list(seq[-(self.max_len - 1):]) + [self.mask_id]185            tok_t = torch.tensor(tokens, dtype=torch.long).unsqueeze(0)  # [1, T]186            hs = self._forward(tok_t)                                     # [1, T, dim]187            return hs[0, -1, :]                                           # [dim]188189    def recommend(self, seq: list[int], top_k: int = 5,190                  exclude_seen: bool = True) -> list[tuple[int, float]]:191        """给定历史序列，预测下一个物品 Top-K。"""192        with torch.no_grad():193            u = self.user_embedding(seq)194            # Tied Embedding 打分195            logits = self.item_emb.weight[:self.n_items] @ u + self.out_bias  # [n_items]196            if exclude_seen:197                for iid in seq:198                    logits[iid] = -1e9199            top_vals, top_idx = logits.topk(top_k)200        return [(int(i), float(s)) for i, s in zip(top_idx.tolist(), top_vals.tolist())]201202    def fill_mask(self, tokens: list[int], top_k: int = 3) -> dict[int, list[tuple[int, float]]]:203        """完形填空：对序列中所有 [MASK] 位置，用双向上下文给出 Top-K 候选。"""204        with torch.no_grad():205            tok_t = torch.tensor(tokens, dtype=torch.long).unsqueeze(0)  # [1, T]206            hs = self._forward(tok_t)                                     # [1, T, dim]207            result: dict[int, list[tuple[int, float]]] = {}208            for t, tok in enumerate(tokens):209                if tok != self.mask_id:210                    continue211                logits = hs[0, t, :] @ self.item_emb.weight[:self.n_items].T + self.out_bias212                probs = logits.softmax(dim=-1)213                top_vals, top_idx = probs.topk(top_k)214                result[t] = list(zip(top_idx.tolist(), top_vals.tolist()))215        return result216217218# ── 快速验证 ──────────────────────────────────────────────────────────────────219if __name__ == "__main__":220    random.seed(42)221    torch.manual_seed(42)222    n_users, n_items, seq_len = 8, 20, 8223    # 合成序列：兴趣簇内按 +1 递进点击，前后文都能提供预测线索（适合双向建模）224    user_seqs: dict[int, list[int]] = {}225    for u in range(n_users):226        start = (u % 4) * 5 + random.randint(0, 2)227        user_seqs[u] = [(start + k) % n_items for k in range(seq_len)]228229    print("合成用户行为序列：")230    for u in range(n_users):231        print(f"  user {u}: {user_seqs[u]}")232233    model = BERT4Rec(n_items, dim=16, max_len=10, n_layers=2, n_heads=2, d_ff=32, lr=0.05)234    print("\nBERT4Rec Cloze 训练（双向注意力，无因果 Mask）：")235    model.train_model(user_seqs, epochs=6)236237    seq, probe, truth = user_seqs[0], user_seqs[0][:-1], user_seqs[0][-1]238    print(f"\n历史序列 {probe}（真实下一个物品 = {truth}）\n末尾追加 [MASK] 后的 Top-5 召回：")239    for iid, s in model.recommend(probe, top_k=5, exclude_seen=False):240        flag = "  <-- 命中" if iid == truth else ""241        print(f"  item {iid}: score={s:.4f}{flag}")242243    # 演示双向能力：遮蔽序列中间位置，用前后文同时推断244    mid, tokens = len(seq) // 2, list(seq)245    tokens[mid] = model.mask_id246    print(f"\n完形填空演示：把位置 {mid}（真实物品 {seq[mid]}）替换为 [MASK]")247    print(f"  输入 tokens（{model.mask_id} 表示 [MASK]）：{tokens}")248    for pos, cands in model.fill_mask(tokens, top_k=3).items():249        pretty = ", ".join(f"item {i}(p={p:.3f})" for i, p in cands)250        print(f"  位置 {pos} 的双向预测：{pretty}")