Support protein/biology models: MaskedLMTask + ESMFold + Evolla

[justinchuby]

Overview

Protein/biology models are among the most downloaded scientific models on HuggingFace. Several are already supported implicitly through existing architectures, but dedicated support would unlock key use cases.

Already Working (implicit support)

• ESM-2 (facebook/esm2_t*) — Uses esm model_type, registered as BertModel + feature-extraction. 1.44M+ monthly downloads for 8M variant.
• BioGPT (microsoft/biogpt) — Uses biogpt model_type, registered as GPT2CausalLMModel. Golden data added in Fix audio encoder bugs for phi4mm #26.
• ProtBert / ProtT5 (Rostlab/*) — Use standard bert/t5 model_types, already supported.
• NVIDIA BioNeMo ESM-2 variants — Same architecture as Meta ESM-2, no new code needed.

P0: MaskedLMTask (~50 LOC, massive impact)

The #1 protein model use case is masked amino acid prediction (like BERT's MLM but for protein sequences). Adding a MaskedLMTask that takes input_ids → logits (with optional attention_mask) would unlock:

• ESM-2 masked prediction (1.44M+ monthly downloads)
• ProtBert masked prediction
• ALL encoder-only models for MLM inference (BERT, RoBERTa, DeBERTa, etc.)
• Estimated ~50 lines of code for a new task class

This is the highest-ROI item: one small task class unlocks the primary inference mode for the entire encoder-only model family.

P1: Evolla (model_type: evolla)

The only unregistered bio-specific model_type on HuggingFace. Architecture details:

• Westlake-REPL Evolla (10B, 80B) — Llama-style transformer (SiLU, RoPE, GQA) with multi-modal protein inputs: amino acid sequences, FoldSeek structural encodings, and natural language queries
• Config: hidden=4096, layers=32, heads=32, vocab=128256, max_positions=8192 — nearly identical to Llama-3
• The text decoder can likely map to CausalLMModel; the protein input processor would need a custom InputMixer
• Effort: Medium (multimodal pipeline, similar to existing VLM support)

P2: ESMFold (high value, high effort)

• 3.3M+ monthly downloads — most downloaded bio model on HuggingFace
• Novel architecture: Folding Trunk with triangular attention, triangle multiplicative updates, invariant point attention (IPA), SE(3)-equivariant structure module
• Uses ESM-2 as backbone, then iterative structure refinement with recycling
• Output: 3D atomic coordinates (not logits), pLDDT confidence scores, distogram
• Estimated 2–4 weeks of novel component work (new TriangularAttention, InvariantPointAttention, StructureModule components)
• Would need a new FoldingTask with non-standard I/O contract

Not Worth Pursuing

• AlphaFold2/3 — Not distributed through HuggingFace Transformers; custom Evoformer + MSA pipeline fundamentally incompatible with mobius's transformer-centric architecture

Priority Summary

Priority	Item	Effort	Impact
P0	MaskedLMTask (new task class)	~50 LOC	Unlocks MLM for ALL encoder-only models
P1	Evolla registration (evolla model_type)	Medium	New multimodal model_type
P2	ESMFold support	2–4 weeks	3.3M downloads, novel architecture

Key Insight

Meta dominates the bio model space. Most bio models use standard transformer architectures (BERT, GPT-2, T5) — the main gap is task-level (MaskedLMTask), not model-level. The highest-value action is adding MaskedLMTask as a reusable task class that benefits all encoder-only models, not just bio models.