Learning Cell-Aware Hierarchical Multi-Modal Representations for Robust Molecular Modeling

AI summary

Overview

Research area: Machine learning for molecular property prediction, specifically multi-modal representation learning that combines molecular structure with cell-level biological responses (cell morphology, gene expression). Published on arXiv (2511.21120v1, cs.LG, 26 Nov 2025).

Technical level: Advanced. The paper assumes familiarity with contrastive learning, vector quantization, graph neural networks, and self-supervised pretraining objectives.

Scope: A single paper introducing CHMR, a framework that handles missing cellular modalities and models molecule–cell–gene hierarchies for molecular property prediction.

What This Paper Is About

Most models that predict molecular properties look only at the molecule's chemical structure, but a drug's effect also shows up in how cells respond — their shape and which genes they switch on. The problem is that this cell-level data is frequently missing for many compounds (the paper reports that for the pretraining set, some external biological modalities are missing for over 90% of molecules), and existing methods treat the molecule, cell, and gene layers as a flat list rather than as a hierarchy. The goal of this work is to build a representation learning framework that stays robust when cell or gene data is absent and that explicitly organizes biological information from shallow molecular structure to deeper cellular and genomic responses.

Key Contributions

1. A unified multi-modal framework that jointly models molecular structures, cellular phenotypes, and gene expression profiles, designed to work when external biological modalities are missing.
2. A tree-structured vector quantization module (Tree-VQ) that encodes hierarchical dependencies among molecules, cells, and genes, instead of aligning modalities in a flat latent space.
3. A structure-aware augmentation and alignment pipeline — modality augmentation by neighbor propagation, semantic consistency alignment (sample-level and distribution-level), and context-propagation reconstruction — that addresses asymmetric modality missingness.
4. Large-scale empirical validation: evaluation on nine benchmark datasets spanning 728 molecular property prediction tasks, with reported average improvements of 3.6% on classification and 17.2% on regression over state-of-the-art baselines. Code is released at https://github.com/limengran98/CHMR.

Main Findings

• Headline benchmark results: On ChEMBL, ToxCast, Broad, and Biogen, CHMR reaches 84.7±0.2, 69.3±0.3, 71.4±0.2 (AUC%, higher is better) and 40.9±0.3 (Biogen MAE×100, lower is better), versus the second-best method InfoAlign at 81.3±0.6, 66.4±1.1, 70.0±0.1, and 49.4±0.2 respectively.
• Regression gains are the largest: On Biogen, the average MAE is reported as reduced by 17.2% compared to InfoAlign; classification improvements are reported in the text as approximately 2.0–4.4%.
• Broader benchmark evaluation: CHMR is reported to achieve the best average performance (82.2% AUC) on an extended set including multi-label and toxicity prediction tasks, surpassing InfoAlign (79.1%) and MOL-Mamba (80.8%), with 2.0–3.0% gains on multi-label datasets.
• Imputation strategy matters: Removing structure-aware augmentation in favor of zero or random imputation drops average relative performance by 5.3% and 4.5%; simple neighborhood imputation drops it by 2.9%.
• Every module contributes: Ablations report drops of 3.6% (no SCA), 3.9% (no Tree-VQ), 3.5% (no CPR), and 2.0% (Tree-VQ replaced with flat VQ, or random walks disabled).
• Alignment components are complementary: Removing distribution-level alignment alone costs 1.7%; removing instance-level alignment alone costs 2.9%.
• Multi-modality helps: Using molecular features only costs 4.9%; adding gene data (Mol+Gene, −3.4%), cell data (Mol+Cell, −3.3%), or expression (Mol+Express, −3.6%) partially recovers performance.
• Hyperparameter sensitivity: Best results use λ₁ = 10 (SCA weight); moderate λ₂ (0.1 or 1) is optimal for Tree-VQ; η = 1 gives peak performance for the commitment weight; tree depth h = 6 balances expressiveness and generalization, while h ≤ 4 underperforms and h ≥ 8 risks overfitting and semantic fragmentation.
• Visualization evidence: t-SNE plots show InfoAlign leaves the four modalities largely separated, SCA alone aligns them but produces flat clusters, Tree-VQ alone creates hierarchy without alignment, while CHMR shows both aligned and hierarchically organized clusters with fuller utilization of tree codes.
• Case study: In a real drug property prediction example, CHMR is reported to outperform InfoAlign on four key pharmacological endpoints, with 1D fingerprints capturing substructures such as fluorotoluene relevant to metabolic clearance, 3D conformations affecting binding affinity, and biological context improving prediction of P-glycoprotein (P-gp)-mediated drug efflux.

Methodology in Plain English

The framework has four parts that are pretrained jointly and then frozen while a lightweight prediction head is trained for each downstream task.

1. Filling in the gaps (modality augmentation). Molecules that are structurally similar probably behave similarly. So for any compound missing cell or gene data, CHMR looks at its top-K most structurally similar compounds and iteratively borrows their biological features, using a propagation scheme inspired by Dirichlet energy minimization. Data that already exists is left untouched.
2. Making the modalities agree (semantic consistency alignment). Molecular and cellular features are projected into one shared space. A contrastive loss (InfoNCE-style) pulls a molecule's pooled anchor vector toward its own cell features and away from other molecules' features. A second, VICReg-style loss checks that the borrowed/imputed features have the same distributional statistics as the original ones, which guards against the bias introduced in step 1.
3. Organizing biology as a tree (Tree-VQ). Instead of quantizing everything into one flat codebook, CHMR builds a binary tree whose levels correspond to biological scales — shallow levels for molecular fingerprints, deeper levels for cellular phenotypes and gene expression. Each feature is routed down the tree by cosine distance, with a mask forcing it to only consider the children of its previously selected parent. A symmetric loss with a stop-gradient and a reverse-commitment weight η keeps the encoder and the tree nodes consistent in both directions.
4. Reconstructing with context (context-propagation reconstruction). Using a biological graph whose edges encode known molecular perturbation–response pairs, functional associations, and shared regulatory pathways, the model performs random walks of length L from each node, then decodes the latent vectors back to molecular and cellular features. Binary cross-entropy is used for discrete features and mean-squared error for continuous ones, weighted by the strength of each relation along the walk.

The combined objective is L_CPR + λ₁·L_SCA + λ₂·L_TreeVQ. Pretraining uses 129,592 molecules built from DrugBank, Cell Painting images, the JUMP-CP multi-omics platform, and L1000 gene expression profiles. Evaluation covers nine benchmark datasets — ChEMBL, ToxCast, Broad, BACE, BBBP, ClinTox, SIDER, HIV, and Biogen — with 728 tasks. Classification is scored with AUC and regression with MAE. Splits follow a 0.6:0.25:0.15 ratio for ChEMBL, Broad, and Biogen, and scaffold-based 0.8:0.1:0.1 for the rest; each experiment is run over five random seeds (0 to 4).

Why This Matters

The paper argues that structure-only models miss the biological consequences of a molecule, and cell-aware models have been held back by sparse and asymmetric cell data and by flattening biology that is inherently hierarchical. CHMR's reported robustness under missing modalities matters because real experimental pipelines rarely produce complete multi-omics readouts for every compound — the paper describes compounds that have morphology but no transcriptomics, or vice versa.

Real-world applications:

• Drug discovery and ADMET screening: predicting absorption, distribution, metabolism, excretion and toxicity endpoints earlier, as demonstrated on the Biogen ADME benchmark.
• Toxicity and safety assessment: the ToxCast dataset covers 617 tasks, and ToxCast is described as a biological response/toxicity prediction benchmark.
• Lead optimization: using hierarchy-aware representations to reason about which structural features drive clearance or binding, as illustrated in the P-gp efflux case study.
• Integrative biomedical modeling: a general template where partial multi-omics data can still be fused rather than discarded.

Industry relevance: Pharmaceutical and biotech organizations routinely hold incomplete paired datasets (imaging for some compounds, transcriptomics for others). A method that degrades gracefully under that asymmetry — rather than requiring complete modality pairs — has direct value in preclinical pipelines, and the released code lowers the barrier to adoption.

Future Directions

• Reduce reliance on structural nearest neighbors: the augmentation step assumes molecular modalities are fully available and uses structural similarity as a proxy for biological similarity; whether that assumption holds across more diverse chemical space is an open question.
• Broaden the biological hierarchy: the tree's levels are described as molecular fingerprints, cellular phenotypes, and gene expression; extending it to organism-level or clinical endpoints is a natural next step.
• Scale and evaluate on more complete multi-omics data: results here come from a pretraining set where some modalities are missing for over 90% of molecules, so behavior on richer paired datasets is untested.
• Interpretability of the learned tree codes: the paper shows code utilization and t-SNE structure, but the biological meaning of individual hierarchy levels could be examined more directly.

Target Audience

Researchers and practitioners in AI-driven drug discovery, computational biology, and multi-modal representation learning — particularly those working on molecular property prediction, self-supervised pretraining, or methods that must cope with missing modalities. It will also interest bioinformatics engineers looking for a drop-in open-source framework (code linked in the paper) and machine learning researchers studying vector quantization and hierarchy-aware contrastive objectives. A reader should be comfortable with contrastive learning, vector quantization, and graph propagation to follow the methodology sections.

Authors’ abstract

Understanding how chemical perturbations propagate through biological systems is essential for robust molecular property prediction. While most existing methods focus on chemical structures alone, recent advances highlight the crucial role of cellular responses such as morphology and gene expression in shaping drug effects. However, current cell-aware approaches face two key limitations: (1) modality incompleteness in external biological data, and (2) insufficient modeling of hierarchical dependencies across molecular, cellular, and genomic levels. We propose CHMR (Cell-aware Hierarchical Multi-modal Representations), a robust framework that jointly models local-global dependencies between molecules and cellular responses and captures latent biological hierarchies via a novel tree-structured vector quantization module. Evaluated on nine public benchmarks spanning 728 tasks, CHMR outperforms state-of-the-art baselines, yielding average improvements of 3.6% on classification and 17.2% on regression tasks. These results demonstrate the advantage of hierarchy-aware, multimodal learning for reliable and biologically grounded molecular representations, offering a generalizable framework for integrative biomedical modeling. The code is in https://github.com/limengran98/CHMR.