利用transformer的decoder构建的语言模型

介绍

经典的语言模型是利用过去的信息，对当前的词进行表征，如何利用过去的信息是有技巧的，之前使用RNN模型，将之前时间步的信息利用隐藏层或者cell状态传递，现在我们可以利用transformer的decoder进行信息传递，因为GPT仍然是预训练语言模型，所以我们依然从预训练和fine tuning两个角度来介绍。
由于GPT的源码只有finetuning的部分，没有pretraing的部分，所以只能从finetuning的代码上面找到pretraning的蛛丝马迹。
finetuning使用的是一个rocstories数据集，是根据故事内容，判断两个结局里面到底是哪一个正确的。这里它是用内容和两个结局分辨拼出了两个长句子，然后对着两个句子做预测，然后根据正确的label构建损失函数。

预训练阶段

因为源码里面没有预训练的部分，只有一个finetuning的部分，因为finetuning的loss包含了语言模型的loss，所以大体上只能明白它是怎么利用语言模型的，但是细节上只能参考，实际上细节上和标准的语言模型差别不少。

# 1. input
X = tf.reshape(X, [-1, n_ctx, 2])

# 2. embedding
we = tf.get_variable("we", [n_vocab+n_special+n_ctx, n_embd], initializer=tf.random_normal_initializer(stddev=0.02))
we = dropout(we, embd_pdrop, train)
h = embed(X, we)

# 3. decoder of transformer
for layer in range(n_layer):
    h = block(h, 'h%d'%layer, train=train, scale=True)
    
# 4. loss
lm_h = tf.reshape(h[:, :-1], [-1, n_embd])
lm_logits = tf.matmul(lm_h, we, transpose_b=True)
lm_losses = tf.nn.sparse_softmax_cross_entropy_with_logits(logits=lm_logits, labels=tf.reshape(X[:, 1:, 0], [-1]))
lm_losses = tf.reshape(lm_losses, [shape_list(X)[0], shape_list(X)[1]-1])
lm_losses = tf.reduce_sum(lm_losses*M[:, 1:], 1)/tf.reduce_sum(M[:, 1:], 1)

输入

对于其中的某一个句子，一般是需要在开头和结尾加上特殊符号,以下是某个句子的举例，输入的序列移除</s>，目标序列移除<s>

$<s>,word1,word2,word3...word_n,</s>$

从它的finetuning代码里面以及后续，GPT的finetuning语料没有移除这个位置标签，只是在语言模型的loss部分做了一些处理，那么输入就是[batch_size，2，n_ctx, 2] 即先对所有的句子分别进行encode，n_ctx是句子长度，第一个2是finetuning预料中拼接出来的两个长句子，第二个2分别代表的是句子中词的id，以及位置的id。
模型拿进来就先做将第一个2合并到batch_size里面了，方便求语言模型的loss

embedding

使用了一个shape为[n_vocab+n_special+n_ctx, n_embd]的矩阵作为词嵌入矩阵

def embed(X, we):
    we = convert_gradient_to_tensor(we)
    e = tf.gather(we, X)
    h = tf.reduce_sum(e, 2)
    return h

其中e = tf.gather(we, X)获取的是[-1,n_ctx,n_embd]维度的tensor，这个操作目的是从we中，切出x指定位置的向量并组合起来，x表示的是词的id和位置的id，所以可以获取词的embedding向量和位置的embedding向量，然后将这两个向量相加作为最终向量，最终embedding的输出就是[-1,n_ctx,n_embd]，至于we的维度加上了n_special，应该是增加了n_special个特殊标记符号进去了，所以要增加词汇表数目。

transformer的decoder

def block(x, scope, train=False, scale=False):
    with tf.variable_scope(scope):
        nx = shape_list(x)[-1]
        a = attn(x, 'attn', nx, n_head, train=train, scale=scale)
        n = norm(x+a, 'ln_1')
        m = mlp(n, 'mlp', nx*4, train=train)
        h = norm(n+m, 'ln_2')
        return h

def attn(x, scope, n_state, n_head, train=False, scale=False):
    assert n_state%n_head==0
    with tf.variable_scope(scope):
        c = conv1d(x, 'c_attn', n_state*3, 1, train=train)
        q, k, v = tf.split(c, 3, 2)
        q = split_heads(q, n_head)
        k = split_heads(k, n_head, k=True)
        v = split_heads(v, n_head)
        a = _attn(q, k, v, train=train, scale=scale)
        a = merge_heads(a)
        a = conv1d(a, 'c_proj', n_state, 1, train=train)
        a = dropout(a, resid_pdrop, train)
        return a

def _attn(q, k, v, train=False, scale=False):
    w = tf.matmul(q, k)

    if scale:
        n_state = shape_list(v)[-1]
        w = w*tf.rsqrt(tf.cast(n_state, tf.float32))

    w = mask_attn_weights(w)
    w = tf.nn.softmax(w)

    w = dropout(w, attn_pdrop, train)

    a = tf.matmul(w, v)
    return a

def mask_attn_weights(w):
    n = shape_list(w)[-1]
    b = tf.matrix_band_part(tf.ones([n, n]), -1, 0)
    b = tf.reshape(b, [1, 1, n, n])
    w = w*b + -1e9*(1-b)
    return w

我们可以看到，这块的decoder代码和标准的transformer的decoder只缺少一个外部attention的结构，因为这里没有encoder的输出，也没有办法同encoder做attention，所以只需要做一个masked self-attention就可以了。
这里的mask_attn_weights里面的n就是指序列长度，b是一个下三角矩阵，用来防止信息泄露，w的shape是[batch_size,num_head,n_ctx,n_ctx]，这样的话，w与b相乘，将未来序列信息权重设置为0。但是paddingmask在这里没用啊。。。。不过反正输入序列的</s>都没有去掉，所以这里就不纠结了，这里绝对不是标准的语言模型

loss

        lm_h = tf.reshape(h[:, :-1], [-1, n_embd])
        lm_logits = tf.matmul(lm_h, we, transpose_b=True)
        lm_losses = tf.nn.sparse_softmax_cross_entropy_with_logits(logits=lm_logits, labels=tf.reshape(X[:, 1:, 0], [-1]))
        lm_losses = tf.reshape(lm_losses, [shape_list(X)[0], shape_list(X)[1]-1])
        lm_losses = tf.reduce_sum(lm_losses*M[:, 1:], 1)/tf.reduce_sum(M[:, 1:], 1)

在进行decoder之后，得到了各位置的编码信息h，shape为[-1,n_ctx,n_embd]，
之后将h转成我们需要预测的，因为最后的位置</s>是不需要的，所以要移除该token并reshape成每个词的位置，另外利用原先embedding lookup table的转置作为构建词汇logit的权重矩阵，这个技巧还是用的不少的，但是在google的官方transformer里面是没有的这样用的。
之后构建标签序列，标签序列就是移除第一个token，也就是<s>，形成这种错位的预测形式，也就是根据之前词的信息预测下一个词。
最后将各词的loss损失重新reshape成[-1，n_ctx]，这样就可以利用原先输入的mask矩阵M，移除padding部分的loss，至于为什么M要移除初始的<s>，这是因为实际我们预测标签是真实序列减1，只是需要移除一个实际位置，只是让padding为1的位置少一个就行了。
最后loss就是[batch_size,]，每个值是这个序列真实值的loss的平均，实际上在最后的loss反向传播中，loss又做了一次平均值。
这样就可以构建语言模型了。
怎么利用模型进行finetuning呢？

        clf_h = tf.reshape(h, [-1, n_embd])
        pool_idx = tf.cast(tf.argmax(tf.cast(tf.equal(X[:, :, 0], clf_token), tf.float32), 1), tf.int32)
        clf_h = tf.gather(clf_h, tf.range(shape_list(X)[0], dtype=tf.int32)*n_ctx+pool_idx)

        clf_h = tf.reshape(clf_h, [-1, 2, n_embd])
        if train and clf_pdrop > 0:
            shape = shape_list(clf_h)
            shape[1] = 1
            clf_h = tf.nn.dropout(clf_h, 1-clf_pdrop, shape)
        clf_h = tf.reshape(clf_h, [-1, n_embd])
        clf_logits = clf(clf_h, 1, train=train)
        clf_logits = tf.reshape(clf_logits, [-1, 2])

        clf_losses = tf.nn.sparse_softmax_cross_entropy_with_logits(logits=clf_logits, labels=Y)

这里是finetuning任务的loss计算，首先得到每个词语的embedding向量，然后得到pool_idx，这个pool_idx是每个长句子的clf_token的位置，最终得到的是[batch_size*2]的位置向量。之后我们要获取这些句子中clf_token位置的向量，得到[batch_size*2，n_embedding]。接下来要进行分类，因为两个长句子一起构成2分类，每个句子只需要得到一个维度为1的logits就行了，然后还原出原先的两个长句子，让这个两个长句子的logits构成一个二分类的logits，就可以拿去做二分类了。这里也算是涨了知识了。。。。
finetuning的loss是lm_loss和clf_loss的合并，加了个自定义的超参lm_coef，设定lm_loss的比例

train_loss = tf.reduce_mean(clf_losses) + lm_coef*tf.reduce_mean(lm_losses)

参考资料
GPT的finetuning