transformer多頭注意力的不一樣框架實現（tensorflow+pytorch）

時間 2020-11-19

標籤 git github 編輯器編碼 spa code orm 對象 blog get 欄目 Git 简体版

原文原文鏈接

多頭注意力能夠用如下一張圖描述：git

一、使用pytorch自帶的庫的實現github

torch.nn.MultiheadAttention(embed_dim, num_heads, dropout=0.0, bias=True, add_bias_kv=False, add_zero_attn=False, kdim=None, vdim=None)

參數說明以下：編輯器

embed_dim：最終輸出的 K、Q、V 矩陣的維度，這個維度須要和詞向量的維度同樣編碼
num_heads：設置多頭注意力的數量。若是設置爲 1，那麼只使用一組注意力。若是設置爲其餘數值，那麼 num_heads 的值須要可以被 embed_dim 整除spa
dropout：這個 dropout 加在 attention score 後面code

如今來解釋一下，爲何 num_heads 的值須要可以被 embed_dim 整除。這是爲了把詞的隱向量長度平分到每一組，這樣多組注意力也可以放到一個矩陣裏，從而並行計算多頭注意力。orm

定義 MultiheadAttention 的對象後，調用時傳入的參數以下。對象

forward(query, key, value, key_padding_mask=None, need_weights=True, attn_mask=None)

query：對應於 Key 矩陣，形狀是 (L,N,E) 。其中 L 是輸出序列長度，N 是 batch size，E 是詞向量的維度blog
key：對應於 Key 矩陣，形狀是 (S,N,E) 。其中 S 是輸入序列長度，N 是 batch size，E 是詞向量的維度get
value：對應於 Value 矩陣，形狀是 (S,N,E) 。其中 S 是輸入序列長度，N 是 batch size，E 是詞向量的維度
key_padding_mask：若是提供了這個參數，那麼計算 attention score 時，忽略 Key 矩陣中某些 padding 元素，不參與計算 attention。形狀是 (N,S)。其中 N 是 batch size，S 是輸入序列長度。

若是 key_padding_mask 是 ByteTensor，那麼非 0 元素對應的位置會被忽略
若是 key_padding_mask 是 BoolTensor，那麼 True 對應的位置會被忽略

attn_mask：計算輸出時，忽略某些位置。形狀能夠是 2D (L,S)，或者 3D (N∗numheads,L,S)。其中 L 是輸出序列長度，S 是輸入序列長度，N 是 batch size。

若是 attn_mask 是 ByteTensor，那麼非 0 元素對應的位置會被忽略
若是 attn_mask 是 BoolTensor，那麼 True 對應的位置會被忽略

須要注意的是：在實際中，K、V 矩陣的序列長度是同樣的，而 Q 矩陣的序列長度能夠不同。

這種狀況發生在：在解碼器部分的Encoder-Decoder Attention層中，Q 矩陣是來自解碼器下層，而 K、V 矩陣則是來自編碼器的輸出。

代碼示例：

## nn.MultiheadAttention 輸入第0維爲length
# batch_size 爲 64，有 12 個詞，每一個詞的 Query 向量是 300 維
query = torch.rand(12,64,300)
# batch_size 爲 64，有 10 個詞，每一個詞的 Key 向量是 300 維
key = torch.rand(10,64,300)
# batch_size 爲 64，有 10 個詞，每一個詞的 Value 向量是 300 維
value= torch.rand(10,64,300)

embed_dim = 299
num_heads = 1
# 輸出是 (attn_output, attn_output_weights)
multihead_attn = nn.MultiheadAttention(embed_dim, num_heads)
attn_output = multihead_attn(query, key, value)[0]
# output: torch.Size([12, 64, 300])
# batch_size 爲 64，有 12 個詞，每一個詞的向量是 300 維
print(attn_output.shape)

二、手動實現計算多頭注意力

在 PyTorch 提供的 MultiheadAttention 中，第 1 維是句子長度，第 2 維是 batch size。這裏咱們的代碼實現中，第 1 維是 batch size，第 2 維是句子長度。代碼裏也包括：如何用矩陣實現多組注意力的並行計算。代碼中已經有詳細註釋和說明。

class MultiheadAttention(nn.Module):
    # n_heads：多頭注意力的數量
    # hid_dim：每一個詞輸出的向量維度
    def __init__(self, hid_dim, n_heads, dropout):
        super(MultiheadAttention, self).__init__()
        self.hid_dim = hid_dim
        self.n_heads = n_heads

        # 強制 hid_dim 必須整除 h
        assert hid_dim % n_heads == 0
        # 定義 W_q 矩陣
        self.w_q = nn.Linear(hid_dim, hid_dim)
        # 定義 W_k 矩陣
        self.w_k = nn.Linear(hid_dim, hid_dim)
        # 定義 W_v 矩陣
        self.w_v = nn.Linear(hid_dim, hid_dim)
        self.fc = nn.Linear(hid_dim, hid_dim)
        self.do = nn.Dropout(dropout)
        # 縮放
        self.scale = torch.sqrt(torch.FloatTensor([hid_dim // n_heads]))

    def forward(self, query, key, value, mask=None):
        # K: [64,10,300], batch_size 爲 64，有 12 個詞，每一個詞的 Query 向量是 300 維
        # V: [64,10,300], batch_size 爲 64，有 10 個詞，每一個詞的 Query 向量是 300 維
        # Q: [64,12,300], batch_size 爲 64，有 10 個詞，每一個詞的 Query 向量是 300 維
        bsz = query.shape[0]
        Q = self.w_q(query)
        K = self.w_k(key)
        V = self.w_v(value)
        # 這裏把 K Q V 矩陣拆分爲多組注意力，變成了一個 4 維的矩陣
        # 最後一維就是是用 self.hid_dim // self.n_heads 來獲得的，表示每組注意力的向量長度, 每一個 head 的向量長度是：300/6=50
        # 64 表示 batch size，6 表示有 6組注意力，10 表示有 10 詞，50 表示每組注意力的詞的向量長度
        # K: [64,10,300] 拆分多組注意力 -> [64,10,6,50] 轉置獲得 -> [64,6,10,50]
        # V: [64,10,300] 拆分多組注意力 -> [64,10,6,50] 轉置獲得 -> [64,6,10,50]
        # Q: [64,12,300] 拆分多組注意力 -> [64,12,6,50] 轉置獲得 -> [64,6,12,50]
        # 轉置是爲了把注意力的數量 6 放到前面，把 10 和 50 放到後面，方便下面計算
        Q = Q.view(bsz, -1, self.n_heads, self.hid_dim //
                   self.n_heads).permute(0, 2, 1, 3)
        K = K.view(bsz, -1, self.n_heads, self.hid_dim //
                   self.n_heads).permute(0, 2, 1, 3)
        V = V.view(bsz, -1, self.n_heads, self.hid_dim //
                   self.n_heads).permute(0, 2, 1, 3)

        # 第 1 步：Q 乘以 K的轉置，除以scale
        # [64,6,12,50] * [64,6,50,10] = [64,6,12,10]
        # attention：[64,6,12,10]
        attention = torch.matmul(Q, K.permute(0, 1, 3, 2)) / self.scale

        # 把 mask 不爲空，那麼就把 mask 爲 0 的位置的 attention 分數設置爲 -1e10
        if mask is not None:
            attention = attention.masked_fill(mask == 0, -1e10)

        # 第 2 步：計算上一步結果的 softmax，再通過 dropout，獲得 attention。
        # 注意，這裏是對最後一維作 softmax，也就是在輸入序列的維度作 softmax
        # attention: [64,6,12,10]
        attention = self.do(torch.softmax(attention, dim=-1))

        # 第三步，attention結果與V相乘，獲得多頭注意力的結果
        # [64,6,12,10] * [64,6,10,50] = [64,6,12,50]
        # x: [64,6,12,50]
        x = torch.matmul(attention, V)

        # 由於 query 有 12 個詞，因此把 12 放到前面，把 5 和 60 放到後面，方便下面拼接多組的結果
        # x: [64,6,12,50] 轉置-> [64,12,6,50]
        x = x.permute(0, 2, 1, 3).contiguous()
        # 這裏的矩陣轉換就是：把多組注意力的結果拼接起來
        # 最終結果就是 [64,12,300]
        # x: [64,12,6,50] -> [64,12,300]
        x = x.view(bsz, -1, self.n_heads * (self.hid_dim // self.n_heads))
        x = self.fc(x)
        return x


# batch_size 爲 64，有 12 個詞，每一個詞的 Query 向量是 300 維
query = torch.rand(64, 12, 300)
# batch_size 爲 64，有 12 個詞，每一個詞的 Key 向量是 300 維
key = torch.rand(64, 10, 300)
# batch_size 爲 64，有 10 個詞，每一個詞的 Value 向量是 300 維
value = torch.rand(64, 10, 300)
attention = MultiheadAttention(hid_dim=300, n_heads=6, dropout=0.1)
output = attention(query, key, value)
## output: torch.Size([64, 12, 300])
print(output.shape)

三、tensorflow實現的多頭注意力

def _multiheadAttention(rawKeys, queries, keys, numUnits=None, causality=False, scope="multiheadAttention"):
        # rawKeys 的做用是爲了計算mask時用的，由於keys是加上了position embedding的，其中不存在padding爲0的值
        # numUnits = 50

        
        numHeads = 6
        keepProb = 1
        
        if numUnits is None:  # 如果沒傳入值，直接去輸入數據的最後一維，即embedding size.
            numUnits = queries.get_shape().as_list()[-1] #300

        # tf.layers.dense能夠作多維tensor數據的非線性映射，在計算self-Attention時，必定要對這三個值進行非線性映射，
        # 其實這一步就是論文中Multi-Head Attention中的對分割後的數據進行權重映射的步驟，咱們在這裏先映射後分割，原則上是同樣的。
        # Q, K, V的維度都是[batch_size, sequence_length, embedding_size]
        Q = tf.layers.dense(queries, numUnits, activation=tf.nn.relu) # [64,10,300]
        K = tf.layers.dense(keys, numUnits, activation=tf.nn.relu) # [64,10,300]
        V = tf.layers.dense(keys, numUnits, activation=tf.nn.relu) # [64,10,300]

        # 將數據按最後一維分割成num_heads個, 而後按照第一維拼接
        # Q, K, V 的維度都是[batch_size * numHeads, sequence_length, embedding_size/numHeads]
        Q_ = tf.concat(tf.split(Q, numHeads, axis=-1), axis=0) # [64*6,10,50]
        K_ = tf.concat(tf.split(K, numHeads, axis=-1), axis=0) # [64*6,10,50]
        V_ = tf.concat(tf.split(V, numHeads, axis=-1), axis=0) # [64*6,10,50]


        # 計算keys和queries之間的點積，維度[batch_size * numHeads, queries_len, key_len], 後兩維是queries和keys的序列長度
        similary = tf.matmul(Q_, tf.transpose(K_, [0, 2, 1])) # [64*6,10,10]

        # 對計算的點積進行縮放處理，除以向量長度的根號值
        scaledSimilary = similary / (K_.get_shape().as_list()[-1] ** 0.5) # [64*6,10,10]

        # 在咱們輸入的序列中會存在padding這個樣的填充詞，這種詞應該對最終的結果是毫無幫助的，原則上說當padding都是輸入0時，
        # 計算出來的權重應該也是0，可是在transformer中引入了位置向量，當和位置向量相加以後，其值就不爲0了，所以在添加位置向量
        # 以前，咱們須要將其mask爲0。雖然在queries中也存在這樣的填充詞，但原則上模型的結果之和輸入有關，並且在self-Attention中
        # queryies = keys，所以只要一方爲0，計算出的權重就爲0。
        # 具體關於key mask的介紹能夠看看這裏： https://github.com/Kyubyong/transformer/issues/3

        # 利用tf，tile進行張量擴張， 維度[batch_size * numHeads, keys_len] keys_len = keys 的序列長度

        # 將每一時序上的向量中的值相加取平均值
        # rawkKeys:[64,10,300]
        keyMasks = tf.sign(tf.abs(tf.reduce_sum(rawKeys, axis=-1)))  # 維度[batch_size, time_step] [64,10]
        #tf.sign()是將<0的值變爲-1，大於0的值變爲1，等於0的值變爲0
        keyMasks = tf.tile(keyMasks, [numHeads, 1]) # [64*6,10]
        # 上面這兩句的意思是找出padding的位置

        # 增長一個維度，並進行擴張，獲得維度[batch_size * numHeads, queries_len, keys_len]
        keyMasks = tf.tile(tf.expand_dims(keyMasks, 1), [1, tf.shape(queries)[1], 1]) # [64*6,10,10] 10個爲1組
        print(keyMasks.shape)
        # tf.ones_like生成元素全爲1，維度和scaledSimilary相同的tensor, 而後獲得負無窮大的值
        paddings = tf.ones_like(scaledSimilary) * (-2 ** (32 + 1)) [64*6,10,10]

        # tf.where(condition, x, y),condition中的元素爲bool值，其中對應的True用x中的元素替換，對應的False用y中的元素替換
        # 所以condition,x,y的維度是同樣的。下面就是keyMasks中的值爲0就用paddings中的值替換
        maskedSimilary = tf.where(tf.equal(keyMasks, 0), paddings, scaledSimilary) # 維度[batch_size * numHeads, queries_len, key_len]

        # 在計算當前的詞時，只考慮上文，不考慮下文，出如今Transformer Decoder中。在文本分類時，能夠只用Transformer Encoder。
        # Decoder是生成模型，主要用在語言生成中
        if causality:
            diagVals = tf.ones_like(maskedSimilary[0, :, :])  # [queries_len, keys_len]
            tril = tf.contrib.linalg.LinearOperatorTriL(diagVals).to_dense()  # [queries_len, keys_len]
            masks = tf.tile(tf.expand_dims(tril, 0), [tf.shape(maskedSimilary)[0], 1, 1])  # [batch_size * numHeads, queries_len, keys_len]

            paddings = tf.ones_like(masks) * (-2 ** (32 + 1))
            maskedSimilary = tf.where(tf.equal(masks, 0), paddings, maskedSimilary)  # [batch_size * numHeads, queries_len, keys_len]

        # 經過softmax計算權重係數，維度 [batch_size * numHeads, queries_len, keys_len]
        weights = tf.nn.softmax(maskedSimilary)

        # 加權和獲得輸出值, 維度[batch_size * numHeads, sequence_length, embedding_size/numHeads]
        outputs = tf.matmul(weights, V_)

        # 將多頭Attention計算的獲得的輸出重組成最初的維度[batch_size, sequence_length, embedding_size]
        outputs = tf.concat(tf.split(outputs, numHeads, axis=0), axis=2)
        
        outputs = tf.nn.dropout(outputs, keep_prob=keepProb)

        # 對每一個subLayers創建殘差鏈接，即H(x) = F(x) + x
        outputs += queries
        # normalization 層
        #outputs = self._layerNormalization(outputs)
        return outputs

輸入是：self.embeddedWords = self.wordEmbedded + self.positionEmbedded，即詞嵌入+位置嵌入

仍是以pytorch的輸入的維度爲例：self.wordEmbedded的維度[64,10,300] self.positionEmbedded的維度是[64,10,300]

使用的時候是：

multiHeadAtt = self._multiheadAttention(rawKeys=self.wordEmbedded, queries=self.embeddedWords,
                                  keys=self.embeddedWords)

例如：（這裏簡化了一下輸入）

wordEmbedded = tf.Variable(np.ones((64,10,300)))
positionEmbedded = tf.Variable(np.ones((64,10,300)))
embeddedWords = wordEmbedded + positionEmbedded
multiHeadAtt = _multiheadAttention(rawKeys=wordEmbedded, queries=embeddedWords, keys=embeddedWords, numUnits=300)

須要注意的是，rawkeys是針對於詞嵌入而言，由於加上了位置嵌入以後的embeddedWords的mask被位置嵌入蓋住了，就找不到須要mask的位置了。

上述pytorch的示例實際上對應的是if causality下面的代碼，由於在編碼階段：Q=K=V（它們之間的維度是相同的），在解碼階段，Q來自於解碼階段的輸入，便可以是[64,12,300]，而K和V來自編碼器的輸出，形狀都是[64,10,300]。也就是Encoder-Decoder Attention。而當QKV都來自同一個輸入的時候，也就是self attention。

參考：https://mp.weixin.qq.com/s/cJqhESxTMy5cfj0EXh9s4w