Redefining cell states with In-Context Module discovery

Motivation

One of Prima Mente’s primary goals is to uncover novel biology for new therapeutic hypotheses for complex diseases, starting with those of the brain. In the brains of individuals with neurodegenerative conditions, many cells take on specific states that both recur and vary across patients, diseases, time, and studies. In our effort to understand, identify, and prioritise novel targets, we use these states as reference points, seeking perturbations (such as genetic or chemical interventions) that shift cells out of disease-associated states or toward protective ones. Experimentally characterizing many perturbations across cell states and in combinations is currently intractable, so predicting perturbation effects is a core modeling problem we seek to advance.

We learn to predict perturbation effects primarily from in vitro cellular models, since large perturbation experiments are usually only feasible in cultured cells^1,2. These cells lack the tissue context of the brain and may not fully capture human disease biology^3,4. These in vitro models may also activate different biological processes prior to a perturbation, causing the same intervention to have different effects across models. Therefore, understanding if and how perturbation effects measured using in vitro models translate to human biology is imperative to target discovery.

Here, we introduce In-Context Module discovery (ICM), a deep learning tool that learns from broad pretraining datasets to decompose cells into interpretable units called gene modules. ICM finds modules that correspond to known cell states. Importantly, by comparing how these modules shift across primary human tissue and in vitro cell models, ICM allows us to compare the realism of in vitro models. ICM makes it possible to link perturbation effects to the underlying biological processes active in each system. Together, these insights guide the selection of in vitro models and improve predictions of their behavior. Here, we focus on our initial cell of interest - the microglia, the resident immune cell of the brain, and RNA - the first modality of many we are modelling

Our approach to modeling RNA

Single-cell RNA-sequencing (scRNA-seq) is one of the most widely utilised methods to generate biological data⁵. Although individual-gene measurements in scRNA-seq are often noisy and sparse^6,7, coordinated variation across genes can reveal underlying biological processes. Instead of studying the cell at the individual gene level, we group genes whose expression varies together across cells into gene modules (Fig. 1a)^8,9. Under this view, a cell’s gene expression profile can be decomposed into a weighted sum of its modules, allowing us to represent cell state by which modules are active in a particular cell, and those that are not. Because gene modules capture coordinated expression patterns that can reflect biological processes, we posit that they can help explain and predict perturbation effects based on which processes are active before treatment (Fig. 1b).

aGene modules as functional units

bPerturbations act on modules

cIn-context module updates

dTissues

Organ groups

Broad cell classes

ICM learns a fixed set of gene modules once from a corpus of 2.85M human cells, and adapts the modules to each new dataset from a sample of the dataset's cells (Fig. 1c-d). We designed ICM to ensure modules are comparable: each adapted module remains linked to its pre-fit counterpart. They are also interpretable: explicit per-gene weights let us link module composition to biological processes.

Together, these properties let us examine which aspects of disease-associated cell states an in vitro model captures and how a perturbation alters their activity.

How ICM discovers and adapts modules

ICM works as follows (Fig. 2):

1. Learn a shared bank of prior modules. ICM initializes a shared bank of gene modules using non-negative matrix factorization (NMF)¹⁰, then refines these modules during pretraining. Each module is represented by a set of non-negative per-gene weights. All datasets share this bank.
2. Adapt prior modules to a new context. Support cells from a new context are summarized through per-gene expression statistics and gene–gene correlations. Iterative slot attention (ISA)¹¹ uses this summary to compute a per-gene adjustment to each module in the pretrained prior bank. The prior bank itself is not updated.
3. Describe query cells as combinations of modules. ICM reconstructs an unseen query cell's expression profile as a non-negative weighted sum of the adapted modules. The weight assigned to each module reflects its activity in that cell; we refer to these weights as module activities throughout this post.

Existing module decomposition methods make different trade-offs between adapting to a new dataset and keeping modules comparable (Fig. 3). NMF and cNMF learn a unique module set for each dataset, so module identities (e.g. mod_1) are not anchored across datasets⁹. Projection methods such as projectR and starCAT transfer fixed programs into new data, but they do not adapt to the new context-specific patterns^12,13. LIGER, Spectra, and expiMap can capture dataset- or cell-type-specific patterns , but they do not combine a shared pretrained module bank with context-dependent gene-weight adaptation without retraining^14,15,16. Stack also uses in-context learning to represent new cells without neural network weight updates, but its learned gene groupings are fixed after training and do not adapt to a new dataset’s co-expression patterns¹⁷. ICM is designed to combine context-specific adaptation with a shared, identifiable module bank.

Comparing cell states across datasets requires modules that can accommodate local expression patterns while retaining a common reference. ICM achieves this using two components: a prior module bank learned during pretraining across populations, and a context network that adapts these modules to a reference population. Each adjusted module retains its position in the prior bank, allowing us to track its gene weights and cellular activities across datasets. In the next section, we test whether adapted modules preserve recognizable gene-weight patterns from their original counterparts.

During pretraining, ICM jointly learns the shared module bank and the context network across populations. NMF initializes the prior bank, after which backpropagation updates both the bank and the context network. The primary reconstruction objective encourages context-adapted modules to explain query-cell expression, while an auxiliary reconstruction objective keeps the shared bank useful independently of adaptation. A balancing term further discourages unused modules. At inference, the learned parameters remain fixed: reference cells determine the adapted module bank, and ICM fits module activities separately for each query cell.

To adapt the shared module bank to a new population, ICM uses reference cells to identify which genes vary together. The context encoder combines gene identity, expression statistics, and gene–gene correlations to create a representation of each gene. Iterative slot attention (ISA)¹¹ then aggregates these representations into module states, repeatedly updating which genes contribute to each state. The context network uses these states to compute gene-level adjustments to the shared module bank, producing modules tailored to the reference population.

These adjustments determine each gene's membership across modules. ICM uses Sparsemax²⁷ to produce a sparse gene–module membership matrix M, allowing some memberships to become exactly zero. This design decision forces the model to decisively allocate genes to modules, rather than having large modules with a long tail of low weighed genes. Positive gene-specific scaling factors then convert these memberships into a weighted module bank W, separating how a gene is distributed across modules from its overall contribution to expression. Once ICM has adapted the modules, it uses NMF-based activity inference⁹ to estimate each query cell's module activities by fitting a non-negative combination of the adapted modules to its observed RNA profile. Cells sharing a reference population therefore use the same adapted modules but can differ in their module activities.

Ablation results

For ICM to support meaningful comparisons across datasets, its modules must capture population-specific expression patterns while preserving their comparability to the prior module bank. We evaluate these properties separately, then test how slot attention contributes to maintaining them.

In summary:

• Context adaptation improves reconstruction of gene expression.
• Modules are comparable before and after fit.
• Slot attention improves module identity preservation.

We utilize these properties in the case study below.

Does context adaptation improve fit to expression?

Across 158 validation contexts, context adaptation reduced mean squared reconstruction error (MSE) by 27.6% and increased variance explained from 23.1% to 45.3%, with lower error in every context. We fitted cell activities separately for each bank, so this comparison tests the benefit of changing the modules themselves.

Evaluation / metric	Context-adapted bank	Shared bank only
Validation: reconstruction MSE ↓	0.0575	0.0794
Validation: variance explained ↑	45.3%	23.1%
Held-out microglia: reconstruction MSE ↓	0.0520	0.0807
Held-out microglia: variance explained ↑	48.9%	20.2%
Held-out microglia: modules enriched for a Fumagalli state ↑	16.9%	8.2%
Gene-list specificity ↑	68.21%	66.16%

The same benefit extended to primary microglia from 18 donors held out of training in Tuddenham et al.²⁹. Adaptation increased variance explained from 20.2% to 48.9%, and it also more than doubled the fraction of modules enriched for a Fumagalli microglial state signature, from 8.2% to 16.9%, with all five states represented in every adapted donor/support sample. Enrichment is akin to recall, while gene-list specificity tests whether modules specifically capture the pathway of interest or whether contain broad lists of highly annotated genes. The context-adapted bank achieves 68.21% from 66.16% in the shared bank (table A1).

This result connects the reconstruction objective to the state comparisons above: modules adjusted to the dataset-specific population can fit its expression more closely and capture more of its known state signatures.

Do modules maintain correspondence after context adaptation?

Our downstream analyses rely on comparing the same modules across datasets. However, better reconstruction would be less useful if adaptation changed modules so substantially that they no longer resembled their original counterparts. We therefore compared adapted modules with the shared bank using both their gene memberships (M) and their final weighted representations (W). Figure A1 brings the two comparisons together: panel a shows the relationship between adapted and prior modules in W for one context; panel b summarizes membership preservation in M across contexts.

Panel b quantifies the result across evaluated contexts: each adapted module’s membership vector in M remained closest to its own prior module. This continuity supports following modules across datasets. In the following microglia case study, we examine how these comparisons reveal conserved and context-specific biological features.

What does slot attention contribute to module comparability over self-attention?

Slot attention was chosen to enable module correspondence across contexts, but we wanted to compare to an alternative method: self-attention. To compare, we replaced the three slot-attention iterations and their readout with gene self-attention followed by a direct gene-to-module projection, keeping the context encoder, shared bank, and activity inference unchanged.

Both architectures preserved the closest prior match in the gene membership matrix M. However, slot attention preserved more original module identities in the final weighted bank under one-to-one matching. Reconstruction performance was similar, with self-attention achieving 0.7% lower validation MSE (Table A2). The benefits of slot attention outweigh the minor MSE improvement of self-attention. Slot attention also required 72.6% fewer refinement-component FLOPs per support set (2.645 versus 9.664 GFLOPs). This comparison counts matrix operations for one forward pass through the refinement component, rather than total model computation or measured runtime.

Metric	Slots (3 iterations)	Self-attention + projection
Refinement-component parameters	156,800	132,608
Full-model parameters	2.215M	2.190M
Validation reconstruction MSE ↓	0.05746	0.05705
Validation weighted module IDs retained ↑	91.1%	84.2%
Component GFLOPs / support set ↓	2.645	9.664
Fractional variance explained ↑	45.278%	45.686%

To test whether this pattern extended beyond the validation contexts, we compared the same frozen checkpoints across seven studies excluded from pretraining. Both models received the same support and query cells. Slot attention retained more weighted module IDs in every study mean (Fig. A2a). In all six studies with at least five donors, the paired 95% confidence interval supported higher retention with slot attention.

Across the same seven studies, both architectures achieved similar overall reconstruction performance, with no statistically significant difference in MSE (paired p = 0.287; Fig. A2b).

ICM identifies biologically meaningful modules and reveals unexpected features of in vitro microglia

To quantify how well our learned modules capture real functional relationships between genes, we tested ICM on microglia, the brain’s resident immune cells, using data from primary human microglia and in vitro models. Here, we show that ICM:

• Recovers known microglia states: In every dataset, ICM finds modules matching all six well-described microglial states, even though the modules were learned without these states as labels.
• Groups genes by a shared regulator: ICM finds a module enriched for genes bound by STAT2, a regulator of the interferon (antiviral) response which is enriched in every dataset. So the modules can capture genes that are co-regulated, not simply co-expressed.
• Reveals unexpected features of in vitro microglia: ICM finds that a subset of in vitro microglia express genes typical of dendritic cells, a different immune cell type. We treat this as a hypothesis for follow-up experiments.
• Explains why a perturbation behaves differently across contexts: Reducing STAT2 strongly shifted gene expression one in vitro model but barely moved another. The STAT2-associated module shows why: low baseline module activity results in minimal perturbation responses, even for highly-effective target knockdowns.

Background: How we define ground truth microglial states

Microglia are the first cell Prima Mente's cell modelling efforts have focussed on. As the brain's resident immune cells, they clear cellular debris, respond to pathogens, and help regulate synapses and inflammatory signaling³¹. Research implicates microglia in neuroinflammatory and neurodegenerative diseases^32,33,34, links them to genetic risk for dementia^35,36,37, and identifies early changes as markers of neurodegeneration^38,39,40. Across a spectrum of disease and stress conditions, microglia exhibit recurring states associated with biological processes such as phagocytosis and cytokine signaling^28,41.

As large perturbation screens are difficult to perform in primary microglia, researchers often use microglia-like cells derived from induced pluripotent stem cells (iPSCs)^2,42. However, these models do not necessarily reproduce the biology of microglia in the human brain. Microglial identity depends upon the brain environment, and primary microglia rapidly lose characteristic gene expression patterns in cell culture^3,4. Different iPSC differentiation methods also produce different mixtures of baseline cell states, and the same perturbation can produce different responses across models⁴². Mapping cells to known microglial states therefore provides a common reference for comparing models and understanding their responses to perturbations^2,42.

To establish a consistent reference for these comparisons, we use a common reference of microglial states as defined by Fumagalli et al.²⁸. By synthesizing studies across species and experimental contexts, the group identified five recurring states, each characterized by a set of marker genes:

• Homeostatic: markers of microglia identity and resting state
• Disease-associated microglia (DAM): an activation state associated with neurodegeneration and TREM2 signaling.
• Cytokine-responsive microglia (CRM): an activation state associated with cytokine signaling and inflammation.
• Interferon-responsive microglia (IFN): an activation state marked by interferon-stimulated genes.
• MHC-II: an activation state with high expression of MHC class II and other antigen-presentation genes.

We add a sixth, Cycling, using the G2/M gene signature from Tirosh et al.⁴³ to capture a proliferative microglial state observed in vivo. Together, these six signatures provide a common reference for evaluating ICM modules across primary, xenograft, and iPSC-derived microglia.

Comparing primary human microglia with iPSC-derived models

To evaluate which of these six microglial states are present across experimental contexts, we chose four datasets that span a range from primary human microglia to in vitro models.

1. Primary: Microglia isolated from fresh brain tissue of donors with neurodegenerative disease, representing a cross-disease human microglia resource²⁹.
2. Xenograft: iPSC-derived microglia transplanted into mouse brains, allowing the cells to develop fully in a brain microenvironment⁴⁴.
3. iPSC-derived perturbations: Two microglia differentiation protocols used by McQuade et al. in the same CRISPRi screen, enabling us to compare perturbation responses across models⁴²:
   1. iMG: Microglia generated through directed iPSC differentiation^42,45.
   2. iTF: Microglia generated through transcription-factor induction of iPSCs^2,42.

In the figures below, the datasets are simply labelled by model context: Primary, Xenograft, iMG and iTF.

ICM-derived modules enrich for known microglial states and established regulatory networks

We first asked whether ICM recovers the six reference microglial states across our four datasets. All six state signatures were enriched in at least one module in each dataset, and several modules consistently matched the same state across all four datasets (Fig. 4). These results suggest that ICM can effectively identify consensus cell states. To do this, we used over-representation analysis (ORA) to test whether a module’s highest-weighted genes overlap with the six conserved microglia state signatures more than expected by chance^46,47.

Recovering known microglial states shows that ICM captures recognizable cell states, but it does not tell us which regulators might influence these states. We therefore asked whether modules are enriched for genes bound by known transcription factors, which could help link module activity to specific perturbations. As an example, we focus on STAT2, a transcription factor that mediates interferon responses⁵⁰. mod_190 was the only module significantly enriched for STAT2-bound genes in both the pretrained bank and all four microglial datasets (Figure 5a, b). Together with its association with the interferon-response state, this suggests that mod_190 captures a STAT2-associated biological program and provides a potential link between module activity and perturbation response.

Following an MHC-II module reveals dendritic-cell-like features in in vitro microglia models

Using ICM, we traced a module associated with a known microglial state and uncovered an unexpected resemblance to another cell type.

Identifying known microglial states across datasets can tell us which biological programs ICM captures, but not whether their underlying gene composition differs between models. Because ICM links each adapted module to the same entry in the prior bank, we can examine these differences directly. We demonstrate this with mod_75, a module enriched for the MHC-II state in primary microglia.

In primary microglia, HLA class II genes account for ~53% of mod_75's weight, consistent with an antigen-presentation-associated program (Figure 6a). The xenograft version retains much of this gene composition, whereas in both iMG and iTF, more than 10% of module weight falls on genes with little or no expression in primary microglia (Figure 6b).

The two in vitro models also differ in which additional genes contribute to mod_75. In iMG, mod_75 retains HLA class II genes but also gains CD207, CLEC10A, and CD1A. In iTF, HLA class II genes account for only 9% of the module's weight, while CD74 accounts for 41%, alongside newly weighted genes RAB7B, MMP9, and SLAMF7 (Figure 6c). In the pretraining corpus, these six newly weighted genes are most widely expressed in dendritic cells and are nearly undetected in microglia.

By comparing the same module across contexts, ICM revealed unexpected dendritic-cell-like features without requiring us to search for dendritic-cell markers in advance. This finding raises a testable question: do these features reflect contaminating cell populations or differentiation-associated changes within microglia-like cells? More broadly, this example demonstrates how ICM can uncover biological differences between experimental models that known state annotations alone might miss.

Modules reveal why a STAT2 knockdown has different effects across in vitro models

In the introduction, we posited that gene modules are a useful resolution for studying perturbations. Here, we test that idea on a concrete observation: the same perturbation has a very different effect in different in vitro models. Understanding these differences points to a practical application: Modules can guide the choice of in vitro model to use for a perturbation screen before that screen is run.

McQuade et al.⁴² performed the same 31-gene CRISPRi screen in iMG and iTF, including knockdown of STAT2, a key regulator of interferon signaling. Yet STAT2 knockdown produced strikingly different responses: it ranked among the strongest perturbations in iTF by E-distance, but only 27th of 31 in iMG, where its effect was much smaller (Figure 7).

As described above, mod_190 is highly enriched for STAT2-bound genes in every dataset (Figure 5a, b), so we would expect STAT2 knockdown to lower its activity. In iTF, we see high mod_190 activity at baseline, and STAT2 knockdown removes almost all of the activity (log2 fold change −2.23; Fig. 8a). In iMG, baseline activity is already low and changes little after knockdown (−0.16). The independent interferon-response state signature shows the same pattern (−0.30 versus −0.02; Fig. 8b). These results suggest that differences in baseline biological activity help explain why the same perturbation has such different effects across models.

We chose STAT2 as a positive control because the interferon response is well characterized, allowed us to evaluate our findings against known biology. However, the broader opportunity is to use this approach when we know less about a target or how it will behave in different experimental models. For an understudied target, we would not know in advance which module matters or which in vitro model to screen in.

Here, we propose a few guiding questions using ICM and public data to prospectively select in vitro models before committing to a costly perturbation screen:

1. Which biological programs might the perturbation affect? We can identify modules associated with a target using independent evidence, such as STAT2 ChIP-seq data linking STAT2 to mod_190.
2. Are the relevant biological programs active in the model? Baseline module activities can help us assess whether a model captures the processes we intend to perturb. For example, iTF has substantially higher baseline activity of the STAT2-associated interferon program than iMG.
3. What cell states are we interested in moving between to alter the disease we’re studying? Are both cell states robustly present in the in vitro model at baseline? (e.g. the interferon microglia state activity is strong in iTF, but not in iMG).

Together, these analyses show how comparable and interpretable modules can help us assess the biological suitability of an in vitro model for a particular perturbation. Using public resources and baseline expression data, we could identify potential mismatches between models and experimental goals before committing to costly perturbation screens.”

Discussion

A central challenge in perturbation modeling is that the same perturbation can produce different effects across cell types, in vitro models, and cell states. ICM makes that context inspectable by learning comparable and interpretable gene modules from a broad pretraining corpus and adapting them to the biological patterns present in each dataset.

Our microglia case study demonstrates why both properties are useful. ICM recovered known microglial states and identified a STAT2-associated interferon module across experimental contexts. Comparing module composition revealed unexpected dendritic-cell-like features in in vitro microglia, while comparing baseline module activity helped explain why STAT2 knockdown produced a strong response in iTF but little response in iMG.

Together, these findings suggest a way to evaluate in vitro models based on the biological processes relevant to a particular experiment, rather than their overall resemblance to primary tissue alone. By combining module composition and baseline activity with existing regulatory evidence, ICM helps researchers select experimental models before running perturbation screens and interpreting their results afterward.

An important next step is to test whether ICM can prospectively identify suitable models and anticipate perturbation responses across a wider range of targets and cellular contexts. To extend this work, we are currently running large scale experiments using in vitro models composed of complex tri-cultures of microglia, neurons, and astrocytes. These experiments span hundreds of millions of cells across different genetic and disease backgrounds. Fortunately, the scale of these projects allows for staged readouts, allowing us to plan future combinatorial perturbations by interrogating contexts from a prior stage — a perfect fit for ICM. Ultimately, we see learned, adaptable, and interpretable modules as fundamental architectural building blocks for the next generation of virtual cell models.

About Parthenon

Parthenon is Prima Mente’s virtual cell platform. Designed to be cell- and indication-agnostic, Parthenon today focusses on the microglia, neuron, and astrocyte. The platform enables our efforts in target biology and mechanism discovery by modeling cellular responses to chemical and genetic perturbations. Parthenon is powered by large scale perturbational data generation efforts.

We apply a first-principles approach to modeling the underlying mechanisms of biology. If you find this type of analysis interesting, please sign up for the Alzheimer’s Translation Challenge!

Author contributions

Y.S., A.P., K.B., and J.H. conceived of and designed the study, trained and evaluated models, analyzed results, and drafted the manuscript. F.M.M.-Z. and K.B. prepared the data corpus. A.P., Y.S., and P.N. contributed to machine learning engineering and infrastructure. P.N. and R.S. managed the project. All authors contributed to the interpretation of the results and critically reviewed and revised the manuscript.

Acknowledgements

We thank Pooja Kathail, Andrey Karailiev, Rachel (Qiurui) Zeng, and Alfred Wong for their kind feedback on the text. We thank Sam Gubbins and Shabbir Khan for enabling this work through infrastructure. We are grateful to Hannah Madan for their mentorship and support, as well as to the rest of the Prima Mente team for their continued support. We thank the Nebius and Nvidia teams for their ongoing support, infrastructure, and compute for model training and inference. Parthenon was originally catalysed by the Alzheimer’s Insights AI Prize, backed by the Alzheimer’s Disease Data Initiative.

Citation

@misc{primamente2026icm,
  title        = {Redefining cell states with In-Context Module discovery},
  author       = {Shirsath, Yash and Patel, Anooj and Blake, Kerrigan and
                   Hoover, Jonathan and Martin Zamora, Francisco Manuel and
                   Niki, Pouya and Solanki, Ravi},
  year         = {2026},
  month        = {Nov},
  url          = {https://www.primamente.com/blog/icm/},
  organization = {Prima Mente}
}

Download the citation: BibTeX (.bib) RIS (.ris) EndNote (.enw) CSL-JSON (.json)

References