ESMC-AR 371M — 5B-token baseline
A 371,368,960-parameter, dense autoregressive protein language model trained from scratch on 5,000,000,000 prediction tokens from QingWY/protein-pretraining-data. This is the completed baseline at optimizer step 19,074, not the PLE experiment. It adapts the ESMC architecture to a causal next-token objective; it does not use the pretrained ESMC-300M weights or the original masked objective.
Load and generate
Log in with a Hugging Face account that has repository access if this repository is private. Requires PyTorch, Transformers, safetensors and tokenizers. Tested with PyTorch2.5.1 and Transformers4.51.0; no training repository or FlashAttention installation required.
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
repo = "IvanHU/esmc-ar-371m-5b"
tokenizer = AutoTokenizer.from_pretrained(repo)
model = AutoModelForCausalLM.from_pretrained(
repo, trust_remote_code=True, torch_dtype=torch.float32,
).to("cuda").eval()
batch = tokenizer(["MALWMRLLPLL", "MKWVTFISLLFLFSSAYS"],
padding=True, return_tensors="pt").to(model.device)
with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16):
ids = model.generate(**batch, max_new_tokens=64, do_sample=True,
temperature=0.8, top_p=0.95)
print(tokenizer.batch_decode(ids, skip_special_tokens=True))
Weights are stored in FP32 safetensors (about1.49GB) to preserve the training
checkpoint. For lower-memory inference, load with torch_dtype=torch.bfloat16.
For CPU, keep FP32 and omit CUDA/autocast. See INFERENCE.md.
Architecture and training
| Setting | Value |
|---|---|
| Layers / hidden size | 32 /1024 |
| Attention | GQA:8 Q heads,2 KV heads,128 dimensions/head |
| Positional embedding | Partial RoPE,64 rotary dimensions,theta10000 |
| MLP | SwiGLU, intermediate2560; smooth product softcap7 |
| Normalization | Low-rank gated zero-centered RMSNorm, gate rank16 |
| QK normalization | Per-head RMSNorm |
| Attention output | Sigmoid gate |
| Vocabulary | 64 slots, tied input/output embeddings |
| PLE / short conv / multi-stream GatedResidual | Disabled |
| Final logits clamp | Disabled |
| Maximum training fragment length | 2048 |
| Global batch | 262,144 prediction tokens (final step shortened to exact budget) |
| Optimizers | Muon (5 PolarExpress NS steps, align-RMS0.2) + AdamW |
| AdamW betas / epsilon | (0.95,0.95) /1e-15 |
| LR | 3e-4,100M-token warmup, cosine to3e-5 |
| Weight decay | 0.1, standard decoupled decay |
| Data order | Protein-row shuffle across shards, rank-disjoint sampling |
| Packing | Enabled; independent causal attention and positions per fragment |
| Hardware | 8×RTX3090 |
Tokenizer IDs use the ESMC-compatible mapping, not raw ProGen3 IDs. Automatic BOS/EOS insertion is disabled. Training fragments contain the20 standard amino acids; generation defaults suppress unused slots and special tokens. Provide a nonempty amino-acid prefix and a bounded generation length: biological end-of-protein termination was not trained.
Verification and scope
verification.json records comparison against the final training checkpoint, left-padded batching, tied embeddings, generation and KV cache validation. FP32 maximum logit discrepancy was1.53e-5; the largest measured BF16-autocast probability total-variation discrepancy was0.00324. These checks establish implementation consistency, not biological quality. No biological functionality benchmark or experimental validation is claimed.
The model has not been evaluated for contexts beyond2048 positions. Training implementation: https://github.com/huyiwen/bio . Optional future architecture extensions: https://github.com/huyiwen/bio-next .