BERT 是如何構建模型的

本文於 2020/05/08 首發於我的個人博客，未經允許，不得轉載。

Good things take time, as they should.

前面我寫了一篇文章來講 BERT 是如何分詞的，現在，輪到該說說 BERT 模型是如何定義的了。

BERT 模型的大致結構可能大家已經很清楚了，實際上核心就是 Transformer encoder。本文主要是結合代碼（modeling.py）實現來看下模型的定義，以及相關輔助函數，帶你解讀整個 modeling.py。

modeling.py 共有 2 個類，16 個函數，我先放一張 modeling.py 中類、方法和函數的總的調用關係圖，大致瞭解一下：

本文先介紹下文件中僅有也比較重要的兩個類：BertConfig 和 BertModel。然後根據構建 BERT 模型「三步走」的順序，分別介紹下這三步，同時介紹一下相關函數。

類 BertConfig

BERT 模型的配置類，BERT 的超參配置都在這裏。其參數（藍色）和方法（黃色）總覽如下：

下面我分別介紹下參數和方法的意義。

參數

• vocab_size：詞彙表大小。
• hidden_size=768：encoder 層和 pooler 層大小。這實際上就是 embedding_size，BERT 乾的事情就是不停地優化 embedding。。。
• num_hidden_layers=12：encoder 中隱層個數。
• num_attention_heads=12：每個 attention 層的 head 個數。
• intermediate_size=3072：中間層大小。
• hidden_act="gelu"：隱層激活函數。
• hidden_dropout_prob=0.1：所有全連接層的 dropout 概率，包括 embedding 和 pooler。
• attention_probs_dropout_prob=0.1：attention 層的 dropout 概率。
• max_position_embeddings=512：最大序列長度。
• type_vocab_size=16：token_type_ids 的詞彙表大小。
• initializer_range=0.02：初始化所有權重時的標準差。

方法

• from_dict(cls, json_object)：從一個字典來構建配置。
• from_json_file(cls, json_file)：從一個 json 文件來構建配置。
• to_dict(self)：將配置保存爲字典。
• to_json_string(self)：將配置保存爲 json 字符串。

類 BertModel

BERT 模型類，主角，BERT 模型的詳細定義就在這裏了。其參數（藍色）、方法（框內黃色）和對其他類、函數的調用關係總覽如下：

下面我分別介紹下參數和方法的意義。

參數

• config：配置，BertConfig 實例。
• is_training：是否開啓訓練模式，否則是評估/預測模式。也控制了是否使用 dropout。
• input_ids：輸入文本對應的 id，大小爲 [batch_size, seq_length]。
• input_mask=None：int32 類型，大小和 input_ids 相同。
• token_type_ids=None：int32 類型，大小和 input_ids 相同。
• use_one_hot_embeddings=False：是否使用 one-hot embedding，否則使用 tf.embedding_lookup()。
• scope=None：變量 scope，默認爲 bert。

方法

• __init__()：重頭戲，模型的構建在此完成，三步走。主要分爲三個模塊：embeddings、encoder 和 pooler。首先構建輸入，包括 input_ids、input_mask 等。其次進入 embeddings 模塊，進行一系列 embedding 操作，涉及 embedding_lookup() 和 embedding_postprocessor() 兩個函數。然後進入 encoder 模塊，就是 transformer 模型和 attention 發揮作用的地方了，主要涉及 transformer_model() 函數，得到 encoder 各層輸出。最後進入 pooler 模塊，只取 encoder 最後一層的輸出的第一個 token 的信息，送入到一個大小爲 hidden_size 的全連接層，得到 pooled_output，這就是最終輸出了。
• get_pooled_output(self)：獲取 pooler 的輸出。
• get_sequence_output(self)：獲取 encoder 最後的隱層輸出，輸出大小爲 [batch_size, seq_length, hidden_size]。
• get_all_encoder_layers(self)：獲取 encoder 中所有層。返回大小應該是 [num_hidden_layers, batch_size, seq_length, hidden_size]。
• get_embedding_output(self)：獲取對 input_ids 的 embedding 結果，大小爲 [batch_size, seq_length, hidden_size]，這是 word embedding、positional embedding、token type embedding（論文中的 segment embedding）和 layer normalization 一系列操作的結果，也是 transformer 的輸入。
• get_embedding_table(self)：獲取 embedding table，大小爲 [vocab_size, embedding_size]，即詞彙表中的詞對應的 embedding。

Embedding

如前所述，構建 BERT 模型主要有三塊：embeddings、encoder 和 pooler。先來介紹下 embeddings。

顧名思義，此步就是對輸入進行嵌入。在初始詞向量的基礎上，BERT 又加入 Token type embedding（即論文中的 segment embedding）和 Position embedding 來增強表示。

計算初始詞向量對應於 embedding_lookup() 函數，參數爲

• input_ids
• vocab_size
• embedding_size=128
• initializer_range=0.02
• word_embedding_name="word_embeddings"
• use_one_hot_embeddings=False

此函數根據 input_ids 找對應的 embedding，輸入大小爲 [batch_size, seq_length]，輸出兩個值：

• output，大小爲 [batch_size, seq_length, embedding_size]。
• embedding_table，大小爲 [vocab_size, embedding_size]。

後續的騷操作對應於 embedding_postprocessor() 函數，參數爲：

• input_tensor
• use_token_type=False
• token_type_ids=None
• token_type_vocab_size=16
• token_type_embedding_name="token_type_embeddings"
• use_position_embeddings=True
• position_embedding_name="position_embeddings"
• initializer_range=0.02
• max_position_embeddings=512
• dropout_prob=0.1

該函數在初始 embedding（input_tensor）的基礎上再進行一頓 embedding 騷操作，然後加上 layer normalization 和 dropout 層。

根據 use_token_type 和 use_position_embeddings 的值，最多會進行兩種騷操作：

• Token type embedding：即論文中的 segment embedding，首先會創建一個 token_type_table，然後拿着 token_type_ids 去查，得到 token type embedding，該 embedding 的 shape 和原 embedding 是一樣的，直接將其加到原 embedding 就行。
• Position embedding：位置信息嵌入。這裏有一個需要注意的地方：max_position_embeddings，這個參數的值必須 ≥ seq_length，因爲代碼中會首先構造一個大小爲 [max_position_embeddings, embedding_size] 的 full_position_embeddings，然後再使用 tf.slice 截取 seq_length 大小，從而得到一個 [1, seq_length, embedding_size] 的 embedding，最後加上原 embedding 即可。這裏注意，爲什麼得到的 embedding 第一維是 1 呢？因爲一個 batch 內的位置嵌入是相同的，假如一個 batch 有兩句話，那麼這兩句話第一個字的位置嵌入都是 1 對應的 embedding，是相同的，所以可以直接 broadcast 到整個 batch 維度上。而 token type embedding 這些是不同的。

在 Token type embedding 代碼實現部分，根據 token_type_table 獲取 token_type_ids 的對應 embedding 的時候，BERT 使用的是 one-hot 方法，即 token_type_ids 的 one-hot 矩陣乘 token_type_table，而不是使用的直接按索引取的方法。根據 BERT 代碼註釋，這是因爲對於該 embedding，token_type_vocab_size 通常很小（一般爲 2），此時 one-hot 方法更快：

This vocab will be small so we always do one-hot here, since it is always faster for a small vocabulary.

我一開始並沒想通這點，於是做了個測試，結果如下：

可見兩者差距甚小，在 vocab size 比較小的時候，one-hot 甚至會比索引方法慢。one-hot 方法需要進行矩陣乘法，而索引方法則是直接按索引取值，所以 one-hot 應該慢點纔對，此問題尚未想清楚，歡迎評論區討論。

OK 回到正題，一頓操作後 embeddings 的維度維持不變，仍然是 [batch_size, seq_length, embedding_size]。歸來仍是少年。

Embeddings 部分結束。

Encoder

Embeddings 部分結束後的輸出大小是 [batch_size, seq_length, embedding_size]，這個將會輸入給 encoder。

該部分首先創建一個 attention_mask，然後作爲參數傳給 transformer encoder 模型，最終得到多層 encoder layer 的輸出。實際傳給下一步 pooler 的時候，使用的是最後一層輸出。

創建 attention_mask 部分使用的是 create_attention_mask_from_input_mask() 函數，參數爲：

• from_tensor
• to_mask

此函數將一個二維的 mask 變成一個三維的 attention mask，最終的輸出 shape 爲 [batch_size, from_seq_length, to_seq_length]。from_tensor 在這裏的唯一作用就是提供一下 batch_size 和 from_seq_length。事實上該 attention_mask 是全 1 的：

We don’t assume that from_tensor is a mask (although it could be). We don’t actually care if we attend from padding tokens (only to padding) tokens so we create a tensor of all ones.

核心 transformer encoder 的部分對應於 transformer_model() 函數，參數爲：

• input_tensor
• attention_mask=None
• hidden_size=768
• num_hidden_layers=12
• num_attention_heads=12
• intermediate_size=3072
• intermediate_act_fn=gelu
• hidden_dropout_prob=0.1
• attention_probs_dropout_prob=0.1
• initializer_range=0.02
• do_return_all_layers=False

要注意的一點是，hidden_size 必須能夠整除 num_attention_heads，因爲每個 head 的大小就是兩者相除得到的，兩者關係如下圖：

和其他函數的調用關係如下圖：

這個函數是重頭戲，大致的整體流程如下圖，我省略了 transpose 之類的轉 shape 的操作：

OK，是不是看起來也沒那麼複雜？核心就是 hidden layer，我下面簡單解釋下一個 hidden layer 的流程：

transformer 的輸入（input_tensor 是初始值，後續的輸入是 layer_input，即上一層的輸出）的 shape 是 [batch_size, seq_length, hidden_size]，而 hidden_size 和 embedding_size（或者叫 input_width）是相等的，即你可以認爲輸入就是 embedding 結果。

1. 輸入送入 attention layer，得到輸出 attention output。
2. 一層線性映射，神經元數量（hidden_size）和 embedding_size 相同。
3. dropout 和 layer normalization，注意後者的輸入是前者 + layer_input。
4. 一層非線性映射，默認情況下神經元數量要遠大於線性映射層的數量。
5. 再來一層線性映射，重新將維度拉回 embedding_size。
6. dropout 和 layer normalization，注意後者的輸入是前者 + 4 的輸出。
7. 完事，得到這一個 hidden layer 的輸出，然後作爲下一層 hidden layer 的輸入。

這樣一來，一個 hidden layer 得到一個輸出，總共會得到 num_hidden_layers 個 輸出，都 append 到一個 list。如果 do_return_all_layers=True 的話，就把這些輸出全都 reshape 成原來的樣子然後返回。否則，直接把最後一層的輸出 reshape 成原來的樣子然後返回。

經過一頓操作，歸來還是少年。

第一步的 attention layer 在這裏有非常重要的作用，後面幾步基本就是對其輸出做一些映射變換，比較好理解。這裏的 attention 其實是 MultiHead self-attention，我們先回顧下其數學形式：

$\text{MultiHead}(Q, K, V) = \text{Concat}(\text{head}_1,\dots,\text{head}_\text{h})W^O$

其中，

$\text{head}_\text{i} = \text{Attention}(QW_i^Q, KW_i^K, VW_i^V)$

而，

$\text{Attention}(Q, K, V) = \text{softmax}\left(\dfrac{QK^T}{\sqrt{d_k}}\right)V$

現在來說下 attention layer 的實現。

attention layer 對應函數爲 attention_layer()，參數爲：

• from_tensor
• to_tensor
• attention_mask=None
• num_attention_heads=1
• size_per_head=512
• query_act=None
• key_act=None
• value_act=None
• attention_probs_dropout_prob=0.0
• initializer_range=0.02
• do_return_2d_tensor=False
• batch_size=None
• from_seq_length=None
• to_seq_length=None

大致的整體流程如下圖，我同樣省略了 transpose 之類的轉 shape 的操作：

看了這個圖之後，相信大家會覺得過程並沒那麼複雜，我來簡單解釋下：

1. 首先得到 Q、K、V 三個矩陣，都是分別經過一層相同大小的線性映射。其中 Q 通過 from_tensor 得到，K、V 通過 to_tensor 得到。
2. Q、K 經過矩陣乘法和 scale 得到初步的 raw attention score，注意 shape 爲 [batch_size, num_attention_heads, from_seq_length, to_seq_length]。如果有 attention mask，那麼將其加到 raw attention score 上。
3. 上面得到的 raw attention score 經過 softmax，得到概率形式的 attention probability。
4. dropout。
5. attention probability 和 V 做矩陣乘法，得到 context layer。這就是最終的返回結果。注意 shape 爲 [batch_size, from_seq_length, num_attention_heads * size_per_head]，通常來說最後一維就是 embedding size。

OK，Encoder 部分到此結束。

Pooler

前面 Encoder 部分得到的多層輸出，最終是取最後一層輸出來輸入給 Pooler 部分。這部分相對簡單，只是取每個 sequence 的第一個 token，即原本的輸入大小爲 [batch_size, seq_length, hidden_size]，變換後大小爲 [batch_size, hidden_size]，去掉了 seq_length 維度，相當於是每個 sequence 都只用第一個 token 來表示。然後接上一層 hidden_size 大小的線性映射即可，激勵函數爲 tf.tanh。

至此就得到了 BertModel 的輸出了。

此外，再插播一個關於第一步實現方面的疑問。原代碼中第一步的實現是這樣的：

first_token_tensor = tf.squeeze(self.sequence_output[:, 0:1, :], axis=1)

此處用切片操作（0:1）來取第一個元素，保持了結果的 rank 和 sequence_output 的 rank 相同，sequence_output 的大小爲 [batch_size, seq_length, hidden_size]，切片操作後的大小變爲 [batch_size, 1, hidden_size]。然後再用 tf.squeeze 來將第二個維度壓縮掉，即大小變爲 [batch_size, hidden_size]。

其實我覺得，這步操作可以簡化，不使用切片操作即可一步到位，即：

first_token_tensor = self.sequence_output[:, 0, :]

不是很明白原代碼那樣寫是有何意圖，此問題尚未想清楚，歡迎評論區討論。

總結

簡而言之，BERT 的大致流程就是：引入配置 BertConfig -> 定義初始化輸入大小等常量 -> 對輸入進行初步 embedding -> 加入 token type embedding 和 position embedding -> 創建 encoder 獲取輸出 -> 獲取 pooled 輸出，就是最終輸出了，在 run_classifier.py 中會將此輸出接上一個 Dropout，然後接上一個 softmax 分類層。

run_classifier.py 中涉及 modeling.py 的地方有三處：modeling.BertModel、modeling.get_assignment_map_from_checkpoint、modeling.BertConfig.from_json_file。

BERT 構建模型部分到此結束。

Reference

• google-research/bert: TensorFlow code and pre-trained models for BERT
• [1810.04805] BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding
• [1706.03762] Attention Is All You Need
• The Illustrated Transformer – Jay Alammar – Visualizing machine learning one concept at a time.
• BERT Explained – A list of Frequently Asked Questions – Let the Machines Learn

END

附錄

這裏是一些正文沒有提到的函數的解釋。

• 函數 gelu(x)：GELU（Gaussian Error Linear Unit）激活函數，是 RELU 的平滑版本，見論文 Gaussian Error Linear Units (GELUs)。
• 函數 get_activation(activation_string)：就是一個映射，將類似 'relu' 這樣的 str 格式的 activation_string，變成 tf 中實際的函數 tf.nn.relu。就是一些 if 判斷。
• 函數 get_assignment_map_from_checkpoint(tvars, init_checkpoint)：獲取 assignment_map，同時也返回 initialized_variable_names。
  • 什麼是 assignment_map？這其實是 tf.compat.v1.train.init_from_checkpoint(ckpt_dir_or_file, assignment_map) 的一個參數，用於指定當前 graph 的哪些變量的值需要從其他 checkpoint 中導入，dict 格式，key 爲 checkpoint 中的變量（即舊變量），value 爲當前 graph 中的變量（即新變量）。
  • initialized_variable_names 和 assignment_map 基本相同，後面 run_classifier.py 中會用其查詢某變量值是否是從外部 checkpoint 中導入的。
• dropout(input_tensor, dropout_prob)：dropout 層。
• layer_norm(input_tensor, name=None)：layer normalization 層，關於 layer normalization 和 batch normalization 的區別，參見 Weight Normalization and Layer Normalization Explained (Normalization in Deep Learning Part 2) | Machine Learning Explained 和 What are the practical differences between batch normalization, and layer normalization in deep neural networks? - Quora，簡而言之就是 layer normalization 是在 feature 維度進行 normalization，而 batch normalization 是在 batch 維度進行。
• layer_norm_and_dropout(input_tensor, dropout_prob, name=None)：先 layer normalization 後 dropout。
• create_initializer(initializer_range=0.02)：創建一個 truncated_normal_initializer 來初始化參數。
• get_shape_list(tensor, expected_rank=None, name=None)：獲取 tensor
  的 shape，list 形式返回。要注意的一點是，如果這個 tensor 有動態維度，即某個維度爲 None ，那麼返回的時候，該維度會是一個 tensor。例如有一個 shape 爲 [None, 3] 的 tensor，調用該函數時的返回就類似於 [<tf.Tensor 'strided_slice:0' shape=() dtype=int32>, 3]。函數內部實現是先獲取動態維度的索引，然後使用 tf.shape() 來取得對應索引的 tensor 形式的維度。
• reshape_to_matrix(input_tensor)：將一個 rank >= 2 的 tensor reshape 成 rank = 2 的 tensor，即矩陣。具體是固定最後一個維度，將剩餘維度都壓縮到一個維度，即 reshape((-1, shape[-1]))。注意 TensorFlow 中的 rank 不同於數學中的 rank 概念，數學中是秩，而 TF 中是 ndims，即 len(shape)。
• reshape_from_matrix(output_tensor, orig_shape_list)：和 reshape_to_matrix() 相反，將已經 reshape 成矩陣的 tensor 重新 reshape 到原來的樣子。
• assert_rank(tensor, expected_rank, name=None)：檢查 tensor 的 rank 是否符合要求（=/in expected_rank），不符合則拋出 ValueError 異常。注意此函數是用 dict 來存儲 expected_rank 的。