[Paper Review] FoldToken: Learning Protein Language via Vector Quantization and Beyond

Zhangyang Gao, Cheng Tan+4 more|arXiv (Cornell University)|Feb 4, 2024

Machine Learning in Bioinformatics5 citations

TL;DR

FoldToken learns a discrete protein language by vector quantization of sequence-structure, enabling a GPT-style model FoldGPT for autoregressive sequence-structure generation and applications like backbone inpainting and antibody design. It introduces SoftCVQ to balance reconstruction and generation.

ABSTRACT

Is there a foreign language describing protein sequences and structures simultaneously? Protein structures, represented by continuous 3D points, have long posed a challenge due to the contrasting modeling paradigms of discrete sequences. We introduce extbf{FoldTokenizer} to represent protein sequence-structure as discrete symbols. This innovative approach involves projecting residue types and structures into a discrete space, guided by a reconstruction loss for information preservation. We refer to the learned discrete symbols as extbf{FoldToken}, and the sequence of FoldTokens serves as a new protein language, transforming the protein sequence-structure into a unified modality. We apply the created protein language on general backbone inpainting and antibody design tasks, building the first GPT-style model ( extbf{FoldGPT}) for sequence-structure co-generation with promising results. Key to our success is the substantial enhancement of the vector quantization module, Soft Conditional Vector Quantization ( extbf{SoftCVQ}).

Motivation & Objective

• Bridge protein sequence and structure into a unified discrete language using FoldTokenizer.
• Improve vector quantization to support high-quality reconstruction and generation in a protein language setting.
• Enable GPT-style autoregressive generation of protein sequence-structure with FoldGPT.
• Demonstrate downstream tasks: backbone inpainting and antibody design.
• Provide analysis and ablations to identify factors contributing to effective SoftCVQ performance.

Proposed method

• Encode sequence and structure into a continuous latent representation with a transformer-based FoldTokenizer.
• Quantize latent codes into a discrete protein language using three proposed VQ strategies (SoftVQ, SoftGVQ, SoftCVQ) with soft querying over the full codebook.
• Decode from discrete codes to reconstruct sequence and angle-based structure representations.
• Train FoldGPT as an autoregressive model operating on the learned protein language for inpainting and co-generation of sequence-structure.

Experimental results

Research questions

• RQ1 Can a discrete protein language faithfully represent sequence-structure information with reconstruction loss?
• RQ2 Do soft-vector quantization methods (SoftVQ, SoftGVQ, SoftCVQ) improve reconstruction and generation compared to vanilla and LFQ methods?
• RQ3 How well does FoldGPT perform on backbone inpainting and antibody design using the learned protein language?
• RQ4 What factors (GM, SI, MC, IB, SGQ) drive performance in the VQ methods, and how does SoftCVQ balance reconstruction and generation?

Key findings

• SoftCVQ achieves strong reconstruction and competitive generation, surpassing baseline VQ methods on reconstruction metrics and enabling effective generation.
• SoftVQ, SoftGVQ, and SoftCVQ significantly improve TMScore and reconstruction success rate over Vanilla and LFQ on CATH4.3 test set.
• FoldGPT delivers sequence inpainting results that outperform baselines and shows improvements in structure inpainting over angle-based methods.
• In antibody design, FoldGPT demonstrates competitive performance, with results indicating the value of the learned protein language for design tasks.
• Discretizing with binary VQ-IDs (SoftCVQ) is crucial for convergence and stable training in large codebook spaces.