fibinet.py

1"""2FiBiNet：SENET 特征重加权 + Bilinear 二阶交互（第15章 §15.5.4）34SENET：5    z_i = (1/k) Σ_t e_i^(t)                 # squeeze：每个 field 压成一个标量6    A = σ(W_2 σ(W_1 Z))                    # excitation：field 权重7    V = [a_i e_i]                          # re-weight8Bilinear（三种共享方式）：9    p_ij = (e_i^T W) ⊙ e_j10    - Field-All：所有 pair 共享 W11    - Field-Each：每个 field 一个 W_i12    - Field-Interaction：每个 pair 一个 W_ij（参数量最大）13两条通路 E→P、E→V→Q 拼接后过 DNN。1415运行：python3 ch04_feature_interaction/fibinet.py16"""1718from __future__ import annotations1920import itertools21import pathlib22import sys2324import torch25import torch.nn as nn2627sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent))2829from _train_utils import (  # noqa: E40230    FIELD_NAMES, N_FIELDS, MLP, BaselineMLP, count_params, embedding_table, print_comparison,31    run_experiment, stack_embeddings,32)333435class SENET(nn.Module):36    def __init__(self, n_fields: int, reduction: int = 2):37        super().__init__()38        self.w1 = nn.Linear(n_fields, max(1, n_fields // reduction))39        self.w2 = nn.Linear(max(1, n_fields // reduction), n_fields)4041    def forward(self, E: torch.Tensor) -> torch.Tensor:42        # E: [B, F, d]43        z = E.mean(dim=2)                       # squeeze -> [B, F]44        a = torch.sigmoid(self.w2(torch.relu(self.w1(z))))  # excitation45        return E * a.unsqueeze(-1)              # re-weight464748class BilinearInteraction(nn.Module):49    """mode: all / each / interaction"""5051    def __init__(self, n_fields: int, dim: int, mode: str = "all"):52        super().__init__()53        self.mode = mode54        if mode == "all":55            self.w = nn.Parameter(torch.empty(dim, dim))56            nn.init.xavier_uniform_(self.w)57        elif mode == "each":58            self.w = nn.Parameter(torch.empty(n_fields, dim, dim))59            nn.init.xavier_uniform_(self.w)60        elif mode == "interaction":61            self.w = nn.Parameter(torch.empty(n_fields, n_fields, dim, dim))62            nn.init.xavier_uniform_(self.w)63        else:64            raise ValueError(mode)65        self.register_buffer("pairs", torch.tensor(66            list(itertools.combinations(range(n_fields), 2)), dtype=torch.long))6768    def forward(self, E: torch.Tensor) -> torch.Tensor:69        # E: [B, F, d] -> 所有 i<j 的 (e_i^T W) ⊙ e_j70        B, F, d = E.shape71        pi, pj = self.pairs[:, 0], self.pairs[:, 1]72        ei, ej = E[:, pi], E[:, pj]              # [B, P, d]73        if self.mode == "all":74            x = ei @ self.w                     # [B, P, d]75        elif self.mode == "each":76            # 每对使用 e_i 自己的 W_i77            x = torch.einsum("bpd,pde->bpe", ei, self.w[pi])78        else:79            x = torch.einsum("bpd,pqde->bpe", ei, self.w[pi, pj])80        return x * ej                           # Hadamard818283class FiBiNet(nn.Module):84    def __init__(self, dim: int = 4, mode: str = "all", hidden=(32,)):85        super().__init__()86        self.emb = embedding_table(dim)87        self.senet = SENET(N_FIELDS)88        self.bi = BilinearInteraction(N_FIELDS, dim, mode)   # 两条通路共享交互参数89        n_pairs = N_FIELDS * (N_FIELDS - 1) // 290        self.mlp = MLP(2 * n_pairs * dim, hidden, out_dim=1, dropout=0.2)9192    def forward(self, x):93        E = stack_embeddings(self.emb, x)     # [B, F, d]94        V = self.senet(E)95        P = self.bi(E).flatten(1)             # 原始通路96        Q = self.bi(V).flatten(1)             # 重加权通路97        return self.mlp(torch.cat([P, Q], dim=1)).squeeze(-1)9899100def main() -> None:101    print("FiBiNet：SENET 重加权 + Bilinear 二阶交互（§15.5.4）\n")102    n_pairs = N_FIELDS * (N_FIELDS - 1) // 2103    d = 4104    print(f"Bilinear 参数量（F={N_FIELDS}, d={d}, pair 数={n_pairs}）：")105    print(f"  Field-All         : {d * d:,}")106    print(f"  Field-Each        : {N_FIELDS * d * d:,}")107    print(f"  Field-Interaction : {N_FIELDS * N_FIELDS * d * d:,}（是 All 的 {N_FIELDS ** 2} 倍）\n")108109    _, a_all, _ = run_experiment(lambda: FiBiNet(mode="all"), "Field-All", steps=600)110    _, a_each, _ = run_experiment(lambda: FiBiNet(mode="each"), "Field-Each", steps=600)111    _, a_base, _ = run_experiment(BaselineMLP, "Baseline", steps=600)112113    print_comparison([114        ("Baseline MLP", a_base, count_params(BaselineMLP())),115        ("FiBiNet Field-All", a_all, count_params(FiBiNet(mode="all"))),116        ("FiBiNet Field-Each", a_each, count_params(FiBiNet(mode="each"))),117    ])118    print("\n参数量从 All 到 Each 涨了 F 倍，AUC 却不一定同步上涨——")119    print("这就是正文'表达能力与成本权衡'的具体表现（§15.5.4）。")120121122if __name__ == "__main__":123    main()