ReLSO: A Transformer-based Model for Latent Space Optimization and Generation of Proteins

• github repo

Model

ReLSO is a regularized VAE protein sequence model, trained to produce a smooth and structured fitness landscape. Gradient ascent can be performed using the fitness prediction of ReLSO for efficient travesal of latent space. Additionally, ReLSO is a guided representation learner of protein sequences focused on a dense local sampling of protein space.

Example usage

1. clone and install ReLSO

2. specify and load checkpoint

import json
from relso.nn.models import relso1

# available datasets: 'gifford', 'gb1', 'gfp'
# model settings: 'all', 'neg', 'interp', 'vanilla', 'alpha0', 'alpha_01', 'alpha_05', 'cnn'

# Load checkpoint metadata
with open("trained_models.json") as f:
    weights_resolver = json.load(f)

dataset = 'gifford'
msetting = 'all'
checkpoint_id = weights_resolver[dataset][msetting]
weight_path = f"trained_models/{checkpoint_id}_model_weights.ckpt"

relso_model = relso1.load_from_checkpoint(weight_path)

Model Details

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.

Citation

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.