foldingdiff generates protein backbones by diffusing internal angles

Angles instead of coordinates, which is the whole idea

foldingdiff is a diffusion model for generating novel protein backbone structures, and the subtitle of the project states the technical bet plainly: trigonometry and attention are all you need. The paper, linked from the project page as a Nature Communications article, and the arXiv preprint, describe a model that works in the space of internal angles rather than in Cartesian coordinates. The consequence for a user is that the generated output is a set of backbone torsions which the pipeline then converts to three dimensional coordinates, and the sampled angles are what the training objective actually sees. That choice is what lets the model use a plain transformer with attention, and it is why the project description leads with trigonometry and attention. The audience is therefore a research audience rather than a drug discovery audience. The project is a reference implementation with a training path, not a product, and the repository mixes a library package with scripts, configs, notebooks and a test suite, which is what you expect from code released alongside a paper.

The environment is pinned hard, and that is a real constraint

The installation is conda based and the required environment is defined in an environment.yml file in the repository, so the setup is a three step process after you have conda or mamba installed and the repository cloned:

conda env create -f environment.yml
conda activate foldingdiff
pip install -e ./  # make sure ./ is the dir including setup.py

The last line is worth reading twice, because the editable install has to run from the directory that contains setup.py. The dependency versions are where the friction is, and they are old. The requirements file pins numpy at 1.22, pandas at 1.1, torch at 1.12, scipy at 1.9.1, scikit-learn at 1.2.1, transformers at 4.11.3, pytorch-lightning at 1.6.4 and biotite at 0.34, while huggingface_hub, seaborn, mpl-scatter-density, astropy, gitpython and requests are unpinned. The setup script lists the same set as install_requires, and the packaging metadata requires Python 3.8 or newer with the project version at 0.0.1. Pinning numpy at 1.22 and pandas at 1.1 in a project whose environment is now several years old means the conda environment file is likely to be the path of least resistance, and that you should expect to work around version conflicts if you try to install this into an existing modern environment. The Dockerfile-based route is not offered here, so the container path is yours to build if you need one.

The CATH data has to be fetched separately, and only training needs it

Some data files are not packaged on GitHub because of their size, and the documentation is clear about when you need them. They are not required for sampling, as long as you are not using the --testcomparison option, and they are required if you want to train your own model. A script in the data directory does the download, and the instructions are explicit that you have to be in that subdirectory first:

# Download the CATH dataset
cd data  # Ensure that you are in the data subdirectory within the codebase
chmod +x download_cath.sh
./download_cath.sh

There is a fallback for when the download link in the shell script is not working, since the tarball is also mirrored on Dropbox, and the file name in that mirror is a non redundant S40 PDB set. The directory listing at the top of the repository shows a .gitmodules file, which suggests some parts of the tree arrive as submodules, so a fresh clone that skips submodules may not be complete. The data requirement and the dataset choice also tell you what the model is and is not. CATH is a curated set of protein domain structures classified by fold, and a model trained on it learns the distribution of folds that appear in that classification. That is a useful research substrate and a narrow design space, which is the second reason not to treat this as a general structure predictor.

Training is a config file and a script, and the dry run flag matters

Training your own model means picking one of the json config files under config_jsons, or writing your own, and running the script at bin/train.py:

python bin/train.py config_jsons/cath_full_angles_cosine.json --dryrun

The flag in that example is the important detail. By default the training script calculates the KL divergence at each timestep before it starts training, which the documentation says can be quite computationally expensive with more timesteps, and the --dryrun flag skips that. The config file name itself also tells you something, since it points at the full angle set and a cosine beta schedule, and the hyperparameters live separately under a hyperparam_jsons directory. The output layout is fixed and worth knowing in advance, because the results folder holds the config.json for the huggingface BERT model itself, the logs from training, the model checkpoints where the best five models by validation loss and the best five by training loss are kept by default, and a training_args.json with the full set of arguments, which the documentation says can be used to reproduce a run. That last file is the one to keep, since a diffusion run is only reproducible if you kept the arguments and the config together.

Loading the pretrained model and running a forward pass

The published weights are on the HuggingFace model hub at wukevin/foldingdiff_cath, and the documented snippet shows the whole loop: download the model, build a dataset, add noise, and run a forward pass. The model is loaded from the hub snapshot with files cached for future calls, the clean dataset is a CathCanonicalAnglesOnlyDataset with a padding of 128 and a randomcrop trim strategy, and the noised dataset wraps it with 1000 timesteps and a cosine beta schedule. The data loader takes a batch size of 32 without shuffling, and the forward pass takes the corrupted angles and the timestep with an attention mask:

from huggingface_hub import snapshot_download
from torch.utils.data.dataloader import DataLoader
from foldingdiff import modelling
from foldingdiff import datasets as dsets

# Load the model (files will be cached for future calls)
m = modelling.BertForDiffusion.from_dir(snapshot_download("wukevin/foldingdiff_cath"))

# Load dataset
# As part of loading, we try to compute internal angles in parallel. This may
# throw warnings like the following; this is normal.
# WARNING:root:Illegal values for omega in /home/*/projects/foldingdiff-main/data/cath/dompdb/2ebqA00 -- skipping
# After computing these once, the results are saved in a .pkl file under the
# foldingdiff source directory for faster loading in future calls.
clean_dset = dsets.CathCanonicalAnglesOnlyDataset(pad=128, trim_strategy='randomcrop')
noised_dset = dsets.NoisedAnglesDataset(clean_dset, timesteps=1000, beta_schedule='cosine')
dl = DataLoader(noised_dset, batch_size=32, shuffle=False)
x = iter(dl).next()

# Forward pass
predicted_noise = m(x['corrupted'], x['t'], x['attn_mask'])

The warning about illegal values for omega is called out as normal, and the internal angles are computed in parallel on first load and then cached in a .pkl file under the source directory for faster loading afterwards. That cache behaviour is worth remembering, because it means the first load of a dataset is slow and subsequent ones are not.

Sampling is a script, and the runtime table tells you what you need

Generating structures uses the script at bin/sample.py, and the documented example samples 10 backbones per length across the range from 50 to 128 with a batch size of 512, which the documentation says reproduces the results in the manuscript:

# To sample 10 backbones per length ranging from [50, 128) with a batch size of 512 - reproduces results in our manuscript
python ~/projects/foldingdiff/bin/sample.py -l 50 128 -n 10 -b 512 --device cuda:0

Omitting the device flag defaults to the first device, cuda:0, and using --device cpu works but is described as very slow. The documentation includes a runtime table for sampling 512 structures, with seven minutes on an Nvidia RTX 2080Ti and two hours on an i9-9960X with 16 physical cores, which is a factor of roughly seventeen and a clear statement about what hardware this expects. The output directory structure is fixed: a plots folder comparing the distribution of training and generated angles, a sampled_angles folder with .csv.gz files, a sampled_pdb folder with the .pdb files produced by converting the sampled angles to cartesian coordinates, and a model_snapshot folder holding a copy of the model used to produce the results. A --fullhistory flag writes an additional sample_history subdirectory under each of the sampled_angles and sampled_pdb folders, with the pdb and csv files for every timestep, which is what makes the folding animation possible.

What you get is a backbone, and the browser demos are the easy on-ramp

The limitation to state plainly is scope. The model produces backbone structures, meaning the torsions that define the fold, and it does not design sequences, predict stability, or evaluate whether a generated fold is worth making. Two related capabilities exist for the research workflow. TM-scores can be calculated to evaluate the similarity of the generated sequences against the originals, using the script at bin/tmscore_training.py, which needs the downloaded data. And the folding process itself can be visualised, since the pdb files written by the full history flag can be passed into the script at foldingdiff/pymol_vis.py to produce a gif of the folding, which is the animation the project shows. For people who do not want to set up the environment at all, the project hosts a trained version of the model on HuggingFace spaces and on SuperBio, so you can generate protein structures from a browser. The result is a sensible adoption ladder: try it in the browser, then sample locally with the pretrained weights, then retrain on the CATH data if your research question needs it.

Editorial conclusion

Adopt foldingdiff if you are doing research on protein structure generation and want a reference implementation with pretrained weights, a documented sampling script and a training path you can rerun on your own data. Do not adopt it as a protein design tool, since the output is a backbone without sequence, folding or stability prediction, and the model was trained on CATH. Before you start, decide whether you need a GPU, since the documented runtime for sampling 512 structures is seven minutes on an RTX 2080Ti against two hours on a 16 core CPU, and check the pinned dependency versions, which are old enough to be a real constraint on a modern Python environment.

Frequently asked questions

Official sources

1. License: MIT
2. microsoft/foldingdiff on GitHub
3. Project website
4. README
5. Releases