Moco训练方式的学习以及在文本表示方面的应用探索

最新推荐文章于 2023-12-26 22:16:58 发布

colourmind

最新推荐文章于 2023-12-26 22:16:58 发布

阅读量1.7k

点赞数 2

分类专栏： # 文本匹配和文本分类

本文链接：https://blog.csdn.net/HUSTHY/article/details/115028741

版权

文本匹配和文本分类专栏收录该内容

16 篇文章 16 订阅

订阅专栏

一、Moco——Momentum Contrast for Unsupervised Visual Representation Learning

二、Moco_bert文本匹配的尝试

最近看到2020年机器视觉里的一篇论文——Momentum Contrast for Unsupervised Visual Representation Learning——用于无监督视觉表示学习的动量对比，发现这样的训练方式很有意思。感觉也可以直接引用到NLP一些任务中，进行微调Bert系列模型，增强它的准确率。为此进行了论文的理解和一些简单的探索，写成一篇博客，记录一下。

一、Moco——Momentum Contrast for Unsupervised Visual Representation Learning

这是FaceBook AI 团队何凯明大神的又一力作，论文地址——https://arxiv.org/pdf/1911.05722.pdf，下面就是读论文的时间了。

首先就是要提一下对比学习——基本思想就是让相似的样本距离比非相似的样本更近，这里构造一种损失来度量；那么神经网络经过这样的对比学习，就模型区分能力就更加强大了。如果负样本对越多，对比效果就越好，模型的训练难度就越大，最后训练得到的模型具备的能力肯定就更强。这里怎么样才能够在训练过程中增加对比度，也就是增加负样本对的数目，同时不会依赖硬件GPU显存太多，就是一个比较难的问题。Moco给出了一个很好的答案。

Moco的整体结构

可以看到Moco整体架构用，对于输入的特征提取采取了2个子模型分别是encoder和momentum encoder；同时构造了一个队列，里面动态的存储K+1个key，其中第一个key与q构成了正样本对，其他的就是负样本对；最后采用对比损失来训练模型，让模型能够把q从K+1个key样本中精确的匹配到对应的正样本——也正是基于这样的思想，在做文本匹配的时候也可以构建这样的下游任务做微调，增强Bert的区分和表示能力。文章中有很多细节值得学习，首先就是这种对比损失；负样本的构建；动量更新key encoder参数，保持key的一致性；queue的设计解耦了mini-batch，能够让小机器也能基于这样的训练方式来进行训练。

对比损失函数

Moco提出的损失函数简称InfoNCE(With similarity measured by dot product, a form of a contrastive loss function),论文给出的公式如下：

表示一个正样本与负样本进行相似度计算，然后对对应的K+1类别进行softmax，最后取负的似然估计——就有点先计算相似度，然后在取交叉熵。

梯度更新参数和动量更新参数

moco框架中有2个encoder，K个样本特征提取的那个encoder不进行梯度回传，而是采用动量的更新的方式，把正样本Q对应的encoder的参数更新给它。这里的m往往取值比较大，能很好的保证k的一致性。

上图中的encoder使用了梯度回传，而momentum encoder这边没有使用梯度更新，直接采用动量赋值更新的方式。

以上就是论文中我任务从思想层面来说比较重要的内容，当然还有很多细节的地方值得深究，比如Batch shuffle, for making use of BatchNorm、以及queue的实现、单一encoder梯度更新的实现。这里也有很多pytorch代码值得学习。

论文给出的算法伪代码——采用的是pytorch和python的风格——太友好了，对于理解文章的思路简直不要太好。

这里主要是关注x_q和x_k其实是同一个x来自不同的增强方式，在本质上就是一样的，提取特征后q和k就构成了正样本对，q和队列中其他的k就构成了负样本对——这样就完美的实现了无监督的训练，这也是一个很好的创新。

损失函数果然就是交叉熵损失函数输入对比的相似度矩阵——对比损失函数。

然后就是encoder参数的两种更新方式以及队列的入队和出队，非常清晰！

二、Moco_bert文本匹配的尝试

之前做过文本匹配——当时模型训练采用的是基于sentence_bert来做分类任务微调的，标注的句子对只有2W条；现在可以把title_a与abstract_a作为正样本，title_a和其他的abstracts作为负样本对，做一个对比学习训练进行微调。

Moco_bert

直接把Moco中的encoder换成了bert，其他的几乎不变，代码如下：

import torch
import torch.nn as nn
from transformers import BertModel
import copy

class Bert_moco(nn.Module):
    def __init__(self,dim=768,K=5120,m=0.999,T=0.07,bert_path='bert_models/roberta',device=None):
        """
        dim: feature dimension (default: 768)
        K: queue size; number of negative keys (default: 5120)
        m: moco momentum of updating key encoder (default: 0.999)
        T: softmax temperature (default: 0.07)
        """
        super(Bert_moco,self).__init__()
        self.device = device
        self.K = K
        self.m = m
        self.T = T
        bert = BertModel.from_pretrained(bert_path)
        self.encoder_title = bert
        self.encoder_abstract = copy.deepcopy(bert)
        for param_title,param_abstract in zip(self.encoder_title.parameters(),self.encoder_abstract.parameters()):
            param_abstract.data.copy_(param_title.data)# initialize
            param_abstract.requires_grad = False #不支持梯度更新
        
        #创建一个queue
        self.register_buffer("queue", torch.randn(dim, K))
        #对queue进行初始化
        self.queue = nn.functional.normalize(self.queue, dim=0)
        #创建一个queue_ptr
        self.register_buffer("queue_ptr", torch.zeros(1, dtype=torch.long))

    def momentum_update_abstract_encoder(self):
        #动量更新encoder_abstract的参数，采用比较大的m可以使得参数更新的比较缓慢，保持队列中的embedding变化的不是那么快速
        for param_title,param_abstract in zip(self.encoder_title.parameters(),self.encoder_abstract.parameters()):
            param_abstract.data = param_abstract.data * self.m + param_title * (1. - self.m)


    def dequeue_and_enqueue(self,keys):
        """
        队列的出队和入队
        """
        batch_size = keys.shape[0]
        ptr = int(self.queue_ptr)
        assert self.K % batch_size == 0  # for simplicity

        #队列入队，这种替换方式很好
        self.queue[:,ptr:ptr+batch_size] = keys.T
        #队列出队
        ptr = (ptr + batch_size) % self.K
        self.queue_ptr[0] = ptr
    
    def forward(self,title_input_ids,title_mask_attentions,abstract_input_ids,abstract_mask_attentions):

        title_embeddings = self.encoder_title(title_input_ids,title_mask_attentions)[0]
        title_embeddings = title_embeddings.mean(dim=1)
        title_embeddings = nn.functional.normalize(title_embeddings,dim=1)

        with torch.no_grad():
            self.momentum_update_abstract_encoder()
            abstract_embeddings = self.encoder_abstract(abstract_input_ids,abstract_mask_attentions)[0]
            abstract_embeddings = abstract_embeddings.mean(dim=1)
            abstract_embeddings = nn.functional.normalize(abstract_embeddings,dim=1)

        #[N,1]，单位向量计算内积a*b就是计算相似度
        l_pso = torch.einsum('nd,nd->n',[title_embeddings,abstract_embeddings]).unsqueeze(-1)
        #[N,K]
        l_neg = torch.einsum('nd,dk->nk',[title_embeddings,self.queue.clone().detach()])
        #logits---[N,1+k]
        logits = torch.cat([l_pso,l_neg],dim=1)
        logits /= self.T

        #[N]---batch_size个0，每个样本对应的匹配abstract都是第一个
        labels = torch.zeros(logits.shape[0],dtype=torch.long).to(self.device)

        self.dequeue_and_enqueue(abstract_embeddings)
        return logits,labels

值得学习的细节就是，这里的队列维持，直接采用了torch中的register_buffer来实现的——维持一个持久缓冲区——在显存中的一个矩阵，可以像tensor一样操作；另一个实现的细节上就是with torch.no_grad()和tensor().detach()联合起来实现不传播梯度；一些超参数都是使用论文默认的那些。

这里有一个问题，就是这里构建的训练任务，随着训练的进行会出现退化的问题。

使用roberta base和large微调过程