dnn_vs_two_tower.py

1"""2精排 DNN 主模型 vs 粗排双塔模型（第 6 章 §6.2 精排模型与 LTR 模型）34对应正文概念：5正文指出，粗排用双塔（用户塔/物品塔分别编码，最后内积），精排则把6「用户侧特征 + 物品侧特征 + 上下文特征统一编码并拼接（Concatenate）后7输入 DNN」，因而能在网络内部自动学习高阶特征交叉；精排主模型采用8多任务学习（MTL）框架联合预估 click / effective_view / like / comment /9follow / share 等多个行为目标。正文同时强调：两者差异源于「效果、算力、10时延之间的工程权衡」，而非建模能力的本质差异。1112本实现做了什么：131. 用 torch 构造一份带「用户年龄 × 内容类目」交叉效应的合成数据——这是正文14   举的例子，双塔的后期内积结构很难拟合，而拼接式 DNN 可以。152. 同时训练双塔模型与拼接式 DNN 多任务模型，对比 AUC 与参数量/推理 FLOPs，16   量化「精排为什么更准、为什么更贵」。173. DNN 侧用共享底座 + 每任务独立塔头的 MTL 结构，输出 6 个 pxtr。1819关键公式：20    双塔：    score = <f_user(x_u), f_item(x_i)>21    精排DNN： h = MLP([x_u ; x_i ; x_c]),  p_k = sigmoid(w_k^T h + b_k)22    MTL 损失：L = sum_k lambda_k * BCE(p_k, y_k)23"""2425import math26import random2728import torch29import torch.nn as nn3031NUM_AGE_BUCKET = 32  # 用户侧离散特征（年龄分桶 × 人群标签的组合桶）32NUM_CATEGORY = 40    # 物品侧离散特征（一级类目 × 二级类目）33NUM_CTX = 4          # 上下文：时段34INTERACT_DIM = 8     # 双塔内积维度 = 粗排的在线算力预算，受工程约束35TASKS = ["click", "effective_view", "like", "comment", "follow", "share"]36# 正文 §6.3.2 给出的稀疏度量级：CTR 3%~10%，LTR <1%，CMTR <1‰，FTR <1‱37TASK_BASE_RATE = [0.08, 0.40, 0.009, 0.0009, 0.0004, 0.002]383940# --------------------------------------------------------------------------41# 合成数据：核心是「用户年龄 × 内容类目」的高阶交叉42# --------------------------------------------------------------------------434445def _calibrate_bias(s: torch.Tensor, target: float) -> float:46    """二分搜索偏置，使 mean(sigmoid(bias + s)) ≈ target（对齐正文给出的稀疏度）。"""47    lo, hi = -30.0, 30.048    for _ in range(60):49        mid = (lo + hi) / 250        if torch.sigmoid(mid + s).mean().item() < target:51            lo = mid52        else:53            hi = mid54    return (lo + hi) / 2555657def make_dataset(n: int, gen: torch.Generator):58    age = torch.randint(0, NUM_AGE_BUCKET, (n,), generator=gen)59    cat = torch.randint(0, NUM_CATEGORY, (n,), generator=gen)60    ctx = torch.randint(0, NUM_CTX, (n,), generator=gen)61    u_dense = torch.rand(n, 4, generator=gen)62    i_dense = torch.rand(n, 4, generator=gen)6364    # (1) 满秩的 (年龄段 × 类目) 交叉矩阵。双塔最终只能用 emb 维内积逼近它，65    #     即只能表达秩 <= EMB_DIM 的双线性形式；当交叉矩阵秩远高于 EMB_DIM 时，66    #     双塔结构性地欠拟合。67    g2 = torch.Generator().manual_seed(7)68    cross = torch.randn(NUM_AGE_BUCKET, NUM_CATEGORY, generator=g2) * 1.569    age_main = torch.randn(NUM_AGE_BUCKET, generator=g2) * 0.5   # 可加主效应70    cat_main = torch.randn(NUM_CATEGORY, generator=g2) * 0.57172    # (2) 稠密特征上的非双线性交互：L1 距离 |u - i| 无法写成任何有限维的73    #     <f(u), g(i)> 形式，因此双塔的「后期内积」结构天然学不到。74    dense_cross = -2.2 * (u_dense - i_dense).abs().sum(1)7576    logit_core = (77        cross[age, cat] + age_main[age] + cat_main[cat]78        + dense_cross + 0.25 * ctx.float()79    )80    labels = torch.zeros(n, len(TASKS))81    for k, base in enumerate(TASK_BASE_RATE):82        s = (0.9 + 0.15 * k) * logit_core83        p = torch.sigmoid(_calibrate_bias(s, base) + s)84        labels[:, k] = torch.bernoulli(p, generator=gen)85    # oracle 不参与训练，只用于报告「理论 AUC 上限」86    return {"age": age, "cat": cat, "ctx": ctx, "oracle": logit_core,87            "u_dense": u_dense, "i_dense": i_dense, "y": labels}888990# --------------------------------------------------------------------------91# 模型92# --------------------------------------------------------------------------939495class TwoTower(nn.Module):96    """粗排双塔：两侧独立编码，最后仅靠 INTERACT_DIM 维内积交互。9798    注意内积维度 INTERACT_DIM 是粗排的工程预算（决定向量检索/打分开销），99    不能随意放大。而 score 对 (age, cat) 的依赖只能写成100        s(age, cat) = <f(age, ...), g(cat, ...)>101    这是一个秩 <= INTERACT_DIM 的双线性型；当真值交叉矩阵秩更高时，102    双塔就是结构性欠拟合，加宽隐层也救不回来。103    """104105    def __init__(self, emb=16, hid=128, n_tasks=len(TASKS)):106        super().__init__()107        self.age_emb = nn.Embedding(NUM_AGE_BUCKET, emb)108        self.cat_emb = nn.Embedding(NUM_CATEGORY, emb)109        self.ctx_emb = nn.Embedding(NUM_CTX, emb)110        self.user_tower = nn.Sequential(111            nn.Linear(emb * 2 + 4, hid), nn.ReLU(),112            nn.Linear(hid, hid), nn.ReLU(), nn.Linear(hid, INTERACT_DIM))113        self.item_tower = nn.Sequential(114            nn.Linear(emb + 4, hid), nn.ReLU(),115            nn.Linear(hid, hid), nn.ReLU(), nn.Linear(hid, INTERACT_DIM))116        # 每个任务一组缩放/偏置，让双塔也有多任务输出能力117        self.scale = nn.Parameter(torch.ones(n_tasks))118        self.bias = nn.Parameter(torch.zeros(n_tasks))119120    def forward(self, b):121        u = self.user_tower(torch.cat(122            [self.age_emb(b["age"]), self.ctx_emb(b["ctx"]), b["u_dense"]], 1))123        i = self.item_tower(torch.cat([self.cat_emb(b["cat"]), b["i_dense"]], 1))124        dot = (u * i).sum(1, keepdim=True)125        return dot * self.scale + self.bias126127128class RankingDNN(nn.Module):129    """精排 DNN：全部特征拼接 → 共享底座 → 每任务独立塔头（MTL）。"""130131    def __init__(self, emb=16, hid=(128, 64), n_tasks=len(TASKS)):132        super().__init__()133        self.age_emb = nn.Embedding(NUM_AGE_BUCKET, emb)134        self.cat_emb = nn.Embedding(NUM_CATEGORY, emb)135        self.ctx_emb = nn.Embedding(NUM_CTX, emb)136        in_dim = emb * 3 + 8137        self.shared = nn.Sequential(138            nn.Linear(in_dim, hid[0]), nn.ReLU(),139            nn.Linear(hid[0], hid[1]), nn.ReLU())140        self.heads = nn.ModuleList(141            [nn.Sequential(nn.Linear(hid[1], 32), nn.ReLU(), nn.Linear(32, 1))142             for _ in range(n_tasks)])143144    def forward(self, b):145        x = torch.cat([self.age_emb(b["age"]), self.cat_emb(b["cat"]),146                       self.ctx_emb(b["ctx"]), b["u_dense"], b["i_dense"]], 1)147        h = self.shared(x)148        return torch.cat([head(h) for head in self.heads], 1)149150151# --------------------------------------------------------------------------152# 训练 / 评估153# --------------------------------------------------------------------------154155156def auc(scores: torch.Tensor, labels: torch.Tensor) -> float:157    """按秩计算 AUC；全正或全负时返回 nan。"""158    pos = labels.sum().item()159    neg = labels.numel() - pos160    if pos == 0 or neg == 0:161        return float("nan")162    order = torch.argsort(scores)163    ranks = torch.empty_like(order, dtype=torch.float)164    ranks[order] = torch.arange(1, len(scores) + 1, dtype=torch.float)165    return ((ranks[labels == 1].sum().item() - pos * (pos + 1) / 2)166            / (pos * neg))167168169def train(model, data, epochs=12, bs=512, lr=3e-3):170    opt = torch.optim.Adam(model.parameters(), lr=lr)171    lossf = nn.BCEWithLogitsLoss()172    n = data["y"].shape[0]173    for _ in range(epochs):174        perm = torch.randperm(n)175        for s in range(0, n, bs):176            idx = perm[s:s + bs]177            batch = {k: v[idx] for k, v in data.items()}178            logits = model(batch)179            loss = sum(lossf(logits[:, k], batch["y"][:, k])180                       for k in range(len(TASKS)))181            opt.zero_grad()182            loss.backward()183            opt.step()184    return model185186187def n_params(m):188    return sum(p.numel() for p in m.parameters())189190191if __name__ == "__main__":192    random.seed(42)193    torch.manual_seed(42)194    gen = torch.Generator().manual_seed(42)195196    train_data = make_dataset(40000, gen)197    test_data = make_dataset(10000, gen)198    print("=" * 72)199    print("精排 DNN vs 粗排双塔：拟合「用户年龄 × 内容类目」高阶交叉")200    print("=" * 72)201    print(f"训练集 {train_data['y'].shape[0]} 条，测试集 {test_data['y'].shape[0]} 条")202    print("各任务正样本率（测试集）：")203    for k, t in enumerate(TASKS):204        r = test_data["y"][:, k].mean().item()205        print(f"  {t:<16}{r * 100:7.3f}%   （设定基准 {TASK_BASE_RATE[k] * 100:.3f}%）")206207    results = {}208    for name, model in (("TwoTower(双塔)", TwoTower()),209                        ("RankingDNN(精排)", RankingDNN())):210        train(model, train_data)211        model.eval()212        with torch.no_grad():213            logits = model(test_data)214        results[name] = [auc(logits[:, k], test_data["y"][:, k])215                         for k in range(len(TASKS))]216        print(f"\n{name}  参数量 = {n_params(model):,}")217218    print("\n" + "=" * 78)219    print("测试集 AUC 对比（Oracle = 用真值 logit 排序的理论上限）")220    print("=" * 78)221    names = list(results)222    print(f"{'任务':<16}" + "".join(f"{k:>18}" for k in names)223          + f"{'ΔAUC':>9}{'Oracle':>9}")224    for k, t in enumerate(TASKS):225        a, b = results[names[0]][k], results[names[1]][k]226        o = auc(test_data["oracle"], test_data["y"][:, k])227        print(f"{t:<16}{a:>18.4f}{b:>18.4f}{b - a:>+9.4f}{o:>9.4f}")228    ma = sum(results[names[0]]) / len(TASKS)229    mb = sum(results[names[1]]) / len(TASKS)230    mo = sum(auc(test_data["oracle"], test_data["y"][:, k])231             for k in range(len(TASKS))) / len(TASKS)232    print(f"{'平均':<16}{ma:>18.4f}{mb:>18.4f}{mb - ma:>+9.4f}{mo:>9.4f}")233    print(f"\n距 Oracle 上限的差距：双塔 {mo - ma:.4f}，精排 DNN {mo - mb:.4f}"234          f"（精排收回了 {(mb - ma) / (mo - ma) * 100:.1f}% 的可弥补差距）")235236    # 在线单候选算力：双塔的物品塔可离线预计算、用户塔每次请求只算一次，237    # 因此每多一个候选只需 INTERACT_DIM 次乘加；精排 DNN 每个候选都要跑全网络。238    dnn = RankingDNN()239    dnn_flops = sum(m.in_features * m.out_features240                    for m in dnn.modules() if isinstance(m, nn.Linear))241    print(f"\n在线单候选打分乘加次数：双塔 ≈ {INTERACT_DIM}（物品向量离线预算好，"242          f"只剩内积）\n                        精排 DNN ≈ {dnn_flops:,}"243          f"（每个候选跑一遍全网络，约 {dnn_flops // INTERACT_DIM} 倍）")244245    print("\n结论：双塔把两侧特征压成 INTERACT_DIM="246          f"{INTERACT_DIM} 维后才内积，只能表达秩 ≤ {INTERACT_DIM} 的双线性"247          "交叉，\n      且对 |u-i| 这类非双线性关系无能为力；拼接式 DNN 在输入层"248          "就看到两侧特征，\n      AUC 更高，代价是每候选在线算力高出三个数量级——"249          "这正是正文所说的\n      「效果 / 算力 / 时延」工程权衡：粗排候选上千只能用双塔，"250          "精排候选几百才用得起 DNN。")