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

Q1: The methodological delta to existing methods is small. E.g., even the early Text+ChemT5 work named as "pioneering" by the authors can already do all four mentioned tasks natively (Text+Chem T5 is not explicitly evaluated on Mol2Num but it can solve it naturally without any architectural modification).

A1: Thanks for the comment. In fact, on our comprehensive general-purpose molecular task dataset, no existing baseline has been able to fully handle the range of tasks. The main challenge lies in the instability of LLM training when faced with highly diverse task types—simply adding more task-specific data does not lead to effective scaling. OmniMol is the first to address this problem by introducing GAL-MoE, a carefully designed mechanism that enables stable and scalable multi-task learning. We also appreciate your mention of Text-Chem T5, and we will provide a detailed comparison with it in the following response.

Q2: Clarity is good overall but can be improved, e.g., the figure ordering is random and they are not referenced in order.

A2: Thanks for the comment. We will adjust the layout of our paper.

Q3: In some tables OmniMol is incorrectly bolded.

A3: Thanks for the comment. We have noticed the typo and revised it.

Q4: Cherry-picking on the competing methods.

A4: Thanks for the comment. We appreciate the reviewer for bringing up this paper, but we respectively disagree with the implication that we are cherry-picking. We had already taken note of this work and reproduced it on our dataset; however, the reproduced results turned out to be quite poor. The detailed experimental results are shown below.

Forward Pred

Reagent Pred

Retrosynthesis

Solvent Pred

Catalyst Pred

Quantum Mechanics Property Prediction Task

Mol2Num

IUPAC2SELFIES

Description Q&A Task

Experimental Procedure Prediction

Molecule Captioning

Text Guided Molecule Generation

Yield Prediction

To ensure fairness and transparency, we have publicly released all model files, hyperparameter configurations, training procedures, and training logs for Text-Chem T5 in our GitHub repository (TextChemT5-v0-main folder). In the above table, Following the Text-Chem T5 paper, we fine-tuned its parameters to ensure the best possible performance. Evidently, Omni-Mol outperforms Text-Chem T5 on almost all tasks by a significant margin. If you have any insights or results from reproducing this work, we would be happy to discuss and compare findings in the subsequent stages.

Q5: Limited reproducibility.

A5: Thanks for the comment and interest. After submitting the paper, we conduct a thorough audit of our code to ensure its correctness and reproducibility. We are highly committed to open-source practices. If you revisit our repository now, you will find that we have updated the code with refined implementation details, added comprehensive comments, provided a complete and well-documented training pipeline, and released our model weights. This allows for a smooth and fully reproducible replication of our results.

Q6: Wrong perception SMILES may often cause errors in RDKit.

A6: Thanks for the comment. The clear statement is that when an LLM outputs SMILES, it is more prone to generating errors, which can result in outputs that cannot be parsed into valid molecules.

Q7: Then, when you tokenize SELFIES with the BPE-based Lama tokenizer, you may distort the structure of the SELFIES such that not all generated SELFIES are valid anymore. E.g., when you tokenize [C][C] into "[C" "][ "C" and "]" you arrive at a vocab that can still be used to produce invalid SELFIES like "][]".

A7: Thanks for pointing out, indeed, there is no perfect 1D representation for molecules, we use SELFIES for its robustness and broader applicability. In fact, many of our baseline results are based on SMILES, meaning they fully leverage the pretrained knowledge embedded within the LLM. In contrast, our method achieves SOTA performance without relying on such pretrained domain knowledge, which further highlights the effectiveness of Omni-Mol.

Q8: P4 L150: Do you mean the hidden size of the LLM's nn.Embedding and the GNN's node embedding size? If really only the sequence lengths differ you just stack horizontally and no padding is necessary (apart from standard padding when you form batches).

A8: Thanks for the question, actually, the padding is due to the varying seuqence length, i.e. differnet number of nodes in molecular embeddings from GNN. We pad them into a regular shape to allow parallel training.

Q9: Figure 4: I am buying the scalability argument but the visualization is unfortunate as the bar seems flat -- I would visualize this on a relative scale (y-axis is the delta compared to the first value)

A9: Thanks for the advice, we will adjust the figure to enable better visualization

Q10: Add useful ablation on SMILES vs SELFIES.

A10: Please refer to Reviewer 8Wa4 A4

Dear Reviewer FtzT,

Thank you for your insightful and constructive reviews. We would like to kindly remind you that the discussion period will conclude in three days. In the meantime, we have provided detailed supplementary experiments and clarifications addressing your concerns.

If you have any additional questions or would like us to conduct further experiments, please do not hesitate to let us know, we would be more than happy to engage in further discussion and provide additional details about our work.

Thank you once again for your time and expertise. We look forward to hearing from you.

Sincerely, Submission 1411 Authors

Q1: No evaluation on generalization ability to diverse tasks.

A1: Thanks for the comment, the main motivation of Omni-Mol is to train a single model capable of solving multiple tasks, in contrast of using dedicated adapters for different tasks. However, we agree with you that generalizing to OOD data is indeed an interesting aspect. Thus, we directly implement the zero-shot evaluation on three different classification datasets. Please refer to Reviewer fT5E A4.

Q2: Doubt on permutation equivariance.

A2: Thanks for the comment. First, the graph features obtained from the GNN encoder are aligned with the textual input. The graph-based model we adopt, MoleculeSTM [1], has undergone extensive pretraining for graph-text alignment, enabling it to provide the LLM with sufficient structural information from the molecular graph. Therefore, permuting the node order does not significantly impact performance. However, to provide more rigorous evidence, we shuffle the order of the nodes and evaluated the performance of Omni-Mol. This was done by applying random permutation to the graph data. The experimental results are as follows:

Forward

Reagent

Retrosynthesis

Solvent

Catalyst

Mol2Num

Quantum Mechanics Property Prediction

Yield

Experiment

Molcap

TextGuidedMolGen

We can draw the conclusion that Omni-Mol is robust to permutation equivariance.

Q3: Utilization of the internal knowledge in LLMs.

A3: Thanks for the comment. The motivation of using SELFIES is to allow higher success rate when converting the 1D representation into 2D graph and 3D graph. There are a line of works using SELFIES as their 1D molecule representation, as suggested by PRESTO [2], we use SELFIES for its broader applicability. In fact, many of our baseline results are based on SMILES, meaning they fully leverage the pretrained knowledge embedded within the LLM. In contrast, our method achieves SOTA performance without relying on such pretrained domain knowledge, which further highlights the effectiveness of Omni-Mol.

Q4: SELFIES could lead misunderstanding.

A4: Thanks for the comment. You’ve made a very insightful point, there is no universally perfect molecular representation. We chose SELFIES primarily due to its higher robustness. In fact, we conduct empirical comparisons and observe that SMILES often leads to parsing failures for a significant number of samples.

Number of failed cases

The experimental results are shown below, which further justifies our decision to use SELFIES.

Q5: Multimodal projector training.

A5: Thanks for the comment. The projector is pre-trained on PubChem for early alignment between LLM and the graph features. In the second stage, projector is still activated and can be updated by multi-task training signals.

Q6: Doubt on the dynamic scaling factor, as the rank is still static.

A6: Thanks for the comment. The rank in LoRA adapters is fixed by design to enable efficient fine-tuning. Our solution is innovative in that it introduces a learnable scaling factor, calculated as $\alpha / r$, which effectively links the scaling factor to the adapter’s rank. When the intrinsic rank varies across tasks, GAL addresses this issue by automatically scaling the learning weights under specific lora rank, thereby achieving better multi-task balance.

Q7: Related works [3, 4, 5] are not experimentally compared.

A7: Thanks for the comment, we have cited all the related works to our paper. Due to time limit, we re-implement Mol-Llama as one additional baseline, we noticed that the 3D GNN used by Mol-Llama supports single molecule input only, we selected tasks that involve only one molecule as input, the results are shown below

DQA

Quantum Mechanics

Mol2Num

Molcap

Omni-Mol slightly falls short on several Mol2Num tasks but exhibits significantly better performance on other tasks, which fully demonstrates the superiority of Omni-Mol as a general-purpose molecular LLM.

Q8: Doubt on novelty of datasets.

A8: Thanks for the comment. Our dataset makes a significant contribution to the field in several key aspects: First, we collect the most comprehensive set of molecular tasks, augmenting instruction-missing data with carefully constructed prompts, and further enhancing all tasks with diverse instructions. Second, we conduct thorough data leakage checks across the entire dataset and remove duplicate samples to build the largest general-purpose molecular dataset to date. Third, our dataset covers the broadest and most complete range of task types, laying a solid foundation for the development of general-purpose molecular AI models.

Q9: Ablation on clipping in GLA.

A9: Thanks for the comment. Clipping is very important in our training. We have added the ablation study results on clipping below. Due to the time limit, we conduct the experiment on a subset of tasks.

Catalyst

Forward

Quantum Mechanics

Molcap

Reagent

Retrosynthesis

Solvent

Yield

From the results, we observe performance drops across several tasks, which aligns with our assumption that different tasks have different optimal rankings.

Q10: Why are other tasks, such as molecular classification, not considered in this work?

A10: Thanks for the comment. First, the DQA task includes approximately 10k entries related to molecular graph classification. For example, there is a sample with instruction: “What is dysidenin classified as chemically?”, and output: “Dysidenin is classified as an organochlorine compound.” Therefore, strictly speaking, we categorize it under the text2text setting.

Q12: What do the left and right y-axes denote in Figure 5 (Middle) and (Right)?

A12: Thanks for the comment. The left y-axis displays average performance excluding yield prediction, while the right y-axis shows yield prediction performance. These were separated due to substantial scale differences that would impair visualization.

[1] Liu, S., Nie, W., Wang, C., Lu, J., Qiao, Z., Liu, L., Tang, J., Xiao, C. and Anandkumar, A., 2023. Multi-modal molecule structure–text model for text-based retrieval and editing. Nature Machine Intelligence, 5(12), pp.1447-1457.

[2] Cao, H., Shao, Y., Liu, Z., Liu, Z., Tang, X., Yao, Y. and Li, Y., 2024. PRESTO: Progressive pretraining enhances synthetic chemistry outcomes. arXiv preprint arXiv:2406.13193.

[3] Park, Jinyoung, et al. "LLaMo: Large Language Model-based Molecular Graph Assistant." NeurIPS, 2024.

[4] Zhang, Juzheng, et al. "UniMoT: Unified molecule-text language model with discrete token representation." ArXiv, 2024.

[5] Kim, Dongki, et al. "Mol-LLaMA: Towards general understanding of molecules in large molecular language model." ArXiv, 2025.

Dear Reviewer 8Wa4,

Thank you for your insightful and constructive reviews. We would like to kindly remind you that the discussion period will conclude in three days. In the meantime, we have provided detailed supplementary experiments and clarifications addressing your concerns.

If you have any additional questions or would like us to conduct further experiments, please do not hesitate to let us know, we would be more than happy to engage in further discussion and provide additional details about our work.

Thank you once again for your time and expertise. We look forward to hearing from you.

Sincerely, Submission 1411 Authors

The OOD performance of LLaMo is coming!

Concern on OOD generalization ability comparison with LLaMo

Thanks for the comment. We re-implement LLaMo and utilize the officially released lamo_checkpoint.ckpt for evaluation. To begin with the conclusion: on classification tasks, it completely fails to generalize, producing a chaotic mix of outputs including numbers, SMILES strings, text, and more. Its OOD generalization ability on regression tasks is also extremely poor.

You may question the fairness of our comparison. However, we have made the entire reproduction process fully transparent — including the reproduction code, datasets, logs, and results — all of which have been uploaded to our codebase. You can directly find the reproduction logs under llamo-reimplement/all_checkpoints, the complete MoleculeNet datasets under llamo-reimplement/bbbp, llamo-reimplement/esol, and llamo-reimplement/lipophilicity, and the environment setup as well as file configurations for the reproduction code in llamo-reimplement/README.md. The process is absolutely reliable.

Specifically, it has no OOD generalization ability on BBBP, a classification task from MoleculeNet. Here is part of the results:

{"prediction": "0.024", "target": "True"}
{"prediction": "2-Methyl-4-phenyl-1,3-benzodioxole", "target": "False"}
{"prediction": "0.035", "target": "True"}
{"prediction": "The molecule is a member of the class of ethylxanthenes that is ethylxanthene in which the hydrogen at position 4 is replaced by an ethyl group. It is a member of the class of ethylxanthenes, an aromatic ether and an organic cation.", "target": "False"}
{"prediction": "The molecule is a metabolite found in or produced by Saccharomyces cerevisiae.", "target": "True"}
{"prediction": "0.038", "target": "True"}
{"prediction": "CC(=O)N1CCN(C(=O)CC2=CC(F)=CC(F)=C2)C(CN3CCC(O)C3)C1", "target": "False"}
{"prediction": "The molecule is a metabolite found in or produced by Saccharomyces cerevisiae.", "target": "False"}
{"prediction": "The molecule is a natural product found in Solanum tuberosum and Triticum aestivum with data available.", "target": "True"}

To ensure the fairness, we even modify the instruction with more detailed information, such as "only output the true or false as the response". This result is shown in llamo-reimplement/all_checkpoints/desc_bbbp_output/lightning_logs/version_1.

We can clearly see that LLaMo just outputs a chaotic mix of outputs including numbers, SMILES strings, text, and more. More details are in llamo-reimplement/all_checkpoints .

As for regression, its outputs still contain 3 mixed samples in ESOL and 24 in Lipophilicity, so we have to remove these noisy entries in order to calculate the MAE. The results are shown below. We can obviously see that the OOD performance is worse than OmniMol.

Table 1: OOD results on ESOL

Table 2: OOD results on Lipophilicity

In conclusion, this further confirms that OmniMol, universally instruction-tuned on 16 molecular tasks, is a powerful generalist molecular model that makes a significant and lasting contribution to the community.

Dear Reviewer 8Wa4,

Thank you very much for your reply to our comments and for recognizing our contributions. We have learned a great deal from your suggestions and feedback, which have been instrumental to improving our work. We sincerely appreciate the time and effort you have devoted to our paper.

Regarding your concerns:

Originality of the proposed dataset

The molecular editing (MolEdit) task is novel in the context of instruction-following datasets, comprising approximately 220K samples. We independently designed and authored the instructions, formulating the dataset explicitly in an instruction-following format. Furthermore, we optimized the computational tool for the MolEdit task’s evaluation metric (see Appendix E.4 and Appendix B.3), ensuring a fairer comparison between our proposed model and the baselines. We also provide case studies in Appendix G.3 to illustrate the behavior of the LLM on this newly introduced task. In addition, we conducted a statistical analysis of the entire dataset, reporting metrics such as Atom Count, Ring Count, Molecular Weight, and Bertz Complexity. This analysis offers a more comprehensive and transparent characterization of the dataset.

Generalizability

We acknowledge OOD performance is an interesting topic. However, Omni-Mol, a generalist model capable of handling diverse tasks (Any-to-Any modality) using a single set of model weights while maintaining competitive performance (as detailed in Appendix I, line 1114), is not designed for OOD generalizability. The architectural design emphasizes balancing multi-task learning and enhancing performance across more than 10 tasks.

Once again, we sincerely thank you for your time, expertise, and constructive engagement during the discussion period.

Dear Reviewer 8Wa4,

Thank you once again for your highly professional comments. Your profound insight into molecular LLMs has been incredibly valuable to us. Once again, we sincerely appreciate the time and effort you devoted to the review.

As the discussion period is scheduled to conclude in 30 minutes, we wish to let you know that we will be standing by. Should you have any final thoughts or questions about our response, we would be more than happy to provide further clarification.

We are truly grateful for your comment!

Q1: Ablation on GNN-based encoder.

A1: Thanks for the comment. We follow a long established line of works [1, 2, 3] that uses GNN features as multi-modal feature, which can enhance the model’s ability to understand molecules’ stuctures. To further validate the role of the GNN component within our framework, we conducted an additional ablation study. Due to time constraints, the experiments were carried out on a subset of the tasks. Specifically, we remove the graph feature derived by projector and re-train a pure text model that uses only features from LLM’s word embedding, the results are shown below.

Molcap

Reagent

Retrosynthesis

Yield

Catalyst

Quantum Mechanics Property Prediction

We can see that Omni-Mol benefits significantly from the GNN encoder on Molcap, Reagent, Yield etc., which particularly rely on the structural information of molecules. For other tasks, we observe similar performance with or without the GNN, indicating that Omni-Mol relies more heavily on textual information to generate responses in those scenarios.

Q2: The experimental results lack error bars or standard deviation.

A2: Thanks for the comment. Following the guidance provided in the cited paper, we varied the sampling temperature and random seed to perform multiple inference runs. The experimental results are shown below.

Q3: Lacks a baseline with pure text.

A3: Thanks for the comment. Please refer to A1 above.

Q4: Runtime to train Omni-Mol

A4: Thanks for the comment. The runtime for 15 tasks in total is 3 * 24 * 8 = 576 A100 80G GPU hours.

[1] Cao, H., Shao, Y., Liu, Z., Liu, Z., Tang, X., Yao, Y. and Li, Y., 2024. PRESTO: Progressive pretraining enhances synthetic chemistry outcomes. arXiv preprint arXiv:2406.13193.

[2] Li, S., Liu, Z., Luo, Y., Wang, X., He, X., Kawaguchi, K., Chua, T.S. and Tian, Q., 2024. Towards 3d molecule-text interpretation in language models. arXiv preprint arXiv:2401.13923.

[3] Cao, H., Liu, Z., Lu, X., Yao, Y. and Li, Y., 2023. Instructmol: Multi-modal integration for building a versatile and reliable molecular assistant in drug discovery. arXiv preprint arXiv:2311.16208.

I want to thank the authors for their rebuttal which has largely addressed my concerns. I am not fully convinced that the GNN encoder is really an essential component of Omni-Mol, given the results above. While the GNN definitely improves scores on some tasks, there are also quantum-chemistry tasks where I would expect the structure to be rather important but Omni-Mol w/o GNN seems to perform better. That said, I do not regard this as a fundamental weakness of Omni-Mol, rather an empirical observation which deserves attention, in particular regarding future work on developing approaches like Omni-Mol further. I encourage the authors to add these results to the paper, accompanied by a nuanced discussion of the importance of the GNN encoder.

Q1: Lack of empirical analysis of robustness to instruction variability.

A1: Thanks for the comment. In our dataset, each input instruction is paraphrased using at least six different prompt templates. During evaluation, we also test the model with various paraphrased prompts rather than relying on a single fixed instruction. You can refer to the repository link for the full set of prompt templates. In fact, each sample is paired with a uniquely crafted prompt, making the model robust to varying "zero-shot prompts". We would be happy to explore further directions and experiments if you have additional suggestions.

Q2: Difficulty in modeling subtle structure-function relationships specific to catalysis.

A2: Thank you for the comment. Our Omni-Mol model has already demonstrated strong performance on 13 datasets. Although it slightly underperforms PRESTO on the catalyst task, it’s important to note that PRESTO benefits from two additional advantages: (1) it undergoes multiple stages of pretraining, including approximately 3 million samples of synthetic procedure descriptions and molecule name conversions, whereas the LLM backbone of Omni-Mol is used without any pretraining; and (2) PRESTO is based on a 7B-parameter model, while Omni-Mol uses a significantly smaller backbone of around 2B parameters, leaving substantial room for future improvement.

Q3: How sensitive is GAL’s performance to hyperparameter tuning?

A3: Thanks for the comment. We tune the rank of GAL and the results are shown in the table below. Due to the time limit, we train the model on 8 tasks.

Catalyst

Forward

Quantum Mechanics Property Prediction

Molcap

Reagent

Retrosynthesis

Solvent

Yield

In the experimental results, we observed that an increase in rank leads to performance improvements, validating the robustness of our method. The growth effect are not particularly evident in Quantum Mechanics Property Prediction task, but significant increases were observed in other tasks.

Q4: How well does the model generalize to truly novel molecular tasks not seen in training?

A4: Thanks for the comment. Generalizing to out-of-domain data is indeed an interesting aspect. Coincidentally, one of the reviewers pointed out that we did not fully consider the molecular graph classification task. This gives us a great opportunity to use it as a testbed to evaluate the generalization capability of our model. We directly implement the zero-shot evaluation on three different classification datasets [1, 2]. Since the original pretrained weights of many baseline models are not publicly available, we use Qwen3, which represents the state-of-the-art on scientific tasks, as our primary comparison.

BBBP

HIV

SIDER

As the experimental results show, Omni-Mol demonstrates certain performance on completely unseen classification tasks, slightly outperforming Qwen3 in some tasks and metrics.

Q5: How graph encoder handle multi-input task?

A5: Thanks for the comment. When multiple molecules involved, we stack them into batch dimensions and feed them to the graph encoder together, the resulting graph features are correspondingly concatenated together.

Q6: Limited explorations on model scaling.

A6: Thanks for the comment. Despite currently using only a 2B model with 5 experts, we have already achieved state-of-the-art performance on the majority of tasks. We believe Omni-Mol has strong potential, and we will include scaling results for larger size models in the next version.

Q7: No mechanism for dynamic prompt generation.

A7: Thanks for the comment. You make a good point regarding the deployment of general-purpose molecular models. In practice, we do apply prompt engineering techniques, for example, calling GPT-4 to refine or polish the instructions. However, these strategies fall outside the scope of the technical contributions of this paper, and thus are not discussed.

[1] Liu, S., Nie, W., Wang, C., Lu, J., Qiao, Z., Liu, L., Tang, J., Xiao, C. and Anandkumar, A., 2023. Multi-modal molecule structure–text model for text-based retrieval and editing. Nature Machine Intelligence, 5(12), pp.1447-1457.

[2] Michael Kuhn, Ivica Letunic, Lars Juhl Jensen, and Peer Bork. "The SIDER database of drugs and side effects". In: Nucleic acids research 44.D1 (2015), pp. D1075–D1079.