README.md

---
library_name: tedbench
license: bsd-3-clause
pipeline_tag: graph-ml
tags:
- protein
- structure
- fold-classification
- tedbench
---

# TEDBench — Pretrained autoencoder (structure only)

**Variant:** `miae_s`  |  **Parameters:** 29M  |  **Layers:** 6  |  **Hidden dim:** 512  |  **Attn heads:** 8

This repository contains a **pretrained MiAE** checkpoint. Use it as a feature extractor or as the starting point for fine-tuning on protein structure tasks.

MiAE (Masked Invariant Autoencoder) is an $\mathrm{SE(3)}$-invariant masked autoencoder that masks up to 90% of backbone frames and reconstructs the full structure with a lightweight decoder. It was introduced in the paper [Protein Fold Classification at Scale: Benchmarking and Pretraining](https://arxiv.org/abs/2605.18552).

- **Paper:** [Protein Fold Classification at Scale: Benchmarking and Pretraining](https://arxiv.org/abs/2605.18552)
- **Repository:** [BorgwardtLab/TEDBench](https://github.com/BorgwardtLab/TEDBench)

## Architecture sizes

| Variant | Params | Layers | Hidden dim | Attn heads |
|---------|-------:|-------:|-----------:|-----------:|
| `miae_s` | 29 M | 6 | 512 | 8 |
| `miae_b` | 102 M | 12 | 768 | 12 |
| `miae_l` | 339 M | 24 | 1 024 | 16 |

Append `+model.use_seq_input=true` to `miae_b` for the **+seq** variant.

## Usage

### Load from the HuggingFace Hub

You can load the model using the `tedbench` library:

```python
from tedbench.utils.io import load_from_hf

model = load_from_hf("TEDBench/miae-s")
model.eval()
```

Alternatively, using the library's high-level API:

```python
import tedbench

model = tedbench.load_model("miae-s")
model.eval()
```

### From a Lightning checkpoint

```python
from tedbench.model import MiAE

model = MiAE.load_from_checkpoint("model.ckpt", weights_only=False)
model.eval()
```

## Citation

```bibtex
@inproceedings{chen2026tedbench,
  title={Protein Fold Classification at Scale: Benchmarking and Pretraining},
  author={Chen, Dexiong and Manolache, Andrei and Niepert, Mathias and Borgwardt, Karsten},
  booktitle={Proceedings of the 43rd International Conference on Machine Learning (ICML)},
  year={2026}
}
```