ProSST

What ProSST does

Most protein language models see only amino-acid sequences. Structure is often injected later — as continuous coordinates, graphs, or hand-crafted features — which can be awkward to scale and hard to fuse with sequence Transformers. ProSST instead turns structure into a parallel discrete token stream and pretrains a Masked Language Model that attends over both sequences jointly.

The result is a family of Hugging Face checkpoints (ProSST-*) that support zero-shot mutant effect prediction and plug into engineering workflows such as VenusFactory2. An MSA-aware follow-up, VenusREM, builds on the same structure-aware foundation and leads ProteinGym’s Substitution track.

Structure quantization

From a PDB (or AlphaFold) structure, a structure quantizer maps local geometric context into discrete structure tokens — analogous to a vocabulary over structural “words.” Vocabularies of size 20 / 128 / 512 / 1024 / 2048 / 4096 are released; each residue gets one structure token aligned with its amino acid.

Quantization keeps structure in the same discrete modeling regime as sequence tokens, so standard Transformer pretraining applies without continuous coordinate heads. In practice, ProSST-2048 is the recommended default for ProteinGym-style zero-shot scoring.

from prosst.structure.get_sst_seq import SSTPredictor

predictor = SSTPredictor(structure_vocab_size=2048)  # 20, 128, 512, 1024, 2048, 4096
result = predictor.predict_from_pdb("example.pdb")
# → [407, 998, 1841, 1421, ...]

Disentangled attention

Amino acids and structure tokens carry related but distinct signals. ProSST uses disentangled attention so the model can separate content interactions along the sequence modality and the structure modality, rather than collapsing everything into a single mixed token stream.

Intuitively: sequence attention captures evolutionary / chemical patterns in the AA string, while structure attention routes geometric neighborhood information through the quantized tokens. Cross-talk still happens — the Masked LM must predict masked residues given both views — but the attention factorization keeps the two sources of evidence easier to learn and ablate.

Zero-shot mutant scoring

Load a checkpoint with Hugging Face Transformers, tokenize the wild-type sequence together with its structure tokens, and score mutants from log-likelihood ratios (see the score_mutant notebook and ProteinGym benchmark script in the repo).

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
)

Citation

If you use ProSST, please cite:

@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}
}