chandar-lab/AMPLIFY_350M_base

AMPLIFY

AMPLIFY is an efficient, state-of-the-art protein language model pre-trained using masked language modeling on UniRef100, OAS, and SCOP (UR100P). AMPLIFY can generate residue and protein embeddings, suggest mutations, differentiate disordered proteins from non-protein sequences, and much more. AMPLIFY is available in two sizes, 120M and 350M parameters, with the _base models not extended beyond 512 residues (Stage 1). The model architecture and pre-training procedure are detailed below. For more details, please refer to the accompanying paper.

• AMPLIFY_350M
• AMPLIFY_350M_base
• AMPLIFY_120M
• AMPLIFY_120M_base

Model Descritpion

	AMPLIFY 120M	AMPLIFY 350M
hidden-size	640	960
num-hidden-layers	24	32
num-attention-heads	10	15
intermediate-size	2560	3840
max-position-embeddings	2048	2048
vocab-size	27	27
rope-theta	10000	10000
dropout-prob	0	0
embedding-init-range	0.02	0.02
norm-eps	1.0e-05	1.0e-05
hidden-act	swiglu	swiglu
pre-activation-layer-norm	true	true
layer-norm-after-embedding	false	false
layer-norm-before-last-layer	true	true
rms-norm	true	true
ffn-bias	false	false
attn-bias	false	false

Training Descritpion

	Stage 1	Stage 2
dataset	UR100P	UR100P
max-steps	1000000	25000 (120M) or 50000 (350M)
max-length	512	2048
optimizer	adamw	adamw
lr	0.001	0.001
betas	(0.9, 0.95)	(0.9, 0.95)
eps	1.0e-08	1.0e-08
weight-decay	0.01	0.01
scheduler	cosinedecay	none
warmup-steps	1,000	none
final-step	900,000	none
warmup-steps	1,000	none
gradient-clipping	1.0	1.0
tf32	true	true
mixed-precision	bf16	bf16
padding	max-length	max-length
random-truncate	true	true
mask-probability	0.15	0.15
total-batch-size	4096	4096
deepspeed	true	true
zero-stage	3	3

Get Started

from transformers import AutoModel
from transformers import AutoTokenizer
from datasets import load_dataset

# Load AMPLIFY and tokenizer
model = AutoModel.from_pretrained("chandar-lab/AMPLIFY_350M", trust_remote_code=True)
tokenizer = AutoTokenizer.from_pretrained("chandar-lab/AMPLIFY_350M", trust_remote_code=True)

# Move the model to GPU (required due to Flash Attention)
model = model.to("cuda")

# Load the UniProt validation set
dataset = load_dataset("chandar-lab/UR100P", data_dir="UniProt", split="test")

for sample in dataset:
    # Protein
    print("Sample: ", sample["name"], sample["sequence"])

    # Tokenize the protein
    input = tokenizer.encode(sample["sequence"], return_tensors="pt")
    print("Input: ", input)

    # Move to the GPU and make a prediction
    input = input.to("cuda")
    output = model(input)
    print("Output: ", output)

    break

Citations

If you find the models useful in your research, we ask that you cite the paper:

@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},
    elocation-id = {2024.09.23.614603},
    eprint       = {https://www.biorxiv.org/content/early/2024/09/23/2024.09.23.614603.full.pdf}
}