dlrm.py

1"""2DLRM 精排模型（第12章 §12.4.1.5）34Deep Learning Recommendation Model（Meta, 2019）：56    e_dense  = bottom_mlp(x_dense)          # 稠密特征 → 同维嵌入7    e_i      = Embedding(feat_i)             # 类别特征 → 嵌入8    s_{ij}   = e_i · e_j                     # 两两内积（显式二阶交互）9    z        = [e_dense; s_{01},…,s_{m-1,m}]10    ŷ        = σ(top_mlp(z))1112特征交互模块仅计算内积，不构造高阶交互（区别于 CIN/xDeepFM）。13顶部 MLP 在二阶交互结果和稠密嵌入的基础上学习高阶非线性。1415本书依据：《第十二章-精排模型.pdf》§12.4.1.5，PDF 第 30–36 页。16依赖：PyTorch。运行：python3 dlrm.py17"""1819import random20import torch21import torch.nn as nn22import torch.nn.functional as F23import torch.optim as optim242526def _auc(scores, labels) -> float:27    pos = [s for s, y in zip(scores, labels) if y == 1]28    neg = [s for s, y in zip(scores, labels) if y == 0]29    if not pos or not neg:30        return 0.531    return sum(1.0 if p > n else 0.5 if p == n else 0.032               for p in pos for n in neg) / (len(pos) * len(neg))333435def _make_mlp(dims: list[int], act=nn.ReLU) -> nn.Sequential:36    layers = []37    for i in range(len(dims) - 1):38        layers.append(nn.Linear(dims[i], dims[i + 1]))39        if i < len(dims) - 2:40            layers.append(act())41    return nn.Sequential(*layers)424344# ── DLRM 模型 ─────────────────────────────────────────────────────────────────4546class DLRM(nn.Module):47    """48    DLRM：底部 MLP + 类别嵌入 → 两两内积交互 → 顶部 MLP。4950    正文（PDF 第 30–36 页）关键设计：51      - bottom_mlp 把稠密数值特征映射到与嵌入相同的 emb_dim 维度；52      - 内积交互层计算所有 C(m+1,2) 对特征向量的内积，输出为一维向量 s；53      - top_mlp 的输入为 [e_dense; s]（保留稠密特征表示 + 交互特征）；54      - 与 DeepFM 将 FM 和 DNN 直接融合不同，DLRM 将交互结果作为 top_mlp 输入。55    """5657    def __init__(self, n_sparse: int, vocab_size: int, emb_dim: int = 8,58                 dense_dim: int = 4,59                 bottom_dims: list[int] | None = None,60                 top_dims: list[int] | None = None):61        """62        参数：63          n_sparse   : 类别（稀疏）字段数64          vocab_size : 每个字段的取值数（统一 vocab，演示简化）65          emb_dim    : 嵌入维度（bottom_mlp 输出维度也是此值）66          dense_dim  : 稠密数值特征维度67          bottom_dims: bottom_mlp 隐藏层维度，最后一层必须 == emb_dim68          top_dims   : top_mlp 各层维度，最后为 1（输出 logit）69        """70        super().__init__()71        self.emb_dim = emb_dim72        self.n_sparse = n_sparse7374        # 类别特征嵌入（每个字段独立嵌入表，真实系统中 vocab 各不相同）75        self.embs = nn.ModuleList([76            nn.Embedding(vocab_size, emb_dim) for _ in range(n_sparse)])7778        # 底部 MLP：稠密特征 → emb_dim79        b_dims = bottom_dims or [dense_dim, emb_dim * 2, emb_dim]80        assert b_dims[-1] == emb_dim, "bottom_mlp 最后一层维度必须等于 emb_dim"81        self.bottom_mlp = _make_mlp(b_dims, nn.ReLU)8283        # 两两内积后的维度：C(n_sparse+1, 2) 个标量84        n_vecs = n_sparse + 1                       # +1 for e_dense85        n_interact = n_vecs * (n_vecs - 1) // 28687        # 顶部 MLP：[e_dense(emb_dim); s(n_interact)] → 188        t_dims = top_dims or [emb_dim + n_interact, 32, 1]89        self.top_mlp = _make_mlp(t_dims, nn.ReLU)9091        for emb in self.embs:92            nn.init.normal_(emb.weight, std=0.01)9394    def _interact(self, vecs: torch.Tensor) -> torch.Tensor:95        """96        vecs : (B, m+1, emb_dim)97        返回  : (B, C(m+1,2))  两两内积98        """99        n = vecs.size(1)100        pairs = [(i, j) for i in range(n) for j in range(i + 1, n)]101        dots = [102            (vecs[:, i, :] * vecs[:, j, :]).sum(dim=1, keepdim=True)103            for i, j in pairs104        ]105        return torch.cat(dots, dim=1)                # (B, n_interact)106107    def forward(self, feat_ids: torch.Tensor,108                dense: torch.Tensor) -> torch.Tensor:109        """110        feat_ids : (B, n_sparse) — 类别特征索引111        dense    : (B, dense_dim) — 稠密数值特征112        """113        # 稠密特征 → 嵌入空间114        e_dense = self.bottom_mlp(dense)             # (B, emb_dim)115116        # 类别特征嵌入117        e_sparse = [emb(feat_ids[:, i])118                    for i, emb in enumerate(self.embs)]  # n_sparse × (B, emb_dim)119120        # 拼成向量组: (B, m+1, emb_dim)121        vecs = torch.stack([e_dense] + e_sparse, dim=1)122123        # 两两内积124        s = self._interact(vecs)                     # (B, n_interact)125126        # top_mlp 输入：保留稠密嵌入 + 交互向量127        z = torch.cat([e_dense, s], dim=1)128        return self.top_mlp(z).squeeze(-1)129130    def predict_proba(self, feat_ids, dense) -> torch.Tensor:131        with torch.no_grad():132            return torch.sigmoid(self.forward(feat_ids, dense))133134    def fit(self, feat_ids, dense, y, epochs: int = 40, lr: float = 1e-2):135        opt = optim.Adam(self.parameters(), lr=lr, weight_decay=1e-5)136        for ep in range(epochs):137            opt.zero_grad()138            loss = F.binary_cross_entropy_with_logits(139                self.forward(feat_ids, dense), y.float())140            loss.backward()141            opt.step()142            if (ep + 1) % 10 == 0:143                print(f"  epoch {ep+1:3d}  loss={loss.item():.4f}")144        return loss.item()145146    def rank(self, feat_ids, dense, top_k):147        scores = self.predict_proba(feat_ids, dense).tolist()148        return sorted(enumerate(scores), key=lambda t: -t[1])[:top_k]149150151# ── 合成数据 ──────────────────────────────────────────────────────────────────152153def make_data(n: int = 500, n_sparse: int = 4, vocab: int = 8,154              dense_dim: int = 4, seed: int = 42):155    """156    n_sparse 个类别字段（如 用户ID桶 / 物品类别 / 设备 / 地域），157    dense_dim 个数值字段（如 历史CTR / 时长统计 / 价格 / 评分）。158    标签依赖类别字段的二阶交互 + 数值主效应，这正是 DLRM 内积层的159    建模目标：e_i · e_j 捕捉类别对之间的二阶关系。160    """161    torch.manual_seed(seed)162    random.seed(seed)163    ids = torch.randint(0, vocab, (n, n_sparse), dtype=torch.long)164    dense = torch.randn(n, dense_dim)165166    local = ids.float()167    # 类别主效应 + 0×1 二阶交互 + 数值主效应168    logit = (0.3 * local[:, 0] + 0.4 * local[:, 1]169             + 0.9 * (local[:, 0] == local[:, 1]).float()170             + 0.7 * (local[:, 2] == local[:, 3]).float()171             + 0.5 * dense[:, 0] - 0.4 * dense[:, 1] - 2.5)172    y = (torch.sigmoid(logit + 0.3 * torch.randn(n)) > 0.5).long()173    return ids, dense, y174175176# ── 主演示 ────────────────────────────────────────────────────────────────────177178if __name__ == "__main__":179    random.seed(42)180    torch.manual_seed(42)181182    N_SPARSE, VOCAB, DENSE_DIM = 4, 8, 4183    ids, dense, y = make_data(500, N_SPARSE, VOCAB, DENSE_DIM)184    split = 380185    ids_tr, d_tr, y_tr = ids[:split], dense[:split], y[:split]186    ids_te, d_te, y_te = ids[split:], dense[split:], y[split:]187188    print("=== DLRM 训练（底部 MLP + 类别嵌入 + 两两内积 + 顶部 MLP）===")189    model = DLRM(n_sparse=N_SPARSE, vocab_size=VOCAB, emb_dim=8,190                 dense_dim=DENSE_DIM,191                 bottom_dims=[DENSE_DIM, 16, 8],192                 top_dims=[8 + 10, 32, 1])   # 10 = C(5,2) 对内积193    model.fit(ids_tr, d_tr, y_tr, epochs=40)194195    p_te = model.predict_proba(ids_te, d_te).tolist()196    auc_te = _auc(p_te, y_te.tolist())197    print(f"\n测试集 AUC={auc_te:.4f}")198199    # ── 交互维度验证 ──────────────────────────────────────────────────────────200    n_vecs = N_SPARSE + 1201    n_interact = n_vecs * (n_vecs - 1) // 2202    print(f"\n特征向量数 m+1={n_vecs}  两两内积维度 C({n_vecs},2)={n_interact}")203    print(f"顶部 MLP 输入维度: emb_dim({model.emb_dim}) + n_interact({n_interact})"204          f" = {model.emb_dim + n_interact}")205    print("  正文（PDF 第 32 页）：顶部 MLP 保留稠密特征表示 e_dense，")206    print("  避免预测过程完全依赖内积交互结果。")207208    # ── 与 DeepFM 架构对比 ────────────────────────────────────────────────────209    print(f"\n与 DeepFM 的关键区别（正文 PDF 第 35 页）：")210    print("  DeepFM: FM 二阶项 + DNN 并联，共享嵌入，直接融合两分支得分")211    print("  DLRM : 内积 → 拼接为 top_mlp 输入，而非与 DNN 直接相加")212    print("  DLRM 更侧重控制交互粒度（每对特征压缩为 1 个标量），适合工业大规模部署。")213214    # ── 精排推理演示 ──────────────────────────────────────────────────────────215    top5 = model.rank(ids_te[:20], d_te[:20], top_k=5)216    print("\n精排 Top-5（候选集 20 条）:")217    for rank_i, (idx, s) in enumerate(top5, 1):218        print(f"  #{rank_i}  item={idx:2d}  score={s:.4f}  label={y_te[idx].item()}")