Structure-Aware Fusion with Progressive Injection for Multimodal Molecular Representation Learning

AI summary

Overview

• Research area: Machine learning for molecular representation learning, sitting at the intersection of graph neural networks, state space sequence models (Mamba), and computational chemistry / AI for drug discovery.
• Technical level: Advanced. The paper assumes familiarity with message passing on graphs, attention mechanisms, state space models, masked language modeling, and standard molecular benchmarks (TDC, MoleculeNet, QM7/8/9, Reaxtica).
• Scope: The paper introduces MuMo, a structured multimodal fusion framework that combines SMILES sequences, 2D molecular graphs, and 3D geometry through a Structured Fusion Pipeline and a Progressive Injection mechanism, and benchmarks it on 29 molecular property prediction tasks.

What This Paper Is About

Molecular property prediction models that rely on 3D conformers face two problems: conformers generated by tools such as RDKit vary even for the same molecule, and treating all modalities as equally important in a single fusion step causes "modality collapse," where noisy 3D signals dominate or distort other modalities. The authors propose MuMo, which first folds 2D topology and 3D geometry into a single stable structural prior and then injects that prior asymmetrically into the SMILES sequence stream rather than fusing modalities symmetrically.

Key Contributions

1. Structured Fusion Pipeline (SFP): A pipeline that aligns and encodes 2D and 3D inputs into a unified structural prior, using a Unified Graph that carries node, edge, and geometric linkage (bond length and angle) information, plus geometry-aware substructure partitioning that extends BRICS rules to spatial graphs for multiscale (local plus global) representations.
2. Progressive Injection (PI): An asymmetric fusion mechanism that injects the fused structural prior into a designated global token of the sequence stream through Injection Enhanced Attention (IEA), rather than concatenating or attending over all modalities uniformly. A companion Structural Prior Evolution mechanism lets structural information propagate independently across layers via state space recurrence.
3. State space backbone: MuMo is built on a Mamba-style backbone to support long-range dependency modeling and robust information propagation, with a staged strategy where early layers process SMILES only and later layers receive the structural prior.
4. Benchmark results and released code: On 29 tasks from TDC, MoleculeNet, and Reaxtica, the authors report an average improvement of 2.7% over the best-performing baseline per task, first place on 22 tasks, and up to a 27% improvement on LD50. Code is available at github.com/selmiss/MuMo.

Main Findings

• Headline benchmark result: Across 29 benchmark tasks from TDC and MoleculeNet, MuMo achieves an average improvement of 2.7% over the best-performing baseline on each task and ranks first on 22 of them, including a 27% improvement on the LD50 task. Note that the results section words this differently, stating an average 2.7% improvement "across 21 benchmark tasks from TDC and MoleculeNet, ranking first on 17 of them" — the paper reports both figures, and the two statements are not reconciled in the provided text.
• TDC classification (AUROC, 5 runs): MuMo records BBB 0.899 ± 0.014, HIA 0.979 ± 0.013, Pgp 0.942 ± 0.019, Bioavailability 0.714 ± 0.021, Tox-Avg 0.840 ± 0.015, and CYP-Avg 0.880 ± 0.017, with a Top2 count of 7 out of 10 columns. MuMo is highest among the listed baselines on BBB, HIA, Pgp, Bioavailability, and CYP-Avg, but not on Tox-Avg, where FPGNN is listed at 0.860.
• MoleculeNet classification (AUROC): MuMo records BACE-R 0.878 ± 0.046, BACE-S 0.849 ± 0.014, BBBP-R 0.962 ± 0.007, BBBP-S 0.957 ± 0.011, CLINTOX 0.985 ± 0.011, SIDER 0.677 ± 0.009, and TOX21 0.834 ± 0.009. It is not best on every column: MolBERT is listed at 0.882 on BACE-R, ChemBERTa-2 at 0.708 on SIDER, and FPGNN at 0.833 on TOX21.
• TDC regression (MAE): LD50 0.426 ± 0.031 versus 0.589 ± 0.006 for DeepMol, the baseline used for the 27% comparison; Caco-2 0.315 ± 0.055; PPBR 7.324 ± 0.323; LIPO 0.448 ± 0.007.
• MoleculeNet regression (RMSE): ESOL 0.536 ± 0.061 and Freesolv 1.082 ± 0.088.
• QM9 benchmarks (MAE, compared against Uni-Mol-v2): MuMo 505M records HOMO/LUMO/GAP 0.0030 ± 1e-05, alpha 0.283 ± 0.003, Cv 0.126 ± 0.003, mu 0.400 ± 0.0018, R² 18.08 ± 0.533, and ZPVE 0.0005 ± 1e-05. It is best on HOMO/LUMO/GAP, alpha, Cv, and ZPVE; on mu and R² the Uni-Mol2 variants are lower.
• QM7/8/9 (MAE, MoleculeNet): MuMo records QM7 42.8 ± 0.6, QM8 0.0111 ± 0.0001, and QM9 0.0030 ± 0.00001. Uni-Mol is listed at 41.8 ± 0.2 on QM7, so MuMo is not first there; the paper states MuMo outperforms baselines on QM7/8/9 for 7 out of 10 tasks.
• Reaction-level generalization: Extended to reaction inputs and evaluated on Reaxtica datasets using official splits, MuMo obtains BHC (R²) 0.952 ± 0.002, CPA (MAE) 0.144 ± 0.000, and HTE (R²) 0.873 ± 0.002, which the authors describe as best on three out of four tasks. The table in the provided text lists three tasks (BHC, CPA, HTE) even though the text refers to four Reaxtica datasets.
• Ablation, fusion components (Table 4): Full SFP (2D + SUG + GSP) is the reference at 0.00% impact. Removing 3D geometry (no SUG) drops to -7.46%; removing geometry-aware substructure partitioning (no GSP) drops to -3.00%; removing both drops to -7.29%; using sequence information alone drops to -13.11%, the largest degradation. Corresponding values: BACE 0.766 ± 0.006 and ESOL 0.719 ± 0.022 for the sequence-only setting versus BACE 0.849 ± 0.014 and ESOL 0.536 ± 0.061 for full MuMo.
• Ablation, injection timing: MuMo at layers 9–16 gives BBBP 0.957 ± 0.011 and ESOL 0.536 ± 0.061 with 0.00% average drop; full injection at 1–16 gives 0.954 / 0.587 (-1.85%); early injection at 1–8 gives 0.946 / 0.597 (-2.96%); late injection at 13–16 gives 0.961 / 0.617 (-4.40%); no injection gives 0.928 / 0.939 (-18.91%). The text reports improvement margins of 14.51%, 15.95%, and 17.06% for late, early, and full injection respectively over the no-injection setting.
• Ablation, progressive versus fixed injection: Progressive injection gives BBBP 0.957 ± 0.011 and ESOL 0.536 ± 0.061 (0.00% drop), while fixed injection gives 0.946 ± 0.008 and 0.597 ± 0.051 (-6.28%), which the authors attribute to the value of independently evolving the structural prior across layers.
• Pretraining dynamics: Comparing modality configurations, adding 2D graphs to SMILES-only consistently accelerates convergence and lowers training and validation loss, and adding geometry produces the lowest loss across all steps, which the authors interpret as each modality contributing complementary inductive signal.
• Similarity analysis: Using 20,000 molecules randomly sampled from ZINC to build molecule pairs, MuMo's embedding distances show stronger Pearson correlation with Tanimoto distance and MCS substructure overlap than MoLFormer's, suggesting embeddings better reflect underlying structural relationships.
• Representation analysis: A layer-wise UMAP of embeddings over 10 selected scaffolds (5,000 molecules) is used to show scaffold-level separation across layers; the provided text is truncated before the detailed findings of this analysis are reported.

Methodology in Plain English

The authors treat a molecule as three things at once: a SMILES string, a 2D graph of atoms and bonds, and a 3D set of coordinates.

First, they build a "Unified Graph" in which atoms are nodes, bonds are edges, and additionally each pair of bonds meeting at a central atom is described by a geometric linkage — the two bond lengths and the angle between them. Because this description uses distances and angles rather than raw coordinates, it stays consistent when a conformer is rotated; the authors prove rotational invariance in an appendix. Molecules in a batch are merged into a single large graph with offset indices so message passing can be vectorized, and a two-step procedure alternates edge-centered and node-centered message updates.

Second, they break the molecule into substructures. Rather than partitioning purely by 2D topology (as standard BRICS does), they cut edges identified from the fused spatial graph, then run message passing over both the whole graph and the substructure graph. A gated fusion scheme adaptively mixes the global and local embeddings, producing the final structural prior.

Third, instead of mixing SMILES and structure at every layer, they let the SMILES sequence run alone through early layers, then start injecting. At each injection layer, an Injection Enhanced Attention module extracts node embeddings from the graph, unbatching them into per-molecule features; it applies self-attention over the sequence, then bidirectional cross-attention (structure attends to sequence, sequence attends to structure); then it pools the graph embeddings and adds them residually to the sequence's global token. The graph representation is re-batched and carried forward so it keeps evolving across layers. The base model uses Mamba state space blocks, whose recurrent latent state is updated by both the sequence state and the injected structural prior — allowing the prior to persist and influence later layers without directly disturbing token-to-token interactions.

Training uses masked language modeling pretraining on ChEMBL-1.6M followed by task-specific fine-tuning. Evaluation uses AUROC for classification and MAE (TDC) or RMSE (MoleculeNet) for regression, with official splits for TDC and five generated splits for MoleculeNet at an 8:1:1 train/validation/test ratio, running each task five times.

Why This Matters

The paper argues that accurate in silico prediction reduces the cost of drug development — it cites a figure of over $2.6 billion on average per new drug, attributed to the Tufts Center for the Study of Drug Development — by eliminating suboptimal candidates early. If 3D-based predictions are unstable because of conformer noise, that reduces the value of using 3D at all; MuMo's contribution is a way to keep the geometric signal while making the pipeline less sensitive to how the conformer happens to be generated. Ablation evidence that removing 3D geometry costs 7.46% while removing the entire structural signal costs 13.11% supports the claim that both topological and geometric priors matter.

Real-world applications:

• ADMET screening: The TDC benchmarks span absorption, distribution, metabolism, excretion, and toxicity endpoints such as BBB permeability, HIA, Pgp, bioavailability, DILI, hERG, Ames, and CYP inhibition — the properties used to triage compounds before expensive experiments.
• Toxicity and safety prediction: The reported LD50 result (MAE 0.426) targets acute toxicity estimation, where the authors report the largest relative gain over baseline.
• Reaction and catalysis optimization: The Reaxtica evaluation covers catalytic activity (CPA) and reaction yield (BHC, HTE), extending the approach beyond single-molecule property prediction to reaction-level inputs.
• Quantum-chemical property estimation: The QM7/8/9 and QM9 evaluations target electronic and thermodynamic properties such as HOMO/LUMO gap, polarizability, heat capacity, and ZPVE, which are relevant to materials and electronic-structure workflows.

Industry relevance: the model is pretrained on ChEMBL-1.6M and released as open code, and it is benchmarked against deployed chemical models (Uni-Mol, ChemBERTa-2, MoLFormer, MolBERT, GROVER, FPGNN) as well as pretrained models with much larger parameter counts — the QM9 comparison includes Uni-Mol2 at 310M, 570M, and 1.1B, against MuMo 505M.

Future Directions

• Resolving the reported benchmark counts: The abstract/introduction and Section 4.2 give different task counts (29 versus 21) and first-place counts (22 versus 17); reconciling and reporting a single consistent set of numbers is a clear next step.
• Confirming reaction-level generalization: The text refers to four Reaxtica datasets while the results table lists three tasks (BHC, CPA, HTE), and MuMo ties rather than wins on CPA (0.144, identical to Reaxtica and MFF), so the scope of the reaction-level claim warrants further study.
• Improving conformer-sensitive property prediction: MuMo is not best

Authors’ abstract

Multimodal molecular models often suffer from 3D conformer unreliability and modality collapse, limiting their robustness and generalization. We propose MuMo, a structured multimodal fusion framework that addresses these challenges in molecular representation through two key strategies. To reduce the instability of conformer-dependent fusion, we design a Structured Fusion Pipeline (SFP) that combines 2D topology and 3D geometry into a unified and stable structural prior. To mitigate modality collapse caused by naive fusion, we introduce a Progressive Injection (PI) mechanism that asymmetrically integrates this prior into the sequence stream, preserving modality-specific modeling while enabling cross-modal enrichment. Built on a state space backbone, MuMo supports long-range dependency modeling and robust information propagation. Across 29 benchmark tasks from Therapeutics Data Commons (TDC) and MoleculeNet, MuMo achieves an average improvement of 2.7% over the best-performing baseline on each task, ranking first on 22 of them, including a 27% improvement on the LD50 task. These results validate its robustness to 3D conformer noise and the effectiveness of multimodal fusion in molecular representation. The code is available at: github.com/selmiss/MuMo.