A 350M-parameter protein language model pre-trained on UR100P using masked language modeling. Trained for 1M steps (~2T tokens) with context length 2,048.

This model was trained using the AMPLIFY training codebase. The original models and code were released under chandar-lab/AMPLIFY. See also flair-bio/amplify-350m.

Property	Value
Architecture	BERT-style encoder (RoPE, SwiGLU, RMSNorm)
Parameters	350M
Training tokens	~2T
Vocabulary size	32 (amino acid alphabet + special tokens)
Context length	2,048
Training steps	1,000,000
License	Apache 2.0

Quick Start

python

from transformers import AutoModelForMaskedLM, AutoTokenizer

model = AutoModelForMaskedLM.from_pretrained("flair-bio/amplify-350m", trust_remote_code=True)
tokenizer = AutoTokenizer.from_pretrained("flair-bio/amplify-350m", trust_remote_code=True)
model.eval()

How to Use

Extract Embeddings

python

import torch
from transformers import AutoModel, AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("flair-bio/amplify-350m", trust_remote_code=True)
model = AutoModel.from_pretrained("flair-bio/amplify-350m", trust_remote_code=True)

sequences = ["MKTAYIAK", "MVLSPADKTNVK"]
inputs = tokenizer(sequences, return_tensors="pt", padding=True, truncation=True, max_length=2048)

with torch.no_grad():
    outputs = model(**inputs)

embeddings = outputs.last_hidden_state  # [batch, seq_len, 960]

Masked Language Modeling

python

from transformers import AutoTokenizer, AutoModelForMaskedLM
import torch

tokenizer = AutoTokenizer.from_pretrained("flair-bio/amplify-350m", trust_remote_code=True)
model = AutoModelForMaskedLM.from_pretrained("flair-bio/amplify-350m", trust_remote_code=True)

sequence = "MKTAY<mask>AKQRQISFVK"
inputs = tokenizer(sequence, return_tensors="pt")

with torch.no_grad():
    logits = model(**inputs).logits

mask_idx = (inputs["input_ids"] == tokenizer.mask_token_id).nonzero(as_tuple=True)[1]
predicted = tokenizer.decode(logits[0, mask_idx].argmax(dim=-1))
print(predicted)

Model Description

Architecture

AMPLIFY 350M is a BERT-style transformer encoder with 32 layers, 960-dimensional hidden states, and 15 attention heads. It uses rotary positional embeddings (RoPE), SwiGLU feed-forward blocks, and RMSNorm. Tokenization is at the amino acid level with a vocabulary of 32 tokens.

Config	Layers	Hidden dim	Heads	FFN dim	Context
AMPLIFY 350M	32	960	15	2,560	512 → 2,048

Intended Use

This model is intended for extracting per-residue or per-sequence representations for downstream tasks, zero-shot variant effect prediction via pseudo-log-likelihood scoring, and fine-tuning on protein fitness, stability, binding, or functional annotation tasks.

Training

Data

Pre-trained on UR100P (chandar-lab/UR100P), a deduplicated union of UniRef100, OAS, and SCOPe.

Training Procedure

Hyperparameter	Value
Hardware	8× H100 80GB
Optimizer	AdamW
Learning rate	1e-3 (peak)
LR schedule	Linear warmup + cosine decay
Batch size (tokens)	~2M per step
Masking rate	15%
Training objective	Masked language modeling
Precision	BF16
Framework	PyTorch + HuggingFace Transformers

Training logs are available on Weights & Biases.

Citation

If you use this model in your work, please cite:

bibtex

@article{Fournier2024.09.23.614603,
  title        = {Protein Language Models: Is Scaling Necessary?},
  author       = {Fournier, Quentin and Vernon, Robert M. and van der Sloot, Almer and Schulz, Benjamin and Chandar, Sarath and Langmead, Christopher James},
  year         = {2024},
  journal      = {bioRxiv},
  publisher    = {Cold Spring Harbor Laboratory},
  doi          = {10.1101/2024.09.23.614603},
  url          = {https://www.biorxiv.org/content/early/2024/09/23/2024.09.23.614603}
}