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

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.