AFM模型原理及Pytorch代碼複現

一、前言

該模型是和NFM模型結構上非常相似， 算是NFM模型的一個延伸，在NFM中， 不同特征域的特征embedding向量經過特征交叉池化層的交叉，将各個交叉特征向量進行“加和”, 然後後面跟了一個DNN網絡， 這裡面的問題是這個加和池化，它相當于“一視同仁”地對待所有交叉特征， 沒有考慮不同特征對結果的影響程度，作者認為這可能會影響最後的預測效果， 因為不是所有的互動特征都能夠對最後的預測起作用。 沒有用的互動特征可能會産生噪聲。

二、AFM模型原理

作者在提出NFM之後， 又對其進行了改進， 把注意力機制引入到了裡面去， 來學習不同交叉特征對于結果的不同影響程度

AFM模型結構如下：

 1、Input和embedding層

和NFM模型的一樣，也是大部分深度學習模型的标配了， 這裡為了簡單，他們的輸入把連續型的特征給省去了， 輸入的是稀疏特征， 然後進入embedding層， 得到相應稀疏特征的embedding向量

2、Pair-wise Interaction Layer

這裡和NFM是一樣的，采用的也是每對Embedding向量進行各個元素對應相乘（element-wise product）互動， 這個和FM有點不太一樣， 那裡是每對embedding的内積， 而這裡是對應元素相乘（不想加），這個要注意一下。 公式長下面這樣子：

 3、Attention based Pooling layer

本篇論文的一個核心創新 — Attention based Pooling layer。

這個想法是不同的特征互動向量在将它們壓縮為單個表示時根據對預測結果的影響程度給其加上不同權重， 然後在對其進行求和。

計算公式如下：

 其中

，表示

對的注意力分數， 表示該互動特征對于預測目标的重要性程度。為了解決泛化問題，這裡才使用了一個多層感覺器（MLP）将注意力得分參數化，就是上面的那個Attention Net。

該注意力網絡的結構是一個簡單的單全連接配接層加softmax輸出層的結構， 數學表示如下：

 4、output

基于注意力的池化層的輸出是一個k 維向量，該向量是所有特征互動向量根據重要性程度進行了區分了之後的一個聚合效果，然後我們将其映射到最終的預測得分中。是以AFM的總體公式如下：

 這個模型也是回歸任務和分類任務皆可， 并且相對于NFM， 目前上面暫時沒有用到DNN網絡來學習高階的互動了， 這個暫定為了作者未來的研究工作。

三、AFM的pytorch代碼實作

老樣子

1、DNN網絡

（可以不加，加上效果更穩點，可以更好的處理高階互動）

class Dnn(nn.Module):
    def __init__(self, hidden_units, dropout=0.):
        """
        hidden_units: 清單， 每個元素表示每一層的神經單元個數， 比如[256, 128, 64], 兩層網絡， 第一層神經單元128， 第二層64， 第一個次元是輸入次元
        dropout = 0.
        """
        super(Dnn, self).__init__()

        self.dnn_network = nn.ModuleList(
            [nn.Linear(layer[0], layer[1]) for layer in list(zip(hidden_units[:-1], hidden_units[1:]))])
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        for linear in self.dnn_network:
            x = linear(x)
            x = F.relu(x)
        x = self.dropout(x)
        return x
           

2、Attention_layer網絡

class Attention_layer(nn.Module):
    def __init__(self, att_units):
        """
        :param att_units: [embed_dim, att_vector]
        """
        super(Attention_layer, self).__init__()

        self.att_w = nn.Linear(att_units[0], att_units[1])     #8*8
        self.att_dense = nn.Linear(att_units[1], 1)               #8*1

    def forward(self, bi_interaction):  # bi_interaction (None, (field_num*(field_num-1)_/2, embed_dim)
        a = self.att_w(bi_interaction)  # (None, (field_num*(field_num-1)_/2, t)  這裡是次元變化32*325*8→ 32*325*8
        a = F.relu(a)  # (None, (field_num*(field_num-1)_/2, t)  非線性激活
        att_scores = self.att_dense(a)  # (None, (field_num*(field_num-1)_/2, 1) 再次進行次元變化 32*325*8→ 32*325*1
        att_weight = F.softmax(att_scores, dim=1)  # (None, (field_num*(field_num-1)_/2, 1)    32*325*1  對分數進行0-1範圍限定

        att_out = torch.sum(att_weight * bi_interaction, dim=1)  # (None, embed_dim)     32*325*8  求和後→32*8
        return att_out
           

3、AFM網絡

class AFM(nn.Module):
    def __init__(self, feature_columns, mode, hidden_units, att_vector=8, dropout=0.5, useDNN=False):
        """
        AFM:
        :param feature_columns: 特征資訊， 這個傳入的是fea_cols array[0] dense_info  array[1] sparse_info
        :param mode: A string, 三種模式, 'max': max pooling, 'avg': average pooling 'att', Attention
        :param att_vector: 注意力網絡的隐藏層單元個數
        :param hidden_units: DNN網絡的隐藏單元個數， 一個清單的形式， 清單的長度代表層數， 每個元素代表每一層神經元個數， lambda文裡面沒加
        :param dropout: Dropout比率
        :param useDNN: 預設不使用DNN網絡
        """
        super(AFM, self).__init__()
        self.dense_feature_cols, self.sparse_feature_cols = feature_columns
        self.mode = mode
        self.useDNN = useDNN

        # embedding
        self.embed_layers = nn.ModuleDict({
            'embed_' + str(i): nn.Embedding(num_embeddings=feat['feat_num'], embedding_dim=feat['embed_dim'])
            for i, feat in enumerate(self.sparse_feature_cols)
        })

        # 如果是注意機制的話，這裡需要加一個注意力網絡
        if self.mode == 'att':
            self.attention = Attention_layer([self.sparse_feature_cols[0]['embed_dim'], att_vector])

        # 如果使用DNN的話， 這裡需要初始化DNN網絡
        if self.useDNN:
            # 這裡要注意Pytorch的linear和tf的dense的不同之處， 前者的linear需要輸入特征和輸出特征次元， 而傳入的hidden_units的第一個是第一層隐藏的神經單元個數，這裡需要加個輸入次元
            self.fea_num = len(self.dense_feature_cols) + self.sparse_feature_cols[0]['embed_dim']  #13*8=21
            hidden_units.insert(0, self.fea_num)    #[21, 128, 64, 32]

            self.bn = nn.BatchNorm1d(self.fea_num)
            self.dnn_network = Dnn(hidden_units, dropout)
            self.nn_final_linear = nn.Linear(hidden_units[-1], 1)
        else:
            self.fea_num = len(self.dense_feature_cols) + self.sparse_feature_cols[0]['embed_dim']
            self.nn_final_linear = nn.Linear(self.fea_num, 1)

    def forward(self, x):
        dense_inputs, sparse_inputs = x[:, :len(self.dense_feature_cols)], x[:, len(self.dense_feature_cols):]
        sparse_inputs = sparse_inputs.long()  # 轉成long類型才能作為nn.embedding的輸入
        sparse_embeds = [self.embed_layers['embed_' + str(i)](sparse_inputs[:, i]) for i in
                         range(sparse_inputs.shape[1])]
        sparse_embeds = torch.stack(sparse_embeds)  # embedding堆起來， (field_dim, None, embed_dim)   26*32*8
        sparse_embeds = sparse_embeds.permute((1, 0, 2))              #32*26*8
        # 這裡得到embedding向量之後 sparse_embeds(None, field_num, embed_dim)
        # 下面進行兩兩交叉， 注意這時候不能加和了，也就是NFM的那個計算公式不能用， 這裡兩兩交叉的結果要進入Attention
        # 兩兩交叉enbedding之後的結果是一個(None, (field_num*field_num-1)/2, embed_dim)
        # 這裡實作的時候采用一個技巧就是組合
        # 比如fild_num有4個的話，那麼組合embeding就是[0,1] [0,2],[0,3],[1,2],[1,3],[2,3]位置的embedding乘積操作
        first = []
        second = []
        for f, s in itertools.combinations(range(sparse_embeds.shape[1]), 2):  #這裡就是從前面的（0-26）  産生2配對  n*（n-1）/2
            first.append(f)  #325
            second.append(s)  #325
        # 取出first位置的embedding  假設field是3的話，就是[0, 0, 0, 1, 1, 2]位置的embedding
        p = sparse_embeds[:, first, :]  # (None, (field_num*(field_num-1)_/2, embed_dim)
        q = sparse_embeds[:, second, :]  # (None, (field_num*(field_num-1)_/2, embed_dim)
        bi_interaction = p * q  # (None, (field_num*(field_num-1)_/2, embed_dim)  32*325*8

        if self.mode == 'max':
            att_out = torch.sum(bi_interaction, dim=1)  # (None, embed_dim)
        elif self.mode == 'avg':
            att_out = torch.mean(bi_interaction, dim=1)  # (None, embed_dim)
        else:
            # 注意力網絡
            att_out = self.attention(bi_interaction)  # (None, embed_dim)  32*8

        # 把離散特征和連續特征進行拼接
        x = torch.cat([att_out, dense_inputs], dim=-1)  #32*21

        if not self.useDNN:
            outputs = F.sigmoid(self.nn_final_linear(x))
        else:
            # BatchNormalization
            x = self.bn(x)
            # deep
            dnn_outputs = self.nn_final_linear(self.dnn_network(x))  #32*1
            outputs = F.sigmoid(dnn_outputs)

        return outputs
           

補充：

1、産生組合排列（兩兩特征交叉）如下部分需要注意

first = []
        second = []
        for f, s in itertools.combinations(range(sparse_embeds.shape[1]), 2):  #這裡就是從前面的（0-26）  産生2配對  n*（n-1）/2
            first.append(f)  #325
            second.append(s)  #325
        # 取出first位置的embedding  假設field是3的話，就是[0, 0, 0, 1, 1, 2]位置的embedding
        p = sparse_embeds[:, first, :]  # (None, (field_num*(field_num-1)_/2, embed_dim)
        q = sparse_embeds[:, second, :]  # (None, (field_num*(field_num-1)_/2, embed_dim)
        bi_interaction = p * q  # (None, (field_num*(field_num-1)_/2, embed_dim)  32*325*8
           

這裡沒有采用這種for循環，才是采用了排列組合的方式，求得組合數，然後選出相應位置的embedding，最後相乘得到的。 假設有4個特征embedding的話，下标位置是[0,1,2,3]， 考慮兩兩位置交叉， 那麼位置就是[0, 1], [0, 2], [0, 3], [1, 2], [1, 3], [2, 3]的6種交叉，這個會發現正好是個組合數， 我們直接用itertools.combinations函數産生上面的組合數，然後把左邊位置上的這些embedding存到一個p矩陣， 右邊位置上的embedding存入一個q矩陣， 然後兩者對應位置的embedding相乘就是交叉結果了。 這樣應該會快不少

2、對于注意力的了解

att_scores = self.att_dense(a)  #32*325*1
           

32是我們的這一批次訓練數量，  325是兩兩互動的特征   ，1是次元

可以了解為 某一個樣本（共32樣本），對這325個排列組合的感興趣程度，1維（n維也一樣）就是表示具體的權重值。

4、模型訓練

 總結

解決的痛點問題是各個特征交叉之後的embedding向量被同等看待，賦予對預測相同重要性的問題， 是以這裡加了一個注意力機制，給各個特征交叉後的embedding向量不同的權重，這樣表示了他們對預測結果的重要程度。