ares-softmoe-4b-consecutive-150K

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 variant, with the last 10 of the 20 layers using MoE feed-forwards and the first 10 staying dense (consecutive 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.

Checkpoint	Routing	MoE placement	Steps	ProteinGym
ares-softmoe-4b-consecutive-150K ← you are here	Soft	Consecutive	150,000	0.341
ares-softmoe-4b-l2-consecutive-225K	Soft + L2	Consecutive	225,000	0.341
ares-softmoe-4b-l2-consecutive-150K	Soft + L2	Consecutive	150,000	0.319
ares-expert-choice-4b-interleaved-150K	Expert choice	Interleaved	150,000	0.126
ares-ec-moe-4b-86k	Expert choice	Consecutive	86,000	not evaluated

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

Parameters	4,236,352,544 (~4.24B)
Weights on disk	~16.9 GB, float32, single model.safetensors
Layers	20 encoder blocks, pre-norm
Hidden size	1024
Feed-forward size	4096, gated SiLU (SwiGLU)
Attention	Grouped-query, 16 heads / 8 KV heads, head dim 64
Position encoding	Rotary (RoPE), base 10000
Normalization	RMSNorm
Vocabulary	31 tokens (20 standard + BXZJUO + 5 special)
Trained context	1024 tokens
Experts	32
Routing	soft_router
MoE placement	Consecutive (layers 10-19)
Objective	Masked language modeling with scheduled masking + mutation noising

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.

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-consecutive-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

Data	UniRef50 (agemagician/uniref50_09012025), 68,346,946 sequences
Steps	150,000 optimizer steps
Objective	MLM, masking rate cycled over 0.15 / 0.20 / 0.25 / 0.30
Corruption	80% <mask>, 10% random-residue mutation, 10% unchanged
Masking schedule	staged_linear, starting at 15% masking and progressively mixing in the higher rates
Optimizer	AdamW, lr 3e-4, betas (0.9, 0.95), eps 1e-8, weight decay 0.01
Schedule	5% warmup, gradient clipping at 1.0
Precision	bfloat16 autocast, float32 master weights
Sequence handling	Multiple sequences packed per batch row with block-diagonal attention masking and per-sequence position IDs
Hardware	Google Cloud TPU via PyTorch/XLA with SPMD sharding, gradient checkpointing enabled

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 150,000 steps: 0.313.

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.

Aggregation	Activity	Binding	Expression	OrganismalFitness	Stability	All
Standard	0.330	0.305	0.364	0.246	0.379	0.325
Fisher-z	0.343	0.323	0.372	0.259	0.403	0.341

Against the rest of the family and public baselines

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

Model	Spearman (Fisher-z)
ElnaggarLab/ankh-large	0.393
ares-softmoe-4b-consecutive-150K	0.341
ares-softmoe-4b-l2-consecutive-225K	0.341
ares-softmoe-4b-l2-consecutive-150K	0.319
ElnaggarLab/ankh-base	0.270
ares-expert-choice-4b-interleaved-150K	0.126

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.
• 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.

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.