MolGraph-xLSTM: A graph-based dual-level xLSTM framework with multi-head mixture-of-experts for enhanced molecular representation and interpretability

TL;DR

MolGraph-xLSTM tackles the challenge of long-range dependencies in molecular property prediction by marrying dual-level molecular graphs (atom- and motif-level) with an xLSTM-based feature extractor and a Multi-Head Mixture of Experts (MHMoE) for refined representations. The approach integrates GNN-based local feature extraction, jumping knowledge, and xLSTM to capture both local structure and distant interactions, then fuses atom- and motif-level signals through MHMoE before final prediction. Empirically, it achieves consistent gains across 10 datasets, including notable improvements on BBBP (AUROC) and ESOL (RMSE), and offers interpretability by highlighting influential substructures and atoms. The work demonstrates a scalable, interpretable framework for molecular representation learning that has practical implications for drug discovery, with future work aimed at enhancing motif initialization, incorporating bond-level information, and extending to broader tasks like drug-target interactions.

Abstract

Predicting molecular properties is essential for drug discovery, and computational methods can greatly enhance this process. Molecular graphs have become a focus for representation learning, with Graph Neural Networks (GNNs) widely used. However, GNNs often struggle with capturing long-range dependencies. To address this, we propose MolGraph-xLSTM, a novel graph-based xLSTM model that enhances feature extraction and effectively models molecule long-range interactions. Our approach processes molecular graphs at two scales: atom-level and motif-level. For atom-level graphs, a GNN-based xLSTM framework with jumping knowledge extracts local features and aggregates multilayer information to capture both local and global patterns effectively. Motif-level graphs provide complementary structural information for a broader molecular view. Embeddings from both scales are refined via a multi-head mixture of experts (MHMoE), further enhancing expressiveness and performance. We validate MolGraph-xLSTM on 10 molecular property prediction datasets, covering both classification and regression tasks. Our model demonstrates consistent performance across all datasets, with improvements of up to 7.03% on the BBBP dataset for classification and 7.54% on the ESOL dataset for regression compared to baselines. On average, MolGraph-xLSTM achieves an AUROC improvement of 3.18\% for classification tasks and an RMSE reduction of 3.83\% across regression datasets compared to the baseline methods. These results confirm the effectiveness of our model, offering a promising solution for molecular representation learning for drug discovery.

Figures (9)

• Figure 1: Comparison of atom-level and motif-level graph representations. (a) The atom-level graph represents each atom as a node and each bond as an edge. (b) The motif-level graph combines substructures into single nodes, resulting in a graph that is less complex than the atom-level graph.
• Figure 2: Architecture of MolGraph-xLSTM. The architecture consists of four main components: (A) Motif graph construction. The atom-level graph is decomposed into motifs to form a motif-level graph. (B) Feature extraction on the atom-level graph. A GCN-based xLSTM framework with jumping knowledge extracts features, followed by pooling to generate the atom-level representation $f_{atom}^{xLSTM}$. (C) Feature extraction on the motif-level graph. Using xLSTM blocks and pooling to produce the motif-level representation $f_{motif}^{xLSTM}$. (D) Multi-Head Mixture-of-Experts (MHMoE) and property prediction. Features ($f_{gcn}$, $f_{atom}^{xLSTM}$ and $f_{motif}^{xLSTM}$) are combined and refined through the MHMoE module for final property prediction.
• Figure 3: Ablation study results on the Sider and FreeSolv datasets. Performance comparison of ablation variants against the full MolGraph-xLSTM model and baseline models.
• Figure 4: Examples of different atom input orders for molecular graphs in xLSTM. (a) RDKit Default Order: Atoms are ordered as per the default output from RDKit. (b) DFS Order: Atoms are ordered based on a Depth-First Search traversal of the molecular graph.
• Figure 5: Performance comparison of MolGraph-xLSTM trained with different node orderings. The RDKit Default Order refers to the node sequence provided by RDKit, while the DFS Random Order generates a sequence by performing a depth-first search starting from a randomly selected node. (a) Results on the Sider dataset (classification) using AUROC and AUPRC as metrics. (b) Results on the FreeSolv dataset (regression) using RMSE and PCC as metrics.