ares-expert-choice-4b-interleaved-150K

A ~4.22B-parameter sparse Mixture-of-Experts protein language model, pretrained from scratch on UniRef50 with a masked-language-modeling objective on TPU.

This is the expert choice variant, with MoE feed-forwards on every other layer, alternating with dense ones (interleaved placement), taken at 150,000 training steps.

Spotlight at the GenBio Workshop, ICML 2026. Ares: Loss-Free Mixture-of-Experts Routing for Bidirectional Protein Encoders.

Pick the right checkpoint

Five Ares checkpoints are published. They share an architecture and a training recipe and differ only in how tokens reach experts and where the MoE layers sit.

ProteinGym numbers are the Fisher-z aggregated Spearman over 217 DMS substitution assays (see Evaluation). Higher is better.

If you just want a protein encoder, start with ares-softmoe-4b-consecutive-150K. It and ares-softmoe-4b-l2-consecutive-225K are the two strongest of the family and are effectively tied on ProteinGym (0.341 vs 0.341; they split the 217 assays 111 to 106). The 150K checkpoint is the simpler default: same score, fewer training steps, no router normalization to reason about.

Model details

How routing works here

Each of the 32 experts selects its own top-k tokens, with k set by a capacity factor of 2.0 over the flattened batch. Load is balanced by construction (every expert receives exactly the same number of tokens), but tokens can be dropped entirely or picked by several experts at once.

Compute cost

~0.44B active parameters per token. A capacity factor of 2.0 across 32 experts means roughly 2 experts of work per token, so about 0.44B of the 4.22B parameters are exercised on any given token. The exact figure varies per token: expert choice makes no guarantee that a token is picked at all.

Usage

Ares is not part of transformers, so install the ares package first. It provides the Ares model class and the tokenizer, and no trust_remote_code is needed.

pip install git+https://github.com/hazemessamm/ares.git

Only the core dependencies are required for inference; the training and evaluation extras are never imported on the model path.

Fill in masked residues

import torch
from ares import Ares, AresProteinTokenizer

model = Ares.from_pretrained("HazemLab/ares-expert-choice-4b-interleaved-150K", dtype=torch.bfloat16).eval()
tokenizer = AresProteinTokenizer()

sequence = "MKTAYIAKQRQISFVKSHFSRQ<mask>ERLEKLLQ"
batch = tokenizer(sequence, return_tensors="pt")

with torch.no_grad():
    logits = model(**batch).logits

mask_pos = (batch["input_ids"] == tokenizer.mask_token_id).nonzero()[0, 1]
top = logits[0, mask_pos].topk(5)
for score, token_id in zip(top.values, top.indices):
    print(tokenizer.decode(token_id), float(score))

Extract residue and sequence embeddings

hidden_states[0] holds the final normalized representation.

with torch.no_grad():
    out = model(**batch)

residue = out.hidden_states[0]                       # (batch, length, 1024)
mask = batch["attention_mask"].unsqueeze(-1)
pooled = (residue * mask).sum(1) / mask.sum(1)       # (batch, 1024), mean-pooled

Both the <cls> token and mean pooling over unpadded positions are reasonable sequence representations; the downstream evaluations in the repository use mean pooling.

Notes on loading

• Weights are stored in float32 (~17 GB). Pass dtype=torch.bfloat16 unless you specifically need float32; every evaluation in the paper was run in bfloat16.
• AresProteinTokenizer builds its vocabulary in code, so it needs no download and is identical across every Ares checkpoint.
• This repo ships no tokenizer files, and does not need any: the tokenizer is constructed in code.
• The model accepts an optional sequence_ids argument for block-diagonal attention over packed sequences. Leave it unset for ordinary batched inference.

Training

Sequence packing means no compute is spent on padding. Correctness of packed training against unpacked inference is asserted in the repository's tests/test_packing_correctness.py.

Evaluation

ProteinGym (zero-shot DMS substitutions)

Spearman correlation between masked-marginal likelihood scores and measured fitness, over 217 DMS substitution assays spanning 200 UniProt entries. Scores are aggregated per UniProt entry and then per selection type, both arithmetically and under a Fisher-z transform. Inference ran in bfloat16 with no length cutoff.

Against the rest of the family and public baselines

Same protocol, same 217 assays, Fisher-z aggregation:

Per-assay and per-UniProt breakdowns for every row above are checked into the repository under evaluation/proteingym_results/.

Downstream tasks

The repository also provides fine-tuning and frozen-embedding evaluations for GB1 epistasis, fluorescence, stability, remote homology, 3- and 8-state secondary structure, and subcellular localization. Those scripts live in evaluation/; results are not included in this model card.

Limitations and known issues

• Encoder only. This is a bidirectional masked LM. It scores and embeds sequences; it does not generate them autoregressively.
• 1024-token training context. RoPE allows longer inputs to run, but nothing beyond 1024 residues was seen during training and quality past that point is untested.
• Single-chain amino acid sequences. No structure, no MSA, no multimer or nucleotide input.
• Weights are float32. Expect a ~17 GB download and load in bfloat16 for inference.
• UniRef50 inherits the biases of the sequence databases it was built from. Well-studied organisms and protein families are heavily over-represented. Zero-shot variant-effect performance varies sharply by assay type; the per-assay CSVs in the repository make this visible.
• The expert-choice router is not faithful to the original formulation. Its routing softmax normalizes over tokens rather than over experts, which changes what the router probabilities mean and what gradient the router receives. This is a known defect, and it is the most likely explanation for the gap between the expert-choice and soft-routing checkpoints. Treat this checkpoint as a research artifact for studying that failure mode rather than as a production encoder.
• Tokens can be dropped. With a capacity factor of 2.0, expert choice offers no guarantee that any particular token is selected by any expert. Dropped tokens pass through the MoE layer with a zero contribution from the feed-forward path.
• moe_num_slots appears in config.json but is inert for expert-choice routing; it only affects the soft router.
• This checkpoint scores poorly on ProteinGym (0.126 Fisher-z, against 0.341 for the soft-routing sibling trained for the same 150,000 steps). It is published for reproducibility and for the MoE interpretability work, not as a recommended encoder.
• Not validated for clinical, diagnostic, or biosafety-relevant decisions. Variant-effect predictions from this model are hypotheses for experimental follow-up, nothing more.

MoE interpretability

Ares ships an analysis pipeline for inspecting what the experts in these checkpoints actually do: per-expert amino-acid and biochemical-property preferences, positional preferences, routing heatmaps, causal expert-knockout importance, and steering interventions. See evaluation/moe_analysis/ and read ANALYSIS_OUTPUTS.md before interpreting any specialization number. It documents every artifact, every metric, and the axis each routing weight normalizes over.

Acknowledgments

We gratefully acknowledge Google's TPU Research Cloud (TRC) program for providing the Cloud TPU resources that made the training of Ares possible. We thank Dr. Ahmed Saleh Mansour for his valuable feedback and careful review of the manuscript.

Citation

@inproceedings{alsamkary2026ares,
  title     = {Ares: Loss-Free Mixture-of-Experts Routing for Bidirectional Protein Encoders},
  author    = {Alsamkary, Hazem},
  booktitle = {ICML 2026 Workshop on Generative AI and Biology (GenBio)},
  year      = {2026},
  note      = {Spotlight},
  url       = {https://openreview.net/forum?id=gq0R7xiPjg}
}

Not the final version; it will be updated when the camera-ready lands.

License

MIT.

Affiliation: Proteinea.