From-scratch
MLX implementation of Meta's
ESM-2 protein language model, optimized for Apple Silicon.
Supports masked language modeling and residue–residue contact prediction across all six ESM-2 model sizes.
Weights are converted from the
official PyTorch checkpoints to safetensors format.
1import mlx.core as mx
2from esm_mlx import ESM2, Tokenizer
3
4model = ESM2.from_pretrained("esm2_t33_650M_UR50D") # auto-downloads weights
5tok = Tokenizer()
6
7tokens = tok.encode("MKTAYIAKQRQISFVKSHFSRQLE")
8out = model(tokens)
9logits = out["logits"] # (1, seq_len, vocab_size)
ESM-2 650M on M2 Pro (16 GB), MLX 0.30.6 vs PyTorch 2.10.0 MPS. Median latency over 50 iterations after 10 warmup passes.
FP16 widens the gap significantly. The 3.39x result at batch=16, seq=1024 likely reflects PyTorch MPS thrashing near its memory ceiling — it OOMs entirely one step later at batch=32. MLX's unified-memory allocation avoids this cliff and continues to scale linearly up to batch=192 at seq=1024, sustaining ~3,784 tok/s on 16 GB.
MIT — same as the original ESM-2 weights.
