unisid_one_stage.py

1r"""对应正文第 9 章 §9.6.3「UniSID 算法」: 一阶段端到端语义 ID 生成。23正文要点与公式(本实现逐条复刻):4  * 输入序列 X_i = [x_instruction, x_image, x_text, x_attribute, x_sid, x_emb],5    经共享编码器 Z_i = MLLM(X_i), 取 SID/Emb 位置隐状态过两个独立投影头:6        z_i^SID = f_SID(Z_i^SID),  z_i^Emb = f_Emb(Z_i^Emb)7  * z_i^SID 切成 L 段, 每段直接取 argmax 维度作为 Token(不查码本、无最近邻):8        s_i^l = argmax(z_i^l),  s_i = (s_i^1, ..., s_i^L)9    因此语义 ID 只依赖维度间的大小顺序, 对数值缩放/平移具有鲁棒性。10  * 多粒度对比学习:11        L_sid = (1/L) Σ_l ( -(1/|P_l|) Σ_{p∈P_l} log12                 exp(sim(z_i^l, z_p^l)/τ) / Σ_{a∈A_l} exp(sim(z_i^l, z_a^l)/τ) )13    浅层正样本按粗粒度(大类)定义, 深层正样本按细粒度(子类)定义。14  * Embedding 分支标准对比学习:15        L_emb = -log exp(sim(z_i^Emb, z_j^Emb)/τ) / Σ_{k∈N_i} exp(sim(z_i^Emb, z_k^Emb)/τ)16  * Summary-based Reconstruction: 冻结 LLM 产出摘要 token, 由17        Z_i = [Z_i^SID; Z_i^Emb] -> h_i^rec = f_rec(Z_i) -> NTP 重建摘要18    对应 L_rec = Σ_t log p(s_{i,t}^sum | h_i^rec, s_{i,<t}^sum)19  * 总损失 L_total = L_sid + L_emb + λ L_rec2021本实现:22  1. [MOCK] 用一个小 Transformer 编码器代替 MLLM, 多模态输入用随机投影的23     "图像 / 文本 / 属性" 三组合成向量 + 可学习的 SID / Emb 占位 Token;24  2. [MOCK] 冻结 LLM 的"摘要"用确定性哈希规则从属性生成一小段 token 序列;25  3. 打印: argmax 取 SID 的鲁棒性验证表(缩放/平移/噪声)、层次语义结构统计26     (浅层与大类的一致性 vs 深层与子类的一致性)、以及三项损失的消融表。27"""2829from __future__ import annotations3031import random3233import torch34import torch.nn as nn35import torch.nn.functional as F3637L_LAYER = 3          # SID 层数38K_CODE = 8           # 每层维度数 = 候选 Token 数39D_MODEL = 3240N_COARSE = 4         # 粗粒度大类数41N_FINE = 3           # 每个大类下的细粒度子类数42SUM_VOCAB = 12       # [MOCK] 摘要 token 词表43SUM_LEN = 444N_ITEM = 180454647class UniSID(nn.Module):48    """[MOCK] 统一编码器: 同一个 Transformer 同时产出 SID 与 Embedding。"""4950    def __init__(self, d_in=20):51        super().__init__()52        self.proj = nn.ModuleDict({53            "image": nn.Linear(d_in, D_MODEL),54            "text": nn.Linear(d_in, D_MODEL),55            "attr": nn.Linear(d_in, D_MODEL),56        })57        # x_instruction / x_sid / x_emb 三个特殊占位 Token58        self.special = nn.Parameter(torch.randn(3, D_MODEL) * 0.1)59        layer = nn.TransformerEncoderLayer(D_MODEL, 4, 4 * D_MODEL, dropout=0.0,60                                           batch_first=True, norm_first=True,61                                           activation="gelu")62        self.enc = nn.TransformerEncoder(layer, 2, enable_nested_tensor=False)63        self.f_sid = nn.Linear(D_MODEL, L_LAYER * K_CODE)      # f_SID64        self.f_emb = nn.Linear(D_MODEL, D_MODEL)               # f_Emb65        self.f_rec = nn.Linear(2 * D_MODEL, D_MODEL)           # f_rec66        self.sum_emb = nn.Embedding(SUM_VOCAB + 1, D_MODEL)67        self.sum_out = nn.Linear(D_MODEL, SUM_VOCAB)6869    def forward(self, img, txt, att):70        b = img.shape[0]71        toks = [self.special[0].expand(b, -1),                  # x_instruction72                self.proj["image"](img),                        # x_image73                self.proj["text"](txt),                          # x_text74                self.proj["attr"](att),                          # x_attribute75                self.special[1].expand(b, -1),                   # x_sid76                self.special[2].expand(b, -1)]                   # x_emb77        z = self.enc(torch.stack(toks, 1))                       # Z_i = MLLM(X_i)78        z_sid = self.f_sid(z[:, 4])                              # z_i^SID79        z_emb = self.f_emb(z[:, 5])                              # z_i^Emb80        h_rec = self.f_rec(torch.cat([z[:, 4], z[:, 5]], -1))    # h_i^rec81        return z_sid.view(b, L_LAYER, K_CODE), z_emb, h_rec8283    def sid_of(self, z_sid):84        """s_i^l = argmax(z_i^l): 只看维度间大小顺序, 不查码本。"""85        return z_sid.argmax(-1)8687    def recon_loss(self, h_rec, summ):88        """[MOCK] 冻结 LLM 的 NTP 重建: p(s_t^sum | h_rec, s_<t^sum)。"""89        b, t = summ.shape90        inp = torch.cat([torch.full((b, 1), SUM_VOCAB, dtype=torch.long),91                         summ[:, :-1]], 1)92        h = self.sum_emb(inp) + h_rec[:, None]93        logits = self.sum_out(torch.tanh(h))94        return F.cross_entropy(logits.reshape(b * t, -1), summ.reshape(-1))959697def multi_granularity_contrastive(z_sid, coarse, fine, tau=0.2):98    """L_sid: 浅层正样本=同大类, 深层正样本=同子类(从粗到细的层次结构)。"""99    total = 0.0100    for l in range(L_LAYER):101        zl = F.normalize(z_sid[:, l], dim=-1)102        sim = zl @ zl.t() / tau103        if l == 0:104            pos = coarse[:, None] == coarse[None, :]              # 粗粒度105        elif l == L_LAYER - 1:106            pos = (coarse[:, None] == coarse[None, :]) & (fine[:, None] == fine[None, :])107        else:                                                      # 中间层: 过渡108            pos = (coarse[:, None] == coarse[None, :]) & \109                  ((fine[:, None] == fine[None, :]) | (torch.rand_like(sim) < 0.3))110        eye = torch.eye(len(zl), dtype=torch.bool)111        pos = pos & ~eye112        logden = torch.logsumexp(sim.masked_fill(eye, -1e9), dim=-1)113        num = (sim - logden[:, None]) * pos114        cnt = pos.sum(-1).clamp_min(1)115        total = total + (-(num.sum(-1) / cnt)).mean()116    return total / L_LAYER117118119def emb_contrastive(z_emb, view2, tau=0.2):120    """L_emb: 标准 InfoNCE, 正样本为同一物品的另一视图(属性 dropout 增强)。"""121    a = F.normalize(z_emb, dim=-1)122    b = F.normalize(view2, dim=-1)123    s = a @ b.t() / tau124    return F.cross_entropy(s, torch.arange(len(a)))125126127def make_summary(coarse, fine, gen):128    """[MOCK] 冻结 LLM 的摘要: 由大类/子类确定性映射到一段 token, 带少量噪声。"""129    table = torch.randint(0, SUM_VOCAB, (N_COARSE, N_FINE, SUM_LEN), generator=gen)130    return table[coarse, fine]131132133def nmi_like(tokens, labels, n_lab):134    """用"条件纯度"衡量 Token 与标签的一致性: Σ_t max_c p(c|t) * p(t)。"""135    pur = 0.0136    for t in tokens.unique():137        m = tokens == t138        cnt = torch.bincount(labels[m], minlength=n_lab).float()139        pur += float(cnt.max()) / len(tokens)140    return pur141142143def train(lam_rec=1.0, use_sid=True, use_emb=True, epochs=40, d_in=20, seed=42):144    torch.manual_seed(seed)145    gen = torch.Generator().manual_seed(seed)146    # ------- 合成多模态物料 -------147    c_cent = torch.randn(N_COARSE, d_in, generator=gen) * 2.0148    f_cent = torch.randn(N_COARSE, N_FINE, d_in, generator=gen) * 0.8149    coarse = torch.randint(0, N_COARSE, (N_ITEM,), generator=gen)150    fine = torch.randint(0, N_FINE, (N_ITEM,), generator=gen)151    base = c_cent[coarse] + f_cent[coarse, fine]152    img = base + torch.randn(N_ITEM, d_in, generator=gen) * 0.5153    txt = base + torch.randn(N_ITEM, d_in, generator=gen) * 0.5154    att = base + torch.randn(N_ITEM, d_in, generator=gen) * 0.5155    att2 = base + torch.randn(N_ITEM, d_in, generator=gen) * 0.5   # 增强视图156    summ = make_summary(coarse, fine, gen)157158    model = UniSID(d_in)159    opt = torch.optim.Adam(model.parameters(), lr=0.01)160    hist = []161    for ep in range(epochs):162        z_sid, z_emb, h_rec = model(img, txt, att)163        _, z_emb2, _ = model(img, txt, att2)164        ls = multi_granularity_contrastive(z_sid, coarse, fine) if use_sid \165            else torch.zeros(())166        le = emb_contrastive(z_emb, z_emb2) if use_emb else torch.zeros(())167        lr = model.recon_loss(h_rec, summ)168        loss = ls + le + lam_rec * lr                      # L_total169        opt.zero_grad()170        loss.backward()171        opt.step()172        if (ep + 1) % 10 == 0:173            hist.append((ep + 1, float(ls), float(le), float(lr), float(loss)))174    with torch.no_grad():175        z_sid, z_emb, _ = model(img, txt, att)176        sid = model.sid_of(z_sid)177    return model, (img, txt, att), sid, coarse, fine, hist178179180def main() -> None:181    random.seed(42)182    torch.manual_seed(42)183184    model, data, sid, coarse, fine, hist = train()185    img, txt, att = data186187    print("=" * 90)188    print(f"配置: 物品 {N_ITEM} 个, 大类 {N_COARSE} / 子类 {N_FINE}, "189          f"SID L={L_LAYER} K={K_CODE}, d_model={D_MODEL}")190    print()191    print("表 1  L_total = L_sid + L_emb + λ L_rec 的训练轨迹(λ=1.0)")192    print("-" * 90)193    print(f"{'epoch':>7}{'L_sid':>12}{'L_emb':>12}{'L_rec':>12}{'L_total':>12}")194    for ep, ls, le, lr, lo in hist:195        print(f"{ep:>7}{ls:>12.4f}{le:>12.4f}{lr:>12.4f}{lo:>12.4f}")196197    # ---------- argmax 取 SID 的鲁棒性 ----------198    print()199    print("表 2  s^l = argmax(z^l) 的鲁棒性(正文: 只关注维度大小顺序, 不关心绝对数值)")200    print("-" * 90)201    print(f"{'扰动方式':<34}{'SID 完全一致比例':>20}{'逐层 Token 一致率':>20}")202    with torch.no_grad():203        z0, _, _ = model(img, txt, att)204        s0 = model.sid_of(z0)205        cases = [("原始输出(基准)", lambda z: z),206                 ("整体乘 10(缩放)", lambda z: z * 10.0),207                 ("整体除 10(缩放)", lambda z: z / 10.0),208                 ("每层加常数 +5(平移)", lambda z: z + 5.0),209                 ("加 N(0, 0.05²) 噪声", lambda z: z + torch.randn_like(z) * 0.05),210                 ("加 N(0, 0.5²) 噪声", lambda z: z + torch.randn_like(z) * 0.5)]211        for name, fn in cases:212            s = model.sid_of(fn(z0.clone()))213            same_all = float((s == s0).all(-1).float().mean())214            same_tok = float((s == s0).float().mean())215            print(f"{name:<34}{same_all:>20.4f}{same_tok:>20.4f}")216    print("-" * 90)217    print("可见缩放/平移完全不改变 argmax 结果; 只有足够大的噪声才会翻转 Token。")218219    # ---------- 层次语义结构 ----------220    print()221    print("表 3  多粒度对比学习形成的层次结构(条件纯度: 越大说明该层 Token 越贴合该粒度)")222    print("-" * 90)223    print(f"{'SID 层':>8}{'对大类(粗粒度)纯度':>22}{'对(大类,子类)纯度':>22}"224          f"{'用到的 Token 数':>16}")225    joint = coarse * N_FINE + fine226    for l in range(L_LAYER):227        pc = nmi_like(sid[:, l], coarse, N_COARSE)228        pf = nmi_like(sid[:, l], joint, N_COARSE * N_FINE)229        print(f"{l + 1:>8}{pc:>22.4f}{pf:>22.4f}{len(sid[:, l].unique()):>16}")230    uniq = len({tuple(s.tolist()) for s in sid})231    print("-" * 90)232    print(f"不同 SID 数 = {uniq} / {N_ITEM} 个物品 (空间上限 K^L = {K_CODE ** L_LAYER})")233234    # ---------- 消融 ----------235    print()236    print("表 4  损失项消融(各配置独立训练 40 epoch)")237    print("-" * 90)238    print(f"{'配置':<34}{'第1层对大类纯度':>18}{'第3层对子类纯度':>18}{'不同 SID 数':>14}")239    for name, kw in (("L_sid + L_emb + L_rec(完整)", {}),240                     ("去掉 L_sid", {"use_sid": False}),241                     ("去掉 L_emb", {"use_emb": False}),242                     ("去掉 L_rec(λ=0)", {"lam_rec": 0.0})):243        _, _, s, c, f, _ = train(**kw)244        j = c * N_FINE + f245        print(f"{name:<34}{nmi_like(s[:, 0], c, N_COARSE):>18.4f}"246              f"{nmi_like(s[:, L_LAYER - 1], j, N_COARSE * N_FINE):>18.4f}"247              f"{len({tuple(x.tolist()) for x in s}):>14}")248    print("-" * 90)249    print("正文强调 L_sid 负责层次解耦、L_emb 约束 Embedding 分支、L_rec 迫使 SID")250    print("编码更抽象的高层语义; 此处为教学规模复现, 数值仅用于说明三者的作用方向。")251252253if __name__ == "__main__":254    main()