All-Atom Novel Protein Sequence Design Using Discrete Diffusion Models

June 18, 2025 · View on GitHub

Abstract

Advancing protein design is crucial for breakthroughs in medicine and biotechnology, yet traditional approaches often fall short by focusing solely on representing protein sequences using the 20 canonical amino acids. This thesis explores discrete diffusion models for generating novel protein sequences with an all-atom representation, specifically SELFIES a widely used molecular string representation. This all-atom approach considers the atomic composition of each amino acid in the protein, enabling the inclusion of non-canonical amino acids and post-translational modifications. Using a modified ByteNet architecture and the D3PM framework, we compare the effects of this all-atom representation to the standard amino acid representation on the generated proteins' quality, diversity and novelty. Additionally, we see how a uniform or absorbing noise process affects the results. While models trained on the all-atom representation struggle to generate fully valid proteins consistently, those successfully designed showed improved novelty and diversity. Moreover, the all-atom representation can achieve comparable structural reliability results from OmegaFold to the amino acid models. Lastly, our results show that the use of an absorbing noise schedule is the most effective for both the all-atom and amino acid representation.

About this project

This project is a modified version of EvoDiff, originally developed by Microsoft. It has been adapted for our research into seeing the functionality of an all-atom protein sequence representation, with significant changes to the original codebase and the addition of new functionality.

This project retains the original LICENSE from EvoDiff and complies with its terms. For more details, refer to the original EvoDiff repository.

License

This project is distributed under the MIT License. See the LICENSE file for details.

Setup

To setup this project, first clone the repository:

git clone https://github.com/Intelligent-molecular-systems/All-Atom-Protein-Sequence-Generation
cd All-Atom-Protein-Sequence-Generation

Makefile

This project includes a Makefile to simplify common tasks. You can use it to run various commands without needing to remember the exact syntax. Running make will display the available targets and their descriptions. All configuration and parameters are defined in the Makefile, allowing you to easily modify them as needed. Running the commands manually is also possible, see the sections below for details.

Docker

This repository provides a Dockerfile to build and run the project in a containerized environment. The Docker image encapsulates all dependencies and provides a reproducible setup.

Build the Docker image

Ensure Docker is installed on your system. From the project root directory, build the Docker image:

docker build -t evodiff:latest .

This command:

Uses the Dockerfile in the current directory.
Creates a Docker image tagged as evodiff:latest.

Run the Docker container

You can start an interactive shell within the container, with the code directory mounted for live development:

docker run -it --rm \
  -v $(pwd):/app \
  -w /app \
  evodiff:latest bash

Options:

-it: Interactive terminal
--rm: Remove container when it exits
-v $(pwd):/app: Mount current directory into /app in container
-w /app: Set working directory inside container

The container’s default CMD activates the Conda environment (evodiff) and drops you into bash.

Run commands in the container

With the container shell:

# Activate environment (if not already):
conda activate evodiff

# Run training locally:
python -u train.py config/config38M.json output_amino_abs/ --representation 'amino acid' --mask 'absorbing'

# Run validation locally:
python -u validate.py config/config38M.json output_amino_abs/ --representation 'amino acid' --mask 'absorbing'

Mount data and outputs

To persist data or results outside the container, mount additional host directories:

docker run -it --rm \
  -v $(pwd):/app \
  -v /path/to/data:/data \
  -v /path/to/results:/results \
  -w /app \
  evodiff:latest bash

Host /path/to/data is available as /data in the container.
Host /path/to/results is available as /results.

Adjust paths and mount points as needed for your workflow.

Apptainer

To run this project in a cluster environment, Apptainer (formerly Singularity) is required for containerization. Follow these steps to build and use the container.

Build the container

Navigate to the apptainer directory:
```
cd apptainer
```
Build the container file evodiff.sif using one of the provided scripts:
```
./build.sh
```
Or, for cluster environments:
```
./build_slurm.sh
```

The container specifications are defined in apptainer/evodiff.def, and the required Python packages are listed in environment.yml.

UniRef50 dataset

The dataset used in this project is the UniRef50 dataset. To download the dataset, follow these steps:

Create a data directory and navigate to it:
```
mkdir data
cd data
```

Download the UniRef50 dataset:

 curl -O ftp://ftp.uniprot.org/pub/databases/uniprot/uniref/uniref50/uniref50.fasta.gz

Extract the downloaded file:
```
tar -xvzf uniref50.fasta.gz
```

Filter the dataset

To filter the dataset, use the create_dataset.py script. This script filters the dataset by sequence length and removes sequences with non-standard amino acids. It also converts the amino acid sequences to the SELFIES representation. Please inspect the file beforehand and run it in stages since it might take a long time to process. The file handles file splitting and creates a unified H5 dataset file from which the MyUniRefDataset.py file can create a Dataset object. To create the dataset, use the following command:

python create_dataset.py

Run data analysis

python data_analysis/data_analysis.py

Training and validation

WANDB setup

Use Weights & Biases to track training:

export WANDB_API_KEY='[your_api_key]'
# Or add to ~/.bashrc:
echo "export WANDB_API_KEY='[your_api_key]'" >> ~/.bashrc
source ~/.bashrc

or create the file wandb.env and add the api key in the file:

touch wandb.env
echo WANDB_API_KEY='[your_api_key]' >> wandb.env

Run training

python -u train.py config/config38M.json output_amino_abs/ \
  --representation 'amino acid' --mask 'absorbing' \
  --warmup --wandb_run_id 'kvrevwsc' --reweighting_term 0.1

Run validation

Run validation of selected checkpoints:

python -u validate.py config/config38M.json output_amino_abs/ \
  --representation 'amino acid' --mask 'absorbing' \
  --reweighting_term 0.1 --checkpoint_dir output_amino_abs/ \
  --checkpoint_list checkpoint1 checkpoint2 checkpoint3

git clone https://github.com/HeliXonProtein/OmegaFold
cd OmegaFold
curl -O https://helixon.s3.amazonaws.com/release1.pt
mkdir gen_sequences

Run generation and evaluation pipeline

python my_pipeline.py --representation="amino" \
  --seq_file="generated_sequences/uniform/amino/true_amino_acid.fasta" \
  --masking="uniform" --results_file_name="output_amino_uni_random"

python my_pipeline.py --representation="amino" \
  --model_checkpoint="${checkpoint}" \
  --model_config="config/config38M.json" --n_sequences=1000 --random \
  --masking="uniform" --results_file_name="output_amino_uni_random"

Per-sequence length analysis

Use per_length_experiments.ipynb for detailed length-based analysis.

Additional Scripts

shell.sh to open container shell for debugging
cleanup_checkpoints.sh to keep only the last checkpoint

Citation

If you find this repository useful in your research, please cite the following works.

@article{villegas2025allatom,
  title={All-Atom Protein Sequence Design using Discrete Diffusion Models},
  author={Amelia Villegas-Morcillo and Gijs J. Admiraal and Marcel J.T. Reinders and Jana M. Weber},
  journal={bioRxiv},
  year={2025},
}

Notes

If you encounter issues, please do not hesitate to contact the project maintainers for assistance.