GPT-2 论文+代码笔记

Paper: Language Models are Unsupervised Multitask Learners

Code:

openai/gpt-2: Code for the paper “Language Models are Unsupervised Multitask Learners”
minimaxir/gpt-2-simple: Python package to easily retrain OpenAI’s GPT-2 text-generating model on new texts

核心思想：基于 Transformer 的更加 General 的语言模型。

What

动机和核心问题

作者认为目前这种对单域数据集进行单任务训练的普遍性，是造成当前系统缺乏普遍性的主要原因。另外，目前最好的方法是预训练模型 + 下游任务的监督训练。本文是将二者结合起来，提供更加 general 的迁移方法，它可以使下游任务能够在 zero-shot 下实施，不需要参数或架构调整。证明了语言模型有在 zero-shot 下执行一系列任务的潜力。所以本文核心点有两个：

更加 general 的预训练模型
zero-shot 实施的多下游任务

前者是后者的基础，后者是前者的应用。

模型和算法

其核心是：多领域文本建模，以实现在不同任务上的迁移。所以 model 是这样的：p(output | input, task).

具体方法是：基于 Transformer，对 GPT-1 的改进：

Layer normalization 移动到每个 sub-block 的 input
最后一个 self-attention block 后面加 layer normalization
在初始化时按 1 /√N 的比例缩放残差层的权重，其中 N 是残差层的数量
Vocabulary 扩展到 50,257
上下文 token 数从 512 调整到 1024
更大的 batch size（512）

看这些貌似没啥感觉（当你对底层不理解时，顶层的东西看似懂了的其实都没懂，又称 “司机的知识”），还是代码来的实在。这部分的代码主要来自官方代码。正式开始前，先把参数放出来：

class HParams:
    n_vocab:int=50257
    n_ctx:int=1024
    n_embd:int=768
    n_head:int=12
    n_layer:int=12

输入

首先是输入的 X：

1	context = tf.fill([batch_size, 1], start_token)

这里的 context 就是 start 不为 None 时的 token，我们假设 batch_size 为 1，tf.fill 就是把 start_token 填充为 [batch_size, 1] 的 shape，后面的 1 是指序列长度，我们是按一个一个 token 依次生成的。

然后是 position 和 token 的 embedding：

# weight of position embedding
wpe = tf.compat.v1.get_variable('wpe', [hparams.n_ctx, hparams.n_embd],
                                initializer=tf.random_normal_initializer(stddev=0.01))
# weight of token embedding
wte = tf.compat.v1.get_variable('wte', [hparams.n_vocab, hparams.n_embd],
                                initializer=tf.random_normal_initializer(stddev=0.02))

n_ctx 是上下文的 token 数，n_vocab 是词表大小。合并两个 embedding：

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

position_for 的目的是给 position 编码，它的结果是和 x 类似的一个 Tensor，可以看出这里采用的是绝对位置编码。tf.gather 就是根据索引从参数轴收集切片，即 X 这里的 one hot token 编码使用 wte 索引 X 的向量，和 lookup 类似。这样就得到了最终的输入 h，它的 shape 为：[batch_size, 1, embedding_shape]，在本例中即为：[1, 1, 768]。

Transformer

接下来就是核心的模型部分了，也就是我们熟知的 Transformer。先看一下整体的流程：

pasts = tf.unstack(past, axis=1) if past is not None else [None] * hparams.n_layer
for layer, past in enumerate(pasts):
    h, present = block(h, 'h%d' % layer, past=past, hparams=hparams)
    presents.append(present)
    
results['present'] = tf.stack(presents, axis=1)
h = norm(h, 'ln_f')

tf.unstack 是从 axis 维对张量从 R 级调整为 R-1 级，我们从 None 开始，暂时不需要考虑。可以看到就是 n_layer 个 block 叠加，最后做了个归一化处理。tf.stack 和 tf.unstack 的作用恰好相反。

在 block 前，可以先简单看一下 norm 函数：

def norm(x, *, axis=-1, epsilon=1e-5):
    """Normalize to mean = 0, std = 1, then do a diagonal affine transform."""
    n_state = x.shape[-1]
    g = tf.compat.v1.get_variable('g', [n_state], initializer=tf.constant_initializer(1))
    b = tf.compat.v1.get_variable('b', [n_state], initializer=tf.constant_initializer(0))
    u = tf.reduce_mean(x, axis=axis, keepdims=True) # 均值
    s = tf.reduce_mean(tf.square(x-u), axis=axis, keepdims=True) # 方差
    x = (x - u) * tf.compat.v1.rsqrt(s + epsilon)
    x = x*g + b
    return x

tf.compat.v1.rsqrt 是计算平方根的倒数，上面这其实就是 z 分数标准化过程。

block

先看 block 是什么：

def block(x, scope, *, past, hparams):
	nx = x.shape[-1]
	a, present = attn(norm(x, 'ln_1'), 'attn', nx, past=past, hparams=hparams)
	x = x + a
	m = mlp(norm(x, 'ln_2'), 'mlp', nx*4, hparams=hparams)
	x = x + m
    return x, present

它的结构如下：

attention layer：注意力层
skip connection：残差连接
mlp：就按字面理解为感知机吧
skip connection

attention

无疑这个地方是最核心的了：

# x: [1, 1, 768], nx: 768, past: None
def attn(x, scope, n_state, *, past, hparams):
    c = conv1d(x, 'c_attn', n_state*3)
    q, k, v = map(split_heads, tf.split(c, 3, axis=2))
    present = tf.stack([k, v], axis=1)
    if past is not None:
        pk, pv = tf.unstack(past, axis=1)
        k = tf.concat([pk, k], axis=-2)
        v = tf.concat([pv, v], axis=-2)
    a = multihead_attn(q, k, v)
    a = merge_heads(a)
    a = conv1d(a, 'c_proj', n_state)
    return a, present

这里出现了好几个操作：conv1d, split_heads, multihead_attn, merge_heads，我们分别来看一下。

# conv1d
def conv1d(x, scope, nf, *, w_init_stdev=0.02):
    # x [1, 1, 768], start is [1,1], nx=768, nf=768*3
    *start, nx = shape_list(x)
    w = tf.compat.v1.get_variable(
        'w', [1, nx, nf], initializer=tf.random_normal_initializer(stddev=w_init_stdev))
    b = tf.compat.v1.get_variable(
        'b', [nf], initializer=tf.constant_initializer(0))
    c = tf.reshape(
        	tf.matmul(tf.reshape(x, [-1, nx]), tf.reshape(w, [-1, nf])) + b, 
        	start+[nf])
    # 1*768 × 768*2304 => 1*2304 reshape to [1, 1, 2304]
    return c

我不知道这里为什么会有这么个操作，不过在 Transformers 源代码中看到这么一句注释：Conv1D layer as defined by Radford et al. for OpenAI GPT (and also used in GPT-2). Basically works like a Linear layer but the weights are transposed. 姑且就这样吧。如果只看输入输出的话，它无非就是把输入的维度做了调整，在这里就是把维度放大 3 倍（真不知道怎么想出来的）。这里的目的是为了后面的 q k v。

def split_heads(x):
    # From [batch, sequence, features] to [batch, heads, sequence, features]
    # 即：[1, 1, 768] ==> [1, 12, 1, 64]
    return tf.transpose(split_states(x, hparams.n_head), [0, 2, 1, 3])
def split_states(x, n):
    """Reshape the last dimension of x into [n, x.shape[-1]/n]."""
    *start, m = shape_list(x) # [1, 1, 768]
    return tf.reshape(x, start + [n, m//n]) # [1, 1, 12, 768//12=64]

最后的 shape：[batch, heads, sequence, features] 也就是 q k v 的 shape，也就是把输入 X 的维度分配到 12 个 head 上。

def multihead_attn(q, k, v):
    # q, k, v have shape [batch, heads, sequence, features], [1, 12, 1, 64]
    w = tf.matmul(q, k, transpose_b=True) # [1, 12, 1, 1]
    w = w * tf.compat.v1.rsqrt(tf.cast(v.shape[-1], w.dtype)) # q*k^T/√d_k
    w = mask_attn_weights(w)
    w = softmax(w)
    a = tf.matmul(w, v)
    return a
def mask_attn_weights(w):
    # w has shape [batch, heads, dst_sequence, src_sequence]
    # where information flows from src to dst.
    _, _, nd, ns = shape_list(w)
    b = attention_mask(nd, ns, dtype=w.dtype)
    b = tf.reshape(b, [1, 1, nd, ns])
    w = w*b - tf.cast(1e10, w.dtype)*(1-b)
    return w
def attention_mask(nd, ns, *, dtype):
    """
    1's in the lower triangle, counting from the lower right corner.
    Same as tf.matrix_band_part(tf.ones([nd, ns]), -1, ns-nd), 
    but doesn't produce garbage on TPUs.
    """
    i = tf.range(nd)[:,None]
    j = tf.range(ns)
    m = i >= j - ns + nd # 只有 nd > ns 时，m 才可能有 False，mask 才生效
    return tf.cast(m, dtype)

tf.cast 是改变类型，w 其实就是 Transformer 中的标准 w 计算方式，后面的这个 mask_attn_weights 是对权重 mask，即 masked self-attention，这里其实没有发生作用，我感觉它的主要作用是 mask 未来的字符，这里的 LM 模型利用的是已生成文本，所以并不需要 mask。softmax 没啥好说的，也是 self-attention 的标准配置。

def merge_heads(x):
    # Reverse of split_heads
    return merge_states(tf.transpose(x, [0, 2, 1, 3]))
def merge_states(x):
    """Smash the last two dimensions of x into a single dimension."""
    *start, a, b = shape_list(x)
    return tf.reshape(x, start + [a*b])

这步是合并 head 操作，最终的 shape 为 [batch_size, sequence, n_state]，即 [1, 1, 768]。然后后面再接一个 conv1d 就是最终的输出。

mlp

接下来是一个 mlp 层：

def mlp(x, scope, n_state, *, hparams):
    nx = x.shape[-1]
    h = gelu(conv1d(x, 'c_fc', n_state))
    h2 = conv1d(h, 'c_proj', nx)
    return h2
def gelu(x):
    return 0.5*x*(1+tf.tanh(np.sqrt(2/np.pi)*(x+0.044715*tf.pow(x, 3))))

gelu 是个激活函数，其实就是两层 conv1d，最终 h2 的 shape 和输入的 x 其实一样，只不过中间的隐层这里为 4 倍的 n_state。

所有的 block 跑完后，最后接一个归一化层，就这样 Transformer 部分就执行完了。简单总结一下大致的结构：

若干个（n_layer）block + 归一化。
每个 block 包括一个 self_attention 层和 mlp 层，以及分别对应了残差连接。
self_attention 和 mlp 的输入都做了归一化。
self_attention 分为：conv1d => split_heads => multihead_attention => merge_heads => conv1d。
mlp 分为就是一个两层的 conv1d，其实就是一个感知机。

Language model loss

Transformer 之后接的是语言模型的损失函数：

1
2
3

h_flat = tf.reshape(h, [batch*sequence, hparams.n_embd])
logits = tf.matmul(h_flat, wte, transpose_b=True)
logits = tf.reshape(logits, [batch, sequence, hparams.n_vocab])

特点和创新

更加 general（更加庞大的数据集和参数）
下游任务 zero-shot

How

官方代码也没提供训练代码，下面这部分的代码部分来自 gpt-2-simple，注意这个并不适用于 Tensorflow 2.0 及以上版本（无法读取预训练模型），其他的倒是没有啥问题。

如何构造数据

输入的数据可以是最原始的纯文本，只不过进来后需要 Token 化，将文本变成 One-Hot 编码。参考官方的实现：

def encode(self, text):
    bpe_tokens = []
    for token in re.findall(pat, text):
        token = ''.join(byte_encoder[b] for b in token.encode('utf-8'))
        bpe_tokens.extend(encoder[bpe_token] for bpe_token in bpe(token).split(' '))
    return bpe_tokens
text = "I'm loving U."
# 原始 tokens: ['I', "'m", ' loving', ' U', '.']
# unicode tokens: ['I', "'m", 'Ġloving', 'ĠU', '.']
# bpe tokens: [40, 1101, 14442, 471, 13]

这个就是把原始的输入 text 转为 one-hot 编码的代码，Tokenize 的模型是 bpe，encoder 类似 word2id，作者这里利用了 unicode 编码，先将 token 换成 unicode 的统一编码，然后再利用 bpe 模型去 token 化。中文的处理方式类似。

如何开始训练

这里主要是定义好 loss function：

context = tf.compat.v1.placeholder(tf.int32, [batch_size, None])
output = model(hparams=hparams, X=context)
loss = tf.reduce_mean(
    input_tensor=tf.nn.sparse_softmax_cross_entropy_with_logits(
        labels=context[:, 1:],
        logits=output['logits'][:, :-1]))

label 其实就是输入 one-hot 的下一个 token，logits 可以看出没有取最后一个。这里 context 可以批量训练，序列的长度在 GPT-2 中是 1024。

如何使用结果

使用 gpt-2-simple 结合官方代码，代码如下：

import gpt_2_simple as gpt2 

models_dir = "/path/to/gpt-2/models/"
model_name = "124M"

sess = gpt2.start_tf_sess()
gpt2.load_gpt2(sess, model_name=model_name, model_dir=models_dir)

output = sample_sequence(
    hparams=hparams, length=50, start_token=enc.encoder['<|endoftext|>'], 
    batch_size=1, temperature=1, top_k=40, top_p=1)[:, 1:]

context_tokens = sampler.sample(1) # 取一个 token，比如：array([198])
out = sess.run(output, feed_dict={context: batch_size * [context_tokens]})
text = enc.decode(out[0])

如果需要进一步理解，可以仔细阅读下 gpt-2-simple 的代码，可以算是一个最佳实践了。在此基础上调整为中文的也非常方便。

数据和实验

文章一共做了八个项目的对比实验，分别是：

Language Modeling：
Children’s Book Test：考察 LM 在命名实体，名词，动词和介词等不同类别单词上的性能。
LAMBADA：测试系统对文本中的远程依存关系建模的能力。
Winograd Schema Challenge：测量系统解决文本歧义的能力来衡量系统执行常识推理的能力。
Reading Comprehension：测试阅读理解能力及模型根据历史对话记录回答问题的能力。
Summarization
Translation
Question Answering

内容着实有点多，就不一一列出来了，重点看一下 LM 的吧。

Discussion

特殊情况

训练集和测试集数据重叠问题。好的去重技术：可伸缩的模糊匹配（Scalable fuzzy matching），目前建议在切分训练集和测试集时，将基于 n-gram 重叠的重复数据删除方法作为重要的验证步骤和健全性检查。另一个确定模型性能是否由重复数据导致的方法是，根据数据本身的测试集（而不是另外的数据集）检查性能（其实这也可能有重复）。

打开脑洞

我：这 GPT-2 怎么看起来好像就是一个 Transformer 的 Language Model 啊？

小 A：可不是么，怎么有点大力出奇迹的感觉呢：更大的词表、更大的 batch size、更长的 context、更多的参数……然后结果就——更加 general……

小 B：也不能这么说，至少它证明了两点，第一，Transformer 的表征能力；第二，大模型的表征能力。而且也让我们的武器库里又多了一件不错的工具不是么？

小 C：确实如此，NLP 领域 LM 是最基本的模型，纵览发展历史，从 N-gram 到 RNN 再到 LSTM 再到 GPT-2，每个模型背后都隐藏着一次大的创新。现在是 Attention 的高光时刻。大家觉得下一个创新点可能在哪里？

小 A：我觉得肯定是某个新东西 + 大力出奇迹！

小 B：那你说这个新东西可能的方向是什么？

小 C：我觉得这样干想不现实，可能还是从人的认知方面去借鉴一些思想。我们再次概览历史，N-gram 可以看作局部建模，RNN 建模范围更广，LSTM 让建模范围更广的同时更加符合人的机制（遗忘、更新等），Transformer 则是关注更广范围内的重点信息，就像人会关注核心词和关键词一样。如果再从人的角度往下推，可能就是背景知识这些外部信息了，因为同样的文本在不同的背景下，我们关注的点可能不一样。也许可以可以借鉴 GloVe 的方式，也许是使用 Knowledge Graph。

小 B：我之前看过一篇关于 AI 建模体系的文章，作者对不同种类的 AI 方法进行了整合，最终的结论就是多种方法融合。不过那个有点大，而且也是一家之言。

小 A：你们想的可真多，我还就喜欢大力出奇迹，简单粗暴，就是干！

众人：……

Appendix

补充两篇我觉得不错的博客文章和本文可以相互补充：

Yam

Feeling, Coding, Thinking