Improved Fitness Optimization Landscapes for Sequence Design

• Description
• Citation
• How to run
• Training models
• Original data source

In recent years, deep learning approaches for determining protein sequence-fitness relationships have gained traction. Advances in high-throughput mutagenesis, directed evolution, and next-generation sequencing have allowed for the accumulation of large amounts of labelled fitness data and consequently, attracted the application of various deep learning methods. Although these methods learn an implicit fitness landscape, there is little work on using the latent encoding directly for protein sequence optimization. Here we show that this latent space representation of a fitness landscape can be made very amenable to latent space optimization through a joint-training process. We also show that this encoding strategy which also provides improvements to generalization over more traditional training strategies. We apply our approach to several biological contexts and show that latent space optimization in a smooth learned folding landscape allows for more accurate and efficient optimization of protein sequences.

This repo accompanies the following publication:

Castro, Egbert, Abhinav Godavarthi, Julian Rubinfien, Kevin Givechian, Dhananjay Bhaskar, and Smita Krishnaswamy. "Transformer-based protein generation with regularized latent space optimization." Nature Machine Intelligence 4, no. 10 (2022): 840-851.

# clone project   
git clone https://github.com/KrishnaswamyLab/ReLSO-Guided-Generative-Protein-Design-using-Regularized-Transformers.git

see PyTorch Installation page for more details. For convenience, here are some common options

# make conda environment
conda create --name relsoenv python=3.9
conda activate relsoenv

# install pytorch
# GPU (linux)
pip3 install torch torchvision torchaudio

# CPU only (linux)
pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu

# CPU only (mac)
pip3 install torch torchvision torchaudio

python -m pip install networkx pytorch-lightning==1.9 wandb scikit-learn pandas matplotlib gdown phate

pip install -e .   

# GPU training
python train_relso.py  --data gifford

# CPU training
python train_relso.py  --data gifford --cpu

*note: if arg option is not relevant to current model selection, it will not be used. See init method of each model to see what's used.

    gifford, GB1_WU, GFP, TAPE

    base_reg

bash download_weights.sh

which will create a directory called relso_model_weights

❯ tree relso_model_weights -L 1
relso_model_weights
├── model_embeddings
├── model_embeddings.zip
├── trained_models
├── trained_models.json
└── trained_models.zip

2 directories, 3 files

1. Loading GIFFORD Model

python run_optim.py --weights <path to ckpt file>/model_state.ckpt --embeddings  <path to embeddings file>train_embeddings.npy --dataset gifford

• GIFFORD: https://github.com/gifford-lab/antibody-2019/tree/master/data/training%20data
• GB1: https://elifesciences.org/articles/16965#data
• GFP: https://figshare.com/articles/dataset/Local_fitness_landscape_of_the_green_fluorescent_protein/3102154
• TAPE: https://github.com/songlab-cal/tape#data