GRASS-MIL: graph-based representation and discovery of phenotype-associated spatial structures with multiple instance learning

Abstract

Tissue function is shaped by both cellular composition and the spatial organization of cells. Spatial single-cell technologies capture this organization across tissue sections, enabling tissue architecture to be linked to disease phenotypes, prognosis and treatment response. However, many current graph-based predictive models depend on assay-specific molecular features, limiting joint analysis of cohorts profiled with incompatible panels, and either aggregate local predictions uniformly or rely on post hoc importance scores that do not directly decompose the basis of a sample-level prediction. Here we introduce GRASS-MIL, a graph multiple-instance learning framework that learns the relative relevance and class-specific evidence of individual cellular neighborhoods, predicts phenotypes across platforms, and decomposes each sample-level prediction exactly into signed contribution from each local neighborhood. These contributions are terms of the prediction itself rather than a post hoc approximation and sum exactly to the predicted logit. In simulations, GRASS-MIL distinguished phenotypes with identical global cell-type composition but different spatial organization, whereas composition-based classifiers performed no better than chance. A single model trained jointly on CosMx spatial molecular imaging, Xenium and MERFISH datasets, harmonized only through a shared cell-type vocabulary, predicted treatment status in high-grade serous ovarian carcinoma with an area under the receiver operating characteristic curve of 0.74 and associated this prediction with two stromal architectures: an immune-rich niche admixed with malignant cells and a desmoplastic niche characterized by malignant-cell exclusion. Applied to spatial proteomics data from non-small-cell lung cancer, GRASS-MIL identified distinct niches associated with relapse (area under the receiver operating characteristic curve, 0.57), including myeloid-rich niches in relapsing tumors and adaptive immune niches in non-relapsing tumors. In a separate lung adenocarcinoma cohort, a similar lymphoid-myeloid polarity was observed when predicting 3-year mortality (area under the receiver operating characteristic curve, 0.62), with a B-cell and CD4+ T-cell niche associated with favorable outcome and a neutrophil and macrophage niche associated with adverse outcome. GRASS-MIL thus combines cross-platform phenotype prediction with an exact, signed decomposition of each prediction into the cellular neighborhoods on which it depends, providing a direct route from sample-level prediction to testable hypotheses linking tissue architecture with patient phenotypes.