| --- |
| license: mit |
| library_name: transformers |
| pipeline_tag: fill-mask |
| tags: [protein, protein-language-model, recAMPLIFY, recursive-transformer] |
| --- |
| |
| # recAMPLIFY |
|
|
| Pretrained models from *recAMPLIFY: Large-Scale Pretraining of Recursive Protein Language Models*: the 11 models of the paper's main table. recAMPLIFY is a masked |
| protein language model whose weight-tied recurrent core refines a residue-aligned latent state over `L` |
| cycles, so compute per sequence can grow at test time without adding parameters; Gated Mixture |
| Recurrence (GMR) makes that depth adaptive. |
|
|
| Each model lives in its own subfolder; the model code (`amplify.py`, `recamplify.py`, `tokenizer.py`) is |
| shared at the repository root and loaded with `trust_remote_code=True`. |
|
|
| | subfolder | paper row | steps | params | L | description | |
| |---|---|---|---|---|---| |
| | `amplify-120m-1m` | AMPLIFY-120M (1M) | 1M | 118M | – | AMPLIFY-120M baseline retrained with the AMPLIFY recipe in this codebase | |
| | `amplify-350m-1m` | AMPLIFY-350M (1M) | 1M | 354M | – | AMPLIFY-350M baseline retrained with the AMPLIFY recipe in this codebase | |
| | `recamplify-1m` | fixed-L recAMPLIFY (1M) | 1M | 155M | 4 | fixed-depth recAMPLIFY, (p,r,c) = (2,10,2), trained and evaluated with L = 4 | |
| | `amplify-120m-100k` | AMPLIFY-120M (100K) | 100K | 118M | – | AMPLIFY-120M baseline at 100K steps (retrained in this codebase) | |
| | `amplify-350m-100k` | AMPLIFY-350M (100K) | 100K | 354M | – | AMPLIFY-350M baseline at 100K steps (retrained in this codebase) | |
| | `recamplify-vary-l-100k` | vary-L recAMPLIFY | 100K | 133M | 6 | recAMPLIFY trained with a variable number of cycles L ~ U{1,...,6} per minibatch (evaluated at L = 6) | |
| | `recamplify-fixed-l-100k` | fixed-L recAMPLIFY | 100K | 133M | 6 | fixed-depth recAMPLIFY, (p,r,c) = (2,8,2), trained with L = 6 | |
| | `gmr-halt-only` | + halt-only GMR | 100K | 133M | 6 | halt-only GMR: a halt head alone allocates depth (PonderNet-style halting) | |
| | `gmr-scalar` | + scalar gated GMR | 100K | 133M | 6 | GMR with a scalar residual gate (one survival channel per sequence) | |
| | `gmr-rank16` | + rank-16 gated GMR | 100K | 133M | 6 | GMR with a rank-16 residual gate (16 survival channels mapped to features) | |
| | `gmr-perdim` | + per-dim gated GMR | 100K | 134M | 6 | GMR with a per-dimension residual gate (one survival channel per feature) | |
|
|
| The AMPLIFY baselines are retrained in this codebase and are **not** the official |
| `chandar-lab/AMPLIFY_120M` / `chandar-lab/AMPLIFY_350M` weights. |
|
|
| ## Usage |
|
|
| ```python |
| import torch |
| from transformers import AutoModelForMaskedLM, AutoTokenizer |
| |
| kw = dict(subfolder="gmr-perdim", trust_remote_code=True) |
| tokenizer = AutoTokenizer.from_pretrained("chandar-lab/recamplify", **kw) |
| model = AutoModelForMaskedLM.from_pretrained("chandar-lab/recamplify", **kw).eval() |
| |
| inputs = tokenizer(["MSKGEELFTGVVPILVELDGDVNGHKFSVSGEGEGDATYGKLTLKFICTTGKLPVPWPTLVTT"], return_tensors="pt") |
| with torch.no_grad(): |
| logits = model(**inputs).logits |
| |
| # Test-time depth: run more (or fewer) recursive cycles than in training. |
| model.set_num_cycles(9) |
| ``` |
|
|
| Replace `subfolder` with any name from the table. Padded batches run on CPU or GPU with PyTorch SDPA; |
| packed (FlashAttention varlen) batches need `flash-attn`, CUDA and half precision. |
|
|
| Training and evaluation code: <https://github.com/chandar-lab/AMPLIFY/tree/recamplify> |
|
|
| ## Citation |
|
|
| ```bibtex |
| @misc{recamplify, |
| title = {recAMPLIFY: Large-Scale Pretraining of Recursive Protein Language Models}, |
| author = {Jha, Saurav and Liao, Qianli and Vernon, Robert McCoy and Chandar, Sarath and Fournier, Quentin}, |
| year = {2026}, |
| } |
| ``` |
|
|
