Enhancing Cross Text-Molecule Learning by Self-Augmentation (2024)

Abstract: The development of Large Language Models (LLMs) has greatly advanced the field of drug discovery, with the belief that natural language can enhance human control over molecule design. However, the scarcity of high-quality labeled data remains a challenge for cross textmolecule learning. Existing datasets are limited due to the difficulty of collecting precise molecule-description pairs. Although recent efforts have utilized pseudo data generated by LLMs for augmentation, the lack of specialized chemistry knowledge of LLMs and the absence of an effective high-quality data selector may introduce noise into the annotations, compromising the models’ robustness. To address these challenges, this paper introduces a novel framework that interweaves model finetuning and data augmentation to overcome the scarcity of high-quality data. The proposed approach involves an iterative procedure where the model plays dual roles in annotating unlabeled data and sampling a subset of high-quality data until convergence is achieved, enhancing the model’s understanding and adaptability. Additionally, a new dataset called SAPubChem-41 is presented, which comprises meticulously curated high-quality parallel molecule-description pairs designed specifically for fine-tuning purposes. This research provides an important contribution to the field by addressing the need for high-quality datasets and presenting an effective framework for cross text-molecule learning.

heading: 1 Introduction
text: The emergence of Large Language Models (LLMs) has significantly propelled the development of drug discovery (Zhang et al., 2024). Recent progress in language models shed light on drug discovery, with the vision that humans can possess a higher-level control over molecule design facilitated by natural
*Corresponding author.
language. Edwards et al. (2022a) introduce two new tasks: molecule captioning (Mol2Cap) and text-based molecule generation (Cap2Mol), and present MolT5 (Edwards et al., 2022b) based on the T5 (Raffel et al., 2020) architecture to translate between molecule and text. Subsequently, diverse endeavors have been undertaken to tackle these challenges. Text+Chem T5 (Christofidellis et al., 2023) and BioT5 (Pei et al., 2024) are also T5-like models, which incorporate multi-task and multidomain pretraining process and enable the bidirectional generation between different modalities in a single model. Some other methods employ the Generative Pretrained Transformer (GPT) architecture. MolXPT (Liu et al., 2023) proposes to wrap molecules in the sentences to make the pertaining corpus. More recently, MolCA(Liu et al., 2024) considers the 2D graph information for molecules and designs a projector to connect molecular content to language models. These studies mark the research efforts in the realm of cross text-molecule learning.
Despite rapid development in cross textmolecule modalities, there is still demand for highquality labeled data in this field. Such data of superior quality is indispensable to capacitate models for downstream tasks. However, the available datasets relevant to the tasks are limited, a consequence of the arduous nature of collecting precise and enlightening parallel molecule-description pairs. This scarcity of annotated data hinders the model’s ability to learn between the modalities of the molecule and the text. As shown in Figure 1, recent efforts by Chen et al. (2023) address the scarcity challenge by utilizing pseudo data generated by LLMs as augmentation. The problem is that the absence of specialized knowledge in chemistry may introduce noise into the data annotated by general Large Language Models, thereby compromising the robustness of models fine-tuned on such datasets. In spite of these commendable advance-
9551
ments, the accessibility of high-quality datasets for fine-tuning remains notably deficient.
To address these challenges, this paper introduces a novel framework that interweaves model fine-tuning and data augmentation (See Fig. 1, aiming to rectify the high-quality data scarcity observed in the realm of cross text-molecule learning. Our proposed approach integrates the model into the process by assigning it dual roles: annotating unlabelled data and sampling a subset of high-quality data. This iterative procedure continues until convergence is achieved, thereby refining the model’s understanding and adaptability. Additionally, we present a new, larger dataset SAPubChem-41. This meticulously curated dataset comprises a wealth of high-quality parallel molecule-description pairs meticulously designed for the specific purpose of fine-tuning. The model optimized using SAPubChem-41 has been evaluated across a total of 13 molecule-text benchmarks, exhibiting superior performance in comparison to the model optimized using the original dataset, ChEBI-20. Our key contributions can be summarized as:
• We present an innovative framework designed to augment cross-text-molecule learning through a self-augmentation strategy. This system leverages the model itself to annotate unlabelled data, choosing only samples that meet a high-quality standard.
• We introduce SAPubChem-41, a novel dataset consisting of both real data and high-quality pseudo-labeled data, augmented by the model itself. This makes a notable contribution to the available datasets in the field of crossmolecule-text modality.
• We show that within our iterative setting, selfaugmented data significantly improves the performance of the model in downstream tasks with each epoch. Furthermore, these results tend to converge as the number of epochs increase. This process thereby successfully validates the efficacy of self-augmented data in enhancing model performance.

heading: 2 Background

heading: 2.1 Cross Text-Molecule Learning
text: Edwards et al. (2021) introduce a new task Text2Mol, which uses descriptions as search queries to retrieve the target molecules. Edwards et al. (2022a) first addresses the problem of crossdomain generation by linking natural language and chemistry, tackling tasks such as text-conditional de novo molecule generation and molecule captioning. Li et al. (2023a) combines retrieval-based prompt paradigm with LLMs like ChatGPT to achieve translation between molecule language and natural language. Christofidellis et al. (2023) propose the first multi-domain, multi-task language model that can solve a wide range of tasks in both the chemical and natural language domains. Liu et al. (2023) propose MolXPT, a GPT-based model pre-trained on molecule SMILES, biomedical text, and wrapped text. BioT5 (Pei et al., 2024) further exploits SELFIES for 100% robust molecular representations and discriminate structured knowledge from unstructured knowledge. To capitalize on these insights, we adopt an iterative framework to interweave model fitting and data augmentation, to mitigate data scarcity and further enhance the alignment between molecule and text representations.

heading: 2.2 Data Augmentation for Cross Text-Molecule Learning
text: Data augmentation is a regularization strategy utilized to enhance model performance by diversifying the available data through various techniques (Hernández-García and König, 2018). In the field of natural language processing (NLP), data augmentation has been widely employed to address data scarcity issues, leading to the proposal of various augmentation methods (Feng et al., 2021; Li et al., 2022). Textual data augmentation can be accomplished through simple rules, such as synonym replacement and word order modification (Zhang et al., 2015). Furthermore, the development of deep learning models has facilitated the generation of new text to augment the existing data (Szegedy et al., 2016; Wu et al., 2019; Anaby-Tavor et al., 2020; Yoo et al., 2021; Zhou et al., 2021; Li et al., 2023b; Dai et al., 2023). Acquiring high-quality annotated molecule-text datasets is prohibitively expensive, for many descriptions of corresponding molecules are scattered and inaccurate, leading to limited fine-tuning datasets for cross text-molecule learning.More data augmentation techniques are desired to further enhance cross text-molecule learning.Chen et al. (2023) pioneers to exploit LLMs’ annotated pseudo label for domain adaption and data augmentation.However, while pseudo label as domain adaptation shows impressive performance compared to existing methods, pseudo label as data augmentation still struggles against label noise interfering model’s performance. To address the problem and further improve the quality of augmented dataset, our work introduces a pioneering iterative framework that interweaves model finetuning and data augmentation.

heading: 3 Method
text: In this section, we describe a computational framework for jointly optimizing model performance and dataset quality.

heading: 3.1 Overall Framework
text: The literature demonstrates a substantial body of research focused on enhancing both model performance and dataset quality concurrently, as depicted in Figure 2. If we assume the existence of a subjective testing environment capable of gathering reliable molecule-caption pairs, then the core challenge becomes how to effectively sample a subset A from a large-scale pseudo-labeled dataset D
that exhibits high confidence. In line with Li et al. (2023b), our approach adopts an iterative strategy, combining model fine-tuning and data augmentation operations defined below. 1. Model fine-tuning: Given a data distribution, find an approximate risk minimizer. 2. Data augmentation: Given a model, augment data and sample a new data distribution. 3. Data combination: Combine a set of distributions into a single distribution.
At the beginning of each iteration, a new model is finetuned on the latest dataset. This newly trained model is subsequently employed to annotate a set of unlabeled data, thereby generating a pseudolabeled dataset. The sampler is then entrusted with the responsibility of discerning samples distinguished by their high quality from this pseudolabeled dataset. The discerned insights are then harnessed to augment and refine the current dataset, thereby facilitating a continuous and iterative improvement process.

heading: 3.2 Self-Annotating
text: We gathered 200,000 unannotated SMILES of molecules from PubChem (Kim et al., 2023) and conducted deduplication to mitigate the risk of data leakage in the validation dataset. For each unlabeled SMILES, we executed inference on the model itself to produce a candidate caption, thereby obtaining the pseudo-labeled dataset.
ChatGPT Annotating vs Specialized Model Annotating Despite the remarkable performance of Large Language Models (LLMs) such as ChatGPT in various cross-modal tasks, they have not succeeded in surpassing specialized models in the molecule captioning task (Li et al., 2023a; Christofidellis et al., 2023). This suggests that specialized models serve as better annotators for unlabeled SMILES strings. Furthermore, as depicted in Figure 3(a), by leveraging the Text2Mol score (Edwards et al., 2022a) as a fundamental metric of data quality, which measures the similarity between SMILES strings and their captions, we evaluated the distribution of the real training set, MolReGPT annotated data, and self-annotated data. It is evident that the self-annotated data yields higherquality data, closely resembling real data more than the MolReGPT annotated data.
Iterative Annotating We employ an iterative training approach to enhance the proficiency of annotators and produce captions of superior quality.
As depicted in Figure 3, during each iteration, the quality of the molecule-caption pairs is enhanced through the implementation of the iterative strategy.

heading: 3.3 Self-Sampling
text: A discerning subset sampling strategy is anticipated to effectively filter pseudo labels, thereby ensuring the attainment of high-quality labels for training the model. This refined approach aims to enhance the overall reliability of the model by mitigating the impact of unreliable or inaccurate labels in the training process. In contrast to prior work (Li et al., 2023b), wherein fine-tuning an external score model was employed to select highquality pairs, we have found this approach to be
both time-consuming and unreliable when applied to the domain of cross text-molecule learning. Notably, the susceptibility of the model to overfitting incorrect pseudo-labels is a concern, as detailed in section 4.4. To tackle this dilemma, we introduce a dual selection mechanism that not only proves to be more computationally efficient but also enhances reliability in ensuring the high quality of our augmented dataset. This mechanism involves two key steps. Firstly, leveraging the generation likelihood as a confidence score to filter out unreliable molecule-caption pairs. Secondly, updating and ranking the remaining molecule-caption pairs based on the Text2Mol score (Edwards et al., 2022a) This refined approach streamlines the sampling process, offering increased efficiency and
heightened confidence in the identification of highquality pairs for dataset augmentation.
Data Filtering First, we employ the model itself to filter out a set of low-quality pairs. Specifically, we utilize the generation likelihood as a measure of the quality of the generated sequence. However, previous studies (Liu and Liu, 2021; Liu et al., 2022) have underscored that the correlation between sequence probability and its quality for Maximum Likelihood Estimation (MLE) trained models can be low, due to the presence of deterministic (one-point) target distribution issues. Furthermore, sequence likelihood estimation becomes noisier when the decoded sequences of models deviate from the exposed training data distribution, thereby exacerbating the problem of exposure bias
(Ranzato et al., 2016). To address these intricate challenges, we introduce the Sequence Likelihood Calibration (SLiC) stage. This stage, depicted in Figure 1, serves to mitigate the aforementioned issues and diminish the gap between sequence likelihood and its associated quality. Following the methodology of Zhao et al. (2022), we conduct further fine-tuning of the model, introducing a new objective that encompasses the following loss functions:
Lcalrank = max ( 0, β − logPθ ( ŷ+ | x ) + logPθ ( ŷ− | x )) ,
Lregce = ∑
t
− logPθ ( ȳt | yt−1,x ) ,
Lregkl = ∑ t Pθ ( ȳt | yt−1,x ) log
Pθ(ȳt|yt−1,x) Pθft(ȳt|yt−1,x) .
Given the context x, target y,and positive and negative candidates pairs ŷ+, ŷ−, Pθ(y | x) denotes the generation sequence likelihood. Rank loss Lcalrank optimizes the ranking order of positive and negative candidate pairs. Cross entropy loss Lregce is the standard fine-tuning MLE objective . KL divergence loss Lregkl directly minimizes the probability distribution distance between the calibrated model and the fine-tuned model at each token on the observed target sequence.
Following the SLiC stage, we regard the generation likelihood as one of our metrics for assessing the quality of pseudo labels. In our experiments, we discard pairs with log-likelihood equal to negative infinity and retain the remaining pairs as highquality pairs.
Update and Rank Secondly, for those highquality pairs, we further refine the data based on the similarity scoreFtext2mol proposed by Edwards et al. (2022a). At iteration t, we update the caption if Ftext2mol (m, ct) > Ftext2mol (m, ct−1) ,then we rank the updated dataset according to Ftext2mol ,selecting the top t× 2k pairs to form the new Augmented Dataset At.

heading: 4 Experiments
text: In order to validate the efficacy of utilizing our framework, we have undertaken a series of comprehensive experiments.

heading: 4.1 Experimental Setup
text: Dataset We employ the ChEBI-20 (Edwards et al., 2021) to serve as our original dataset. Additionally, we gather 200,000 unannotated molecules
from the PubChem database (Kim et al., 2023). Subsequently, we conduct a rigorous filtering procedure to exclude any molecules present in downstream datasets, ensuring the elimination of potential overlaps between the newly acquired molecules and those already present in the ChEBI-20 dataset. For our experimental dataset augmentation, we set the iteration step to 2000. At each iteration stage, the dataset size expands by 2000 from the preceding iteration. A detailed comparison between our final augmented dataset, SAPubChem-41, and the original dataset, ChEBI-20, is presented in Table 1.
Base Model To prove the effectiveness of our framework, we use the Text+Chem T5 (Christofidellis et al., 2023) as the base model for simulating the iterative process.
Metrics Following the previous studies (Edwards et al., 2022a; Li et al., 2023a; Christofidellis et al.,
2023; Chen et al., 2023), we evaluate the results with following metrics:
• Molecular Captioning: BLEU-2 and BLEU4(Papineni et al., 2002) are metrics used to evaluate the quality of machine-generated text by comparing it to a reference text, with a higher score on the BLEU metric indicating a higher level of similarity between the generated text and reference text. ROUGE1, ROUGE-2, and ROUGE-L(Lin, 2004) is similar to BLEU, while computing the recalloverlap of unigrams, bigrams, and longest common subsequences between the generated and reference texts. METEOR (Banerjee and Lavie, 2005) is a metric that uses a combination of unigram precision, recall, and a synonym-matching component to evaluate the generated text against the reference text, which is designed to be more sensitive to fluency, meaning, and structure than BLEU.
• Text-based Molecule Generation: BLEU and the Exact Match scores are calculated as basic assessments. Moleculespecified metrics including Levenshtein distance(Levenshtein et al., 1966), validity(Edwards et al., 2022a), and three molecule fingerprints scores - MACCS FTS(Durant et al., 2002), RDK FTS(Schneider et al., 2015), and Morgan FTS(Rogers and Hahn, 2010) are calculated to provide valuable insights into the quality, validity, and structural characteristics of the generated molecules.
Implementation We conducted our experiments on Ubuntu 22.04 using RTX 4090(24GB) * 4 with CUDA 11.7.1. Our primary dependencies are Python 3.11.4, PyTorch 2.0.1, Transformers 4.31.0, and Numpy 1.24.3. We initialize the three pretrained models using public checkpoints. For finetuning, we adopted the configuration in MolT5 (Edwards et al., 2022b), with a learning rate of 1e-3, 50,000 fine-tuning steps, weight decay of 0.1, batch size of 32, random seed of 42, and 1000 warm-up steps. For evaluation, we use a greedy search with a maximum generation length of 512 during generation. The evaluation metrics script is derived from MolT5 (Edwards et al., 2022b). For other hyperparameters, we relied on the default settings of the T5ForConditionalGeneration class in Huggingface.

heading: 4.2 Performance Comparison of Molecule-Caption Translation
text: Molecular Captioning Table 2 presents the results of the molecule captioning task, utilizing the original ChEBI-20 dataset and the augmented SAPubChem-41 dataset to optimize three distinct pre-trained models. Across all measurement metrics, it is observable that the performance of all models using our augmented SAPubChem-41 dataset is generally superior to that utilizing the original ChEBI-20 dataset.
Text-based Molecule Generation Table 3 presents the results of the text-based molecule generation task, utilizing the original ChEBI-20 dataset and the augmented SAPubChem-41 dataset to optimize three distinct pre-trained models.Notably, the performance of all models that utilize our augmented SAPubChem-41 dataset generally surpasses those that use the original ChEBI-20 dataset. The accuracy of the generated molecule, as well as its similarity to the ground truth molecule, have been significantly improved following the selfaugmentation of the dataset. This evidences substantial enhancements in performance when integrated with our self-rewarding framework.

heading: 4.3 Iterative Procedure
text: We utilize our framework on the Text+Chem T5, taking the ChEBI-20 as the original dataset. This approach facilitates an augmented dataset while simultaneously enhancing the model’s performance. It is critical to note that our intention in these experiments is not necessarily to achieve a state-of-the-art model. Rather, our prime objective is to examine the efficacy of our framework, which is designed
to optimize both the model and the pseudo-labeled dataset concurrently. Figure 4 illustrates the iterative process in which dashed lines represent the molecule captioning task and solid lines signify the text-based molecule generation task. We report the complete results in Table 4 and 5. There are a few observations: First, it is noticeable that for both tasks the outcome of iteration 5 mostly falls short of that of iteration 4 in all metrics. This suggests that the training process reaches convergence on iteration 4 for both tasks. Consequently, we adopt the augmented dataset obtained at the fourth iteration as our conclusive dataset. Second, for the text-based molecule generation task, there is a consistent upward trend across all pertinent metrics. This suggests that an increased quantity of highquality pseudo-labeled data contributes to the optimization of the model in the molecule generation process. This trend underscores the necessity of sufficient high-quality data in cross text-molecule learning. Lastly, when it comes to the molecule captioning task, the general trend in relevant metrics indicates a rise. Initially, there is a minor decrease in the BLEU and Meteor metrics, which is counterweighed by an enhancement in Rouge. As the volume of high-quality pseudo-labeled data increases, all metrics demonstrate improvement. This trend suggests that the quantity of high-quality pseudolabeled data is also a significant factor.

heading: 4.4 Ablation Study
text: We perform further ablation studies to validate the effectiveness of the key component in our framework, i.e., the self-sampling strategy.
To demonstrate the significance of pseudo-
labeled data selection, we conducted a comparative analysis of the performance of all intermediate models. We compared the use of randomlyselected data of the same size with our refined data. As depicted in Figure 5 and 6, our refined data consistently outperformed the randomly-selected data across all metrics at each iteration for both tasks.
Notably, when optimizing the model on randomlyselected data, the results were generally inferior to the original model on most metrics. This indicates that the model tends to overfit to noisy pseudolabels. These findings emphasize the critical role of the self-sampling strategy in our framework.

heading: 5 Conclusion
text: In this paper, we have explored and grappled with the challenge of a deficiency in the accessibility of high-quality datasets for model fine-tuning in the field of cross text-molecule learning, by developing a novel framework. This innovative framework efficiently combines model fine-tuning with data augmentation, integrating the model in such a way that it expands its roles to include annotating unlabeled data and sampling a subset of high-quality data. The paper further introduces a high-quality dataset, the SAPubChem-41, expressly designed for fine-tuning purposes. The experiments’ results were impressive as the model optimized using SAPubChem-41 consistently outperformed the model optimized using the original dataset, ChEBI-20. Going forward, this paper’s findings present crucial steps toward addressing high-quality data scarcity for fine-tuning and ensuring further advancement in the exciting field of cross text-molecule learning.

heading: 6 Limitations and Future Works
text: One limitation of our framework is that it only considers a 1D representation of molecules. The integration of additional representations, such as molecule graphs, is left as future work. Additionally, our data augmentation module currently only supports sample-level labels. It is expected that more quality measurements from different levels will be incorporated to provide a comprehensive perspective on quantifying the dataset quality.

heading: Acknowledgements
text: This work is supported by National Natural Science Foundation of China (62302433, U23A20496), Zhejiang Provincial “Jianbing” “Lingyan” Research and Development Program of China (2024C01135), Zhejiang Provincial Natural Science Foundation of China (LQ24F020007), New Generation AI Development Plan for 2030 of China (2023ZD0120802) and CCF-Tencent Rhino-Bird Fund (RAGR20230122).

heading: A Datasets Information
text: SAPubChem-41 The SAPubChem-41 datasets comprises 33k real data sourced from CHEBI-20 and 8k high-quality pseudo labelled data sourced from self-augmentation on Text+ChemT5. Representative examples of pseudo-labelled part are provided in Table 7. In addition, to further elucidate the advantages of iterative strategy on model performance, we have documented representative examples of updated captions in Table 6.
Comparison with Existing Datasets In Table 8 we present a detailed comparison of our SAPubChem-41 datasets with existing limited datasets with parallel molecule-description pairs,including:
• ChEBI-20 (Edwards et al., 2021): consists of 33k molecule-caption pairs with caption exploiting ChEBI(Degtyarenko et al., 2007)annotations and more than 20 words.
• PCdes (Zeng et al., 2022): consists of 15k substances in PubChem which have names, SMILES and corresponding paragraphs of property descriptions.
• DrugBank-23(Chen et al., 2023): consists of 23k compounds from DrugBank (Wishart et al., 2018) with corresponding description’s length longer than 20 words.

heading: B Pre-trained Model Information
text: Table 9 presents the information of three different pre-trained models utilized in our experiments.

heading: C Experiments
text: C.1 Performance Comparison Here we present the complete results of performance comparison for both molecular captioning task and molecule generation task in Table 10&Table 11.
C.2 Ablation Study Here we present the complete results of both molecular captioning task and molecule generation task using different sampling strategy in Table 12 & Table 13.
C.3 Additional Experiments To further substantiate the efficacy of our framework and the exceptional quality of the SAPubChem-41 dataset, we have conducted a series of experiments. These were performed with the goal of optimizing three distinct pre-trained models, utilizing both our augmented dataset, SAPubChem41, and other prevalent public datasets which include ChEBI-20 and PCdes. The findings from these experiments are reported in Table 14 & Table 15. It can be noted that across both task categories, models that were optimized using our self-augmented dataset, SAPubChem-41, generally exhibited superior performance in comparison to those optimized with other existing datasets.

heading: Future Work
text: The integration of additional representations, such as molecule graphs, is left as future work. Additionally, our data augmentation module currently only supports sample-level labels. It is expected that more quality measurements from different levels will be incorporated to provide a comprehensive perspective on quantifying the dataset quality