can.py

1"""2CAN：Co-Action Network，由特征生成参数（第15章 §15.5.7）34Co-Action Unit 把一对特征分成两个角色：5    Induction Side（参数生成侧）：P_induction 被 reshape 成 Micro-MLP 的 W 与 b6    Feed Side（输入侧）        ：P_feed 作为 Micro-MLP 的输入7        h_0 = P_feed8        h_i = σ(w_{i-1} h_{i-1} + b_{i-1})9        F(u, m) = [h_1 ‖ h_2 ‖ … ‖ h_L]1011Multi-order Enhancement：把 (P_feed)^c（c=1..C）分别送进同一个 Micro-MLP 再求和，12用 tanh 约束高阶项的数值范围（正文 §15.5.7）。1314参数量对比是本节的重点：15    笛卡尔积交互 Embedding：O(N²D)16    CAN：O(N(D' + D))17本脚本用一个 400×40 的 induction 参数表，与 400×200 的 pair 表直接对照。1819运行：python3 ch04_feature_interaction/can.py20"""2122from __future__ import annotations2324import pathlib25import sys2627import torch28import torch.nn as nn2930sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent))3132from _train_utils import (  # noqa: E40233    FIELD_NAMES, FIELD_SIZES, N_FIELDS, MLP, BaselineMLP, count_params, embedding_table,34    print_comparison, run_experiment, stack_embeddings,35)3637DIM = 438INDUCTION_LAYERS = 2394041class CoActionUnit(nn.Module):42    """43    一个 Micro-MLP：每层都是 d×d 的 Weight 加 d 维 Bias，参数由 Induction 侧生成。44    multi_order=True 时，Feed 侧额外提供 P² 作为第二阶输入。45    """4647    def __init__(self, dim: int, n_layers: int = INDUCTION_LAYERS, multi_order: bool = False):48        super().__init__()49        self.dim = dim50        self.n_layers = n_layers51        self.multi_order = multi_order52        self.param_per_layer = dim * dim + dim53        self.param_dim = n_layers * self.param_per_layer54        self.n_orders = 2 if multi_order else 15556    def forward(self, p_induction: torch.Tensor, p_feed: torch.Tensor) -> torch.Tensor:57        # p_induction: [B, param_dim]，p_feed: [B, d]58        B = p_feed.shape[0]59        d = self.dim60        outs = []61        for c in range(1, self.n_orders + 1):62            # (P_feed)^c，用 tanh 做有界约束（§15.5.7 Multi-order Enhancement）63            h = torch.tanh(p_feed) if c == 1 else torch.tanh(p_feed ** c)64            layer_outs = []65            offset = 066            for _ in range(self.n_layers):67                w = p_induction[:, offset:offset + d * d].view(B, d, d)68                offset += d * d69                b = p_induction[:, offset:offset + d].view(B, d)70                offset += d71                h = torch.tanh(torch.bmm(h.unsqueeze(1), w).squeeze(1) + b)72                layer_outs.append(h)              # 每层输出都要保留73            outs.append(torch.cat(layer_outs, dim=1))   # Concat(h_1..h_L)74        return torch.stack(outs, dim=0).sum(0)   # Σ_c H(P_ind, (P_feed)^c)757677class CAN(nn.Module):78    """79    Induction 侧：用户 ID（词表大、组合多，适合用参数生成）80    Feed 侧    ：物品 ID81    正文 §15.5.7 提到可以反过来选，取决于两侧规模；这里把选择写在构造参数里。82    """8384    def __init__(self, dim: int = DIM, multi_order: bool = False,85                 induction_field: str = "user", feed_field: str = "item"):86        super().__init__()87        self.emb = embedding_table(dim)88        self.induction_field = induction_field89        self.feed_field = feed_field90        self.unit = CoActionUnit(dim, INDUCTION_LAYERS, multi_order)91        # 只为 induction 侧的特征建一张"参数表"，而不是为每个 pair 建 Embedding92        self.induction_table = nn.Embedding(FIELD_SIZES[induction_field], self.unit.param_dim)93        # 初始化偏大一些：权重太小会让 Micro-MLP 输出接近常数、梯度消失（实测94        # std=0.02 时 AUC 掉约 0.006，std=0.5 才与 Baseline 持平）。95        nn.init.normal_(self.induction_table.weight, std=0.5)96        inter_dim = self.unit.n_layers * dim97        self.mlp = MLP(dim * N_FIELDS + inter_dim, (32,), out_dim=1, dropout=0.2)9899    def forward(self, x):100        E = stack_embeddings(self.emb, x)                    # [B, m, d]101        i_idx = FIELD_NAMES.index(self.induction_field)102        f_idx = FIELD_NAMES.index(self.feed_field)103        p_ind = self.induction_table(x[:, i_idx])            # Parameter Lookup104        p_feed = E[:, f_idx]105        inter = self.unit(p_ind, p_feed)106        return self.mlp(torch.cat([E.flatten(1), inter], dim=1)).squeeze(-1)107108109def main() -> None:110    print("CAN：Induction Side 生成 Micro-MLP 参数，Feed Side 作为输入（§15.5.7）\n")111    m = N_FIELDS112    d = DIM113    n_items = FIELD_SIZES["item"]114    n_users = FIELD_SIZES["user"]115116    print("参数量对照（用户侧 × 物品侧）：")117    print(f"  笛卡尔积交互 Embedding : {n_users} × {n_items} × {d} = "118          f"{n_users * n_items * d:,}")119    print(f"  CAN（induction 参数表）: {n_users} × {CoActionUnit(d).param_dim} = "120          f"{n_users * CoActionUnit(d).param_dim:,}"121          f"（相当于前者的 1/{n_items}）")122    print(f"  另一侧物品仍走常规 Embedding：{n_items} × {d} = {n_items * d:,}\n")123124    _, a_plain, _ = run_experiment(lambda: CAN(multi_order=False), "CAN", steps=600)125    _, a_multi, _ = run_experiment(lambda: CAN(multi_order=True), "CAN+Multi-order", steps=600)126    _, a_base, _ = run_experiment(BaselineMLP, "Baseline", steps=600)127128    print_comparison([129        ("Baseline MLP", a_base, count_params(BaselineMLP())),130        ("CAN", a_plain, count_params(CAN())),131        ("CAN + Multi-order", a_multi, count_params(CAN(multi_order=True))),132    ])133    print("\nCAN 的核心不是'多加一层网络'，而是把过去固定的交互函数换成由一侧特征")134    print("现场生成的 Micro-MLP：不同 user 对应不同的交互函数（§15.5.7）。")135    print("micro-MLP 只有 1 层、无 bias、无激活时，它退化为内积，也就是 FM。")136137138if __name__ == "__main__":139    main()