miae-l / README.md
dexiongc's picture
Improve model card metadata and links (#1)
371fab5
metadata
library_name: tedbench
license: bsd-3-clause
pipeline_tag: graph-ml
tags:
  - protein
  - structure
  - fold-classification
  - tedbench

TEDBench — Pretrained autoencoder (structure only)

Variant: miae_l  |  Parameters: 339M  |  Layers: 24  |  Hidden dim: 1024  |  Attn heads: 16

This is a pretrained MiAE checkpoint. Use it as a feature extractor or as the starting point for fine-tuning.

This model was presented in the paper Protein Fold Classification at Scale: Benchmarking and Pretraining.

Part of the TEDBench benchmark for protein fold classification (ICML 2026). MiAE is an SE(3)-invariant masked autoencoder that masks up to 90% of backbone frames and reconstructs the full structure with a lightweight decoder.

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

from tedbench.utils.io import load_from_hf

model = load_from_hf("TEDBench/miae-l")
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},
  year={2026}
}