谷歌|128张GPU炼出中国版AlphaFold2：训练代码开源，浏览器就可体验

杨净梦晨发自凹非寺量子位报道 | 公众号 QbitAI中国版AlphaFold2，来了！
Uni-Fold，官宣即开源，还为科研者提供了测试接口和在线试用。
半年前，AlphaFold2和RoseTTAFold两个顶级AI算法齐齐引发生物学界、AI学界大地震。
半年后，这个横空出世的Uni-Fold，不光能完整复现AlphaFold2。
其预测精度超过RoseTTAFold、与AlphaFold2相当，效率还超AlphaFold2 2-3倍。

文章插图
△Uni-Fold预测的蛋白结构
更重要的是，背后的团队不是如谷歌这样的互联网巨头，也不是数十年深耕于此的顶尖高校团队。
而是一家成立刚3年的AI创业公司——深势科技。
这个Uni-Fold究竟有何看头？接下来一看究竟。
128张GPU炼出中国版AlphaFold2AlphaFold2开源的消息曾一度火爆全场，等稍微冷静下来大家才发现，真正把它用起来并不容易。
用谷歌官方提供的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全规模训练的团队？