128张GPU炼出中国版AlphaFold2：训练代码全开源，打开浏览器就可体验，北大元培系AI公司打造

杨净 梦晨 发自 凹非寺

量子位 报道 | 公众号 QbitAI

中国版AlphaFold2，来了！

Uni-Fold，官宣即开源，还为科研者提供了测试接口和在线试用。

半年前，AlphaFold2和RoseTTAFold两个顶级AI算法齐齐引发生物学界、AI学界大地震。

半年后，这个横空出世的Uni-Fold，不光能完整复现AlphaFold2。

其预测精度超过RoseTTAFold、与AlphaFold2相当，效率还超AlphaFold2 2-3倍。

△Uni-Fold预测的蛋白结构

更重要的是，背后的团队不是如谷歌这样的互联网巨头，也不是数十年深耕于此的顶尖高校团队。

而是一家成立刚3年的AI创业公司——深势科技。

这个Uni-Fold究竟有何看头？接下来一看究竟。

128张GPU炼出中国版AlphaFold2

AlphaFold2开源的消息曾一度火爆全场，等稍微冷静下来大家才发现，真正把它用起来并不容易。

用谷歌官方提供的Colab资源吗？可以，请排队。

自己有算力资源想部署？也可以，不过AlphaFold2只开源了模型推理部分的代码，并没有开源训练代码。

这意味着部署好后只能按AlphaFold2设计好的流程使用，难以针对特定的研究项目做一些迁移调整。

比如拿去预测蛋白质复合物结构，研究蛋白质与小分子相互作用，又或者与电镜实验相结合，这些都无能为力。

再加上，该代码基于谷歌自研的可微分计算框架JAX及TPU硬件平台开发，高度依赖谷歌生态系统。

最后，AlphaFold2虽然模型开源，但调好的参数组合却不开放商业使用。

谷歌母公司Alphabet对AlphaFold2有着自己的商业化打算，11月新成立的子公司Isomorphic Laboratories便是其第一步。

鉴于这样的局限性，像深势科技这样的创业团队得自己想办法。

他们的办法，便是根据开源模型自己复现AlphaFold2的训练部分。

经过几个月的攻关，他们终于在128块英伟达V100上复现了AlphaFold2的全规模训练。

其中遇到最大的困难是GPU资源有限，难以同时进行多次训练来比较效果。于是他们只好在设计上下功夫，尽量减少试错次数。

没想到因此有了一些额外的收获。

经过混合精度、MPI 并行以及数据异步读取等一系列优化后，复现模型做到了在相同硬件下推理效率还比AlphaFold2公开版高上2-3倍。

另外复现的模型也不再依赖谷歌生态，增加了与英伟达GPU、国产DPU等硬件体系的适配性。

模型搞定了，那接下来最关键的问题就是，预测精度如何？

团队用CASP14蛋白预测任务中的绝大部分序列进行了测试。

其中包括基于模版建模的两种难度（TBM-easy/hard），以及最具挑战性的无模版建模（FM），发布时Uni-Fold平均Cα-lDDT达到82.6。目前模型还在持续改进中，精度有望再上一层楼。

这个成绩超过了RoseTTAFold，与公开版AlphaFold2模型接近。

对于这个精度媲美原版、推理效率还更高的蛋白质结构预测工具，他们决定将其推理代码、训练代码全部开源，命名为Uni-Fold。

深势科技公布训练代码的举动受到业界广泛好评，连深度学习预测蛋白质结构的先驱、芝加哥丰田计算技术研究所的许锦波教授都评价道：

Uni-Fold的出现，让每个研究者都可以训练自己的模型。这无疑将进一步推动技术的发展。对整个蛋白质折叠领域来说，都是令人激动的消息。

深势科技这家成立仅3年的公司，为什么能成为国内首家复现Alphafold2全规模训练的团队？

复现后又为何要率先公开训练代码？

下面就来一探究竟。

北大元培系创业公司，一年完成三轮融资

实际上，在推出Uni-Fold之前，深势科技这个初创团队就已经颇受业内关注：

一年狂揽三轮融资，最近一轮达数千万美元，由高瓴创投领投。

从团队背景看，创始人兼首席科学家张林峰，北京大学元培学院出身，在普林斯顿大学获应用数学系博士学位。

另一位创始人兼CEO孙伟杰同样来自北大元培，后继续在北京大学深造，获管理学硕士学位。

孙伟杰在科技领域分析及投资经验丰富，张林峰则偏重技术。他的研究方向包括统计物理、分子模拟和机器学习及其应用。

此外担任首席科学顾问的是中科院院士、北京大学教授鄂维南。

北大元培系AI创业公司这个标签，算是深势科技受到资本青睐的原因之一。

再看技术实力，由鄂维南院士领衔的数十人科研队伍，包含物理建模、数值算法、机器学习、高性能计算及药物和材料计算等多个领域人才。

他们大多来自世界一流高校、科研机构和企业。

其中许多人拥有交叉背景，深势科技特别介绍了一位高中搞过生物竞赛，大学主修人工智能又参加过超算研究的同学。

这样的人了解不同学科的基础知识和术语体系，对团队协作沟通大有好处。

有着多学科交叉这样的特点，深势科技自述致力于“以新一代分子模拟方法打造微尺度工业设计平台”。

其中的新一代分子模拟方法，指「多尺度建模+机器学习+高性能计算」新研究范式。

新范式由深势科技首创，在保证了计算效率和精度的基础上还具备通用性，可为药物、材料等领域带来新的计算模拟及设计工具。

比如鄂维南、张林峰共同参与研究的一项技术，用机器学习方法将分子动力学极限提升至10亿原子规模同时保持高精度。

这项成果在2020年获得有「超算届诺贝尔奖」之称的戈登·贝尔奖，还与「量子优越性」和「人造太阳」等成果一起当选2020中国十大科技进展。

基于这样的团队背景和科研实力，也就不难理解深势科技为何能在这么短的时间内复现AlphaFold2全规模训练。

但作为一家创业公司，光有上述两种因素还不够，创造行业价值才是当前大环境下的题中之义。

有算法工具的突破做基础，Unif-Fold配套解决方案也已集成到他们自己的药物设计平台Hermite，免费注册即可测试使用。

Hermite基于云计算，无需操心模型部署、环境配置问题，只要打开浏览器就能使用Uni-Fold的各种功能。

除此之外，作为一站式药物设计平台的Hermite还集成了更多功能。

如预测得到蛋白质结构后还可以继续进行蛋白质的动力学模拟、结构精修以及环区优化。

在药物优化阶段也提供了比较不同分子与蛋白靶点结合活性差别的工具。

这也是深势科技打破蛋白质结构预测工具现状，率先公布训练代码背后的考量——

他们做的不只是蛋白质结构预测这一个工具，而是为药物设计领域打造一套模拟-AI-实验三者驱动的解决方案和开发平台。

而开源开放，吸引更多人来到开源社区，所积累的数据、汇聚的多学科人才、实现的应用都是平台生态建设的重要力量。

深势科技的微尺度工业设计平台不仅涉及生物医药领域，在新能源材料研发中也有布局。

最近，深势科技与宁德时代达成战略合作，双方将共建联合实验室，以计算驱动能源材料设计研发。

计算生物热潮

这时候推出Uni-Fold，倒也并不是心血来潮，而是聚集了天时地利人和，选择谋定而后动。

这当中既有现实需求的驱动，也有行业大势的推动。

对深势自身而言，过去一直致力于以AI+物理模型驱动方式为行业创造价值。

简单来说，就是以第一性原理的方式，将实际面临的问题或场景演化为定义明确的物理模型。

像蛋白质的结构以及构象变化、配体与蛋白质的相互作用 (包括亲和力与动力学参数的评估）、以及药物分子晶型和剂型的溶解度等这些药物研发过程中的关键问题，都可以用物理化学模型进行定义。

这些问题对应到药物设计平台Hermite上，就是诸如药物优化模块Uni-FEP、蛋白动力学采样与结构优化模块RiD这类的功能环节。

在推出Uni-Fold之前，Hermite就已经开启商业化进程，被多家药企研发人员使用。

而Uni-Fold——作为药物设计的第一环的身份出现，直接为Hermite填补了数据驱动引擎。

它有助于深入了解生物机制，随后将更多“黑盒子”问题转变为物理问题。

中科院院士、北京大学教授、深势科技首席科学顾问鄂维南评价道：

物理模型和数据驱动结合的解决方案将为蛋白结构和药物设计行业提供更好的基础。

事实上放眼整个行业，计算生物正在引领一场科技革命。

其核心代表就是以AlphaFold2为代表的新型蛋白质结构预测技术。

利用原有的实验手段（比如，冷冻电镜、X光晶体衍射等）可能需要经年累月的研究，才能解析一个蛋白质结构，而利用计算结构模型最快只需要10分钟。

而若再向前推进，到应用端——生物医药领域，新药研发的效率有望大幅提高；科学家能借助AI系统设计出自然界不存在的蛋白质，催生各种新材料，用于能源、化工、环保等行业。

因此，选择在这时候完整复现AlphaFold2，并将其开源、开放给更多的研发人员，无疑是最快实现物尽其用、创造价值的方式之一。

在人们固有的印象中，以往生命科学的研究思路都是由表及里、从现象到本质。

当中也有不少科学家身先士卒，从本质出发探索微观分子世界，但始终受限于实验工具、设备等外在条件。

直到人工智能的出现，更多自下而上的问题逐步得到解决。

算力、算法的发展能够真正缩短理论研究与解决实际问题之间的距离。

鄂维南院士曾在多个场合呼吁：“科学是人工智能下一个主战场。”

AI+Science的范式，价值已经显现。

而深势科技，正是站在这一科技潮头上的先行者。