Tokenized and Continuous Embedding Compressions of Protein Sequence and Structure (CHEAP)

Tuesday, October 1st, 4-5pm EST | Amy Lu, PhD student (UC Berkeley)

Abstract: Existing protein machine learning representations typically model either the sequence or structure distribution, with the other modality implicit. The latent space of sequence-tostructure prediction models such as ESMFold represents the joint distribution of sequence and structure; however, we find these embeddings to exhibit massive activations, whereby some channels have values 3000× higher than others, regardless of the input. Further, on continuous compression schemes, ESMFold embeddings can be reduced by a factor of 128× along the channel and 8× along the length, while retaining structure information at <2Å scale accuracy, and performing competitively on protein function and localization benchmarks. On discrete compression schemes, we construct a tokenized all-atom structure vocabulary that retains high reconstruction accuracy, thus introducing a tokenized representation of all-atom structure that can be obtained from sequence alone. We term this series of embeddings as CHEAP (Compressed Hourglass Embedding Adaptations of Proteins) embeddings, obtained via the HPCT (Hourglass Protein Compression Transformer) architecture. CHEAP is a compact representation of both protein structure and sequence, sheds light on information content asymmetries between sequence and structure, democratizes representations captured by large models, and is designed to have flexible downstream applications such as generation, search, and prediction.

Preprint: https://www.biorxiv.org/content/10.1101/2024.08.06.606920v1

 

Amy is a Computer Science PhD student at UC Berkeley and BAIR advised by Pieter Abbeel, and a part-time researcher at Prescient Design (Genentech). Previously, she was a part-time Student Researcher at Google Brain and Machine Learning Engineer at insitro. She completed her Masters in Computer Science from the University of Toronto and her undergrad in Honours Science from the University of Waterloo. Her research interests include multimodal generative modeling for protein engineering, and broadly bringing task-generalization capabilities from vision-language models to biology.