深入理解transformer源码
參考資料:
transformer原論文
深入理解transformer及源碼
圖解Transformer(完整版)
The Annotated Transformer
The Annotated Transformer的中文注釋版
前言
??本來今年(2020)七月份就想寫這篇博客,因為畢業、工作入職等一系列事情一直拖到了現在(主要是因為懶)。transformer的實現不止一個版本的源碼,本文主要講解哈佛大學利用torch實現的版本。相對更“高級”的源碼,這個用框架實現的版本顯得不是那么底層,但我們大多數人閱讀源碼的主要目的還是為了更多的理解這個算法的細節然后去更好的使用它。
??對transformer的入門級理解我個人認為可以分為三個層次,第一個層次是僅閱讀了原論文,對transformer有了宏觀的認識。讀過后你會知道transformer的宏觀結構、創新點和牛逼之處。但這遠遠不夠,因為transformer的論文省略了非常多的細節;第二個層次是讀過那篇傳播很廣的博客——圖解transformer,這篇博客以圖解的方式深入淺出的從微觀角度講解了什么是自注意力機制、多頭等概念,讀過后結合原論文你能大致理順tranformer的前向傳播機制,每一個解構開來的功能單元都能大致理解它的原理,此時你會有豁然開朗的感覺,但細想一段時間之后又會有很多疑惑,因為很多細節還是只知道模糊的原理,不知細節;第三個層次便是閱讀源碼了,如果讀的足夠認真,或者自己跟著實現一遍,你會知道transformer的前向傳播中矩陣是怎么拼接、變換的,理解設計精巧的mask機制是怎么與損失函數協同工作的,這時候你才算真正入門了這個算法。
??本文重點在于理解transformer的源碼,在源碼批注中寫清楚了每一步計算前后矩陣的size,這樣可以保證充分理解每一步數學運算的結果。但本文不是從零開始介紹transformer,且強烈依賴“圖解transformer”中的大量圖例,所以強烈推薦仔細讀完原論文和“圖解transformer”(開頭參考資料有鏈接)后再來閱讀,沒有基礎直接閱讀會引起強烈不適。
transformer介紹
??雖說已經推薦了大家有基礎后再閱讀,但這里還是簡單的介紹一下。在深度學習領域一直有兩種主流的特征抽取器結構,一種是CNN,一種是RNN。這兩種特征抽取器以其獨特的優勢稱霸了很久。尤其是RNN,其循環結構與nlp天然契合,在nlp中以RNN為基礎的各種算法一直都是舞臺的主角。transformer的出現打破了這個格局,transformer獨特的自注意力機制既與自然語言任務契合,又解決了RNN一直被詬病的并行困難、梯度消失等問題,可謂是從各方面革了RNN的命。當然,transformer優異的性能并不僅僅是自注意力機制帶來的,具體我就不展開聊了,感興趣的推薦看看《放棄幻想,全面擁抱Transformer:自然語言處理三大特征抽取器(CNN/RNN/TF)比較》,這篇類似綜述一樣的知乎文章可以說講解的很清晰了。
fig.1
??閑話就扯到這兒,下面我們就來進入transformer的源碼解讀。transformer的結構也是一個編碼-解碼結構。輸入序列先進行Embedding,經過Encoder之后結合上一次output再輸入Decoder,最后用softmax計算序列下一個單詞的概率。下面以前向傳播的順序解讀各個模塊的源碼。
1、Word Embedding and Positional Encoding
fig.2
??文本進入到模型需要編碼,transformer的編碼方式包括Word Embedding和Positional Encoding。
??word embedding一般有兩種方式:
??1、使用預訓練模型直接生成,相當于提供一個lookup table;
??2、 隨機初始化,在模型訓練過程中作為一個可訓練的參數參與梯度優化,這樣在模型訓練完成之時詞向量也會隨之生成。
??這兩種方法沒有孰優孰劣,只有適用場景。若針對一個特定的nlp任務,其訓練集詞庫有限,即預測過程中可能出現詞庫外的詞語,此時顯然應該借助預訓練模型。但拋開一個特定的nlp任務,若一個算法模型可以作為預訓練模型,那word embedding層必然不能少。transformer作為一個算法而非任務顯然屬于后者,word embedding的代碼實現如下:
class Embeddings(nn.Module):def __init__(self, d_model, vocab):""":param d_model: word embedding維度:param vocab: 語料庫詞的數量"""super(Embeddings, self).__init__()self.lut = nn.Embedding(vocab, d_model)self.d_model = d_modeldef forward(self, x):""":param x: 一個batch的輸入,size = [batch, L], L為batch中最長句子長度"""return self.lut(x) * math.sqrt(self.d_model) #這里乘了一個權重,不改變維度. size = [batch, L, d_model]??d_model是此向量的維度,理論上輸入以后就是個定值了,那么最后乘的權重math.sqrt(self.d_model)顯然也是個常數。所有隨機向量乘一個常數,相當于同時放縮,理論上不會有什么影響。所以為什么要乘個常數?先記住這個問題。
??transformer沒有RNN的循環結構,無法感知一個句子中詞語出現的先后順序,而詞語的位置是相當重要的一個信息。作者提出了位置編碼,即Positional Encoding,來解決這個問題。
??關于位置編碼,作者提供了兩種方法:
??1、通過訓練學習positional encoding 向量;
??2、編了個公式來計算positional encoding向量;
??作者經過試驗后發現兩種方式的結果是相似的,所以選擇了第二種。畢竟要簡單一點,減少了訓練參數,而且在訓練集中么有出現過的句子長度上也能用。公式如下:
PE(pos,2i)=sin?(pos/100002i/dmodel)P E_{(p o s, 2 i)}=\sin \left(p o s / 10000^{2 i / d_{\mathrm{model}}}\right)PE(pos,2i)?=sin(pos/100002i/dmodel?) PE(pos,2i+1)=cos?(pos/100002i/dmodel?)P E_{(p o s, 2 i+1)}=\cos \left(p o s / 10000^{2 i / d_{\text {model }}}\right)PE(pos,2i+1)?=cos(pos/100002i/dmodel??)
其中,pospospos指的是這個word在這個句子中的位置;2i2i2i指的是embedding詞向量的偶數維度,2i+12i+12i+1指的是embedding詞向量的奇數維度。具體代碼如下:
??為什么上面的兩個公式能體現單詞的相對位置信息呢?作者也舉了一個例子來說明,寫一段代碼取位置編碼向量的四個維度進行可視化:
plt.figure(figsize=(15, 5)) pe = PositionalEncoding(20, 0) y = pe.forward(Variable(torch.zeros(1, 100, 20))) plt.plot(np.arange(100), y[0, :, 4:8].data.numpy()) plt.legend(["dim %d"%p for p in [4,5,6,7]])fig.3
??從圖中可以看出,位置編碼基于不同位置添加了正弦波,對于每個維度,波的頻率和偏移都有不同。也就是說對于序列中不同位置的單詞,對應不同的正余弦波,可以認為他們有相對關系。
??代碼中的構造函數是Positional Encoding公式的代碼實現,但從forward函數中可以看出,詞語最終的編碼是word_embedding + positional_embedding,所以positional_embedding的維度與word_embedding一致。此時就直到為什么word_embedding末尾要乘一個常量權重了,就是為了放大,防止因與positional_embedding的量級相差過大,做加法后被忽視掉。
2 Encoder
??一個batch的數據經過embedding后,就要進入transformer的編碼器了。transformer的編碼器主要包含幾個功能點:多頭自注意力機制、求和與歸一化、前饋神經網絡。
fig.4
??其中“多頭自注意力機制”和“mask掩碼機制”是頗具創新的兩個亮點,也是后續重點關注的部分。
MultiHeadedAttention
??多頭自注意力機制就是多個自注意力機制。多頭注意力機制的源碼如下:
def clones(module, N):"工具人函數,定義N個相同的模塊"return nn.ModuleList([copy.deepcopy(module) for _ in range(N)])class MultiHeadedAttention(nn.Module):def __init__(self, h, d_model, dropout):"""實現多頭注意力機制:param h: 頭數:param d_model: word embedding維度:param dropout: drop out"""super(MultiHeadedAttention, self).__init__()assert d_model % h == 0 #檢測word embedding維度是否能被h整除# We assume d_v always equals d_kself.d_k = d_model // hself.h = h # 頭的個數self.linears = clones(nn.Linear(d_model, d_model), 4) #四個線性變換,前三個為QKV三個變換矩陣,最后一個用于attention后self.attn = Noneself.dropout = nn.Dropout(p=dropout)def forward(self, query, key, value, mask=None):""":param query: 輸入x,即(word embedding+postional embedding),size=[batch, L, d_model] tips:編解碼器輸入的L可能不同:param key: 同上,size同上:param value: 同上,size同上:param mask: 掩碼矩陣,編碼器mask的size = [batch , 1 , src_L],解碼器mask的size = [batch, tgt_L, tgt_L]"""if mask is not None:# 在"頭"的位置增加維度,意為對所有頭執行相同的mask操作mask = mask.unsqueeze(1) # 編碼器mask的size = [batch,1,1,src_L],解碼器mask的size= = [batch,1,tgt_L,tgt_L]nbatches = query.size(0) # 獲取batch的值,nbatches = batch# 1) 利用三個全連接算出QKV向量,再維度變換 [batch,L,d_model] ----> [batch , h , L , d_model//h]query, key, value = \[l(x).view(nbatches, -1, self.h, self.d_k).transpose(1, 2) # view中給-1可以推測這個位置的維度for l, x in zip(self.linears, (query, key, value))]# 2) 實現Scaled Dot-Product Attention。x的size = (batch,h,L,d_model//h),attn的size = (batch,h,L,L)x, self.attn = attention(query, key, value, mask=mask,dropout=self.dropout)# 3) 這步實現拼接。transpose的結果 size = (batch , L , h , d_model//h),view的結果size = (batch , L , d_model)x = x.transpose(1, 2).contiguous().view(nbatches, -1, self.h * self.d_k)return self.linears[-1](x) # size = (batch , L , d_model)??clones函數不用多說,就是一個工具人函數,負責定義N個相同的模塊。
??MultiHeadedAttention類即實現多頭注意力機制的類。其構造函數中首先檢測word embedding的維度能否被頭數h整除,其原因是因為必須整除才能講word embedding均分到每個頭上。說到這里可能有讀者會疑惑,因為與圖解transformer中的實現方式不同。圖解transformer中即=每個word embedding是不變的,變得是WqW_qWq?,WkW_kWk?,WvW_vWv?三個權值矩陣,每一組不同的權值矩陣與word embedding相乘后會得到不同的Query、Key和Value矩陣,即對應不同的頭。也就是說,圖解transformer中,多頭的實現是靠多個權值矩陣實現的。其實這里的實現方式其實也是與之等效的,后面會具體講到。
構造函數中除了定義了一些后面用到的變量之外,還定義了四個相同的線性變換。這里說的相同,不是指完全相同,是指輸入維度和輸出維度是相同的,其權值是不同的。這里的線性變換切勿與前饋神經網絡混淆,如果加了激活函數(即引入非線性)才是前饋神經網絡。四個線性變換中,前三個代表與權值矩陣WqW_qWq?,WkW_kWk?,WvW_vWv?相乘的過程,對應圖解transformer中的下圖。
fig.5
??最后一個線性變化表示attention求解的結果拼接后的線性變化,用來約束結果的維度。對應圖解transformer的下圖。而約束維度的目的就是為了講維度控制至與輸入的word embedding一致,方便實現下一步的“求和”操作。
fig.6
??forward函數中,我們首先關注第35行,這里是利用前三個線性變換,與word embedding相乘,得到,對應fig.5中的操作。仔細的讀者可能會發現,這個forward函數的輸入參數是三個,即query、key、value,為什么又說三個線性變換都與word embedding相乘呢?原因是,多頭注意力機制,它的Query、Key、Value三個矩陣并不一定全是同樣的輸入word embedding與WqW_qWq?,WkW_kWk?,WvW_vWv?相乘得到的。編碼器中要求解每個詞與其他詞之間的attention值,也就是“自注意力”,所以輸入全是word embedding。在解碼器中就不是如此,所以這里索性留出三個參數以便函數在多種情況下的復用。
??然而“自注意力”的訓練過程中,并不是一條一條數據進去的,而是一個batch的數據以矩陣的形式進入。既然是矩陣運算這里就涉及了一個問題,一個batch的句子長度并不是一致的,在進入之前會有一個pad token將句子補齊,全部補齊成當前batch中最長句子的長度,pad token也會有自己的word embedding。這樣在求解的過程中,詞語與pad token之間的attention值就變成了多余的東西,對后續計算一定會有影響,這里就是mask機制發揮作用的地方。其實這里還有一個問題,為什么是補齊成當前batch的最長長度,而不是所有樣本最長長度?原因是batch之間長度L的不同并不影響模型權值的數量,權值是作用在詞編碼的維度上的。先帶著這個答案繼續往下看。
??第35行,求Query、Key、Value三個矩陣的過程可以表示成下圖。
fig.7
??fig.7雖然只示意了Query矩陣的求解過程和size,但其實Key和Value矩陣都是一個道理。fig.7對應第35行代碼,其中紅色的部分表示pad token的word embedding,可以看出,在線性變換的過程中,pad token也被帶到了結果中,并且對應的位置沒變。第35行代碼除了線性變換,還要將變換后結果的size從 [batch,L,d_model] ----> [batch , h , L , d_model//h],其實就是將結果的d_model均分成h份,來代表多頭。這與“圖解transformer”中的利用多個權值矩陣來生成結果其實是一樣的,因為transformer并沒有規定權值矩陣的size必須是(d_model,d_model),只要是(d_model,?)能保證矩陣相乘即可。在這套源碼中,為了便利將線性變換的輸入輸出維度都設置成了d_model,將結果h份均分,其實等價于h個size為(d_model, d_model/h)的權值矩陣與輸入進行了線性相乘,即h個頭。這也是為什么構造函數中檢驗能否整除的原因。
??得到Query、Key、Value三個矩陣后,就到了正式計算attention的時候了,論文中稱之為“Scaled Dot-Product Attention”(fig.8),對應代碼中的第40行。
fig.8
??對應公式:Attention?(Q,K,V)=softmax?(QKTdk)V\text { Attention }(Q, K, V)=\operatorname{softmax}\left(\frac{Q K^{T}}{\sqrt{d_{k}}}\right) V?Attention?(Q,K,V)=softmax(dk??QKT?)V
??第40行代碼引用的函數源碼如下:
??attention函數中第11行即計算Query、Key矩陣相乘并縮放,縮放的目的是作者認為過大的值會影響softmax函數,將其推入一個梯度很小的空間。具體為什么會這樣,我也不清楚…這步對應的圖解如圖fig.9
fig.9
??圖9中忽略了縮放,因為縮放不影響矩陣的形狀。從圖9中我們可以看到,結果為L*L的一個方陣,表示這個句子中,每個詞都與其他詞打了個照面。對于pad token,最終傳播到結果中,行列都有,列表示別的詞與pad token的結果,行表示pad token與別的詞的結果。按步驟下一步應該softmax了,但圖中紅色部分是與pad token相關的,理論上是不應參與softmax的,所以需要一點手段將其屏蔽,即mask矩陣。觀察attention函數的第14行,mask矩陣的作用是將圖中紅色部分替換一個極小的數字,這樣就使得softmax幾乎忽略掉它。生成mask矩陣的代碼如下:
class Batch:def __init__(self, src, trg=None, pad=0):""":param src: 一個batch的輸入,size = [batch, src_L]:param trg: 一個batch的輸出,size = [batch, tgt_L]"""self.src = srcself.src_mask = (src != pad).unsqueeze(-2) #返回一個true/false矩陣,size = [batch , 1 , src_L]if trg is not None:self.trg = trg[:, :-1] # 用于輸入模型,不帶末尾的<eos>self.trg_y = trg[:, 1:] # 用于計算損失函數,不帶起始的<sos>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):""":param tgt: 一個batch的target,size = [batch, tgt_L]:param pad: 用于padding的值,一般為0:return: mask, size = [batch, tgt_L, tgt_L]"""tgt_mask = (tgt != pad).unsqueeze(-2) # 返回一個true/false矩陣,size = [batch , 1 , tgt_L]tgt_mask = tgt_mask & Variable(subsequent_mask(tgt.size(-1)).type_as(tgt_mask.data)) # 兩個mask求和得到最終mask,[batch, 1, L]&[1, size, size]=[batch,tgt_L,tgt_L]return tgt_mask # size = [batch, tgt_L, tgt_L]def subsequent_mask(size):""":param size: 輸出的序列長度:return: 返回下三角矩陣,size = [1, size, size]"""attn_shape = (1, size, size)subsequent_mask = np.triu(np.ones(attn_shape), k=1).astype('uint8') #返回上三角矩陣,不帶軸線return torch.from_numpy(subsequent_mask) == 0 #返回==0的部分,其實就是下三角矩陣??Batch類中生成了兩種mask,一種用于編碼器,一種用于解碼器。這里我們暫時僅需要關注第一種,在Batch類構造函數的第8行。這里的mask的size為[batch , 1 , src_L],表示針對batch中的每一條分詞后的句子,都有一個形為[1,src_L]布爾矩陣與每個詞一一對應,這里的src_L即本batch輸入的最長句子長度。mask矩陣要處理的矩陣形狀為[L,L],算上多頭,形狀為[h,L,L],再算上batch,mask矩陣最終要遮蓋的矩陣形狀為[batch,h,L,L],那么目前mask矩陣的形狀肯定還是對應不上的?,F在回到MultiHeadedAttention代碼塊的第29行,這里增加了一個維度,值為1,意為在這個維度上執行相同的操作,即對每個頭執行相同的掩碼操作。
掩碼的最細粒度,就是利用size為[1,src_L]的mask對一個size為[L,L]的矩陣進行掩碼,mask中的維度1意為在此維度上進行相同的操作。具體過程可以看下圖。
fig.10
??執行完mask就該softmax了,對應attention函數代碼塊的第15行,由于每行表示同一個詞與其他詞“打照面”的結果,所以對行進行softmax。圖11中的softmax,每行可以看作一個詞自己與其他所有詞的權重。比如第一行的L個權重,歸屬第一個詞,意為第一個詞與其他L個詞(包括自己)的“相關性“,當然這里的”相關性“只是比喻。與Value矩陣進行矩陣乘法,就將L個權重作用到了L個詞的Value值上。其中softmax矩陣右側紅色豎列是經過mask處理的,下端紅色的行為pad token引入,這兩者都是我們不考慮的。圖11的過程對應第18行,完成這步后至此完成Scaled Dot-Product Attention。
fig.11
??結果如圖10所示,一個批次中帶有pad token,那么最終會傳播到結果中,還是原位置。得到結果后就要執行圖6的拼接和線性變換的操作,對應MultiHeadedAttention代碼塊的第43-45行。最終h個頭拼接后的size為[batch,L,d_model],正好對應MultiHeadedAttention的18行定義的輸入輸出維度都為d_model的最后一個線性變換。所以看到這里可以明白,克隆4個線性變換、要求d_model整除h,都是作者刻意設計的結果,這些設計讓整個代碼更加簡潔。這部分結束后transformer最復雜的地方也就清楚了,下面是一些不那么創新的組件,但仍然發揮了重要的作用。
SublayerConnection and LayerNorm
自注意力求解完后就輪到了求和與歸一化。這部分原理比較簡單,直接上代碼。
class LayerNorm(nn.Module):def __init__(self, features, eps=1e-6):"""實現層歸一化"""super(LayerNorm, self).__init__()self.a_2 = nn.Parameter(torch.ones(features)) # 類似BN層原理,對歸一化后的結果進行線性偏移,feature=d_model,相當于每個embedding維度偏移不同的倍數self.b_2 = nn.Parameter(torch.zeros(features)) # 偏置。系數和偏置都為可訓練的量self.eps = eps # 保證歸一化分母不為0def forward(self, x):""":param x: 輸入size = (batch , L , d_model):return: 歸一化后的結果,size同上"""mean = x.mean(-1, keepdim=True) # 最后一個維度求均值std = x.std(-1, keepdim=True) # 最后一個維度求方差return self.a_2 * (x - mean) / (std + self.eps) + self.b_2 #歸一化并線性放縮+偏移class SublayerConnection(nn.Module):"""實現殘差連接"""def __init__(self, size, dropout):super(SublayerConnection, self).__init__()self.norm = LayerNorm(size) # 讀入層歸一化函數self.dropout = nn.Dropout(dropout)def forward(self, x, sublayer):""":param x: 當前子層的輸入,size = [batch , L , d_model]:param sublayer: 當前子層的前向傳播函數,指代多頭attention或前饋神經網絡"""return self.norm(x + self.dropout(sublayer(x))) #這里把歸一化已經封裝進來,size = [batch , L , d_model]??LayerNorm是一個層歸一化的函數,比較基礎就不多聊了。SublayerConnection是殘差連接,其構造函數把層歸一化集成了進來,forward函數中的sublayer指代一個傳播函數,它可以是多頭注意力機制也可以是前饋神經網絡,因為在transformer中這兩個部分之后都用到了求和與歸一化,這樣寫可以讓其通用。無論是歸一化還是殘差連接,他們都不會改變輸入的維度,所以輸出結果的size還是[batch,L,d_model]。
PositionwiseFeedForward
??完成求和與歸一化后,下面要經過一個特別的前饋神經網絡,稱為Position-wise Feed-Forward Networks,公式如下:
FFN(x)=max?(0,xW1+b1)W2+b2\mathrm{FFN}(x)=\max \left(0, x W_{1}+b_{1}\right) W_{2}+b_{2}FFN(x)=max(0,xW1?+b1?)W2?+b2?代碼如下:
??Position-wise Feed-Forward Networks由兩個線性變換和一個relu激活函數組成,構造函數中定義了兩個線性變換,第14行實現了公式。完成Position-wise Feed-Forward后仍然是一個“求和與歸一化”,直接復用上一小節的代碼即可。至此編碼器的組建就全部介紹完了,下一步利用這些組件組成編碼器層。
EncoderLayer
??這里之所以叫編碼器層而不叫編碼器,是因為編碼器可能含有N個編碼器層。編碼器的代碼如下:
class EncoderLayer(nn.Module):"""Encoder層整體的封裝,由self attention、殘差連接、歸一化和前饋神經網絡組成"""def __init__(self, size, self_attn, feed_forward, dropout):super(EncoderLayer, self).__init__()self.self_attn = self_attn #定義多頭注意力,即傳入一個MultiHeadedAttention類self.feed_forward = feed_forward #定義前饋,即傳入一個PositionwiseFeedForward類self.sublayer = clones(SublayerConnection(size, dropout), 2) #克隆兩個殘差連接,一個用于多頭attention后,一個用于前饋神經網絡后self.size = sizedef forward(self, x, mask):""":param x: 輸入x,即(word embedding+postional embedding),size = [batch, L, d_model]:param mask: 掩碼矩陣,編碼器mask的size = [batch , 1 , src_L],解碼器mask的size = [batch, tgt_L, tgt_L]:return: size = [batch, L, d_model]"""x = self.sublayer[0](x, lambda x: self.self_attn(x, x, x, mask)) #實現多頭attention和殘差連接,size = [batch , L , d_model]return self.sublayer[1](x, self.feed_forward) #實現前饋和殘差連接, size = [batch , L , d_model]??構造函數中定義了多頭注意力、前饋網絡和殘差連接,其中的size參數是word embedding的維度,即d_model。在forward函數中,由于SublayerConnection類中提前留了函數接口,所以按先后順序把多頭自注意力和前饋網絡放進去即可,這段代碼我認為是極其清晰的,也非常佩服作者的設計、整合能力。
看到這里我們注意到,在最開始我們的輸入size為[batch,L,d_model],這里最終輸出的size也是一樣,這為無限套娃多層累加提供了極大便利,也是transformer算法設計上的一個小妙處。
Encoder
??接下來就是真正的編碼器了,編碼器就是多層編碼器層。代碼如下:
class Encoder(nn.Module):"""Encoder最終封裝,由若干個Encoder Layer組成"""def __init__(self, layer, N):super(Encoder, self).__init__()self.layers = clones(layer, N) # N層Encoder Layerself.norm = LayerNorm(layer.size) # 最終輸出進行歸一化def forward(self, x, mask):for layer in self.layers:x = layer(x, mask)return self.norm(x)??編碼器就是多個編碼器層組成,這里還是非常清晰的,唯一有疑問的是,不知道作者最后為什么還要進行歸一化,按我的理解在最后一層編碼器中已經歸一化過了,不過這個倒是無傷大雅。
3 Decoder
??我們再把圖1搬過來,看一下解碼器的結構。
??解碼器的輸入是目標label的embedding,同樣包括word embedding和position embedding。編碼后的label經過一個帶有特殊mask操作的“多頭注意力機制”,然后再經過一個與編碼器結果有交互的“多頭注意力機制”與“求和歸一化”,最后經過“前饋神經網絡”。當然,這幾個模塊兩兩之間都夾雜了“求和與歸一化”。
??解碼器的輸入同樣是一個batch的label,包含了label的所有信息。但預測時,我們顯然是沒有輸入的,這里的處理方式就如同RNN一次樣,以SOS起始符為第一個詞輸入,下一個詞預測出來,再與當前輸入連接作為新的輸入。為了模擬這個過程,在解碼器的“多頭注意力機制”中采用了一種特殊的mask方式。接下來解讀一下解碼器解碼的整個過程。
Decoder Layer
??Decoder Layer的宏觀結構還是比較簡單的,就是將上文講過的幾個模塊堆疊起來,代碼如下:
class DecoderLayer(nn.Module):"解碼器由 self attention、編碼解碼self-attention、前饋神經網絡 組成"def __init__(self, size, self_attn, src_attn, feed_forward, dropout):super(DecoderLayer, self).__init__()self.size = size # embedding的維度self.self_attn = self_attnself.src_attn = src_attnself.feed_forward = feed_forwardself.sublayer = clones(SublayerConnection(size, dropout), 3) #克隆3個sublayer分別裝以上定義的三個部分def forward(self, x, memory, src_mask, tgt_mask):""":param x: target,size = [batch, tgt_L, d_model]:param memory: encoder的輸出,size = [batch, src_L, d_model]:param src_mask: 源數據的mask, size = [batch, 1, src_L]:param tgt_mask: 標簽的mask,size = [batch, tgt_L, tgt_L]"""m = memoryx = self.sublayer[0](x, lambda x: self.self_attn(x, x, x, tgt_mask)) # self-atten、add&norm,和編碼器一樣, size = [batch, tgt_L, d_model]x = self.sublayer[1](x, lambda x: self.src_attn(x, m, m, src_mask)) # 編碼解碼-attention、add&norm,Q來自target,KV來自encoder的輸出,size = [batch, tgt_L, d_model]return self.sublayer[2](x, self.feed_forward) # 前饋+add&norm, size = [batch, tgt_L, d_model]??Decoder Layer的構造函數相當簡單,就是傳入幾個模塊,forward函數也和編碼器的類似,這里就不展開細聊了,唯一不同的是,在第20行傳入的mask矩陣稍有不同,這里的mask矩陣同樣在Batch類中生成。為了讀者閱讀的方便,再次將Batch類的代碼搬運過來。
class Batch:def __init__(self, src, trg=None, pad=0):""":param src: 一個batch的輸入,size = [batch, src_L]:param trg: 一個batch的輸出,size = [batch, tgt_L]"""self.src = srcself.src_mask = (src != pad).unsqueeze(-2) #返回一個true/false矩陣,size = [batch , 1 , src_L]if trg is not None:self.trg = trg[:, :-1] # 用于輸入模型,不帶末尾的<eos>self.trg_y = trg[:, 1:] # 用于計算損失函數,不帶起始的<sos>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):""":param tgt: 一個batch的target,size = [batch, tgt_L]:param pad: 用于padding的值,一般為0:return: mask, size = [batch, tgt_L, tgt_L]"""tgt_mask = (tgt != pad).unsqueeze(-2) # 返回一個true/false矩陣,size = [batch , 1 , tgt_L]tgt_mask = tgt_mask & Variable(subsequent_mask(tgt.size(-1)).type_as(tgt_mask.data)) # 兩個mask求和得到最終mask,[batch, 1, L]&[1, size, size]=[batch,tgt_L,tgt_L]return tgt_mask # size = [batch, tgt_L, tgt_L]def subsequent_mask(size):""":param size: 輸出的序列長度:return: 返回下三角矩陣,size = [1, size, size]"""attn_shape = (1, size, size)subsequent_mask = np.triu(np.ones(attn_shape), k=1).astype('uint8') #返回上三角矩陣,不帶軸線return torch.from_numpy(subsequent_mask) == 0 #返回==0的部分,其實就是下三角矩陣??同樣的,label輸入的size也為[batch,tgt_L],label的mask由兩個部分生成。第一個部分是第21行,生成方式與輸入的mask系統,也是統計出哪些地方是pad token,size也為[batch,1,tgt_L]。第二部分是由subsequent_mask函數生成的一個size為[1,tgt_L,tgt_L]下三角矩陣,如下所示,其中0代表false,即在后面要掩住的地方。
[[[1, 0, 0, 0, 0, 0, 0, 0, 0, 0],[1, 1, 0, 0, 0, 0, 0, 0, 0, 0],[1, 1, 1, 0, 0, 0, 0, 0, 0, 0],[1, 1, 1, 1, 0, 0, 0, 0, 0, 0],[1, 1, 1, 1, 1, 0, 0, 0, 0, 0],[1, 1, 1, 1, 1, 1, 0, 0, 0, 0],[1, 1, 1, 1, 1, 1, 1, 0, 0, 0],[1, 1, 1, 1, 1, 1, 1, 1, 0, 0],[1, 1, 1, 1, 1, 1, 1, 1, 1, 0],[1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]] size = [1,tgt_L,tgt_L]??然后在代碼第23行對兩個mask矩陣進行了求和操作,可以理解為求并集。[batch, 1, L] & [1, size, size] = [batch,tgt_L,tgt_L],size中為1的部分,表示在對方此維度上進行“廣播”,也就是進行相同的操作。最終結果的size為[batch, tgt_L, tgt_L],也就是說對于batch中的每一個句子,都對應一個size為[tgt_L, tgt_L]的掩碼矩陣。
fig.12
??在編碼器的講解流程中,先講了傳播過程,在傳播過程中遇到mask需求時候,才自然引出mask的介紹。這里之所以沒有按照這個順序是因為,解碼器mask的設計不光針對transformer的訓練階段,與模型預測階段也息息相關甚至更相關。但不清楚訓練階段又無法理解測試階段,所以我們先陳述“現象”,帶著問題,當所有的“現象”陳述完畢之后,自然就能明白解碼器mask為何這樣設計。
??我們先來推導在訓練階段label的正向傳播。label經過embedding后,size與編碼器的輸入是一致的,也是[batch,tgt_L,d_model]。labe的第一個“多頭注意力”機制傳播如圖13。
fig.13
??圖13簡化了傳播流程,忽視了batch維度和多頭h的維度,僅在最小粒度表示一條label的傳播過程。其中Query、Key、Value矩陣生成的過程,完全與圖7一致,最終結果的size依然是[batch,tgt_L,d_model]。然后Query矩陣與Key矩陣相乘生成size為[batch,h,tgt_L,tgt_L]的矩陣。圖13在這里mask也進行了簡化,著重突出解碼器特有的下三角矩陣,而忽略了pad token進行掩碼的情況。
??從結果往前反推,結果的第一行,只受到了第一個詞與自身attention值的影響;第二行只受到了第二個詞與第一個詞的attention值、第二個詞與自身attention值的影響。舉個全網都在舉的例子,假設這里輸入的lael為[我,愛,機器,學習]四個詞語(忽略與),結果的第一行只與attention(我,我)有關系,后面的部分被mask掉了。第二行僅與attention(愛,我)、attention(愛,愛)有關系,其余部分被mask掉了,結果第三、四、五行等等也是同樣的道理。這么設計的原因,很多博客會說到這是在模擬預測時的情況,因為預測結果是迭代生成的,這是為了不讓模型“偷看”到未來的內容。emmm…其實這樣的解釋,也不能說錯,但終究還是沒說透徹。我們帶著問題先繼續往下看。
??經過解碼器第一個“自注意力機制”后是一個“求和與歸一化”步驟,對應這一步不影響矩陣的size,且與編碼器完全相同,就不展開講了。緊接著又是一個“自注意力機制”,對應DecoderLayer的第21行。前面在編碼器首次解釋“自注意力機制”代碼的時候,我就說過這份代碼之所以把輸入設計成三個,是因為Query、Key、Value矩陣并不一定是由一個東西線性變換生成的,這里就體現到了。解碼器這里的輸入為解碼器上一步的結果和編碼器最終的結果,其傳播流向如圖14。
fig.14
??這次的“自注意力機制”與編碼器就類似了,雖然輸入編碼器的結果引入了src_L(輸入句子的長度),但這并不影響整個傳播流程,這也得益于編碼器的mask矩陣是以廣播的形式對矩陣進行處理的。
??在解碼器第一個“自注意力機制”末尾我們講到,結果第一行只受到了第一個詞自身的attention值的影響,第二個詞只受到了與第一個詞的attention值、第二個詞與自身attention值的影響。反映到圖14中,圖14的結果,第一行只受到了解碼器上一步結果中第一行的影響,第二行只受到了第二行的影響,也就是說,經過兩個“自注意力機制”,解碼器特殊的mask矩陣帶來的影響,是無縫傳遞過來的。
Decoder
??解碼器層與編碼器層有一個相同的特點,那就是輸出的size與輸入的size會保持相同,這同樣為無限套娃帶來了方便。解碼器就是N個解碼器層組成的,代碼如下:
class Decoder(nn.Module):"解碼器的高層封裝,由N個Decoder layer組成"def __init__(self, layer, N):super(Decoder, self).__init__()self.layers = clones(layer, N)self.norm = LayerNorm(layer.size)def forward(self, x, memory, src_mask, tgt_mask):for layer in self.layers:x = layer(x, memory, src_mask, tgt_mask)return self.norm(x) # size = [batch, tgt_L, d_model]??經過N層解碼器層,輸入的size依然是[batch, tgt_L, d_model]。
Encoder-Decoder
??上面將編碼器與解碼器都介紹過了,整體的編碼解碼結構就是對兩者的封裝,代碼如下:
class EncoderDecoder(nn.Module):"""編碼解碼架構"""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 #源的embeddingself.tgt_embed = tgt_embed #目標的embeddingself.generator = generator # 定義最后的線性變換與softmaxdef forward(self, src, tgt, src_mask, tgt_mask): # 編碼解碼過程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)??經過編碼解碼之后,最后還有一個線性變換與softmax,這步是為了計算損失函數做準備,代碼如下:
class Generator(nn.Module):"""定義一個全連接層+softmax"""def __init__(self, d_model, vocab):super(Generator, self).__init__()self.proj = nn.Linear(d_model, vocab) # vocab為整個詞袋的詞數def forward(self, x):""":param x: 輸入的 size = [batch, tgt_L, d_model]"""return F.log_softmax(self.proj(x), dim=-1) #dim=-1在最后一維上做softmaxfig.15
??圖15表示了這個過程。vocab表示整個詞袋的詞數,將其線性變換為這個維度的原因就是為了在這個維度softmax,每一行會產生一組和為1的概率,共tgt_L行。例如,第一行會產生一組長為vocab的概率組,表示預測結果為詞袋中任意一次的概率。同樣的,第一行也會對應一個真實標簽的one-hot編碼。這樣總共tgt_L個預測結果,對應tgt_L個真實標簽的one-hot編碼,兩者做交叉熵,完成了損失函數的求解。
??還記得前面說過,結果的第一行只受第一個詞自身attention值的影響,第二個詞只受到了與第一個詞的attention值、第二個詞與自身attention值的影響。。。這也會反映到圖15最終的結果上。那么還是那個問題,為什么這么做模擬了預測過程?現在可以來解答這個問題。
??首先講一下transformer的預測過程。transformer雖然很復雜,但它仍然是一個編碼解碼的結構,其預測過程與其他編碼解碼結構的算法也是一致的。即首先輸入,預測出第一個詞,然后將預測值拼接回來,輸入[,預測結果1]去預測第二個詞,以此類推。在預測結果不斷返回來拼接然后再次輸入的過程中,有一個步驟是一直在重復經歷的——解碼器下三角矩陣mask機制。我們畫個圖來模擬一下這個過程。
fig.16
??圖16模擬了兩步,我們發現了兩個事實:
??1、輸入詞起始符和詞1的時候,仍然預測了詞1,只不過我們選擇性忽略,只抽取出詞2然后繼續拼接而已。
??2、在圖中第二步中,對應的attention,仍然只有自身與自身,這便是下三角矩陣掩碼的結果。
??那么為什么在預測階段還要進行下三角矩陣的掩碼?我們試想一下,如果不進行掩碼,在第二步中,的attention除了與自身,還有與詞1的,這就會導致在第二步中預測的詞1與第一步不同。遵循這個規律,第三步中預測的詞1、2與第二步中的也不同,而我們的做法是每次選擇性忽略重復的預測的詞,只摘取出最新預測的詞語拼接然后繼續預測,所以我們一定要保持每一步中重復預測的詞語是一致的,否則我們不斷拼接、不斷迭代預測將變成一種腦癱行為。
??在訓練階段,label全量只進去一次,不是一個一個進去的,如果想要在訓練階段進去一次就能模擬出預測階段多次預測、拼接的過程,將mask設計成下三角矩陣恰好能實現。
??綜上,總結一下解碼器mask設計思路的原因:1、預測階段要保持重復預測詞一致——>必須保持每步attention的值不變——>掩碼掉未來詞——>mask下三角矩陣;2、恰好也可以使模型在訓練階段的傳播過程與預測階段一致。
結尾
??至此,transformer的源碼就介紹完了,個人感覺寫得邏輯還算清晰,哪里有不明白或錯誤的地方歡迎留言。
總結
以上是生活随笔為你收集整理的深入理解transformer源码的全部內容,希望文章能夠幫你解決所遇到的問題。
- 上一篇: SQL SERVER 服务器名称的更改
- 下一篇: qq背景音乐外链