transformer 翻译模型及简明源码

yichudu

5832人浏览 · 2019-04-21 23:33:36

yichudu · 2019-04-21 23:33:36 发布

google 出品, 基于 encoder-decoder 架构的 seq2seq 自然语言翻译模型. paper 见参考[1].
创新之处是弃用 RNN,CNN 这样的结构, 采用 self-attention 与 positional-encoding 来加速训练并提升效果, 取得了 SOTA 成绩. 是后续 GPT 及 LLM 系列的基石之作!

1. BPE 分词

Byte Pair Encoding, 是一种子词（subword）分词算法. 特点是数据驱动, 先通过语料库构建词汇表(vocabulary), 然后用它做输入文本的贪心最长匹配.
优点是词典大小可控, 可有效处理未知词（OOV）和词形变化, 因此广泛用于 GPT, Bert 及后续的 LLM 中.

1.1 训练

从字符级别开始，逐步合并高频的符号对，形成常见的子词单元。

所有单词从字符级别开始。
统计相邻符号（symbol）对的出现频率。
将最频繁的符号对合并为一个新的符号。
重复该过程，直到达到预定的合并次数或词汇表大小。

1.2 预测

贪心最长匹配（Greedy Longest Match）.
从左到右, 找到存在于词汇表中的最长子串.

#例子: 
from transformers import GPT2Tokenizer
tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
tokens = tokenizer.tokenize("drawing badminton")
print(tokens) # 输出: ['draw', 'ing', 'bad', 'm', 'inton']

2. encoder 结构

在这里插入图片描述
灰色的框叫 block, decoder block 比 encoder block 多了一次 masked-attn 处理.

2.1 inputs

word -> token_id -> lookup ->emb

2.2 Positional Encoding

将token所在seq的位置信息加以编码表示. 可以是 {sin/cos三角函数, embedding中学习} 两种方案, ablation study 显示差异不大.

2.3 self-attention

source 与 target 是两个 sequence, self-attention 是指 对同一 sequence 内, 各个 token 间的关联做计算.
论文中把提出的 attention 叫 Scaled Dot-Product Attention.
引入三个共享参数矩阵, Query,Key,Value, 对 embedding 后的 token 做表达的学习.
所谓 attention, 就是这么一个映射关系, $output=f(query, \{key , value\} )$ .
key 和 value 的长度必须一致, query 则不要要求.
$Attention(Q,K,V)=softmax(\frac{QK^T}{\sqrt{d_k}})V \tag 1$

图2. $output=f(query, \{key\_value\} )$ 的直观图示

$MultiHead(Q,K,V)=Concat(head_1,...,head_h)W^O$
where $head_i=Attention(QW_i^Q,KW_i^K,VW_i^V)$ , $W^O\in \mathbb R^{hd_v\times d_{model}}$

2.4 feed-forward (add&norm)

是图片中 feed-forward 的超集, 也包含了 add&norm.
其中 add 残差的操作与时机可以参考

# feed-forward, add & norm
    def call(self, inputs, **kwargs):
        # Inner layer, ReLU
        outputs = Dense(self.num_units[0], activation=tf.nn.leaky_relu)(inputs)
        # Outer layer, linear
        outputs = Dense(self.num_units[1])(outputs)
        # Residual connection
        outputs += inputs
        # Normalize
        outputs = utils.modules.layer_norm(outputs)

        return outputs

3. decoder 结构

3.1 cross attention

除了 source_seq 与 target_seq 各自的 self_attrntion 之外, 还有二者之间的 cross attention (旧称 vanilla_attention). 对应图1 transformer 架构图中, decoder内的第二个 sub-layer. 此时有 vanilla_attention(decoder_Q, encoder_memory_K, encoder_memory_V).

3.2 causal mask

旧称 future mask.
decoder中有个特别点就是masking, masking 的作用就是防止在训练的时候使用未来时刻才输出的单词。比如训练时，第一个单词是不能参考第二个单词的生成结果的(但可以参考自己)。Masking就会把这个信息变成 0，用来保证预测位置 i 的信息只能基于比 i 小的输出。
下面章节的tf实现部分, 会有下三角矩阵的api说明.

3.3 logits 计算

与 vocab_emb 计算内积. 这种做法有以下几个优点：

参数共享：共享这两个层之间的权重可以减少模型中的总参数数量，有助于缓解过拟合，并且在实践中有时能带来更好的性能。
对称性：输入和输出使用相同的表示空间，这在某种程度上保持了模型结构的对称性。
初始化优势：如果使用预训练的嵌入，则这些预训练的权重也可以用于初始化最后的线性层，从而可能加速收敛。

四. train 与 predict 不同的 inference

4.1 train 与 loss

train 阶段, 直接输出整体的 logits #shape=(N, T2, vocab_size) , 与y_ = label_smoothing(tf.one_hot(y, depth=self.hp.vocab_size)) # (N, T2, vocab_size) 计算交叉熵.

4.2 predict 与 beam-search

在预测阶段, 根据 source sentence 的中间语义表示C和之前已经生成的历史信息, predicts $\hat y$ auto-regressively.
$y_i=g(C,y_1,y2_,...,y_{i−1})$
逐个位置选 token 时, 会选中对数概率最大的那个 token, 属于贪心策略. 贪心能保证局部最优但可能错过全局最优. 可以类比动态规划的背包问题.

所以引入 beam-search . 可以改善局部最优, 且增加翻译风格的多样性. 具体来说, 每个位置挑选token时, 选 top-n, 然后剪枝, 保留 n 个候选句子.
Q:那么最终如何从 n 个句子里挑最优呢?
A: 根据每个句子生成每个 token的对数概率, 作连乘或平均, 当作句子整体的得分来比较.

五. tf 实现

没有纯净的官方 tf 实现. 因为官方基于 tf, 又搞了 tensor2tensor 库, 这里面才有 transformer 的实现.

不过, github 上有人基于 tf 给出了实现, 见参考 [3]. 但它是完全的面向过程风格, 没有 keras 面向对象的风格. 不好用.

数据集
Germany_2_English 翻译任务. 为了单机体验完整流程, 控制数据量. 训练集有 19万样本, 验证集有 993 条样本. 英德的word加起来构成了 size=8000 的 vocabulary.
超参
$vocab\_size$ =32000, $d\_model$ =512, $max\_len\_source=100$ , $max\_len\_target=100$ .

input/output

 	xs: tuple of
        x: int32 tensor. (N, T1) . T1为当前batch样本的最长原文长度.
        x_seqlens: int32 tensor. (N,)
        sents1: str tensor. (N,)
    ys: tuple of
        decoder_input: int32 tensor. (N, T2). T2为当前batch样本的最长译文长度.
        y: int32 tensor. (N, T2)
        y_seqlen: int32 tensor. (N, )
        sents2: str tensor. (N,)

主干逻辑

def train(self, xs, ys):
    '''
    Returns
    loss: scalar.
    train_op: training operation
    global_step: scalar.
    summaries: training summary node
    '''
    # forward
    memory, sents1 = self.encode(xs)
    logits, preds, y, sents2 = self.decode(ys, memory)

    # train scheme
    y_ = label_smoothing(tf.one_hot(y, depth=self.hp.vocab_size))
    ce = tf.nn.softmax_cross_entropy_with_logits_v2(logits=logits, labels=y_)
    nonpadding = tf.to_float(tf.not_equal(y, self.token2idx["<pad>"]))  # 0: <pad>
    loss = tf.reduce_sum(ce * nonpadding) / (tf.reduce_sum(nonpadding) + 1e-7)

future mask

diag_vals = tf.ones_like(inputs[0, :, :])  # (T_q, T_k)
# 详见, https://www.tensorflow.org/api_docs/python/tf/linalg/LinearOperatorLowerTriangular?version=stable, 对角线的元素被保留, 不置0.
# 但<=1.4 版本的tf, 是没有这个api的.  可以用 tf.matrix_band_part(input, -1, 0) 替代
tril = tf.linalg.LinearOperatorLowerTriangular(diag_vals).to_dense()  # (T_q, T_k)
masks = tf.tile(tf.expand_dims(tril, 0), [tf.shape(inputs)[0], 1, 1])  # (N, T_q, T_k)

paddings = tf.ones_like(masks) * padding_num
outputs = tf.where(tf.equal(masks, 0), paddings, inputs)

六. torch 实现

一个流行的实现见参考 [8].
logits 计算有两个选择: 1)与 vocab_emb 计算内积; 2)新开一个 linear. 可通过下面的配置选择.

class Transformer(nn.Module):
    def __init__():
        self.trg_word_prj = nn.Linear(d_model, n_trg_vocab, bias=False)

        if trg_emb_prj_weight_sharing:
            # Share the weight between target word embedding & last dense layer
            self.trg_word_prj.weight = self.decoder.trg_word_emb.weight