UniMoT: Unified Molecule-Text Language Model with Discrete Token Representation

W1. Limited novelty.

We would like to clarify that the Q-Former is only a part of our VQ-driven tokenizer. UniMoT does not follow the BLIP-2 structure in its entirety. UniMoT adopts a tokenizer-based architecture, which uses discrete tokens, fundamentally different from adapter-based architectures that use continuous embeddings, such as MolCA and 3D-MoLM. We would like to clarify the novel aspects of our work to address your concerns. Due to space constraints, please refer to the global rebuttal for details on the novelty of our work.

W2. Different backbone LLMs.

We have extended our experiments to include additional LLMs to validate the generalizability of UniMoT. The performance is shown in Table 1 of the global rebuttal PDF. Our experiments show that UniMoT performs well across multiple LLMs, including Galactica and Mistral series, demonstrating its robustness and generalizability. This confirms that UniMoT can be successfully applied to other SOTA LLMs.

Llama-2 was chosen initially due to its versatility and performance in handling diverse data types. While it may not be pre-trained specifically on chemistry or biomedical corpora, our pre-training and fine-tuning processes on relevant datasets make sure that the model learns the necessary domain-specific information.

W3 & Q2. Selection of datasets.

The PubChem dataset is significantly larger than the ChEBI-20 dataset. This allows for more comprehensive pretraining. The table below summarizes the size of both datasets:

The ChEBI-20 dataset replaces molecular names with “the molecule,” focusing the model on properties rather than names. However, accurately predicting molecular names demonstrates the model’s understanding of molecular structures. Thus, we conducted the main experiments on PubChem.

For molecule generation tasks, we utilized the Mol-Instructions benchmark. The caption-guided molecule generation task within this benchmark actually uses the PubChem dataset.

W4 & Q1. Models with comparable sizes.

We provide a detailed performance comparison between UniMoT and MolCA using models of comparable sizes. The performance is shown in Table 2 of the global rebuttal PDF. UniMoT consistently outperforms MolCA when using the same Galactica-125M, Galactica-1.3B, and Llama-2-7B backbones. This demonstrates the effectiveness of our proposed UniMoT.

For the modalities, adapter-based architectures require the LLM to directly output SMILES strings to perform molecule generation tasks. This approach relies on strong alignment between SMILES strings and molecule captions during pretraining. In practice, this alignment is challenging to establish, leading to suboptimal performance in text-to-molecule generation tasks.

Our UniMoT treats molecular and textual data equally under a unified token representation. This enables the LLM to unify the modalities under a shared autoregressive training paradigm. The molecule tokens encapsulate high-level molecular and textual information, providing a richer representation than SMILES strings alone.

Q3. BioT5 and BioT5+ baselines.

We have included BioT5 and BioT5+ in our comparative analysis. We selected the molecule captioning task from the molecule comprehension tasks with the molecule generation tasks. The results are presented in Tables 3 and 4 of the global rebuttal PDF. These comparisons demonstrate the comparable performance of our model with BioT5 and BioT5+ in comprehension and generation tasks. In future revisions, we will incorporate performances of BioT5 and BioT5+ to provide a more comprehensive comparison.

W5 & Q4. Ablation studies on fine-tuning strategy and codebook size.

We would like to highlight that we have indeed provided the performance of full fine-tuning of Llama-2-7B in Table 6 of our manuscript. The pre-trained Llama-2-7B model was directly fine-tuned on the task-specific dataset without any specialized techniques.

The choice of 2048 for the molecule codebook size is based on a balance between model complexity and performance. A larger codebook could potentially capture more subtle interactions between molecules and text. However, there may be some codes that are not often used on large codebooks. A smaller codebook may result in nearby embeddings being assigned the same code, which reduces the granularity of the representation. The performance with different codebook sizes on the molecule captioning task is shown in Table 5 of the global rebuttal PDF. The results demonstrate that the codebook size of 2048 consistently provides the best performance.

W6. Data leakage.

We appreciate the concern regarding potential data leakage. We take steps to ensure that data leakage is avoided throughout our training process. We list the datasets employed at each stage:

• Stage 1: Causal Q-Former Pretraining
  • Datasets: PubChem pretrain subset
• Stage 2: Molecule Tokenizer Pretraining
  • Datasets: PubChem pretrain subset, CheBI-20 train subset
• Stage 3: Unified Molecule-Text Pretraining
  • Datasets: PubChem pretrain subset, CheBI-20 train subset
• Stage 4: Task-Specific Instruction Tuning
  • Datasets: PubChem, PCdes, MoMu, MoleculeNet, QM9, USPTO (all train subsets)

W7. The terminology.

Thank you for your feedback regarding the terminology. When we refer to “a unified molecule-text LLM,” we mean that the LLM can process input sequences that may be purely molecular, purely textual, or multi-modal (combining both molecules and text). This capability is achieved through treating molecular and textual data equally under a unified token representation. The LLM learns to predict the next token in a sequence, regardless of whether the sequence consists of molecules, text, or a combination of both.

Thanks for the clarification and extended experiments.

• W1: As the author explained, the model structure is exactly similar to MolCA and 3D-MoLM. In this case, the novelty should be further discussed and I personally don’t see enough novelty in the model structure.
• W3: I do not agree with the statement of the dataset size. The training part of PubChem is only 12000 samples, while ChEBI-20 is even double. The performance comparison is unfair because you introduce more related data.
• W6: Still not guarantee that the data will not leak during the pretraining stages.
• W7: I don’t think so.

Thank you for your thorough review and valuable feedback. Below, we address the further concerns you have raised in detail.

I personally don’t see enough novelty in the model structure.

We respectfully disagree with the assessment that the model structure is similar to MolCA and 3D-MoLM. The architecture of UniMoT is fundamentally different from these models. Specifically:

1. Model Architecture:

   • MolCA and 3D-MoLM: These models follow the Q-Former architecture, as depicted in Figure 1(b) of our paper. In this architecture, the input to the LLM is molecule embeddings.
   • UniMoT: Our model employs a VQ-VAE architecture, as illustrated in Figure 1(c). The key difference lies in the input to the LLM: UniMoT uses molecule tokens instead of embeddings.

2. Supervision Signal for Molecule Modality:

   • MolCA and 3D-MoLM: These models do not incorporate any supervision signal for the molecule modality. This limitation significantly constrains their capacity and effectiveness.
   • UniMoT: In contrast, UniMoT provides direct supervision for the molecule tokens through autoregressive pretraining. This allows for stronger generation abilities and enables the molecule tokens to be decoded into molecules.

3. Text-to-Molecule Generation:

   • MolCA and 3D-MoLM: These models rely solely on SMILES strings for text-to-molecule generation tasks. This heavily relies on strong alignment between SMILES strings and molecule captions during pretraining. In practice, achieving this alignment is challenging, leading to suboptimal performance in text-to-molecule generation tasks.
   • UniMoT: Our model not only supports text-to-molecule generation but also benefits from the direct supervision of molecule tokens, leading to improved accuracy in generation tasks. The molecule tokens encapsulate high-level molecular and textual information, providing a richer representation than SMILES strings alone.

4. Cross-Modal Alignment:

   • MolCA and 3D-MoLM: These models rely heavily on the Q-Former for cross-modal alignment, which can be restrictive in certain scenarios.

   • UniMoT: Our approach allows for the Q-Former to be discarded entirely. Instead, we quantize the molecule features into molecule tokens using the VQ-VAE pipeline, with a codebook size of 1024. We report the performance of the molecule captioning task on the PubChem dataset. This demonstrates that the performance remains comparable even without the Causal Q-Former.

     	BLEU-2	BLEU-4	ROUGE-1	ROUGE-2	ROUGE-L	METEOR
     UniMoT w/o Causal Q-Former	28.1	20.8	33.2	20.7	30.1	30.8
     UniMoT w/ Causal Q-Former	31.3	23.8	37.5	23.7	33.6	34.8

I do not agree with the statement of the dataset size.

• The PubChem training set is divided into two distinct subsets: a pretraining subset and a training subset. These subsets do not intersect, which ensures that data leakage is avoided. Specifically, the pretraining subset (301,658 molecule-text pairs) is used exclusively for pretraining the model. And the training subset (12,000 molecule-text pairs) is reserved for instruction tuning.

• ChEBI-20 only has a single training set, which consists of 26,407 pairs. This set is used solely for pretraining.

When comparing dataset sizes, it is important to consider the relevant subsets. The PubChem pretraining subset (301,658 pairs) is significantly larger than the ChEBI-20 training set (26,407 pairs). This is why we chose to use PubChem for pretraining, following a practice as used in InstructMol, MolCA, and 3D-MoLM.

The performance comparison is unfair because you introduce more related data.

We would like to clarify why the performance comparison is indeed fair.

The ChEBI-20 dataset replaces molecular names with “the molecule,” which shifts the model's focus toward predicting molecular properties rather than specific names. However, accurately predicting molecular names is crucial for demonstrating a model’s understanding of molecular structures. To address this, we conducted experiments on both PubChem and ChEBI-20, and the results for ChEBI-20 are presented in the Appendix, while the main experiments on PubChem are included in the main text.

All baseline models were evaluated on the same PubChem test set on the experiments in the main text. Since recent baselines like MolCA, 3D-MoLM, and InstructMol also use PubChem for pretraining, our comparison is consistent and fair.

Still not guarantee that the data will not leak during the pretraining stages.

We can assure that data leakage is effectively prevented during the pretraining and instruction tuning stages. Here's how we ensure this:

• The PubChem dataset is carefully divided into two distinct subsets. The pretraining subset is exclusively used during the pretraining stages. And the training subset is reserved for instruction tuning. These subsets do not intersect, ensuring that data used for instruction tuning is entirely separate from the data used during pretraining.

• The additional datasets used during the instruction tuning stage (PCdes, MoMu, MoleculeNet, QM9, USPTO) are also separate and distinct from the datasets used in the pretraining stages. These datasets do not appear in the pretraining phase, ensuring that there is no overlap or data leakage.

Thank you for your interest in our work. Your insightful comments have contributed to strengthening our paper. Should you have any questions or require additional information, please do not hesitate to contact us.

Thanks to the authors for explaining my concerns.

However, I hope to clarify my views again considering the model structure.

1. The model structure and model components proposed in this paper have already been used in the field of Computer Vision. The authors brought nothing new. In other words, the authors assemble the Q-Former with the VQ-VAE.
2. Although the authors claim they have better "cross-modal alignment", the benefits are not satisfying at all. We can call it a small improvement, but I don't see enough potential for this work to further benefit the community, as well as the development of computational biology. Frankly speaking, this work still does not align molecular graphs well with texts, and it can not help chemists save their workload.

The dataset size

I never see anyone compare their pretraining dataset with the downstream training set.

Unfair comparison

The responses did not address the problem well. Let's change the question. When comparing the downstream performance between Mistral-7B and Llama-2-7B, Mistral-7B shows better performance. Then we could say that the pre-training of Mistral-7B is better. However, when both the two models are adopted in your model structure, is it still fair to say it is your model structure that brings the improvement? No. The performance gain is partly from pre-training. When you compare your method with the previous baselines, how could you ensure that your model structure is better?

data leakage

I don't think the authors answered my question. Yes, the datasets are not used in pertaining. However, there are molecules in downstream tasks that appear in your pre-training set. When you pre-train the model, the information has already been leaked. It is worth noting that in the work of BioT5, the authors did a special operation that removed all the overlapping molecules from the pre-training set.

Overall, I strongly believe this work is not qualified to be accepted, and these issues should be further addressed. I will maintain my score.

W1. Text-to-molecule generation tasks.

As Reviewer d9hA pointed out, adapter-based architectures can also perform text-to-molecule generation tasks. We will revise the manuscript accordingly. However, our contention is that such methods typically do not perform as well as our approach, as demonstrated in Table 4 of our paper. Here, we elaborate on the key differences and advantages of tokenizer-based architecture over adapter-based architectures:

• Limitations of Adapter-Based Architectures: Adapter-based architectures require the LLM to directly output SMILES strings to perform molecule generation tasks. This approach relies heavily on strong alignment between SMILES strings and molecule captions during pretraining. In practice, this alignment is challenging to establish, leading to suboptimal performance in text-to-molecule generation tasks.

• Advantages of Tokenizer-Based Architecture: Our method leverages a tokenizer to convert molecular features and text captions into molecule tokens. These tokens encapsulate high-level molecular and textual information, providing a richer representation than SMILES strings alone. By linking molecule tokens generated by the tokenizer to the molecule captions during molecule-to-text and text-to-molecule pretraining, our approach ensures that the model learns to understand and generate molecule tokens in an autoregressive manner.

W2. Alignment methods.

We outline the purpose and necessity of each alignment method below:

• Molecule-Text Contrastive Learning: Molecule-Text Contrastive Learning is employed solely in Stage 1: Causal Q-Former Pretraining. The primary objective here is to align the molecule features with the text captions, ensuring that the queries incorporate aligned molecular and textual information. This contrastive learning helps the Causal Q-Former to understand and relate molecular structures with their corresponding textual descriptions.

• Aligning MSE Loss: The aligning MSE loss is used exclusively in Stage 2: Molecule Tokenizer Pretraining. Its purpose is to align the discrete latent space of molecule tokens with the continuous latent space of a molecular generative model. The alignment ensures that the tokenized representation can be mapped to the decoder's latent space.

• SMILES Decoder Reconstruction: The reconstruction is primarily relevant during inference. During training, we focus on obtaining the reconstruction loss (the aligning MSE loss) to perform back-propagation. The actual reconstruction of the molecule via the SMILES decoder is done during inference to generate the final molecular output.

W3. Fixed size of queries.

The design of our Causal Q-Former effectively addresses the concern regarding the fixed size of the queries through attention mechanisms:

• Self-Attention Mechanism: Although the query size is fixed, the Causal Q-Former employs a dynamic self-attention mechanism for the queries. This allows the model to adaptively capture molecular details necessary for understanding and generating textual descriptions.

• Cross-Attention Mechanism: The Causal Q-Former also utilizes a cross-attention mechanism where the queries selectively attend to different aspects of the molecule features. This allows the model to adaptively focus on different parts of the molecule based on its complexity and size.

We also conducted an ablation study to evaluate the performance of UniMoT with different query sizes, as presented in Table 6 of the global rebuttal PDF. The results show that increasing the query size improves performance, with the best results observed at a query size of 32. However, since a query size of 32 requires significantly more training time and memory, we still opt to use a query size of 16.

Q1. Datasets usage.

Below, we provide a list of the specific datasets employed at each stage:

• Stage 1: Causal Q-Former Pretraining
  • Datasets: PubChem pretrain subset
• Stage 2: Molecule Tokenizer Pretraining
  • Datasets: PubChem pretrain subset, CheBI-20 train subset
• Stage 3: Unified Molecule-Text Pretraining
  • Datasets: PubChem pretrain subset, CheBI-20 train subset
• Stage 4: Task-Specific Instruction Tuning
  • Datasets: PubChem, PCdes, MoMu, MoleculeNet, QM9, USPTO (all train subsets)

Q2. Different outputs in Figures 1 and 3.

Yes, it is consistent. Figure 1 is designed to convey the core idea of our method. It shows the autoregressive prediction of molecule and text tokens. When data is organized as interleaved molecule and text tokens, the model can predict the next token (whether it is a molecule or text) indiscriminately. Each token prediction is supervised, and learning occurs across both types of tokens.

Figure 3 illustrates the separate supervision of molecule and text tokens adopted in Stage-3, which is the practical implementation of our training process. Instead of using an interleaved organization of molecule and text tokens, we use the dataset as provided, consisting of molecule-text pairs. This involves two tasks:

• Molecule-to-text autoregression: Using molecule tokens as a prompt to supervise and generate text.

• Text-to-molecule autoregression: Using text as a prompt to supervise and generate molecule tokens.

We utilize the PubChem and CheBI-20 datasets for the unified molecule-text pretraining.

Q3. Datasets with baselines.

We provide a list of the datasets used for pretraining and fine-tuning UniMoT and the baselines:

• UniMoT:

  • Pretraining: PubChem and CheBI-20

  • Fine-Tuning: Specific task datasets

• MolT5:

  • Pretraining: Colossal Clean Crawled Corpus (C4) and ZINC

  • Fine-Tuning: CheBI-20

• InstructMol:

  • Pretraining: PubChem

  • Fine-Tuning: Specific task datasets

• MolCA:

  • Pretraining: PubChem

  • Fine-Tuning: PubChem

• 3D-MoLM:

  • Pretraining: PubChem

  • Fine-Tuning: 3D-MoIT

W1. Limited novelty.

While we use components from existing works like Q-Former, the molecule encoder, and the SMILES encoder and decoder, we adopt a tokenizer-based architecture that uses discrete tokens. This is fundamentally different from adapter-based architectures that use continuous embeddings. Figure 4 is directly adapted from Figure 2 of Q-Former's original paper to illustrate related concepts and objectives of Causal Q-Former. We would like to clarify the novel aspects of our work to address your concerns. Due to space constraints, please refer to the global rebuttal for details on the novelty of our work.

W2 & Q1. The validity of generated molecules.

Our model is pre-trained on extensive datasets of valid molecules, such as PubChem, which contain 0.3M of structurally diverse and chemically valid SMILES strings. This large-scale pre-training helps the model learn a robust representation of valid molecular structures. After the initial pre-training, we fine-tune our model on datasets composed exclusively of valid molecules. This fine-tuning process reinforces the learned characteristics of valid SMILES strings.

Our empirical validation shows that while a few invalid molecules are generated when testing, these occurrences are rare and do not affect the overall validity rate.

W3 & Q2. Ablation study on codebook size and query size.

The choice of 2048 for the molecule codebook size is based on a balance between model complexity and performance. A larger codebook could potentially capture more subtle interactions between molecules and text. However, there may be some codes that are not often used on large codebooks. A smaller codebook may result in nearby embeddings being assigned the same code, which reduces the granularity of the representation.

We conducted experiments with different codebook sizes and report the performance of the molecule captioning task on the Pubchem dataset. The performance with different codebook sizes is shown in Table 5 of the global rebuttal PDF. The results demonstrate that the codebook size of 2048 consistently provides the best performance.

We also conducted an ablation study to evaluate the performance of UniMoT with different query sizes, as presented in Table 6 of the global rebuttal PDF. The results show that increasing the query size improves performance, with the best results observed at a query size of 32. However, since a query size of 32 requires significantly more training time and memory, we still opt to use a query size of 16.

W4 & Q3. Comparison with different tokenizers.

Difference from Existing Tokenizers:

• The cited works [4] and [5] focus on SMILES tokenizers, which primarily tokenize the SMILES strings of molecules. While SMILES strings capture sequential information, they do not fully encapsulate the structural intricacies of molecules necessary for understanding and generation tasks.

• Our molecule tokenizer, on the other hand, is designed to tokenize molecules into molecule tokens. This involves incorporating textual information and structural information embedded in molecule features. Our tokenizer is specifically designed to generate tokens that LLMs can understand and generate, similar to text tokens.

We report the performance of the molecule captioning task on the Pubchem dataset. Note that in the SMILES tokenizer baselines, we do not incorporate molecule features; we only use the SMILES strings as input for the LLM.

W5 & Q4 & Q5. The synonym problem.

For our pretraining, we utilize the molecule-text pairs as provided by the PubChem dataset without any modification.

• Molecule-to-Text Autoregression:

  • For molecule-to-text autoregression, our model uses the canonical SMILES representation, which uniquely defines the chemical structure of a molecule.

  • We also incorporate contextual tokens derived from contextual molecule embeddings. These tokens help capture the meaning of molecule names within their specific context, allowing the model to understand that different names can refer to the same chemical entity.

• Text-to-Molecule Autoregression:

  • For text-to-molecule autoregression, our model generates molecule tokens that have been aligned with textual descriptions during pretraining.

  • We conducted several tests, which showed that the model's output remains consistent regardless of the synonyms used in the input. For instance, when we replace “Acetylcarnitine” with “L-Carnitine, acetyl ester” or “O-Acetyl-L-carnitine,” the output remains consistent.

  • While the model generally produces consistent outputs regardless of the synonym used, slight variations may occur due to differences in how the model processes each synonym.

W1. Including CLAMP and fingerprints baselines.

Thank you for your feedback and for suggesting additional baselines such as CLAMP and standard molecular fingerprints on the molecular property prediction task. We will include these baselines in our revised manuscript to provide a more comprehensive evaluation. We will also update the ToxCast results to reflect the performance of KV-PLM linear probing.

Q1. The reconstruction rate of the SMILES decoder.

We measured the reconstruction rate on a test set of SMILES strings. The reconstruction rate is defined as the percentage of SMILES strings that are exactly reproduced by the decoder from their encoded representations. Our experiments show that the SMILES decoder achieves a reconstruction rate of 96% on the PubChem test set of 2,000 valid molecules.