🤖 machine learning

ScreenShot: A Foundation Model for Few-Shot Combination Drug Screening

ScreenShot is a hierarchical transformer foundation model that leverages in-context learning on 40 drug screening datasets to accurately predict patient responses to combination therapies using only a few-shot functional context, eliminating the need for molecular profiling or fine-tuning while enabling cost-effective experimental design through active learning.

Original authors: Antoine de Mathelin, Christopher Tosh, Wesley Tansey

Published 2026-08-13

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Antoine de Mathelin, Christopher Tosh, Wesley Tansey

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are a detective trying to solve a mystery, but instead of fingerprints or DNA, your clues are tiny drops of liquid and the suspects are thousands of different medicines. In the world of cancer treatment, doctors often use a "cocktail" of drugs rather than just one, because cancer cells are tricky and can learn to resist a single attacker. To find the perfect cocktail, scientists usually have to test every possible mix of drugs on a patient's cells in a lab. This is like trying to find a needle in a haystack by pulling out every single piece of hay one by one. It takes forever, costs a fortune, and often, there isn't enough hay (or patient cells) to go around.

This is where the idea of "predicting the future" comes in. Scientists have been building computer models to guess which drug mixes will work without having to test them all. But most of these models are like students who only pass a test if they've memorized the specific textbook for that class. They need a massive amount of detailed biological data (like a patient's genetic code) to make a guess. If that data is missing or messy—which happens often with real patients—these models get confused. The big question is: Can we build a super-smart AI that learns from all the drug experiments ever done, so it can look at just a few new clues from a single patient and instantly guess the best treatment, without needing any extra biological data?

Enter ScreenShot, a new AI tool developed by researchers at Memorial Sloan Kettering Cancer Center that tries to answer "yes" to that question. Think of ScreenShot not as a student who memorized a textbook, but as a master chef who has tasted millions of different dishes. Instead of needing to know the exact ingredients list (the genetic code) of a new dish to guess how it tastes, ScreenShot looks at the few bites you've already taken and says, "Ah, based on how this tasted with salt and pepper, I bet it would go great with garlic and thyme."

The researchers built this AI on a massive "foundation" of data, training it on about 30 million measurements from 40 different drug screening datasets. This training included roughly 3,700 different drugs tested on 6,000 biological samples (like cancer cells grown in a lab or tiny organ "mini-organs" from patients). The AI learned the patterns of how drugs interact with cells just by looking at the results: which drug, at what dose, made the cells live or die.

Here is the magic trick: ScreenShot uses a technique called in-context learning. Imagine you show the AI a few test results from a new patient (say, 50 or 100 experiments). The AI doesn't need to be retrained or "fine-tuned" for that specific patient. Instead, it instantly uses those few examples as a context window to predict the results of the thousands of drug combinations it hasn't seen yet. It does this directly on the functional results (did the cell survive?) without needing any genetic profiling.

The paper finds that ScreenShot is incredibly good at this. When tested on four new, unseen datasets, it outperformed all the previous best methods in both accuracy and in finding the "hits"—the drug combinations that actually kill cancer cells while leaving healthy ones alone. In fact, on some tests, it was 5 to 50 times faster than the other methods.

But the researchers didn't stop at just predicting; they also taught ScreenShot how to be a smart experimentalist. They created a strategy called active learning. Instead of testing drugs randomly, ScreenShot uses its own predictions to decide which experiments to run next. It's like a detective who, after interviewing a few suspects, decides to focus on the most likely culprits rather than interviewing everyone in town. The paper shows that using this smart strategy, ScreenShot can find the same number of effective drug hits as a random search, but it only needs to run one-third of the experiments. This means saving a huge amount of time, money, and precious patient cells.

One of the most impressive things about ScreenShot is that it works even when the drug is new or unknown to the system. If a drug wasn't in its training data, ScreenShot uses a special "unknown" token to guess its behavior based on the few clues it has from the current patient, rather than failing completely.

In short, the paper suggests that ScreenShot is a powerful new tool that can make drug screening much more efficient. It suggests that we can move away from expensive, slow, and data-hungry methods toward a system that learns from the past to guide the future, potentially helping doctors find the right drug cocktail for a patient much faster. While the results are strong in these computer simulations and held-out tests, the paper presents this as a significant step forward in the lab, paving the way for more efficient experiments in the real world.

Technical Summary: ScreenShot

Problem Statement

The identification of effective combination drug therapies is critical for reducing resistance, yet it is hindered by the prohibitively large search space of potential combinations. Traditional combinatorial screens are often technically infeasible, time-consuming, and expensive, particularly in clinical settings where patient-derived samples yield limited cell counts. Existing predictive models attempt to bridge this gap but typically rely on genomic or transcriptomic profiling (omics data) and require per-cohort or per-patient fine-tuning. These requirements limit applicability when tissue is scarce, molecular profiling is unavailable, or inconsistent. Furthermore, many existing methods fail to model continuous dose-response relationships for combinations or rely on flat attention mechanisms that scale quadratically with context size, making them inefficient for few-shot scenarios.

Methodology

ScreenShot Architecture

ScreenShot is a hierarchical transformer foundation model designed for few-shot prediction in combination drug screens. It operates directly on functional measurements (drug identity, dose, and viability) without requiring molecular profiling or fine-tuning on new samples. The architecture mirrors the nested structure of screening data through three levels:

Level 1: Drug Combination Encoder (DCE): Encodes individual drug-dose pairs. Drug identities are mapped to learned embeddings based on functional behavior across diverse samples rather than chemical structure. Doses are encoded via learned Fourier features. A transformer encoder processes the sequence of drugs within a combination using self-attention to capture interaction effects, aggregating them into a single combination embedding.
Level 2: Response Encoder (RE): Performs cross-attention between a query (untested drug-dose condition) and the few-shot context (observed drug-dose-viability tuples). Context observations are augmented with their measured viability. This level allows the query to retrieve evidence from dose-response curves of similar perturbations observed in the context.
Level 3: Sample Encoder (SE): Aggregates information across all observed perturbations for a specific sample using self-attention. This propagates evidence across different drug combinations (e.g., using response patterns of one drug to inform predictions for related drugs). A linear head followed by a sigmoid outputs the final viability estimate.

This hierarchical design reduces computational complexity from $O(n^2)$ (flat attention) to $O(n + p^2)$ , where $n$ is the number of observations and $p$ is the number of unique perturbations, significantly improving efficiency.

Pretraining and In-Context Learning

ScreenShot is pretrained on a pan-cancer atlas comprising 40 drug screening datasets (approx. 30 million viability measurements, 3,700 drugs, 6,000 samples). The training objective follows an in-context learning paradigm:

Input: A context set $D_s$ (labeled drug-dose-viability tuples) and a query set $Q_s$ .
Task: Predict viability for queries in $Q_s$ conditioned on $D_s$ .
Optimization: The model minimizes Mean Absolute Error (MAE) over the query set.
Mechanism: At inference, the model generalizes to new, unseen samples entirely through its pretrained attention mechanism. It conditions on the provided few-shot context via cross-attention, eliminating the need for gradient updates or fine-tuning.

Active Learning Strategy

To address the bottleneck of deciding which experiments to run, the authors propose a multi-round active learning pipeline:

Cold-Start: An initial batch is selected using $k$ -medoids clustering on Level 1 drug-dose embeddings to ensure coverage of the treatment space before any viability is observed.
Adaptive Rounds: Subsequent batches are selected using weighted $k$ -means++ seeding on Level 3 context-conditioned embeddings. Candidates are weighted by their predicted differential viability ( $\Delta_t = \hat{y}_{target} - \hat{y}_{control}$ ). This prioritizes likely "hits" (selectively effective treatments) while maintaining diversity in the selection.

Key Contributions

Foundation Model for Functional Drug Response: ScreenShot is the first foundation model for drug screening that operates purely on drug-dose-viability measurements, requiring no omics features and no per-patient fine-tuning. It supports monotherapy and arbitrary combination sizes.
Embedding-Driven Experimental Design: The authors introduce a practical screening policy that leverages the model's pretrained embeddings for cold-start selection and context-conditioned embeddings for active learning, explicitly balancing the trade-off between testing promising treatments and covering the design space.
Efficient Hierarchical Architecture: The model utilizes a hierarchical attention mechanism that aligns with the data's compositional structure, offering significant computational advantages over flat attention baselines while capturing multi-granularity representations.

Experimental Results

The model was evaluated on four held-out datasets: PDO-Breast (monotherapy ex vivo), BATCHIE, NCI-ALMANAC, and GDSC-SQ (combination screens).

Prediction Accuracy: ScreenShot achieved state-of-the-art Pearson correlation across all datasets and budgets (10 to 300 few-shot observations), outperforming baselines including XGBoost, TabPFN, and fine-tuned MLPs.
Hit Detection: In terms of identifying selectively effective treatments (Top-10% hits in Top-20% predictions), ScreenShot consistently outperformed all baselines. For example, on the BATCHIE dataset with 100 shots, ScreenShot achieved 72.6% hit recall compared to 63.9% for the next-best method.
Active Learning Efficiency: The proposed adaptive selection strategy significantly improved hit detection over random or cold-start-only selection. On BATCHIE, the adaptive strategy achieved the same hit recall as random selection at a budget of 300 using only one-third of the budget.
Speed: ScreenShot is 5–50 times faster than baselines (e.g., 8–68 seconds per experiment vs. minutes for MLPs) because it requires only a single forward pass without fine-tuning.

Significance and Claims

The paper claims that ScreenShot addresses the critical limitations of current drug screening methods by removing the dependency on molecular profiling and per-sample training, thereby enabling application to patient-derived samples where such data is scarce or unavailable. By leveraging in-context learning, the model effectively amortizes the adaptation step, allowing it to generalize to new patients immediately.

The authors emphasize that the internal representations learned by ScreenShot are directly useful for experimental design. The active learning strategy demonstrates that it is possible to identify top therapeutic hits with significantly reduced experimental costs (budget) compared to uniform screening. The work positions ScreenShot as a practical tool for precision oncology, capable of guiding combination therapy selection in resource-constrained clinical settings.

The paper modestly notes limitations, such as the reliance on drugs present in the pretraining corpus (though it handles novel compounds via UNK tokens) and the observation that global prediction accuracy does not always correlate directly with hit detection ability, suggesting that uncertainty-guided selection may not be optimal for finding top hits.