整理：马瑞军

校对：裘捷中

智源导读：3月20日，在「智源大规模预训练模型交流论坛」，由智源研究院副院长、清华大学教授唐杰带头，发布了「悟道1.0」四大模块的进展工作，分别包括

• 悟道·文源——以中文为核心的超大规模预训练语言模型；

• 悟道·文澜——超大规模多模态预训练模型；

• 悟道·文溯——超大规模蛋白质序列预测预训练模型；

• 悟道·文汇——面向认知的超大规模新型预训练模型。

其中，「悟道·文溯」是「超大规模蛋白质序列预测预训练模型」，最终目标是研发出超大参数规模、可以处理超长蛋白质序列的超大规模预训练模型，在基本性能、可解释性和鲁棒性等多个方面达到世界领先水平。目前，「悟道·文溯」已在蛋白质方面完成基于Pfam数据库训练的BERT模型，同时搭建训练软件框架并验证其可扩展性。团队成员来自清华计算机系KEG实验室，清华药学院鲁白实验室和腾讯量子实验室。

清华大学计算机系五年级博士生、青源研究组成员裘捷中在论坛详细介绍了「悟道·文溯」项目中的突破性进展。以下为「悟道·文溯」报告总结：

01

蛋白质预训练模型

预训练，即在大量无标注数据集上训练，在少量有标注数据集上微调。其始于2018年google开发的BERT，发展到如今已经在NLP以及其他诸多领域取得了大量进展。

图1：预训练语言模型家族谱

图 2所示为语言模型的进步图。可以看到在预训练模型出来前，模型的表现相较于人类差距较大，但是预训练模型出现后，模型的表现逐步超越人类，表现出令人惊叹的结果。

图2：预训练模型表现发展

在预训练模型如火如荼的发展时代背景下，一部分学者将其引入蛋白质检测序列任务中。

图3：蛋白质序列特征

如图3所示，蛋白质包含一级结构，二级结构和三级结构。从一级结构上看，蛋白质和语言具有相似的序列特性，这为将语言预训练模型引入蛋白质序列检测提供了有利条件。于是，科学家们直接将Transformer模型，转移到蛋白质序列检测的预训练模型中，发现具有很高的可行性。同时，科学家还发现，在蛋白质二级结构和三级结构预测的下游任务中，例如Contact Prediction等任务中和蛋白质折叠中，Transformer模型依旧有效，这为蛋白质序列检测模型的发展打开了新的大门。

此外，文溯团队的成员还联合清华大学计算机系 KEG 和 PACMAN 实验室一起开发了超大规模分布式预训练框架FastMoE

（github.com/laekov/fastmoe），未来计划将FastMoE和蛋白质预训练模型结合起来。

02

文溯模型发展

在蛋白质序列检测领域中有三大模型，分别是：

1） 由Berkley主导完成的TAPE；

2） 由FAIR主导完成的ESM-1b；

3） 由慕尼黑工大和Google主导的ProTrans。

三者均采用了预训练模型，不同之处在于模型使用的训练数据集、模型的规模以及蛋白质序列个数有所不同，具体信息如图4所示。

图4：蛋白质序列预测模型信息

悟道-文溯在蛋白质预测的规模选择上，选取了与Berkley相同的38M模型规模。在性能方面，将要面世的文溯V1模型在TAPE的下游任务预测实验中表现良好，复现了TAPE的大部分功能，其中在Secondary Structure、Contact Prediction、Remote Homology Detection等子任务中超越了TAPE，实现了新的性能高度，具体信息如图5所示。

图5：文溯对比TAPE

预训练模型的规模，近年来正朝着大规模化方向快速发展。蛋白质预训练也正是如此。

目前蛋白质预训练模型的规模与语言预训练模型在2019年2月规模相当，但是在未来，蛋白质预测也将朝大规模方向发展。在此前景条件下，悟道-文溯团队的成员还联合清华大学计算机系 KEG 和 PACMAN 实验室一起开发了超大规模分布式预训练框架FastMoE（github.com/laekov/fastmoe），未来计划将FastMoE和蛋白质预训练模型结合起来，开发出可以与现有万亿规模语言预训练模型相当的蛋白质预训练模型。

图6：预训练模型规模发展

03

大规模混合专家预训练框架FastMoE

MoE（Mixture of Experts）是在模型计算量不显著增加的前提下训练大规模预训练模型的核心技术，例如Google发布的具有1.6万亿参数的预训练模型Switch Transformer，其核心技术就是MoE结构。

MoE结构是一种由n个专家（expert）组成的专家池技术，每个专家池可以是任意的神经网络，结构示意图如图7所示。

图7：MoE结构

对于每个输入X，由门网络gating network对专家的适用性进行打分，再根据打分结果选择具体选用哪个/些专家来计算，例如图8所示的的例子中，门网络最终选用2号和n-1号专家作为计算网络。

Switch Transformer使用MoE改造了Transformer的FFN模块，在计算量基本不变的情况下，大幅度增加了参数量，但是Google的分布式训练框架需要依赖于Mesh-Tensorflow分布式训练框架和定制的TPU硬件，相当程度上限制了学术界和业界的使用和研究。

为了降低训练万亿规模模型的经济成本，打破Google在大规模预训练平台的垄断，弥补现有平台对于框架和硬件选择具有严重缺陷等问题，开源分布式MoE模型训练系统FastMoE应运而生。为广大GPU用户和PyTorch用户提供了研究平台，使得部分学者可以研究MoE结构和基于MoE结构开发的大规模预训练模型。

FastMoE系统具有以下4个优异的特性：

（1）易用：只需一行代码即可完成MoE改造。FastMoE系统可以作为PyTorch网络中的一个模块，快速插入模型中。代码可以如下所示：

model = ...
from fmoe.megatron import fmoefy
model = fmoefy(model, num_experts=<number of experts per worker>)
train(model, ...)

（2）灵活：支持自定义专家网络、自定义门网络。FastMoE系统的专家网络配置十分灵活，研究者可以根据自己的需求和实际要求对专家网络搭建和改造。

（3）高效：针对常见专家网络进行专门性能优化。FastMoE相对于用朴素PyTorch实现的MoE结构，可以实现47倍的大幅提升。

图8：FastMoE实战表现

（4）分布式支持：提供多种并行模式可供选择。FastMoE可以为用户提供数据并行模式和模型并行两种方式，采用两种策略实现模型的大规模化。

图9：FastMoE 的并行模式

数据并行模式可以将包含少量专家的专家池（Worker）复制多份并行计算，这种特性可以与PyTorch的数据并行模式和分布式数据并行模式无缝契合，方便PyTorch用户开发，但是缺陷在于专家池容量不能太大。

如图9左图所示为数据并行模式情况下的两个MoE结构，该结构中包含两个Worker，每个Worker包含三个专家和一个门网络，且不同Worker上的专家和门网络共享参数。

FastMoE采用模型并行模式时，gate门网络仍然在每个Worker上复制，但是专家在专家池之间被分开放置，通过引入通信机制，FastMoE拥有一个与Worker数量成比例的大型专家库，当算力的数量增加时，专家的数量会线性增加，以此可以实现模型参数的快速叠加达到大规模化。

如图9右图所示为模型并行的6专家MoE的正向传递结构，其中专家1-3位于Worker 1中，专家4-6位于Worker 2中，同时Worker 1和2均包含共享参数的gate门网络。

目前，FastMoE已经深度集成到了业界常用的大规模预训练模型Megatron-LM框架中。相对于原始的GPT版本，在使用64个专家的FastMoE模型并行模式下，使用FastMoE修改后的GPT在相同浮点运算量的情况下，收敛速度更快，具体如图10所示。

图10：FastMoE分布式训练收敛图

此外，由于Megatron-LM中没有提供FastMoE模型并行模式下需要使用的平行策略，FastMoE提供了相应的模块用以代替PyTorch中的DDP模块，具体可以查看fmoe.DistributedGroupedDataParallel文件。

04

未来展望

在未来，悟道-文溯团队将从以下5方面出发，促进文溯模型的快速发展：

（1）发布悟道-文溯V1模型，为广大用户提供蛋白质预训练平台；

（2）在设计训练任务时，更加偏向于有生物学意义的预训练任务；

（3）将模型朝更大规模方向发展，将模型的性能持续提高；

（4）开发基于FastMoE的超大蛋白质预训练模型；

（5）欢迎广大学者，厂商使用文溯模型以及共同参与相关的开发工作。

更多内容请关注后续项目发布

推荐阅读：

悟道 · 文汇详解：少样本学习等近十个数据集取得第一

悟道 · 文澜详解：目前最大的中文多模态通用预训练模型

【关于智源社区】

智源社区隶属于北京智源人工智能研究院，我们致力于创建一个AI领域内行人的交流平台。

在这里你有机会参与全年线上线下百场专题论坛，与顶尖学者零距离接触；也可以与同行探讨领域前沿，碰撞思想火花。如果你更想进入微信群与更多同行人发起实时讨论，或者加入「青源会」结识更多研究伙伴，或者成为智源社区编辑参与更多文字工作，欢迎填写以下表单（扫描二维码）进行申请。