NeurIPS 2024

ProSST

带量化结构与解耦注意力的蛋白质语言建模——面向表征学习与零样本突变评分的序列–结构预训练 Transformer。

Mingchen Li^*, Yang Tan^*, Xinzhu Ma, Bozitao Zhong, Huiqun Yu, Ziyi Zhou, Wanli Ouyang, Bingxin Zhou, Pan Tan, Liang Hong
（* 共同一作）

论文代码模型排行榜

ProSST：结构量化与氨基酸–结构 token 解耦注意力 — ProSST 在 Masked LM 目标下，将离散结构 token 与氨基酸序列联合建模，并采用解耦注意力。

方法

多数蛋白质语言模型只看见氨基酸序列；结构常在后期以连续坐标、图或手工特征注入，既难扩展也不易与序列 Transformer 融合。 ProSST 将结构转为并行的离散 token 流，并预训练同时关注两条序列的 Masked Language Model。

由 PDB（或 AlphaFold）结构出发，结构量化器把局部几何映射为离散结构 token——相当于结构“词表”。已发布词表规模 20 / 128 / 512 / 1024 / 2048 / 4096；每个残基对应一个与氨基酸对齐的结构 token。实践中推荐 ProSST-2048 作为 ProteinGym 风格零样本评分的默认检查点。

解耦注意力与零样本评分

氨基酸与结构 token 携带相关但不同的信号。ProSST 使用解耦注意力，分别沿序列模态与结构模态建模内容交互，而非混入单一 token 流。加载 Hugging Face 检查点，将野生型序列与结构 token 一并分词，即可用对数似然比进行突变评分。

from transformers import AutoModelForMaskedLM, AutoTokenizer

model = AutoModelForMaskedLM.from_pretrained(
    "AI4Protein/ProSST-2048", trust_remote_code=True
)
tokenizer = AutoTokenizer.from_pretrained(
    "AI4Protein/ProSST-2048", trust_remote_code=True
)

检查点系列见 AI4Protein/ProSST-*。后续工作 VenusREM 在同一结构感知基础上引入 MSA 检索，并领跑 ProteinGym Substitution；亦可在 VenusFactory2 中直接调用。

引用

若使用 ProSST，请引用：

@inproceedings{li2024prosst,
  title={ProSST: Protein Language Modeling with Quantized Structure and Disentangled Attention},
  author={Li, Mingchen and Tan, Yang and Ma, Xinzhu and Zhong, Bozitao and Yu, Huiqun and Zhou, Ziyi and Ouyang, Wanli and Zhou, Bingxin and Tan, Pan and Hong, Liang},
  booktitle={Advances in Neural Information Processing Systems},
  year={2024}
}

谭扬 (Yang Tan)

方法

解耦注意力与零样本评分

引用