Molecular Generation with State Space Sequence Models

Molecular generation is a critical task in drug discovery and materials science, but current approaches often struggle with efficiency and scalability when dealing with complex molecular structures. This paper aims to address these challenges by training and evaluating models for molecular generation using the MAMBA State Space Model architecture.We develop models with approximately 20M and 90M parameters trained on MOSES and ZINC datasets, respectively, using the Sequential Attachment-based Fragment Embedding (SAFE) representation. We compare MAMBA models against the prevailing Transformer architecture in terms of generation quality and computational efficiency. Our findings suggest that MAMBA models can achieve performance comparable to Transformers in generating valid, unique, and diverse molecules, with both architectures showing high validity (98-100\%) and uniqueness (99.9-100\%) scores. MAMBA models consistently demonstrates lower perplexity and reduced GPU power consumption (up to 30\% reduction) compared to Transformer models. These results indicate that State Space Models may offer a computationally efficient alternative for molecular generation tasks, potentially enabling more efficient processing of larger datasets and complex molecular structures. Notably, the efficiency gains of MAMBA models become more pronounced with longer sequences, suggesting that this architecture could enable the modeling and generation of more complex molecules, such as longer peptides or even DNA sequences. This capability could significantly expand the scope of AI-driven molecular design in drug discovery and materials science.Our study contributes to the exploration of architectural approaches in AI-driven molecular design, highlighting the potential of State Space Models for accelerating drug discovery processes and materials development through improved molecular generation capabilities.