Transformer解讀（論文 + PyTorch源碼）

Transformer模型早在2017年就出現了，當時實驗室的分享也有關於這個的。但我當時沒有意識到這篇論文的厲害之處，聽名字感覺像是那種曇花一現的論文，也沒有關注它。直到最近出現了BERT這一神物之後，方纔後知後覺此時Transformer已然這麼有用！因此，這才仔仔細細地擼了這篇“古老”的論文和源碼，這裏將主要對照論文和相應的PyTorch源碼進行逐一對照解讀。因筆者能力有限，如有不詳實之處，讀者可移步至文末的傳送門去看更多細節，並歡迎指出~

前言

2017年6月，Google發佈了一篇論文《Attention is All You Need》，提出了Transformer模型。正如論文的名稱所說，其旨在全部利用Attention方式來替代掉RNN的循環機制，從而能並行化計算並實現提速。同時，在特定的任務上，這個模型也超過了當時Google神經機器翻譯模型。筆者主要閱讀了論文及兩篇博客（鏈接見文末的傳送門），這裏主要是對這些內容做一個整合和提煉~

一. 背景

在Transformer出現之前，LSTM、GRU等RNN系列網絡以及encoder-decoder+attention架構基本上鑄就了所有NLP任務的鐵桶江山。但RNN的一個缺陷在於是自迴歸的模型，只能串行的一步一步進行計算，無法並行化。因此有一些網絡如ByteNet和ConvS2S都是以此爲切入點，使用CNN作爲基本構建模塊，這樣可以並行計算所有輸入和輸出位置的隱層表示。但在這些模型中，關聯來自兩個任意輸入或輸出位置的信號所需的操作數量會隨着位置之間的距離而增長，如ConvS2S呈線性增長、ByteNet呈對數增長，這使得學習較遠位置之間的依賴變得更加困難。而在Transformer中，兩個輸入之間的距離對其計算來說沒有影響，都是一樣的，它沒有使用RNN和卷積，可以進行並行計算。

二. Transformer整體架構

下面是從論文中截出的Transformer整體結構圖：

這個圖乍一看非常唬人，但實際上仔細看的話仍舊是熟悉的Encoder-Decoder架構，左邊的是Encoder，右邊的是Decoder。下面將一一進行剖析。

三. Transformer細節剖析

1. 編碼器

首先來看Encoder部分（左半部分），它是由N層方框裏面的內容堆疊起來的。對於每一層來說，都由兩部分構成：一部分是multi-head self-attention機制，另一部分是一個簡單的全連接前饋網絡。在每一部分上，都使用殘差+layer normalization來進行處理。論文中，這樣的方框有6個，即 $N=6$ ，模型的隱層單元數 $d_{model} = 512$ 。

2. 自注意力機制

Encoder內部沒有使用RNN，取而代之的是一種self-attention（自注意力）機制。

一般我們用的attention機制，可以抽象爲輸入一個查詢（query），去查詢鍵值對（key-value pair）中的key，然後得到一個概率分佈，再據此對value進行加權相加，獲取當前query下的注意力表徵。而我們的query，往往是Decoder中某一個step的輸出，key-value pair往往是encoder的輸出。

論文裏面使用的也是這種attention機制，只不過其query、key、value都是由encoder的輸出經過不同的變換而來，也即self-attention，所有的東西都是自己。他們定義了一種叫“Scaled Dot-Product Attention”的計算方式，用於計算給定query、key和value下的注意力表徵，如下圖（左）所示：

這裏的 $Q$ 、 $K$ 和 $V$ 分別表示query、key和value矩陣，它們的維度分別爲 $L_q * d_k$ 、 $L_k * d_k$ 、 $L_k * d_v$ 。計算公式爲：

一般我們經常使用的attention計算方式有兩種：一種是乘性attention，即使用內積的方式；另一種是加性attention，即使用額外一層隱藏層來計算。這兩種計算方式理論上覆雜度是差不多的，但乘性attention因爲可以用矩陣運算，會更節省時間和空間。對照着上圖（左）和公式來看，這個公式與乘性attention計算方式的唯一不同就在於使用了一個縮放因子 $\frac{1}{\sqrt{d_k}}$ 。這裏爲何要進行縮放呢？論文中給出瞭解釋：在 $d_k$ 比較小的時候，不加縮放的效果和加性attention的效果差不多，但當 $d_k$ 比較大的時候，不加縮放的效果就明顯比加性attention的效果要差，懷疑是當 $d_k$ 增長的時候，內積的量級也會增長，導致softmax函數會被推向梯度較小的區域，爲了緩解這個問題，加上了這個縮放項進行量級縮小。

論文裏面還提到，只使用一個attention的計算方式未免太過單薄，所以他們提出了multi-head（多頭）注意力機制。將注意力的計算分散到不同的子空間進行，以期能從多方面進行注意力的學習，具體做法如上圖（右）所示。並行地將 $Q$ 、 $K$ 和 $V$ 通過不同的映射矩陣映射到不同的空間（每個空間是一個頭），再分別在這些空間中對應着進行單個“Scaled Dot-Product Attention”的學習，最後將得到的多頭注意力表徵進行拼接，經過一個額外的映射層映射到原來的空間。其公式如下：

這裏的 $W_i^Q \in R^{d_{model} * d_k}$ ， $W_i^K \in R^{d_{model} * d_k}$ ， $W_i^V \in R^{d_{model} * d_v}$ ， $W^O \in R^{hd_v * d_{model}}$ 。表示第 $i$ 個頭的變換矩陣， $h$ 表示頭的個數。

在論文裏面， $h = 8$ ，並且 $d_k = d_v = d_{model} / h = 64$ 。可見這裏雖然分了很多頭去計算，但整體的維度還是不變的，因此計算量還是一樣的。

3. 前饋網絡

這部分是整體架構圖中的Feed Forward模塊，其實就是一個簡單的全連接前饋網絡。它由兩層全連接及ReLU激活函數構成，計算公式如下：

這裏的全連接是Position-wise逐位置的，即設前面的attention輸出的維度爲 $B * Length * d_{model}$ ，則變換時，實際上是隻針對 $d_{model}$ 進行變換，對於每個位置（Length維度）上，都使用同樣的變換矩陣。

在論文中，這裏的 $d_{model}$ 仍然是512，兩層全連接的中間隱層單元數爲 $d_{ff} = 2048$ 。

4. add & norm

在整體架構圖中，還有一個部分是add&norm，這其實是借鑑了圖像中的殘差思想。在self-attention和feed forward計算之後都會加上一個殘差變換，同時也會加上Layer Normalization（參見： https://arxiv.org/pdf/1607.06450.pdf ，用在有循環機制的網絡裏面效果較好）。設輸入爲 $x$ ，則輸出爲 $LayerNorm(x+SubLayer(x))$ ，這裏的 $SubLayer$ 即是self-attention或feed forward層。

5. 解碼器

接着來看Decoder部分（右半部分），它同樣也是由N層（在論文中，仍取 $N = 6$ ）堆疊起來，對於其中的每一層，除了與Encoder中相同的self-attention及Feed Forward之外，還在中間插入了一層傳統encoder-decoder框架中的attention層，即將decoder的輸出作爲query去查詢encoder的輸出，同樣用的是multi-head attention，使得在decode的時候能看到encoder的所有輸出。

同時，作爲decoder，在預測當前步的時候，是不能知道後面的內容的，即attention需要加上mask，將當前步之後的分數全部置爲 $-\infty$ ，然後再計算softmax，以防止發生數據泄露。

6. 位置編碼層

細心的讀者可能發現了，在整體架構圖中，還有一個叫Positional Encoding的東西，這是個啥？

Transformer雖然摒棄了RNN的循環結構和CNN的局部相關性，但對於序列來說，最重要的其實還是先後順序。看前面self-attention的處理方式，實際上與“詞袋”模型沒什麼區別，這樣忽略了位置信息的缺陷肯定是要通過一定的手段來彌補。

論文中提出了一個非常“smart”的方式來加入位置信息，就是這裏的Positional Encoding，它對於每個位置 $pos$ 進行編碼，然後與相應位置的word embedding進行相加，構成當前位置的新word embedding。它採用如下的公式爲每個 $pos$ 進行編碼：

其中， $i$ 表示embedding向量中的位置，即 $d_{model}$ 中的每一維。選擇這種sin函數有兩種好處：1）可以不用訓練，直接編碼即可，而且不管什麼長度，都能直接得到編碼結果；2）能表示相對位置，根據 $sin(\alpha+\beta)=sin\alpha cos\beta + cos\alpha sin\beta$ ， $PE_{pos+k}$ 可以表示爲 $PE_{pos}$ 的線性變換，這爲表達相對位置信息提供了可能性。

7. 一些Tricks和技巧

Embedding和Softmax：論文中將embedding層的參數與最後的Softmax層之前的變換層參數進行了共享（參見：https://arxiv.org/pdf/1608.05859.pdf ），並且在embedding層，將嵌入的結果乘上 $\sqrt{d_{model}}$ 。
初始化：看代碼裏面的初始化方式採用的是PyTorch裏面的nn.xavier_uniform，不知道是不是必須的，這個還是要具體問題具體嘗試？
優化器：論文裏面提到了他們用的優化器，是以 $\beta_1=0.9$ ， $\beta_2=0.98$ 和 $\epsilon=10^{-9}$ 的Adam爲基礎，而後使用一種warmup的學習率調整方式來進行調節。具體公式如下：基本上就是先用一個固定的warmup_steps進行學習率的線性增長（熱身），而後到達warmup_steps之後會隨着step_num的增長而逐漸減小，他們用的 $warmup\_steps = 4000$ ，這個可以針對不同的問題自己嘗試。
$l_{rate} = d_{model}^{-0.5}·min(step\_num^{-0.5}, step\_num · warmup\_steps^{-1.5})$
正則化：論文在訓練的時候採用了兩種正則化的方式。1）dropout：主要用在每個SubLayer計算結束之後，比如self-attention或feed forward，然後再與輸入進行add & norm，同時也作用在經過了位置編碼後的embedding上，他們取的 $P_{drop}=0.1$ ；2）標籤平滑：即Label Smoothing（參見： https://arxiv.org/pdf/1512.00567.pdf ），其實還是從圖像上搬過來的，具體操作可以看下一節的代碼實現。這裏論文取的 $\epsilon_{ls}=0.1$ ，他們發現會損失困惑度，但能提升準確率和BLEU值！

四. PyTorch實現

對於PyTorch實現部分，主要參考的是 http://nlp.seas.harvard.edu/2018/04/03/attention.html 。這裏將針對核心部分進行剖析和解讀：

1. 編解碼器

這是一個通用的Encoder-Decoder架構：

class EncoderDecoder(nn.Module):
    """
    A standard Encoder-Decoder architecture. Base for this and many 
    other models.
    """
    def __init__(self, encoder, decoder, src_embed, tgt_embed, generator):
        super(EncoderDecoder, self).__init__()
        self.encoder = encoder
        self.decoder = decoder
        self.src_embed = src_embed
        self.tgt_embed = tgt_embed
        self.generator = generator
        
    def forward(self, src, tgt, src_mask, tgt_mask):
        "Take in and process masked src and target sequences."
        return self.decode(self.encode(src, src_mask), src_mask,
                            tgt, tgt_mask)
    
    def encode(self, src, src_mask):
        return self.encoder(self.src_embed(src), src_mask)
    
    def decode(self, memory, src_mask, tgt, tgt_mask):
        return self.decoder(self.tgt_embed(tgt), memory, src_mask, tgt_mask)

要注意的是，src和tgt都需要傳入mask進行計算。

2. 多頭注意力機制

multi-head attention可用於三個地方，分別是Encoder和Decoder中各自的self-attention部分，還有Encoder-Decoder之間的attention部分。但其實這三個地方的不同僅僅在於query、key、value和mask的不同，因此當將這4部分作爲參數傳入時，模型的計算方式便可抽象爲如下的形式：

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)  # (3 + 1)
        self.attn = None
        self.dropout = nn.Dropout(p=dropout)
        
    def forward(self, query, key, value, mask=None):
        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)

參數裏面的h和d_model分別表示注意力頭的個數，以及模型的隱層單元數。注意到在__ini__函數中，定義了self.linears = clones(nn.Linear(d_model, d_model), 4)，clone(x, N)即爲深拷貝N份，這裏定義了4個全連接函數，實際上是3+1，其中的3個分別是Q、K和V的變換矩陣，最後一個是用於最後將多頭concat之後進行變換的矩陣。

在forward函數中，是首先將query、key和value進行相應的變換，然後需要經過attention這個函數的計算，這個函數實際上就是“Scaled Dot Product Attention”這個模塊的計算，如下所示：（注意這裏面的mask方式）

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

3. 前饋網絡

前面說了，前饋網絡實際上就是兩層全連接，其代碼如下：

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))))

4. add & norm

這是殘差模塊+LayerNormalization的實現方式：

class SublayerConnection(nn.Module):
    """
    A residual connection followed by a layer norm.
    Note for code simplicity the norm is first as opposed to last.
    """
    def __init__(self, size, dropout):
        super(SublayerConnection, self).__init__()
        self.norm = LayerNorm(size)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x, sublayer):
        "Apply residual connection to any sublayer with the same size."
        return x + self.dropout(sublayer(self.norm(x)))

在forward函數裏面，彷佛與前面的 $LayerNorm(x+SubLayer(x))$ 不太一樣，其實這裏都可以的，主要是看任務，自己實驗。

下面是LayerNormalization的實現，其實PyTorch裏面已經集成好了nn.LayerNorm，這裏列出來只是方便讀者看清其原理，爲了代碼簡潔，可以直接使用PyTorch裏面實現好的函數。

class LayerNorm(nn.Module):
    "Construct a layernorm module (See citation for details)."
    def __init__(self, features, eps=1e-6):
        super(LayerNorm, self).__init__()
        self.a_2 = nn.Parameter(torch.ones(features))
        self.b_2 = nn.Parameter(torch.zeros(features))
        self.eps = eps

    def forward(self, x):
        mean = x.mean(-1, keepdim=True)
        std = x.std(-1, keepdim=True)
        return self.a_2 * (x - mean) / (std + self.eps) + self.b_2

5. 位置編碼

位置編碼相關的代碼如下所示：

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)

可見，這裏首先是按照最大長度max_len生成一個位置，而後根據公式計算出所有的向量，在forward函數中根據長度取用即可，非常方便。注意要設置requires_grad=False，因其不參與訓練。

6. 關於mask

在Transformer裏面，Encoder和Decoder的attention計算都需要相應的mask處理，但功能卻不同。在Encoder中，mask主要是爲了讓那些在一個batch中長度較短的序列的padding不參與attention的計算，而在Decoder中，還要考慮不能發生數據泄露。那這些具體是怎麼實現的呢？看下面的代碼：

class Batch:
    "Object for holding a batch of data with mask during training."
    def __init__(self, src, trg=None, pad=0):
        self.src = src
        self.src_mask = (src != pad).unsqueeze(-2)
        if trg is not None:
            self.trg = trg[:, :-1]
            self.trg_y = trg[:, 1:]
            self.trg_mask = \
                self.make_std_mask(self.trg, pad)
            self.ntokens = (self.trg_y != pad).data.sum()
    
    @staticmethod
    def make_std_mask(tgt, pad):
        "Create a mask to hide padding and future words."
        tgt_mask = (tgt != pad).unsqueeze(-2)
        tgt_mask = tgt_mask & Variable(
            subsequent_mask(tgt.size(-1)).type_as(tgt_mask.data))
        return tgt_mask

對於src的mask方式就比較簡單，直接把pad給mask掉即可。對於trg的mask計算略微複雜一些，不僅需要把pad給mask掉，還需要進行一個subsequent_mask的操作，其代碼如下：

def subsequent_mask(size):
    "Mask out subsequent positions."
    attn_shape = (1, size, size)
    subsequent_mask = np.triu(np.ones(attn_shape), k=1).astype('uint8')
    return torch.from_numpy(subsequent_mask) == 0

這裏是給定一個序列長度size，生成一個下三角矩陣，在主對角線右上的都是false，其示意圖如下：

經過&得到的mask即爲最終trg需要的mask。

7. 優化器

這裏其實是手動實現了上一節提到的帶warmup的學習率調節公式，代碼比較簡單：

class NoamOpt:
    "Optim wrapper that implements rate."
    def __init__(self, model_size, factor, warmup, optimizer):
        self.optimizer = optimizer
        self._step = 0
        self.warmup = warmup
        self.factor = factor
        self.model_size = model_size
        self._rate = 0
        
    def step(self):
        "Update parameters and rate"
        self._step += 1
        rate = self.rate()
        for p in self.optimizer.param_groups:
            p['lr'] = rate
        self._rate = rate
        self.optimizer.step()
        
    def rate(self, step = None):
        "Implement `lrate` above"
        if step is None:
            step = self._step
        return self.factor * \
            (self.model_size ** (-0.5) *
            min(step ** (-0.5), step * self.warmup ** (-1.5)))
        
def get_std_opt(model):
    return NoamOpt(model.src_embed[0].d_model, 2, 4000,
            torch.optim.Adam(model.parameters(), lr=0, betas=(0.9, 0.98), eps=1e-9))

主要調節是在rate這個函數中，model_size即爲 $d_{model}$ ，warmup即爲 $warmup\_steps$ ，factor可以理解爲初始的學習率。

8. 標籤平滑

Label Smoothing這裏，我看到的PyTorch版本是用KL散度損失，對於輸出的分佈，從原始的one-hot分佈轉爲在groundtruth上使用一個confidence值，而後其他的所有非groudtruth標籤上採用 $\frac{1 - confidence}{odim - 1}$ 作爲概率值進行平滑。具體代碼如下：

class LabelSmoothing(nn.Module):
    "Implement label smoothing."
    def __init__(self, size, padding_idx, smoothing=0.0):
        super(LabelSmoothing, self).__init__()
        self.criterion = nn.KLDivLoss(size_average=False)
        self.padding_idx = padding_idx
        self.confidence = 1.0 - smoothing
        self.smoothing = smoothing
        self.size = size
        self.true_dist = None
        
    def forward(self, x, target):
        assert x.size(1) == self.size
        true_dist = x.data.clone()
        true_dist.fill_(self.smoothing / (self.size - 2))
        true_dist.scatter_(1, target.data.unsqueeze(1), self.confidence)
        true_dist[:, self.padding_idx] = 0
        mask = torch.nonzero(target.data == self.padding_idx)
        if mask.dim() > 0:
            true_dist.index_fill_(0, mask.squeeze(), 0.0)
        self.true_dist = true_dist
        return self.criterion(x, Variable(true_dist, requires_grad=False))

這裏的size是輸出詞表的大小，smoothing是用於分攤在非groundtruth上面的概率值。

五. 爲啥要用Self-Attention？

論文專門開了一個章節來闡釋爲什麼選用self-attention這種方式來代替RNN和CNN，我這也開一個章節專門講一下吧，以示尊重。

論文從計算複雜度、序列操作數以及最大路徑長度三個角度比較了不同的層，包括Self-Attention、Recurrent、Convolutional等，如下表所示：

表裏面的 $n$ 代表序列長度， $d$ 代表向量維度， $k$ 表示kernel的大小， $r$ 表示受限的memory的長度（主要是針對過長序列，直接使用self-attention未免太大）。從表中的數據看起來，好像Self-Attention確實比較優良。

六. 總結

模型特點：採用全attention的方式，完全摒棄了RNN和CNN的做法。
優勢：訓練速度更快，在兩個翻譯任務上取得了SoTA。
不足：在decode階段還是自迴歸的，即還是不能並行，而且對於每個step的計算，都是要重新算一遍，沒有前面的記憶。

傳送門

論文：https://arxiv.org/pdf/1706.03762.pdf
源碼：https://github.com/tensorflow/tensor2tensor （TensorFlow）
https://github.com/OpenNMT/OpenNMT-py （PyTorch）
https://github.com/awslabs/sockeye (MXNet)
參考：https://jalammar.github.io/illustrated-transformer 一個優質的英文博客，有很好的可視化圖例，適合不進行原理深究或只關注實現的入門級博客。其後面還有很多好的資源可以用來參考，一些是googleblog，還有視頻等，可以收藏後慢慢研讀！
http://nlp.seas.harvard.edu/2018/04/03/attention.html PyTorch實現的核心源碼博客，有原理，也有對應的代碼段，非常適合對照學習！

Transformer解讀（論文 + PyTorch源碼）

文章目錄

前言

一. 背景

二. Transformer整體架構

三. Transformer細節剖析

1. 編碼器

2. 自注意力機制

3. 前饋網絡

4. add & norm

5. 解碼器

6. 位置編碼層

7. 一些Tricks和技巧

四. PyTorch實現

1. 編解碼器

2. 多頭注意力機制

3. 前饋網絡

4. add & norm

5. 位置編碼

6. 關於mask

7. 優化器

8. 標籤平滑

五. 爲啥要用Self-Attention？

六. 總結

傳送門

【SQL進階】CASE語句的使用

npm error Cannot read properties of null (reading 'isDescendantOf')

《推薦系統實踐》算法純享（附代碼鏈接）（二）—— 協同過濾篇

XLM解讀（論文 + PyTorch源碼）

Transformer解讀（論文 + PyTorch源碼）

《推薦系統實踐》算法純享（附代碼鏈接）（一）—— 評價指標篇

Transformer-XL解讀（論文 + PyTorch源碼）

https://yachay.unat.edu.pe/blog/index.php?comment_area=format_blog&comment_component=blog&comment_co

linux以太網驅動總結