fscd.py

1"""2FSCD 粗排模型（第11章 §11.4.3）34FSCD = Feature Selection method based on feature Complexity and variational Dropout5论文出处：Ma et al., "Towards a Better Tradeoff between Effectiveness and Efficiency in6          Pre-Ranking: A Learnable Feature Selection based Approach", SIGIR 2021（阿里巴巴）。78核心思想（对应书中 §11.4.3）：9    FSCD 不重新设计粗排网络，而是直接继承精排的交互式 DNN 结构，只在「输入特征」这一层做裁剪，10    把特征选择变成一个可学习、且把在线算力成本写进目标函数的过程。11    1) 为每个特征域 j 引入伯努利门控   z_j ~ Bern(θ_j)12    2) 特征复杂度                      c_j = G(o_j, e_j, n_j) = γ1·o_j + γ2·e_j + γ3·n_j13       （o_j 在线计算复杂度、e_j Embedding 维度、n_j 特征取值数量）14    3) 复杂度决定先验保留概率           θ_j = 1 - σ(c_j)      → 越贵越不容易被保留15    4) 统一训练目标                     L = -1/N·logP(D|v,w,z) + λ/N·(‖v‖² + ‖w‖²) + 1/N·Σ_j α_j·z_j16       其中特征选择正则权重             α_j = log(1-θ_j) - log(θ_j)17    5) 伯努利不可导 → Gumbel-Sigmoid 连续松弛18       z_j = σ( ( logδ_j - log(1-δ_j) + log u_j - log(1-u_j) ) / t ),  u_j ~ U(0,1)19    6) 两阶段：第一阶段联合学 δ_j 与模型参数；收敛后按 δ_j 取 Top-K 特征，20       仅保留这些特征重新微调（目标退化为纯 CTR 损失）。2122PyTorch 实现（autograd 自动微分），可直接 `python3 fscd.py` 运行。23"""24import math25import random26import collections27from typing import Dict, List, Tuple2829import torch30import torch.nn as nn31import torch.nn.functional as F32import torch.optim as optim333435# ── 工具函数 ──────────────────────────────────────────────────────────────────3637def _sigmoid_py(x: float) -> float:38    return 1.0 / (1.0 + math.exp(-max(-30.0, min(30.0, x))))3940def _dot_py(a: List[float], b: List[float]) -> float:41    return sum(x * y for x, y in zip(a, b))4243def _auc(pairs: List[Tuple[float, int]]) -> float:44    pos = [s for s, y in pairs if y == 1]45    neg = [s for s, y in pairs if y == 0]46    if not pos or not neg:47        return 0.548    hit = sum(1.0 if p > n else (0.5 if p == n else 0.0) for p in pos for n in neg)49    return hit / (len(pos) * len(neg))505152# ── 特征域定义：(名称, 所属侧, o_j 在线计算复杂度, e_j Embedding 维度, n_j 取值数量) ──53FEATURE_FIELDS: List[Tuple[str, str, float, int, int]] = [54    ("u_profile", "user",  0.4, 4,     500),55    ("u_city",    "user",  0.4, 2,    3000),56    ("u_id",      "user",  1.0, 4,   50000),57    ("u_stat_7d", "user",  2.0, 4,  100000),58    ("i_cate",    "item",  0.5, 4,     800),59    ("i_price",   "item",  0.5, 2,     100),60    ("i_id",      "item",  1.0, 4,   20000),61    ("x_ui_seq",  "cross", 4.0, 6, 1000000),62]636465class FSCD(nn.Module):66    """67    FSCD：复杂度感知的可学习特征选择 + 交互式 DNN 粗排模型。68    网络结构与精排一致（所有特征域 Embedding 拼接 → MLP → CTR），69    唯一的区别是在 Embedding 与 DNN 之间插入一层可学习门控 z_j。7071    梯度：PyTorch autograd 自动微分，无手写梯度。72    """7374    def __init__(self, fields=FEATURE_FIELDS, hidden: int = 24, lr: float = 0.05,75                 reg: float = 1e-5, gamma=(0.5, 0.3, 0.3), temperature: float = 0.5,76                 sel_weight: float = 0.05):77        super().__init__()78        self.fields    = list(fields)79        self.names     = [f[0] for f in self.fields]80        self.dims      = [f[3] for f in self.fields]81        self.M         = len(self.fields)82        self.total_dim = sum(self.dims)83        self.t         = temperature84        self.sel_weight = sel_weight8586        # 特征复杂度 c_j、先验保留概率 θ_j、选择正则权重 α_j（Python 常量，不参与训练）87        g1, g2, g3 = gamma88        self.c     = [g1 * o + g2 * (e / 4.0) + g3 * (math.log10(n) / 2.0)89                      for _, _, o, e, n in self.fields]90        self.theta = [1.0 - _sigmoid_py(cj) for cj in self.c]91        self.alpha = [math.log(1 - th) - math.log(th) for th in self.theta]9293        # 可学习参数 δ_j（logit 形式）：nn.Parameter → autograd 自动追踪94        self.logit_delta = nn.Parameter(torch.zeros(self.M))9596        # 打分 DNN97        self.dnn = nn.Sequential(98            nn.Linear(self.total_dim, hidden),99            nn.ReLU(),100            nn.Linear(hidden, 1),101        )102103        # 第二阶段使用的硬门控（不是参数，只是 buffer）104        self.register_buffer("mask", torch.ones(self.M))105106        self.optimizer = optim.Adam(self.parameters(), lr=lr, weight_decay=reg)107108    def gumbel_sigmoid(self) -> torch.Tensor:109        """Gumbel-Sigmoid 连续松弛（训练阶段可微）：110           z_j = σ((logit δ_j + gumbel_noise) / t)111           gumbel = -log(-log(u))，u ~ U(0,1)。112        """113        u = torch.rand(self.M).clamp(1e-6, 1 - 1e-6)114        gumbel = -torch.log(-torch.log(u))            # Gumbel(0,1) 样本115        logit_u = torch.log(u) - torch.log(1 - u)    # logit(u)116        return torch.sigmoid((self.logit_delta + logit_u) / self.t)117118    def delta(self) -> torch.Tensor:119        """返回当前特征保留强度 δ_j = σ(logit_delta_j)。"""120        return torch.sigmoid(self.logit_delta)121122    def _apply_gate(self, feat: Dict[str, torch.Tensor], z: torch.Tensor) -> torch.Tensor:123        """逐特征域缩放后拼接：x = [z_1·v_1, ..., z_M·v_M]。"""124        parts = []125        for j, name in enumerate(self.names):126            parts.append(z[j] * feat[name])127        return torch.cat(parts)128129    def forward(self, feat: Dict[str, torch.Tensor], z: torch.Tensor) -> torch.Tensor:130        """返回 logit（未过 Sigmoid）。"""131        flat = self._apply_gate(feat, z)132        return self.dnn(flat)133134    # ── 第一阶段：联合优化模型参数与特征选择 ──135    def train_stage1(self, samples, epochs: int = 5) -> None:136        data = list(samples)137        alpha_t = torch.tensor(self.alpha, dtype=torch.float32)138        for epoch in range(epochs):139            random.shuffle(data)140            c_sum = s_sum = 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                z = self.gumbel_sigmoid()145                logit = self.forward(feat_t, z).squeeze()146                loss_ctr = F.binary_cross_entropy_with_logits(logit, torch.tensor(float(y)))147                # 特征选择正则：Σ_j α_j · z_j（α_j 为正则权重，复杂度越高越大）148                loss_sel = self.sel_weight * (alpha_t * z).sum()149                loss = loss_ctr + loss_sel150                loss.backward()151                self.optimizer.step()152                c_sum += loss_ctr.item()153                s_sum += loss_sel.item()154            n = max(1, len(data))155            print(f"  [stage1] epoch {epoch}  CTR_loss={c_sum / n:.4f}  "156                  f"选择正则={s_sum / n:.4f}  总损失={(c_sum + s_sum) / n:.4f}")157158    # ── 第二阶段：Top-K 特征微调 ──159    def select_top_k(self, k: int) -> List[str]:160        """按 δ_j 降序选 Top-K，更新硬门控 mask。"""161        with torch.no_grad():162            d = self.delta()163            order = torch.argsort(d, descending=True)[:k].tolist()164        self.mask = torch.zeros(self.M)165        for j in order:166            self.mask[j] = 1.0167        return [self.names[j] for j in sorted(order)]168169    def set_mask(self, names: List[str]) -> None:170        keep = set(names)171        self.mask = torch.tensor(172            [1.0 if n in keep else 0.0 for n in self.names])173174    def train_stage2(self, samples, epochs: int = 3) -> None:175        """仅保留选中特征，目标退化为纯 CTR 损失（以第一阶段参数作为初始化）。"""176        data = list(samples)177        for epoch in range(epochs):178            random.shuffle(data)179            total = 0.0180            for feat, y in data:181                self.optimizer.zero_grad()182                feat_t = {k: torch.tensor(v, dtype=torch.float32) for k, v in feat.items()}183                logit = self.forward(feat_t, self.mask).squeeze()184                loss = F.binary_cross_entropy_with_logits(logit, torch.tensor(float(y)))185                loss.backward()186                self.optimizer.step()187                total += loss.item()188            print(f"  [stage2] epoch {epoch}  CTR_loss={total / max(1, len(data)):.4f}")189190    def online_cost(self) -> float:191        """被保留特征的在线计算复杂度之和。"""192        return sum(self.fields[j][2] for j in range(self.M) if self.mask[j].item() > 0)193194    # ── 推理接口 ──195    def score(self, feat: Dict[str, List[float]]) -> float:196        with torch.no_grad():197            feat_t = {k: torch.tensor(v, dtype=torch.float32) for k, v in feat.items()}198            return float(torch.sigmoid(self.forward(feat_t, self.mask)).item())199200    def pre_rank(self, uid: int, candidates: List[int], top_k: int = 10) -> List[Tuple[int, float]]:201        """粗排：只用筛选后的特征对候选集打分并返回 Top-K。"""202        scored = [(iid, self.score(FEATURE_STORE[(uid, iid)])) for iid in candidates]203        scored.sort(key=lambda x: x[1], reverse=True)204        return scored[:top_k]205206207FEATURE_STORE: Dict[Tuple[int, int], Dict[str, List[float]]] = {}208209210def make_data(n_users: int, n_items: int):211    """212    合成特征：u_profile / i_cate 廉价且携带完整兴趣信号；u_id / i_id / u_stat_7d 信号被噪声稀释213    且词表大、复杂度高；x_ui_seq 信号最强但在线复杂度最高（FSCD 重点裁剪对象）；214    u_city / i_price 为纯噪声但复杂度极低。215    """216    lat_u = [[random.gauss(0, 1) for _ in range(4)] for _ in range(n_users)]217    lat_i = [[random.gauss(0, 1) for _ in range(4)] for _ in range(n_items)]218    samples = []219    for u in range(n_users):220        for i in range(n_items):221            s = _dot_py(lat_u[u], lat_i[i]) / 2.0222            y = 1 if random.random() < _sigmoid_py(s) else 0223            feat = {224                "u_profile": list(lat_u[u]),225                "u_city":    [random.gauss(0, 1) for _ in range(2)],226                "u_id":      [x + random.gauss(0, 0.6) for x in lat_u[u]],227                "u_stat_7d": [x + random.gauss(0, 1.2) for x in lat_u[u]],228                "i_cate":    list(lat_i[i]),229                "i_price":   [random.gauss(0, 1) for _ in range(2)],230                "i_id":      [x + random.gauss(0, 0.6) for x in lat_i[i]],231                "x_ui_seq":  [s + random.gauss(0, 0.2)] * 3 + [s, s, s],232            }233            FEATURE_STORE[(u, i)] = feat234            samples.append((feat, y))235    return samples236237238if __name__ == "__main__":239    random.seed(42)240    torch.manual_seed(42)241    n_users, n_items = 8, 20242    samples = make_data(n_users, n_items)243    model = FSCD(lr=0.05, temperature=0.5)244    print(f"FSCD 粗排：特征域数 M={model.M}  拼接维度={model.total_dim}  样本数={len(samples)}")245246    print("\n[1] 特征复杂度与选择正则（复杂度越高 → θ 越小 → α 越大 → 越容易被裁掉）")247    print(f"  {'特征域':<12}{'o_j':>6}{'e_j':>5}{'n_j':>10}{'c_j':>8}{'θ_j':>8}{'α_j':>8}")248    for j, (name, _, o, e, n) in enumerate(model.fields):249        print(f"  {name:<12}{o:>6.1f}{e:>5d}{n:>10d}{model.c[j]:>8.3f}"250              f"{model.theta[j]:>8.3f}{model.alpha[j]:>8.3f}")251252    print("\n[2] 第一阶段：Gumbel-Sigmoid 变分 Dropout，联合学习 δ_j 与 DNN")253    model.train_stage1(samples, epochs=5)254    with torch.no_grad():255        deltas = model.delta().tolist()256    print("  学到的特征保留强度 δ_j（降序）：")257    for j in sorted(range(model.M), key=lambda j: deltas[j], reverse=True):258        print(f"    {model.names[j]:<12}δ={deltas[j]:.4f}")259260    full_auc = _auc([(model.score({k: v for k, v in f.items()}) if False261                      else float(torch.sigmoid(model.dnn(262                          torch.cat([torch.tensor(f[n], dtype=torch.float32)263                                     for n in model.names]))).item()), y)264                     for f, y in samples])265    # 全特征（mask=1）的 AUC266    with torch.no_grad():267        orig_mask = model.mask.clone()268        model.mask = torch.ones(model.M)269        full_auc = _auc([(model.score(f), y) for f, y in samples])270        model.mask = orig_mask271272    chosen = model.select_top_k(4)273    fscd_auc  = _auc([(model.score(f), y) for f, y in samples])274    fscd_cost = model.online_cost()275276    # 对照组：人工按「复杂度最高」挑 4 个特征277    heavy = sorted(model.names, key=lambda n: model.c[model.names.index(n)], reverse=True)[:4]278    model.set_mask(heavy)279    heavy_auc  = _auc([(model.score(f), y) for f, y in samples])280    heavy_cost = model.online_cost()281    model.set_mask(chosen)282283    total_cost = sum(f[2] for f in model.fields)284    print(f"\n[3] 效果-成本权衡（全量在线复杂度={total_cost:.1f}）")285    print(f"  全特征            AUC={full_auc:.4f}  在线复杂度=9.8")286    print(f"  FSCD Top-4        AUC={fscd_auc:.4f}  在线复杂度={fscd_cost:.1f}   {chosen}")287    print(f"  高复杂度 Top-4    AUC={heavy_auc:.4f}  在线复杂度={heavy_cost:.1f}   {heavy}")288289    print("\n[4] 第二阶段：仅用筛选后的特征微调（目标退化为纯 CTR 损失）")290    model.train_stage2(samples, epochs=3)291    final_auc = _auc([(model.score(f), y) for f, y in samples])292    print(f"  微调后 FSCD Top-4 AUC：{final_auc:.4f}"293          f"（≈ 全特征效果，但在线复杂度只有 {fscd_cost / 9.8:.0%}）")294295    uid = 0296    print(f"\n[5] 用户 {uid} 粗排 Top-5（仅使用 {len(chosen)} 个被保留的特征域）")297    for iid, s in model.pre_rank(uid, list(range(n_items)), top_k=5):298        print(f"  item {iid:2d}: {s:.4f}")