Open-AlphaFold

基于论文《使用 AlphaFold3 准确预测生物分子相互作用的结构》的 AlphaFold 开源实现，采用 PyTorch 框架。本人及本仓库的贡献者与 Google 或 DeepMind 均无任何关联。

安装

$ pip install alphafold3

输入张量尺寸示例

import torch

# 定义批次大小、节点数和特征数
batch_size = 1
num_nodes = 5
num_features = 64

# 使用 torch.randn 生成随机的成对表示
# 形状：(batch_size, num_nodes, num_nodes, num_features)
pair_representations = torch.randn(
    batch_size, num_nodes, num_nodes, num_features
)

# 使用 torch.randn 生成随机的单体表示
# 形状：(batch_size, num_nodes, num_features)
single_representations = torch.randn(
    batch_size, num_nodes, num_features
)

遗传扩散

尚需审核，但其核心操作是基于原子坐标进行的。

import torch
from alphafold3.diffusion import GeneticDiffusion

# 创建 GeneticDiffusionModuleBlock 实例
model = GeneticDiffusion(channels=3, training=True)

# 生成随机输入坐标
input_coords = torch.randn(10, 100, 100, 3)

# 生成随机真实坐标
ground_truth = torch.randn(10, 100, 100, 3)

# 将输入坐标和真实坐标输入模型
output_coords, loss = model(input_coords, ground_truth)

# 打印输出坐标
print(output_coords)

# 打印损失值
print(loss)

全模型示例前向传播

import torch 
from alphafold3 import AlphaFold3

# 创建随机张量
x = torch.randn(1, 5, 5, 64)  # 形状：(batch_size, seq_len, seq_len, dim)
y = torch.randn(1, 5, 64)  # 形状：(batch_size, seq_len, dim)

# 初始化 AlphaFold3 模型
model = AlphaFold3(
    dim=64,
    seq_len=5,
    heads=8,
    dim_head=64,
    attn_dropout=0.0,
    ff_dropout=0.0,
    global_column_attn=False,
    pair_former_depth=48,
    num_diffusion_steps=1000,
    diffusion_depth=30,
)

# 进行前向传播
output = model(x, y)

# 打印输出张量的形状
print(output.shape)

笔记

-> 成对表示 -> 显式原子位置

-> 在主干网络中，MSA 处理被弱化，仅使用一个更简单的 MSA 模块，共 4 层。

-> MSA 处理 -> 对成对表示进行加权平均。

-> Pairformer：取代了 Evoformer，直接作用于成对表示和单体表示。

-> Pairformer 共 48 层。

-> 成对表示、单体表示以及输入表示一同传递给扩散模块。

-> 扩散模块接收 3 个张量 [成对表示、单体表示，以及新的 Pairformer 表示]。

-> 扩散模块直接操作原始原子坐标。

-> 采用标准扩散方法，模型训练时接收带噪声的原子坐标，并预测真实的原子坐标。

-> 网络在不同长度尺度上学习蛋白质结构，其中小噪声下的去噪任务更强调系统的宏观结构。

-> 推理时，先采样随机噪声，然后逐步去噪以生成最终结构。

-> 扩散模块会生成一组可能的结构。

-> 对于每种可能的结构，局部细节都会非常清晰。

-> 扩散模型容易产生“幻觉”，即生成看似合理但实际上不准确的结构。

-> 为应对这一问题，他们采用了一种新颖的交叉蒸馏方法，在训练数据中加入了 AlphaFold Multimer v2.3 预测的结构。

-> 置信度评估可以预测最终结构中的原子级误差和成对误差，这是通过在训练过程中回归结构模块输出的误差来实现的。

-> 训练时采用扩散展开过程生成完整结构（步长比正常情况更大）。

-> 使用扩散生成的预测结构来排列真实结构和配体，从而计算指标以训练置信度头。

-> 置信度头利用成对表示预测 LDDT（pLDDT），并生成类似于 AlphaFold 2 中使用的对齐误差矩阵，以及预测结构与真实结构的距离误差矩阵。

-> 置信度评估还能预测原子级和成对误差。

-> 根据上述各项指标的加权平均进行早停。

-> AF3 可以根据输入的聚合物序列、修饰信息和配体 SMILES 预测结构。

-> 适用于少于 1000 个残基的结构。

-> AlphaFold3 还能预测包含数千个残基的蛋白质核结构。

-> 共价修饰（如结合配体、糖基化以及修饰后的蛋白质残基和 202 个核酸碱基）也能被 AF 准确预测。

-> 蒸馏了 AlphaFold2 的预测结果。

-> 蛋白质结构预测的一个关键问题是，目前的方法只能预测静态结构，而无法捕捉动态行为。

-> 即使使用多个随机种子运行扩散头或整个网络，也无法近似得到解的集合。

-> 未来计划：生成大量预测结果并对其进行排序。

-> 推理时：从 5 次种子运行和每个模型种子的 5 次扩散运行中，选择置信度最高的样本，总共得到 25 个样本。

-> 接口精度通过接口 LDDT 来衡量，该指标基于接口处不同链之间原子间的距离计算得出。

-> 使用一种针对聚合物的 LDDT 指标，该指标考虑实体中的每个原子与半径范围内的任意 C 或 C1 聚合物原子之间的差异。

待办事项

模型架构

• 实现来自 AlphaFold2 openfold 的输入嵌入模块 链接
• 实现来自 openfold 的模板模块 链接
• 实现来自 openfold 的 MSA 嵌入 链接
• 修复残差连接，并确保成对表示和生成的输出能够正确进入扩散模型。
• 实现循环机制以修复残差连接。

训练流程

• 将所有数据集上传至 Hugging Face。

资料

• EvoFormer 论文

• Pairformer

• AlphaFold 3 论文

• OpenFold

参考文献

@article{Abramson2024-fj,
  title    = {使用{AlphaFold} 3准确预测生物分子相互作用的结构},
  author   = {阿布拉姆森，乔什；阿德勒，乔纳斯；邓格，杰克；埃文斯，理查德；格林，蒂姆；普里茨尔，亚历山大；罗内贝格，奥拉夫；威尔莫尔，林赛；巴拉德，安德鲁·J；班布里克，乔舒亚；博登斯坦，塞巴斯蒂安·W；埃文斯，大卫·A；洪，贾春；奥尼尔，迈克尔；雷曼，大卫；图尼亚苏瓦纳库尔，凯瑟琳；吴，扎卡里；泽姆古利特，阿克维莱；阿尔瓦尼蒂，艾琳妮；比蒂，查尔斯；贝尔托利，奥塔维娅；布里奇兰德，亚历克斯；切列帕诺夫，阿列克谢；康格里夫，迈尔斯；科温-里弗斯，亚历山大·I；科伊，安德鲁；菲古尔诺夫，米哈伊尔；福克斯，法比安·B；格拉德曼，汉娜；贾因，里舒布；汗，优素福·A；洛，卡罗琳·M·R；佩尔林，库巴；波塔彭科，安娜；萨维，帕斯卡尔；辛格，苏赫迪普；斯泰库拉，阿德里安；蒂拉孙达拉姆，阿肖克；通，凯瑟琳；亚克宁，谢尔盖；钟，艾伦·D；齐耶林斯基，米哈尔；日代克，奥古斯丁；巴普斯特，维克托；科利，普什米特；雅德伯格，马克斯；哈萨比斯，德米斯；詹珀，约翰·M},
  journal  = {自然},
  month    =  "五月",
  year     =  2024
}

@inproceedings{Darcet2023VisionTN,
    title   = {视觉Transformer需要寄存器},
    author  = {蒂莫泰·达尔塞、马克西姆·欧卡布、朱利安·梅拉尔、皮奥特尔·博扬诺夫斯基},
    year    = {2023},
    url     = {https://api.semanticscholar.org/CorpusID:263134283}
}

@article{Arora2024SimpleLA,
    title   = {简单的线性注意力语言模型平衡召回率与吞吐量的权衡},
    author  = {西姆兰·阿罗拉、萨布里·埃尤博格鲁、迈克尔·张、阿曼·蒂马尔西纳、西拉斯·阿尔贝蒂、迪伦·津斯利、詹姆斯·邹、阿特里·鲁德拉、克里斯托弗·R'e},
    journal = {ArXiv},
    year    = {2024},
    volume  = {abs/2402.18668},
    url     = {https://api.semanticscholar.org/CorpusID:268063190}
}

@article{Puny2021FrameAF,
    title   = {用于不变性和等变网络设计的帧平均},
    author  = {奥姆里·普尼、马坦·阿茨蒙、赫利·本-哈穆、爱德华·詹姆斯·史密斯、伊桑·米斯拉、阿迪提亚·格罗弗、亚龙·利普曼},
    journal = {ArXiv},
    year    = {2021},
    volume  = {abs/2110.03336},
    url     = {https://api.semanticscholar.org/CorpusID:238419638}
}

@article{Duval2023FAENetFA,
    title   = {FAENet：用于材料建模的帧平均等变图神经网络},
    author  = {亚历山大·杜瓦尔、维克托·施密特、亚历克斯·埃尔南德斯·加西亚、圣地亚哥·米雷特、弗拉吉斯科斯·D·马利亚罗斯、约书亚·本吉奥、大卫·罗尔尼克},
    journal = {ArXiv},
    year    = {2023},
    volume  = {abs/2305.05577},
    url     = {https://api.semanticscholar.org/CorpusID:258564608}
}

@article{Wang2022DeepNetST,
    title   = {DeepNet：将Transformer扩展到1,000层},
    author  = {洪宇·王、舒明·马、李东、绍韩·黄、董东·张、富鲁·魏},
    journal = {ArXiv},
    year    = {2022},
    volume  = {abs/2203.00555},
    url     = {https://api.semanticscholar.org/CorpusID:247187905}
}

@inproceedings{Ainslie2023CoLT5FL,
    title   = {CoLT5：利用条件计算加速长距离Transformer},
    author  = {乔舒亚·艾恩斯利、陶雷、米希尔·德·容、圣地亚哥·昂塔翁、西达尔塔·布拉玛、尤里·泽姆良斯基、戴维·乌瑟斯、曼迪·郭、詹姆斯·李-索普、易泰、云轩·宋、苏米特·桑盖},
    year    = {2023}
}

引用

@article{Abramson2024-fj,
  title    = {使用{AlphaFold} 3准确预测生物分子相互作用的结构},
  author   = {阿布拉姆森，乔什；阿德勒，乔纳斯；邓格，杰克；埃文斯，理查德；格林，蒂姆；普里茨尔，亚历山大；罗内贝格，奥拉夫；威尔莫尔，林赛；巴拉德，安德鲁·J；班布里克，乔舒亚；博登斯坦，塞巴斯蒂安·W；埃文斯，大卫·A；洪，贾春；奥尼尔，迈克尔；雷曼，大卫；图尼亚苏瓦纳库尔，凯瑟琳；吴，扎卡里；泽姆古利特，阿克维莱；阿尔瓦尼蒂，艾琳妮；比蒂，查尔斯；贝尔托利，奥塔维娅；布里奇兰德，亚历克斯；切列帕诺夫，阿列克谢；康格里夫，迈尔斯；科温-里弗斯，亚历山大·I；科伊，安德鲁；菲古尔诺夫，米哈伊尔；福克斯，法比安·B；格拉德曼，汉娜；贾因，里舒布；汗，优素福·A；洛，卡罗琳·M·R；佩尔林，库巴；波塔彭科，安娜；萨维，帕斯卡尔；辛格，苏赫迪普；斯泰库拉，阿德里安；蒂拉孙达拉姆，阿肖克；通，凯瑟琳；亚克宁，谢尔盖；钟，艾伦·D；齐耶林斯基，米哈尔；日代克，奥古斯丁；巴普斯特，维克托；科利，普什米特；雅德伯格，马克斯；哈萨比斯，德米斯；詹珀，约翰·M},
  journal  = {自然},
  month    =  "五月",
  year     =  2024
}

Open-AF3 快速上手指南

Open-AF3 是基于 PyTorch 实现的 AlphaFold3 开源项目，旨在复现论文《Accurate structure prediction of biomolecular interactions with AlphaFold3》中的核心算法。本项目由社区贡献者维护，与 Google 或 DeepMind 无官方关联。

环境准备

在开始之前，请确保您的开发环境满足以下要求：

• 操作系统：Linux 或 macOS（Windows 用户建议使用 WSL2）
• Python 版本：推荐 Python 3.8 及以上版本
• 深度学习框架：PyTorch（需支持 CUDA 加速以获得最佳性能）
• 依赖管理：pip

提示：国内开发者在安装大型依赖时，建议配置清华或阿里镜像源以加速下载： pip config set global.index-url https://pypi.tuna.tsinghua.edu.cn/simple

安装步骤

通过 pip 直接安装预发布的 alphafold3 包：

pip install alphafold3

基本使用

以下是初始化模型并执行一次前向传播（Forward Pass）的最小可用示例。该示例展示了如何构建随机输入张量并获取模型输出。

完整模型前向传播示例

import torch 
from alphafold3 import AlphaFold3

# 创建随机输入张量
# x: 成对表示 (batch_size, seq_len, seq_len, dim)
x = torch.randn(1, 5, 5, 64)
# y: 单体表示 (batch_size, seq_len, dim)
y = torch.randn(1, 5, 64)

# 初始化 AlphaFold3 模型
# 注意：此处参数为示例配置，实际训练或推理可根据需求调整
model = AlphaFold3(
    dim=64,
    seq_len=5,
    heads=8,
    dim_head=64,
    attn_dropout=0.0,
    ff_dropout=0.0,
    global_column_attn=False,
    pair_former_depth=48,
    num_diffusion_steps=1000,
    diffusion_depth=30,
)

# 执行前向传播
output = model(x, y)

# 打印输出张量的形状
print(output.shape)

核心模块说明

• 输入表示：模型接收“成对表示”（pair representations）和“单体表示”（single representations）。
• PairFormer：取代了 AlphaFold2 中的 Evoformer，专门处理成对和单体表示。
• 扩散模块 (Diffusion Module)：直接作用于原始原子坐标。通过去噪过程，从随机噪声中生成最终的蛋白质结构分布。