AlphaFlow

AlphaFlow是AlphaFold的一个修改版本，通过流匹配目标进行了微调，旨在对蛋白质构象集合进行生成建模。具体来说，AlphaFlow旨在模拟：

实验集合，即可能被存储在PDB中的潜在构象状态
生理温度下的分子动力学集合

我们还提供了一个类似微调的ESMFold版本，称为ESMFlow。技术细节和全面的基准测试结果可以在我们的论文《AlphaFold遇上流匹配：生成蛋白质集合》中找到，作者为Bowen Jing、Bonnie Berger和Tommi Jaakkola。该存储库包含运行该方法所需的所有代码、说明和模型权重。如果您有任何问题，请随时提出issue或通过bjing@mit.edu联系我们。

2024年6月更新： 我们训练了一个12层版本的AlphaFlow-MD+Templates（基础版和蒸馏版），其运行速度比48层版本快2.5倍，但性能略有损失。如果有参考结构（PDB或AlphaFold）可用且运行时间是高优先级，我们建议考虑使用这个模型。

安装

在Python 3.9环境中（例如，conda create -n [名称] python=3.9），运行：

pip install numpy==1.21.2 pandas==1.5.3
pip install torch==1.12.1+cu113 -f https://download.pytorch.org/whl/torch_stable.html
pip install biopython==1.79 dm-tree==0.1.6 modelcif==0.7 ml-collections==0.1.0 scipy==1.7.1 absl-py einops
pip install pytorch_lightning==2.0.4 fair-esm mdtraj wandb
pip install 'openfold @ git+https://github.com/aqlaboratory/openfold.git@103d037'

我们在CUDA 11.6的机器上进行了安装，并在A100和A6000 GPU上进行了测试。（有关如何在Conda中安装特定CUDA版本的说明，请参见此链接。）

模型权重

我们提供了几个版本的AlphaFlow（以及类似命名的ESMFlow版本）。

AlphaFlow-PDB——在PDB结构上训练，用于模拟X射线晶体学或冷冻电镜在不同条件下的实验集合
AlphaFlow-MD——在300K温度下的全原子、显式溶剂MD轨迹上训练
AlphaFlow-MD+Templates——额外训练以接受PDB结构作为输入，并模拟相应的300K温度下的MD集合

对于所有模型，蒸馏版本运行速度显著更快，但会牺牲一些准确性（在论文中进行了基准测试）。

对于AlphaFlow-MD+Templates，12l版本具有12个而不是48个Evoformer层，运行速度快2.5倍，但性能略有损失。

AlphaFlow模型

模型	版本	权重
AlphaFlow-PDB	基础	https://alphaflow.s3.amazonaws.com/params/alphaflow_pdb_base_202402.pt
AlphaFlow-PDB	蒸馏	https://alphaflow.s3.amazonaws.com/params/alphaflow_pdb_distilled_202402.pt
AlphaFlow-MD	基础	https://alphaflow.s3.amazonaws.com/params/alphaflow_md_base_202402.pt
AlphaFlow-MD	蒸馏	https://alphaflow.s3.amazonaws.com/params/alphaflow_md_distilled_202402.pt
AlphaFlow-MD+Templates	基础	https://alphaflow.s3.amazonaws.com/params/alphaflow_md_templates_base_202402.pt
AlphaFlow-MD+Templates	蒸馏	https://alphaflow.s3.amazonaws.com/params/alphaflow_md_templates_distilled_202402.pt
AlphaFlow-MD+Templates	12l-基础	https://alphaflow.s3.amazonaws.com/params/alphaflow_12l_md_templates_base_202406.pt
AlphaFlow-MD+Templates	12l-蒸馏	https://alphaflow.s3.amazonaws.com/params/alphaflow_12l_md_templates_distilled_202406.pt

ESMFlow模型

模型	版本	权重
ESMFlow-PDB	基础	https://alphaflow.s3.amazonaws.com/params/esmflow_pdb_base_202402.pt
ESMFlow-PDB	蒸馏	https://alphaflow.s3.amazonaws.com/params/esmflow_pdb_distilled_202402.pt
ESMFlow-MD	基础	https://alphaflow.s3.amazonaws.com/params/esmflow_md_base_202402.pt
ESMFlow-MD	蒸馏	https://alphaflow.s3.amazonaws.com/params/esmflow_md_distilled_202402.pt
ESMFlow-MD+Templates	基础	https://alphaflow.s3.amazonaws.com/params/esmflow_md_templates_base_202402.pt
ESMFlow-MD+Templates	蒸馏	https://alphaflow.s3.amazonaws.com/params/esmflow_md_templates_distilled_202402.pt

训练检查点（可从中恢复微调）可根据要求提供；如果您想合作，请联系我们！

运行推理

准备输入文件

准备一个包含每行name和seqres条目的输入CSV文件。参见splits/atlas_test.csv获取示例。
如果运行AlphaFlow模型，准备一个MSA目录并将.a3m格式的比对放置在以下路径：{alignment_dir}/{name}/a3m/{name}.a3m。如果您没有MSA，有两种方法可以生成它们：
1. 使用python -m scripts.mmseqs_query --split [路径] --outdir [目录]查询ColabFold服务器。
2. 根据https://github.com/sokrypton/ColabFold/blob/main/setup_databases.sh下载UniRef30和ColabDB，并运行`python -m scripts.mmseqs_search_helper --split [路径] --db_dir [目录] --outdir [目录]`。
如果运行MD+Templates模型，将模板PDB文件放入模板目录，文件名与输入CSV中的名称匹配。PDB文件应仅包含一个链，且没有残基间隙。

运行模型

使用AlphaFlow运行推理的基本命令是：

python predict.py --mode alphafold --input_csv [路径] --msa_dir [目录] --weights [路径] --samples [N] --outpdb [目录]

如果运行PDB模型，我们建议附加--self_cond --resample以提高性能。

使用ESMFlow运行推理的基本命令是：

python predict.py --mode esmfold --input_csv [路径] --weights [路径] --samples [N] --outpdb [目录]

任一模型的其他命令行参数：

使用--pdb_id参数选择CSV中的一行或多行。如果未指定参数，将对所有行进行推理。
如果运行带模板的MD模型，附加--templates_dir [目录]。
如果运行任何蒸馏模型，附加参数--noisy_first --no_diffusion。
要截断推理过程以提高精度并减少多样性，附加（例如）--tmax 0.2 --steps 2。默认推理设置对应于--tmax 1.0 --steps 10。有关更多详细信息，请参见论文附录B.1。

评估脚本

我们的集合评估可以通过以下步骤重现：

从所需的根目录运行bash scripts/download_atlas.sh下载ATLAS数据集
准备集合目录，为每个ATLAS目标提供一个PDB文件，每个文件包含250个结构（参见下面的压缩AlphaFlow集合示例）。对于具有不同结构数量的评估，某些结果可能无法直接比较。
运行python -m scripts.analyze_ensembles --atlas_dir [目录] --pdb_dir [目录] --num_workers [N]。这将在pdb_dir中生成一个名为out.pkl的分析文件。
运行python -m scripts.print_analysis [路径] [路径] ...，其中包含任意数量的out.pkl文件路径。将打印格式化的比较表。

训练

下载数据集

要下载和预处理PDB，

从所需目录运行aws s3 sync --no-sign-request s3://pdbsnapshots/20230102/pub/pdb/data/structures/divided/mmCIF pdb_mmcif。
运行find pdb_mmcif -name '*.gz' | xargs gunzip以提取MMCIF文件。
从存储库根目录运行python -m scripts.unpack_mmcif --mmcif_dir [目录] --outdir [目录] --num_workers [N]。这将预处理所有链为NPZ文件并创建一个pdb_mmcif.csv索引。
从所需目录运行aws s3 sync --no-sign-request s3://openfold/ openfold下载OpenProteinSet。
运行python -m scripts.add_msa_info --openfold_dir [目录]以生成带有OpenProteinSet MSA查找的pdb_mmcif_msa.csv索引。
运行python -m scripts.cluster_chains以生成40%序列相似性的pdb_clusters文件（需要Mmseqs安装）。
根据上一节的说明为PDB验证集（splits/cameo2022.csv）创建MSA。

要下载和预处理ATLAS MD轨迹数据集，

从所需目录运行bash scripts/download_atlas.sh。
从存储库根目录运行python -m scripts.prep_atlas --atlas_dir [目录] --outdir [目录] --num_workers [N]。这将预处理ATLAS轨迹为NPZ文件。
根据上一节的说明为splits/atlas.csv中的所有条目创建MSA。

运行训练

在运行训练之前，通过以下命令将预训练的AlphaFold和ESMFold权重下载到存储库根目录：

wget https://storage.googleapis.com/alphafold/alphafold_params_2022-12-06.tar
tar -xvf alphafold_params_2022-12-06.tar params_model_1.npz
wget https://dl.fbaipublicfiles.com/fair-esm/models/esmfold_3B_v1.pt

AlphaFlow的基本训练命令是：

python train.py --lr 5e-4 --noise_prob 0.8 --accumulate_grad 8 --train_epoch_len 80000 --train_cutoff 2018-05-01 --filter_chains \
    --train_data_dir [目录] \
    --train_msa_dir [目录] \
    --mmcif_dir [目录] \
    --val_msa_dir [目录] \
    --run_name [名称] [--wandb]

其中需要指定PDB NPZ目录、OpenProteinSet目录、PDB mmCIF目录和验证MSA目录。这个训练运行会产生AlphaFlow-PDB基础版本。所有其他模型都是基于这个检查点构建的。

要继续在ATLAS上训练，运行：

python train.py --normal_validate --sample_train_confs --sample_val_confs --num_val_confs 100 --pdb_chains splits/atlas_train.csv --val_csv splits/atlas_val.csv --self_cond_prob 0.0 --noise_prob 0.9 --val_freq 10 --ckpt_freq 10 \
    --train_data_dir [目录] \
    --train_msa_dir [目录] \
    --ckpt [路径] \
    --run_name [名称] [--wandb]

其中需要指定ATLAS MSA和NPZ目录以及AlphaFlow-PDB检查点。

如果要在ATLAS上使用模板进行训练，添加以下额外参数运行：--first_as_template --extra_input --lr 1e-4 --restore_weights_only --extra_input_prob 1.0。

蒸馏：要蒸馏模型，附加--distillation并提供要蒸馏的模型的--ckpt [路径]。对于PDB训练，我们移除--accumulate_grad 8，并建议使用更短的--train_epoch_len 16000进行蒸馏。注意，--self_cond_prob和--noise_prob将被忽略，可以省略。

ESMFlow：使用相同的命令运行，但加上--mode esmfold和--train_cutoff 2020-05-01。

集成

我们提供了从模型中采样的集成，这些集成用于论文中报告的分析和结果。

AlphaFlow集成

模型	版本	样本
AlphaFlow-PDB	基础	https://alphaflow.s3.amazonaws.com/samples/alphaflow_pdb_base_202402.zip
AlphaFlow-PDB	蒸馏	https://alphaflow.s3.amazonaws.com/samples/alphaflow_pdb_distilled_202402.zip
AlphaFlow-MD	基础	https://alphaflow.s3.amazonaws.com/samples/alphaflow_md_base_202402.zip
AlphaFlow-MD	蒸馏	https://alphaflow.s3.amazonaws.com/samples/alphaflow_md_distilled_202402.zip
AlphaFlow-MD+Templates	基础	https://alphaflow.s3.amazonaws.com/samples/alphaflow_md_templates_base_202402.zip
AlphaFlow-MD+Templates	蒸馏	https://alphaflow.s3.amazonaws.com/samples/alphaflow_md_templates_distilled_202402.zip
AlphaFlow-MD+Templates	12l-基础	https://alphaflow.s3.amazonaws.com/samples/alphaflow_12l_md_templates_base_202406.zip
AlphaFlow-MD+Templates	12l-蒸馏	https://alphaflow.s3.amazonaws.com/samples/alphaflow_12l_md_templates_distilled_202406.zip

ESMFlow集成

模型	版本	样本
ESMFlow-PDB	基础	https://alphaflow.s3.amazonaws.com/samples/esmflow_pdb_base_202402.zip
ESMFlow-PDB	蒸馏	https://alphaflow.s3.amazonaws.com/samples/esmflow_pdb_distilled_202402.zip
ESMFlow-MD	基础	https://alphaflow.s3.amazonaws.com/samples/esmflow_md_base_202402.zip
ESMFlow-MD	蒸馏	https://alphaflow.s3.amazonaws.com/samples/esmflow_md_distilled_202402.zip
ESMFlow-MD+Templates	基础	https://alphaflow.s3.amazonaws.com/samples/esmflow_md_templates_base_202402.zip
ESMFlow-MD+Templates	蒸馏	https://alphaflow.s3.amazonaws.com/samples/esmflow_md_templates_distilled_202402.zip

许可证

MIT。其他许可证可能适用于文件头中注明的第三方源代码。

引用

@misc{jing2024alphafold,
      title={AlphaFold Meets Flow Matching for Generating Protein Ensembles}, 
      author={Bowen Jing and Bonnie Berger and Tommi Jaakkola},
      year={2024},
      eprint={2402.04845},
      archivePrefix={arXiv},
      primaryClass={q-bio.BM}
}