[论文解读] Learning the Language of Protein Structure

Benoit Gaujac, Jérémie Donà+4 more|arXiv (Cornell University)|May 24, 2024

Biomedical Text Mining and Ontologies被引用 5

一句话总结

引入向量量化自编码器，将蛋白质结构离散化为令牌式编码，从而实现基于序列模型的设计蛋白骨架生成。展示了在学习得到的码本上训练的GPT模型进行de novo结构生成的竞争性效果。

ABSTRACT

Representation learning and \emph{de novo} generation of proteins are pivotal computational biology tasks. Whilst natural language processing (NLP) techniques have proven highly effective for protein sequence modelling, structure modelling presents a complex challenge, primarily due to its continuous and three-dimensional nature. Motivated by this discrepancy, we introduce an approach using a vector-quantized autoencoder that effectively tokenizes protein structures into discrete representations. This method transforms the continuous, complex space of protein structures into a manageable, discrete format with a codebook ranging from 4096 to 64000 tokens, achieving high-fidelity reconstructions with backbone root mean square deviations (RMSD) of approximately 1-5 Å. To demonstrate the efficacy of our learned representations, we show that a simple GPT model trained on our codebooks can generate novel, diverse, and designable protein structures. Our approach not only provides representations of protein structure, but also mitigates the challenges of disparate modal representations and sets a foundation for seamless, multi-modal integration, enhancing the capabilities of computational methods in protein design.

研究动机与目标

• 将三维蛋白质结构转换为离散的序列表示，以实现序列建模。
• 学习一个离散的、可能是低维的潜在空间，具有限定的码本用于蛋白质结构。
• 在提供可扩展的令牌词汇表的同时，实现高重建保真度（RMSD ~1-5 Å）。
• 通过在学习得到的码本上训练GPT模型以生成新颖、可设计的蛋白结构，展示下游应用价值。

提出的方法

• 使用基于3层MPNN的图编码器对骨架坐标(N-Cα-C-O)进行编码，生成下采样残基的嵌入。
• 使用有限标量量化(FSQ)将潜在嵌入离散化，形成码本（码数高达64k）。
• 从量化潜在码解码回3D骨架坐标，使用受AlphaFold帧框架启发的结构模块，并用Frame Align Point Error (FAPE)损失进行优化。
• 在约31万条PDB条目上进行训练，采用类增强的下采样和旋转不变特征以确保帧对齐与鲁棒重建。
• 通过RMSD和TM-score评估重建，并通过对令牌化结构训练解码器为主的GPT来展示生成能力。
• 提供公开代码并将生成质量与专门的扩散模型(FrameDiff、RFDiffusion)进行对比。

实验结果

研究问题

• RQ1 蛋白质结构是否能够在不发生灾难性信息丢失的情况下，有效离散化为有限的令牌词汇？
• RQ2 离散潜在空间是否能够有效地让标准序列模型（如GPT）用于de novo蛋白质结构生成？
• RQ3 在向量量化自编码器中改变码本大小和下采样时，重建保真度与可设计性之间的权衡是什么？
• RQ4 生成的结构在可设计性、新颖性和多样性方面与参考数据相比如何？

主要发现

• 在无下采样（r=1）下，64k码本在测试结构上实现约1.59 Å RMSD和TM-score 0.95。
• 量化不会使码本坍缩；更大的码本在提高重建的同时保持具有竞争力的保真度。
• 下采样和较小的码本会增加重建误差，但平均TM-score仍保持在0.5以上，表明全局折叠保持良好。
• 一个简单的仅解码器GPT，在令牌化结构上训练，可以生成新颖、可设计的骨架，并与扩散基线相比具有竞争力的自一致性指标。
• 生成的结构表现出非平凡的二级结构元素，并与ESM预测的结构在设计性评估中对齐良好。
• 与FrameDiff和RFDiffusion相比，基于GPT的方法在可设计性和新颖性/多样性方面具有竞争力的特征，取决于采样策略存在权衡。