IvanHU/esmc-ar-371m-5b

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 .