Protein Fold Classification at Scale: Benchmarking and Pretraining

Summary

The paper makes two coupled contributions. First, it introduces TEDBench, a large-scale, non-redundant, domain-level fold-classification benchmark built by projecting the automated TED annotation of the AlphaFold Database onto the Foldseek-clustered AFDB (pLDDT > 80), together with an experimentally-derived CATH v4.4 external test set. At 462K training/val/test proteins across 965 aggregated CATH T-level labels plus a 27,638-structure external test, TEDBench is an order of magnitude larger than prior fold benchmarks (DeepSF, TAPE, ProteinShake) — a scale the reproducibility check confirms exactly. Second, it proposes MiAE (Masked Invariant Autoencoders), a masked-autoencoding recipe for backbone frames that inherits the asymmetric encoder–decoder pattern of vision MAEs: an SE(3)-invariant encoder processes only ~10% of visible residue frames, and a lightweight decoder reconstructs the full backbone from latent tokens and mask tokens, trained with an ESM3-style composite loss.

The conceptual move — treating fold classification as an ImageNet-scale recognition task on backbone geometry and showing vision-style masked pretraining transfers to it — is timely and cleanly executed. MiAE numbers on Table 2's TED test split reproduce essentially exactly against the released checkpoints (max 0.02 F1 deviation on the rows that were run), and the ablations broadly reinforce the MAE analogy: a 90% optimal masking ratio, a deep/narrow decoder helping mean pooling, and a load-bearing inverse-folding loss.

The main limitations concern what the benchmark measures and how confidently methods can be compared on it. Labels come from an automated Foldseek + Merizo-search pipeline whose noise floor is never quantified, and the most natural baseline — Foldseek k-NN, which produced the labels — is not evaluated as a classifier. Cross-model comparisons mix architecture, corpus, and modality without controlled variants; every headline number is single-seed; and several summary sentences depart from the tables. The released repo is well-organised and reproduces exactly, but the shipped datamodule hangs on test-only eval, ablation checkpoints and figure reproducers are absent, and non-default cells still require a full pretraining rerun to verify.

Strengths

• Timely, scale-appropriate benchmark. TEDBench is ~30× larger than prior structural fold benchmarks and, unlike TAPE-fold or DeepSF, exposes long-tail structural diversity across 965 CATH-topology classes with an explicit rare-class merging protocol. Every scale claim in the paper (462,175 / 27,638 / 965 / 749,679 / 880) reproduced exactly in the code-execution check.
• External-test-set discipline. Curating a CATH v4.4 40% non-redundant, experimentally determined test set that is disjoint from the training corpus is a strong evaluation choice and gives a real signal about transfer from predicted to experimental structures; ID-level disjointness with TEDBench also reproduces exactly.
• Clean adaptation of MAE to backbone geometry. MiAE's SE(3)-invariant encoder-on-visible-frames design is a natural port of the vision MAE recipe and, in the ablations, recovers the same high-masking-ratio phenomenon at even more extreme values (90%) than in images (75%).
• Compute-efficient recipe. Even accounting for caveats in the ESM2/SaProt compute comparison, MiAE reaches SaProt-650M-competitive representations at roughly 100M parameters and without a sequence-scale pretraining corpus — a useful efficiency result.
• Ablation coverage of MiAE-side design choices. Decoder depth/width, mask ratio, sequence conditioning, inverse-folding auxiliary loss, mask sampling strategy, pretraining data size, and model size are all systematically ablated, and the released code exposes each lever as a configurable knob.
• Faithful released model code. The published repository (afdb_ted.py pipeline, main_pretrain.py/main_finetune_ted.py entry points, HuggingFace-hosted MiAE-S/B/B+seq/L checkpoints) mirrors the paper's methodology; the MiAE forward pass reproduces the Table 2 MiAE rows on the TED test split to 0.02 F1 points.

Weaknesses

• Label-noise floor never quantified. TEDBench labels come from automated Foldseek + Merizo-search assignment. No per-sample agreement with expert CATH annotation is reported, so the reader cannot tell whether the 71–75 macro-F1 ceiling is model-limited or label-limited. This is a foundational concern for a benchmark that other authors will optimise against.
• Missing Foldseek nearest-neighbor baseline. The paper motivates its problem against retrieval-based fold assignment and uses Foldseek to build the benchmark, yet never runs a Foldseek k-NN classifier as a baseline in Tables 2/3. Its absence makes it impossible to say whether MiAE is genuinely learning transferable structural priors or approximating the retrieval pipeline that generated the labels.
• Confounded cross-model comparisons. ESM2 (UniRef50 sequences), SaProt (paired sequence + 3Di tokens on ~40M structures), and MiAE (749K backbones aligned with TEDBench distribution) differ along multiple axes at once. The one-sentence caveat is not carried into the abstract's stronger "outperforms state-of-the-art" framing.
• Single-seed measurements throughout. Tables 2, 3, 4, 5 and Figures 4, 5 report single-run numbers. Many bold-vs-underline distinctions are ~1 F1 point; several ablation swings are ≤3 points. Without per-seed variance, the bold entries are not statistically defensible.
• Under-ablated architectural choices. The paper credits the recipe's strength to MAE-style pretraining but does not ablate two design choices that could carry as much weight: (i) the loss operates on all backbone atoms rather than just masked residues — a real departure from MAE that is asserted but not tested against a masked-only variant; (ii) the ESM3-style geometric-attention blocks in front of the encoder are never compared to a plain-Transformer variant, so "geometric encoder helps" and "masked pretraining helps" are entangled.
• Sequence-conditioning fairness. MiAE-B+seq consumes both structure and sequence and beats sequence-only ESM2 on some F1 rows; the information asymmetry is never called out. The strongest sequence-vs-structure comparisons in the paper are effectively "structure + sequence" vs "sequence only".
• Baseline fine-tuning uses MiAE's optimiser recipe unchanged. ESM2 and SaProt are fine-tuned with the same layer-wise LR decay (0.8) and epoch budget as MiAE; nothing establishes that schedule is competitive for those architectures at TEDBench scale, so the fine-tuning-side comparisons risk under-serving the baselines.
• Coarse aggregate metrics. Only overall accuracy and macro-F1 are reported. Given the deliberately long-tailed 965-class label space and the '.x' aggregation, per-Architecture and per-support-quantile breakdowns would clarify whether MiAE's advantage sits on head classes or on the tail.
• Ambiguities and small errors in the numeric claims. "490K proteins" conflates TEDBench and the disjoint external set; "over 15 points" for MiAE-L vs MiAE-S is actually 14.07; "59 to 70" silently refers to external-test F1 rather than the primary split; "~10% higher accuracy on external test" holds for accuracy but the macro-F1 delta is only 2–3 points; "consistent 4%" fine-tuning gain is variant-dependent (some variants show ≤2 points). Individually minor, but the density of small mismatches erodes trust in the tables.
• Under-argued interpretive claims. The rationale that masking beyond 90% "leads to more informative representations" is contradicted by the F1 drop at 0.95 in Figure 4; the "higher label fidelity" claim for CATH v4.4 is asserted rather than shown; the non-masking fine-tune underperforming from-scratch is explained away by an untested optimiser-mismatch hypothesis; the attention visualisations in Figure 7 are used to make biological-interpretability claims without a documented reproducer or identified proteins; the "MiAE representations already separate fold classes effectively" claim is stretched by the MiAE-S linear-probe F1 (49.43) sitting well below ESM2/SaProt at similar or smaller scales.
• Under-specified pipeline details. "Largest domain" is not defined and ties are not addressed, contradicting the "single, unambiguous target" claim; the relationship between the external test's 880 T-level classes and the 965 training labels is not laid out; the Figure 1 split percentages sum across two disjoint corpora.
• Contributions bullet overstates the ranking. "Consistently outperform" is bested by ESM2-3B/15B on linear probing at higher scale; the qualifier "at comparable parameter budgets" that Section 5.3 uses is dropped in the contributions list.

Reproducibility & code

• Headline aggregates and MiAE Table 2 rows reproduce exactly. Running the released code against the public HuggingFace checkpoints reproduces the MiAE-B/L rows of Table 2 on the TED test split within 0.02 F1 points (MiAE-B scratch 78.36/71.61 vs paper 78.36/71.60; MiAE-L scratch 78.34/70.95 vs 78.34/70.95; MiAE-L fine-tuned 80.08/73.49 vs 80.10/73.47). Table 1's split counts, the 965-class label space, the 749,679-structure pretraining subset, the 27,638/880 CATH external set, and ID-level disjointness all match exactly. The external-vs-test-split accuracy gap (~10 points) also holds across every replicated MiAE variant (+9.8 to +10.8).
• Table 3 (linear probing), Tables 4/5/9 (ablations), and all baseline rows are not verifiable from this run. The linear-probing driver was prepared but never executed; the equivariant baselines (GotenNet/E3NN/MACE), ESM2, and SaProt were never trained or evaluated. All of the paper's cross-model rankings and the entire ablation section therefore rest on numbers this run could neither confirm nor contradict. The obstacles are practical: baselines require significant compute, and each ablation cell requires a fresh ~100k-step (~768 GPU-hour) MiAE pretraining run because only the four default checkpoints are published.
• Undocumented datamodule hang in the shipped code. The evaluator had to patch tedbench/data/hf_datasets.py because it built the full 369,740-row train split even when only the 46,218-row test split was requested, indefinitely stalling the shipped main_test_ted.py orchestration. This was the blocking root cause across two run attempts. Custom parallel-eval drivers and a linear-probing driver had to be added around the shipped entry points; tedbench/utils/linear_clf.py was also patched. The MiAE model forward pass, the HF-loaded test data, and the torchmetrics-based metrics are the shipped code, so the reproduced numbers are faithful — but a first-time reproducer running the shipped pipeline end-to-end will hit the same hang with no warning.
• Sweep, visualisation, and results-aggregation scripts absent. Figure 4 (masking-ratio sweep), Figure 5 (pretraining-data-size sweep), Figures 6/8 (t-SNE), Figure 7 (attention heatmaps), and Figure 9 (reconstructions) all lack reproducers. In particular, the encoder blocks in minesm/layers/blocks.py do not expose attention weights, and the four proteins in Figure 7 are not identified. No script assembles Tables 2/3 from per-model outputs; a reproducer must run ~14 configurations and stitch results by hand.
• Baseline compute figures cited, not measured. The Table 8 pretraining numbers for ESM2/SaProt are copied from the SaProt paper rather than measured on the authors' hardware, so "much fewer computational resources" is not a matched-hardware wall-clock claim.
• CATH v4.4 external download is not snapshotted. The download points at ftp://orengoftp.biochem.ucl.ac.uk/cath/releases/latest-release/, a rolling URL. Once CATH updates, the exact 27,638-structure set will no longer be recoverable.
• Appendix pseudocode diverges from the released code. The Appendix A process_cath_labels defaults cts_cutoff=20, while the shipped tedbench/data/afdb_ted.py uses 10 (which is what Section 3.2 and Table 1 assume). The appendix also implements an undocumented fallback that folds the smallest non-rare class into .x when the summed rare count is still below cutoff.

Recommended Changes

Essential

• Add a Foldseek k-NN classifier baseline to Tables 2 and 3 (nearest-training-neighbor label transfer, and ideally with an alignment-score rejection variant), so readers can see how much of MiAE's F1 sits above the retrieval pipeline that produced the TED labels — addresses "Missing Foldseek nearest-neighbor baseline".
• Quantify the TED label-noise floor on a small hand-checked (or CATH-consortium overlap) subset and report it explicitly as the label-noise floor for TEDBench — addresses "Label-noise floor never quantified".
• Rerun the headline MiAE rows plus the top 1–2 baselines per section across ≥3 seeds and report mean ± std, so the bold/underline ordering has statistical footing — addresses "Single-seed measurements throughout".
• Fix the datamodule hang and add a test-split-only code path so main_test_ted.py runs on the shipped repo without user patches — addresses "Undocumented datamodule hang in the shipped code".
• Release the two most load-bearing ablation checkpoints (mask_ratio=0.0 and w/o inverse-folding loss) on HuggingFace so those rows can be validated without a 768 GPU-hour re-pretraining — addresses "Non-default ablation cells require full re-pretraining".
• Reconcile the cts_cutoff default and document the fallback merging rule by setting the Appendix A pseudocode default to 10 and adding a sentence to Section 3.2 explaining the "include the smallest non-rare class" step — addresses "Appendix pseudocode diverges from the released code".
• Fix the numeric mismatches: replace "490K proteins" with "462K", "over 15 points" with "14 points", disclose that "59 to 70" is external-test F1, restrict the "~10% higher accuracy on external test" claim to accuracy (and add the much smaller macro-F1 delta), and rewrite "consistent 4%" as a range or as per-variant deltas — addresses "Ambiguities and small errors in the numeric claims".
• Add a per-Architecture (or head-vs-tail) macro-F1 breakdown to at least Table 3 — addresses "Coarse aggregate metrics".
• Soften the abstract and contributions bullet to acknowledge the ESM2-3B/15B linear-probe scaling gap and the parameter-budget qualifier that Section 5.3 already uses — addresses "Contributions bullet overstates the ranking".

Suggested

• Add a masked-only reconstruction loss ablation (loss on masked residues, or on pairs involving at least one masked residue), to isolate the effect of the paper's departure from vanilla MAE — addresses "Reconstruction on all backbone atoms is unablated".
• Ablate the geometric-attention blocks against plain Transformer blocks so "geometric encoder helps" and "masked pretraining helps" can be separated — addresses "Geometric-attention architecture choice is not ablated".
• Actually rerun mask_ratio=0.0 fine-tuning with a matched (non-MAE) schedule instead of explaining the loss away — addresses the non-masking rationalisation weakness.
• Rewrite the >90% masking rationale to describe 0.9 as the empirical optimum and drop the claim that pushing further improves representations — addresses the ">90% masking" weakness.
• Verify or sweep baseline fine-tuning hyperparameters (ESM2, SaProt) instead of reusing MiAE's schedule verbatim, and either report a small sweep or a matched literature comparison — addresses the "Baseline fine-tuning uses MiAE's optimiser recipe unchanged" weakness.
• Define "largest domain" and its tie-breaking rule in Section 3.1, and state the fraction of proteins where ties or near-ties occur — addresses "Largest domain is not defined".
• Describe the mapping between the 880 external T-classes and the 965 training labels and report how many external proteins map to atomic (non-.x) training labels — addresses "External test label vocabulary vs. training vocabulary is not laid out".
• Redraw Figure 1's split pie so percentages are computed within TEDBench (or shown as two adjacent panels), not summed with the disjoint external set — addresses "Figure 1's split pie is computed over TEDBench + the disjoint external set".
• Publish figure-generation and results-aggregation scripts: masking-ratio sweep, pretraining-data-size sweep, t-SNE, reconstruction visualiser, attention heatmap (with an attention-return hook in minesm/layers/blocks.py and the identities of the depicted proteins), and a scripts/aggregate_results.py that assembles Tables 2 and 3 from per-model outputs — addresses "Sweep, visualisation, and results-aggregation scripts absent" and "Attention-weight interpretability figure lacks a reproducer".
• Snapshot the CATH v4.4 40% non-redundant download on the HuggingFace/Zenodo TEDBench mirror and pin the URL to a versioned path — addresses "CATH v4.4 external download is not snapshotted".
• Restate the compute comparison in pretraining tokens or FLOPs, or run a matched-hardware throughput micro-benchmark, rather than citing SaProt-paper numbers — addresses "Baseline compute figures cited, not measured".
• Split the "up to 15K proteins" citation so TAPE (sequence-first) is not grouped with structure-first fold benchmarks — addresses the "15K citation grouping" weakness.
• Frame the MiAE+seq comparison against SaProt (not ESM2) when making "outperforms" claims, and note the structure-plus-sequence information advantage over sequence-only baselines — addresses "Sequence-conditioning fairness".