autoint.py

1"""2AutoInt：基于 Self-Attention 的自动特征交互（第15章 §15.5.3）34把每个 field 的 Embedding 当成一个 token：5    Q = E W_Q, K = E W_K, V = E W_V6    A = softmax(Q K^T / sqrt(d_k)),  H = A V7堆叠多头注意力，最后把各 field 的表示展平送入 MLP。相比 DCN 的固定交叉，8Attention 的权重随样本变化，是"动态特征交互"的代表。910运行：python3 ch04_feature_interaction/autoint.py11"""1213from __future__ import annotations1415import math16import pathlib17import sys1819import torch20import torch.nn as nn2122sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent))2324from _train_utils import (  # noqa: E40225    N_FIELDS, MLP, BaselineMLP, count_params, embedding_table, print_comparison,26    run_experiment, stack_embeddings,27)282930class MultiHeadSelfAttention(nn.Module):31    def __init__(self, n_fields: int, dim: int, n_heads: int = 2):32        super().__init__()33        assert dim % n_heads == 034        self.n_heads = n_heads35        self.d_head = dim // n_heads36        self.wq = nn.Linear(dim, dim)37        self.wk = nn.Linear(dim, dim)38        self.wv = nn.Linear(dim, dim)39        self.wo = nn.Linear(dim, dim)4041    def forward(self, E: torch.Tensor) -> torch.Tensor:42        # E: [B, F, d]43        B, F, d = E.shape44        def split(t):45            return t.view(B, F, self.n_heads, self.d_head).transpose(1, 2)46        q, k, v = split(self.wq(E)), split(self.wk(E)), split(self.wv(E))47        att = torch.softmax(q @ k.transpose(-1, -2) / math.sqrt(self.d_head), dim=-1)48        h = (att @ v).transpose(1, 2).reshape(B, F, d)49        return self.wo(h)505152class AutoInt(nn.Module):53    def __init__(self, dim: int = 4, n_blocks: int = 1, n_heads: int = 2, hidden=(32,)):54        super().__init__()55        self.emb = embedding_table(dim)56        self.blocks = nn.ModuleList([MultiHeadSelfAttention(N_FIELDS, dim, n_heads)57                                     for _ in range(n_blocks)])58        self.norms = nn.ModuleList([nn.LayerNorm(dim) for _ in range(n_blocks)])59        self.mlp = MLP(dim * N_FIELDS, hidden, out_dim=1, dropout=0.2)6061    def forward(self, x):62        E = stack_embeddings(self.emb, x)63        for block, norm in zip(self.blocks, self.norms):64            E = norm(E + block(E))   # 残差 + LayerNorm65        return self.mlp(E.flatten(1)).squeeze(-1)6667    def attention_map(self, x):68        """返回第一个 block 的平均注意力矩阵，用于观察'谁和谁交互'。"""69        E = stack_embeddings(self.emb, x)70        block = self.blocks[0]71        B, F, d = E.shape72        q = block.wq(E).view(B, F, block.n_heads, block.d_head).transpose(1, 2)73        k = block.wk(E).view(B, F, block.n_heads, block.d_head).transpose(1, 2)74        att = torch.softmax(q @ k.transpose(-1, -2) / math.sqrt(block.d_head), dim=-1)75        return att.mean(1).mean(0)767778def main() -> None:79    print("AutoInt：A = softmax(QK^T/√d_k)，H = AV（§15.5.3）\n")80    model, auc, _ = run_experiment(lambda: AutoInt(), "AutoInt", steps=600)81    _, base_auc, _ = run_experiment(BaselineMLP, "Baseline", steps=600)8283    # 观察某一个 batch 的注意力图84    from _train_utils import FIELD_NAMES, make_dataset85    X, _ = make_dataset(64, seed=3)86    with torch.no_grad():87        att = model.attention_map(torch.from_numpy(X).long())88    print("\nAutoInt 第一个 block 的平均注意力（行=query field，列=key field）：")89    print("        " + "  ".join(f"{n[:4]:>4}" for n in FIELD_NAMES))90    for i, name in enumerate(FIELD_NAMES):91        row = "  ".join(f"{att[i, j].item():.2f}" for j in range(N_FIELDS))92        print(f"{name[:8]:>8}  " + row)93    print("\n  对角线较大表示 field 保留自身信息；非对角元越大，说明该 field 越依赖其他 field。")94    print("  权重随样本变化正是 Attention 与固定 Cross 结构的关键区别。")9596    print_comparison([97        ("Baseline MLP", base_auc, count_params(BaselineMLP())),98        ("AutoInt（1 block, 2 heads）", auc, count_params(AutoInt())),99    ])100101102if __name__ == "__main__":103    main()