ProtWave-VAE: Integrating autoregressive sampling with latent-based inference for data-driven protein design

Deep generative models (DGMs) have shown great success in the understanding of data-driven design of proteins. Variational autoencoders (VAEs) are a popular DGM approach that can learn the correlated patterns of amino acid mutations within a multiple sequence alignment (MSA) of protein sequences and distill this information into a low-dimensional latent space to expose phylogenetic and functional relationships and guide generative protein design. Autoregressive (AR) models are another popular DGM approach that typically lack a low-dimensional latent embedding but do not require training sequences to be aligned into an MSA and enable the design of variable length proteins. In this work, we propose ProtWave-VAE as a novel and lightweight DGM employing an information maximizing VAE with a dilated convolution encoder and autoregressive WaveNet decoder. This architecture blends the strengths of the VAE and AR paradigms in enabling training over unaligned sequence data and the conditional generative design of variable length sequences from an interpretable low-dimensional learned latent space. We evaluate the models ability to infer patterns and design rules within alignment-free homologous protein family sequences and to design novel synthetic proteins in four diverse protein families. We show that our model can infer meaningful functional and phylogenetic embeddings within latent spaces and make highly accurate predictions within semi-supervised downstream fitness prediction tasks. In an application to the C-terminal SH3 domain in the Sho1 transmembrane osmosensing receptor in bakers yeast, we subject ProtWave-VAE designed sequences to experimental gene synthesis and select-seq assays for osmosensing function to show that the model enables de novo generative design, conditional C-terminus diversification, and engineering of osmosensing function into SH3 paralogs.