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

ScreenShot is a hierarchical transformer-based foundation model for few-shot prediction in combination drug screening. It is pretrained on a pan-cancer drug screening atlas of 40 datasets, covering approximately 30 million viability measurements across about 3,700 drugs and 6,000 biological samples. The model performs in-context learning at inference: given a small set of drug-dose-viability measurements from a new patient sample, it predicts the remaining untested responses without fine-tuning and without requiring molecular profiling of the sample. The architecture mirrors the nested structure of screening data with three levels: a Drug Combination Encoder (Level 1) that processes drug and dose embeddings via self-attention within each combination, a Response Encoder (Level 2) that uses cross-attention between the query and context observations augmented with observed viability, and a Sample Encoder (Level 3) that aggregates information across perturbations via self-attention. The model is trained by minimizing mean absolute error over query sets, and it uses learned Fourier features for doses and viabilities, with dedicated UNK embeddings for novel compounds.

On four held-out datasets (NCI-ALMANAC, GDSC-SQ, BATCHIE, and PDO-Breast), ScreenShot outperforms all baselines—including XGBoost, TabPFN, MLP trained from scratch, and MLP fine-tuned on the same pretraining corpus—in both prediction accuracy (Pearson correlation) and identification of selectively effective treatments (hit recall). The advantage is largest at low budgets; for example, at n=10 on BATCHIE, ScreenShot reaches 63.9% hit recall compared to 45.9% for the next-best method. ScreenShot is also 5–50× faster than all baselines. The paper states: ScreenShot achieves the highest Pearson correlation on every dataset and budget, and the highest hit recall on 16 out of 20 dataset/budget combinations.

Beyond prediction, the authors develop a multi-round active learning strategy for experimental design. The strategy uses k-medoids clustering on Level 1 drug-dose embeddings for cold-start batch selection, then uses weighted k-means++ seeding on Level 3 context-conditioned embeddings for adaptive rounds, weighting candidates by predicted differential viability (therapeutic index) to prioritize likely hits while ensuring diversity. The paper reports: our adaptive selection algorithm achieves the same hit recall as uniform selection with only a third of the budget. On BATCHIE at budget 300, adaptive selection reaches approximately 85% hit recall compared to about 76% for both random and cold-start selection. The ablation study shows that delta weighting is the most important component, that diversity is essential (pure exploitation strategies perform substantially worse), and that uncertainty weighting improves overall prediction accuracy but not hit detection, with 6 percentage points lower hit recall on BATCHIE compared to delta weighting.

Improvements to AI Systems:

1. In-Context Learning for Personalized Prediction Without Fine-Tuning: Build AI systems that can take a small set of labeled examples from a new user, patient, or environment at inference time and immediately predict outcomes for unlabeled cases—without retraining or needing rich metadata (e.g., genomic profiles). This enables rapid personalization in domains like drug response, credit risk, or sensor calibration.

2. Hierarchical Attention for Nested Data Structures: Design model architectures with multiple levels of self- and cross-attention that mirror the natural hierarchy of the data (e.g., combinations → individual responses → overall sample). This improves generalization from sparse observations by explicitly modeling interactions at each level, applicable to multi-agent systems, supply chains, or clinical trial design.

3. Learned Fourier Features for Continuous Inputs: Replace raw continuous inputs (doses, times, concentrations) with learned Fourier embeddings to capture periodic and multi-scale patterns. This improves prediction accuracy for any system dealing with continuous variables, such as pharmacokinetics, material properties, or financial time-series.

4. Cold-Start Active Learning via Embedding Clustering: For new tasks with no prior data, use k-medoids clustering on learned input embeddings to select a diverse, representative initial batch. This reduces the number of experiments or queries needed to achieve a target performance, applicable to autonomous labs, survey design, or recommendation systems.

5. Adaptive Batch Selection with Weighted Diversity and Exploitation: Implement an iterative selection algorithm that combines predicted outcome value (e.g., therapeutic index) with diversity constraints (weighted k-means++). This balances exploration and exploitation, achieving the same hit rate with one-third the budget—useful for drug discovery, A/B testing, or hyperparameter optimization.

6. Delta-Weighted Prioritization for Rare but High-Impact Events: Weight candidates by the predicted difference from a baseline (e.g., differential viability) rather than absolute value. This improves recall of rare, selectively effective outcomes (e.g., identifying patients who respond exceptionally to a treatment) and can be applied to anomaly detection, personalized marketing, or rare disease diagnosis.

7. Uncertainty-Aware Selection with Separate Objectives: Decouple uncertainty weighting from outcome prioritization—use uncertainty for improving overall regression accuracy, but use delta weighting for hit detection. This modular approach allows AI systems to optimize multiple, sometimes conflicting, objectives (e.g., accuracy vs. recall of rare events) without compromising either.

8. Fast Inference via Pretrained Universal Embeddings: Use a single pretrained model that maps any new input (e.g., novel drug, unseen sample) to a shared embedding space via dedicated UNK tokens. This enables zero-shot transfer to unseen categories, reducing the need for retraining and enabling real-time predictions in dynamic environments.

What the Improved AI System Can Do:

• Predict a patient’s full response profile to thousands of drug combinations from just 10–30 initial measurements, enabling personalized cancer therapy selection in days rather than months.

• Design adaptive experiments that automatically choose the most informative next batch of tests, achieving high hit rates with 3× fewer resources than random selection.

• Operate in low-data regimes (e.g., rare diseases, new markets) without fine-tuning, by leveraging a large pretraining corpus and in-context learning.

• Identify selectively effective interventions (e.g., drugs that work only for a subset of patients) with high recall, even when those events are rare.

• Provide both accurate continuous predictions and high-precision classification of extreme outcomes simultaneously, using separate weighting mechanisms.

• Scale to millions of observations with 5–50× faster inference than traditional ML baselines, enabling real-time decision support in clinical or industrial settings.

Treating patients with combinations of drugs reduces the risk of resistance to any individual drug. Finding effective combinations is difficult because the large search space makes combinatorial screens prohibitively expensive, time consuming, and often technically infeasible. Predictive models can fill this gap, yet existing methods typically require molecular profiling of each sample and per-cohort training, limiting their applicability when time and tissue are scarce. To address this challenge, we introduce ScreenShot, a hierarchical transformer pretrained on 40 drug screening datasets covering 3,700 drugs and 6,000 biological samples, whose architecture mirrors the nested structure of screening data. Given a few-shot context of observations from a new patient, ScreenShot predicts the response of the sample to combination therapies through in-context learning, operating directly on functional measurements with no fine-tuning and no molecular profiling. On four held-out datasets, ScreenShot outperforms all baselines in both prediction accuracy and identification of selectively effective treatments. ScreenShot's internal representations are directly useful for experimental design: we use them to drive a weighted k-means++ active learning strategy that selects which experiments to run, achieving the same hit detection as uniform screening with a third of the budget. Source code and interactive dashboard: https://github.com/tansey-lab/screenshot.