AFM模型原理及Pytorch代码复现

AFM (Attentional Factorization Machine) 是一种用于CTR (Click-Through Rate) 预测的模型，其中的特征注意力部分对于模型的性能至关重要。在 PyTorch 中，可以通过以下方式控制特征添加注意力权重的比例： 1. 定义注意力网络 ```python import torch.nn as nn class AttentionNet(nn.Module): def __init__(self, feature_dim, attention_dim): super(AttentionNet, self).__init__() self.attention_layer = nn.Sequential( nn.Linear(feature_dim, attention_dim), nn.ReLU(inplace=True), nn.Linear(attention_dim, 1) ) def forward(self, x): """ :param x: feature tensor, shape: [batch_size, feature_dim] :return: attention weights tensor, shape: [batch_size, 1] """ attention_weights = self.attention_layer(x) return attention_weights ``` 2. 定义 AFM 模型 ```python class AFM(nn.Module): def __init__(self, feature_dim, attention_dim): super(AFM, self).__init__() self.embedding_layer = nn.Embedding(feature_dim, embedding_dim) self.attention_net = AttentionNet(embedding_dim, attention_dim) self.linear_layer = nn.Linear(embedding_dim, 1) def forward(self, x): """ :param x: input tensor, shape: [batch_size, num_features] :return: output tensor, shape: [batch_size] """ embeddings = self.embedding_layer(x) # [batch_size, num_features, embedding_dim] attention_weights = self.attention_net(embeddings) # [batch_size, num_features, 1] attention_weights = F.softmax(attention_weights.squeeze(dim=2), dim=1) # [batch_size, num_features] embeddings_weighted = embeddings * attention_weights.unsqueeze(dim=2) # [batch_size, num_features, embedding_dim] y_interactions = torch.sum(torch.sum(embeddings_weighted, dim=1), dim=1) # [batch_size, embedding_dim] y_linear = self.linear_layer(torch.mean(embeddings_weighted, dim=1)) # [batch_size, 1] y = y_interactions + y_linear.squeeze() # [batch_size] return y ``` 其中，`attention_weights` 表示特征注意力权重，通过 softmax 函数对其进行归一化，然后与特征向量相乘得到加权特征向量 `embeddings_weighted`，最终求和得到交互向量 `y_interactions`。在 `forward()` 方法中，还包含了一个线性层，用于学习偏置项。可以通过调整注意力网络的结构和超参数来控制特征注意力权重的比例。