cl4srec.py

1"""2CL4SRec 对比学习序列召回（第10章 §10.4.5「对比学习序列召回 - CL4SRec召回」）34论文：Xie et al., "Contrastive Learning for Sequential Recommendation", ICDE 20225核心思想：6- 以 SASRec（因果 Self-Attention）作为序列编码器；7- 对同一条用户行为序列做两次随机数据增强，得到两个视图，互为正样本对；8  同一个 Mini-Batch 内其它用户的增强序列作为负样本；9- 三种序列数据增强算子：10    * Item Crop   ：随机截取一段连续子序列（学习局部兴趣，适配不同序列长度）11    * Item Mask   ：随机把部分物品替换为 [MASK] 标记（抗行为缺失 / 噪声）12    * Item Reorder：随机打乱一段连续子序列的顺序（降低对严格顺序的依赖）13- 对比损失采用 InfoNCE：14      L_cl = -log( exp(sim(h_i, h_j)/τ) / Σ_k exp(sim(h_i, h_k)/τ) )15      实现为 F.cross_entropy(sim_matrix / tau, torch.arange(B))16- 多任务联合训练：L_total = L_rec + λ·L_cl （L_rec 为下一物品预测损失）1718本实现：PyTorch autograd，增强操作用 Python 列表操作，编码器用19nn.TransformerEncoderLayer + causal mask，完整反向传播由 autograd 自动完成。20"""2122import math23import random2425import torch26import torch.nn as nn27import torch.nn.functional as F28import torch.optim as optim293031# ── 三种序列数据增强算子 ───────────────────────────────────────────────────────3233class SeqAugment:34    """CL4SRec 的三种数据增强算子：Crop / Mask / Reorder。"""3536    def __init__(self, mask_token: int, eta: float = 0.6, gamma: float = 0.3, beta: float = 0.5):37        self.mask_token = mask_token   # [MASK] 占位 id（= n_items）38        self.eta = eta                 # crop 保留比例39        self.gamma = gamma             # mask 比例40        self.beta = beta               # reorder 子序列比例4142    def item_crop(self, seq: list[int]) -> list[int]:43        """随机截取长度为 ⌊η·T⌋ 的连续子序列。"""44        T = len(seq)45        sub_len = max(1, int(self.eta * T))46        start = random.randint(0, T - sub_len)47        return seq[start:start + sub_len]4849    def item_mask(self, seq: list[int]) -> list[int]:50        """随机将 ⌊γ·T⌋ 个位置替换为 [MASK]。"""51        out = list(seq)52        n_mask = max(1, int(self.gamma * len(seq)))53        for idx in random.sample(range(len(seq)), n_mask):54            out[idx] = self.mask_token55        return out5657    def item_reorder(self, seq: list[int]) -> list[int]:58        """随机选取长度为 ⌊β·T⌋ 的连续片段并打乱其内部顺序。"""59        out = list(seq)60        T = len(seq)61        sub_len = max(2, int(self.beta * T)) if T >= 2 else T62        sub_len = min(sub_len, T)63        start = random.randint(0, T - sub_len)64        chunk = out[start:start + sub_len]65        random.shuffle(chunk)66        out[start:start + sub_len] = chunk67        return out6869    def random_augment(self, seq: list[int]) -> tuple[str, list[int]]:70        """等概率随机挑选一种增强算子。"""71        name = random.choice(["crop", "mask", "reorder"])72        return name, {"crop": self.item_crop,73                      "mask": self.item_mask,74                      "reorder": self.item_reorder}[name](seq)757677# ── CL4SRec 模型 ──────────────────────────────────────────────────────────────7879class CL4SRec(nn.Module):80    """81    CL4SRec = SASRec 编码器 + 自监督对比学习辅助任务。8283    组件：84      item_emb : nn.Embedding [(n_items + 1) × dim]，最后一行为 [MASK] 的 Embedding85      pos_emb  : nn.Embedding [max_len × dim]86      transformer_layer: nn.TransformerEncoderLayer（因果 mask，单头）87    """8889    def __init__(self, n_items: int, dim: int = 16, max_len: int = 10,90                 lr: float = 0.05, reg: float = 1e-4,91                 tau: float = 0.5, lam: float = 0.3, n_heads: int = 1,92                 dropout: float = 0.1):93        super().__init__()94        self.n_items = n_items95        self.dim = dim96        self.max_len = max_len97        self.lr = lr98        self.reg = reg99        self.tau = tau                 # 温度系数 τ100        self.lam = lam                 # 对比损失权重 λ101        self.mask_token = n_items102103        scale = math.sqrt(1.0 / dim)104        self.item_emb = nn.Embedding(n_items + 1, dim)   # +1 for [MASK]105        self.pos_emb = nn.Embedding(max_len, dim)106        nn.init.uniform_(self.item_emb.weight, -scale, scale)107        nn.init.uniform_(self.pos_emb.weight, -scale, scale)108109        self.transformer_layer = nn.TransformerEncoderLayer(110            d_model=dim,111            nhead=n_heads,112            dim_feedforward=dim * 4,113            dropout=dropout,114            activation="relu",115            batch_first=True,116            norm_first=False,117        )118        self.aug = SeqAugment(self.mask_token)119120    def _encode_batch(self, seqs: list[list[int]], pool: str = "last") -> torch.Tensor:121        """122        编码一批序列，自动左填充到同长度。123        pool="last"：取最后位置（下一物品预测任务）；124        pool="mean"：均值池化（对比学习视图，对增强后长度不一的序列更稳健）。125        返回 [B, dim]。126        """127        T = max(len(s) for s in seqs)128        T = min(T, self.max_len)129        B = len(seqs)130131        padded = torch.zeros(B, T, dtype=torch.long)132        for i, seq in enumerate(seqs):133            s = seq[-T:]134            padded[i, T - len(s):] = torch.tensor(s, dtype=torch.long)135136        positions = torch.arange(T).unsqueeze(0)                   # [1, T]137        x = self.item_emb(padded) + self.pos_emb(positions)        # [B, T, dim]138139        causal_mask = nn.Transformer.generate_square_subsequent_mask(T, device=x.device)140        out = self.transformer_layer(x, src_mask=causal_mask, is_causal=True)  # [B, T, dim]141142        if pool == "mean":143            # Mean over non-padding positions only144            lengths = torch.tensor([min(len(s), T) for s in seqs], dtype=torch.float)145            mask = torch.zeros(B, T)146            for i, seq in enumerate(seqs):147                l = min(len(seq), T)148                mask[i, T - l:] = 1.0149            return (out * mask.unsqueeze(-1)).sum(1) / lengths.unsqueeze(-1)150        else:151            return out[:, -1, :]152153    # ── 主任务：下一物品预测 ─────────────────────────────────────────────154    def _rec_loss(self, contexts: list[list[int]], pos_ids: list[int],155                  neg_ids: list[int]) -> torch.Tensor:156        """BCE 形式的下一物品预测损失。"""157        h = self._encode_batch(contexts, pool="last")               # [B, dim]158        v_pos = self.item_emb(torch.tensor(pos_ids, dtype=torch.long))159        v_neg = self.item_emb(torch.tensor(neg_ids, dtype=torch.long))160        s_pos = (h * v_pos).sum(-1)161        s_neg = (h * v_neg).sum(-1)162        return (F.binary_cross_entropy_with_logits(s_pos, torch.ones_like(s_pos))163                + F.binary_cross_entropy_with_logits(s_neg, torch.zeros_like(s_neg)))164165    # ── 辅助任务：InfoNCE 对比学习 ──────────────────────────────────────166    def _cl_loss(self, view_seqs_a: list[list[int]],167                 view_seqs_b: list[list[int]]) -> torch.Tensor:168        """169        Batch 内 InfoNCE：view_seqs_a[i] 与 view_seqs_b[i] 来自同一序列，互为正样本。170        L_cl = F.cross_entropy(sim_matrix / τ, torch.arange(B))171        """172        ha = F.normalize(self._encode_batch(view_seqs_a, pool="mean"), dim=-1)  # [B, dim]173        hb = F.normalize(self._encode_batch(view_seqs_b, pool="mean"), dim=-1)  # [B, dim]174175        # sim_matrix[i, j] = cosine(ha[i], hb[j]) / τ176        sim = (ha @ hb.T) / self.tau                                # [B, B]177        labels = torch.arange(sim.size(0), dtype=torch.long)178        # Both directions: ha predicts hb and hb predicts ha179        return (F.cross_entropy(sim, labels) + F.cross_entropy(sim.T, labels)) / 2.0180181    def train_model(self, user_seqs: dict[int, list[int]], epochs: int = 6,182                    batch_size: int = 4) -> None:183        """多任务联合训练：L_total = L_rec + λ·L_cl。"""184        uids = list(user_seqs.keys())185        optimizer = optim.Adam(self.parameters(), lr=self.lr, weight_decay=self.reg)186187        for epoch in range(epochs):188            random.shuffle(uids)189            rec_sum, cl_sum, n_batch, n_rec = 0.0, 0.0, 0, 0190191            for b in range(0, len(uids), batch_size):192                batch_uids = uids[b:b + batch_size]193                contexts_rec, pos_ids, neg_ids = [], [], []194                view_seqs_a, view_seqs_b = [], []195196                for uid in batch_uids:197                    seq = user_seqs[uid]198                    # 主任务：序列前缀预测下一物品（每个位置都作为监督信号）199                    for t in range(1, len(seq)):200                        contexts_rec.append(seq[:t])201                        pos_ids.append(seq[t])202                        neg = random.randrange(self.n_items)203                        while neg == seq[t]:204                            neg = random.randrange(self.n_items)205                        neg_ids.append(neg)206                        n_rec += 1207                    # 辅助任务：同一序列两次独立增强 → 正样本对208                    _, va = self.aug.random_augment(seq[:-1])209                    _, vb = self.aug.random_augment(seq[:-1])210                    view_seqs_a.append(va)211                    view_seqs_b.append(vb)212213                optimizer.zero_grad()214                rec_loss = self._rec_loss(contexts_rec, pos_ids, neg_ids)215                cl_loss = self._cl_loss(view_seqs_a, view_seqs_b)216                loss = rec_loss + self.lam * cl_loss217                loss.backward()218                optimizer.step()219220                rec_sum += rec_loss.item()221                cl_sum += cl_loss.item()222                n_batch += 1223224            print(f"  epoch {epoch:2d}  L_rec={rec_sum / max(n_batch, 1):.4f}  "225                  f"L_cl={cl_sum / max(n_batch, 1):.4f}  "226                  f"L_total={rec_sum / max(n_batch, 1) + self.lam * cl_sum / max(n_batch, 1):.4f}")227228    def recommend(self, seq: list[int], top_k: int = 5,229                  exclude: set[int] | None = None) -> list[tuple[int, float]]:230        """用序列表示 h 对全量物品打分（线上等价于 ANN 检索）。"""231        exclude = exclude if exclude is not None else set(seq)232        with torch.no_grad():233            h = self._encode_batch([seq], pool="last")[0]           # [dim]234            scores = self.item_emb.weight[:self.n_items] @ h        # [n_items]235            for iid in exclude:236                scores[iid] = -1e9237            top_vals, top_idx = scores.topk(top_k)238        return [(int(i), float(s)) for i, s in zip(top_idx.tolist(), top_vals.tolist())]239240241# ── 快速验证 ──────────────────────────────────────────────────────────────────242if __name__ == "__main__":243    random.seed(42)244    torch.manual_seed(42)245    n_items, n_users = 20, 8246    user_seqs = {u: random.sample(range(n_items), k=8) for u in range(n_users)}247248    demo_seq = user_seqs[0]249    print("三种数据增强算子作用在同一条序列上的示例：")250    print(f"  原始序列      : {demo_seq}")251    aug = SeqAugment(mask_token=n_items)252    print(f"  Item Crop     : {aug.item_crop(demo_seq)}")253    print(f"  Item Mask     : {aug.item_mask(demo_seq)}   （{n_items} 表示 [MASK]）")254    print(f"  Item Reorder  : {aug.item_reorder(demo_seq)}")255256    model = CL4SRec(n_items, dim=16, max_len=10, lr=0.05, tau=0.5, lam=0.3)257    print(f"\nCL4SRec 训练（{n_users} 用户，λ={model.lam}，τ={model.tau}）：")258    model.train_model(user_seqs, epochs=6, batch_size=4)259260    recs = model.recommend(demo_seq, top_k=5)261    print(f"\n用户 0 序列 {demo_seq} 的 Top-5 召回：")262    for iid, sc in recs:263        print(f"  item {iid:2d}  score={sc:.4f}")264265    # 对比学习效果观察：同一序列两视图相似度 vs 不同用户序列相似度266    pos_sims, neg_sims = [], []267    with torch.no_grad():268        reps = {u: F.normalize(model._encode_batch([s], pool="mean")[0], dim=-1)269                for u, s in user_seqs.items()}270        for u, s in user_seqs.items():271            _, va = aug.random_augment(s)272            _, vb = aug.random_augment(s)273            ha = F.normalize(model._encode_batch([va], pool="mean")[0], dim=-1)274            hb = F.normalize(model._encode_batch([vb], pool="mean")[0], dim=-1)275            pos_sims.append((ha @ hb).item())276            for v in user_seqs:277                if v != u:278                    neg_sims.append((ha @ reps[v]).item())279280    print(f"\n同序列两视图平均余弦相似度={sum(pos_sims) / len(pos_sims):.4f}  "281          f"不同用户序列平均余弦相似度={sum(neg_sims) / len(neg_sims):.4f}")