Atomas: Hierarchical Adaptive Alignment on Molecule-Text for Unified Molecule Understanding and Generation

I want to first thank the authors for their responses. The rebuttal has already adressed some of my concerns.

However, the performance on the molecule-caption translation task is only on the ChEBI-20 dataset, which can not reveal the generalization of the proposed method regarding this downstream task. Meanwhile, the performance of Atomas is still close to the previous baselines. I would like to recommend that the authors add the Mol-Instruction [1] and PubChem [2] datasets for further demonstration. I am happy to increase my score if we see similar improvements on these datasets.

References

[1] Fang, Y., Liang, X., Zhang, N., Liu, K., Huang, R., Chen, Z., ... & Chen, H. (2023). Mol-instructions: A large-scale biomolecular instruction dataset for large language models. arXiv preprint arXiv:2306.08018. [2] Liu, Z., Li, S., Luo, Y., Fei, H., Cao, Y., Kawaguchi, K., ... & Chua, T. S. (2023). MolCA: Molecular graph-language modeling with cross-modal projector and uni-modal adapter. arXiv preprint arXiv:2310.12798.

W4&Q1: The compared baseline models.

Reply: Thank you for suggesting. We believe the baseline models we have included represent the latest published methods and demonstrate strong performance. However, we agree that ICMA is a relevant baseline and have incorporated a comparison with ICMA, as shown in the table below. We will also include a citation to ICMA in the revised version of the paper.

Based on the data presented, although Atomas-large has only approximately 12% of the parameters of ICMA, it still achieves superior performance.

Regarding BioT5 and BioT5+, we acknowledge their use of SELFIES as a molecular representation, which ensures 100% validity and contributes to their reported exactness (0.413) and validity (1.00) metrics. However, we attribute these advantages to the choice of molecular representation rather than methodological differences. Since SMILES is the dominant molecular representation, Atomas focuses on SMILES and compares only with methods that also use SMILES. Despite the exactness and validity metrics of SELFIES-based methods, Atomas achieves superior performance in quantitative metrics compared to the BioT5 and BioT5+, as shown in the table below:

Additionally, we have already included qualitative comparisons with BioT5 in Figures 6 and 7, which analyze molecule captioning and generation tasks. These figures show that global alignment methods like BioT5 struggle to differentiate between enantiomers such as "D-glutamate" and "L-glutamate." In contrast, Atomas generates more accurate and detailed molecular descriptions, highlighting the effectiveness of our hierarchical alignment model.

Q3: Scaling Atomas to the decoder-only models.

Reply: Thank you for your insightful feedback. We believe that the Atomas Hierarchical Adaptive Alignment can indeed be applied to decoder-only models as a plug-in module. The hierarchical alignment can be performed directly on the representations of the decoder, further highlighting the versatility and adaptability of our approach. In alignment with the goal expressed in our title, our aim was to develop a unified model capable of addressing both understanding tasks (via the encoder) and generative tasks (via the encoder-decoder combination).

Regarding the choice of backbone LLM, we respectfully believe that our comparison remains fair, as each method employs distinct strategies for molecule-text alignment, input modalities, and model architectures. Moreover, the parameter sizes of the models vary substantially. For example, ICMA employs a 7B parameter backbone, while Atomas utilizes fewer than 1B parameters. Given these differences, ensuring identical backbone LLMs across all methods in the comparison may not always be feasible or practical. To provide additional clarity, we have discussed the variations among existing molecule-text alignment methods in Appendix A.1 (Table 10), detailing their input formats, alignment strategies, training approaches, and application tasks.

To ensure complete fairness, we have used the same backbone LLM as MolT5 in our experiments. As shown in Table 3 of the paper, our Atomas-base model outperforms the MolT5-large model while using only 35% of its parameters and without initial pretraining. This highlights the efficiency and effectiveness of our proposed framework. We will further emphasize this point in the revised version of the paper.

We thank reviewer 1BNJ for the valuable feedback and for recognizing the innovation and effectiveness of our method. Please kindly find the detailed responses point-by-point below. Any further comments and discussions are welcomed!

W1: The description of atom level.

Reply: Thank you for your insightful feedback. We appreciate your concern regarding the terminology used to describe the "atom level" alignment. Our intention with the term "atom level" was to refer to alignment at the smallest unit of division within the molecule. However, we acknowledge that the inclusion of numbers in SMILES strings, which are not strictly atoms, may lead to potential confusion.

To address this, we will revise the terminology in our paper to "word level," which we believe more accurately reflects the nature of the alignment process while avoiding ambiguity.

W2: The report metrics of MolCA.

Reply: Thank you for your careful review and for bringing up the performance metrics of [1] in Table 3. We would like to clarify that the metrics you referred to are from an earlier version of [1]. Our reported metrics are based on the updated versions of [1], as reflected in their arXiv versions V3 or V4 and ACL Anthology version V2. If it is convenient, we encourage you to check these updated sources for the latest results. We have included the latest MolCA paper below for your reference.

With this clarification, we would like to emphasize that our method achieves superior performance while relying solely on SMILES, in contrast to models that depend on both SMILES and graph representations for certain tasks.

[1] MolCA: Molecular Graph-Language Modeling with Cross-Modal Projector and Uni-Modal Adapter. EMNLP, 2023.

W3&Q2: Confusion about the alignment process.

Reply: Thank you for your thoughtful feedback. We would like to clarify how the three-level alignment works in our Hierarchical Adaptive Alignment (HAA) model. As illustrated in Figure 2 and the Algorithm provided in Appendix D of our paper, the alignment between molecules and text is performed at three levels: atom, fragment, and molecule.

At the fragment level, the process consists of three steps:

1. Step1: Assignment Step: Tokens at the atom level are assigned to form clusters.
2. Step2: Merge Step: These clusters are further merged to generate the fragment representations.
3. Step3: Alignment Step: The Weighted Alignment Module (WAM) aligns the fragments between the two modalities.

At the molecule level, fragments are further clustered into higher-level representations using the same Assignment and Merge steps. The WAM then aligns these higher-level molecule representations between the text and molecule modalities.

In summary:

• Steps 1 and 2 (Assignment and Merge) are executed twice—once at the fragment level and once at the molecule level.
• Step 3 (Alignment) is executed three times—once at each level: atom, fragment, and molecule.

By iteratively applying steps 1 and 2 along with the WAM, we construct a hierarchical alignment structure that spans three levels of abstraction, enabling the model to capture fine-grained alignments across varying levels of detail.

We hope this explanation clarifies how our three-level alignment works and addresses your concerns.

We are encouraged by Reviewer 1BNJ's positive feedback and agree that adding additional experiments can better showcase Atomas' generalization ability.

In response to your suggestion, we conducted an additional experiment to compare Atomas' performance with the state-of-the-art model, ICMA, on the PubChem dataset provided by MolCA. Specifically, we trained Atomas-large using only the PubChem324K training dataset, aligning the experimental settings with ICMA by ignoring pretaining data. The results are presented in the table below.

Supplementary experiment for Cap2Mol results on PubChem324k test dataset

Based on the data presented, Atomas achieves significant improvements on most metrics over ICMA, including a 63.11% improvement on the BLEU metric and a 46.54% improvement on Levenshtein compared to ICMA. This demonstrates Atomas' superior performance on the PubChem dataset, highlighting its robust generalization capabilities. Given the extensive size of the Mol-Instruction dataset (297K training samples vs. 12K training samples in PubChem) challenge for us during the rebuttal period, we prioritized the PubChem dataset during the rebuttal to demonstrate Atomas's generalization within the available timeframe.

We hope this response addresses your concerns. Please let us know if you have any further questions or concerns.

W3: The hierarchical feature construction is effective but needs more explicit rationale explaining.

Reply: Our approach is motivated by the observation that text-molecule multimodal data naturally exhibits a hierarchical structure. For instance, textual descriptions and molecular representations can both be understood at multiple levels of granularity, such as word-atom, phrase-fragment, and paragraph-molecule. Recognizing this inherent hierarchy inspired us to adopt hierarchical clustering [1] for fine-grained alignment. we employ a "bottom-up" strategy, where each observation begins in its own cluster, and clusters are progressively merged as the hierarchy ascends.

To further validate the rationale for our approach, we have provided a visual example in Figure 5 of the paper. Figure 5 illustrates how our method constructs hierarchical features. The molecule shown in Figure 5 is formed by combining atoms at positions 0–15 and 16–26 through dehydration condensation. At the fragment level, our model clusters atoms (or words) into functional groups (or phrases) using the Adaptive Polymerization Module. As the hierarchy progresses to the molecule level, fragments are further clustered to form monomer-like structures.

Additionally, the performance improvements demonstrated in the ablation study in Table 7 further validate the effectiveness and practicality of this hierarchical feature construction. These empirical results underscore the value of modeling the hierarchical structure inherent in multimodal molecular data.

[1] Algorithms for hierarchical clustering: an overview. Wiley Interdisciplinary Reviews: Data Mining and Knowledge Discovery 2011.

Q1: The reason for using both global alignment and molecule level alignment.

Reply: Thank you for your thoughtful feedback. We would like to clarify that these two alignments serve distinct purposes in our architecture. The Global Alignment and Hierarchical Adaptive Alignment ensures that the model captures both local features and global information effectively.

In Global Alignment, we aggregate all tokens to obtain the global feature $(bsz, 1)$, where $bsz$ represents the batch size. This is akin to the use of the $[CLS]$ token in BERT-like models to capture global information.

In contrast, Molecule Level Alignment operates at a higher level of granularity. As illustrated in Figure 2 of the paper, molecular fragments, and textual paragraphs are clustered into higher-level groups. At this level, tokens are aggregated into $(bsz, n_{s})$ and $(bsz, n_{t})$, where $n_{s}$ and $n_{t}$ denote the number of molecular and text clusters, respectively. For example, as shown in Figure 2, alignment is performed between 2 molecular clusters and 1 text cluster.

We name this the "Molecule Level" alignment because, at this stage, the model is capable of capturing relationships between complex molecules composed of multiple submolecular units. For example, as depicted in Figure 5, Atomas effectively clusters the trans-sinapic acid and beta-D-glucose into distinct classes, reflecting their independent molecular identities despite being part of a larger structure. This ability to recognize and align such complex molecular components is a unique feature of our approach.

Q2: Try using other clustering algorithms.

Reply: Thank you for your thoughtful question. Yes, we have explored general clustering algorithms, including general approaches such as k-means clustering and molecular-specific methods like BRICS decomposition.

1. BRICS Decomposition: We applied the BRICS decomposition method to fragment the molecules shown in Figure 5. The fragments produced by BRICS include:

• [1*]C(=O)C[7*]
• [13*][C@@H]1OC@HC@@HC@H[C@H]1O
• [16*]c1cc([16*])c(O)c([16*])c1
• [3*]OC
• [3*]O[3*]
• [7*]C[8*]
• [8*]CO

However, BRICS introduces characters not present in the original SMILES strings and does not allow for hierarchical decomposition, limiting its direct replacement with the Adaptive Polymerization Module in Atomas. Additionally, as discussed in [1], BRICS typically generates only a few fragments for each molecule, resulting in coarse-grained representations that may lack the level of detail required for effective language modeling.

2. K-Means Clustering: We also considered k-means clustering. However, this method requires specifying the number of clusters in advance. Unlike images with fixed dimensions (e.g.,$512*512$), molecules and text exhibit variable lengths in their SMILES and textual representations, leading to varying cluster requirements. Therefore, predefining a fixed number of clusters is not practical for molecular data. Instead, we chose to use the density peak-based clustering algorithm, which dynamically assigns tokens to clusters without requiring a predefined cluster count.

[1] Multi-objective drug design based on graph-fragment molecular representation and deep evolutionary learning. Frontiers in Pharmacology 2022.

We sincerely thank Reviewer N9hg for their valuable feedback and for recognizing the comprehensive methodology and experiments, the intuitive and clear presentation, and the extensive ablation studies contributed by our method. Please find our detailed point-by-point responses below. We welcome any further comments and discussions!

W1: Contrastive-autoregressive synergy boosts representation has been proved.

Reply: We sincerely thank the reviewer for their insightful comment and for referencing [1]. We have cited [1] in our paper, as it inspired us to adopt the use of soft labels. However, we respectfully note that [1] primarily demonstrates the effectiveness of the combined optimization of contrastive learning alignment and autoencoder learning, rather than autoregressive learning. Specifically, the pre-training objectives in [1] are contrastive learning and masked language modeling, distinguishing its focus from the contrastive-autoregressive optimization explored in our work.

To further clarify the effectiveness of the contrastive-autoregressive synergy in boosting representation, we provide the following in-depth discussion of this joint optimization strategy:

• Model Perspective: As suggested by the existing study [2], attention-based generation tasks essentially perform a form of soft alignment. During the generation process, the attention mechanism facilitates a mutual translation between text and SMILES, reinforcing the semantic consistency between the textual description and the molecular structure it represents. Concurrently, representation learning bridges the domain gap between text and SMILES, enhancing the caption/generation task.

• Data Perspective: The captioning and generation tasks may provide complementary information for learning molecular representations. These tasks necessitate the model to learn the mapping between text and molecular domains, which allows the model to grasp the intricate relationship between textual descriptions and molecular structures, thereby enriching the quality of the learned molecular representations. Hierarchical alignment further aids in capturing local data pair relationships, benefiting the generation process.

In summary, we believe that Atomas provides new insights into molecule generation tasks through its joint optimization strategy. This approach enables the model to learn more effective molecular representations by leveraging the complementary information provided by each task. The experimental results presented in Table 8 and Figure 4 of the paper further substantiate the effectiveness of this strategy.

[1] Align before fuse: Vision and language representation learning with momentum distillation. NeurIPS 2021, Junnan Li.

[2] Neural Machine Translation by Jointly Learning to Align and Translate. ICLR 2015, oral, Yoshua Bengio.

W2: Using all tokens to calculate similarity and the use of the weighted alignment module.

Reply: Thank you for your thoughtful feedback. In our method, similarity is calculated using all tokens only at the first level: the atom level. As acknowledged in our paper, we employ the Weighted Alignment Module (WAM) inspired by [1] to perform alignment.

However, we would like to highlight an important distinction. The weighted alignment in [1, 2] primarily focus on alignment at a single level, utilizing token-wise interactions. In contrast, our approach extends beyond token-wise interactions by introducing a hierarchical structure. This enables us to perform alignment in a set-wise manner, where groups of semantically similar tokens interact. Furthermore, to the best of our knowledge, we are the first to apply such an alignment manner specifically in the text-to-molecule domain. We are also actively exploring alternative alignment modules tailored for set-wise interactions to further enhance the novelty and effectiveness of our framework. Thank you again for your insightful comments!

[1] Disentangled Representation Learning for Text-Video Retrieval. arxiv 2022.

[2] FILIP: Fine-grained Interactive Language-Image Pre-Training. ICLR 2022.

W2: Adding theoretically computational comeplexity.

Reply: Thank you for your valuable suggestion. The time complexity for the global alignment is $\mathcal{O}(D)$, while for fine-grained alignment it is $\mathcal{O}(N_{t}*N_{s}*D)$, where $N_{t}$ and $N_{s}$ represent the number of text tokens and SMILES tokens, respectively, and $D$ is the dimensionality of the representations. We will incorporate this computational complexity analysis into the revised version of the paper.

To provide more clarity, the following table (also presented as Table 7 in the paper) compares global alignment (GA) and hierarchical alignment (HA) in Atomas-base. We report the average training time on the Chebi-20 dataset using an NVIDIA A100 GPU, along with performance metrics on the molecule generation task.

We acknowledge that fine-grained alignment introduces additional computational overhead. However, we believe the observed performance improvements shown in the table above justify this marginal increase in computational cost.

Additionally, we have made every effort to include statistics on the pre-training computational cost for baseline methods, as shown in the table below. However, we note that most methods do not explicitly report pre-training time costs. Variations in batch sizes, GPU types, and other factors across different setups make it challenging to estimate these costs with high accuracy.

Supplementary experiment for computational time comparisons:

Based on the data presented, we believe the pre-training stage of our framework remains computationally efficient and affordable. We will add these supplementary experiments to the revised version of the paper.

We thank reviewer 1nxB for the valuable and positive feedback and for acknowledging the novelty, clear present, and adequate ablation contributed by our method. We have addressed all the comments. Please kindly find the detailed responses below. Any further comments and discussions are welcomed!

W1: Missing the standard deviation or significance t-test in Retrieval Tasks.

Reply: Thanks for your valuable comment. In the retrieval task, as all models are evaluated without fine-tuning, the varying seeds do not influence the experimental results, leading to the standard deviation equal to 0.

We greatly appreciate your suggestion and acknowledge the importance of including such details to improve the rigor of our work. To address your concern, we have re-pretrained Atomas-base and Atomas-large three times and report the standard deviation results in below tables:

Supplementary experiment for Text to Molecule Task in Table 1:

Supplementary experiment for Molecule to Text Task in Table 1:

Regarding Table 8 and Figure 4, we would like to clarify that these are ablation studies specifically designed to highlight the influence of model setting choices rather than variability in task performance. Nonetheless, we have included the standard deviation results for Table 8 in the tables below:

Supplementary experiment for Text to Molecule Task in Table 8:

Supplementary experiment for Molecule to Text Task in Table 8:

We will also include these standard deviation results in our revised version of the paper.