Alan Lee

BERT 是如何构建模型的

2020/05/08 Share

Photo by Damian Patkowski on Unsplash

Good things take time, as they should.

前面我写了一篇文章来讲 BERT 是如何分词的,现在,轮到该说说 BERT 模型是如何定义的了。

BERT 模型的大致结构可能大家已经很清楚了,实际上核心就是 Transformer encoder。本文主要是结合代码(modeling.py)实现来看下模型的定义,以及相关辅助函数,带你解读整个 modeling.py

modeling.py 共有 2 个类,16 个函数,我先放一张 modeling.py 中类、方法和函数的总的调用关系图,大致了解一下:

modeling-call-graph.png

本文先介绍下文件中仅有也比较重要的两个类:BertConfigBertModel。然后根据构建 BERT 模型「三步走」的顺序,分别介绍下这三步,同时介绍一下相关函数。

BertConfig

BERT 模型的配置类,BERT 的超参配置都在这里。其参数(蓝色)和方法(黄色)总览如下:

bert-config.png

下面我分别介绍下参数和方法的意义。

参数

  • 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=16token_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 模型的详细定义就在这里了。其参数(蓝色)、方法(框内黄色)和对其他类、函数的调用关系总览如下:

bert-model.png

下面我分别介绍下参数和方法的意义。

参数

  • 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_idsinput_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_typeuse_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.

我一开始并没想通这点,于是做了个测试,结果如下:

one-hot-vs-index.png

可见两者差距甚小,在 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_sizefrom_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 的大小就是两者相除得到的,两者关系如下图:

annotation-on-attention.png

和其他函数的调用关系如下图:

transformer-model.png

这个函数是重头戏,大致的整体流程如下图,我省略了 transpose 之类的转 shape 的操作:

bert-transformer-model.png

OK,是不是看起来也没那么复杂?核心就是 hidden layer,我下面简单解释下一个 hidden layer 的流程:

transformer 的输入(input_tensor 是初始值,后续的输入是 layer_input,即上一层的输出)的 shape 是 [batch_size, seq_length, hidden_size],而 hidden_sizeembedding_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,我们先回顾下其数学形式:

其中,

而,

现在来说下 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 的操作:

bert-attention-layer.png

看了这个图之后,相信大家会觉得过程并没那么复杂,我来简单解释下:

  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 的输出了。

此外,再插播一个关于第一步实现方面的疑问。原代码中第一步的实现是这样的:

1
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]

其实我觉得,这步操作可以简化,不使用切片操作即可一步到位,即:

1
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.BertModelmodeling.get_assignment_map_from_checkpointmodeling.BertConfig.from_json_file

BERT 构建模型部分到此结束。

Reference

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_namesassignment_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 ExplainedWhat 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 的。
CATALOG
  1. 1. 类 BertConfig
    1. 1.1. 参数
    2. 1.2. 方法
  2. 2. 类 BertModel
    1. 2.1. 参数
    2. 2.2. 方法
  3. 3. Embedding
  4. 4. Encoder
  5. 5. Pooler
  6. 6. 总结
  7. 7. Reference
  8. 8. END
  9. 9. 附录