TEDBench/miae-s

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.

• Paper: Protein Fold Classification at Scale: Benchmarking and Pretraining
• Repository: 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:

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:

import tedbench

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

From a Lightning checkpoint

from tedbench.model import MiAE

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

Citation

@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}
}