3位牛津本科生学霸,4个月复现AlphaFold 3直接开源!

新智元报道

编辑:乔杨 好困

【新智元导读】5月发布的AlphaFold 3有论文、没代码,让许多机构和团队纷纷开启了「复现AF3」的工作。率先做出成果的,是一家成立不到一年的初创Ligo,3位创始人全都是牛津大学的本科生。

成就如此卓著、意义如此重大,AF3的成果公开自然让众多科学家翘首以盼。

然而,DeepMind团队却泼下了一盆巨大的冷水。他们只放出了论文,并没有公布任何相关的代码或模型权重。

如今刚刚过去不到4个月,初创公司Ligo宣布——他们已经完成了AlphaFold 3开源复现的工作。

项目地址:https://github.com/Ligo-Biosciences/AlphaFold3

这个成果相当激动人心,也得到了Figure创始人Brett Adcock的转发。

Ligo团队表示,他们正在使用AlphaFold 3的想法来进行酶设计,于是顺便开启了支线任务——复现AF3。

作为生物分子的结构预测模型,AF3主要可以用于三类任务:

预测蛋白质结构

预测药物-蛋白质相互作用结构

预测核酸-蛋白质复合物结构

这是结构建模技术的根本进步,整个生物科技行业理应从中受益。其应用范围广泛,包括:

CRISPR基因编辑技术:科学家可以准确看到DNA如何与「剪刀」Cas蛋白相互作用

癌症研究:预测潜在药物如何与癌症靶标结合,AF3论文的亮点之一就是预测KRAS抑制剂与其靶标的复合物

抗体/纳米抗体的靶向预测:AlphaFold3在这一类分子上的准确性比现有的最佳工具提高了两倍

而此次Ligo发布的模型是在单链蛋白质上训练的,可以完成上述三项功能中的第一项,即预测蛋白质结构,其他两个功能将在不久后完成训练并发布。

GitHub仓库中目前仅公开了代码,但团队表示,一旦训练和基准测试完成就会发布权重,而且会使用Apache 2.0许可证,实现「真正的开源」!

项目地址:https://github.com/Ligo-Biosciences/AlphaFold3

如何「复刻」AlphaFold 3?

由于DeepMind在论文中发布了模型的完整架构,以及每个组件的伪代码,因此任何团队都有权复现。

但复现这件事,说起来容易,做起来难。

Ligo选择将其完全翻译为PyTorch代码,其中涉及到不少逆向分析和重构工作,远多出他们的想象。

复现过程中,他们也发现了原始论文中存在的多个问题,会干扰训练,恰好是深度学习领域的关注点,因此一并放出,供社区参考。

1. MSE损失缩放公式有误

如下图所示,公式中使用的是加号而非乘号;如果使用加法,就无法在高噪声水平下正确降低权重,且MSE在初始化时不是单位化的。

这与Karras等人2022年发表的论文不同,可能只是一个笔误。

原文地址:https://arxiv.org/abs/2206.00364

2. 论文中省略了原DiT包含的残差层

代码中将其添加了回来,并对两种情况进行了对比实验,发现引入残差层可以改善梯度流和收敛性。

3. 当前形式的MSA模块中存在无效层,

如果使用论文所述的MSA模块的通信步骤,最后的配对加权平均和转换层无法对配对表示(pair representation)做出贡献,因此没有梯度。

代码使用了AlphaFold2中ExtraMsaStack的顺序;另一种解决方案是使用权重共享,但论文中是否有此操作尚不明确

除了复刻模型,Ligo团队也在探索更加快速高效的实现方式。

比如,重用了OpenFold的三角注意力(triangular attention),还将初代AF提出的MSARowAttentionWithPairBias重用于DiT,这是AF3论文中没有提及的操作。

下面这个动画由Ligo复现的模型生成(未使用模板),模型仅使用了8个A100 GPU训练 10 小时。

分析显示,尽管进行了优化,但模型超过60%的操作仍然受限于内存。因此除了扩展到更多功能,团队还在借鉴ScaleFold的理念,致力于实现一个更具扩展性、更高效的方案。

「闭源」惹众怒,3位本科生率先复现

作为一种生物分子结构预测模型,AlphaFold 3大大加速了蛋白质结构的绘制,解决了一个70多年来一直无法破解的巨大问题,因而对科学进步具有极其重要的意义。

研究人员可能需要用读完一整个博士的时间(4~6年)才能建模出一个结构,但AlphaFold 3只需几分钟,即可获得与实验精度相当的预测结果。

但除了博客和论文外,DeepMind只开放了一个服务器,允许科学家们进行非商业用途的使用,每天调用上限为20次。

网站地址:https://gdm-alphafold.corp.google.com/welcome

之所以如此一反常态,「捂紧」AF3的成果,很可能是有商业盈利用途的考量。

他们和新成立的子公司Isomorphic Labs正在进军药物研发领域,以满足大型制药公司的需求。

今年年初,Isomorphic Labs已经与诺华和礼来签订了价值30亿美元的合同,而最新的AF3的开发也有Isomorphic Labs的贡献,可被用于加速药物发现。

虽然为了商业盈利也无可厚非,但这种做法引起了科学家们的强烈不满。

文章地址:https://undark.org/2024/06/06/opinion-alphafold-3-open-source/

同时也是在批评Nature——为什么在没有公开代码时接收了论文。

在反对的声浪下,DeepMind没有顶住压力,随后改口表示,会在6个月内公布模型及权重。

Nature也不得不站出来回应,原原本本地向研究者们交代接收论文的考量和依据。

文章地址:https://www.nature.com/articles/d41586-024-01463-0

然而,很多反对者依旧不买账,毕竟在争分夺秒的研究领域,6个月实在太久;而且DeepMind出于合作方Isomophic Lab的盈利目的,很可能只会放出「阉割版」。

因此,除了Ligo,很多实验室和机构都开始了破解、复现AlphaFold 3的工作。

文章地址:https://www.nature.com/articles/s41586-024-07487-w

Ligo所借鉴的OpenFold团队就是其中之一,由哥伦比亚大学助理教授Mohammed AlQuraishi领头。

他们此前就对AlphaFold 2进行过复现和重新训练,成果在今年5月刚刚被Nature接收,代码也如数公布。

原文地址:https://www.nature.com/articles/s41592-024-02272-z

仓库地址:https://github.com/aqlaboratory/openfold

就在发推宣传这项成果时,Alquraishi表示,AlphaFold 3项目在进行时了。

此外,GitHub上的开源大佬、旧金山的独立软件工程师Phil Wang也组织起了一个众包开源项目,同样是用PyTorch复现AlphaFold 3,正进行得如火如荼。

项目地址:https://github.com/lucidrains/alphafold3-pytorch

本科生 VS DeepMind

有如此多的竞争者,能在不到4个月的时间率先复现AF3,Ligo这家初创究竟是什么来头?

这家YC系初创成立于去年12月,总部位于伦敦,3位创始人都来自牛津,且有丰富的研究经历。

CEO Edward Harris曾在普林斯顿大学入读计算机科学系,之后在2021年转入牛津医学院,目前正辍学全职创业。

进入牛津前,Harris曾在墨西哥创办外卖平台Abas2Go,目前年营业额超过120万美元。

CSO Emily Egerton-Warburton目前在牛津分子和细胞生物化学系学习,曾获得英国皇家化学学会颁布的化学奥赛金奖。

CTO Arda Goreci是牛津大学生物医学系的学生,2023年凭借计算生物学方面研究成果入选Google Cloude Research Innovator计划,他也是AF3开源项目的主要参与者和领导者。

参考资料:

https://x.com/adcock_brett/status/1832812049421807898

https://github.com/Ligo-Biosciences/AlphaFold3

https://news.ycombinator.com/item?id=41448439

https://x.com/ArdaGoreci/status/1830744265007480934