Adaptive Protein Tokenization

TL;DR

This paper introduces a global, adaptive protein tokenization method that enhances protein structure understanding, improves generative and representation tasks, and supports advanced applications like protein shrinking and affinity maturation.

Contribution

It presents a novel global tokenization approach for protein structures that overcomes limitations of local methods, enabling better generative, representation, and task-specific adaptation.

Findings

Outperforms existing local tokenizers in reconstruction and generative tasks.

Enables inference based on information content, improving designability.

Supports zero-shot protein shrinking and affinity maturation.

Abstract

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 how adaptive tokens enable inference criteria based on information content, which boosts designability. 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.