afn.py

1"""2AFN：自适应阶数的对数神经元（第15章 §15.5.5）34对数神经元：5    y = exp(Σ_i w_i ln x_i) = Π_i x_i^{w_i}6对第 j 个神经元、第 i 个 field：7    z_j = exp(Σ_i w_ij ln(e_i)) = Π_i e_i^{w_ij}8    O_j = Σ_i |w_ij|            # 该神经元学到的有效交互阶数910AFN 的关键是"阶数本身是参数"，不需要预先指定 2 阶还是 3 阶。Embedding 可能11为负，工程上先取绝对值并加 eps（正文提到的"数值处理"），并在前后加 BN。1213运行：python3 ch04_feature_interaction/afn.py14"""1516from __future__ import annotations1718import pathlib19import sys2021import torch22import torch.nn as nn2324sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent))2526from _train_utils import (  # noqa: E40227    FIELD_NAMES, N_FIELDS, MLP, BaselineMLP, count_params, embedding_table, print_comparison,28    run_experiment, stack_embeddings,29)303132class LogTransformationLayer(nn.Module):33    """J 个对数神经元，每个向量维度各自做一次 Π e_i^{w_ij}。"""3435    def __init__(self, n_fields: int, dim: int, n_neurons: int = 12):36        super().__init__()37        self.w = nn.Parameter(torch.empty(n_fields, n_neurons))38        # 初始化到 0 附近：起步接近"不交互"，让模型自己长阶数39        nn.init.normal_(self.w, std=0.1)40        self.bn_in = nn.BatchNorm1d(n_fields * dim)41        self.bn_out = nn.BatchNorm1d(n_neurons * dim)42        self.dim = dim43        self.n_neurons = n_neurons4445    def forward(self, E: torch.Tensor) -> torch.Tensor:46        # E: [B, F, d]47        B, F, d = E.shape48        e = self.bn_in(E.flatten(1)).view(B, F, d)49        log_e = torch.log(e.abs() + 1e-3)              # [B, F, d]50        # Σ_i w_ij ln(e_i)  -> [B, J, d]51        z = torch.einsum("bfd,fj->bjd", log_e, self.w)52        out = torch.exp(torch.clamp(z, -6, 6))         # 防溢出53        return self.bn_out(out.flatten(1)).view(B, self.n_neurons, d)5455    def orders(self) -> list[float]:56        return [float(v) for v in self.w.abs().sum(0)]575859class AFN(nn.Module):60    def __init__(self, dim: int = 4, n_neurons: int = 12, hidden=(32,)):61        super().__init__()62        self.emb = embedding_table(dim)63        self.log_layer = LogTransformationLayer(N_FIELDS, dim, n_neurons)64        self.mlp = MLP(n_neurons * dim, hidden, out_dim=1, dropout=0.2)6566    def forward(self, x):67        z = self.log_layer(stack_embeddings(self.emb, x))68        return self.mlp(z.flatten(1)).squeeze(-1)697071def main() -> None:72    print("AFN：z_j = exp(Σ_i w_ij ln e_i)，阶数 O_j = Σ_i |w_ij|（§15.5.5）\n")73    model, auc, _ = run_experiment(lambda: AFN(), "AFN", steps=600)74    _, base_auc, _ = run_experiment(BaselineMLP, "Baseline", steps=600)7576    with torch.no_grad():77        orders = model.log_layer.orders()78    print("\n各对数神经元学到的有效交互阶数 O_j：")79    for j, o in enumerate(orders):80        bar = "█" * max(1, int(o * 8))81        print(f"  neuron {j:2d}: O={o:5.2f}  {bar}")82    print(f"\n  阶数分布范围 {min(orders):.2f} ~ {max(orders):.2f}："83          "模型自己选择了从近 0 阶（几乎不交互）到多阶的混合，")84    print("  这就是'自适应阶数'——不需要人工指定最大交互阶数。")8586    print_comparison([87        ("Baseline MLP", base_auc, count_params(BaselineMLP())),88        ("AFN（12 神经元）", auc, count_params(AFN())),89    ])90    print("\n注：本实现按正文对 Embedding 取绝对值加 eps 再取对数（§15.5.5 的'数值处理'）。")919293if __name__ == "__main__":94    main()