diger_differentiable_sid.py

1r"""对应正文第 9 章 §9.6.1「DIGER 算法」的 DRIL(可微语义 ID + 探索式学习)范式。23正文公式(本实现逐条复刻):4  * logits:        l_{v,j,i} = sim(r_{v,j}, e_i)5  * Gumbel 噪声:   g = -log(-log u), u ~ U(0,1)      (逆变换采样)6  * Gumbel-Softmax: y~_{v,j,i} = exp((l+g)/τ) / Σ_k exp((l+g)/τ)7  * 前向硬选择:    c_{v,j} = argmax_i (l_{v,j,i} + g_{v,j,i})8  * 反向软码字:    ē_{v,j} = Σ_i y~_{v,j,i} · e_i      ("Forward Hard, Backward Soft")9  * SDUD:          L_σ = L_gen / (2(σ+λ)²) + log(σ+λ),  σ* = max(0, sqrt(L_gen) - λ)10  * FrqUD:         f_i^(e) = β f_i^(e-1) + (1-β) f^_i^(e);  γ = r/K;11                   高频码字(f_i > γ)保留 Gumbel 噪声, 低频码字用确定性 Softmax1213本实现:14  1. 用一个小的「物品表示 -> 可微语义 ID -> Next-SID 预测」联合模型做端到端训练;15  2. 对比三种离散化方式: 纯 STE(argmax 最近邻) / DRIL(Gumbel, 固定 τ) / DRIL+FrqUD;16  3. 打印 Gumbel-Max Trick 的数值验证表、σ* 随 L_gen 衰减的轨迹表、17     以及三种方式的码本利用率(perplexity)与生成准确率对比表。18  4. 生成式推荐模型为教学规模的单层 GRU + 线性头, 不是正文中的大模型。19"""2021from __future__ import annotations2223import math24import random2526import torch27import torch.nn as nn28import torch.nn.functional as F2930L_LAYER = 3          # 语义 ID 层数 m31K_CODE = 8           # 每层码字数 K32D_LAT = 16           # 码字维度33N_ITEM = 120343536# ---------------------------------------------------------------- DRIL37def gumbel_noise(shape, gen, sigma=1.0):38    """g = -log(-log u), u ~ Uniform(0,1); sigma 为正文 SDUD 中的噪声尺度。"""39    u = torch.rand(shape, generator=gen).clamp_(1e-9, 1 - 1e-9)40    return sigma * (-torch.log(-torch.log(u)))414243def dril_quantize(r, codebook, tau, gen, sigma=1.0, noise_mask=None):44    """返回 (硬 Token, 可微软码字 ē, 软概率 y~)。4546    sim 用负欧氏距离(正文仅写 sim(·,·), 此处与 RQ-VAE 保持一致取 -||r-e||²)。47    noise_mask: FrqUD 用, shape (K,), 1 表示该码字加 Gumbel 噪声, 0 表示关闭。48    """49    logits = -torch.cdist(r, codebook) ** 2                 # (B, K)50    g = gumbel_noise(logits.shape, gen, sigma)51    if noise_mask is not None:52        g = g * noise_mask[None, :]53    y_soft = F.softmax((logits + g) / tau, dim=-1)54    hard = torch.argmax(logits + g, dim=-1)                 # 正文: 硬 argmax55    e_bar = y_soft @ codebook                               # 正文: 软码字56    # Forward Hard / Backward Soft: 前向用硬码字数值, 梯度走软码字57    e_hard = codebook[hard]58    e_out = e_hard + (e_bar - e_bar.detach())59    return hard, e_out, y_soft606162def ste_quantize(r, codebook):63    """对照组: 普通 VQ-VAE 的 STE(最近邻 + 直通梯度), 无探索。"""64    d = torch.cdist(r, codebook) ** 265    hard = torch.argmax(-d, dim=-1)66    e = codebook[hard]67    return hard, r + (e - r).detach(), F.softmax(-d, dim=-1)686970class DigerModel(nn.Module):71    """物品塔(Tokenizer) + 生成式推荐塔, 两者通过可微语义 ID 相连。"""7273    def __init__(self, d_in=24, mode="dril"):74        super().__init__()75        self.mode = mode76        self.enc = nn.Sequential(nn.Linear(d_in, 32), nn.GELU(),77                                 nn.Linear(32, L_LAYER * D_LAT))78        self.codebooks = nn.ParameterList(79            [nn.Parameter(torch.randn(K_CODE, D_LAT) * 0.5) for _ in range(L_LAYER)])80        self.gru = nn.GRU(L_LAYER * D_LAT, 48, batch_first=True)81        self.heads = nn.ModuleList([nn.Linear(48, K_CODE) for _ in range(L_LAYER)])8283    def tokenize(self, feat, tau, gen, sigma=1.0, masks=None):84        """feat: (N, d_in) -> tokens (N, L), 可微 item 表示 (N, L*D)。8586        正文中 DIGER 的连续表示直接是 r_v=(r_{v,1},...,r_{v,m}), 每个位置 j 独立地87        与第 j 个码本做相似度比较, 因此这里按层并行量化(不做残差累减)。88        """89        r = self.enc(feat).view(-1, L_LAYER, D_LAT)90        toks, embs, probs = [], [], []91        for l in range(L_LAYER):92            cur = r[:, l]93            if self.mode == "ste":94                t, e, p = ste_quantize(cur, self.codebooks[l])95            else:96                m = None if masks is None else masks[l]97                t, e, p = dril_quantize(cur, self.codebooks[l], tau, gen, sigma, m)98            toks.append(t)99            embs.append(e)100            probs.append(p)101        return torch.stack(toks, 1), torch.cat(embs, -1), probs102103    def gen_loss(self, seq_emb, target_tok):104        """Next-SID Prediction: L_gen = -Σ_l log P(c_l | 历史序列)。"""105        h, _ = self.gru(seq_emb)106        last = h[:, -1]107        loss = 0.0108        correct = torch.ones(last.shape[0], dtype=torch.bool)109        for l in range(L_LAYER):110            logit = self.heads[l](last)111            loss = loss + F.cross_entropy(logit, target_tok[:, l])112            correct &= (logit.argmax(-1) == target_tok[:, l])113        return loss, correct.float().mean()114115116def perplexity(tokens):117    """码本 perplexity = exp(H(p)); 越接近 K 说明利用率越高。"""118    out = []119    for l in range(L_LAYER):120        cnt = torch.bincount(tokens[:, l], minlength=K_CODE).float()121        p = cnt / cnt.sum()122        h = -(p[p > 0] * p[p > 0].log()).sum()123        out.append(float(h.exp()))124    return out125126127def run(mode, feat, seqs, tgt, gen, epochs=18, tau=1.0, frqud=False,128        sdud=True, lam=0.3, beta=0.9, r_hot=1.5):129    torch.manual_seed(42)130    model = DigerModel(feat.shape[1], mode="ste" if mode == "ste" else "dril")131    opt = torch.optim.Adam(model.parameters(), lr=0.02)132    freq = [torch.full((K_CODE,), 1.0 / K_CODE) for _ in range(L_LAYER)]133    sigma, traj = 1.0, []134    for ep in range(epochs):135        masks = None136        if frqud:137            gamma = r_hot / K_CODE138            masks = [(f > gamma).float() for f in freq]139        toks, emb, _ = model.tokenize(feat, tau, gen, sigma, masks)140        seq_emb = emb[seqs]                        # (B, T, L*D)141        tgt_tok = toks[tgt]142        loss, acc = model.gen_loss(seq_emb, tgt_tok)143        opt.zero_grad()144        loss.backward()145        opt.step()146        lg = float(loss)147        if sdud and mode != "ste":                 # SDUD: σ* = max(0, sqrt(L_gen)-λ)148            sigma = max(0.0, math.sqrt(max(lg, 0.0)) - lam)149        if frqud:                                  # 频率 EMA150            for l in range(L_LAYER):151                cnt = torch.bincount(toks[:, l], minlength=K_CODE).float()152                freq[l] = beta * freq[l] + (1 - beta) * cnt / cnt.sum()153        traj.append((ep + 1, lg, sigma, float(acc)))154    with torch.no_grad():155        toks, emb, _ = model.tokenize(feat, 1e-6, gen, 0.0)   # 推理: 确定性 argmax156        seq_emb = emb[seqs]157        loss, acc = model.gen_loss(seq_emb, toks[tgt])158    uniq = len({tuple(t.tolist()) for t in toks})159    return traj, perplexity(toks), float(loss), float(acc), uniq, freq160161162def main() -> None:163    random.seed(42)164    torch.manual_seed(42)165    gen = torch.Generator().manual_seed(42)166167    # ---------- 1. Gumbel-Max Trick 数值验证 ----------168    logits = torch.tensor([2.0, 1.0, 0.0, -1.0])169    N = 200000170    g = gumbel_noise((N, 4), gen)171    emp = torch.bincount((logits + g).argmax(-1), minlength=4).float() / N172    ref = F.softmax(logits, -1)173    print("=" * 84)174    print("表 1  Gumbel-Max Trick 验证: argmax_i(l_i + g_i) ~ Softmax(l)")175    print(f"      采样 {N} 次, logits = {logits.tolist()}")176    print("-" * 84)177    print(f"{'类别 i':>8}{'logit':>10}{'Softmax(l)':>14}{'argmax 经验频率':>18}"178          f"{'绝对误差':>12}")179    for i in range(4):180        print(f"{i:>8}{logits[i]:>10.1f}{ref[i]:>14.4f}{emp[i]:>18.4f}"181              f"{abs(ref[i] - emp[i]):>12.4f}")182183    # ---------- 2. 合成数据 ----------184    D_IN, T, B = 24, 6, 256185    centers = torch.randn(6, D_IN, generator=gen) * 2.0186    item_grp = torch.randint(0, 6, (N_ITEM,), generator=gen)187    feat = centers[item_grp] + torch.randn(N_ITEM, D_IN, generator=gen) * 0.4188    seqs, tgt = [], []189    for _ in range(B):190        grp = int(torch.randint(0, 6, (1,), generator=gen))191        pool = (item_grp == grp).nonzero().flatten()192        pick = pool[torch.randint(0, len(pool), (T + 1,), generator=gen)]193        seqs.append(pick[:T])194        tgt.append(int(pick[T]))195    seqs = torch.stack(seqs)196    tgt = torch.tensor(tgt)197198    # ---------- 3. SDUD 轨迹 ----------199    traj_d, ppl_d, loss_d, acc_d, uq_d, _ = run("dril", feat, seqs, tgt, gen)200    print()201    print("表 2  SDUD 噪声尺度自适应衰减(λ=0.3): σ* = max(0, sqrt(L_gen) - λ)")202    print("-" * 84)203    print(f"{'epoch':>7}{'L_gen':>12}{'sqrt(L_gen)':>14}{'σ*':>10}{'训练 SID 准确率':>18}")204    for ep, lg, sg, ac in traj_d:205        if ep % 3 == 1 or ep == len(traj_d):206            print(f"{ep:>7}{lg:>12.4f}{math.sqrt(max(lg, 0)):>14.4f}{sg:>10.4f}"207                  f"{ac:>18.4f}")208    print("-" * 84)209    print("可见生成损失下降 -> 最优噪声尺度自动减小 -> 从探索平滑过渡到利用。")210211    # ---------- 4. 三种离散化方式对比 ----------212    traj_s, ppl_s, loss_s, acc_s, uq_s, _ = run("ste", feat, seqs, tgt, gen)213    traj_f, ppl_f, loss_f, acc_f, uq_f, _ = run("dril", feat, seqs, tgt, gen,214                                                frqud=True)215    print()216    print("表 3  离散化方式对比(推理时统一关闭噪声, 走确定性 argmax)")217    print("-" * 84)218    print(f"{'方式':<22}{'L_gen':>10}{'SID 全命中率':>14}"219          f"{'码本 perplexity(L1/L2/L3)':>28}{'不同 SID 数':>12}")220    for name, lo, ac, pp, uq in (221            ("STE(最近邻直通)", loss_s, acc_s, ppl_s, uq_s),222            ("DRIL(Gumbel+SDUD)", loss_d, acc_d, ppl_d, uq_d),223            ("DRIL+FrqUD", loss_f, acc_f, ppl_f, uq_f)):224        pps = "/".join(f"{x:.2f}" for x in pp)225        print(f"{name:<22}{lo:>10.4f}{ac:>14.4f}{pps:>28}{uq:>12}")226    print("-" * 84)227    print(f"注: 每层码本上限 perplexity = K = {K_CODE}; Gumbel 噪声的随机探索提高了")228    print("    码字访问覆盖, 对应正文所述'缓解码本坍塌、提高语义空间利用率'。")229230    # ---------- 5. FrqUD 的高/低频划分 ----------231    print()232    print("表 4  FrqUD 码字划分示例(β=0.9, r=1.5, γ=r/K={:.4f})".format(1.5 / K_CODE))233    print("-" * 84)234    *_, freq_f = run("dril", feat, seqs, tgt, gen, frqud=True)235    gamma = r_hot = 1.5 / K_CODE236    for l in range(L_LAYER):237        f = freq_f[l]238        hi = (f > gamma).nonzero().flatten().tolist()239        lo = [i for i in range(K_CODE) if i not in hi]240        print(f"  第 {l + 1} 层 EMA 频率 = [" + ", ".join(f"{x:.3f}" for x in f) + "]")241        print(f"         I_high(保留 Gumbel 探索) = {hi};  I_low(确定性 Softmax) = {lo}")242    print("-" * 84)243    print("正文: 对过度使用的热门码字保留探索, 对低频码字采用稳定的确定性分配。")244245246if __name__ == "__main__":247    main()