ESMC-AR 848M + PLE, trained on 5B tokens
A protein autoregressive model trained from scratch for 5,000,000,000 prediction
tokens, with 848,446,464 parameters and large per-layer n-gram embeddings (PLE).
This is the completed stage1-ple-large-shuffled experiment, final checkpoint
19,074, finished on September 10, 2026 (Asia/Shanghai). It does not use
GatedResidual. The separate PLE + GatedResidual experiment is a different model.
The architecture adapts ESMC-style blocks to causal next-token prediction and adds Qwen4-Exp-style PLE. It is not the original ESMC-300M masked language model and does not load its pretrained weights. Training code: https://github.com/huyiwen/bio-next
Load and generate
Tested with PyTorch 2.5.1 and Transformers 4.51.0. This repository includes the configuration, portable inference Python files, tied FP32 safetensors weights, tokenizer, generation configuration, training configuration and export verification. The tokenizer accepts amino-acid strings directly.
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
repo = "IvanHU/esmc-ar-848m-ple-5b"
tokenizer = AutoTokenizer.from_pretrained(repo)
model = AutoModelForCausalLM.from_pretrained(
repo, trust_remote_code=True, torch_dtype=torch.float32,
).cuda().eval()
inputs = tokenizer("MKWVTFISLLFLFSSAYS", return_tensors="pt").to("cuda")
with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16):
output = model.generate(**inputs, max_new_tokens=64, do_sample=True,
temperature=0.8, top_p=0.95, use_cache=False)
print(tokenizer.decode(output[0], skip_special_tokens=True))
PLE generation currently recomputes the full context and requires use_cache=False;
KV caching is not implemented for its n-gram/convolution history. Keep prefix plus
continuation within the 2048-token training context. Default generation suppresses
all output slots except the 20 standard amino acids, including EOS, so use an
explicit max_new_tokens. Use a nonempty amino-acid prefix. Generation quality
and biological functionality have not been evaluated; exported example generation
is a software check, not a biological validation.
Architecture and training
- 32 layers, hidden size 1024, 8 query heads and 2 KV heads, head dimension 128.
- Partial RoPE dimension 64; per-head QK RMSNorm; gated attention.
- SwiGLU intermediate size 2560 with softcap 7; rank-16 sigmoid LoRA GatedNorm.
- Zero-centered RMSNorm gamma; tied input/output embeddings, 64 model vocabulary slots.
- PLE in all 32 layers: dimension 384, n-gram orders 2/3/4, 8 heads per order, prime-table base 32768, convolution width 4, seed 1234.
- Muon plus AdamW; peak learning rates 3e-4, 100M-token warmup, cosine decay to 3e-5.
- Eight RTX 3090 GPUs; shuffled real data with sequence packing and independent fragment attention/position boundaries; ordinary global batch 262,144 prediction tokens. The final batch is shortened to finish at exactly 5B tokens.
training_config.json records the complete settings; export_info.json records
checkpoint provenance. FP32 checkpoint versus HF logits had maximum absolute error
1.15e-5 on the verification prompts; generation with padded batches matched separate
unpadded generation. BF16 autocast probability total variation was at most 0.00733
on those prompts. See verification.json; this is a limited inference parity test.
Dataset token IDs versus model token IDs
The source dataset QingWY/protein-pretraining-data stores ProGen3 token IDs.
Before training, those IDs were converted to ESMC token IDs by amino-acid identity.
The bundled tokenizer matches the actual trained embedding/output rows. For example,
ProGen3 A:8 maps to model A:5, and L:19 maps to L:4.
Do not feed raw ProGen3 IDs directly to this model or replace only its tokenizer.
Raw amino-acid strings work directly with the bundled tokenizer, without adding
synthetic BOS/EOS. The exact mapping is included in progen3_to_esmc_mapping.json.
To convert existing tokenized data, use convert_progen3_to_esmc.py with an explicit
--mapping progen3_to_esmc_mapping.json; see --help and TOKEN_CONVERSION.md.