ESM2 Model

A protein language model surrogate using ESM-2 as the backbone. Accepts amino acid sequences as inputs and supports three operating modes, selected via ESM2TrainConfig.mode together with freeze_backbone:

• mode=’linear_head’ (default): a trainable linear head is stacked on top of pooled sequence embeddings. With freeze_backbone=True (default) only the head is trained on top of the frozen backbone. With freeze_backbone=False the backbone is fine-tuned jointly with the head, using backbone_learning_rate (typically lower than learning_rate) for the backbone parameters — usually slower to train but often more accurate. The head is configured via loss_fn and output_dim:

  • loss_fn='mse': mean-squared-error regression. Set output_dim=1. predict() returns raw scalar values.

  • loss_fn='cross_entropy': multi-class cross-entropy classification. Set output_dim=N for N classes. Labels must be integers in [0, N); float labels are truncated to int with a warning. predict() returns the argmax class index as a float.

  Call train() to fit the head on labelled data (loss_fn must be set). Sequence embeddings can also be extracted via embed for use with downstream models.

• mode=’likelihoods’ with freeze_backbone=True: zero-shot scoring. predict() returns per-sequence pseudo-log-likelihood (PLL) scores — each non-special token is masked one at a time and the log-probability of the correct residue at that position is accumulated, and the final score is the mean log-probability across all non-special positions. Nothing is trained, so train() raises NotImplementedError and loss_fn must be left None.

• mode=’likelihoods’ with freeze_backbone=False: MLM fine-tuning. train() fine-tunes the full ESM-2 backbone via masked language modelling on the input sequences; loss_fn must be 'mlm'. Masking is controlled by mask_probability (fraction of eligible tokens masked per sequence) and mask_splitting (the (p_mask, p_random, p_unchanged) 3-way replacement probabilities, which must sum to 1.0). Labels in train_data are ignored, but LabelledCandidates still requires them — pass placeholder values. After fine-tuning, predict() returns updated PLL scores. Training reports perplexity and token_accuracy over masked positions; validation masking is seeded from ESM2ModelConfig.seed so metrics are comparable across epochs.

Key properties:

• Input: Amino acid sequences stored in Candidate.data

• Output: Mean-only scalar predictions (no uncertainty estimates)

• Backbone: Configurable ESM-2 checkpoint — specify via ESM2ModelConfig(model_id=...), e.g. facebook/esm2_t6_8M_UR50D

• Sampling: Not supported — raises NotImplementedError

• Tokenisation: featurise() converts sequences to input_ids and attention_mask tensors for the ESM-2 tokeniser. It does not produce embeddings — use embed for that.

Note

This model requires the optional esm2 dependency. Install it with the [esm2] extra:

pip install "alf-tools[esm2]"

class alf_tools.models.esm2.ESM2Model(name, model_config, train_config=None, device=None)

Bases: BaseModel

ESM-2 protein language model wrapper.

Loads a pre-trained ESM-2 checkpoint from HuggingFace and exposes it as a BaseModel. predict() returns per-sequence masked-marginal scores (mode=’likelihoods’) or passes embeddings through a trainable linear head (mode=’linear_head’). embed() returns per-sequence embeddings.

embed(candidate_points)

Compute sequence embeddings using the configured pooling strategy.

Parameters:

candidate_points (list[Candidate]) – List of candidates to embed.

Return type:

ndarray

Returns:

Numpy array of shape (n_candidates, hidden_dim) for mean or cls pooling, or (n_candidates, seq_len, hidden_dim) for last_hidden_state pooling. Returns shape (0, hidden_dim) for mean/cls pooling, or (0, 0, hidden_dim) for last_hidden_state pooling, if candidate_points is empty.

Note

All sequences are tokenized in one pass before batching the forward pass. batch_size_inference controls the embed() forward pass batch size but has no effect on zero-shot PLL predict(). For very large candidate lists, consider chunking externally.

For last_hidden_state pooling, the returned array has shape (n_candidates, max_padded_seq_len, hidden_dim). Positions beyond each sequence’s EOS token are padding and have non-zero values. Use the attention_mask from featurise() to identify valid positions.

featurise(inputs)

Tokenize sequences into input tensors for the ESM-2 model.

Parameters:

inputs (LabelledCandidates | list[Candidate]) – Either LabelledCandidates or a list of Candidates to featurise.

Return type:

dict[str, Tensor]

Returns:

Dictionary with keys ‘input_ids’ and ‘attention_mask’ as tensors.

Raises:

ValueError – If the input is not LabelledCandidates or list of Candidates.

get_epoch_metrics()

Return per-epoch metrics from the most recent train() call.

Return type:

list[SurrogateEpochMetrics]

Returns:

List of SurrogateEpochMetrics, one per logged epoch.

get_training_summary_metrics()

Return summary metrics from the most recent train() call.

Return type:

dict[str, float | int | number]

Returns:

Dictionary of training metrics, e.g. final train and validation losses.

predict(candidate_points)

Compute predictions for the given candidates.

When mode=’likelihoods’: computes pseudo-log-likelihood (PLL) by masking one residue at a time and recording log P(token_i | all other tokens). Returns the mean PLL over residue positions per sequence (higher = more probable). Sequences with ≤_PLL_BATCH_THRESHOLD residues are scored in a single batched forward pass; longer sequences are scored position-by-position to bound memory usage. When mode=’linear_head’: embeds sequences through the backbone (frozen or fine-tuned) and passes them through the linear head. Returns regression values (output_dim=1) or argmax class indices (output_dim>1).

Parameters:

candidate_points (list[Candidate]) – List of candidates to score. Must be non-empty.

Return type:

Predictions

Returns:

Predictions whose means are shape (n_candidates,). variances is always None.

Raises:

• ValueError – If candidate_points is empty.

• ValueError – If any sequence has no scoreable residue positions.

• RuntimeError – If mode=’linear_head’ but the head is uninitialised (should not happen if __init__ ran without error).

Note

All sequences are tokenised in one pass before batching. For very large candidate lists, consider calling predict() on smaller chunks externally.

sample(*args, **kwargs)

Not implemented for ESM-2.

Raises:

NotImplementedError – Always, as sampling is not supported.

Return type:

list[Candidate]

train(train_data, val_data=None)

Fine-tune ESM-2 using the configured mode and loss function.

When mode=’likelihoods’ and freeze_backbone=True, train() raises NotImplementedError (zero-shot PLL only; nothing to train). When mode=’likelihoods’ and freeze_backbone=False, fine-tunes the full ESM-2 backbone via masked language modelling. When mode=’linear_head’, trains the linear head (with frozen backbone) or fine-tunes the backbone jointly with the head when freeze_backbone=False.

Parameters:

• train_data (LabelledCandidates) – Training data containing sequences and (for linear_head) labels. MLM fine-tuning ignores labels, but LabelledCandidates still requires them — pass placeholder values.

• val_data (LabelledCandidates | None) – Optional validation data for monitoring training loss.

Raises:

• NotImplementedError – If mode=’likelihoods’ and freeze_backbone=True.

• ValueError – If mode=’linear_head’ and loss_fn is None.

• AssertionError – If optimizer_type is invalid (unreachable if __post_init__ ran).

• RuntimeError – If mode=’linear_head’ but head is uninitialised.

Return type:

None

Configuration

class alf_tools.models.esm_utils.config.ESM2ModelConfig(model_id, pooling='mean', repr_layer=-1, max_length=None, seed=42)

Bases: object

Configuration for ESM-2 model architecture.

Parameters:

• model_id (str) – HuggingFace model identifier, e.g. ‘facebook/esm2_t6_8M_UR50D’.

• pooling (Literal['mean', 'cls', 'last_hidden_state']) – Strategy for reducing per-token hidden states to a sequence embedding. ‘mean’ averages over all non-padding positions (CLS and EOS included). ‘cls’ uses only the first [CLS] token representation. ‘last_hidden_state’ returns the full (seq_len, hidden_dim) tensor per sequence.

• repr_layer (int) – Transformer layer index to extract embeddings from. -1 = final layer.

• max_length (int | None) – Maximum tokenisation length. Defaults to the tokeniser’s model_max_length.

• seed (int) – Random seed for reproducible linear head initialisation.

max_length: int | None = None

model_id: str

pooling: Literal['mean', 'cls', 'last_hidden_state'] = 'mean'

repr_layer: int = -1

seed: int = 42

class alf_tools.models.esm_utils.config.ESM2TrainConfig(learning_rate=0.0001, log_frequency=1, normalise_inputs_strategy=None, standardise_outputs=False, label_dtype=None, mode='linear_head', freeze_backbone=True, loss_fn=None, output_dim=1, backbone_learning_rate=1e-05, mask_probability=0.15, mask_splitting=(0.8, 0.1, 0.1), optimizer_type='adamw', batch_size=8, batch_size_inference=None, num_epochs=10, max_grad_norm=None)

Bases: BaseTrainConfig

Configuration for ESM-2 training.

Set mode first — it determines the architecture and which other fields are active:

• ‘linear_head’: trainable linear head on top of pooled embeddings (supervised). freeze_backbone=True -> backbone frozen; only the head is trained. freeze_backbone=False -> backbone fine-tuned jointly with the head, using

  backbone_learning_rate for the backbone parameters.

  Active fields: loss_fn (‘mse’/’cross_entropy’), output_dim, and (when unfrozen) backbone_learning_rate.

• ‘likelihoods’: Base ESM-2 backbone, no linear head. freeze_backbone=True -> zero-shot PLL only; train() unavailable. freeze_backbone=False -> MLM fine-tuning; loss_fn must be ‘mlm’. Active fields (when unfrozen): loss_fn (‘mlm’), mask_probability, mask_splitting.

Parameters:

• mode (Literal['linear_head', 'likelihoods']) – Architecture mode. ‘linear_head’ adds a trainable head on top of the backbone. ‘likelihoods’ uses the base backbone directly; predict() always returns PLL scores.

• freeze_backbone (bool) – Whether to freeze ESM-2 backbone parameters. For ‘linear_head’: True = train head only; False = fine-tune backbone + head. For ‘likelihoods’: True = zero-shot PLL only; False = MLM fine-tuning.

• loss_fn (Optional[Literal['mse', 'cross_entropy', 'mlm']]) – Training objective. None = no training intended. ‘mse’ and ‘cross_entropy’ are only for mode=’linear_head’. ‘mlm’ is only for mode=’likelihoods’ + freeze_backbone=False.

• output_dim (int) – Output dimension of the linear head. Only for mode=’linear_head’.

• backbone_learning_rate (float) – Learning rate applied to the ESM-2 backbone parameters when mode=’linear_head’ + freeze_backbone=False. The linear head still uses learning_rate. Has no effect when the backbone is frozen.

• mask_probability (float) – Fraction of eligible tokens to mask per sequence. Only active when mode=’likelihoods’ + freeze_backbone=False.

• mask_splitting (tuple[float, float, float]) – (p_mask, p_random, p_unchanged) 3-way replacement probabilities. Must sum to 1.0. Only active when mode=’likelihoods’ + freeze_backbone=False.

• learning_rate (float) – Learning rate for the optimizer.

• optimizer_type (Literal['adam', 'adamw']) – Which optimizer to use (‘adam’ or ‘adamw’).

• batch_size (int) – Batch size for training.

• batch_size_inference (int | None) – Batch size for embed() and linear-head predict(). None defaults to batch_size.

• num_epochs (int) – Number of epochs to train for.

• log_frequency (int) – Record epoch metrics every N epochs.

• max_grad_norm (float | None) – Maximum norm for gradient clipping. None disables clipping.

backbone_learning_rate: float = 1e-05

batch_size: int = 8

batch_size_inference: int | None = None

freeze_backbone: bool = True

learning_rate: float = 0.0001

log_frequency: int = 1

loss_fn: Literal['mse', 'cross_entropy', 'mlm'] | None = None

mask_probability: float = 0.15

mask_splitting: tuple[float, float, float] = (0.8, 0.1, 0.1)

max_grad_norm: float | None = None

mode: Literal['linear_head', 'likelihoods'] = 'linear_head'

num_epochs: int = 10

optimizer_type: Literal['adam', 'adamw'] = 'adamw'

output_dim: int = 1

Loss and scoring

alf_tools.models.esm_utils.loss.compute_supervised_loss(preds, targets, loss_fn)

Compute the linear-head loss between predictions and targets.

Parameters:

• preds (Tensor) – Linear-head outputs of shape (batch, output_dim).

• targets (Tensor) – Target values of shape (batch,).

• loss_fn (str | None) – Either ‘mse’ (regression) or ‘cross_entropy’ (classification).

Return type:

Tensor

Returns:

Scalar loss tensor.

Raises:

AssertionError – If loss_fn is not ‘mse’ or ‘cross_entropy’ (should be unreachable given ESM2TrainConfig.__post_init__ validation).

alf_tools.models.esm_utils.loss.mask_tokens(input_ids, tokeniser, mask_probability, mask_splitting, generator=None)

Apply random token masking for MLM.

Non-masked positions in labels are set to -100 so CrossEntropyLoss ignores them. Special tokens (cls, eos, pad, unk) are never masked.

Parameters:

• input_ids (Tensor) – Token IDs of shape (batch, seq_len).

• tokeniser (Any) – A HuggingFace tokeniser exposing mask_token_id and vocab_size.

• mask_probability (float) – Fraction of eligible tokens to mask per sequence.

• mask_splitting (tuple[float, float, float]) – (p_mask, p_random, p_unchanged) 3-way replacement probabilities.

• generator (Generator | None) – Optional torch.Generator for reproducible masking. Pass a seeded generator (e.g. during validation) to keep the masking fixed across calls; None uses the global RNG.

Return type:

tuple[Tensor, Tensor]

Returns:

Tuple of (masked_input_ids, labels), both of shape (batch, seq_len).

Raises:

ValueError – If the tokeniser does not have a mask token.

alf_tools.models.esm_utils.loss.special_tokens_mask(input_ids, tokeniser)

Return a boolean tensor True at positions occupied by CLS, EOS, PAD, or UNK tokens.

Parameters:

• input_ids (Tensor) – Token IDs of any shape.

• tokeniser (Any) – A HuggingFace tokeniser exposing the special-token id attributes.

Return type:

Tensor

Returns:

A boolean tensor matching input_ids shape, True at special-token positions.

alf_tools.models.esm_utils.scoring_function.compute_pll(esm_model, tokeniser, device, all_input_ids, all_attention_mask, batch_threshold)

Compute zero-shot pseudo-log-likelihood (PLL) scores for sequences.

Scores each sequence by masking one residue at a time and accumulating the log-probability the model assigns to the correct residue at that position. The per-sequence score is the mean (average) log-probability across all non-special (i.e. amino-acid) positions. Two execution modes trade off memory and speed:

• If the number of scoreable residues is ≤ batch_threshold, residues are masked

  in a single batched forward pass (one masked position per batch row) to leverage GPU parallelism. This creates a forward pass of shape (n_residues × padded_seq_len), which can spike GPU memory for sequences near the threshold on large models.

• For longer sequences, positions are masked and scored one-at-a-time to avoid

  excessive memory usage.

Parameters:

• esm_model (Any) – A HuggingFace masked-LM model returning .logits.

• tokeniser (Any) – A HuggingFace tokeniser exposing mask_token_id.

• device (str | device) – Device to run the forward passes on.

• all_input_ids (Tensor) – Tensor of shape (n_candidates, seq_len) with token IDs.

• all_attention_mask (Tensor) – Tensor of shape (n_candidates, seq_len) with 1 for non-padding tokens.

• batch_threshold (int) – Maximum number of scoreable residues to score in a single batched forward pass before falling back to position-by-position scoring.

Returns:

means is a float32 numpy array of per-sequence PLL

scores (average log-likelihood per residue).

Return type:

Predictions

Raises:

ValueError – If any sequence has no scoreable residue positions (e.g. all special tokens).