cold.py

1"""2COLD 粗排模型（第11章 §11.4.2）34COLD = Computing power cost-aware Online and Lightweight Deep pre-ranking system5论文出处：Wang et al., "COLD: Towards the Next Generation of Pre-Ranking System", 阿里巴巴 2020。67核心思想（对应书中 §11.4.2）：8    (1) 粗排不再被「必须用向量内积」的双塔结构绑死，而是允许直接使用表达能力更强的9        全连接 DNN（用户特征、物品特征、交叉特征统一 Embedding 后拼接 → MLP → CTR）；10    (2) 用「算力代价感知」的手段把复杂 DNN 的在线开销压下来，核心是基于 SE11        （Squeeze-and-Excitation）模块的特征组选择（Feature Group Selection）：12            squeeze    : z_i = mean(e_i)                        （每个特征组压成一个标量统计量）13            excitation : s   = σ(W2·ReLU(W1·z + b1) + b2)       （书中写作 s = σ(W[e_1..e_M] + b)）14            re-weight  : v_i = s_i · e_i                        （逐特征组缩放）15        随后按重要性 s_i 挑 Top-K 特征组，并结合 GAUC / QPS / RT 不断调整 K，16        在「模型效果」和「算力代价」之间找平衡点。1718本文件在书中公式基础上补充了一张显式的「特征组算力代价表」，并实现了预算约束下的19贪心特征组选择（按 收益/代价 比贪心），用于直观演示 COLD 的算力-效果权衡。2021PyTorch 实现（autograd 自动微分），可直接 `python3 cold.py` 运行。22"""23import math24import random25import collections26from typing import Dict, List, Tuple2728import torch29import torch.nn as nn30import torch.nn.functional as F31import torch.optim as optim323334# ── 工具函数 ──────────────────────────────────────────────────────────────────3536def _sigmoid_py(x: float) -> float:37    return 1.0 / (1.0 + math.exp(-max(-30.0, min(30.0, x))))3839def _dot_py(a: List[float], b: List[float]) -> float:40    return sum(x * y for x, y in zip(a, b))4142def _auc(pairs: List[Tuple[float, int]]) -> float:43    pos = [s for s, y in pairs if y == 1]44    neg = [s for s, y in pairs if y == 0]45    if not pos or not neg:46        return 0.547    hit = sum(1.0 if p > n else (0.5 if p == n else 0.0) for p in pos for n in neg)48    return hit / (len(pos) * len(neg))495051# ── 特征组算力代价表 ─────────────────────────────────────────────────────────52# (名称, 所属侧, Embedding 维度, 单次在线计算代价)53# 代价的量纲可理解为「相对 RT 开销」：用户静态画像类最便宜（可缓存），54# 物品/用户实时行为序列类最贵（需要对每个候选物品实时拉取并计算）。55FEATURE_GROUPS: List[Tuple[str, str, int, float]] = [56    ("u_id",           "user", 4, 0.5),57    ("u_profile",      "user", 4, 0.6),58    ("u_stat_7d",      "user", 4, 1.8),59    ("u_realtime_seq", "user", 4, 4.5),60    ("i_id",           "item", 4, 0.6),61    ("i_cate",         "item", 4, 0.8),62    ("i_stat_7d",      "item", 4, 1.7),63    ("i_realtime_ctr", "item", 4, 4.0),64]656667class COLD(nn.Module):68    """69    COLD 粗排模型：SE 特征组选择 + 全连接 DNN 打分。7071    与双塔模型的差别：用户/物品/交叉特征在底层就拼接进同一个 MLP，72    因此能够学习任意形式的非线性交互；代价是无法像双塔那样预计算物品向量，73    必须靠特征组裁剪 + 工程优化（列式计算、Float16、多级并行）把 RT 压回毫秒级。7475    梯度：PyTorch autograd 自动微分，无手写梯度。76    """7778    def __init__(self, groups=FEATURE_GROUPS, hidden: int = 24, lr: float = 0.05, reg: float = 1e-5):79        super().__init__()80        self.groups = list(groups)81        self.names  = [g[0] for g in self.groups]82        self.dims   = [g[2] for g in self.groups]83        self.costs  = {g[0]: g[3] for g in self.groups}84        self.M      = len(self.groups)85        self.total_dim = sum(self.dims)86        self.active = set(self.names)8788        # SE excitation 网络：M → M//2 → M，输出过 Sigmoid 得到各组重要性 s89        self.se = nn.Sequential(90            nn.Linear(self.M, max(2, self.M // 2)),91            nn.ReLU(),92            nn.Linear(max(2, self.M // 2), self.M),93        )94        # 打分 DNN：拼接后的全部特征 → hidden → 195        self.dnn = nn.Sequential(96            nn.Linear(self.total_dim, hidden),97            nn.ReLU(),98            nn.Linear(hidden, 1),99        )100101        self.gate_sum = [0.0] * self.M102        self.gate_cnt = 0103104        self.optimizer = optim.Adam(self.parameters(), lr=lr, weight_decay=reg)105106    def _squeeze(self, fields: List[torch.Tensor]) -> torch.Tensor:107        """Squeeze：对每个特征组做全局平均池化 z_i = (1/d)·Σ_t e_{i,t}。"""108        return torch.stack([f.mean() for f in fields])109110    def se_forward(self, fields: List[torch.Tensor]):111        """SE 全流程：squeeze → excitation → re-weight（v_i = s_i·e_i）。"""112        z = self._squeeze(fields)                 # shape: (M,)113        s = torch.sigmoid(self.se(z))             # shape: (M,)114        v = [s[i] * fields[i] for i in range(self.M)]115        return v, s116117    def forward(self, feat: Dict[str, torch.Tensor]) -> torch.Tensor:118        """返回 logit（未过 Sigmoid）。"""119        fields = []120        for name, _, dim, _ in self.groups:121            if name in self.active:122                fields.append(feat[name])123            else:124                fields.append(torch.zeros(dim))125        v, s = self.se_forward(fields)126        # 更新门控统计（detach 避免影响梯度图）127        with torch.no_grad():128            for i in range(self.M):129                self.gate_sum[i] += float(s[i].detach())130            self.gate_cnt += 1131        flat = torch.cat(v)132        return self.dnn(flat)133134    def train_model(self, samples, epochs: int = 5) -> None:135        data = list(samples)136        for epoch in range(epochs):137            random.shuffle(data)138            self.gate_sum = [0.0] * self.M139            self.gate_cnt = 0140            total = 0.0141            for feat, y in data:142                self.optimizer.zero_grad()143                feat_t = {k: torch.tensor(v, dtype=torch.float32) for k, v in feat.items()}144                logit = self.forward(feat_t).squeeze()145                loss = F.binary_cross_entropy_with_logits(logit, torch.tensor(float(y)))146                loss.backward()147                self.optimizer.step()148                total += loss.item()149            print(f"  epoch {epoch}  logloss={total / max(1, len(data)):.4f}")150151    # ── 算力代价感知的特征组选择 ──152    def importance(self) -> Dict[str, float]:153        """特征组重要性 = 训练过程中 SE 门控权重 s_i 的均值。"""154        c = max(1, self.gate_cnt)155        return {self.names[i]: self.gate_sum[i] / c for i in range(self.M)}156157    def select_by_budget(self, budget: float) -> List[str]:158        """159        预算约束下的贪心特征组选择：160            按「收益密度 = 重要性 / 算力代价」从高到低贪心加入，直到耗尽代价预算。161        对应书中「按权重大小选 Top-K，并结合 GAUC/QPS/RT 调整 K」的工程做法。162        """163        imp = self.importance()164        order = sorted(self.names, key=lambda n: imp[n] / self.costs[n], reverse=True)165        chosen, used = [], 0.0166        for name in order:167            if used + self.costs[name] <= budget:168                chosen.append(name)169                used += self.costs[name]170        return chosen171172    def apply_selection(self, names: List[str]) -> None:173        self.active = set(names)174175    def cost_of(self, names=None) -> float:176        names = self.names if names is None else names177        return sum(self.costs[n] for n in names)178179    def score(self, feat: Dict[str, List[float]]) -> float:180        with torch.no_grad():181            feat_t = {k: torch.tensor(v, dtype=torch.float32) for k, v in feat.items()}182            return float(torch.sigmoid(self.forward(feat_t)).item())183184    def pre_rank(self, uid: int, candidates: List[int], top_k: int = 10) -> List[Tuple[int, float]]:185        """粗排：对候选集逐个跑 DNN 打分并返回 Top-K（仅计算激活的特征组）。"""186        scored = [(iid, self.score(FEATURE_STORE[(uid, iid)])) for iid in candidates]187        scored.sort(key=lambda x: x[1], reverse=True)188        return scored[:top_k]189190191FEATURE_STORE: Dict[Tuple[int, int], Dict[str, List[float]]] = {}192193194def make_data(n_users: int, n_items: int, dim: int = 4):195    """合成特征：u_id / i_id 承载真实兴趣信号，实时类特征信号更强但贵，画像/统计类掺大量噪声。"""196    lat_u = [[random.gauss(0, 1) for _ in range(dim)] for _ in range(n_users)]197    lat_i = [[random.gauss(0, 1) for _ in range(dim)] for _ in range(n_items)]198    samples = []199200    def nz(v, sd):201        return [x + random.gauss(0, sd) for x in v]202203    for u in range(n_users):204        for i in range(n_items):205            s = _dot_py(lat_u[u], lat_i[i]) / 2.0206            y = 1 if random.random() < _sigmoid_py(s) else 0207            feat = {208                "u_id":           list(lat_u[u]),209                "u_profile":      [random.gauss(0, 1) for _ in range(dim)],210                "u_stat_7d":      nz(lat_u[u], 0.9),211                "u_realtime_seq": nz(lat_u[u], 0.2),212                "i_id":           list(lat_i[i]),213                "i_cate":         [1.0 if x > 0 else -1.0 for x in lat_i[i]],214                "i_stat_7d":      [random.gauss(0, 1) for _ in range(dim)],215                "i_realtime_ctr": nz(lat_i[i], 0.2),216            }217            FEATURE_STORE[(u, i)] = feat218            samples.append((feat, y))219    return samples220221222if __name__ == "__main__":223    random.seed(42)224    torch.manual_seed(42)225    n_users, n_items = 8, 20226    samples = make_data(n_users, n_items)227    model = COLD(lr=0.05)228    print(f"COLD 粗排：特征组数 M={model.M}  拼接维度={model.total_dim}  "229          f"全量算力代价={model.cost_of():.1f}  样本数={len(samples)}")230231    print("\n[1] 全特征组训练（SE 模块同时学习各特征组重要性）")232    model.train_model(samples, epochs=5)233234    print("\n[2] SE 门控给出的特征组重要性 / 算力代价")235    imp = model.importance()236    print(f"  {'特征组':<16}{'重要性 s_i':>12}{'代价':>8}{'收益密度':>12}")237    for n in sorted(model.names, key=lambda n: imp[n] / model.costs[n], reverse=True):238        print(f"  {n:<16}{imp[n]:>12.4f}{model.costs[n]:>8.1f}{imp[n] / model.costs[n]:>12.4f}")239240    print("\n[3] 不同算力预算下的特征组选择与粗排效果")241    uid = 0242    cands = list(range(n_items))243    for budget in [2.0, 5.0, 14.5]:244        chosen = model.select_by_budget(budget)245        model.apply_selection(chosen)246        auc = _auc([(model.score(f), y) for f, y in samples])247        top = model.pre_rank(uid, cands, top_k=5)248        print(f"\n  预算={budget:>5.1f}  实际代价={model.cost_of(chosen):>4.1f}  AUC={auc:.4f}")249        print(f"    选中特征组：{chosen}")250        print(f"    用户 {uid} 粗排 Top-5：" + "  ".join(f"item{i}:{s:.3f}" for i, s in top))