Back to feed
Book Open access

UniEdit: A Graph-based MoE Alternative to Sequence Generation for Molecular Editing

Aug 2026 · Proceedings of the 32nd ACM SIGKDD Conference on Knowledge Discovery and Data Mining V.2 · 0 citations · 9 references

TL;DR

This paper proposes UniEdit, a Unified Graph-based Mixture-of-Experts (MoE) Molecular Editing model that offers a robust alternative to LLMs and incorporates a Mixture-of-Experts architecture that dynamically routes tasks to specialized components.

Abstract

Recently, Large Language Models (LLMs) have become the dominant paradigm for molecular editing due to their strong generalization capabilities across diverse tasks. However, treating molecules as 1D text strings (SMILES) introduces significant challenges in controllability and structural validity. In this paper, we question whether sequence generation is truly optimal for this topological task. We propose UniEdit, a Unified Graph-based Mixture-of-Experts (MoE) Molecular Editing model that offers a robust alternative to LLMs. Diverging from the generative approach, UniEdit reformulates molecular editing as a hierarchical node-level classification task. By predicting discrete edit actions (e.g., Add, Remove, Replace) directly on the graph, our model ensures topological precision by design. To handle conflicting multi-objective constraints within a single framework, we incorporate a Mixture-of-Experts architecture that dynamically routes tasks to specialized components. Extensive experiments across 28 diverse tasks demonstrate that UniEdit significantly outperforms sequence-based baselines. Furthermore, a preliminary scaling study reveals that our graph-based approach benefits consistently from increased model capacity, suggesting a scalable path toward general-purpose molecular editing. Our code is available at https://github.com/jiajunyu1999/GraphEditing.

Read PDF

Similar papers

Preprint Jul 2026

Do LLMs Truly Generalize in the Molecular Domain? A Perturbation-Based Analysis

Large Language Models (LLMs) have recently shown promise in molecular discovery, yet a gap remains between their probabilistic nature over discrete sequential tokens and the rigid topological constraints of chemical space. This raises the question of whether molecular LLMs can generalize beyond the local neighborhoods induced by their sequence-based representations. To systematically investigate this question, we introduce a Molecular Perturbation framework that generates syntax-valid structural variants of training molecules under controlled Graph Edit Distance (GED) to probe the manifold regularity of molecular LLMs. Our analysis shows that even a single edit can cause substantial performance drops on common molecular tasks, revealing a narrow local trust region and fragile sensitivity to structural changes. Since similar molecules tend to exhibit similar properties, In-Context Tuning (ICT), which anchors predictions on structurally similar molecules, offers a natural way to mitigate such fragility. Our experiments also examine whether ICT confers robustness under controlled structural perturbations, and the results suggest that it can partially expand the local trust region and offer a promising direction for stabilizing molecular LLMs against structural variation.

Jiatong Li, Weida Wang, Changmeng Zheng et al. · 0 citations
Book Open access Aug 2026

Caduceus: MoE Foundation Models for Unifying Biological and Natural Language

This paper introduces Caduceus, a family of MoE-enhanced foundation models built with a hierarchical pre-training paradigm to jointly integrate biological and natural language, and incorporates a multi-task instruction tuning phase, enabling robust protein parsing and natural language question answering.

Mingze Yin, Yiheng Zhu, Jialu Wu et al. · 0 citations
Jun 2026

EvoOptiGraph: Weakness-Driven Coevolution via Graph-Based Structural Generation for Optimization Modeling

Automating optimization modeling from natural language with large language models (LLMs) faces two key challenges. First, training corpora lack structural diversity. Second, data generation pipelines remain static and decoupled from model learning. To address these challenges, we propose EvoOptiGraph, a novel framework where data and model co-evolve, driven by model weaknesses. EvoOptiGraph represents each mixed-integer linear program (MILP) as an attributed bipartite graph and applies validity-preserving evolutionary operators to generate structurally diverse instances. The evolved graphs are converted into solver code and natural language via deterministic compilation and verified back-translation. Training proceeds in two stages: supervised fine-tuning (SFT) on an initial dataset, followed by reinforcement learning with verifiable rewards (RLVR), where graph-derived weakness signals guide the generation of new instances targeting the model's failures. This forms a closed loop that continuously updates the training distribution. Empirical results on six public datasets show that EvoOptiGraph significantly outperforms larger generalist models, agentic methods, and specialized baselines in accuracy, executability, and generalization. These results demonstrate that targeted data-model coevolution is an effective strategy for improving LLMs on optimization modeling tasks.

Qingcan Kang, Mingyang Liu, Xiaojin Fu et al. · 0 citations
Preprint Aug 2026

Multi-Granular Rationale-Guided Molecular LLM for Property Prediction

Large language models (LLMs) are widely applied across chemical tasks, such as molecular property prediction, which underpins drug discovery. Molecular LLMs represent a molecule through several modalities, notably a 1D SMILES sequence or a 2D molecular graph. Both encode molecular information implicitly, so the contribution of individual substructures remains opaque. Retrieval and augmentation methods add context, but from external sources. However, the cues chemists reason over are the internal substructures that drive a property up or down. We propose MR-MoL, a multi-granular rationale-guided molecular LLM that supplies this evidence directly. A fine-tuned GNN scores each substructure through masking, and the most influential ones are serialized as a ranked, direction-tagged rationale that the LLM reads alongside the SMILES sequence and molecular graph. The rationale spans three levels of granularity: Murcko scaffolds with their side chains, BRICS fragments, and functional groups. This is, to our knowledge, the first method to expose GNN-derived attributions to an LLM as evidence for property prediction. On eight MoleculeNet tasks, MR-MoL achieves the best overall results among generalist models and narrows the gap to specialist models tuned for each task. Five diagnostics further confirm that the model reads the rationale rather than merely benefiting from its presence. Its direction, rank, and substructure each shape the prediction, and its attributions reproduce known structure-property relationships.

Junwoo Park, Minyoung Shin, C. Lee et al. · 0 citations
Open access Jul 2026

FragBERTa: a fragment-aware molecular representation model with sequential attachment-based fragment embeddings

Automated exploration of chemical space through machine learning has become a central strategy in modern drug discovery, where the ability to learn transferable molecular representations directly impacts predictive modeling performance. Recent advances in chemical language models have primarily relied on atom-level string encodings such as SMILES or SELFIES, treating molecules as sequences of individual atoms or characters. However, medicinal chemistry reasoning is often fragment-centric, emphasizing scaffolds, functional groups, and modular substructures rather than isolated atoms. In this study, we introduce FragBERTa, a molecular fragment-aware transformer-based representation language model pretrained using masked language modeling on Sequential Attachment-based Fragment Embedding (SAFE) representations. SAFE decomposes molecules into chemically meaningful fragments while preserving attachment information, enabling the model to learn fragment-level structural patterns within a sequence framework. FragBERTa is pretrained on 11 million unique SAFE sequences extracted from multiple public chemical databases and subsequently finetuned on diverse MoleculeNet benchmarks covering both regression (ESOL, FreeSolv, Lipophilicity, PDBbind) and classification (BACE, BBBP, HIV, Tox21, SIDER) tasks under both random and scaffold splits. Our evaluation shows that FragBERTa achieves competitive performance across tasks, often matching or surpassing prior models. In regression, it demonstrates strong gains in lipophilicity and the best scaffold-split performance on protein-ligand binding affinity within our comparison. In classification, FragBERTa attains the strongest reported random-split ROC-AUC results on BACE, HIV, and Tox21 and the strongest reported scaffold-split ROC-AUC performance on BBBP among the compared baselines, and delivers robust performance across bioactivity, permeability, and toxicity benchmarks, indicating that fragment-based representations effectively capture biologically driven and interaction-sensitive endpoints. Collectively, these findings suggest that fragment-based string representations offer advantages over atom-level representations such as SMILES and SELFIES for scaffold-sensitive and interaction-driven tasks, highlighting the value of fragment-level inductive bias for future molecular foundation models in drug discovery.

Neerav Kaushal, Ajay Mnv Penmatsa · 0 citations