AlphaFold3 开源实现

简介

这是 Ligo 对 AlphaFold3 的开源实现，是一项持续进行的研究项目，旨在推动开源生物分子结构预测的发展。本次发布实现了完整的 AlphaFold3 模型及其训练代码。我们首先开放单链预测功能，待配体、多聚体和核酸预测功能训练完成后，将陆续加入。在此报名参与测试。

本仓库旨在加速构建忠实且完全开源的 AlphaFold3 实现，供整个生物科技社区免费使用。

演示视频

我们发现该模型的训练动态非常迅速。以下视频展示的是在 8 张 A100 GPU 上无模板条件下训练了 4,000 步、历时 10 小时的模型输出示例。

动画制作：Matthew Clark

致谢

本项目离不开以下项目及个人的贡献：

• Google DeepMind 的 AlphaFold3 团队，感谢他们开创性的研究工作以及核心算法的公开。

• OpenFold 项目（https://github.com/aqlaboratory/openfold），该项目为开源蛋白质结构预测奠定了基础。我们复用了其许多核心模块，如三角注意力机制和乘法更新机制，以及数据处理流程。

• ProteinFlow 库（https://github.com/adaptyvbio/ProteinFlow），尤其是 ProteinFlow 的架构师 Liza Kozlova（@elkoz），她在整个过程中发挥了至关重要的作用。我们的大部分原型模型都是基于 ProteinFlow 进行训练的，因为它提供了一个简洁且文档完善的蛋白质数据处理管道。我们已与 AdaptyvBio 合作，基于 ProteinFlow 构建了支持完整配体和核酸的 AlphaFold3 数据管道。目前，@elkoz 和 @igor-krawczuk 正在开发 ProteinFlow 的下一个版本，以全面支持这些数据模态。

• @alexzhang13，感谢他在 Triton 中实现的自定义 MSA 对权重平均核，对于较长序列而言，其内存效率比 PyTorch 实现高出 10 倍以上。这一改进有效缓解了大规模序列训练中的关键性能瓶颈。衷心感谢 Alex 对本项目的贡献！

项目状态

本项目仍处于早期阶段，属于一项活跃的研究工作。我们正致力于为社区准备一个稳定的版本。尽管我们对这项工作的潜力充满期待，但仍需强调，目前它尚未达到生产级工具的标准。我们已针对单链蛋白质训练了一个 AlphaFold3 版本以测试实现效果；下一版将包含完整的配体和核酸支持。我们正在招募少量测试用户，协助我们测试实现并提供反馈。如果您有意参与测试，请 加入我们的候补名单。

与 AlphaFold3 伪代码的差异

在项目推进过程中，我们发现 AlphaFold3 补充资料中描述的一些算法特性与现有的深度学习文献并不一致。具体如下：

• MSA 模块顺序：补充资料中，MSA 模块的通信步骤位于 MSA 堆栈之前。这导致最后一个区块的 MSA 堆栈无法参与结构预测，因为所有信息都通过成对表示流出。按照伪代码中的顺序，最后一层的 MSA 堆栈没有机会更新成对表示。为此，我们将 OuterProductMean 操作与 MSA 堆栈的位置互换，以确保所有区块都能对结构预测做出贡献。值得注意的是，这一调整与 AlphaFold2 的 ExtraMSAStack 中的操作顺序一致。DeepMind 提到这些 MSA 模块是“同质的”，但并未明确是指共享权重还是具有相同的架构。若各层确实共享权重，则梯度会流经所有模块，然而最终的 MSA 堆栈计算却未被利用——这部分可以安全跳过（补充资料中未提及）。我们将在收到 DeepMind 的回复后进一步澄清这一问题。

• 损失缩放：补充资料中描述的损失缩放因子在初始化时无法产生单位损失。而单位损失作为初始条件，是 Karras 等人（2022）在训练扩散模型时提出的理想损失函数特性之一；Max Jaderberg 在一次演讲中也提到，他们之所以选择 Karras 等人的框架，正是基于这一特性。我们认为，补充资料中的这一问题可能是一个简单的笔误，即将加法误写为乘法。在我们的实现中，我们采用了与 Karras 等人（2022）一致的损失缩放因子。实测表明，这样可以在初始化时获得单位均方误差损失，而补充资料中的缩放因子在初始化时则大了两到三个数量级。此外，论文中的损失缩放因子在 t=16.0 处存在局部最小值，随后会随噪声水平的增加而增大。这与 Karras 等人（2022）提出的损失函数特性不符——后者强调应在高噪声水平下降低损失权重。我们在仓库中添加了一个 Jupyter 笔记本，展示了我们的实验结果。

• DiT 块设计：AttentionPairBias 和 DiffusionTransformer 块的设计似乎紧密遵循 Peebles & Xie（2022）提出的 DiT 块设计。然而，这些模块缺少残差连接。论文并未解释 DeepMind 为何选择省略残差连接。我们分别进行了实验，发现在我们训练的步数范围内，带有残差连接的 DiT 块能够显著加快收敛速度，并改善网络中的梯度流动。需要注意的是，这是我们最不确定的一处差异，如果原始实现确实未使用残差连接，只需在我们的代码中简单修改几行即可。

以上内容旨在提高透明度，并邀请社区共同探讨解决这些问题的最佳方案。

模型效率

本次实现的一个重要侧重点是优化模型各组件的速度和内存效率。AlphaFold3 包含许多类似 Transformer 的组件，但由于 AlphaFold3 中的成对偏置机制，诸如 FlashAttention2 等高效的硬件感知注意力实现无法直接与这些模块无缝集成。所有注意力操作都会从成对表示中投影出一个成对偏置，该偏置会在键—查询点积之后被添加，并且需要将梯度反向传播回去。这并非 FlashAttention2 无法处理的范围，因为该偏置的梯度与缩放后的 QK^T 点积梯度相同，然而当前的实现并不支持这一点。更近期的注意力实现，如 FlexAttention，前景十分可观，但它们目前同样不支持偏置梯度，因为在前向传播过程中对偏置张量执行的广播操作，在反向传播时会变为约简操作，而这一功能尚未在 FlexAttention 的首个版本中实现。

• 我们尽可能复用 OpenFold 项目中经过实战检验的组件，例如 TriangularAttention 和 TriangularMultiplicativeUpdate。OpenFold 项目的模块化设计使我们能够轻松地将这些模块导入到我们的代码库中。我们正利用 Triton 对这些模块进行效率改进，通过融合操作来提升性能并减少中间张量的分配。

• 我们观察到，在 PyTorch 中对扩散模块进行朴素实现时，经常会因内存不足而报错，这是因为每个批次中扩散模块会被复制 48 次。为解决这一问题，我们重新利用 Deepspeed4Science 中的 MSARowAttentionWithPairBias 内核，实现了一个内存高效的扩散模块版本，将不同噪声水平下的批次副本视为额外的批次维度。对于 AtomAttentionEncoder 和 AtomAttentionDecoder 模块，我们尝试了自定义的 PyTorch 原生实现，以将内存占用从二次方降低到线性，但与简单地重用 AttentionPairBias 内核相比，收益并不显著。我们在仓库中同时提供了这两种实现，不过为了保持代码简洁，我们仍采用朴素的实现方式。

尽管进行了上述优化，我们的性能分析实验表明，模型超过 60% 的操作都受内存限制。我们正在基于 ScaleFold 的思路，开发一种更为高效且可扩展的实现方案，这将使我们能够达到原始 AlphaFold3 的训练规模。

MSA 成对平均效率

@alexzhang13 使用 Triton 实现了一个快速且内存高效的自定义 MSA 成对加权平均内核。

• 我们曾注意到，MSA 成对加权平均操作是关键的内存瓶颈之一。AlphaFold3 论文中指出，该操作用一种“更廉价”的成对加权平均替代了 MSARowAttentionWithPairBias 操作，但在 PyTorch 中以朴素方式实现时，其内存使用量却比 Deepspeed4Science 的 MSARowAttentionWithPairBias 内核高出四倍。我们推测，这是由于 FlashAttention 中的分块和重计算技巧所带来的内存效率提升，而这些技巧也被整合到了 Deepspeed4Science 的 MSARowAttentionWithPairBias 内核中。

• 朴素的成对加权平均实现会分配一个 (*, N_seq, N_res, N_res, heads, c_hidden) 形状的中间张量，即使对于中等长度的序列，这个张量也大到无法容纳在 GPU 内存中。

• Alex 的内核使得网络能够在单个 GPU 上扩展到数千个 token！

PyTorch 与 Triton 内核的内存使用对比

PyTorch 与 Triton 内核的运行时间对比

快速入门

我们目前尚未提供采样代码，因为配体—蛋白质及核酸预测功能仍有待训练。研究人员可以使用 PyTorch Lightning 的检查点加载功能来加载预训练权重，进行实验和模型修改。当前模型仅能预测单链蛋白质，其功能与原始 AlphaFold2 相同。模型组件的设计具有高度的可重用性和模块化，便于研究者将其轻松集成到自己的项目中。 如需参与配体—蛋白质及核酸预测功能的测试，请加入我们的候补名单：加入候补名单

使用说明

目前，本仓库的主要用途是科研与开发。待配体—蛋白质及核酸预测功能准备就绪后，我们将逐步增加面向用户的实用功能。

贡献指南

我们欢迎社区的贡献！我们的代码中很可能存在许多 bug 和细微的实现错误。深度学习训练往往会出现静默失败的情况，即虽然网络仍能收敛，但性能会因此略有下降。如果您有意贡献代码，可以通过 GitHub 提交包含问题描述的 issue，或者直接 fork 本仓库，创建新分支并提交带有清晰变更说明的 pull request。

如有其他意见或建议，请通过电子邮件 alphafold3@ligo.bio 与我们联系。

引用

如果您在研究中使用此代码，请引用以下论文：

@article{Abramson2024-fj,
  title    = {利用{AlphaFold} 3 准确预测生物分子相互作用的结构},
  author   = {Abramson, Josh 和 Adler, Jonas 和 Dunger, Jack 和 Evans, Richard 和 Green, Tim 和 Pritzel, Alexander 和 Ronneberger, Olaf 和 Willmore, Lindsay 和 Ballard, Andrew J 和 Bambrick, Joshua 和 Bodenstein, Sebastian W 和 Evans, David A 和 Hung, Chia-Chun 和 O'Neill, Michael 和 Reiman, David 和 Tunyasuvunakool, Kathryn 和 Wu, Zachary 和 {\v Z}emgulyt{\.e}, Akvil{\.e} 和 Arvaniti, Eirini 和 Beattie, Charles 和 Bertolli, Ottavia 和 Bridgland, Alex 和 Cherepanov, Alexey 和 Congreve, Miles 和 Cowen-Rivers, Alexander I 和 Cowie, Andrew 和 Figurnov, Michael 和 Fuchs, Fabian B 和 Gladman, Hannah 和 Jain, Rishub 和 Khan, Yousuf A 和 Low, Caroline M R 和 Perlin, Kuba 和 Potapenko, Anna 和 Savy, Pascal 和 Singh, Sukhdeep 和 Stecula, Adrian 和 Thillaisundaram, Ashok 和 Tong, Catherine 和 Yakneen, Sergei 和 Zhong, Ellen D 和 Zielinski, Michal 和 {\v Z}{\'\i}dek, Augustin 和 Bapst, Victor 和 Kohli, Pushmeet 和 Jaderberg, Max 和 Hassabis, Demis 和 Jumper, John M},
  journal  = {Nature},
  month    = {5月},
  year     =  2024
}

@article {Ahdritz2022.11.20.517210,
	author = {Ahdritz, Gustaf 和 Bouatta, Nazim 和 Floristean, Christina 和 Kadyan, Sachin 和 Xia, Qinghui 和 Gerecke, William 和 O{\textquoteright}Donnell, Timothy J 和 Berenberg, Daniel 和 Fisk, Ian 和 Zanichelli, Niccol{\`o} 和 Zhang, Bo 和 Nowaczynski, Arkadiusz 和 Wang, Bei 和 Stepniewska-Dziubinska, Marta M 和 Zhang, Shang 和 Ojewole, Adegoke 和 Guney, Murat Efe 和 Biderman, Stella 和 Watkins, Andrew M 和 Ra, Stephen 和 Lorenzo, Pablo Ribalta 和 Nivon, Lucas 和 Weitzner, Brian 和 Ban, Yih-En Andrew 和 Sorger, Peter K 和 Mostaque, Emad 和 Zhang, Zhao 和 Bonneau, Richard 和 AlQuraishi, Mohammed},
	title = {{O}pen{F}old: 重新训练 {A}lpha{F}old2 能够带来对其学习机制和泛化能力的新见解},
	elocation-id = {2022.11.20.517210},
	year = {2022},
	doi = {10.1101/2022.11.20.517210},
	publisher = {Cold Spring Harbor Laboratory},
	URL = {https://www.biorxiv.org/content/10.1101/2022.11.20.517210},
	eprint = {https://www.biorxiv.org/content/early/2022/11/22/2022.11.20.517210.full.pdf},
	journal = {bioRxiv}
}

@article{kozlova_2023_proteinflow,
  author = {Kozlova, Elizaveta 和 Valentin, Arthur 和 Khadhraoui, Aous 和 Gutierrez, Daniel Nakhaee-Zadeh},
  month = {09},
  title = {ProteinFlow：用于深度学习应用的蛋白质结构数据预处理 Python 库},
  doi = {https://doi.org/10.1101/2023.09.25.559346},
  year = {2023},
  journal = {bioRxiv}
}

@misc{ahdritz2023openproteinset,
      title={{O}pen{P}rotein{S}et：大规模结构生物学的训练数据}, 
      author={Gustaf Ahdritz 和 Nazim Bouatta 和 Sachin Kadyan 和 Lukas Jarosch 和 Daniel Berenberg 和 Ian Fisk 和 Andrew M. Watkins 和 Stephen Ra 和 Richard Bonneau 和 Mohammed AlQuraishi},
      year={2023},
      eprint={2308.05326},
      archivePrefix={arXiv},
      primaryClass={q-bio.BM}
}

@article{Peebles2022DiT,
  title={基于 Transformer 的可扩展扩散模型},
  author={William Peebles 和 Saining Xie},
  year={2022},
  journal={arXiv 预印本 arXiv:2212.09748},
}

@inproceedings{Karras2022edm,
  author    = {Tero Karras 和 Miika Aittala 和 Timo Aila 和 Samuli Laine},
  title     = {阐明基于扩散的生成模型的设计空间},
  booktitle = {NeurIPS 会议论文集},
  year      = {2022}
}

许可证

本项目采用 Apache License 2.0 许可证授权——详情请参阅 LICENSE 文件。

AlphaFold3 (Ligo 开源版) 快速上手指南

注意：本项目目前处于早期研究阶段，仅支持单链蛋白质结构预测（功能等同于 AlphaFold2）。配体、多聚体及核酸预测功能正在训练中，尚未开放。代码主要用于研究和实验，非生产环境就绪工具。

1. 环境准备

系统要求

• 操作系统: Linux (推荐 Ubuntu 20.04+)
• GPU: NVIDIA GPU (推荐 A100 或同等算力卡)，显存建议 24GB+ 以处理较长序列
• CUDA: 兼容 PyTorch 版本的 CUDA 驱动

前置依赖

• Python 3.8+
• PyTorch (支持 CUDA)
• PyTorch Lightning
• Triton (用于高性能算子优化)
• Git

2. 安装步骤

克隆仓库并安装基础依赖：

git clone https://github.com/ligo/alphafold3.git
cd alphafold3

# 创建虚拟环境 (可选但推荐)
python -m venv venv
source venv/bin/activate

# 安装依赖
pip install -r requirements.txt

提示：若需加速下载，可配置国内镜像源（如清华源）：

pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple

关键优化组件： 本项目集成了由 @alexzhang13 开发的自定义 Triton 内核，用于 MSA 配对加权平均操作，相比原生 PyTorch 实现内存效率提升 10 倍以上，显著降低长序列训练的显存瓶颈。确保安装过程中 triton 包成功编译。

3. 基本使用

由于采样代码（Sampling Code）尚未发布（需等待配体和核酸模型训练完成），当前主要用途为模型加载、实验验证及架构研究。

加载预训练权重进行实验

你可以使用 PyTorch Lightning 加载检查点文件进行模型测试或“模型手术”（Model Surgery）：

import pytorch_lightning as pl
from src.model import AlphaFold3  # 假设的模块路径，具体请参考源码结构

# 加载检查点
checkpoint_path = "path/to/your/checkpoint.ckpt"
model = AlphaFold3.load_from_checkpoint(checkpoint_path)

# 将模型设置为评估模式
model.eval()

# 此时可输入单链蛋白数据进行前向传播测试
# 注意：具体的数据预处理流程需参考 ProteinFlow 管道或源码中的 data loaders

开发者须知

• 模块化设计：模型组件（如 TriangularAttention, DiffusionTransformer）设计为可复用模块，研究人员可将其集成到自己的项目中。
• Beta 测试：如需参与后续配体 - 蛋白及核酸预测功能的 Beta 测试，请访问 等待列表 报名。
• 贡献代码：欢迎提交 Issue 报告 Bug 或提交 PR 修复实现细节。深度学习训练中的静默错误可能导致收敛变差，社区反馈对完善此开源实现至关重要。