iona-contrastive-50m

iona-contrastive-50m embeds tandem mass spectra (MS/MS) so that spectra of the same peptide and charge state lie close together. It is the Iona 50m encoder fine-tuned with a supervised contrastive objective, and is suited to spectrum-to-spectrum search and clustering.

Model details

  • Architecture: Iona encoder (10 layers, hidden size 640, 10 heads), with no projection head.
  • Input: a centroided spectrum given as parallel mz and intensity arrays, with at most 512 peaks.
  • Embedding: the mean and the max of the encoder's last hidden state over the spectrum's peaks, concatenated and L2-normalized (1280 dimensions). Compare embeddings with cosine similarity.

Usage

The checkpoint includes its modeling code, so it loads with trust_remote_code=True. You need torch, transformers, sentence-transformers and pytorch-metric-learning installed.

import torch
from transformers import AutoModelForPreTraining, AutoProcessor

repo = "path/to/iona-contrastive-50m"  # this repository's Hub ID, or a local copy
processor = AutoProcessor.from_pretrained(repo, trust_remote_code=True)
model = AutoModelForPreTraining.from_pretrained(repo, trust_remote_code=True).eval()


def embed(mz, intensity):
    x = processor(mz, intensity, return_tensors="pt")
    with torch.no_grad():
        h = model.iona(mz=x["mz"], log_intensity=x["log_intensity"],
                       attention_mask=x["attention_mask"]).last_hidden_state.float()
    m = x["attention_mask"].bool().unsqueeze(-1)
    mean = (h * m).sum(1) / m.sum(1).clamp_min(1)
    mx = h.masked_fill(~m, float("-inf")).max(1).values
    return torch.nn.functional.normalize(torch.cat([mean, mx], -1), dim=-1)


a = embed([175.119, 262.151, 363.198, 476.282], [1200.0, 5300.0, 800.0, 2400.0])
b = embed([175.119, 262.152, 363.197, 476.283], [1100.0, 5000.0, 900.0, 2600.0])
similarity = (a @ b.T).item()

Pass raw intensity values; the processor normalizes them.

Training

  1. Pretraining: iona-base-50m, trained with masked-intensity modeling on about 100 million consensus tandem mass spectra (checkpoint at step 220,000).
  2. Replicate contrastive training: supervised contrastive loss (SupCon, temperature 0.002) on a dataset of replicate spectra grouped by peptide for 24 epochs, with batches of 64 peptides ร— 4 replicate spectra. A KL term (weight 10) against the frozen pretrained intensity head keeps the encoder close to its pretrained behavior.
  3. Experimental fine-tuning: one epoch of the same objective on the training split of a peptide-labeled retrieval dataset of experimental and consensus spectra (batches of 85 peptides ร— 3 experimental spectra), learning rate 1e-4, seed
    1. Peptides in the replicate dataset were excluded.

The released checkpoint is the best of 3 seeds by validation experimental MAP@R. The test split wasn't used for selection.

Evaluation

Spectrum-to-spectrum retrieval on the validation and test splits of the retrieval dataset: a hit is a spectrum of the same peptide and charge. Experimental ranks experimental spectra against each other (about 25k queries); All also includes the consensus spectra in the ranked set.

Metric Validation Test
Experimental MAP@R 0.8357 0.8387
Experimental Hit@1 0.8901 0.8934
All MAP@R 0.6786 0.6794
All Hit@1 0.8309 0.8314

Baselines on the same test split (experimental MAP@R / Hit@1):

Method MAP@R Hit@1
Binned cosine (0.1 Da) 0.730 0.819
PCA of binned spectra 0.723 0.814
GLEAMS (pretrained) 0.646 0.746
This model after steps 1โ€“2 only 0.655 0.764

This model was trained on the training split of the retrieval dataset; GLEAMS and the step 1โ€“2 model were not.

Limitations

  • Inputs must be centroided spectra with between 1 and 512 peaks.
  • The final training stage and the evaluation use splits of the same dataset. Retrieval quality drops substantially on low-resolution ion-trap MS2 spectra, which weren't in training.
Downloads last month
36
Safetensors
Model size
49.8M params
Tensor type
F32
ยท
Inference Providers NEW
This model isn't deployed by any Inference Provider. ๐Ÿ™‹ Ask for provider support