{"id":51169413,"url":"https://github.com/borgwardtlab/tedbench","last_synced_at":"2026-06-26T23:03:21.274Z","repository":{"id":359934709,"uuid":"1236707507","full_name":"BorgwardtLab/TEDBench","owner":"BorgwardtLab","description":"Protein Fold Classification at Scale: Benchmarking and Pretraining (ICML 2026 oral)","archived":false,"fork":false,"pushed_at":"2026-05-24T06:39:33.000Z","size":1024,"stargazers_count":1,"open_issues_count":0,"forks_count":0,"subscribers_count":0,"default_branch":"main","last_synced_at":"2026-05-24T08:23:17.183Z","etag":null,"topics":["benchmark","geometric-deep-learning","icml-2026","protein-fold-prediction","protein-representation-learning"],"latest_commit_sha":null,"homepage":"https://arxiv.org/abs/2605.18552","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":"CHANGELOG.md","contributing":null,"funding":null,"license":"LICENSE","code_of_conduct":null,"threat_model":null,"audit":null,"citation":"CITATION.cff","codeowners":null,"security":null,"support":null,"governance":null,"roadmap":null,"authors":null,"dei":null,"publiccode":null,"codemeta":null,"zenodo":null,"notice":null,"maintainers":null,"copyright":null,"agents":null,"dco":null,"cla":null}},"created_at":"2026-05-12T13:52:06.000Z","updated_at":"2026-05-24T06:39:37.000Z","dependencies_parsed_at":null,"dependency_job_id":null,"html_url":"https://github.com/BorgwardtLab/TEDBench","commit_stats":null,"previous_names":["borgwardtlab/tedbench"],"tags_count":2,"template":false,"template_full_name":null,"purl":"pkg:github/BorgwardtLab/TEDBench","repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/BorgwardtLab%2FTEDBench","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/BorgwardtLab%2FTEDBench/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/BorgwardtLab%2FTEDBench/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/BorgwardtLab%2FTEDBench/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/BorgwardtLab","download_url":"https://codeload.github.com/BorgwardtLab/TEDBench/tar.gz/refs/heads/main","sbom_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/BorgwardtLab%2FTEDBench/sbom","scorecard":null,"host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":286080680,"owners_count":34835782,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2026-05-26T15:22:16.424Z","status":"online","status_checked_at":"2026-06-26T02:00:06.560Z","response_time":106,"last_error":null,"robots_txt_status":"success","robots_txt_updated_at":"2025-07-24T06:49:26.215Z","robots_txt_url":"https://github.com/robots.txt","online":true,"can_crawl_api":true,"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":["benchmark","geometric-deep-learning","icml-2026","protein-fold-prediction","protein-representation-learning"],"created_at":"2026-06-26T23:03:15.262Z","updated_at":"2026-06-26T23:03:21.268Z","avatar_url":"https://github.com/BorgwardtLab.png","language":"Python","funding_links":[],"categories":[],"sub_categories":[],"readme":"\u003cp align=\"center\"\u003e\n  \u003cimg src=\"./resources/tedbench_logo_transp_trimmed.png\" width=\"50%\" alt=\"TedBench logo\" /\u003e\n\u003cp\u003e\n\n# TEDBench\n\n[![arXiv](https://img.shields.io/badge/arXiv-2605.18552-b31b1b.svg)](https://arxiv.org/abs/2605.18552)\u0026nbsp;\n[![Hugging Face](https://img.shields.io/badge/Hugging%20Face-TEDBench-blue.svg)](https://huggingface.co/TEDBench)\u0026nbsp;\n\n**TEDBench** is a large-scale, non-redundant benchmark for protein fold classification,\ntogether with **MiAE** (Masked Invariant Autoencoders), a self-supervised pretraining\nframework for protein structure representations.\n\n\u003e **Paper:** *Protein Fold Classification at Scale: Benchmarking and Pretraining*  \n\u003e Dexiong Chen, Andrei Manolache, Mathias Niepert, Karsten Borgwardt (ICML 2026 oral)\n\n---\n\n## Overview\n\nTEDBench is built from the [Encyclopedia of Domains (TED)](https://zenodo.org/records/13908086)\nannotations projected onto the Foldseek-clustered AlphaFold Database.\n\n| Split | Structures |\n|---|---|\n| Train | 369,740 |\n| Val | 46,217 |\n| Test | 46,218 |\n| External test (CATH 4.4 experimental) | 27,638 |\n\nAll structures are classified into **965 CATH topology (T-level) classes**.\n\nMiAE is an SE(3)-invariant masked autoencoder that masks up to 90 % of backbone\nframes, processes only the visible residues with a geometric encoder, and\nreconstructs the full backbone structure with a lightweight decoder.\n\n---\n\n## Installation\n\n**From PyPI (recommended):**\n\n```bash\npip install tedbench\n```\n\nFor running ESM2 / SaProt baselines, add the `baselines` extra:\n\n```bash\npip install \"tedbench[baselines]\"\n```\n\n**From source** (for training, baselines, or development):\n\n```bash\n# 1. Create and activate environment\nmicromamba create -n tedbench python=3.12 -y\nmicromamba activate tedbench\n\n# 2. Install dependencies\nuv pip install -r requirements.txt\n\n# 3. Install the tedbench package (editable)\nuv pip install -e .\n```\n\n---\n\n## Datasets\n\nDatasets are available from two sources:\n\n| Dataset | HuggingFace | Direct download |\n|---|---|---|\n| TEDBench (AFDB + CATH labels) | [`TEDBench/ted`](https://huggingface.co/datasets/TEDBench/ted) | [MPCDF datashare](https://datashare.mpcdf.mpg.de/s/m4owC3SQbd2r6rk) |\n| AFDB pretraining corpus | [`TEDBench/afdb`](https://huggingface.co/datasets/TEDBench/afdb) | [MPCDF datashare](https://datashare.mpcdf.mpg.de/s/m4owC3SQbd2r6rk) |\n| CATH 4.4 experimental test set | [`TEDBench/cath`](https://huggingface.co/datasets/TEDBench/cath) | [MPCDF datashare](https://datashare.mpcdf.mpg.de/s/pjXMpff7GsYTR22) |\n\nThe HuggingFace repos require no local setup; the MPCDF archives are auto-downloaded and cached the first time a local dataset class is instantiated (default roots: `./datasets/ted/` and `./datasets/cath/`).\n\nEach sample contains: `coords` `[L, 3, 3]` (backbone N/Cα/C, float32), `plddt` `[L]`, `residue_index` `[L]`, `seq_ids` `[L]`, `sequence`, and `label` (integer CATH topology index).\n\n### Load directly with `datasets`\n\n```python\nfrom datasets import load_dataset\nimport torch\n\n# TEDBench — train / val / test with CATH labels\nted = load_dataset(\"TEDBench/ted\")\nsample = ted[\"train\"][0]\ncoords    = torch.tensor(sample[\"coords\"])   # [L, 3, 3]\nlabel     = sample[\"label\"]                  # int index\ncath_code = ted[\"train\"].features[\"label\"].int2str(label)  # e.g. \"3.40.50.300\"\n\n# CATH 4.4 external test set\ncath = load_dataset(\"TEDBench/cath\", split=\"test\")\n\n# AFDB pretraining corpus\nafdb = load_dataset(\"TEDBench/afdb\", split=\"train\")\n```\n\n### Use with `LightningStructureDataset`\n\n**From HuggingFace** (`dataset_name=\"hf_ted\"` / `\"hf_cath4.4\"` / `\"hf_afdb\"`):\n\n```python\nfrom tedbench.data import LightningStructureDataset\n\ndm = LightningStructureDataset(\n    root=\"TEDBench/ted\",   # HF repo ID\n    dataset_name=\"hf_ted\",\n    batch_size=32,\n    num_workers=4,\n)\ndm.setup(\"fit\")\nfor batch in dm.train_dataloader():\n    print(batch.keys()) \n    # dict_keys(['coords', 'residue_index', 'seq_ids', 'protein_chain', 'mask', 'label'])\n```\n\n**Auto-download from MPCDF** (`dataset_name=\"ted\"` / `\"cath4.4\"` / `\"afdb_stream\"`): the archive is fetched from the MPCDF datashare and cached under `root` on first use — no manual download needed:\n\n```python\ndm = LightningStructureDataset(\n    root=\"./datasets/ted\",   # local cache directory\n    dataset_name=\"ted\",\n    batch_size=32,\n    num_workers=4,\n)\ndm.setup(\"fit\")\nfor batch in dm.train_dataloader():\n    print(batch.keys()) \n    # dict_keys(['coords', 'residue_index', 'seq_ids', 'protein_chain', 'mask', 'label'])\n```\n\nPass `datamodule=hf_ted` (or `datamodule=hf_cath_test`, `datamodule=hf_afdbfs`) to any\ntraining script to use HuggingFace; omit it (or use the default config) for the\nauto-downloading local variant.\n\n---\n\n## Pretrained Models\n\nAll models are available on HuggingFace and can be loaded with a single call:\n\n```python\nimport tedbench\n\nmodel = tedbench.load_model(\"miae-b\")     # pretrained MiAE-B (short name)\nmodel = tedbench.load_model(\"miae-b-ft\")  # fine-tuned on TEDBench\n\n# List all available models\nfor m in tedbench.list_models():\n    print(m[\"name\"], m[\"type\"], m[\"params\"])\n```\n\n### Pretrained MiAE (feature extractor / fine-tuning starting point)\n\n| Model | HF repo | Params |\n|---|---|---|\n| MiAE-S | [`TEDBench/miae-s`](https://huggingface.co/TEDBench/miae-s) | 29 M |\n| MiAE-B | [`TEDBench/miae-b`](https://huggingface.co/TEDBench/miae-b) | 102 M |\n| MiAE-B+seq | [`TEDBench/miae-b-seq`](https://huggingface.co/TEDBench/miae-b-seq) | 102 M |\n| MiAE-L | [`TEDBench/miae-l`](https://huggingface.co/TEDBench/miae-l) | 339 M |\n\n### Fine-tuned on TEDBench (fold classifier)\n\n| Model | HF repo | TEDBench test acc | CATH 4.4 test acc |\n|---|---|---|---|\n| MiAE-S (ft) | [`TEDBench/miae-s-ft`](https://huggingface.co/TEDBench/miae-s-ft) | 72.28 | 76.08 |\n| MiAE-B (ft) | [`TEDBench/miae-b-ft`](https://huggingface.co/TEDBench/miae-b-ft) | 73.71 | 75.72 |\n| MiAE-B+seq (ft) | [`TEDBench/miae-b-seq-ft`](https://huggingface.co/TEDBench/miae-b-seq-ft) | 74.56 | 77.34 |\n| MiAE-L (ft) | [`TEDBench/miae-l-ft`](https://huggingface.co/TEDBench/miae-l-ft) | 73.47 | 76.46 |\n\n### Trained from scratch on TEDBench (no pretraining)\n\n| Model | HF repo |\n|---|---|\n| MiAE-S (sc) | [`TEDBench/miae-s-sc`](https://huggingface.co/TEDBench/miae-s-sc) |\n| MiAE-B (sc) | [`TEDBench/miae-b-sc`](https://huggingface.co/TEDBench/miae-b-sc) |\n| MiAE-B+seq (sc) | [`TEDBench/miae-b-seq-sc`](https://huggingface.co/TEDBench/miae-b-seq-sc) |\n| MiAE-L (sc) | [`TEDBench/miae-l-sc`](https://huggingface.co/TEDBench/miae-l-sc) |\n\n---\n\n## Evaluation\n\nEvaluate any model from the HuggingFace Hub without any local data setup:\n\n```bash\n# Test fine-tuned MiAE-B on TEDBench test split\npython main_test_ted.py \\\n    datamodule=hf_ted \\\n    pretrained_model_path=TEDBench/miae-b-ft\n\n# Test on the CATH 4.4 external experimental test set\npython main_test_ted.py \\\n    datamodule=hf_cath_test \\\n    pretrained_model_path=TEDBench/miae-b-ft\n\n# Test fine-tuned MiAE-B+seq on TEDBench test split\npython main_test_ted.py \\\n    datamodule=hf_ted \\\n    +model.use_seq_input=true \\\n    pretrained_model_path=TEDBench/miae-b-seq-ft\n\n# Test supervised-from-scratch MiAE-B\npython main_test_ted.py \\\n    pretrained_model_path=TEDBench/miae-b-sc\n\n# Linear probing with pretrained MiAE-B\npython main_linprobe_ted.py \\\n    pretrained_model_path=TEDBench/miae-b\n```\n\n---\n\n## Model Variants\n\n| Name | Params | Layers | Hidden dim | Attn heads |\n|---|---|---|---|---|\n| `miae_s` | 29 M | 6 | 512 | 8 |\n| `miae_b` | 102 M | 12 | 768 | 12 |\n| `miae_l` | 339 M | 24 | 1 024 | 16 |\n\nPass `model.name=\u003cvariant\u003e` to any training script to select a size.\nAdd `model.use_seq_input=true` to enable the **+seq** variant (structure + sequence).\n\n---\n\n## Training and Reproducing Paper Results\n\nSee [TRAINING.md](TRAINING.md) for full pretraining, fine-tuning, linear probing,\nand baseline reproduction commands with hyperparameter tables.\n\nThe `baselines/` directory contains scripts for ESM2, SaProt, and ProteinMPNN baselines.\nSee [TRAINING.md](TRAINING.md#baselines) for usage.\n\n---\n\n## Citation\n\n```bibtex\n@inproceedings{chen2026tedbench,\n  title={Protein Fold Classification at Scale: Benchmarking and Pretraining},\n  author={Chen, Dexiong and Manolache, Andrei and Niepert, Mathias and Borgwardt, Karsten},\n  booktitle={Proceedings of the 43rd International Conference on Machine Learning (ICML)},\n  year={2026}\n}\n```\n\n---\n\n## License\n\nBSD-3-Clause\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fborgwardtlab%2Ftedbench","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Fborgwardtlab%2Ftedbench","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fborgwardtlab%2Ftedbench/lists"}