triprorep-3B

triprorep-3B

Structure-aware protein encoder, 3B parameters. ELECTRA-style corrective MLM pre-training on protein structures. The encoder reads three per-residue token streams (seq / bb / fa) and outputs a per-residue embedding of dimension 2560 (fp16).

Architecture: embed_dim=2560, encoder_depth=33, encoder_heads=40.

Part of the TriProRep release. Paper: Atom-level Protein Representation Learning Improves Protein Structure Prediction (arXiv:2605.22133).

Files

• 3B.ckpt: full Lightning checkpoint.
• 3B_encoder.pt: encoder-only state dict, loads strict=True for inference.
• config.yaml: model + data config. The path fields are placeholders, point them at your local data.
• backbone_tokenizer.pt, fullatom_tokenizer.pt: structure tokenizers (PDB to token IDs). See Acknowledgements.

Usage

pip install torch huggingface_hub omegaconf numpy lmdb biotite
git clone https://github.com/hsjang0/TriProRep.git
cd TriProRep

import sys; sys.path.insert(0, "code/triprorep")
from inference import load_encoder, embed_pdb

encoder  = load_encoder("3B", hf_repo="k-fold-structure/triprorep-3B")
features = embed_pdb(encoder, "your_protein.pdb",
                     hf_repo="k-fold-structure/triprorep-3B")
print(features.shape)   # (L, 2560) fp16

embed_pdb downloads the bundled tokenizers from this repo on first call, then runs PDB to (seq, bb, fa) tokens to encoder. If you already have token IDs (e.g. from k-fold-structure/repsp-triprorep-tokens), call encode(encoder, seq, bb, fa) directly. For CPU, pass device="cpu" to load_encoder.

Acknowledgements

backbone_tokenizer.pt (aminoaseed VQ-VAE) is from StructTokenBench.

Citation

@misc{triprorep,
  title  = {Atom-level Protein Representation Learning Improves Protein Structure Prediction},
  author = {Kim, Taewon and Jang, Hyosoon and Seo, Hyunjin and Seo, Seonghwan and Kim, Hyeongwoo and Zhung, Wonho and Shin, Mingyeong and Kim, Wooyoun and Ahn, Sungsoo},
  year   = {2026},
  eprint = {2605.22133},
  archivePrefix = {arXiv},
  primaryClass = {cs.LG},
  url    = {https://arxiv.org/abs/2605.22133}
}

License

MIT