{"id":21280593,"url":"https://github.com/borgwardtlab/pst","last_synced_at":"2025-07-11T10:32:26.130Z","repository":{"id":219248128,"uuid":"736241303","full_name":"BorgwardtLab/PST","owner":"BorgwardtLab","description":"Protein Structure Transformer (PST): Endowing pretrained protein language models with structural knowledge","archived":false,"fork":false,"pushed_at":"2024-10-17T06:42:43.000Z","size":2406,"stargazers_count":31,"open_issues_count":1,"forks_count":1,"subscribers_count":4,"default_branch":"main","last_synced_at":"2024-10-19T08:36:59.444Z","etag":null,"topics":["evolutionary-scale-modeling","graph-transformer","protein-representation-learning","protein-structure"],"latest_commit_sha":null,"homepage":"https://arxiv.org/abs/2401.14819","language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"bsd-3-clause","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/BorgwardtLab.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"funding":null,"license":"LICENSE","code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null,"governance":null,"roadmap":null,"authors":null,"dei":null,"publiccode":null,"codemeta":null}},"created_at":"2023-12-27T11:13:43.000Z","updated_at":"2024-10-13T17:00:39.000Z","dependencies_parsed_at":"2024-07-17T16:31:42.439Z","dependency_job_id":"f4759194-d248-4205-be28-02f15d7b4cad","html_url":"https://github.com/BorgwardtLab/PST","commit_stats":null,"previous_names":["borgwardtlab/pst"],"tags_count":0,"template":false,"template_full_name":null,"repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/BorgwardtLab%2FPST","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/BorgwardtLab%2FPST/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/BorgwardtLab%2FPST/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/BorgwardtLab%2FPST/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/BorgwardtLab","download_url":"https://codeload.github.com/BorgwardtLab/PST/tar.gz/refs/heads/main","host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":225715514,"owners_count":17512903,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2022-07-04T15:15:14.044Z","host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":["evolutionary-scale-modeling","graph-transformer","protein-representation-learning","protein-structure"],"created_at":"2024-11-21T10:37:16.805Z","updated_at":"2024-11-21T10:37:17.526Z","avatar_url":"https://github.com/BorgwardtLab.png","language":"Python","funding_links":[],"categories":[],"sub_categories":[],"readme":"# Protein Structure Transformer\n\nThe repository implements the Protein Structure Transformer (PST). The PST model endows the pretrained protein sequence model [ESM-2][6] with structural knowledge, allowing for extracting representations of protein structures. Full details of PST can be found in the [paper][1].\n\n## Citation\n\nPlease use the following to cite our work:\n\n```bibtex\n@misc{chen2024endowing,\n\ttitle={Endowing Protein Language Models with Structural Knowledge}, \n\tauthor={Dexiong Chen and Philip Hartout and Paolo Pellizzoni and Carlos Oliver and Karsten Borgwardt},\n\tyear={2024},\n\teprint={2401.14819},\n\tarchivePrefix={arXiv},\n\tprimaryClass={q-bio.QM}\n}\n```\n\n\n## Overview of PST\n\nPST uses a structure extractor to incorporate protein structures into existing pretrained protein language models (PLMs) such as [ESM-2][6].\nThe structure extractor adopts a GNN to extract subgraph representations of the 8Å-neighborhood protein structure graph at each residue (i.e., nodes on the graph). The resulting residue-level subgraph representations are then add to the $Q$, $K$ and $V$ matrices of **each** self-attention block of any (pretrained) transformer model (here we use **ESM-2**) pretrained on larger corpuses of sequences. We name the resulting model PST, which can be trained on any protein structure dataset, by either updating the full model weights or only the weights in the structure extractor. The pretraining dataset could be much smaller than the pretraining dataset used for the base sequence model, e.g., SwissProt with only 542k protein structures. \n\nBelow you can find an overview of PST with ESM-2 as the sequence backbone. The ESM-2 model weights were frozen during the training of the structure extractor. The structure extractor was trained on AlphaFold SwissProt, a dataset of 542K proteins with predicted structures. The resulting PST model can then be finetuned on a downstream task, e.g., [torchdrug][5] or [proteinshake][4] tasks. PST can also be used to simply extract representations of protein structures.\n\n![Overview of PST](assets/overview.png)\n\n## Pretrained models\n\n| Model name   | Sequence model        | #Layers | Embed dim | Notes             | Model URL                                                  |\n| ------------ | --------------------- | ------- | --------- | ----------------- | ---------------------------------------------------------- |\n| `pst_t6`     | `esm2_t6_8M_UR50D`    | 6       | 320       | Standard          | [link](https://datashare.biochem.mpg.de/s/ac9ufZ0NB2IrkZL) |\n| `pst_t6_so`  | `esm2_t6_8M_UR50D`    | 6       | 320       | Train struct only | [link](https://datashare.biochem.mpg.de/s/ARzKycmMQePvLXs) |\n| `pst_t12`    | `esm2_t12_35M_UR50D`  | 12      | 480       | Standard          | [link](https://datashare.biochem.mpg.de/s/fOSIwJAIKLYjFe3) |\n| `pst_t12_so` | `esm2_t12_35M_UR50D`  | 12      | 480       | Train struct only | [link](https://datashare.biochem.mpg.de/s/qRvDPTfExZkq38f) |\n| `pst_t30`    | `esm2_t30_150M_UR50D` | 30      | 640       | Standard          | [link](https://datashare.biochem.mpg.de/s/a3yugJJMe0I0oEL) |\n| `pst_t30_so` | `esm2_t30_150M_UR50D` | 30      | 640       | Train struct only | [link](https://datashare.biochem.mpg.de/s/p73BABG81gZKElL) |\n| `pst_t33`    | `esm2_t33_650M_UR50D` | 33      | 1280      | Standard          | [link](https://datashare.biochem.mpg.de/s/RpWYV4o4ka3gHvX) |\n| `pst_t33_so` | `esm2_t33_650M_UR50D` | 33      | 1280      | Train struct only | [link](https://datashare.biochem.mpg.de/s/xGpS7sIG7k8DZX0) |\n\n## Usage\n\n### Installation\n\nThe dependencies are managed by [mamba][2] or [conda](https://docs.conda.io/projects/miniconda/en/latest/)\n\n```\nmamba env create -f environment.yaml \nmamba activate pst\npip install -e .\n```\n\nOptionally, you can install the following dependencies to run the experiments:\n\n```\npip install torchdrug\n```\n\n### Quick start: extract representations of protein structures using PST\n\nYou can PST to simply extract representations of protein structures stored in PDB files. Just run\n\n```bash\npython scripts/pst_extract.py --help\n```\n\nIf you want to work with your own dataset, just create a `my_dataset` directory in `scripts` and put all the PDB files into `my_dataset/raw/`,\nand run:\n\n```bash\npython scripts/pst_extract.py --datadir ./scripts/my_dataset --model pst_t33_so --include_seq\n```\n\n### Use PST for protein function prediction\n\nYou can use PST to perform Gene Ontology prediction, Enzyme Commission Number prediction and any other protein function prediction tasks.\n\n#### Fixed representations\n\nTo train an MLP on top of the representations extracted by the pretrained PST models for Enzyme Commission prediction, run:\n\n```bash\npython experiments/fixed/predict_gearnet.py dataset=gearnet_ec # dataset=gearnet_go_bp, gearnet_go_cc or gearnet_go_mf for GO prediction\n```\n\n#### Finetune PST\n\nTo finetune the PST model for function prediction tasks, run:\n\n```bash\npython experiments/finetune/finetune_gearnet.py dataset=gearnet_ec # dataset=gearnet_go_bp, gearnet_go_cc or gearnet_go_mf for GO prediction\n```\n\n### Pretrain PST on AlphaFold Swissprot\n\nRun the following code to train a PST model based on the 6-layer ESM-2 model by only training the structure extractor:\n\n```bash\npython train_pst.py base_model=esm2_t6 model.train_struct_only=true\n```\n\nYou can replace `esm2_t6` with `esm2_t12`, `esm2_t30`, `esm2_t33` or any pretrained ESM-2 model.\n\n### Reproducibility datasets\n\nWe have folded structures that were not available in the PDB for our VEP datasets. You can download the dataset from [here](https://datashare.biochem.mpg.de/s/2UgA8kBwmCAVEsL), and unzip it in `./datasets`, provided your current path is the root of this repository. Similarly, download the SCOP dataset [here](https://datashare.biochem.mpg.de/s/2yUwpK7pt2TMQ5E).\n\n\n[1]: https://arxiv.org/abs/2401.14819\n[2]: https://mamba.readthedocs.io/en/latest/installation/mamba-installation.html\n[3]: https://arxiv.org/abs/2202.03036\n[4]: https://proteinshake.ai/\n[5]: https://torchdrug.ai/\n[6]: https://github.com/facebookresearch/esm/\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fborgwardtlab%2Fpst","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Fborgwardtlab%2Fpst","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fborgwardtlab%2Fpst/lists"}