[论文解读] FoldToken: Learning Protein Language via Vector Quantization and Beyond

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

Machine Learning in Bioinformatics被引用 5

一句话总结

FoldToken 通过对序列-结构进行向量量化，学习一个离散蛋白质语言，进而产生类似 GPT 的模型 FoldGPT，用于自回归序列-结构生成以及骨架修复（backbone inpainting）和抗体设计等应用。它引入 SoftCVQ 以平衡重建与生成。

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}).

研究动机与目标

• 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.

提出的方法

• 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.

实验结果

研究问题

• 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?

主要发现

• 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.