ProTokens: Probabilistic Vocabulary for Compact and Informative Encodings of All-Atom Protein Structures

Understanding functions of proteins and designing proteins committed to specific functions in silico are highly valuable for science, industry and therapeutics. However, there is a long-standing divergence in how to present function-related protein structures to the machine learning models: Although the 1-dimensional (1D) representation of proteins via Anfinsens tokens (i.e., amino acids) is sufficient in principle and more machine-friendly, it is less successful in structure-oriented protein design compared to symmetry-constrained 3D representation (i.e., atom coordinates). Aiming to bridge the gap between 1D and 3D protein representations and harvest the advantages of the two, we develop probabilistic tokenization theory for metastable protein structures. We present an unsupervised learning strategy, which conjugates inverse folding with structure prediction, to encode protein structures into artificial amino-acid tokens (ProTokens) and decode them back to atom coordinates. We show that tokenizing protein structures variationally can lead to compact and informative representations. Compared to amino acids -- the Anfinsens tokens -- ProTokens are easier to detokenize and more descriptive of finer conformational ensembles. Therefore, protein structures can be efficiently compressed, stored, aligned and compared in the form of ProTokens. By unifying the discrete and continuous representations of protein structures, ProTokens also enable all-atom protein structure design via various generative models without the concern of symmetry or modality mismatch, and allows scalable foundation models to perceive, process and explore the microscopic structures of biomolecules effectively.