Reliable Molecular Retrieval from Mass Spectra using Conformal Prediction

This paper demonstrates that applying conformal prediction to LC-MS/MS molecular retrieval generates spectrum-specific candidate sets with guaranteed coverage probabilities, effectively balancing reliability and efficiency across both in-distribution and shifted data scenarios.

The Gist

Imagine you are a detective trying to identify a mysterious substance found at a crime scene. You have a massive library of fingerprints (chemical databases) containing millions of potential suspects. Your high-tech scanner (the Mass Spectrometer) gives you a "molecular fingerprint" of the mystery substance, but it's not a perfect match.

Your computer assistant runs a search and gives you a ranked list of suspects. It says, "Suspect A is 90% likely, Suspect B is 89%, Suspect C is 88%..."

The Problem:
In the past, scientists would just look at the top of the list. If the computer said "Suspect A," they would assume it was A. But what if the computer was wrong? What if Suspect A, B, and C all look almost identical to the scanner? The computer doesn't tell you how sure it is for this specific case. It just gives a list. If you pick the wrong one, your whole investigation could go down the wrong path.

The Solution: The "Confidence Net"
This paper introduces a new method called Conformal Prediction. Think of this not as picking a single suspect, but as casting a safety net.

Instead of saying, "It's definitely Suspect A," the new method says:

"Based on how hard this case is, I am 90% confident that the real culprit is inside this small group of 3 suspects. If I'm wrong, it's only 1 out of 10 times."

Here is how the paper breaks it down, using simple analogies:

1. The Three Scenarios (The Crime Scenes)

The researchers tested their method in three different "worlds" to see how it holds up:

• Scenario 1 (The Familiar Neighborhood): The mystery substance is very similar to things the computer has seen before. The computer is a pro here. It can easily pick the right guy.
  • Result: The "safety net" is tiny. It might only contain 1 or 2 suspects. Very efficient!
• Scenario 2 (The Foreign Country): The mystery substance is from a different chemical family than what the computer was trained on. The computer is confused. The top suspects all look very similar.
  • Result: The "safety net" has to get huge to be safe. It might need to include 80% of the entire suspect list to ensure it catches the real one.
• Scenario 3 (The Wild Card): The computer is tested on a completely different type of crime scene than it was trained on. It's a total mismatch.
  • Result: The computer struggles even more. The safety net gets big, and sometimes it misses the target entirely because the rules of the game changed.

2. The "Grouping" Trick (The Detective's Intuition)

The researchers realized that not all crime scenes are the same. Some are easy; some are hard.

• The Old Way (Marginal): They used one giant rule for every case. "For 90% of all cases, our net works." This is like saying, "My umbrella works for 90% of days in the year." But if you get caught in a sudden downpour, that average doesn't help you.
• The New Way (Conditional): They grouped the cases.
  • Group A: "Easy cases where the computer is very confident." -> Use a tiny net.
  • Group B: "Hard cases where the computer is confused." -> Use a big net.
  • Group C: "Weird cases with weird chemical masses." -> Use a medium net.

The Magic Ingredient: They found that the computer's own confidence score was the best way to group these cases. If the computer says, "I'm 99% sure," you give it a tiny net. If it says, "I'm only 30% sure," you give it a massive net. This ensures that no matter how hard the case is, the "90% safety guarantee" actually holds true for that specific type of case.

3. The Trade-Off (Safety vs. Efficiency)

There is always a trade-off in detective work:

• Small Net: You have a short list of suspects. It's easy to investigate, but you might miss the real criminal (low reliability).
• Big Net: You have a huge list. You are almost guaranteed to have the criminal in there, but now you have to interview hundreds of people (low efficiency).

The Paper's Conclusion:
The researchers showed that by using this "Confidence Net" method:

1. When things are easy: You get a very short, manageable list of suspects with a high guarantee of being right.
2. When things are hard: The list gets longer, but the method tells you it's getting longer. It doesn't trick you into thinking you have a short list when you don't.
3. Grouping helps: By tailoring the net size to the difficulty of the specific case (using the computer's confidence score), they made the system much more reliable for the "hard" cases without making the "easy" cases unnecessarily large.

Why This Matters

In the real world of medicine and chemistry, guessing wrong can be dangerous. If a doctor thinks a drug is safe because the computer said "Top Match," but the computer was actually unsure, it could be a disaster.

This paper provides a tool that says: "Here is the list of possibilities, and here is exactly how sure we are about this list." It turns a black-box computer guess into a transparent, trustworthy recommendation, ensuring that scientists know when to trust the answer and when to dig deeper.

Technical Summary

1. Problem Statement

In metabolomics, Liquid Chromatography–Tandem Mass Spectrometry (LC–MS/MS) is used to identify small molecules. The standard approach involves molecular retrieval, where a measured spectrum is matched against a chemical database to rank candidate molecules based on predicted fingerprint similarities.

Key Challenges:

• Lack of Spectrum-Specific Reliability: Standard evaluation metrics (e.g., Top-k accuracy) provide dataset-level performance but fail to quantify the reliability of a prediction for a specific spectrum. They do not indicate how many candidates should be retained to ensure the true molecule is included with high probability.
• Heterogeneity of Difficulty: Retrieval difficulty varies drastically between spectra. Some spectra have clear score separations (easy cases), while others have ambiguous rankings with large candidate sets (hard cases).
• Distribution Shift: In real-world scenarios, test data often differs from training/calibration data (e.g., novel molecular clusters, instrument shifts). Standard methods often fail to maintain reliability under these shifts.
• Marginal vs. Conditional Coverage: Standard Conformal Prediction (CP) guarantees marginal coverage (average over the dataset) but does not ensure stable coverage across specific subgroups (e.g., spectra with large candidate sets or low confidence scores).

2. Methodology

The authors propose a framework using Conformal Prediction (CP) to convert ranked candidate lists into prediction sets with user-specified coverage guarantees (e.g., 90% probability that the true molecule is in the set).

A. Retrieval Setup

• Input: A spectrum $x$ and a predefined candidate set $A(x)$ (filtered by precursor mass/formula, max 256 candidates).
• Model: A Multi-Layer Perceptron (MLP) predicts a fingerprint embedding $f(x)$.
• Scoring: Candidates are scored via cosine similarity between $f(x)$ and candidate fingerprints $\phi(c)$, then converted to normalized probabilities $\pi(x, c)$ via Softmax.

B. Conformal Prediction Framework

The goal is to construct a prediction set $\hat{C}_\alpha(x) \subseteq A(x)$ such that $P(Y \in \hat{C}_\alpha(X)) \geq 1-\alpha$.

• Non-Conformity Scores: Three scores were evaluated to measure how "atypical" a candidate is:
  1. LAC (Least Ambiguous set-valued Classifier): $1 - \pi(x, c)$.
  2. APS (Adaptive Prediction Sets): Cumulative sum of probabilities along the ranked list.
  3. RAPS (Regularized APS): APS with a regularization penalty to reduce sensitivity to noise in low-ranked candidates.
• Thresholding: A threshold $\tau$ is determined from a calibration set to define the prediction set.

C. Conditional Conformal Prediction

To address heterogeneity, the authors implement Group-Conditional CP (Mondrian CP) using two strategies:

1. Cluster-Conditional CP (CCCP): Spectra are grouped via clustering based on conditioning variables. A separate threshold is computed for each cluster.
2. Nearest-Neighbor CP (CCP-NN): For each test spectrum, a local neighborhood of calibration samples is defined, and a local threshold is computed.

Conditioning Variables:
The study analyzed four variables to define sub-populations:

1. Precursor Mass.
2. Candidate Set Size ($|A(x)|$).
3. Maximum Softmax Probability (Model Confidence).
4. Candidate-Set Similarity (Average Tanimoto similarity between candidate sets).

D. Experimental Scenarios

Evaluated on the MassSpecGym benchmark under three scenarios:

• S1 (IID): Training, calibration, and test data are from the same distribution.
• S2 (Partial Shift): Training/Validation are from different molecular clusters than Test/Calibration, but Calibration and Test are aligned.
• S3 (Full Shift): Calibration and Test data are drawn from different molecular clusters (violating exchangeability).

3. Key Contributions

1. Application of CP to Molecular Retrieval: First application of conformal prediction to LC–MS/MS candidate retrieval, providing spectrum-specific reliability statements.
2. Evaluation of Non-Conformity Scores: Comprehensive comparison of LAC, APS, and RAPS across varying difficulty levels and distribution shifts.
3. Conditional Strategies: Systematic evaluation of clustering-based (CCCP) and nearest-neighbor (CCP-NN) conditioning to improve subgroup reliability.
4. Conditioning Variable Analysis: Identification of Maximum Softmax Probability as the most effective conditioning variable for reducing coverage disparities.
5. Trade-off Analysis: Quantification of the reliability-efficiency trade-off under different distribution shift scenarios.

4. Key Results

A. Baseline Retrieval Difficulty

• S1: High confidence (Top-1 accuracy ~87%, max softmax ~0.72).
• S2/S3: Significant drop in performance (Top-1 accuracy ~10%, max softmax ~0.30). Rankings become highly ambiguous, and the true molecule is buried deep in the list (Avg Rank ~62).

B. Marginal Conformal Prediction

• S1 (Aligned): All methods achieved ~90% coverage with very small sets (1.5–3.1 candidates, ~2–3% of the candidate pool).
• S2/S3 (Shifted): Prediction sets expanded dramatically to cover 80–83% of the candidate pool.
  • Reason: When the model cannot distinguish the true molecule from competitors (flat softmax distributions), the conformal threshold must be high to maintain coverage, forcing the inclusion of most candidates.
  • Robustness: RAPS showed slightly better robustness in S3 (coverage 0.895 vs. 0.881 for LAC/APS) due to its conservative regularization.

C. Conditional Conformal Prediction

• Best Conditioning Variable: Maximum Softmax Probability consistently yielded the lowest Mean Absolute Coverage Gap (MACG), indicating the most uniform coverage across subgroups.
  • Why: It directly reflects model confidence.
  • Worst Variable: Candidate-Set Similarity performed poorly due to unbalanced cluster sizes leading to unstable quantile estimates.
• Algorithm Performance:
  • S1/S2 (Aligned): CCCP (Clustering) outperformed CCP-NN, providing lower MACG and better stability.
  • S3 (Shifted): CCP-NN (Nearest-Neighbor) became superior for APS and RAPS. Local neighborhoods adapted better to the distribution mismatch than fixed clusters.
• Efficiency vs. Reliability:
  • In S1, conditional CP increased set sizes slightly (e.g., from 1.5 to 6.2 candidates for LAC) to ensure subgroup reliability, but sets remained small.
  • In S2/S3, conditional CP improved coverage uniformity (lower MACG) without significantly increasing set sizes (which were already near-maximal due to shift).

5. Significance and Conclusion

• Practical Impact: The framework transforms a ranked list of candidates into a reliable, spectrum-specific shortlist with an explicit confidence level. This allows practitioners to know exactly how many candidates to inspect for a given spectrum to achieve a desired confidence (e.g., 90%).
• Model Agnostic: The method operates on the output scores of any retrieval model (fingerprint-based, joint embedding, etc.) without requiring model retraining.
• Handling Uncertainty: It explicitly addresses the "unknown unknowns" in metabolomics. When the model is uncertain (high ambiguity), the prediction set expands naturally; when confident, it remains small.
• Limitations: The method cannot recover efficiency if the underlying retrieval scores are fundamentally uninformative (as seen in S2/S3 where sets grow large). It also assumes the true molecule is within the candidate set (a constraint of the benchmark).

Final Takeaway: Conformal prediction is a vital tool for making LC–MS/MS identification trustworthy. By using Maximum Softmax Probability for conditioning and choosing CCCP for aligned data or CCP-NN for shifted data, researchers can achieve robust, subgroup-specific reliability guarantees in molecular identification workflows.