ares-softmoe-4b-l2-consecutive-225K

ares-softmoe-4b-l2-consecutive-225K

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

This is the soft routing with an L2-normalized router variant, with the last 10 of the 20 layers using MoE feed-forwards and the first 10 staying dense (consecutive placement), taken at 225,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 owns 64 learned slots. Every token contributes to every slot through a softmax dispatch, the experts run over the 2048 slots, and the results are recombined per token. No token is dropped and no load-balancing loss is needed.

This checkpoint sets moe_normalize: true: both the hidden states and the routing projection phi are L2-normalized before the routing logits are computed, with a single learned scalar restoring the logit scale. The intent is to keep routing logits bounded as depth and training length grow.

Compute cost

No parameter sparsity. Soft routing touches every expert on every forward pass, so all 4.24B parameters are active. Unlike top-k MoE, no expert is ever skipped.

What the routing changes is where the cost comes from. A dense layer's cost grows with every token; the experts here always run over a fixed number of slots, however long the input is. So this layer is relatively cheaper on long sequences and pricier on short ones: roughly 2.5x a dense layer at the 1024-token training context, breaking even around 4096 tokens.

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-softmoe-4b-l2-consecutive-225K", 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.

Validation masked-token accuracy at 225,000 steps: 0.307.

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.
• L2 normalization bought no accuracy here; it cost sample efficiency. This checkpoint needed 225,000 steps to reach what the unnormalized run reached at 150,000 (0.341 vs 0.341, a difference well inside the noise: the two split the 217 assays 106 to 111). The motivation for normalizing was to keep routing logits bounded over long training, not to raise the score, and nothing in these results argues it is needed at this scale. Prefer ares-softmoe-4b-consecutive-150K unless you specifically want to study the normalized router.
• 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.