GPT-2代码解读[1]：Overview和Embedding

GPT-2代码解读[1]：Overview和Embedding

Abstract

随着Transformer结构给NLU和NLG任务带来的巨大进步，GPT-2也成为当前（2020）年顶尖生成模型的泛型，研究其代码对于理解Transformer大有裨益。可惜的是，OpenAI原始Code基于tensorflow1.x，不熟悉tf的同学可能无从下手，这主要是由于陌生环境¹导致的。

本文的意愿是帮助那些初次接触GPT-2代码且对tensorflow不熟悉的同学，也适合那些在阅读代码过程中遇到困难的同学。我会逐句解析。

本文结构如下：

分三部分解析gpt-2的基础模块和函数，即embedding，attn，mlp，这些模块可以独立使用在任何模型中，具有隽永的价值，我尽量做到模型代码（即/src/model.py）的全覆盖，但经验使然的部分（batch normalization，gelu等）不做解析。

注意事项

1. 在阅读tensorflow代码或考虑模型结构的时候，应该学会区分trainable的部分，和仅仅用于形状变化的部分，前者是真正指导信息流动的部分，是重点，抓住这些重点能够促进减少阅读花费，并且有助于理解本质。
2. 我使用X=[a,b,c,…]来标注X的形状。

Model Overview

def model(hparams, X, past=None, scope='model', reuse=tf.AUTO_REUSE):

模型的输入信息分两种：X和past，X是语言模型的输入，past是已生成上文的状态，实作分四种情况：

1. 训练时，X为一组训练数据²，past为空。
2. 条件生成初始阶段，X为条件语句，past为空
3. 无条件生成初始阶段，X为[end]，past为空
4. 生成过程中，X为上一次生成的词语，past为之前所有K和V。

如果你不太熟悉生成模型，请只考虑第一种情况！我在解析时也会主要考虑第一种情况！

此时，X=[batch,seq]

模型由12个基本块构成，每一块由三部分构成，首先考虑< Embedding Or h >。

当模型处于i=1块时，需要将输入的词语Embedding为语义向量。

当模型处于i>1块时，直接以上一层的输出作为输入。

Embedding

embedding层的作用是将单词嵌入为语义向量，它的输入是模型的输入X。输出单词的语义信息。

在gpt使用的Transformer中，语义分为两种，一是单词本身语义，二是单词所处位置的语义。

换句话说，上次的预测结果提供两种信息

1. 词语是什么？
2. 词语的位置是什么？

信息的流动是本质性的，怎么利用这些信息则见仁见智。

wpe = tf.get_variable('wpe', [hparams.n_ctx, hparams.n_embd],
                     initializer=tf.random_normal_initializer(stddev=0.01))

wpe是位置嵌入矩阵³，size为最大序列长度X嵌入维数，定义为tf.variable⁴，并用正态分布初始化。

wte = tf.get_variable('wte', [hparams.n_vocab, hparams.n_embd],
                             initializer=tf.random_normal_initializer(stddev=0.02))

wte是词语嵌入矩阵，size为单词表长度X嵌入维数，定义为tf.variable⁴，并用正态分布初始化。

past_length = 0 if past is None else tf.shape(past)[-2]

past_length是之前已生成上文的长度，如果past为空，则为0，否则为past的倒数第二维⁵。

h = tf.gather(wte, X) + tf.gather(wpe, positions_for(X, past_length))

h是由输入X提供的信息，也是12层Transformer block的初始输入，它由词语嵌入向量和位置嵌入向量加和⁶而成，shape为[批次大小，输入长度⁷，嵌入维度]。

词语嵌入向量由tf.gather()⁸在词语嵌入矩阵wte中取X对应的值，位置嵌入向量由tf.gather()⁸在位置嵌入矩阵wpe中取X对应的值。

positions_for(X, past_length)

position_for函数输入X和已产生的上文长度past_length，输出为要生成的下标用于选择wpe。

例如，训练过程中，X=“hey you”，past_length=0，返回[[0,1]]，意味着要生存的是第一和第二个词。

生成过程中，已经生成“hey you”，past_length=2，X=“you”，返回[[2]]，意味着要生成的是第三个词。

batch_size = tf.shape(tokens)[0]
nsteps = tf.shape(tokens)[1]

batch_size是批次大小，nsteps是输入X的长度。

return expand_tile(past_length + tf.range(nsteps), batch_size)

返回值是batch_size个X的下标，shape为**[批次大小，输入长度]**。

X的下标由past_length+tf.range(nsteps)提供，tf.range(nsteps)->[0,1,…nsteps-1]，past_length+tf.range(nsteps)->[past_length…n_steps+past_length-1] ，shape为[输入长度]。

然后，使用expand_tile扩展下标到二维。

expand_tile(value, size)

expand_tile输入下标value和批次大小size，返回二维[批次大小，下标]

value = tf.convert_to_tensor(value, name='value')
ndims = value.shape.ndims

value是一个tensor，例如X是已产生序列中的第三个，则是[3]。

ndims是1。

return tf.tile(tf.expand_dims(value, axis=0), [size] + [1]*ndims)

返回X的下标用于选择wpe，shape为[批次大小，下标]。

返回值由tf.tile(input,multiple)产生，该函数将input的第i维拓展multiple[i]次，在这里，是将shape[1,输入长度]的tensor拓展[batch批次大小,1]次，即[批次大小，输入长度] ⁹。

用于拓展的[1,输入长度]tensor，由tf.expand_dims（value）产生，这个函数将原本一维（[输入长度]）的value增加第0维，变为[1,输入长度]。

---

1. 读代码和读书一样，如果遇到太多不熟悉的术语，或者太多不认识的单词，不仅读起来佶屈聱牙，也会同时感到巨大的挫败感。如果你是第一次读，最好的办法是硬着头皮把所有的知识点攻破，然后将其翻译为你所熟知的语言，以便第二次阅读。 ↩︎

2. GPT-2训练使用teacher forcing，将ground truth而非预测词作为语言模型的输入。 ↩︎

3. 对位置信息的利用是见仁见智的，gpt-2里使用可训练的位置嵌入矩阵，原始transformer里则使用sin/cos公式计算位置嵌入。 ↩︎

4. varibale是tensorflow里的几种数据格式，表示可训练的参数，区别于constant（常量），和placehoulder（占位量） ↩︎ ↩︎

5. past为[batch,past_length,embedding]，这一点会在下面详细讲解。 ↩︎

6. 加法看似将不同信息糅混，实则起到融合的作用。加法，拼接，linear都是融合多种tensor的运算，重要的不是运算本身的含义，而是信息的结合。 ↩︎

7. 如果每次预测一个单词（默认情况），就是1，如果X是在条件生成里的条件输入，就会大于1。 ↩︎

8. tf.gather( params, indices, validate_indices=None, axis=None, batch_dims=0, name=None)，功能为取params中indices 对应的axis切片。默认取params中第一个非batch维的切片，在该处即第一维。 ↩︎ ↩︎

9. 每个batch相同， ↩︎