ProtMMLM is a multimodal protein learning framework for protein representation learning and downstream prediction. It jointly uses protein sequence information, native-structure information, and molecular-dynamics-derived inputs to learn structure- and dynamics-aware representations.
Create a Python environment and install the required dependencies,Please ensure you use a newer Python version (e.g., Python 3.12 or above) to properly obtain sequence embeddings:
Additional runtime dependencies may be required for specific workflows:
- ESMC runtime and weights for sequence encoding;
- MMseqs2 for sequence-clustered splits and Conotoxin negative-set filtering;
- CUDA-enabled PyTorch for GPU training and inference;
- external pretraining checkpoints for downstream initialization and probing.
Run all commands from the repository root:
export PYTHONPATH=.You can view the full list of command-line arguments for any script with the --help flag. For example:
PYTHONPATH=. python scripts/run_pretrain.py --help.
├── configs/
│ ├── pretrain/
│ │ └── pretrain.yaml
│ └── downstream/
│ ├── toxteller.yaml
│ ├── prmftp.yaml
│ ├── ppikb.yaml
│ └── conotoxin.yaml
├── data/
│ └── examples/
│ ├── pretrain_manifest.example.csv
│ └── downstream_manifest.example.csv
├── scripts/
│ ├── prepare_pretrain_manifest.py
│ ├── prepare_downstream_manifests.py
│ ├── prepare_structure_token_cache.py
│ ├── filter_conotoxin_negative_set.py
│ ├── run_pretrain.py
│ ├── run_downstream_finetuning.py
│ ├── probe_global_structural_properties.py
│ ├── probe_total_energy.py
│ ├── probe_energy_terms.py
│ └── analyze_rmsd_embedding_consistency.py
├── src/
│ ├── analysis/
│ ├── datasets/
│ ├── evaluation/
│ ├── losses/
│ ├── models/
│ ├── training/
│ └── utils/
├── requirements.txt
└── README.md
Only lightweight example manifests are included in this repository:
data/examples/pretrain_manifest.example.csv
data/examples/downstream_manifest.example.csv
First, prepare a pretraining manifest that links sequence, native-structure, and MD-derived inputs:
PYTHONPATH=. python scripts/prepare_pretrain_manifest.py \
--input-fasta data/pretrain/all_sequences.fasta \
--nature-dir data/pretrain/nature \
--md-dir data/pretrain/MD \
--output data/pretrain/pretrain_manifest.csvSecond, run multimodal pretraining with the formal pretraining configuration:
PYTHONPATH=. python scripts/run_pretrain.py \
--config configs/pretrain/pretrain.yamlThird, prepare downstream manifests. The following example prepares the ToxTeller task:
PYTHONPATH=. python scripts/prepare_downstream_manifests.py \
--task toxteller \
--input-path data/downstream/ToxTeller \
--pretrain-manifest data/pretrain/pretrain_manifest.csv \
--output-dir data/downstream/processed/toxtellerSupported downstream tasks are:
toxteller
prmftp
ppikb
conotoxin
For PPIKB, the default raw-data root is data/downstream/processed/regression, with run_id selecting a run-specific split directory such as run_1. For Conotoxin, the negative-set filtering helper writes by default to data/downstream/conotoxin/filtered_id70, matching configs/downstream/conotoxin.yaml.
Fourth, fine-tune a downstream model:
PYTHONPATH=. python scripts/run_downstream_finetuning.py \
--config configs/downstream/toxteller.yamlThe pretraining workflow is implemented in scripts/run_pretrain.py and controlled by configs/pretrain/pretrain.yaml.
PYTHONPATH=. python scripts/run_pretrain.py \
--config configs/pretrain/pretrain.yamlDownstream behavior is controlled by the YAML files in configs/downstream/.
For full fine-tuning:
sequence_encoder_trainable: true
structure_encoder_trainable: true
fusion_transformer_trainable: true
train_only_task_head: falseExample command:
PYTHONPATH=. python scripts/run_downstream_finetuning.py \
--config configs/downstream/toxteller.yamlProbe Rg, Re, and asphericity from learned embeddings:
PYTHONPATH=. python scripts/probe_global_structural_properties.py \
--config configs/pretrain/pretrain.yaml \
--checkpoint path/to/pretrain_checkpoint.pth \
--output-dir results/global_structural_propertiesPYTHONPATH=. python scripts/probe_total_energy.py \
--config configs/pretrain/pretrain.yaml \
--checkpoint path/to/pretrain_checkpoint.pth \
--energy-csv path/to/frame_energy.csv \
--output-dir results/total_energy_probePYTHONPATH=. python scripts/probe_energy_terms.py \
--config configs/pretrain/pretrain.yaml \
--checkpoint path/to/pretrain_checkpoint.pth \
--energy-csv path/to/frame_energy_terms.csv \
--output-dir results/energy_terms_probeAnalyze whether learned embeddings preserve trajectory-level structural geometry:
PYTHONPATH=. python scripts/analyze_rmsd_embedding_consistency.py \
--config configs/pretrain/pretrain.yaml \
--checkpoint path/to/pretrain_checkpoint.pth \
--output-dir results/rmsd_embedding_consistency