Adaptive Protein Tokenization

Tokenization is a promising path to multi-modal models capable of jointly understanding protein sequences, structure, and function. Existing protein structure tokenizers create tokens by pooling information from local neighborhoods, an approach that limits their performance on generative and representation tasks. In this work, we present a method for global tokenization of protein structures in which successive tokens contribute increasing levels of detail to a global representation. This change resolves several issues with generative models based on local protein tokenization: it mitigates error accumulation, provides embeddings without sequence-reduction operations, and allows task-specific adaptation of a tokenized sequence's information content. We validate our method on reconstruction, generative, and representation tasks and demonstrate that it matches or outperforms existing models based on local protein structure tokenizers. We show that our adaptive approach enables inference criteria based on the information content of the generated proteins. We validate representations generated from our tokenizer on CATH classification tasks and demonstrate that non-linear probing on our tokenized sequences outperforms equivalent probing on representations from other tokenizers. Finally, we demonstrate how our method supports zero-shot protein shrinking and affinity maturation.

How should computers understand data? For some fields, like natural language, it's easy: just one word at a time. For images, audio, or more exotic data (like proteins in biology), the answer isn't as obvious. The traditional method is to separate the data into pieces and feed them to the model one at a time, just like in language. We tried something different, where each piece we gave to the model provided more and more details. It's akin to first describing the Mona Lisa as a picture of a woman, then describing the colors, then finally describing small details like her smile. We made a model that does the same for biology. We found that this helps a lot -- we can encode data much more efficiently and generate new proteins more effectively. This is important, since one of the most exciting ways AI can impact the world is by improving human health, discovering new drugs, etc.