mvke.py

1"""2MVKE 粗排多兴趣模型（第11章 §11.5.1）34MVKE = Mixture of Virtual-Kernel Experts（虚拟核专家混合模型）5论文出处：Xu et al., "Mixture of Virtual-Kernel Experts for Multi-Objective User Profile6          Modeling", KDD 2022（腾讯）。78核心思想（对应书中 §11.5.1）：9    传统双塔把用户多样化兴趣压成一个固定 Embedding，表达能力不足。MVKE 保持双塔的高效10    架构，但把用户塔换成 K 个虚拟核专家（VKE），再用虚拟核门控（VKG）按目标物品动态聚合：11    1) 第 k 个 VKE 维护一个可学习虚拟核 W_VK^k（隐式的「兴趣查询」）。虚拟核经非线性映射12       得到 Query，用户字段 Embedding 映射为 Key / Value：13           Q = σ(W_Qᵀ·W_VK^k + b_Q)，K = σ(W_Kᵀ·E_uf + b_K)，V = σ(W_Vᵀ·E_uf + b_V)14       注意力聚合：C_VKE^k = softmax(Q·Kᵀ / √d_K)·V15       兴趣表示：    E_ui^k = f_u^k(C_VKE^k)16       所有 VKE 共享同一份用户字段 Embedding，但虚拟核不同 → 关注兴趣的不同侧面。17    2) VKG 以目标物品 Embedding E_Ti 为 Query，以虚拟核 {W_VK^k} 为 Key，以 VKE 输出18       {E_ui^k} 为 Value：19           E_ui = Σ_k softmax( Q(E_Ti)·K(W_VK^k)ᵀ / √d_K ) · V(E_ui^k)20       同一用户面对不同目标物品时，门控权重不同 → 动态多兴趣表征。21    3) 在线效率：K 组 VKE 兴趣表示可离线预计算并存储，线上只需算 K 个门控权重再加权求和，22       相比双塔只增加存储开销，不显著增加计算复杂度。2324PyTorch 实现（autograd 自动微分），可直接 `python3 mvke.py` 运行。25"""26import math27import random28import collections29from typing import List, Tuple3031import torch32import torch.nn as nn33import torch.nn.functional as F34import torch.optim as optim353637# ── 工具函数 ──────────────────────────────────────────────────────────────────3839def _dot_py(a, b):40    return sum(x * y for x, y in zip(a, b))4142def _sigmoid_py(x):43    return 1.0 / (1.0 + math.exp(-max(-30.0, min(30.0, x))))4445def _auc(pairs: List[Tuple[float, int]]) -> float:46    pos, neg = [s for s, y in pairs if y == 1], [s for s, y in pairs if y == 0]47    if not pos or not neg:48        return 0.549    hit = sum(1.0 if p > n else (0.5 if p == n else 0.0) for p in pos for n in neg)50    return hit / (len(pos) * len(neg))515253class VKE(nn.Module):54    """虚拟核专家：以可学习虚拟核 W_VK^k 为 Query，对用户各字段 Embedding 做注意力聚合。5556    梯度：PyTorch autograd 自动微分，无手写梯度。57    """5859    def __init__(self, n_fields: int, field_dim: int, d_vk: int = 8, d_k: int = 8,60                 d_v: int = 8, out_dim: int = 8, init_offset: float = 0.0):61        super().__init__()62        self.m = n_fields63        self.d_k = d_k64        self.scale = math.sqrt(d_k)6566        # 可学习虚拟核 W_VK^k（一维向量）；init_offset 使不同专家初始化错开，打破对称性67        w = torch.empty(d_vk).uniform_(-1/math.sqrt(d_vk), 1/math.sqrt(d_vk))68        w[0] = w[0].abs() * (1.0 if init_offset >= 0 else -1.0) + init_offset69        self.w_vk = nn.Parameter(w)7071        # Q = ReLU(W_Q · W_VK + b_Q)72        self.q_proj = nn.Linear(d_vk, d_k)73        # K = ReLU(W_K · e + b_K)，V = ReLU(W_V · e + b_V)  （逐字段共享权重）74        self.k_proj = nn.Linear(field_dim, d_k)75        self.v_proj = nn.Linear(field_dim, d_v)76        # 兴趣建模网络 f_u^k：线性输出（保持正负极性）77        self.f_u = nn.Linear(d_v, out_dim)7879    def forward(self, fields: torch.Tensor):80        """81        fields: (n_fields, field_dim)82        Returns (E_ui_k: (out_dim,), attn: (n_fields,))83        """84        Q = torch.tanh(self.q_proj(self.w_vk))                    # (d_k,)85        K = F.relu(self.k_proj(fields))                            # (n_fields, d_k)86        V = F.relu(self.v_proj(fields))                            # (n_fields, d_v)87        logits = (K @ Q) / self.scale                              # (n_fields,)88        a = torch.softmax(logits, dim=0)                           # (n_fields,)89        C = (a.unsqueeze(1) * V).sum(dim=0)                        # (d_v,)90        E = self.f_u(C)                                            # (out_dim,)91        return E, a929394class VKG(nn.Module):95    """虚拟核门控：目标物品 Embedding 作 Query，虚拟核作 Key，VKE 兴趣表示作 Value。9697    梯度：PyTorch autograd 自动微分，无手写梯度。98    """99100    def __init__(self, vkes: nn.ModuleList, item_dim: int, out_dim: int, d_k: int = 8):101        super().__init__()102        self.vkes = vkes103        self.K = len(vkes)104        self.d_k = d_k105        self.scale = math.sqrt(d_k)106        # Q(E_Ti): ReLU 激活，保持与 VKE 中 Q 网络一致的风格107        self.q_proj = nn.Linear(item_dim, d_k)108        # K(W_VK^k): 线性投影（无激活），使门控 logits 可正可负，不同虚拟核对比度更高109        d_vk = vkes[0].w_vk.shape[0]110        self.k_proj = nn.Linear(d_vk, d_k)111        # V(E_ui^k): 线性投影112        self.v_proj = nn.Linear(out_dim, out_dim)113114    def forward(self, item_e: torch.Tensor, interests: List[torch.Tensor]):115        """116        item_e: (item_dim,)117        interests: list of K tensors, each (out_dim,)118        Returns (E_ui: (out_dim,), gate: (K,))119        """120        Q = torch.tanh(self.q_proj(item_e))                        # (d_k,)121        # Stack virtual kernel keys122        vk_stack = torch.stack([vke.w_vk for vke in self.vkes])    # (K, d_vk)123        Ks = self.k_proj(vk_stack)                                 # (K, d_k)124        logits = (Ks @ Q) / self.scale                             # (K,)125        gate = torch.softmax(logits, dim=0)                        # (K,)126        int_stack = torch.stack(interests)                         # (K, out_dim)127        Vs = self.v_proj(int_stack)                                # (K, out_dim)128        E_ui = (gate.unsqueeze(1) * Vs).sum(dim=0)                 # (out_dim,)129        return E_ui, gate130131132class MVKE(nn.Module):133    """MVKE 粗排模型：K 个 VKE 构成用户塔 + VKG 目标相关门控 + 轻量目标塔。134135    梯度：PyTorch autograd 自动微分，无手写梯度。136    """137138    def __init__(self, n_users: int, n_items: int, n_u_fields: int = 4, n_i_fields: int = 2,139                 field_dim: int = 4, n_vke: int = 3, d_vk: int = 8, d_k: int = 8,140                 out_dim: int = 8, lr: float = 0.02, reg: float = 1e-5, emb_scale: float = 0.8):141        super().__init__()142        self.n_items = n_items143        self.mu = n_u_fields144        self.mi = n_i_fields145        self.fd = field_dim146        self.K = n_vke147        self.out_dim = out_dim148149        self.user_emb = nn.Parameter(150            torch.empty(n_users, n_u_fields, field_dim).uniform_(-emb_scale, emb_scale))151        self.item_emb = nn.Parameter(152            torch.empty(n_items, n_i_fields, field_dim).uniform_(-emb_scale, emb_scale))153154        # init_offset staggers virtual kernels so experts break symmetry from step 0155        offsets = [k * 1.0 - (n_vke - 1) / 2.0 for k in range(n_vke)]156        self.vkes = nn.ModuleList([157            VKE(n_u_fields, field_dim, d_vk, d_k, out_dim, out_dim, init_offset=offsets[k])158            for k in range(n_vke)159        ])160        # 目标塔：两层 MLP，输出层线性（使 E_Ti 可正可负，内积可区分负样本）161        self.item_l1 = nn.Linear(n_i_fields * field_dim, out_dim)162        self.item_l2 = nn.Linear(out_dim, out_dim)163164        self.vkg = VKG(self.vkes, out_dim, out_dim, d_k)165        self.bias = nn.Parameter(torch.zeros(1))166167        self.optimizer = optim.Adam(self.parameters(), lr=lr, weight_decay=reg)168169    def forward(self, uid: int, iid: int):170        """Returns (logit scalar, gate (K,), attn list of (n_u_fields,))"""171        fields = self.user_emb[uid]                                # (n_u_fields, field_dim)172        interests, attns = [], []173        for vke in self.vkes:174            E, a = vke(fields)175            interests.append(E)176            attns.append(a)177178        item_flat = self.item_emb[iid].reshape(-1)                 # (n_i_fields * field_dim,)179        h = F.relu(self.item_l1(item_flat))180        E_T = self.item_l2(h)                                      # (out_dim,) — linear output181182        E_ui, gate = self.vkg(E_T, interests)183        logit = torch.dot(E_ui, E_T) / math.sqrt(self.out_dim) + self.bias.squeeze()184        return logit, gate, attns185186    def score(self, uid: int, iid: int) -> float:187        with torch.no_grad():188            logit, _, _ = self.forward(uid, iid)189            return float(torch.sigmoid(logit).item())190191    def gate_of(self, uid: int, iid: int) -> List[float]:192        with torch.no_grad():193            _, gate, _ = self.forward(uid, iid)194            return gate.tolist()195196    def field_attention(self, uid: int) -> List[List[float]]:197        """各 VKE 对用户字段的注意力分布（体现不同虚拟核关注兴趣的不同侧面）。"""198        with torch.no_grad():199            fields = self.user_emb[uid]200            return [vke(fields)[1].tolist() for vke in self.vkes]201202    def train_model(self, pos_pairs, epochs: int = 6) -> None:203        data, pos_set = list(pos_pairs), set(pos_pairs)204        for epoch in range(epochs):205            random.shuffle(data)206            total = 0.0207            for uid, iid in data:208                self.optimizer.zero_grad()209                # Positive sample210                logit_pos, _, _ = self.forward(uid, iid)211                loss = F.binary_cross_entropy_with_logits(logit_pos, torch.tensor(1.0))212                # Negative sample213                nid = random.randrange(self.n_items)214                while (uid, nid) in pos_set:215                    nid = random.randrange(self.n_items)216                logit_neg, _, _ = self.forward(uid, nid)217                loss = loss + F.binary_cross_entropy_with_logits(logit_neg, torch.tensor(0.0))218                loss.backward()219                self.optimizer.step()220                total += loss.item()221            print(f"  epoch {epoch}  logloss={total / max(1, len(data) * 2):.4f}")222223    def pre_rank(self, uid: int, candidates: List[int], top_k: int = 10) -> List[Tuple[int, float]]:224        """粗排：K 组兴趣表示可预计算，线上逐候选只需算门控权重 + 加权求和 + 一次内积。"""225        scored = [(iid, self.score(uid, iid)) for iid in candidates]226        scored.sort(key=lambda x: x[1], reverse=True)227        return scored[:top_k]228229230if __name__ == "__main__":231    random.seed(42)232    torch.manual_seed(42)233    n_users, n_items, n_vke = 10, 30, 3234235    # 合成多兴趣数据：每个用户有 3 个兴趣分量（与 K=3 簇对应），用户偏好分量 k 的物品簇 k。236    # 这样不同簇的物品需要激活不同专家，驱动 VKG 产生簇相关的门控分布。237    n_clusters = 3238    field_dim = 4239    cluster = [i % n_clusters for i in range(n_items)]240241    # 用户在每个簇方向上有独立的兴趣强度（4维）；物品完全由其簇方向决定242    interest_dirs = [[random.gauss(0, 1) for _ in range(field_dim)] for _ in range(n_clusters)]243    user_modes = [[[random.gauss(0, 1) for _ in range(field_dim)]244                   for _ in range(n_clusters)] for _ in range(n_users)]245    item_vec = [interest_dirs[cluster[i]] for i in range(n_items)]246247    pos_pairs = []248    pos_set = set()249    for u in range(n_users):250        for i in range(n_items):251            c = cluster[i]252            s = _dot_py(user_modes[u][c], item_vec[i]) / field_dim253            if s > 0.2:254                pos_pairs.append((u, i))255                pos_set.add((u, i))256257    model = MVKE(n_users, n_items, n_u_fields=4, n_i_fields=2, n_vke=n_vke)258    print(f"MVKE 粗排：n_users={n_users}  n_items={n_items}  VKE 数 K={n_vke}  "259          f"正样本对={len(pos_pairs)}")260261    print("\n[1] 训练（用户塔 = K 个虚拟核专家，门控由目标物品决定）")262    model.train_model(pos_pairs, epochs=12)263    auc = _auc([(model.score(u, i), 1 if (u, i) in pos_set else 0)264                for u in range(n_users) for i in range(n_items)])265    print(f"  训练集 AUC = {auc:.4f}")266267    print("\n[2] 各 VKE 内部的字段注意力 softmax(Q·Kᵀ/√d_K)：不同虚拟核关注不同字段")268    for u in [0, 3]:269        for k, a in enumerate(model.field_attention(u)):270            print(f"  用户{u} VKE{k}: " + "  ".join(f"字段{i}={v:.3f}" for i, v in enumerate(a)))271272    print("\n[3] VKG 门控权重分布（Query=目标物品 Embedding，故门控由目标侧驱动）")273    sample_items = [i * (n_items // (n_clusters * 2)) for i in range(n_clusters * 2)][:6]274    for i in sample_items:275        g = model.gate_of(0, i)276        print(f"  目标 item{i:2d}（簇{cluster[i]}）：门控=[" + ", ".join(f"{v:.3f}" for v in g)277              + f"]  最偏好 VKE{g.index(max(g))}")278    spread = [max(model.gate_of(0, i)) - min(model.gate_of(0, i)) for i in range(n_items)]279    print(f"  全部 {n_items} 个目标物品上的门控极差均值={sum(spread) / len(spread):.3f}"280          f"（>0 说明门控确实在按目标区分专家）")281282    uid = 0283    cands = [i for i in range(n_items) if (uid, i) not in pos_set]284    print(f"\n[4] 用户 {uid} 粗排 Top-5（候选 {len(cands)} 个）")285    for iid, s in model.pre_rank(uid, cands, top_k=5):286        print(f"  item {iid:2d}（簇{cluster[iid]}）: {s:.4f}  门控="287              f"{[round(v, 3) for v in model.gate_of(uid, iid)]}")