mT5模型是T5模型在多语言数据集C4上的继续训练,T5本身是比较早了,是2019年的一个模型,但是后来又有很多次不同的升级。
mT5模型论文发布自2020年10月。
论文:mT5: A massively multilingual pre-trained text-to-text transformer
代码:https://github.com/google-research/multilingual-t5
根据预训练模型进行Prompt-tuning是最近两年的研究热点,从开始的各种搜索词表的hard-prompt-tuning,到只训练一部分的soft-prompt-tuning,各种如何调整prompt的方法,可以参考这篇论文综述: Pre-train, Prompt, and Predict: A Systematic Survey of Prompting Methods in Natural Language Processing
本文主要介绍其中根据叠加prompt embedding进行soft-prompt-tuning如何实现。
具体实现的主要思路参考了这篇2021年4月的论文The Power of Scale for Parameter-Efficient Prompt Tuning而代码主要参考了https://github.com/kipgparker/soft-prompt-tuning
首先是prompt tuning,就是我们现在已经有了很多类似bert/gpt/t5这样的预训练模型,但是它们都很大,可能有几十亿或者几百亿个参数。
那这些模型进行一般的fine-tune本身就很耗费资源,而一个公司或者项目可能需要很多不同的任务基于这些模型,每个任务又需要单独部署,也很耗费资源。
有没有什么办法能让训练和部署的资源压力更小呢?
我们通过gpt-2/gpt-3的一些zero-shot/few-shot learning为启发点,就是我们给不同任务不同的提示(prompt),就可以很少的训练就可以在同一个模型完成不同的任务。
最开始,这些prompt就是一些特殊的句子,比如说我们给gpt3的提示是:“1+1=2;2+2=4;4+5=9;5+6=”这样的提示,让模型继续生成,希望能输出5+6的正确答案。
这种基于句子或者具体token的prompt,被称为hard-prompt。
那么我们怎么找到每个任务的最好的prompt呢?当然我们可以人工设计去一点一点尝试,或者干脆穷举,当然也有很多基于不同方法的测试,可以参考上面提到的综述论文。
除了hard-prompt以外,假设我们想要更好的结果,是不是我们可以设计一些特殊的类似token的东西,我们称之为prompt embedding,把它和原来模型本身的embedding对齐。 这样利用self-attention的机制,让模型本身可以读取我们加入的embedding。
而训练,就是只更新我们加入的prompt embedding的梯度,也就是只训练模型的非常小的一部分参数,而不去更新整体模型。
这样显然模型训练用的资源就会更好,部署的时候我们也只需要部署一个原版模型,只需要根据不同任务,插入不同的prompt embedding就好了。
所以我们需要:
我们说了我们要做中文的分类任务,mT5这样的encoder-decoder结构其实天然的做的是sequence-to-sequence结构,类似机器翻译/对话聊天之类的
那么分类任务怎么设计的
我们先定位任务为,输入一句(段)中文文本,输出一个三分类的标签,0,1,2。
首先输入中文文本,也就是把中文文本作为encoder的输入肯定没问题。
decoder的输入,也没什么好说的,毕竟我们不是seq2seq任务,不需要特殊的输入。
在代码中的encoder和decoder的输入中,加入了prompt embedding的占位符,是任意token id都可以,反正都会被我们的代码替换掉的。
我们要获取的是decoder最后的输出,并把输出中的一些特殊字符位置,当作我们的三分类结果。
decoder默认肯定会输出一个词表长度的向量,我们只拿其中3个使用,实际代码中我是使用3,4,5,三个特殊token id作为判定三分类的结果。
我们具体计算loss也只计算最后decoder输出的这三个token的概率比较,比如3的概率最大,那么就是分类0,4的大就是分类1,5的大就是分类2。
首先使用tansformers就可以很方便的去下载和调用谷歌的T5/mT5模型
安装pytorch和transformers,以及分词器(tokenizer)所需的sentencepiece
pip install transformers SentencePiece torch
然后就可以执行代码自动从网上下载模型:
from transformers import MT5ForConditionalGeneration, T5Tokenizer
model = MT5ForConditionalGeneration.from_pretrained("google/mt5-base")
tokenizer = T5Tokenizer.from_pretrained("google/mt5-base")
然后我们构建一个替换原版模型的输入器,用来把用于训练的prompt embedding加入到模型。
下面代码主要参考https://github.com/kipgparker/soft-prompt-tuning这个repo进行修改,因为这个repo的训练时基于GPT的,而我们是基于mT5的,所以整体代码上略有区别,而且这个repo的训练代码也不太完整。
import torch
import torch.nn as nn
class SoftEmbedding(nn.Module):
def __init__(self,
wte: nn.Embedding,
n_tokens: int = 10,
random_range: float = 0.5,
initialize_from_vocab: bool = True):
"""这个类用来给模型附加一个用于学习的embedding
Args:
wte (nn.Embedding): 这个参数,是预训练模型的embedding,载入进来用来提取一些参数。
n_tokens (int, optional): number of tokens for task. Defaults to 10.
random_range (float, optional): range to init embedding (if not initialize from vocab). Defaults to 0.5.
initialize_from_vocab (bool, optional): initalizes from default vocab. Defaults to True.
"""
super(SoftEmbedding, self).__init__()
self.wte = wte
self.n_tokens = n_tokens
self.learned_embedding = nn.parameter.Parameter(self.initialize_embedding(wte,
n_tokens,
random_range,
initialize_from_vocab))
def initialize_embedding(self,
wte: nn.Embedding,
n_tokens: int = 10,
random_range: float = 0.5,
initialize_from_vocab: bool = True):
"""初始化学习向量
Args:
same as __init__
Returns:
torch.float: initialized using original schemes
"""
# 有两种初始化方式,一种是从预训练模型copy一部分token,进行训练
# 另一种是随机生成一部分训练
# 结果上来说区别不大
if initialize_from_vocab:
return self.wte.weight[:n_tokens].clone().detach()
return torch.FloatTensor(n_tokens, wte.weight.size(1)).uniform_(-random_range, random_range)
def forward(self, tokens):
"""run forward pass
Args:
tokens (torch.long): input tokens before encoding
Returns:
torch.float: encoding of text concatenated with learned task specifc embedding
"""
input_embedding = self.wte(tokens[:, self.n_tokens:])
learned_embedding = self.learned_embedding.repeat(input_embedding.size(0), 1, 1)
# 把我们新加入的固定长度的,用于代表任务的prompt embedding,和实际的embedding合并
return torch.cat([learned_embedding, input_embedding], 1)
然后我们把我们设计的这个类,载入到预训练模型里面去:
n_tokens = 100
s_wte = SoftEmbedding(model.get_input_embeddings(),
n_tokens=n_tokens,
initialize_from_vocab=True)
# 用我们设计的类,替换原来的embedding层
model.set_input_embeddings(s_wte)
if torch.cuda.is_available():
model = model.cuda()
# 把除了第0个,就是我们要训练的prompt embedding以外的参数,都设置为不需要梯度
parameters = list(model.parameters())
for x in parameters[1:]:
x.requires_grad = False
主要的模型部分代码就是上面的部分,训练过程和详细代码参考Repo https://github.com/qhduan/mt5-soft-prompt-tuning
我们期望的结果是: