关于论文
论文名:Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer
机构:Google
发表于2020年,所谓迁移学习(transfer learning)就是把基于超大数据集预训练的模型在特定任务上用相对小得多的数据集再做一次fine-tuning训练,进而把预训练模型里的“知识”迁移到了给下游任务用的模型里去了,区别于zero-shot learning(即直接用预训练模型而不做fine-tuning)。作者提出了一个基于Transformer encoder-decoder结构的被命名为Text-to-Text Transfer Transformer(简称T5)的统一框架(unified framework)(也是一个模型)来把所有文本处理问题都转换为一个text-to-text形式(format)的问题。实际上论文里并没有太多介绍T5模型或T5框架,而更多的是从模型结构、训练目标、数据集、训练策略等方面用各种对比实验探索(explore)如何更好的做迁移学习,即论文名本身那句话。
模型训练
baseline模型
作者选择的baseline模型的参数量为220M,约为BERT-base的2倍。
预训练模型的训练目标和BERT一样也是denoising objective,也叫masked language modeling。训练target就是要重构出被破坏掉的token。
从数学上看是要利用标准的最大似然(standard maximum likelihood)来让模型预测target。
损失函数使用的是交叉熵(cross-entropy loss),优化器用的是AdaFactor。
Max seq len为512,batch size为128即每个batch包含128个sequence。
使用了packing技术,貌似由于TPU不能很好的支持动态seq length,因此为了尽量减少padding提高有效计算率,会尝试把多个短seq拼成(pack)一个总长不超过max seq len的single sequence,这样满打满算(无任何padding)一个batch里就包含了512*128=65536个token,即2^9 * 2^7 = 2^16。
模型预训练步数(steps)总共是2^19 = 524288步,warm-up步数是10000步,一步一个batch,总共用于预训练的token数达到了2^16 * 2^19 = 2^35 = 34,359,738,368 = 34B个token,这也仅仅是整个C4数据集的一部分(不到1个epoch,甚至2^23步也不到1个epoch),而且这个预训练数据量相比BERT(137B个token)和RoBERTa(2.2T个token)都少了很多,但是训练出来的模型的性能表现也足够好,也说明计算效率更高了。
学习率(lr)用了一种叫"inverse square root"的schedule方式,其中n是当前的step数,k是warm-up步数即10000,也就是在warm-up阶段学习率始终为0.01,随后开始按指数递减(指数为1/2)。
lr = 1/sqrt(max(n,k))
得到的预训练模型在所有任务上做fine-tuning的步数为2^18=262144步,和预训练一样max seq len为512,batch size为128,学习率固定为0.001,每5000步保存一个ckpt。
做模型推理(即生成文本)时使用的是greedy decoding方式,即每一步都选择最高概率的logit对应的token。但是作者也发现了可以使用beam search方式在翻译任务和新闻摘要任务中获得更好的性能表现。
还用baseline模型对比了不做预训练直接用下游任务的数据集跑有监督训练后得到的模型的性能,结果是对于除机器翻译外的任务还是预训练后再对下游任务做fine-tuning的效果更好。
模型结构
从参数量和encoder结构上重点参照了BERT模型。
这个论文里把Vaswani的Transformer结构非常清晰的描述了一遍。
autoregressive就是所谓causal self-attention。
数据集
从网上爬取的约20TB的Common Crawl数据集曾经被直接用于很多模型的训练,比如RoBERTa。作者基于Common Crawl数据集清洗整理了一个叫Colossal Clean Crawled Corpus的无标签纯英文数据集,简称C4,约750GB大小,作为T5模型的预训练数据集,并可在TensorFlow Datasets网站下载。T5模型预训练实际使用的数据量只占C4数据集的一小部分(不到1个epoch),而且相比其他模型的预训练,数据量仅仅是BERT的1/4,XLNet的1/16,RoBERTa的1/64。
下游任务与benchmark包括了机器翻译,QA,新闻摘要,文本分类。相关的数据集(GLUE, SuperGLUE, SQuAD, CNN/Daily Mail, WMT等)也都可以从TensorFlow Datasets网站下载。需要注意的是由于T5的预训练数据全是英文,因此做机器翻译任务前需要用机器翻译相关的数据集重新学习即fine-tuning一下。
其他无标签数据集
除了C4数据集,作者还实验了使用其他无标签数据集对T5模型做预训练,结果如下表。
其中,Wikipedia + Books的数据集组合既包含了百科全书的文字风格也包含了书籍类的文字风格,在很多模型的训练中把这两个数据集同时使用后确实在下游任务中得到了相比使用单一数据集更好的效果。
数据集大小与迭代次数
作者还尝试了缩小数据集大小并增加迭代次数(即多个epoch),结果发现迭代次数并不是越多越好,反而容易导致过拟合(loss过小)。
多任务学习
上述提到的GLUE和SuperGLUE其实都是由多个不同子任务组合起来构成的benchmark任务集,但是在T5模型的fine-tuning训练中会把GLUE当成一个单独的任务(single task)训出来一个模型,把SuperGLUE也当成一个单独的任务并额外加入了DPR(Definite Pronoun Resolution)子任务然后训出来一个模型。
多任务学习可以是multi-task training(不做预训练而直接做多个下游任务的有监督训练),也可以是multi-task pre-training(使用无监督训练数据集和有监督训练数据集混合起来做预训练,随后还可以做或不做fine-tuning)。
如果想把T5模型fine-tuning成一个能完成各种任务的单一的模型(single model),可以把T5模型的fine-tuning训练变成一个多任务学习(multi-task learning)过程,需要把所有子任务concat起来,并在各个子任务的数据集中给输入文本添加代表这个任务的文本前缀(task-specific text prefix),同时保证模型输出的也是一个文本结果,不管是什么任务都想办法变换成这种text-to-text形式,这就是所谓的text-to-text format/framework(实际上T5的预训练也是text-to-text的)。
计算平台
对T5模型做了模型并行和数据并行,使用的加速器是TPU POD,一个POD包含了1024个TPU芯片。
we use a combination of model and data parallelism and train models on “slices” of Cloud TPU Pods. TPU pods are are multi-rack ML supercomputers that contain 1,024 TPU v3 chips connected via a high-speed 2D mesh interconnect with supporting CPU host machines.
对比实验
不同模型结构
这里的模型结构指的是model structure或model achitecture。
严格说来,像GPT这种decoder-only的能够生成文本的才叫语言模型(LM),而BERT这种encoder-only的按本论文的定义貌似不算是语言模型而更算是一种分类模型。
作者对比了这么三种模型结构:
(1) encoder-decoder Transformer,完全fully-visible。
(2) decoder-only Transformer,也叫语言模型(LM),用了causal masking。
(3) prefix LM,改进了causal masking的缺点,对prefix部分不做mask了即让prefix部分fully-visible了。如果做分类任务,prefix LM生成的分类文本相当于BERT里表示分类的[C] token。
这三种模型结构对比如图所示。
不同训练目标
这里的训练目标指的是预训练模型的训练目标,并且是unsupervised objective。
Denoising vs. LM
通常decoder-only结构(比如GPT)的训练目标是Language Modeling,即LM objective,而encoder-only结构(比如BERT)和encoder-decoder结构(比如T5)的训练目标通常是Denoising,即Denoising objective。但是为了对比不同训练目标对最后下游任务的影响,作者做了一些交叉实验,如下表所示,最终结果证明了Denoising objective从整体上看还是要优于LM objective的。
Finally, we confirm the widely-held conception that using a denoising objective always results in better downstream task performance compared to a language modeling objective.
BERT-style
随后作者还进一步实验对比了其他训练目标(objective),模型都是用的T5(encoder-decoder)。
(1) 这个prefix LM就是之前提到的模型结构,其objective就是简单的把一段文本分成两部分,一部分作为input给encoder,另一部分作为target让decoder来预测。
(2) BERT-style objective参照了BERT预训练用的MLM(Masked Language Modeling) objective,即把一段文本中15%的token破坏掉(corrupt),被破坏的token中90%是被mask了,10%是用随机token替换了。由于BERT是encoder-only的结构,预训练时是要尽可能重构(reconstruct)出被破坏的token作为encoder的输出。而T5是encoder-decoder结构,这个实验参照了但没有完全照搬BERT的训练目标,实际在预训练中是把整个原始文本(original text)作为target,即要在decoder的输出位置重构出整个原始文本(注意T5 baseline模型中的target是要重构出被破坏的token,和这里的target也是不同的)。由于要重构出整个原始文本(uncorrupted sequence),计算量相对比较大,训练时间也比较长一些。
(3) 乱序重排,即deshuffling objective。
This approach takes a sequence of tokens, shuffles it, and then uses the original deshuffled sequence as a target.
前3种objective预训练得到的模型在下游任务中的性能表现对比如下表所示,发现还是第2种即BERT-style objective效果最好,实际上当时研究BERT模型的初衷之一就是想尝试找到一种好的objective能够训练出比基于LM objective性能表现更好的模型。
BERT-style的简化
(4) 一个匪夷所思的名字MASS-style,因为在另一篇论文里用的是这个名字,其实就是简单修改了一下上述第2个BERT-style objective,去掉了随机替换部分。verbose表述就是基于BERT的"把一段文本中15%的token破坏掉,被破坏的token中90%是被mask了,10%是用随机token替换了",把10%随机替换省略了,变成"把一段文本中15%的token破坏掉,被破坏的token全部被mask掉",然后仍旧"把整个原始文本(original text)作为target,即要在decoder的输出位置重构出整个原始文本"。
由于重构整个原始文本的代价太高,又想尝试能不能不重构整个原始文本(可降低计算量,缩短训练时间),为此又实验了两种方法:replace spans和drop tokens,即Table 3中第5行和第6行。
We were interested to see if it was possible to avoid predicting the entire uncorrupted text span since this requires self-attention over long sequences in the decoder.
(5) replace spans objective,其中i.i.d代表独立同分布(independent identically distributed)。
(6) drop tokens objective。
破坏率与破坏半径
即尝试修改corruption rate和corruption spans。
不同训练策略
要训练一个大模型,最直接的想法就是先用庞大的无标签数据集做自监督训练得到预训练模型,然后针对不同下游任务用对应的有监督数据集做fine-tuning训练得到最终可用的模型,即完成了一个迁移学习的过程。
同时有很多研究也在寻找新的不同的训练策略以达到减少fine-tuning训练时间、减少模型参数更新量等效果,比如在CV领域经常用到的backbone+classifier的模型结构组合,其中backbone可用当做是预训练模型,classifier一般就是一个参数量较小的线性层,做fine-tuning训练时固定backbone权重(freeze),从backbone出来的feature作为classifier的输入,整个训练过程只对classifier的权重进行更新。
在NLP领域也有类似的方法,即在预训练语言模型后面添加一个classifier,固定预训练模型的权重(freeze),从预训练模型出来的是原始输入文本的embedding或hidden states,作为classifier的输入,整个训练过程只对classifier的权重进行更新。不过这种训练策略对T5这种encoder-decoder结构的模型不太适用,因为需要对整个decoder部分做权重更新,而这部分参数量是很大的,并不能明显的缩短整个fine-tuning训练时间,因此需要尝试使用其他的训练策略。
LoRA的雏形
一种想法就是在预训练好的模型中插入一些adapter layers,这些adapter layers是一个或多个dense-ReLU-dense结构的block,插入位置一般在预训练模型中Transformer结构的FF(feed-forward)网络的后面,同时保证adapter layer的输入和输出shape一样,即插入adapter layer后不改变插入点原来的Tensor(中间计算结果)的shape。做fine-tuning训练时预训练模型中大部分权重不做更新(freeze),而只更新adapter layer和layer norm中的权重。实际上这个方法看起来就很像是LoRA方法(Low-Rank Adaptation)的雏形了。此外,这个方法有一个超参d,即FF网络的中间维度(inner dimensionality),这个d的取值会影响adapter layer的参数量,也会一定程度影响finetuned模型在下游任务中的性能表现,实测结果如下表所示。
逐步解冻
即gradual unfreezing,即一开始freeze大部分模型权重,随后逐渐unfreeze越来越多的权重让它们在fine-tuning训练中得到更新。此外,由于T5模型中input embedding矩阵(token to embedding)和output classification矩阵(embedding to token)的权重是共享的,需要让这部分权重在整个fine-tuning过程中一直更新。具体如何对模型权重逐渐解冻需要有一个schedule,作者这里直接用了一个简单粗暴的schedule: 由于encoder和decoder各有12个layers,把总训练步数2^18除以12得到12个小阶段(episode),每个小阶段在encoder和decoder同时unfreeze一个layer。
这种gradual unfreezing训练策略与上述adapter layers训练策略的效果对比如Table 10所示,发现其效果看起来要更好一些(相比全参数更新还是有一些minor degradation),而二者的fine-tuning训练步数貌似是相同的(2^18 steps)。
多任务学习
多任务学习可以是multi-task training(不做预训练而直接做多个下游任务的有监督训练),也可以是multi-task pre-training(使用无监督训练数据集和有监督训练数据集混合起来做预训练,随后还可以做或不做fine-tuning)。
多任务学习的一个重要的考虑因素是如何对多个数据集做mixing,即从各个数据集中分别选取多少比例的数据,作者探索(explore)了三种策略:
(1) Examples-proportional mixing
p = min(e,K)/sum(min(e,K))
# p代表probability of sampling an example from the mth task, 论文里写的是r_m
# K代表artificial data set size limit
(2) Temperature-scaled mixing
p = renorm(power(p, 1/T)) # T=1时就是Examples-proportional mixing
(3) Equal mixing
p = renorm(power(p, 1/T)) # 把T调大,比如T=100
注意在多任务学习中使用Equal mixing策略其实并不好,因为很容易发生过拟合与欠拟合。
This is most likely a suboptimal strategy, as the model will overfit quickly on low-resource tasks and underfit on high-resource tasks.
对于以上三种策略,自己举例简单计算了一下:
这里的多任务学习其实就是把预训练和fine-tuning训练合起来了,因此作者为了对比公平,把训练步数设置为baseline模型的预训练+fine-tuning的训练步数之和,即2^19 + 2^18 = 786432步。以上三种策略的实测结果如下表所示。
多任务学习已经把下游任务提前在预训练中进行了学习,实际上还可以把多任务学习再与fine-tuning结合起来。
Scaling
扩展4倍的算力,尝试了各种算力分配方案。
除了T5-base,还训练了其他版本的T5模型,包括small, large, 3B, 11B,汇总如下:
大模型 vs. 集成学习
用集成学习(ensemble)的方法把N个小模型的推理结果综合起来往往能得到比单个小模型更好的性能表现,表面上看起来小模型推理比大模型推理还更节省计算资源,但实际上是在每个小模型都做了一遍推理,算力的消耗实际上是N倍于单个小模型推理的,甚至比大模型推理更消耗算力。