10分鐘帶你深入理解Transformer原理及實(shí)現(xiàn)

點(diǎn)擊上方“小白學(xué)視覺”，選擇加"星標(biāo)"或“置頂”

重磅干貨，第一時間送達(dá)

本文轉(zhuǎn)自|深度學(xué)習(xí)這件小事

基于 Transformer《Attention Is All You Need》構(gòu)建的模型（比如 Bert ），在多個自然語言處理任務(wù)上都取得了革命性的效果，目前已取代 RNN 成為默認(rèn)選項(xiàng)，可見 Transformer 的厲害之處。

結(jié)合 Harvard 的代碼《Annotated Transformer》分享一下這個 encoder-decoder 與 attention 機(jī)制結(jié)合的方法。代碼鏈接：The Annotated Transformer。

目錄：

• 整體架構(gòu)描述

• Input & Output Embedding

• OneHot Encoding

• Word Embedding

• Positional Embedding

• Input short summary

• Encoder

• Encoder Sub-layer 1: Multi-Head Attention Mechanism

• Step 1

• Step 2

• Step 3

• Encoder Sub-layer 2: Position-Wise fully connected feed-forward

• Encoder short summary

• Decoder

• Diff_1:“masked” Multi-Headed Attention

• Diff_2:encoder-decoder multi-head attention

• Diff_3:Linear and Softmax to Produce Output Probabilities

• greedy search

• beam search

• Scheduled Sampling

0.模型架構(gòu)

今天的示例任務(wù)為中譯英:中文輸入為“我愛你”，通過 Transformer 翻譯為 “I Love You”。

Transformer 中對應(yīng)的超參數(shù)包括：

這些也是函數(shù) make_model(src_vocal, tgt_vocab, N=6, d_model=512, d_ff = 2048, h=8, dropout=0.1) 使用的超參數(shù)。

整個架構(gòu)猛一看是挺嚇人的，首先還是需要將整個 Transformer 拆分進(jìn)行描述：

• Embedding 部分

• Encoder 部分

• Decoder 部分

1. 對 Input 和 Output 進(jìn)行 representation

1.1 對 Input 的 represent

首先用常用來表達(dá) categorical 特征的方法即one-hot encoding 對句子進(jìn)行表達(dá)。one-hot 指的是一個向量只有一個元素是1，其余的都為0。很直接的，vector 的長度就是由詞匯表 vocabulary 的長度決定的。如果想要表達(dá)10000個word，那么就需要10000維的向量。

1.2 word embedding

但我們不直接給 Transformer 輸入簡單的one-hot vector，原因包括這種表達(dá)方式的結(jié)果非常稀疏，非常大，且不能表達(dá) word 與 word 之間的特征。所以這里對詞進(jìn)行 embedding，用較短的向量表達(dá)這個 word 的屬性。一般在 Pytorch 中，我們都是用 nn.Embedding 來做，或者直接用 one-hot vector 與權(quán)重矩陣 W 相乘得到。

nn.Embedding 包含一個權(quán)重矩陣 W，對應(yīng)的 shape 為 ( num_embeddings，embedding_dim )。num_embeddings 指的是詞匯量，即想要翻譯的 vocabulary 的長度。embedding_dim 指的是想用多長的 vector 來表達(dá)一個詞，可以任意選擇，比如64，128，256，512等。在 Transformer 論文中選擇的是512(即 d_model =512)。

其實(shí)可以形象地將 nn.Embedding 理解成一個 lookup table，里面對每一個 word 都存了向量 vector 。給任意一個 word，都可以從表中查出對應(yīng)的結(jié)果。

處理 nn.Embedding 權(quán)重矩陣有兩種選擇：

• 使用 pre-trained 的 embeddings 并固化，這種情況下實(shí)際就是一個 lookup table。

• 對其進(jìn)行隨機(jī)初始化(當(dāng)然也可以選擇 pre-trained 的結(jié)果)，但設(shè)為 trainable。這樣在 training 過程中不斷地對 embeddings 進(jìn)行改進(jìn)。

Transformer 選擇的是后者。

在 Annotated Transformer 中，class “Embeddings“ 用來生成 word 的embeddings，其中用到 nn.Embedding。具體實(shí)現(xiàn)見下：

   

    class Embeddings(nn.Module):
    def __init__(self, d_model, vocab):
        super(Embeddings, self).__init__()
        self.lut = nn.Embedding(vocab, d_model)
        self.d_model = d_model

    def forward(self, x):
        return self.lut(x) * math.sqrt(self.d_model)
   

1.3 Positional Embedding

我們對每一個 word 進(jìn)行 embedding 作為 input 表達(dá)。但是還有問題，embedding 本身不包含在句子中的相對位置信息。

那 RNN 為什么在任何地方都可以對同一個 word 使用同樣的向量呢？因?yàn)?RNN 是按順序?qū)渥舆M(jìn)行處理的，一次一個 word。但是在 Transformer 中，輸入句子的所有 word 是同時處理的，沒有考慮詞的排序和位置信息。

對此，Transformer 的作者提出了加入 ”positional encoding“ 的方法來解決這個問題?！眕ositional encoding“ 使得 Transformer 可以衡量 word 位置有關(guān)的信息。

positional encoding 與 word embedding 相加就得到 embedding with position。

那么具體 ”positional encoding“ 怎么做？為什么能表達(dá)位置信息呢？作者探索了兩種創(chuàng)建 positional encoding 的方法：

• 通過訓(xùn)練學(xué)習(xí) positional encoding 向量

• 使用公式來計(jì)算 positional encoding向量

試驗(yàn)后發(fā)現(xiàn)兩種選擇的結(jié)果是相似的，所以采用了第2種方法，優(yōu)點(diǎn)是不需要訓(xùn)練參數(shù)，而且即使在訓(xùn)練集中沒有出現(xiàn)過的句子長度上也能用。

計(jì)算 positional encoding 的公式為：

在這個公式中:

• pos 指的是這個 word 在這個句子中的位置

• i指的是 embedding 維度。比如選擇 d_model=512，那么i就從1數(shù)到512

為什么選擇 sin 和 cos ？positional encoding 的每一個維度都對應(yīng)著一個正弦曲線，作者假設(shè)這樣可以讓模型相對輕松地通過對應(yīng)位置來學(xué)習(xí)。

在 Annotated Transformer 中，使用 class ”Positional Encoding“ 來創(chuàng)建 positional encoding 并加入到 word embedding 中：

   

    class PositionalEncoding(nn.Module):
    "Implement the PE function."
    def __init__(self, d_model, dropout, max_len=5000):
        super(PositionalEncoding, self).__init__()
        self.dropout = nn.Dropout(p=dropout)

        # Compute the positional encodings once in log space.
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_model, 2) *
                             -(math.log(10000.0) / d_model))
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        pe = pe.unsqueeze(0)
        self.register_buffer('pe', pe)

    def forward(self, x):
        x = x + Variable(self.pe[:, :x.size(1)],requires_grad=False)
        return self.dropout(x)
   

波的頻率和偏移對于每個維度是不同的：

1.4 Input 小總結(jié)

經(jīng)過 word embedding 和 positional embedding 后可以得到一個句子的 representation，比如”我愛你“這個句子，就被轉(zhuǎn)換成了三個向量，每個向量都包含 word 的特征和 word 在句子中的位置信息：

我們對輸出的結(jié)果做同樣的操作，這里即中英翻譯的結(jié)果 ”I Love You“。使用word embedding 和 positional encoding 對其進(jìn)行 represent。

Input Tensor 的 size 為 [nbatches, L, 512]：

• nbatches 指的是定義的 batch_size

• L 指的是 sequence 的長度,(比如“我愛你”，L = 3)

• 512 指的是 embedding 的 dimension

目前完成了模型架構(gòu)的底層的部分：

2. Encoder

Encoder 相對 Decoder 會稍微麻煩一些。Encoder 由 6 個相乘的 Layer 堆疊而成（6并不是固定的，可以基于實(shí)際情況修改），看起來像這樣：

每個 Layer 包含 2 個 sub-layer：

• 第一個是 ”multi-head self-attention mechanism“

• 第二個是 ”simple，position-wise fully connected feed-forward network“

Annotated Transformer 中的 Encoder 實(shí)現(xiàn)代碼：

   

    class Encoder(nn.Module):
    "Core encoder is a stack of N layers"

    def __init__(self, layer, N):
        super(Encoder, self).__init__()
        self.layers = clones(layer, N)
        self.norm = LayerNorm(layer.size)

    def forward(self, x, mask):
        "Pass the input (and mask) through each layer in turn."
        for layer in self.layers:
            x = layer(x, mask)
        return self.norm(x)

class EncoderLayer(nn.Module):
    "Encoder is made up of self-attn and feed forward (defined below)"
    def __init__(self, size, self_attn, feed_forward, dropout):
        super(EncoderLayer, self).__init__()
        self.self_attn = self_attn
        self.feed_forward = feed_forward
        self.sublayer = clones(SublayerConnection(size, dropout), 2)
        self.size = size

    def forward(self, x, mask):
        "Follow Figure 1 (left) for connections."
        x = self.sublayer[0](x, lambda x: self.self_attn(x, x, x, mask))
        return self.sublayer[1](x, self.feed_forward)
   

1. class “Encoder” 將 <layer> 堆疊N次。是 class “EncoderLayer” 的實(shí)例。

2. “EncoderLayer” 初始化需要指定<size>,<self_attn>,<feed_forward>,<dropout>：

1. <size> 對應(yīng) d_model，論文中為512

2. <self_attn> 是 class MultiHeadedAttention 的實(shí)例，對應(yīng)sub-layer 1

3. <feed_forward> 是 class PositionwiseFeedForward 的實(shí)例，對應(yīng)sub-layer 2

4. <dropout> 對應(yīng) dropout rate

2.1 Encoder Sub-layer 1: Multi-Head Attention Mechanism

理解 Multi-Head Attention 機(jī)制對于理解 Transformer 特別重要，并且在 Encoder 和 Decoder 中都有用到。

概述：
我們把 attention 機(jī)制的輸入定義為 x。x 在 Encoder 的不同位置，含義有所不同。在 Encoder 的開始，x 的含義是句子的 representation。在 EncoderLayer 的各層中間，x 代表前一層 EncoderLayer 的輸出。

使用不同的 linear layers 基于 x 來計(jì)算 keys，queries和values：

• key = linear_k(x)

• query = linear_q(x)

• value = linear_v(x)

linear_k, linear_q, linear_v 是相互獨(dú)立、權(quán)重不同的。

計(jì)算得到 keys(K), queries(Q)和values(V) 值之后，按論文中如下公式計(jì)算 Attention：

矩陣乘法表示：

這里的比較奇怪的地方是為啥要除以 sqrt(d_k) 對吧？

作者的解釋是說防止   增大時，   點(diǎn)積值過大，所以用   對其進(jìn)行縮放。引用一下原文”We suspect that for large values of dk, the dot products grow large in magnitude, pushing the softmax function into regions where it has extremely small gradients” 對  取 softmax 之后值都介于0到1之間，可以理解成得到了 attention weights。然后基于這個 attention weights 對 V 求 weighted sum 值 Attention(Q, K, V)。

詳細(xì)解釋：Annotated Transformer 中 Multi-Headed attention 的實(shí)現(xiàn)為

   

    class MultiHeadedAttention(nn.Module):
    def __init__(self, h, d_model, dropout=0.1):
        "Take in model size and number of heads."
        super(MultiHeadedAttention, self).__init__()
        assert d_model % h == 0
        # We assume d_v always equals d_k
        self.d_k = d_model // h
        self.h = h
        self.linears = clones(nn.Linear(d_model, d_model), 4)
        self.attn = None
        self.dropout = nn.Dropout(p=dropout)

    def forward(self, query, key, value, mask=None):
        "Implements Figure 2"
        if mask is not None:
            # Same mask applied to all h heads.
            mask = mask.unsqueeze(1)
        nbatches = query.size(0)

        # 1) Do all the linear projections in batch from d_model => h x d_k
        query, key, value = \
            [l(x).view(nbatches, -1, self.h, self.d_k).transpose(1, 2)
             for l, x in zip(self.linears, (query, key, value))]

        # 2) Apply attention on all the projected vectors in batch.
        x, self.attn = attention(query, key, value, mask=mask,
                                 dropout=self.dropout)

        # 3) "Concat" using a view and apply a final linear.
        x = x.transpose(1, 2).contiguous() \
            .view(nbatches, -1, self.h * self.d_k)
        return self.linears[-1](x)
   

這個 class 進(jìn)行實(shí)例化時需要指定:

• <h> = 8，即 “heads” 的數(shù)目。在 Transformer 的 base model 中有8 heads

• <d_model> = 512

• <dropout> = dropout rate = 0.1

keys 的維度 d_k 是基于   計(jì)算來的。在上面的例子中 d_k = 512 / 8 = 64。

下面分3步詳細(xì)介紹一下 MultiHeadedAttention 的 forward() 函數(shù)：

從上面的代碼看出，forward 的 input 包括：query，key，values和mask。這里先暫時忽略 mask。query，key和value 是哪來的？實(shí)際上他們是 “x” 重復(fù)了三次得來的，x 或者是初始的句子 embedding或者是前一個 EncoderLayer 的輸出，見 EncoderLayer 的代碼黃色劃線部分，self.self_atttn 是 MultiHeadedAttention 的一個實(shí)例化：

“query” 的 shape 為 [nbatches, L, 512] ,其中:

• nbatches 對應(yīng) batch size

• L 對應(yīng) sequence length ，512 對應(yīng) d_mode

• “key” 和 “value” 的 shape 也為 [nbatches， L， 512]

Step 1）

1. 對 “query”，“key”和“value”進(jìn)行 linear transform ，他們的 shape 依然是[nbatches， L， 512]。

2. 對其通過 view() 進(jìn)行 reshape，shape 變成 [nbatches, L, 8, 64]。這里的h=8對應(yīng) heads 的數(shù)目，d_k=64 是 key 的維度。

3. transpose 交換 dimension1和2，shape 變成 [nbatches， 8, L 64]。

Step 2）

前面提到我們計(jì)算 attention 的公式：

Annotated Transformer 中的 attention() 代碼為：

   

    def attention(query, key, value, mask=None, dropout=None):
    "Compute 'Scaled Dot Product Attention'"
    d_k = query.size(-1)
    scores = torch.matmul(query, key.transpose(-2, -1)) \
             / math.sqrt(d_k)
    if mask is not None:
        scores = scores.masked_fill(mask == 0, -1e9)
    p_attn = F.softmax(scores, dim = -1)
    if dropout is not None:
        p_attn = dropout(p_attn)
    return torch.matmul(p_attn, value), p_attn
   

query 和 key.transpose(-2,-1) 相乘，兩者分別對應(yīng)的 shape 為 [nbatches, 8, L 64] 和 [nbatches， 8， 64， L]。這樣相乘得到的結(jié)果 scores 的 shape為[nbatches, 8, L, L]。

對 scores 進(jìn)行 softmax，所以 p_attn 的 shape 為 [nbatches, 8, L, L]。values的 shape 為 [nbatches, 8, L, 64]。所以最后 p_attn 與 values 相乘輸出的 result 的 shape 為 [nbatches, 8, L, 64]。

在我們的輸入與輸出中，有8個 heads 即 Tensor 中的 dimension 1，[ nbatches, 8, L, 64 ]。8個 heads 都進(jìn)行了不同的矩陣乘法，這樣就得到了不同的 “representation subspace”。這就是 multi-headed attention 的意義。

Step 3）

x的初始shape為 [ nbatches, 8, L, 64 ]，x.transpose(1,2) 得到 [ nbatches，L, 8,64 ]。然后使用 view 進(jìn)行 reshape 得到 [ nbatches, L, 512 ]?？梢岳斫鉃?個heads結(jié)果的 concatenate 。最后使用 last linear layer 進(jìn)行轉(zhuǎn)換。shape仍為 [ nbatches, L, 512 ]。與input時的shape是完全一致的。

可視化見論文中的圖例：

2.2 Encoder Sub-layer 2: Position-Wise fully connected feed-forward network

SubLayer-2 只是一個 feed-forward network。比較簡單。

在 Annotated Transformer 中對應(yīng)的實(shí)現(xiàn)為：

   

    class PositionwiseFeedForward(nn.Module):
    "Implements FFN equation."
    def __init__(self, d_model, d_ff, dropout=0.1):
        super(PositionwiseFeedForward, self).__init__()
        self.w_1 = nn.Linear(d_model, d_ff)
        self.w_2 = nn.Linear(d_ff, d_model)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        return self.w_2(self.dropout(F.relu(self.w_1(x))))
   

2.3 Encoder short summary

Encoder 總共包含6個 EncoderLayers 。每一個 EncoderLayer 包含2個 SubLayer：

• SubLayer-1 做 Multi-Headed Attention

• SubLayer-2 做 feedforward neural network

3. The Decoder

Encoder 與 Decoder 的交互方式可以理解為：

Decoder 也是N層堆疊的結(jié)構(gòu)。被分為3個 SubLayer，可以看出 Encoder 與 Decoder 三大主要的不同：

• Diff_1：Decoder SubLayer-1 使用的是 “masked” Multi-Headed Attention 機(jī)制，防止為了模型看到要預(yù)測的數(shù)據(jù)，防止泄露。

• Diff_2：SubLayer-2 是一個 encoder-decoder multi-head attention。

• Diff_3：LinearLayer 和 SoftmaxLayer 作用于 SubLayer-3 的輸出后面，來預(yù)測對應(yīng)的 word 的 probabilities 。

3.1 Diff_1 : “masked” Multi-Headed Attention

mask 的目標(biāo)在于防止 decoder “seeing the future”，就像防止考生偷看考試答案一樣。mask包含1和0：

Attention 中使用 mask 的代碼中：

   

    if mask is not None:
        scores = scores.masked_fill(mask == 0, -1e9)
   

引用作者的話說， “We […] modify the self-attention sub-layer in the decoder stack to prevent positions from attending to subsequent positions. This masking, combined with fact that the output embeddings are offset by one position, ensures that the predictions for position i can depend only on the known outputs at positions less than i.”

3.2 Diff_2 : encoder-decoder multi-head attention

Annotated Transformer 中的 DecoderLayer 的實(shí)現(xiàn)為：

   

    class DecoderLayer(nn.Module):
    "Decoder is made of self-attn, src-attn, and feed forward (defined below)"
    def __init__(self, size, self_attn, src_attn, feed_forward, dropout):
        super(DecoderLayer, self).__init__()
        self.size = size
        self.self_attn = self_attn
        self.src_attn = src_attn
        self.feed_forward = feed_forward
        self.sublayer = clones(SublayerConnection(size, dropout), 3)

    def forward(self, x, memory, src_mask, tgt_mask):
        m = memory
        x = self.sublayer[0](x, lambda x: self.self_attn(x, x, x, tgt_mask))
        x = self.sublayer[1](x, lambda x: self.src_attn(x, m, m, src_mask))
        return self.sublayer[2](x, self.feed_forward)
   

重點(diǎn)在于 x = self.sublayer1 self.src_attn 是 MultiHeadedAttention 的一個實(shí)例。query = x，key = m, value = m, mask = src_mask，這里x來自上一個 DecoderLayer，m來自 Encoder的輸出。

到這里 Transformer 中三種不同的 Attention 都已經(jīng)集齊了：

3.3 Diff_3 : Linear and Softmax to Produce Output Probabilities

最后的 linear layer 將 decoder 的輸出擴(kuò)展到與 vocabulary size 一樣的維度上。經(jīng)過 softmax 后，選擇概率最高的一個 word 作為預(yù)測結(jié)果。

假設(shè)我們有一個已經(jīng)訓(xùn)練好的網(wǎng)絡(luò)，在做預(yù)測時，步驟如下：

1. 給 decoder 輸入 encoder 對整個句子 embedding 的結(jié)果 和一個特殊的開始符號 </s>。decoder 將產(chǎn)生預(yù)測，在我們的例子中應(yīng)該是 ”I”。

2. 給 decoder 輸入 encoder 的 embedding 結(jié)果和 “</s>I”，在這一步 decoder 應(yīng)該產(chǎn)生預(yù)測 “Love”。

3. 給 decoder 輸入 encoder 的 embedding 結(jié)果和 “</s>I Love”，在這一步 decoder 應(yīng)該產(chǎn)生預(yù)測 “China”。

4. 給 decoder 輸入 encoder 的 embedding 結(jié)果和 “</s>I Love China”, decoder應(yīng)該生成句子結(jié)尾的標(biāo)記，decoder 應(yīng)該輸出 ”</eos>”。

5. 然后 decoder 生成了 </eos>，翻譯完成。

但是在訓(xùn)練過程中，decoder 沒那么好時，預(yù)測產(chǎn)生的詞很可能不是我們想要的。這個時候如果再把錯誤的數(shù)據(jù)再輸給 decoder，就會越跑越偏：

這里在訓(xùn)練過程中要使用到 “teacher forcing”。利用我們知道他實(shí)際應(yīng)該預(yù)測的 word 是什么，在這個時候喂給他一個正確的結(jié)果作為輸入。

相對于選擇最高的詞 (greedy search)，還有其他選擇是比如 “beam search”，可以保留多個預(yù)測的 word。Beam Search 方法不再是只得到一個輸出放到下一步去訓(xùn)練了，我們可以設(shè)定一個值，拿多個值放到下一步去訓(xùn)練，這條路徑的概率等于每一步輸出的概率的乘積，具體可以參考李宏毅老師的課程：

或者 “Scheduled Sampling”：一開始我們只用真實(shí)的句子序列進(jìn)行訓(xùn)練，而隨著訓(xùn)練過程的進(jìn)行，我們開始慢慢加入模型的輸出作為訓(xùn)練的輸入這一過程。

這部分對應(yīng) Annotated Transformer 中的實(shí)現(xiàn)為：

   

    class Generator(nn.Module):
    "Define standard linear + softmax generation step."
    def __init__(self, d_model, vocab):
        super(Generator, self).__init__()
        self.proj = nn.Linear(d_model, vocab)

    def forward(self, x):
        return F.log_softmax(self.proj(x), dim=-1)
   

對著圖，再回顧一下 Encoder 與 Decoder 的結(jié)構(gòu)吧。

參考鏈接：

https://arxiv.org/pdf/1706.03762.pdf

https://glassboxmedicine.com/2019/08/15/the-transformer-attention-is-all-you-need/

https://jalammar.github.io/illustrated-transformer/

—完—

下載1：OpenCV-Contrib擴(kuò)展模塊中文版教程

在「小白學(xué)視覺」公眾號后臺回復(fù)：擴(kuò)展模塊中文教程，即可下載全網(wǎng)第一份OpenCV擴(kuò)展模塊教程中文版，涵蓋擴(kuò)展模塊安裝、SFM算法、立體視覺、目標(biāo)跟蹤、生物視覺、超分辨率處理等二十多章內(nèi)容。

下載2：Python視覺實(shí)戰(zhàn)項(xiàng)目52講

在「小白學(xué)視覺」公眾號后臺回復(fù)：Python視覺實(shí)戰(zhàn)項(xiàng)目，即可下載包括圖像分割、口罩檢測、車道線檢測、車輛計(jì)數(shù)、添加眼線、車牌識別、字符識別、情緒檢測、文本內(nèi)容提取、面部識別等31個視覺實(shí)戰(zhàn)項(xiàng)目，助力快速學(xué)校計(jì)算機(jī)視覺。

下載3：OpenCV實(shí)戰(zhàn)項(xiàng)目20講

在「小白學(xué)視覺」公眾號后臺回復(fù)：OpenCV實(shí)戰(zhàn)項(xiàng)目20講，即可下載含有20個基于OpenCV實(shí)現(xiàn)20個實(shí)戰(zhàn)項(xiàng)目，實(shí)現(xiàn)OpenCV學(xué)習(xí)進(jìn)階。

交流群

歡迎加入公眾號讀者群一起和同行交流，目前有SLAM、三維視覺、傳感器、自動駕駛、計(jì)算攝影、檢測、分割、識別、醫(yī)學(xué)影像、GAN、算法競賽等微信群（以后會逐漸細(xì)分），請掃描下面微信號加群，備注：”昵稱+學(xué)校/公司+研究方向“，例如：”張三 + 上海交大 + 視覺SLAM“。請按照格式備注，否則不予通過。添加成功后會根據(jù)研究方向邀請進(jìn)入相關(guān)微信群。請勿在群內(nèi)發(fā)送廣告，否則會請出群，謝謝理解~