Structure-aware protein encoders pre-trained with an ELECTRA-style corrective MLM objective on protein structures, plus a per-residue representation benchmark covering apo-conditioned co-folding, REPA-supervised folding, and frozen probing.
- Paper: Atom-level Protein Representation Learning Improves Protein Structure Prediction (arXiv:2605.22133)
- Project page: https://holymollyhao.github.io/TriProRep/
- Models + benchmark:
k-fold-structureon HuggingFace
Recent advances in generative modeling show that pretrained representations can improve generation as conditioning features or alignment targets. Motivated by this, we study protein representations for predicting structures beyond conventional function annotation. We propose TriProRep, a structure-aware pretraining method that jointly models three aligned residue-level views: amino-acid identity, backbone geometry, and local full-atom geometry, discretely encoded via VQ-VAE tokenizers. By pretraining to recover original tokens from generator-corrupted views, TriProRep learns to distinguish plausible but incorrect cross-view augmentations from the original protein. We further introduce RepSP, a benchmark for evaluating protein representations in structure-predictive settings. RepSP tests three uses of representations: homodimer co-folding from apo-chain representations, residue-level prediction of homodimer-derived interaction properties, and representation-aligned monomer structure prediction. Across these tasks, TriProRep improves over sequence-only and prior structure-aware representation models, while maintaining competitive performance on conventional benchmarks.
pip install torch huggingface_hub omegaconf numpy lmdb biotite einops
git clone https://github.com/hsjang0/TriProRep.git
cd TriProRep
python examples/quickstart_demo.pyThe demo loads the 650M encoder from HuggingFace and embeds the bundled
example PDB (_assets/example_monomer.pdb), printing an embedding shape
of (119, 1280) in fp16. First run downloads the encoder + tokenizers
(about 3 GB, cached under HF_HOME).
Under the hood:
import sys; sys.path.insert(0, "code/triprorep")
from inference import load_encoder, embed_pdb
encoder = load_encoder("650M", hf_repo="k-fold-structure/triprorep-650M")
features = embed_pdb(encoder, "_assets/example_monomer.pdb")
print(features.shape) # (119, 1280) fp16For CPU inference, pass device="cpu" to load_encoder. If you already
have (seq, bb, fa) token IDs, call encode(encoder, seq, bb, fa) directly.
Four common entry points. Pick the one that matches your goal.
| Goal | Section |
|---|---|
| Embed your PDBs with our encoder | Embed your own PDBs |
| Get our seq/bb/fa tokens for your PDBs | Tokenize your own PDBs |
| Run your encoder on our benchmark | Score your encoder on the benchmark |
| Reproduce our benchmark numbers | Reproduce the benchmark |
One protein:
import sys; sys.path.insert(0, "code/triprorep")
from inference import load_encoder, embed_pdb
encoder = load_encoder("650M", hf_repo="k-fold-structure/triprorep-650M")
features = embed_pdb(encoder, "path/to/your.pdb") # (L, 1280) fp16Whole directory into one LMDB:
python examples/extract_features_from_pdbs.py \
--pdbs_dir ./my_pdbs --pdb_glob "*.pdb" \
--model 650M --output ./features.lmdbLMDB key = PDB stem, value = pickle.dumps(np.ndarray[L, D], dtype=fp16).
Chain A only; if your PDB has multiple chains, pass chain="X" to embed_pdb.
If you just want the three discrete token streams (and not the encoder
forward), use tokenize_pdb. The tokens are arbitrary downstream input
for your own classifier, a generation model, a custom probe.
import sys; sys.path.insert(0, "code/triprorep")
from inference import tokenize_pdb
tok = tokenize_pdb("path/to/your.pdb",
hf_repo="k-fold-structure/triprorep-650M")
tok["seq"] # (L,) int, ESM2-style AA token IDs
tok["bb"] # (L,) int, backbone codebook IDs (vocab = 512)
tok["fa"] # (L,) int, full-atom codebook IDs (vocab = 512)The two tokenizers (backbone_tokenizer.pt, fullatom_tokenizer.pt) are
downloaded from the model repo on first call and cached. The encoder
itself is not loaded. CPU works fine and the GPU footprint stays small.
# (a) Pull the small probing subset (splits + labels + ~2 GB of monomer PDBs).
bash examples/setup_probing.sh
# (b) Build features.lmdb with your encoder (schema: chain-A [L, D] fp16
# per <af-id>, keyed by lowercased AFid. See "Bring your own encoder"
# below for the exact contract).
python my_extract.py \
--pdbs ./REPSP_PDB/monomer \
--splits ./benchmark/splits/probing \
--out ./work/features_theirs.lmdb
# (c) Stage per-split .pt and run the four probing tasks.
LABELS=./benchmark/probing/labels.pkl
python code/repsp/probing/homomer/__lib/extract_probing_features.py \
--features_lmdb ./work/features_theirs.lmdb \
--splits_dir ./benchmark/splits/probing \
--target_pkl $LABELS \
--out_dir ./work/probing_features_theirs
for TASK in binding_site delta_sasa_mean levy_tier bond_type_plip; do
python code/repsp/probing/homomer/__lib/probe_residue_homomer.py \
--features_dir ./work/probing_features_theirs \
--target_pkl $LABELS \
--task $TASK \
--run_name theirs_$TASK
donebash examples/setup_probing.sh # about 2 GB
bash examples/run_probing.sh # extract features + 4 tasks with the 650M encoderOutputs land under ./work/results_650M/<task>.json for the four tasks
(binding_site, delta_sasa_mean, levy_tier, bond_type_plip).
Set MODEL_SIZE=35M | 150M | 650M | 3B to pick a different encoder,
e.g. MODEL_SIZE=3B bash examples/run_probing.sh.
| Model | Params | Hidden dim | Encoder layers | Heads | HuggingFace |
|---|---|---|---|---|---|
triprorep-35M |
35M | 480 | 10 | 20 | k-fold-structure/triprorep-35M |
triprorep-150M |
150M | 640 | 28 | 20 | k-fold-structure/triprorep-150M |
triprorep-650M |
650M | 1280 | 30 | 20 | k-fold-structure/triprorep-650M |
triprorep-3B |
3B | 2560 | 33 | 40 | k-fold-structure/triprorep-3B |
Each model repo ships the full Lightning checkpoint, the sanitized
config.yaml, and the two structure tokenizers
(backbone_tokenizer.pt, fullatom_tokenizer.pt). The 3B repo additionally
ships 3B_encoder.pt, an encoder-only state dict for inference.
All data lives under the k-fold-structure
HuggingFace org.
# Benchmark inputs: splits, probing labels, Boltz-tokenized apo + holo
hf download k-fold-structure/repsp-benchmark --repo-type dataset --local-dir ./benchmark
cd ./benchmark/boltz_holo_tokens && for t in shard*.tar; do tar xf "$t"; done && rm shard*.tar
# Per-protein structure tokens (apo + holo), keyed by AF-{id}
hf download k-fold-structure/repsp-triprorep-tokens --repo-type dataset --local-dir ./tokens.lmdb
cd ./tokens.lmdb && cat data.mdb.part_* > data.mdb && rm data.mdb.part_*
# Pre-training corpus (654 GB, only if re-training the encoder)
hf download k-fold-structure/triprorep-pretrain --repo-type dataset --local-dir ./pretrain.lmdb
cd ./pretrain.lmdb && cat data.mdb.part_* > data.mdb && rm data.mdb.part_*Sharded tar.gz under k-fold-structure/repsp-benchmark/REPSP_PDB/.
Untarring reproduces this layout:
REPSP_PDB/
├── monomer/<AF-id>_monomer.pdb # chain A (apo prediction)
└── homodimer/<AF-id>.pdb # chain A + chain B (holo dimer)
The AFid in both filenames is the same homodimer identifier taken from
splits/{folding,probing}/{train,valid,test}.txt. Sizes:
| Shard | Records | Compressed |
|---|---|---|
monomer/valid.tar.gz |
400 | 18 MB |
monomer/test.tar.gz |
1,000 | 45 MB |
monomer/train_000.tar.gz |
390,861 | 18 GB |
homodimer/valid.tar.gz |
400 | 34 MB |
homodimer/test.tar.gz |
1,000 | 85 MB |
homodimer/train_000.tar.gz |
390,861 | 34 GB |
For probing only, monomer/test.tar.gz alone is enough (about 130 MB).
For folding / co-folding training, grab the train shards as well.
Split identity is preserved as-committed: folding 390,627 / 400 / 1,000, probing 39,100 / 400 / 1,000. The splits are LMDB-cleaned (a small number of AFids that fail Boltz tokenization are already dropped).
The PDBs come from the AFDB-Multimer homodimer collection (predicted with
AlphaFold-Multimer / ColabFold; see
collaborations/nvda
on the EBI FTP) and the corresponding apo monomers we generated with
AlphaFold-2. Both are AlphaFold Protein Structure Database content and are
redistributed here under the same CC BY 4.0
terms as the upstream AlphaFold DB, with attribution to DeepMind and
EMBL-EBI (see Acknowledgements). If you prefer, you
can also rebuild REPSP_PDB/ from a local checkout of the Boltz-format
structures our folding pipeline writes with scripts/npz_to_pdb.py.
The boltz_* dirs use the SimpleFold on-disk format
(manifest.json + records/ + tokens/ or structures/). The
workflow commands point at the extracted folders.
Per-encoder feature LMDBs and probing per-split tensors are not shipped. Both are produced locally in minutes from the structure tokens and an encoder (see Bring your own encoder).
| Workflow | Reads | Build first |
|---|---|---|
| Pre-training | triprorep-pretrain, a model repo for config.yaml |
n/a |
| Folding (REPA) | repsp-benchmark (boltz_apo_* + splits/folding), a model |
features.lmdb |
| Co-folding | repsp-benchmark (boltz_holo_* + splits/folding), a model |
features.lmdb |
| Probing | repsp-benchmark (probing/labels.pkl + splits/probing), repsp-triprorep-tokens, a model |
per-split .pt |
Train the ELECTRA encoder on the provided tokenized LMDB.
cd code/triprorep
# Edit lmdb_dir: in configs/pretrain_650M/pretrain_electra.yaml to point at ./pretrain.lmdb
torchrun --nproc_per_node=8 experiments/train_multinode.py \
--config configs/pretrain_650M/pretrain_electra.yaml
# Sizes: pretrain_{35M, 150M, 650M, 3B}/pretrain_electra.yaml
# Multi-node template: scripts/pretrain/650M.shThree complementary evaluations of the encoder representation on protein
structure tasks: folding (REPA-supervised apo monomer prediction),
co-folding (apo-conditioned homodimer prediction), and probing
(frozen per-residue prediction of four homodimer-derived interaction
properties). All three use the same shipped
k-fold-structure/repsp-benchmark
assets: splits/, probing/labels.pkl, boltz_{apo,holo}_{tokens,targets},
and REPSP_PDB/.
Apo single-chain folding with REPA supervision: a per-token cosine alignment
loss pulls the trunk's mid-block hidden state toward the encoder's frozen
features. Set up the folding env once with
bash code/repsp/folding/scripts/setup_env.sh.
bash examples/setup_folding.sh # ~40 GB: monomer PDBs + Boltz apoBuild features.lmdb (per-encoder features consumed as the REPA target):
python examples/extract_features_from_pdbs.py \
--pdbs_dir ./REPSP_PDB/monomer --pdb_glob "*.pdb" \
--model 650M --output ./features.lmdbLaunch training:
cd code/repsp/folding
APO_TOKENS=./benchmark/boltz_apo_tokens
APO_TARGETS=./benchmark/boltz_apo_targets
SPLIT=./benchmark/splits/folding
FEATURES=../../features.lmdb
REPA_DIM=1280 # = features.lmdb __metadata__.output_dim
REPA_W=2.0
python src/simplefold/train.py experiment=folding_v1_full \
data.feature_paths.repa_target_s=$FEATURES \
model.repa_target_dim=$REPA_DIM \
model.repa_weight=$REPA_W \
++data.datasets.0.tokenized_dir=$APO_TOKENS \
++data.datasets.0.target_dir=$APO_TARGETS \
++data.datasets.0.manifest_path=$APO_TOKENS/manifest.json \
++data.datasets.0.record_list=$SPLIT/train.txt \
++data.datasets.0.val_record_list=$SPLIT/valid.txt \
++data.datasets.0.test_record_list=$SPLIT/test.txtREPA_DIM matches the encoder hidden dim: 35M→480, 150M→640, 650M→1280, 3B→2560.
For the no-REPA baseline, drop data.feature_paths.repa_target_s and set
model.repa_target_dim=0 model.repa_weight=0.
Apo conditioning of a SimpleFold
trunk folds the holo dimer. The folding code under code/repsp/folding/ is
adapted from SimpleFold and Boltz
(both MIT, see code/repsp/folding/LICENSE).
bash examples/setup_cofolding.sh # ~150 GB: monomer PDBs + Boltz holoBuild features.lmdb (per-encoder features consumed as the apo conditioning input):
python examples/extract_features_from_pdbs.py \
--pdbs_dir ./REPSP_PDB/monomer --pdb_glob "*.pdb" \
--model 650M --output ./features.lmdbcd code/repsp/folding
HOLO_TOKENS=./benchmark/boltz_holo_tokens
HOLO_TARGETS=./benchmark/boltz_holo_targets
SPLIT=./benchmark/splits/folding
FEATURES=./features/features.lmdb
APO_DIM=1280 # = features.lmdb __metadata__.output_dim
python src/simplefold/train.py experiment=cofolding_v1_full \
data.feature_paths.apo_s=$FEATURES \
model.architecture.apo_repr_dim=$APO_DIM \
++data.datasets.0.tokenized_dir=$HOLO_TOKENS \
++data.datasets.0.target_dir=$HOLO_TARGETS \
++data.datasets.0.manifest_path=$HOLO_TOKENS/manifest.json \
++data.datasets.0.record_list=$SPLIT/train.txt \
++data.datasets.0.val_record_list=$SPLIT/valid.txt \
++data.datasets.0.test_record_list=$SPLIT/test.txtSame train.txt / valid.txt / test.txt lists are used for both
folding and co-folding. The AFid in each line resolves to two on-disk
files: REPSP_PDB/monomer/<AF-id>_monomer.pdb for the apo side and
REPSP_PDB/homodimer/<AF-id>.pdb for the dimer side.
Frozen-representation, per-residue probing on four homodimer tasks:
binding_site, delta_sasa_mean, levy_tier, bond_type_plip.
For the full benchmark in one command see Reproduce the benchmark.
The block below is the underlying loop, useful when plugging a different
encoder. Build per-split .pt features first
(see Bring your own encoder).
cd code/repsp/probing/homomer
FEATURES_DIR=./probing_features # contains train.pt / valid.pt / test.pt
LABELS=./benchmark/probing/labels.pkl
for TASK in binding_site delta_sasa_mean levy_tier bond_type_plip; do
python __lib/probe_residue_homomer.py \
--features_dir $FEATURES_DIR \
--target_pkl $LABELS \
--task $TASK \
--run_name ours_$TASK
doneDefaults reproduce the benchmark (10 epochs, residue batch 16,824, lr 5e-4).
Setup details are in code/repsp/probing/homomer/PROBING_SETTING.md.
Per-encoder features are built locally in two flavors. Pick the one that matches your encoder's input:
- From our model: feed structure tokens
(
repsp-triprorep-tokens, reassembled to./tokens.lmdb) through our released encoder withinference.load_encoder+encode. - From raw PDBs and other model: bring any non-ours encoder
(ESM-2, SaProt, MIF-ST, ...) that reads raw apo monomer PDBs. The AFDB
monomers are the same ones that ship as
k-fold-structure/repsp-benchmark/REPSP_PDB/monomer/*.tar.gz; untar and pointPDB_DIRat the resulting directory. Replace theyour_encoder_forward(...)body.
Both variants produce the same on-disk schema, so the workflow commands above do not change.
From our model
import sys, pickle
import lmdb
sys.path.insert(0, "code/triprorep")
from inference import load_encoder, encode
APO_LMDB = "./tokens.lmdb"
OUT_LMDB = "./features/features.lmdb"
encoder = load_encoder("650M", hf_repo="k-fold-structure/triprorep-650M")
# Local ckpt alternative:
# encoder = load_encoder("650M", ckpt="./models/650M/650M.ckpt")
apo_env = lmdb.open(APO_LMDB, readonly=True, lock=False, readahead=False)
out_env = lmdb.open(OUT_LMDB, map_size=int(1e12))
n, D = 0, None
with out_env.begin(write=True) as out_txn, apo_env.begin() as apo_txn:
for key, raw in apo_txn.cursor():
if key in (b"__metadata__", b"__keys__"):
continue
rec = pickle.loads(raw)
feats = encode(encoder, rec["apo_seq_A"], rec["apo_bb_A"], rec["apo_fa_A"])
D = feats.shape[1]
out_txn.put(key, pickle.dumps(feats))
n += 1
out_txn.put(b"__metadata__", pickle.dumps({
"n_samples": n, "output_dim": D, "encoder": "ours_650M", "axis": "chain_A_only",
}))
apo_env.close(); out_env.close()From raw PDBs and other model
import pickle
from pathlib import Path
import lmdb
import numpy as np
PDB_DIR = Path("./REPSP_PDB/monomer") # from repsp-benchmark
SPLITS_DIR = Path("./benchmark/splits/folding") # <split>.txt (monomer IDs)
OUT_LMDB = "./features/features.lmdb"
ENCODER = "your_encoder_v1"
D = 1024
# model = YourEncoder.from_pretrained(...).eval().to("cuda")
def pid_to_pdb(pid: str) -> Path:
return PDB_DIR / f"{pid}.pdb" # adjust to your AFDB layout
def your_encoder_forward(pdb_path: Path) -> np.ndarray:
"""Return chain-A features [L_A, D] fp16 from one apo PDB."""
# feats = model.embed(str(pdb_path)).cpu().numpy()
feats = ...
return feats.astype(np.float16)
pids = []
for s in ("train", "valid", "test"):
pids += [l.strip() for l in (SPLITS_DIR / f"{s}.txt").read_text().splitlines() if l.strip()]
out_env = lmdb.open(OUT_LMDB, map_size=int(1e12))
with out_env.begin(write=True) as out_txn:
for pid in pids:
feats = your_encoder_forward(pid_to_pdb(pid))
out_txn.put(pid.lower().encode(), pickle.dumps(feats))
out_txn.put(b"__metadata__", pickle.dumps({
"n_samples": len(pids), "output_dim": D,
"encoder": ENCODER, "axis": "chain_A_only",
}))
out_env.close()The co-folding loader tiles A∥A inline at load. Folding reads [L_A, D]
as-is.
The .pt schema:
{
"X": torch.float16 [N_res, D], # all chain-A residues, stacked
"pid_slices": list of (pid, start, end), # X[start:end] = residues of `pid`
}From our model
import sys, pickle
from pathlib import Path
import lmdb
import numpy as np
import torch
sys.path.insert(0, "code/triprorep")
from inference import load_encoder, encode
APO_LMDB = "./tokens.lmdb"
SPLITS_DIR = Path("./benchmark/splits/probing") # {train,valid,test}.txt
OUT_DIR = Path("./probing_features"); OUT_DIR.mkdir(parents=True, exist_ok=True)
encoder = load_encoder("650M", hf_repo="k-fold-structure/triprorep-650M")
apo_env = lmdb.open(APO_LMDB, readonly=True, lock=False, readahead=False)
with apo_env.begin() as apo_txn:
for split in ("train", "valid", "test"):
pids = [l.strip() for l in
(SPLITS_DIR / f"{split}.txt").read_text().splitlines() if l.strip()]
per_protein = []
for pid in pids:
rec = pickle.loads(apo_txn.get(pid.lower().encode()))
per_protein.append(
encode(encoder, rec["apo_seq_A"], rec["apo_bb_A"], rec["apo_fa_A"])
)
X = torch.from_numpy(np.concatenate(per_protein, axis=0))
pid_slices, offset = [], 0
for pid, feats in zip(pids, per_protein):
pid_slices.append((pid, offset, offset + feats.shape[0]))
offset += feats.shape[0]
torch.save({"X": X, "pid_slices": pid_slices}, OUT_DIR / f"{split}.pt")
print(f"{split}: {len(pid_slices)} proteins, N_res={X.shape[0]}, D={X.shape[1]}")
apo_env.close()From raw PDBs and other model
from pathlib import Path
import numpy as np
import torch
PDB_DIR = Path("./REPSP_PDB/monomer") # from repsp-benchmark
SPLITS_DIR = Path("./benchmark/splits/probing")
OUT_DIR = Path("./probing_features"); OUT_DIR.mkdir(parents=True, exist_ok=True)
# model = YourEncoder.from_pretrained(...).eval().to("cuda")
def pid_to_pdb(pid: str) -> Path:
return PDB_DIR / f"{pid}.pdb"
def your_encoder_forward(pdb_path: Path) -> np.ndarray:
"""Return chain-A features [L_A, D] fp16 from one apo PDB."""
# feats = model.embed(str(pdb_path)).cpu().numpy()
feats = ...
return feats.astype(np.float16)
for split in ("train", "valid", "test"):
pids = [l.strip() for l in
(SPLITS_DIR / f"{split}.txt").read_text().splitlines()
if l.strip()]
per_protein = [your_encoder_forward(pid_to_pdb(pid)) for pid in pids]
X = torch.from_numpy(np.concatenate(per_protein, axis=0))
pid_slices, offset = [], 0
for pid, feats in zip(pids, per_protein):
pid_slices.append((pid, offset, offset + feats.shape[0]))
offset += feats.shape[0]
torch.save({"X": X, "pid_slices": pid_slices}, OUT_DIR / f"{split}.pt")A working reference is code/repsp/probing/homomer/__lib/extract_probing_features.py.
- Backbone tokenizer (
backbone_tokenizer.pt): the aminoaseed VQ-VAE from StructTokenBench. - Folding trunk + Boltz data pipeline: code under
code/repsp/folding/is adapted from SimpleFold (Apple, MIT) and Boltz (Wohlwend et al., MIT). Seecode/repsp/folding/docs/ATTRIBUTION.mdandcode/repsp/folding/LICENSE. - Structures under
REPSP_PDB/: the homodimer PDBs are AFDB-Multimer predictions provided by NVIDIA to the AlphaFold Protein Structure Database (Jumper et al., 2021; Varadi et al., 2022, 2024); the apo monomers are AlphaFold-2 single-chain predictions we generated. Both are redistributed under CC BY 4.0 with attribution to DeepMind and EMBL-EBI, per the AFDB terms of use.
@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}
}MIT. See LICENSE.


