Skip to content

Omni-Mol: Multitask Molecular Model for Any-to-any Modalities

🕒 Published (v1): 2025-01-01 · Source: NeurIPS · Venue: NeurIPS 2025 · link

Ask a follow-up

Open an assistant pre-loaded with this paper's context.

💬 Ask ChatGPT✦ Ask Claude

TL;DR

Omni-Mol is a unified multimodal LLM for small-molecule tasks categorized into four any-to-any modality types (Mol2Mol, Mol2Text, Mol2Num, Text2Mol), covering 16 tasks with 1.4 million samples. It introduces Gradient Adaptive LoRA (GAL), which dynamically adjusts its scaling factor per task, and Mixture-of-GAL-Experts (MoGE), a MoE expansion that separates specialized and shared knowledge across tasks. With 2.2B parameters, it achieves state-of-the-art on 13 of 16 tasks while outperforming specialist models using 33% fewer parameters.

Problem

Existing molecular LLMs cover only subsets of the four modality-pair types — e.g., PRESTO handles Mol2Num and Mol2Mol but not Mol2Text or Text2Mol — leaving no truly universal generalist. Three compounding difficulties: (1) instruction-tuning datasets for molecular tasks are small and domain-narrow; (2) tasks from different subfields exhibit large distributional shifts that destabilize joint learning; (3) different tasks and even different instances within a task require different intrinsic dimensionalities in representation space, making fixed-rank LoRA adapters suboptimal for multi-task learning.

Method

Data. Omni-Mol constructs a unified instruction-tuning dataset of >1.4 million samples across 16 tasks in four modality-pair types. Molecules are represented in SELFIES (preferred over SMILES for guaranteed RDKit-decodability); graph inputs are derived from SELFIES via RDKit and encoded by MoleculeSTM GNN into tokens projected to the LLM hidden space.

GAL (Gradient Adaptive LoRA). Standard LoRA uses a fixed scaling γ = α/r. GAL replaces this with γ_θ = α/r^p + β, where {α, p, β} are learnable per-adapter parameters. The exponent p models rank-effect curvature, and β provides an additive offset; together they let the effective update magnitude adapt dynamically to the task's intrinsic dimensionality during training, modifying the backward gradient accordingly.

MoGE (Mixture-of-GAL-Experts). To handle inter-task and cross-modal interference, the FFN layers in the last L/4 transformer blocks are replaced by a MoE layer comprising N routed GAL experts (each targeting specialized knowledge) plus one shared GAL expert (capturing cross-task general knowledge). A randomly-initialized router (Kaiming uniform) selects experts per token; a sequence-level load balancing auxiliary loss prevents expert collapse. The MHA layers throughout the network are wrapped with standard GAL (not MoE).

Training. Two stages: Stage 1 trains only the multimodal projector on PubChem molecule-caption alignment. Stage 2 freezes base LLM weights (wrapped by GAL), training GAL parameters, expert router, and projector jointly with the combined language modeling + auxiliary load-balancing loss. Backbone is LLaMA 3.2-1B.

Key Contributions

  • Novel four-way modality taxonomy for small-molecule tasks: Mol2Mol, Mol2Text, Mol2Num, Text2Mol, covering all input-output combinations
  • Largest molecular instruction-tuning dataset: 1.4M samples across 16 tasks with comprehensive preprocessing (3D structure via RDKit, data-leakage filtering)
  • GAL: learnable dynamic LoRA scaling factor (α, p, β) that adapts to task intrinsic dimensionality without increasing adapter rank
  • MoGE: MoE over GAL adapters with shared + routed experts applied to the top-quarter of LLM layers for cross-task/modal specialization
  • Representation convergence analysis showing Omni-Mol's multi-task representations converge (mutual KNN similarity increases with task count) while InstructMol's diverge

Results

  • Forward reaction prediction: Exact match 0.73 vs. PRESTO 0.69, DeepSeekV3 (685B) 0.35
  • Reagent prediction: Exact match 0.23 vs. PRESTO 0.21
  • Retrosynthesis: Exact match 0.57 vs. PRESTO 0.53
  • Catalyst prediction: Exact match 0.52 vs. PRESTO 0.42
  • Mol2Num (MAE): 11.07/0.49/5.89 (Weight/LogP/TPSA) vs. 3D-MoLM 16.58/0.78/10.90 — reductions of 33%, 37%, 46%
  • Description Q&A (BLEU-2): 0.52 vs. 3D-MoLM 0.32 (+63%)
  • Molecular captioning (BLEU-2): 0.529 vs. InstructMol 0.498
  • Yield prediction (B-H/S-M R²): 0.94/0.68 vs. PRESTO 0.94/0.65
  • Achieves SOTA on 13/16 tasks with 2.2B params, outperforming 6.7B specialist models and DeepSeekV3 (685B) in-context on most tasks
  • Clear data and parameter scaling trends observed across LLaMA 1B/3B/8B and 20%–100% data subsets

Limitations

  • Constrained to small molecules; does not cover proteins or protein–small molecule interactions, limiting applicability to broader drug discovery pipelines
  • Computational budget prevented exploration of model scales beyond LLaMA 8B, leaving performance ceiling unknown
  • MoGE is applied only to the top-quarter of layers and a subset of experts (2 routed + 1 shared), a design choice that may not generalize to different architectures without tuning

Relevance to Vision-Language Models

Omni-Mol's architecture directly mirrors the VLM paradigm: a domain-specific encoder (GNN, analogous to a vision encoder) → projector → frozen LLM, trained with instruction tuning. The core technical challenges it addresses — how to adapt a pretrained LLM to heterogeneous modalities with different intrinsic dimensionalities, and how to prevent task conflict during joint training — are identical to those in multi-task VLMs handling diverse visual reasoning tasks. GAL's dynamic rank-effect adaptation is a direct response to the intrinsic dimensionality literature motivating LoRA, offering a transferable technique for multi-task VLM fine-tuning where tasks like OCR, grounding, and captioning similarly differ in representation complexity. The MoGE approach parallels MoE-based VLMs (e.g., Mixture-of-Experts vision-language models) and provides empirical evidence that shared + routed expert decomposition improves cross-modal generalization under distributional shift.