The Annotated Transformer注釋加量版，讀懂代碼就真的懂了Transformer 原創(chuàng)
本文是在The Annotated Transformer這篇文章基礎(chǔ)上的二次加工。
1.給代碼加了更詳細的注釋。
2.輸出詳細日志跟蹤數(shù)據(jù)。
原文地址：https://nlp.seas.harvard.edu/annotated-transformer/
或者后臺回復(fù)taf獲取pdf下載鏈接。
The Andnotated Transformer
Attention is All You Need
- v2022: Austin Huang, Suraj Subramanian, Jonathan Sum, Khalid Almubarak, and Stella Biderman.
- Original: Sasha Rush
閱讀方法
由于原文內(nèi)容過長，我沒有把原文拷貝過來，閱讀本文時，請打開原文鏈接或者我添加注釋的notebook。
1、給代碼加了更詳細的注釋。
原文基于pytorch從0開始復(fù)現(xiàn)了transformer模型，我在原文代碼基礎(chǔ)上追加了更詳細的注釋，代碼可以在下面鏈接找到。https://github.com/AIDajiangtang/annotated-transformer/blob/master/AnnotatedTransformer_comment.ipynb
另外，我還在模型結(jié)構(gòu)上加了注釋，我將代碼中重要的類名或者函數(shù)名標注在Transforner結(jié)構(gòu)的圖片上，閱讀代碼時請結(jié)合圖片上的名稱，這樣有助于快速理解代碼。
2、輸出日志跟蹤數(shù)據(jù)。
原文提供了一個訓(xùn)練德譯英模型的代碼，我在此基礎(chǔ)上加了一些日志，打印數(shù)據(jù)的維度來輔助對Transformer的理解。
我將按照圖片上標注數(shù)字順序來跟蹤數(shù)據(jù)。
原始論文中，Transformer是一種Encoder-Decoder架構(gòu)，左邊是Encoder，用于提取源語言的表征，右邊是Decoder，根據(jù)表征結(jié)合目標語言語法生成目標語言。
先從Encoder這邊開始。
0、Inputs：
假設(shè)batch size為2，所以每個batch包含兩個樣本，每個樣本由（德語，英語）文本對組成。
[
('Eine gro?e Gruppe Jugendlicher in einem kleinen Unterhaltungsbereich.', 'A large group of young adults are crammed into an area for entertainment.'),
('Zwei Arbeiter stellen Laternen auf.', 'Two workers working on putting up lanterns.')
]（batch size的意義：模型每次都是基于batch size個樣本的損失來更新參數(shù)，batch size需要根據(jù)內(nèi)存，顯存大小確定）
對于Encoder而言，它只需要源語言，也就是德語。
'Eine gro?e Gruppe Jugendlicher in einem kleinen Unterhaltungsbereich.'
'Zwei Arbeiter stellen Laternen auf'
1、Embedding：
1.1.先將文本轉(zhuǎn)換成tokens，并添加起始和結(jié)束符token。
(load_tokenizers函數(shù),
tokenize函數(shù),
build_vocabulary函數(shù)}
["<s>", "</s>", "<blank>", "<unk>"]起始符token id：0，結(jié)束符token id：1，padding token id：2
'Eine gro?e Gruppe Jugendlicher in einem kleinen Unterhaltungsbereich.'的tokens如下
torch.Size([11])
tensor([ 0, 14, 176, 38, 683, 7, 6, 116, 7147, 4, 1],
device='cuda:0')（通過結(jié)果看是基于詞的tokenization方法）
1.2.因為文本長度不一致，通過padding的方式將序列長度統(tǒng)一為72。
{collate_batch函數(shù)}
（padding不是必須的，只是出于方便和效率考慮，72是個經(jīng)驗值，通過對訓(xùn)練數(shù)據(jù)的統(tǒng)計得出）
'Eine gro?e Gruppe Jugendlicher in einem kleinen Unterhaltungsbereich.'padding后的tokens如下
torch.Size([72])
[tensor([ 0, 14, 176, 38, 683, 7, 6, 116, 7147, 4, 1, 2,
2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2,
2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2,
2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2,
2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2,
2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2],
device='cuda:0')]一個batch下有兩個樣本，對另一個樣本的德語進行同樣的轉(zhuǎn)換最終得到編碼器輸入：X，維度[2, 72]。
在訓(xùn)練過程中，無論是計算注意力還是交叉注意力，每個樣本是相互獨立的，所以可以將一個batch下所有數(shù)據(jù)組織成矩陣的形式輸入到模型進行并行計算。
1.3.最后將上一步的tokens通過一個Embedding線性層轉(zhuǎn)換成詞嵌入，設(shè)置d_model=512，所以詞嵌入維度為512。
{Embeddings類}
Embedding層輸入就是前面的X；維度是torch.Size([2, 72])。
Embedding層的輸出維度是torch.Size([2, 72，512])，也就是每個token id都被轉(zhuǎn)換成512維的向量。
tensor([[[-0.6267, -0.0099, 0.3444, ..., 0.5949, -0.4107, -0.6037],
[ 0.4183, -0.1788, -0.3128, ..., 0.5363, -0.5519, 0.4621],
[ 0.4645, -0.2748, -0.4109, ..., -0.6270, 0.4595, -0.4259],
...,
[-0.1489, 0.6431, -0.0301, ..., -0.0163, 0.4261, 0.3066],
[-0.1489, 0.6431, -0.0301, ..., -0.0163, 0.4261, 0.3066],
[-0.1489, 0.6431, -0.0301, ..., -0.0163, 0.4261, 0.3066]],
[[-0.6267, -0.0099, 0.3444, ..., 0.5949, -0.4107, -0.6037],
[-0.2121, 0.4323, -0.0869, ..., 0.1337, -0.2679, -0.4689],
[ 0.0751, -0.1048, -0.1263, ..., -0.5541, -0.4463, 0.5209],
...,
[-0.1489, 0.6431, -0.0301, ..., -0.0163, 0.4261, 0.3066],
[-0.1489, 0.6431, -0.0301, ..., -0.0163, 0.4261, 0.3066],
[-0.1489, 0.6431, -0.0301, ..., -0.0163, 0.4261, 0.3066]]],
device='cuda:0', grad_fn=<MulBackward0>)（Embedding過程相當于用512個屬性值表示單詞的語義信息，經(jīng)過每個EncoderLayner時屬性值會被修改，使其充分吸收上下文信息，屬性越多，能表示的語音信息越豐富，但計算量和參數(shù)也會增加）
2、PositionalEncoding
{PositionalEncoding類}
在計算注意力分數(shù)時，如果調(diào)整單詞的位置，注意力的輸出結(jié)果不變，也就是自注意力這種計算方式?jīng)]有考慮單詞的位置信息。
所以需要通過一個額外的位置編碼，位置編碼與詞嵌入維度相同，也是512維向量，最后與詞嵌入相加。
前面Embedding層輸出維度torch.Size([2, 72, 512])，將其與位置編碼相加，輸出也是torch.Size([2, 72, 512])。
（位置編碼可以通過訓(xùn)練方法得到，也可以采用固定計算方式，本例采用固定計算方式）
所有樣本共用同一個位置編碼，本例序列長度為72，可以提前計算好位置編碼備用。
pos表示位置，第一個詞位置是0，第二個詞位置是1....本例中就是0-71。
對于512維向量，偶數(shù)位置和奇數(shù)位置的值分別用上面兩個公式計算。
tensor([[[ 0.0000e+00, 1.0000e+00, 0.0000e+00, ..., 1.0000e+00,
0.0000e+00, 1.0000e+00],
[ 8.4147e-01, 5.4030e-01, 8.2186e-01, ..., 1.0000e+00,
1.0366e-04, 1.0000e+00],
[ 9.0930e-01, -4.1615e-01, 9.3641e-01, ..., 1.0000e+00,
2.0733e-04, 1.0000e+00],
...,
[-8.9793e-01, 4.4014e-01, 3.6763e-01, ..., 9.9997e-01,
7.0490e-03, 9.9998e-01],
[-1.1478e-01, 9.9339e-01, -5.5487e-01, ..., 9.9997e-01,
7.1527e-03, 9.9997e-01],
[ 7.7389e-01, 6.3332e-01, -9.9984e-01, ..., 9.9997e-01,
7.2564e-03, 9.9997e-01]]], device='cuda:0')可視化出來就是下面效果。
(上圖每一行都是一個位置編碼向量，一共生成50個位置編碼，每個位置編碼是128維向量，而本例需要生成72個，每個512維）
3.MultiHeadedAttention
{MultiHeadedAttention類,
attention函數(shù)}
MultiHeadedAttention類的輸入是query, key, value，維度都是torch.Size([2, 72, 512])，其實他們的內(nèi)容也是一樣的，就是上一步輸出的Embedding+位置編碼。
然后query, key, value分別經(jīng)過一個獨立的線性層，線性層的維度[512, 512]，兩個樣本的[72, 512]分別與[512, 512]矩陣乘法，所以線性層的輸出維度仍是[2, 72, 512]，最后經(jīng)過reshape和轉(zhuǎn)置將[2, 72, 512]轉(zhuǎn)換成torch.Size([2, 8, 72, 64])，8代表有8個頭，其實就是將512轉(zhuǎn)換成了8*64來實現(xiàn)多頭注意力機制。
(雖然是8個頭，但與一個頭的情況相比，參數(shù)并沒有增加)
接下來計算單個頭的注意力，Attention函數(shù)的輸入query, key, value的維度都是torch.Size([2, 8, 72, 64])，注意力分數(shù)矩陣維度torch.Size([2, 8, 72, 72])，輸出torch.Size([2, 8, 72, 64])。
最后將多個頭的輸出拼接在一起，也就是通過reshape和轉(zhuǎn)置將torch.Size([2, 8, 72, 72])轉(zhuǎn)換成[2, 72, 512]，最后經(jīng)過一個[512, 512]的線性層輸出[2, 72, 512]。
4、SublayerConnection
{SublayerConnection類}
將多頭注意力的輸出經(jīng)過層歸一化和輸入進行殘差鏈接，不改變維度，輸入輸出都是[2, 72, 512]。
5、PositionwiseFeedForward
{PositionwiseFeedForward類}
這其實是一個MLP層，輸入維度512，隱藏層維度2048，輸出層維度512，也就是2*72個tokens并行與[512, 2048]矩陣乘升維至[2, 72, 2048]，然后再與矩陣[2048，512]乘恢復(fù)到原來維度[2, 72, 512]。最后再經(jīng)過層歸一化和殘差鏈接。
6、EncoderLayer
{EncoderLayer類}
將3，4，5重復(fù)6次，這里需要注意下，這6個EncoderLayer只是結(jié)構(gòu)一致，但參數(shù)是獨立的，原始的Embedding經(jīng)過6個EncoderLayer后維度是不變的，仍然是[2, 72, 512]，只不過內(nèi)容被改變了。
7、LaynerNorm
{LayerNorm類}
為了計算穩(wěn)定，整個Encoder的輸出會再次經(jīng)過層歸一化處理，然后輸入到Decoder層作為key和value，維度仍然是[2, 72, 512]。
Encoder把key和value傳遞給Decoder，它的使命就算完成了。剩下的就是根據(jù)那邊的損失等著更新參數(shù)了。
讓我們來到Decoder這邊。
0、Inputs：
[
('Eine gro?e Gruppe Jugendlicher in einem kleinen Unterhaltungsbereich.', 'A large group of young adults are crammed into an area for entertainment.'),
('Zwei Arbeiter stellen Laternen auf.', 'Two workers working on putting up lanterns.')
]對于Decoder，除了Encoder的key和value，還要有query，這個query就是目標語言，也就是英語。
'A large group of young adults are crammed into an area for entertainment.'
'Two workers working on putting up lanterns.'
1、Embedding
Decoder和Encoder的Embedding幾乎一致，也是先轉(zhuǎn)換成tokens。
'A large group of young adults are crammed into an area for entertainment.'->tokens
torch.Size([16])
tensor([ 0, 6, 62, 39, 13, 25, 348, 17, 5318, 71, 28, 179,
55, 4285, 5, 1], device='cuda:0')然后進行padding。
'A large group of young adults are crammed into an area for entertainment.'->padding tokens
torch.Size([72])
[tensor([ 0, 6, 62, 39, 13, 25, 348, 17, 5318, 71, 28, 179,
55, 4285, 5, 1, 2, 2, 2, 2, 2, 2, 2, 2,
2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2,
2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2,
2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2,
2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2],
device='cuda:0')]對另一個樣本進行同樣的操作得到編碼器的輸入Y，維度[2, 72]。
最后將其轉(zhuǎn)換成Embedding，維度是torch.Size([2, 72, 512])。
但有一點需要注意。
Decoder在訓(xùn)練時輸入的是整個batch的英語文本，也就是torch.Size([2, 72, 512])。
但在訓(xùn)練過程中預(yù)測當前token的輸出時，為了讓其只能看到當前以及之前位置的輸入，避免看到后面的內(nèi)容，需要采用遮罩的方式，也就是要構(gòu)造一個mask。
torch.Size([2, 72, 72])
tensor([[[ True, False, False, ..., False, False, False],
[ True, True, False, ..., False, False, False],
[ True, True, True, ..., False, False, False],
...,
[ True, True, True, ..., False, False, False],
[ True, True, True, ..., False, False, False],
[ True, True, True, ..., False, False, False]],
[[ True, False, False, ..., False, False, False],
[ True, True, False, ..., False, False, False],
[ True, True, True, ..., False, False, False],
...,
[ True, True, True, ..., False, False, False],
[ True, True, True, ..., False, False, False],
[ True, True, True, ..., False, False, False]]], device='cuda:0')
2、PositionalEncoding
與Encoder一樣，輸入輸出都是[2, 72, 512]
3、MultiHeadedAttention
Decoder中的DecoderLayner有兩個MultiHeadedAttention，第一個是Mask MultiHeadedAttention，與Encoder中的計算一致，只不過使用了上一步計算的Mask。
另一個MultiHeadedAttention中的key和value來自Encoder，我們稱之為交叉注意力，與自注意力要區(qū)分開，query來自前一層的輸出，維度都是[2, 72, 512]。
4，5，9，7和Encoder都是一樣的。
同樣輸入Embedding經(jīng)過6個DecoderLayner后維度不變[2, 72, 512]。
4、Generator
{Generator類}
這其實是一個沒有隱藏層的MLP，輸入維度512，輸出維度vocab，2*72個token的Embedding與矩陣[512,vocab]相乘,輸出[2, 72, vocab]，vocab為詞表的單詞個數(shù)，本例中英語單詞個數(shù)為6291。經(jīng)過softmax后輸出一個概率分布，最大概率對應(yīng)的位置的詞就是模型預(yù)測的下一個詞。
這樣就得到了Decoder的最終輸出，輸出可以是[2, 72]，里面是英語詞表下的id。也可以是[2, 72, vocab]直接輸出概率分布，輸出形式不同，損失函數(shù)也是不同的。
對于其中一個樣本，訓(xùn)練過程中Decoder的輸入是：
'A large group of young adults are crammed into an area for entertainment.'
torch.Size([72])
[tensor([ 0, 6, 62, 39, 13, 25, 348, 17, 5318, 71, 28, 179,
55, 4285, 5, 1, 2, 2, 2, 2, 2, 2, 2, 2,
2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2,
2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2,
2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2,
2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2],
device='cuda:0')]如果想更新參數(shù)就必須計算損失，計算損失就必須有標簽，那標簽是什么？
對于Decoder，輸入也是輸出，標簽就是將輸入向左移動了一位：
[tensor([ 6, 62, 39, 13, 25, 348, 17, 5318, 71, 28, 179,
55, 4285, 5, 1, 2, 2, 2, 2, 2, 2, 2, 2,
2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2,
2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2,
2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2,
2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2，2],
device='cuda:0')]也就是起始符0對應(yīng)的標簽是A：6，輸入A對應(yīng)的標簽是large：62，Decoder輸出維度[2, 72]，標簽維度也是[2, 72]，最后通過均方誤差計算損失，或者輸出概率分布，通過KL損失函數(shù)計算損失來更新Decoder和Encoder的參數(shù)。
再強調(diào)一下，整個batch下所有數(shù)據(jù)是一起輸入到模型的，也就是通過將數(shù)據(jù)組織成矩陣實現(xiàn)了整個batch的數(shù)據(jù)并行計算。
訓(xùn)練完成后，就可以用它進行德譯英翻譯了。
假設(shè)輸入這么一句德語。
'Eine gro?e Gruppe Jugendlicher in einem kleinen Unterhaltungsbereich.'
德語先經(jīng)過Encoder進行并行編碼，輸出[1, 72, 512]作為Decoder的value和key。
在推理過程中就Deocder就不能并行計算了，只能自回歸的方式每次前向計算只產(chǎn)生一個token。
剛開始只有一個起始符token 0輸入到Deocder。
[tensor([ 0, 1, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2,
2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2,
2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2,
2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2,
2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2,
2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2],
device='cuda:0')]decoder輸出6，將6加到0后面再次輸入到decoder。
[tensor([ 0, 6, 1, 2, 2, 2, 2, 2, 2, 2, 2, 2,
2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2,
2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2,
2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2,
2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2,
2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2],
device='cuda:0')]decoder輸出62，以此類推，直到輸出終止符token 1。
人的大腦在學(xué)習(xí)復(fù)雜事物時，往往習(xí)慣使用一種整體到細節(jié)，抽象到具體的漸進的方式。
雖然我在作者的源代碼添加了更多的注釋和維度信息，但它仍然是細節(jié)，為了更好地理解大模型的工作原理，我建議先閱讀我之前的圖解和動畫Transformer系列，以次獲得對Transformer有一個高層次的認知。
另外，如果你如果弄明白了Encoder-Decoder架構(gòu)，那么就能輕松搞懂GPT和BERT了，因為它們一個只用了Encoder，另一個只用了Decoder。
本文轉(zhuǎn)載自公眾號人工智能大講堂
