從理論到實現，手把手實現Attention網絡
作者 | 梁唐
出品 | 公眾號：Coder梁（ID：Coder_LT）
大家好，我是老梁。
我們之前介紹了Transformer的核心——attention網絡，我們之前只是介紹了它的原理，并且沒有詳細解釋它的實現方法。光聊理論難免顯得有些空洞，所以我們來談談它的實現。
為了幫助大家更好地理解， 這里我選了電商場景中的DIN模型來做切入點。
一方面可以幫助大家理解現在電商系統中的推薦和廣告系統中的商品排序都是怎么做的，另外我個人感覺DIN要比直接去硬啃transformer容易理解一些。
我們可以先從attention網絡的數據入手，它的輸入數據有兩個：一個是用戶的歷史行為序列，一個是待打分的item（以下稱為target item）。用戶的歷史行為序列本質上其實就是一個用戶歷史上有過交互的item的數組。這里為了簡化，我們假設已經完成了從item到embedding的轉換。
首先是target item，它的shape應該是[B, E]。這里的B指的是batch_size，即訓練時候一個批量的大小。這里的E指的是embedding的長度。也有一些文章里使用別的字母表示，這也沒有一個硬性的標準，能看懂就行。
我們再來看用戶行為序列，除了batch_size和embedding長度之外，還需要一個額外的參數來表示行為序列的長度，通常我們用字母T。對于所有的樣本，我們都需要保證它的行為序列長度是T，如果不足T的，則使用默認值補足。如果超過T的，則進行截斷。如此，它的shape應該是[B, T, E]。
根據attention網絡的原理，我們需要根據行為序列中的每個item與target item的相似度，再根據相似度計算權重。最后對這T個item的embedding進行加權求和。求和之后，這T個item根據計算得到的權重合并得到一個embedding。論文中說這個集成T個行為序列的embedding就是用戶興趣的表達，我們只需要將它和目標item拼接在一起發送到神經網絡即可，就可以幫助模型更好地決策了。這里用戶興趣的shape應該和item是一樣的，也是[B, E]。
簡單總結一下，我們現在需要一個模塊，它接收兩個輸入，一個是item的embedding，一個是用戶行為序列的embedding。它的輸出應該是[B, T]，對應行為序列中T個item的權重。剩下的問題就是怎么生成這個結果。
原理講完了，接下來講講實現，我們可以結合一下下面這兩張論文中的結構圖幫助理解。
圖片
首先，我們來統一一下輸入的維度，手動將item的embedding這個二維的向量變成三維，即shape變成[B, 1, E]。
這里一種做法是，手動循環T次，每次從行為序列中拿出一個item embedding，和目標item的embedding拼接在一起丟進一個神經網絡中得到一個分數。
這種做法非常不推薦，一般在神經網絡當中，我們不到萬不得已，不手動循環，因為循環是線性計算，沒辦法利用GPU的并行計算來加速。
對于當前問題來說，我們完全可以使用矩陣運算來代替。通過使用expand/tile函數，將[B, 1, E]的item embedding復制T份，形狀也變成[B, T, E]。這樣一來，兩個輸入的shape都變成了[B, T, E]，我們就可以把它們拼接到一起變成[B, T, 2E]。
然后經過一個輸入是2E，輸出是1的神經網絡，最終得到[B, T, 1]的結果，我們把它調換一下維度，變成[B, 1, T]，這個就是我們想要的權重了。
這里我找來一份Pytorch的代碼，大家代入一下上面的邏輯去看一下，應該不難看懂。
class LocalActivationUnit(nn.Module):
def __init__(self, hidden_units=(64, 32), embedding_dim=4, activation='sigmoid', dropout_rate=0, dice_dim=3,
l2_reg=0, use_bn=False):
super(LocalActivationUnit, self).__init__()
self.dnn = DNN(inputs_dim=4 * embedding_dim,
hidden_units=hidden_units,
activation=activation,
l2_reg=l2_reg,
dropout_rate=dropout_rate,
dice_dim=dice_dim,
use_bn=use_bn)
self.dense = nn.Linear(hidden_units[-1], 1)def forward(self, query, user_behavior): # query ad : size -> batch_size * 1 * embedding_size # user behavior : size -> batch_size * time_seq_len * embedding_size user_behavior_len = user_behavior.size(1)
queries = query.expand(-1, user_behavior_len, -1)
attention_input = torch.cat([queries, user_behavior, queries - user_behavior, queries * user_behavior],
dim=-1) # as the source code, subtraction simulates verctors' difference
attention_output = self.dnn(attention_input)
attention_score = self.dense(attention_output) # [B, T, 1]
return attention_score其實這一段代碼就是attention網絡的核心，它生成的是attention中最重要的權重。權重有了之后，我們只需要將它和用戶行為序列的embedding相乘。利用矩陣乘法的特性，一個[B, 1. T]的矩陣乘上一個[B, T, E]的矩陣，會得到[B, 1, E]的結果。這個相乘之后的結果其實就是我們需要的加權求和，只不過是通過矩陣乘法來實現了。
我們再看下源碼加深一下理解：
class AttentionSequencePoolingLayer(nn.Module): """The Attentional sequence pooling operation used in DIN & DIEN.
Arguments
- **att_hidden_units**:list of positive integer, the attention net layer number and units in each layer.
- **att_activation**: Activation function to use in attention net.
- **weight_normalization**: bool.Whether normalize the attention score of local activation unit.
- **supports_masking**:If True,the input need to support masking.
References
- [Zhou G, Zhu X, Song C, et al. Deep interest network for click-through rate prediction[C]//Proceedings of the 24th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining. ACM, 2018: 1059-1068.](https://arxiv.org/pdf/1706.06978.pdf)
"""
def __init__(self, att_hidden_units=(80, 40), att_activation='sigmoid', weight_normalization=False,
return_score=False, supports_masking=False, embedding_dim=4, **kwargs):
super(AttentionSequencePoolingLayer, self).__init__()
self.return_score = return_score
self.weight_normalization = weight_normalization
self.supports_masking = supports_masking
self.local_att = LocalActivationUnit(hidden_units=att_hidden_units, embedding_dim=embedding_dim,
activation=att_activation,
dropout_rate=0, use_bn=False)[docs] def forward(self, query, keys, keys_length, mask=None): """ Input shape - A list of three tensor: [query,keys,keys_length]
- query is a 3D tensor with shape: ``(batch_size, 1, embedding_size)``
- keys is a 3D tensor with shape: ``(batch_size, T, embedding_size)``
- keys_length is a 2D tensor with shape: ``(batch_size, 1)``
Output shape
- 3D tensor with shape: ``(batch_size, 1, embedding_size)``.
"""
batch_size, max_length, _ = keys.size()
# Mask
if self.supports_masking:
if mask is None:
raise ValueError("When supports_masking=True,input must support masking")
keys_masks = mask.unsqueeze(1)
else:
keys_masks = torch.arange(max_length, device=keys_length.device, dtype=keys_length.dtype).repeat(batch_size,1) # [B, T]
keys_masks = keys_masks < keys_length.view(-1, 1) # 0, 1 mask
keys_masks = keys_masks.unsqueeze(1) # [B, 1, T]
attention_score = self.local_att(query, keys) # [B, T, 1]
outputs = torch.transpose(attention_score, 1, 2) # [B, 1, T]
if self.weight_normalization:
paddings = torch.ones_like(outputs) * (-2 ** 32 + 1)
else:
paddings = torch.zeros_like(outputs)
outputs = torch.where(keys_masks, outputs, paddings) # [B, 1, T]
# Scale
# outputs = outputs / (keys.shape[-1] ** 0.05)
if self.weight_normalization:
outputs = F.softmax(outputs, dim=-1) # [B, 1, T]
if not self.return_score:
# Weighted sum
outputs = torch.matmul(outputs, keys) # [B, 1, E]
return outputs這段代碼當中加入了mask以及normalization等邏輯，全部忽略掉的話，真正核心的主干代碼就只有三行：
attention_score = self.local_att(query, keys) # [B, T, 1] outputs = torch.transpose(attention_score, 1, 2) # [B, 1, T] outputs = torch.matmul(outputs, keys) # [B, 1, E] 到這里我們關于attention網絡的實現方法就算是講完了，對于DIN這篇論文也就理解差不多了，不過還有一個細節值得聊聊。就是關于attention權重的問題。
在DIN這篇論文當中，我們是使用了一個單獨的LocalActivationUnit來學習的兩個embedding拼接之后的權重，也就是上圖代碼當中這個部分：
圖片
我們通過一個單獨的神經網絡來對兩個向量打分給出權重，這個權重的運算邏輯并不一定是根據向量的相似度來計算的。畢竟神經網絡是一個黑盒，我們無從猜測內部邏輯。只不過從邏輯上或者經驗上來說，我們更傾向于它是根據向量的相似度來計算的。
而Transformer當中也有attention結構，它就是正兒八經地利用向量之間的相似度來計算的。常理上來說，按照向量相似度來計算權重，這種做法應該更容易理解一些。但實際上學習的過程當中的感受卻并不一定如此，這也是為什么我先來分享DIN而不是直接上transformer self-attention的原因。
