MT-DNN解读（论文 + PyTorch源码）

前一段时间，看到微软发布了用于学习通用语言嵌入的多任务深度神经网络模型MT-DNN，可谓是紧随BERT之后，结合他们之前的MTL工作以及BERT的优势，在10项NLU任务上的表现都超过了BERT。

PS：预感到BERT的起飞又将给NLP领域带来新一波论文热潮（把之前的东西加上BERT再train一遍？效果飞起？发论文！当然这里没有任何diss这个MT-DNN模型的意思哈哈，毕竟也得底子好是吧，不然真废柴加上BERT也没用~）

一. 前言

预训练（加源头）和多任务学习（加目标）都是能用来提升效果的手段。

预训练近期的进展，如BERT、GPT等，大家都有目共睹，就不多说了~

对于MTL（Multi-task Learning，多任务学习）来说，其优点有两个：1）弥补了有些任务的数据不足问题；2）有正则的作用，防止模型过拟合。

论文中作者认为，MTL和pretrain有很好的互补作用，那么是不是可以结合一下，发挥两者的作用。更具体的就是，先用BERT进行pretrain，然后用MTL进行finetune，这就形成了MT-DNN。可见，与BERT的不同在于finetune的过程，这里用MTL作为目标。

换个角度来想，其实是在BERT没有出来的时候，是直接训练MTL的模型，现在BERT出来了，那就拿这个初始化试试？

二. MT-DNN原理

1. 多任务介绍

在讲MT-DNN之前，先来唠一唠有哪些任务？因为毕竟是Multi-Task的学习，总要知道Task是哪些吧！

MT-DNN是结合了4种类型的NLU任务：单句分类、句子对分类、文本相似度打分和相关度排序。下面举一些GLUE中的例子：

• 单句分类：比如CoLA是判断这句话是否语法合适、SST-2是判断这句话的情感
• 文本相似度：比如STS-B是为两句话进行相似度打分
• 句子对分类：比如RTE和MNLI是文本蕴含任务，QQP和MRPC是判断两句话是否语义上一致
• 相关性排序：比如QNLI，它其实在原始的GLUE任务中定义为二分类问题，但这里论文给它定义成排序问题，自己采样了一堆负样本，然后用softmax学习排序

2. 模型结构

看下面这个图：

其实很直观，下面的shared layers是BERT，上面的task specific layers是MTL。BERT的部分就不多说了，还不了解的读者戳笔者之前的博客。这里主要讲一下用于适配各个任务的MTL部分。

1. 单句的分类任务

用[CLS]的表征作为特征，设为$x$，则对于单句的分类任务，直接在后面接入一个分类层即可，以SST-2任务为例：

$$P_r(c|X)=softmax(W_{SST}^T\cdot x)$$

loss就是分类的交叉熵，即：

$$-\sum _{c}I(X,c)log(P_r(c|X))$$

2. 句子相似度

以STS-B任务为例，将两句话pack后送进去，得到的[CLS]的表征，可拿出来计算分数：

$$Sim(X_1,X_2)=sigmoid(w_{STS}^T\cdot x)$$

loss采用MSE损失，即：

$$(y-Sim(X_1,X_2){)}^2$$

3. 句子对分类

以NLI任务为例，这里接的是SAN网络，一个在这个任务上表现得比较好的网络。

SAN的计算流程为：

• 输入premise $P=(p_1,...,p_m)$ 和 hypothesis $H=(h_1,...,h_n)$
• 通过BERT得到premise和hypothesis的表示，$M^p={\mathbb{R}}^{d\ast m}$和$M^h={\mathbb{R}}^{d\ast n}$
• 开始K步的推理，初始状态$s^0$是$M^h$的self-attention，$s^0={\sum }_j{\alpha }_j{M}_j^h$，这里的${\alpha }_j=\frac{exp(w_1^T\cdot M_j^h)}{{\sum }_{i}exp(w_1^T\cdot M_i^h)}$，然后对于第k步，状态变化为$s^k=GRU(s^{k-1},x^k)$，这里的$x^k$计算方式为$x^k={\sum }_j{\beta }_j{M}_j^p,{\beta }_j=softmax(s^{k-1}{W}_2^T{M}^p)$
• 最后接一层分类层，用于捕捉每一步的推理结果：$P_r^k=softmax(W_3^T[s^k;x^k;|s^k-x^k|;s^k\cdot x^k])$
• 最后分数是所有推理结果的平均：$P_r=avg([P_r^0,P_r^1,...,P_r^{K-1}])$

loss仍采用分类的交叉熵，即：

$$-\sum _{c}I(X,c)log(P_r(c|X))$$

4. 相关性排序

以QNLI为例，这里主要是先计算两个句子之间的相似度，输入两个句子pack，采用[CLS]的输出作为表征。

$$Rel(Q,A)=g(w_{QNLI}^T\cdot x)$$

loss采用排序损失：

$$-\sum _{Q,A^{+}}{P}_r(A^{+}|Q)$$

$$P_r(A^{+}|Q)=\frac{exp(\gamma Rel(Q,A^{+}))}{{\sum }_{A^{\prime }\in A}exp(\gamma Rel(Q,A^{\prime }))}$$

3. 训练流程

见下图：

这里的Eq.6~Eq.8，读者就对应到前面相应任务的loss公式即可。

三. 实验

1. 数据集

评估的数据集还是很多的，与BERT论文重合的也比较多，详细信息见下表：

这里除了WNLI、SNLI和SciTail，都是BERT里面评估的GLUE数据集。

2. 实验结果

在GLUE上面的表现如下：（这里应该是用了9个GLUE任务进行的MTL？）

这里的MT-DNN使用的是BERT_LARGE作为base model。

在SNLI和SciTail上的表现如下：（这里应该是用了这11个任务进行的MTL？）

PS：感觉论文里面并没有说清楚，在做各个实验的时候，是用了哪些任务进行MTL，上面是笔者根据蛛丝马迹猜测的，若有大神看到了，还烦请解答~

3. 对比实验

首先是对比了使用MTL的作用：

这里的ST-DNN是用BERT_BASE作为预训练，然后单独finetune各个任务的结果，但在某些任务（如MNLI、QQP和MRPC）上使用了SAN，与BERT那么简单的finetune模型不一样。

然后是对比了一些domain adaption的结果，主要还是SNLI和SciTail这两个数据集，首先在8个GLUE任务上进行finetune（没有用WNLI，为了防止与这两个数据集有重合任务），然后再用各自的数据进行finetune，最后评估的效果如下：

这里不是ZSL的设定，只是域适应，所以是在预训练的时候没有用到目标域的数据，但finetune的时候，是允许用目标域数据进行finetune的。

作者还列举出了在进行域适应的时候，不同的模型需要的数据量的比较：

四. PyTorch实现

看了MT-DNN的源码，也是较为详细的。包括MT-DNN的pretrain，以及一些任务的finetune，同时也提供了pretrain好的模型。下面将分为pretrain和finetune两部分来剖析：

1. pretrain

首先来看MT-DNN模型的构建：（这里只列举了核心的代码）

class SANBertNetwork(nn.Module):
    def __init__(self, opt, bert_config=None):
        super(SANBertNetwork, self).__init__()
        self.bert_config = BertConfig.from_dict(opt)
        self.bert = BertModel(self.bert_config)
        mem_size = self.bert_config.hidden_size
        self.decoder_opt = opt['answer_opt']
        self.scoring_list = nn.ModuleList()
        labels = [int(ls) for ls in opt['label_size'].split(',')]

        for task, lab in enumerate(labels):
            decoder_opt = self.decoder_opt[task]
            if decoder_opt == 1:
                out_proj = SANClassifier(mem_size, mem_size, lab, opt, prefix='answer', dropout=dropout)
                self.scoring_list.append(out_proj)
            else:
                out_proj = nn.Linear(self.bert_config.hidden_size, lab)
                self.scoring_list.append(out_proj)

    def forward(self, input_ids, token_type_ids, attention_mask, premise_mask=None, hyp_mask=None, task_id=0):
        all_encoder_layers, pooled_output = self.bert(input_ids, token_type_ids, attention_mask)
        sequence_output = all_encoder_layers[-1]
        decoder_opt = self.decoder_opt[task_id]
        if decoder_opt == 1:
            max_query = hyp_mask.size(1)
            assert max_query > 0
            assert premise_mask is not None
            assert hyp_mask is not None
            hyp_mem = sequence_output[:,:max_query,:]
            logits = self.scoring_list[task_id](sequence_output, hyp_mem, premise_mask, hyp_mask)
        else:
            pooled_output = self.dropout_list[task_id](pooled_output)
            logits = self.scoring_list[task_id](pooled_output)
        return logits

可见，是首先用了BERT的结构作为底层的encoder，然后在上面接task-specific的结构，这里的labels表示类别数，decoder_opt表示是否使用SAN网络。

PS：这里发现一个在PyTorch中调用BERT的方法，huggingface的pytorch-pretrained-bert居然是可以直接pip安装使用的，也太方便了吧，而且还有gpt,gpt2和transformer-xl！感谢神人~

下面是SAN网络的实现：

class SANClassifier(nn.Module):
    """Implementation of Stochastic Answer Networks for Natural Language Inference, Xiaodong Liu, Kevin Duh and Jianfeng Gao
    https://arxiv.org/abs/1804.07888
    """
    def __init__(self, x_size, h_size, label_size, opt={}, prefix='decoder', dropout=None):
        super(SANClassifier, self).__init__()
        if dropout is None:
            self.dropout = DropoutWrapper(opt.get('{}_dropout_p'.format(self.prefix), 0))
        else:
            self.dropout = dropout
        self.prefix = prefix
        self.query_wsum = SelfAttnWrapper(x_size, prefix='mem_cum', opt=opt, dropout=self.dropout)
        self.attn = FlatSimilarityWrapper(x_size, h_size, prefix, opt, self.dropout)
        self.rnn_type = '{}{}'.format(opt.get('{}_rnn_type'.format(prefix), 'gru').upper(), 'Cell')
        self.rnn =getattr(nn, self.rnn_type)(x_size, h_size)
        self.num_turn = opt.get('{}_num_turn'.format(prefix), 5)
        self.opt = opt
        self.mem_random_drop = opt.get('{}_mem_drop_p'.format(prefix), 0)
        self.mem_type = opt.get('{}_mem_type'.format(prefix), 0)
        self.weight_norm_on = opt.get('{}_weight_norm_on'.format(prefix), False)
        self.label_size = label_size
        self.dump_state = opt.get('dump_state_on', False)
        self.alpha = Parameter(torch.zeros(1, 1), requires_grad=False)
        if self.weight_norm_on:
            self.rnn = WN(self.rnn)

        self.classifier = Classifier(x_size, self.label_size, opt, prefix=prefix, dropout=self.dropout)

    def forward(self, x, h0, x_mask=None, h_mask=None):
        h0 = self.query_wsum(h0, h_mask)
        if type(self.rnn) is nn.LSTMCell:
            c0 = Variable(h0.new(h0.size()).zero_())
        scores_list = []
        for turn in range(self.num_turn):
            att_scores = self.attn(x, h0, x_mask)
            x_sum = torch.bmm(F.softmax(att_scores, 1).unsqueeze(1), x).squeeze(1)
            scores = self.classifier(x_sum, h0)
            scores_list.append(scores)
            # next turn
            if self.rnn is not None:
                h0 = self.dropout(h0)
                if type(self.rnn) is nn.LSTMCell:
                    h0, c0 = self.rnn(x_sum, (h0, c0))
                else:
                    h0 = self.rnn(x_sum, h0)
        if self.mem_type == 1:
            mask = generate_mask(self.alpha.data.new(x.size(0), self.num_turn), self.mem_random_drop, self.training)
            mask = [m.contiguous() for m in torch.unbind(mask, 1)]
            tmp_scores_list = [mask[idx].view(x.size(0), 1).expand_as(inp) * F.softmax(inp, 1) for idx, inp in enumerate(scores_list)]
            scores = torch.stack(tmp_scores_list, 2)
            scores = torch.mean(scores, 2)
            scores = torch.log(scores)
        else:
            scores = scores_list[-1]
        if self.dump_state:
            return scores, scores_list
        else:
            return scores

与论文中列举的公式一致。

接下来看loss的构建：

logits = self.mnetwork(*inputs)
if batch_meta['pairwise']:
    logits = logits.view(-1, batch_meta['pairwise_size'])
if task_type > 0:
    loss = F.mse_loss(logits.squeeze(), y)
else:
    loss = F.cross_entropy(logits, y)

这里是用到了两种loss，实际上论文中是有三种loss，这里是将分类的交叉熵loss和rank-loss都实现为多分类的交叉熵。

总结起来看，源码中一共是用mnli,rte,qqp,qnli,mrpc,sst,cola和stsb这8个任务进行MTL，其中mnli,rte,qqp,mrpc都是输入为pair的分类任务，它们用的是SAN网络+交叉熵损失的形式；qnli是输入为pair的rank任务，它用的也是SAN网络+交叉熵损失的形式，但需要对输出做一个变换，才能将rank_loss转为交叉熵的形式；stsb是输入为pair的分数预测（回归）任务，它用的是普通的Linear映射+mse损失的形式；sst,cola是输入为single的分类任务，它用的也是普通的Linear映射+交叉熵损失的形式。

同时，在MT-DNN的预训练过程中，是每次先根据各个任务中的数据量分布对任务进行采样，而后取那个任务的一个batch数据出来进行训练。

2. finetune

finetune的部分就是用各自任务的数据，对应各自任务上已经训好的MT-DNN模型进行微调即可，模型和损失等和前面的pretrain一样，这里就不再赘述。

整体来看，流程就是：

1. 先用BERT预训练好的模型做一个大的打底
2. 然后在上面接各种任务自己的task-model，并进行多任务的训练，得到MT-DNN
3. 最后再分别利用各个任务的数据，在前面得到的MT-DNN上进行finetune，得到最终的结果（如果不是用的类似之前多任务的数据进行finetune的话，就可以看做是域适应了）

五. 总结

优势

1. 同时结合了BERT和MTL的思想，两者互补，效果确实好！
2. 整个流程都十分的清晰，没有特别难以理解的地方
3. 方便扩展，可以用更多的task，进行MTL
4. 源码、预训练模型都比较充足

不足

1. 有堆砌的感觉，，虽然确实非常有用
2. 在做实验的时候，其实对标BERT有些不公平，因为BERT就是在单纯的凸显它的预训练模型效果，在接入下游任务的时候，基本没有加太多额外的参数，而MT-DNN在接入下游任务的时候，不管是从模型结构还是目标函数上，都有做相应的优化，虽然也有对比实验，但还是感觉整体思路不那么clean。而且后面大家随便接入一个更好的下游模型，提升了效果，是不是就可以又发一篇论文？（一些浅见而已，不要太当真。。）

传送门

论文：https://arxiv.org/pdf/1901.11504.pdf
源码：https://github.com/namisan/mt-dnn
博客：https://mp.weixin.qq.com/s?__biz=MzAwMTA3MzM4Nw==&mid=2649447160&idx=2&sn=db35c6a7b2c5e7c45fb09a137adf6b17&chksm=82c0b37cb5b73a6a1f57fee2c281c299799508b803b49e5350828c2a203a14157dc6815445c6&mpshare=1&scene=1&srcid=#rd