Proto is not affiliated with Meta AI and Biohub. This toolkit is open source and builds on the implementations produced by these organizations. Product names, logos, and trademarks are the property of their respective owners.

` symbol. The model is then trained to predict the original amino acid from the surrounding bidirectional context. No structural, functional, or alignment supervision is used. ESM-2 has since become a de facto sequence representation model for protein engineering. Its direct successor, ESM3 ([Hayes et al., 2025](https://doi.org/10.1126/science.ads0018)), extends the recipe at [EvolutionaryScale](https://www.evolutionaryscale.ai) into a multimodal generative model that jointly handles sequence, structure, and function tracks via discrete diffusion. ESM-2 still remains the lightest and most widely deployed protein language model. Within this toolkit, the 650M checkpoint (`esm2_t33_650M_UR50D`) is a standard quality/speed tradeoff and is the default for every tool. ## Tools

### ESM2 Embeddings (`esm2-embedding`) Runs a single forward pass over ESM-2 to extract contextualized per-residue hidden states. The hidden states are mean-pooled across valid positions to produce a fixed-length sequence descriptor. Per-position 20-way amino-acid logits over the canonical order `ACDEFGHIKLMNPQRSTVWY` are also returned on request. #### API Reference

Source Protein sequence(s) to process. Each must be ≤ 1022 residues (ESM-2's positional-encoding cap); over-length inputs raise `ValueError`.

Source ESM2 weights variant. Sizes range from 8M (320-dim, fastest) to 15B (5120-dim, highest quality). The 650M variant offers a good speed/quality trade-off. Available options: `esm2_t6_8M_UR50D`, `esm2_t12_35M_UR50D`, `esm2_t30_150M_UR50D`, `esm2_t33_650M_UR50D`, `esm2_t36_3B_UR50D`, `esm2_t48_15B_UR50D` Include per-position logits in the output (large; disable to save memory). Transformer layer index for embeddings. `-1` selects the last (top) layer; uses HuggingFace `hidden_states` indexing where `0` is the embedding-layer output and `N` is transformer layer N. Print status messages during model execution. Device to run the model on. Maximum execution time in seconds. `None` waits indefinitely. Random seed. When set, tools run reproducibly up to small GPU float noise (see `BaseToolOutput.approx_equal`), and the seed participates in cache keys. When None, cacheable seed-sensitive tools skip cache until seeded. Number of sequences to process in parallel. Larger batches improve throughput but require more GPU memory.

Source Per-sequence embedding results. Each `SequenceEmbedding` contains: Mean-pooled embedding vector for one sequence. Binary mask indicating valid positions (1) vs padding (0). Optional per-position amino acid logits for one sequence. Optional 2D coordinate from a UMAP projection of all embeddings in the same call. Populated when `n_sequences >= 4`; `None` otherwise (single-point or 2-3-point UMAP is meaningless).

#### Applications The mean-pooled embedding is a standard learned protein representation for downstream supervised tasks like clustering, classification, and regression on protein properties. The same embeddings also power similarity search through cosine similarity on the mean vector. The per-position logits support variant-effect screening by comparing wild-type and mutant log-probabilities at each position. The underlying attention maps are themselves rich enough to recover residue-residue contacts without explicit supervision. #### Usage Tips * **The last transformer layer carries the richest bidirectional context.** `repr_layer` chooses which layer to read for the mean-pooled embedding; the default `-1` selects the last layer and is the standard pick for downstream classification, regression, and variant-effect work. Earlier layers can outperform the top on certain probes (contact prediction is the canonical example). * **Per-position logits are large and slow to materialize.** Enabling `return_logits` adds a `seq_len × 20` float tensor per sequence to the output, dominating wall time on long inputs. Leave it `False` unless you actually need the per-position distribution.

### ESM2 Sampling (`esm2-sample`) Selects positions to mutate via a specifiable masking strategy, replaces them with ``, and resamples from ESM-2's predicted distribution. Two decoding modes are available. The `single_pass` mode fills every masked position in one forward pass with independent draws. The `iterative_refinement` mode instead runs a [MaskGIT](https://arxiv.org/abs/2202.04200)-style multi-round commit loop. Each round of that loop uses a cosine or linear unmask schedule with optional temperature annealing. To target specific positions directly, pre-mask them yourself with `_` in the input string. The tool will then fill exactly those positions and skip the masking strategy entirely. #### API Reference

Source Protein sequence(s) with `_` at positions to sample. Each must be ≤ 1022 residues (ESM-2's positional-encoding cap); over-length inputs raise `ValueError`.

Source Positions to mask before sampling. Temperature for position selection. \< 1.0 is greedy, 1.0 uses scores as-is, > 1.0 is more uniform. Only affects model-based methods. Scoring method for position selection. `"random"`: uniform random, `"entropy"`: highest model uncertainty, `"max-logit"`: lowest model confidence. Available options: `random`, `entropy`, `max-logit` Exact number of positions to mask per sequence. Fraction of designable positions to mask (e.g. 0.15 for \~15%). 1-indexed positions that must NOT be masked. Applied uniformly to all sequences. ESM2 weights variant. Available options: `esm2_t6_8M_UR50D`, `esm2_t12_35M_UR50D`, `esm2_t30_150M_UR50D`, `esm2_t33_650M_UR50D`, `esm2_t36_3B_UR50D`, `esm2_t48_15B_UR50D` "single\_pass" fills every mask in one forward; "iterative\_refinement" runs an iterative MaskGIT-style loop driven by the five settings below. Available options: `single_pass`, `iterative_refinement` Softmax temperature. Nucleus threshold (iterative only). Refinement steps (iterative only). Unmask schedule (iterative only). Available options: `cosine`, `linear` Per-round commit selection (iterative only). Available options: `random`, `entropy` Anneal toward 0 across rounds (iterative only). Include per-position logits. Verbosity level (0=quiet, 1=info, 2=debug, 3=raw subprocess stderr). `True` is coerced to `1` and `False` to `0`. Device to run on. Maximum execution time in seconds. `None` waits indefinitely. Random seed. When set, tools run reproducibly up to small GPU float noise (see `BaseToolOutput.approx_equal`), and the seed participates in cache keys. When None, cacheable seed-sensitive tools skip cache until seeded. Sequences per GPU forward pass.

Source One entry per input sequence, in input order. Each holds the sampled sequence (the input with masked positions replaced by model-predicted alternatives) and its optional per-position logits. The sampled or restored protein sequence. Per-position amino acid logits for this sequence, shape `[seq_len, 20]`. Present only when the tool's config sets `return_logits=True`.

#### Applications This tool drives guided point mutation, variant generation, and infilling at designable sites for protein engineering work. Resampling masked positions from a protein language model is the core operation behind directed-evolution proposals and antibody affinity maturation, which was demonstrated at experimental scale in [Hie et al., 2024](https://www.nature.com/articles/s41587-023-01763-2). It is also the inner loop of MaskGIT-style iterative refinement schemes adapted from image generation ([Chang et al., 2022](https://arxiv.org/abs/2202.04200)) for biological sequences. #### Usage Tips * **`iterative_refinement` produces more coherent joint samples than `single_pass`.** It is a multi-round MaskGIT-style commit loop (each round uses a cosine or linear unmask schedule) and is roughly `num_steps×` slower than the one-shot `single_pass` mode. Default to it whenever you mask more than a handful of sites. * **`masking_strategy` controls which positions get masked before sampling.** See the [masking strategy README](https://github.com/evo-design/proto-tools/blob/main/proto_tools/transforms/masking/README.md) for the available selection methods and tuning parameters. As an alternative to passing a strategy, pre-mask exact positions yourself with `_` directly in the input string and the masking strategy is skipped entirely. * **`temperature` scales the per-position logits before sampling.** Values of 0.5 to 0.7 yield conservative mutations close to the input; values above 1.0 broaden exploration of the model's distribution. * **Long-range coherence is weak.** ESM-2 has no global coherence beyond its local context window, so very long-range dependencies between distant residues are not well captured even in iterative mode. * **ESM-2 was trained as a masked language model, not with a generative objective.** Resampling masked positions works for local edits, but the model was optimized for representation rather than de novo generation. For generative workloads (large-scale infilling, sequence design), [ESM3](https://bio-pro.mintlify.app/tools/masked-models/esm3) adds an explicit generative training objective and is the better fit.

### ESM2 Scoring (`esm2-score`) Computes the masked-language-model pseudo-perplexity for each input sequence. Each position is masked individually, and the model's log-probability of the true amino acid under bidirectional context is recorded. The per-position scores are then aggregated into per-sequence log-likelihood, average log-likelihood, and perplexity metrics. #### API Reference

Source Protein sequence(s) to score. Each must be ≤ 1022 residues (ESM-2's positional-encoding cap); over-length inputs raise `ValueError`.

Source ESM2 weights variant. Available options: `esm2_t6_8M_UR50D`, `esm2_t12_35M_UR50D`, `esm2_t30_150M_UR50D`, `esm2_t33_650M_UR50D`, `esm2_t36_3B_UR50D`, `esm2_t48_15B_UR50D` Print status messages during scoring. Device to run the model on. Maximum execution time in seconds. `None` waits indefinitely. Random seed. When set, tools run reproducibly up to small GPU float noise (see `BaseToolOutput.approx_equal`), and the seed participates in cache keys. When None, cacheable seed-sensitive tools skip cache until seeded. Masked variants per forward pass, pooled across all input sequences. Larger batches improve throughput but use more memory. Include per-position logits in the output (large; disable to save memory).

Source List of scoring outputs, one per input sequence. Each entry is a `Metrics` subclass with scalar metrics (accessed via `score.perplexity` or `score["perplexity"]`) plus declared `logits` / `vocab` fields that carry raw model outputs when requested. Per-position logits array `(seq_len, vocab_size)`. `None` unless `return_logits=True`. Token ordering for `logits`. Name of the metric that best summarizes the result overall (e.g. `"avg_plddt"` for AlphaFold2). Used by downstream UI and reporting to pick a headline value. Concrete Metrics subclass tag; enables typed reconstruction after a serialization round-trip. **Metrics** (one set per `scores` item) | Metric | Type | Range | Availability | | -------------------- | ----- | ----- | ------------ | | `log_likelihood` | float | ≤ 0.0 | always | | `avg_log_likelihood` | float | ≤ 0.0 | always | | `perplexity` | float | ≥ 1.0 | always |

#### Applications ESM2 pseudo-perplexity is a standard fitness proxy when ranking variants, filtering generated sequences for naturalness, or comparing engineered constructs against wild type. The same masked log-likelihood difference between wild-type and mutant residues is a canonical zero-shot baseline for variant-effect prediction. #### Usage Tips * **Pseudo-perplexity is a relative score, not an absolute fitness.** It is measured against ESM-2's training distribution, which is UniRef50 (the natural proteins it saw during pretraining), which can bias it to proteins that are more heavily represented. The metric is also sensitive to length, so it is most useful for comparing closely related sequences of similar length. * **Ambiguous residues are excluded.** Perplexity is computed only over the 20 canonical amino acids; `X`, `B`, `Z`, and similar are dropped from both the log-likelihood sum and the position count.

### ESM2 Gradient (`esm2-gradient`) Computes the gradient of the mean masked negative log-likelihood with respect to a relaxed `(L, 20)` input distribution over the canonical amino-acid order `ACDEFGHIKLMNPQRSTVWY`. The ESM-2 weights are kept frozen throughout. The relaxed distribution is mixed against ESM-2's per-residue token embeddings to form a soft input. Each amino-acid position is then masked in turn, and a per-chunk backward pass accumulates the gradient. An optional Straight-Through Estimator runs the forward on hard one-hot tokens while still routing gradients through the soft probabilities. #### API Reference

Source Relaxed protein sequence state with shape `(L, 20)` in canonical amino-acid order `ACDEFGHIKLMNPQRSTVWY`. `L` must be ≤ 1022 (ESM-2's positional-encoding cap); over-length inputs raise `ValueError`. Optional softmax temperature. When set, applies `softmax(input / temperature)` before computing the gradient. When `None` (default), the input is used as-is.

Source ESM2 weights variant. Available options: `esm2_t6_8M_UR50D`, `esm2_t12_35M_UR50D`, `esm2_t30_150M_UR50D`, `esm2_t33_650M_UR50D`, `esm2_t36_3B_UR50D`, `esm2_t48_15B_UR50D` Straight-Through Estimator: hard one-hot in the forward pass with gradients flowing through soft probabilities. When `False`, uses soft blended embeddings directly. Run backward pass and return gradient. Set `False` for forward-only log-likelihood scoring. AA positions per forward pass for batched PLL. `None` selects the backend default. Verbosity level (0=quiet, 1=info, 2=debug, 3=raw subprocess stderr). `True` is coerced to `1` and `False` to `0`. Device to run the model on. Maximum execution time in seconds. `None` waits indefinitely. Random seed. When set, tools run reproducibly up to small GPU float noise (see `BaseToolOutput.approx_equal`), and the seed participates in cache keys. When None, cacheable seed-sensitive tools skip cache until seeded.

Source Gradient w\.r.t. input logits, or `None` when `compute_gradient=False`. Mean negative log-likelihood over AA positions. Log-likelihood, perplexity, sequence length, and objective details. Amino-acid column ordering for the input logits.

#### Applications This tool exposes ESM-2 as a differentiable, structure-free protein-language-model loss for use inside MCMC, gradient descent, or any other optimization loop over relaxed protein sequences. It is most often used as a naturalness prior in continuous design pipelines, including latent Bayesian optimization frameworks and discrete walk-jump sampling approaches for de novo protein design. #### Usage Tips * **`temperature` controls how the raw input is converted into a distribution.** With a value set, the tool applies `softmax(logits / T)` before the forward pass; leave it `None` (the default) if the input already sums to 1 per position. * **`use_ste` enables the Straight-Through Estimator.** The forward then runs on hard one-hot tokens while gradients still route through the soft probabilities, giving stronger guidance toward discrete sequences. Leave it off for smooth optimization over the relaxed simplex. * **`compute_gradient` toggles whether the backward pass runs.** When set to `False`, the `gradient` field is `None`, but `loss` and `metrics` (log-likelihood, perplexity, and so on) are still populated. Useful for ranking MCMC proposals without paying the backward cost.

## Toolkit Notes These apply to every ESM-2 tool in this toolkit (`esm2-embedding`, `esm2-sample`, `esm2-score`, `esm2-gradient`). * **Different ESM-2 checkpoints produce different embedding sizes.** Downstream tasks built on one checkpoint will not transfer to another without re-fitting; pick one and stick with it for an analysis. * **Smaller checkpoints run faster.** The 150M and 35M variants are significantly faster than the 650M default, with drops in representation quality. * **Max sequence length is 1022 residues.** ESM-2's positional encoding caps inputs at 1022 residues, and will raise `ValueError` on longer inputs rather than truncating. * **`batch_size` controls memory usage across the toolkit.** Lower it if you OOM; raise it for short-sequence throughput. One nuance: for `esm2-score`, `batch_size` counts masked variants pooled across all input sequences rather than sequences themselves (each input contributes `L` masked variants). **Example notebook:** See the [full working example](https://github.com/evo-design/proto-tools/blob/47e34afa5ea240a3b406e323dc38aa5dc85f223e/proto_tools/tools/masked_models/esm2/examples/example.ipynb) for a copy-paste-ready walkthrough. ## Infrastructure Guides The following guides cover how to run tools efficiently and at scale. Keep a tool's model warm across calls instead of reloading it every invocation. How GPUs are allocated to tools and how to target specific devices. Fan a batch of inputs out across multiple GPUs.