ALMol: Aligned Language-Molecule Translation LLMs through Offline Preference Contrastive Optimisation (2024)

Abstract: The field of chemistry and Artificial Intelligence (AI) intersection is an area of active research that aims to accelerate scientific discovery. The integration of large language models (LLMs) with scientific modalities has shown significant promise in this endeavour. However, challenges persist in effectively addressing training efficacy and the out-of-distribution problem, particularly as existing approaches rely on larger models and datasets. In this context, we focus on machine language-molecule translation and deploy a novel training approach called contrastive preference optimisation, which avoids generating translations that are merely adequate but not perfect. To ensure generalisability and mitigate memorisation effects, we conduct experiments using only 10% of the data. Our results demonstrate that our models achieve up to a 32% improvement compared to counterpart models. Finally, we introduce a fine-grained, domain-agnostic evaluation method to assess hallucination in LLMs and promote responsible use.

heading: 1 Introduction
text: The world is facing unprecedented complexity in the form of global challenges such as climate change, healthcare, and pandemics. Innovative scientific solutions are urgently needed to address these challenges. Chemistry has been at the forefront of developing such solutions, pioneering new drugs (Ferguson and Gray, 2018), creating advanced materials (Kippelen and Brédas, 2009), or enhancing chemical processes (Zhong et al., 2023). However, these frontiers are vast and require the involvement of Artificial Intelligence (AI) technology to navigate them effectively.
Large language models (LLMs) have shown promising potential for accelerating scientific discovery across various domains, including chemistry, biology, and materials science (Zhang et al.,
2023; AI4Science and Quantum, 2023). Existing work has applied successful paradigms from natural language processing (NLP) and multimodal representation learning to the chemistry domain. One common approach involves converting the inherent three-dimensional structures of molecules into SMILES, which provide a mapping to symbolic character-level representations. Subsequently, researchers have explored learning language-molecule representations either in separate yet coordinated spaces (Edwards et al., 2022, 2021; Liu et al., 2023a), in a joint space (Liu et al., 2023b), or through hybrid approaches (Luo et al., 2023; Christofidellis et al., 2023). In light of the recent significant advancements in the field, none of the above approaches effectively tackle the inherent challenges in training such models. Instead, they rely on sparse or noisy synthetic data, often necessitating exponentially more data than is typically used in NLP tasks (Edwards et al., 2024).
However, training on larger models and datasets does not necessarily guarantee higher performance. A successful paradigm that augments the capabilities of LLMs across multiple NLP tasks is Reinforcement Learning with Human Feedback (RLHF) (Ouyang et al., 2022). Although initially challenged by issues of slowness and instability, recent research has addressed many of these challenges by shifting towards closed-form losses that operate directly on offline preference data (Rafailov et al., 2024). RLHF has demonstrated superior performance compared to standard minimising crossentropy optimisation approaches.
In this context, we address challenges related to effectively training robust language models when integrated with scientific modalities. We deploy a novel way of training LLMs for language-molecule translation that avoids generating translations that are only adequate but not perfect, called contrastive preference optimisation (CTO) (Xu et al., 2024). CTO is based on offline preferences instead of su-
22
pervised fine-tuning, mimicking reference translations. To ensure that our models can effectively generalise instead of memorising patterns, we conduct experiments using only 10% of the L+M-24 dataset (Edwards et al., 2024). Our contributions have as follows: • Our models achieve significant performance im-
provements across various evaluation metrics compared to models trained on extensive indistribution and out-of-distribution data (§ 4.4). • We showcase their robustness through experiments comparing pivot and minor cross-modals. Our empirical results demonstrate that our models consistently outperform the leading baseline, Meditron, which is trained on the entire dataset, even in agnostic cross-modal scenarios (§ 4.4). • We propose a fine-grained evaluation method that is domain-independent, assessing factual consistency in generated captions using a questionanswering evaluation metric and measuring overlaps of unigrams in generated molecules against references (§ 3.3). Our analysis shows that our models achieve improved factual consistency and character-level unigram overlaps for caption and molecule generation (§ 4.5).

heading: 2 Background
text: Reinforcement Learning with Human Feedback (RLHF) optimisation (Ouyang et al., 2022) operates with a triple dataset D = {x, yw, yl}, where yw and yl represent preferred and dis-preferred outputs, corresponding to input x, such that yw ≻ yl for x. The probability of yw over yl in pairwise comparisons is typically computed using the Bradley-Terry model (Bradley and Terry, 1952):
p∗(yw ≻ yl|x) = σ(r∗(x, yw)− r∗(x, yl)) (1)
where σ is the logistic function, and r∗ denotes the reward function that underlies the preferences.
As obtaining the reward directly from a human would be prohibitively expensive, a reward model rϕ is trained to act as a surrogate by minimising the negative log-likelihood of the preference data:
L(rϕ) = −E(x,yw,yl)∼D[log σ(rϕ(x, yw)−rϕ(x, yl))] (2) Additionally, the Kullback-Leibler (KL) divergence between the outputs generated by πref and the parameterised πθ models serves as an additional reward signal, ensuring that the generated responses closely align with the reference model. Conse-
quently, an optimal model πθ is one that maximises:
E(x∈D,y∈πθ)[rϕ(x, y)]−βDKL(πθ(y|x)||πref(y|x)) (3) where β is the temperature parameter typically ∈ [0.1, 0.5].
RLHF can present challenges due to its inherent slowness and instability, especially in distributed settings (Zheng et al., 2024). Recent work has shifted towards closed-form losses to align LLMs with human preferences. Here, we experiment with contrastive preference optimisation that adopts a closed-form loss for RLHF.

heading: 3 Methodology

heading: 3.1 Task Formulation
text: Let (x, y) be a pair of source and target sequences mapped to X and Y spaces, respectively. We cast the problem of language-molecule translation as a cross-modal translation task that operates on offline preference data D = {x(i), y(i)w , y(i)l }Ni=1, where x is an input, yw are preferred (e.g. human gold standard) and yl dis-preferred outputs (typically synthetic, obtained from an appropriate translation model), and N is the total number of pairs. The goal is to learn an optimal function f : X ↔ Y through a model πθ parameterised by θ. We coordinate the two spaces through instructional modelling to regulate the translation process in both directions. Specifically, for LMolT, we use instructions for language-to-molecule and molecule-to-language translation (see Appx. A).

heading: 3.2 Contrastive Preference Optimisation
text: Contrastive preference optimisation (CTO) (Xu et al., 2024) addresses challenges stemming from the inherent limitation in RLHF, as discussed in § 2, and from the necessity of high-quality data. CTO is a general approximation of Eq. 3 using a uniform reference model, which assumes equal likelihood for all possible generated outputs:
L(πθ;U) = −E(x,yw,yl)∼D[ log σ ( β log πθ(yw|x)− β log πθ(yl|x) )] (4)
where πθ is parameterised model by θ and β hyperparameter (please refer § 2). Eq. 4 implies that the loss is calculated based on how well the generated translations match this uniform distribution of possible translations, rather than being biased towards
any particular translation. To maintain πθ close to the preferred data distribution, a behaviour cloning (BC) (Hejna et al., 2023) regulariser is introduced:
min θ
L(πθ, U) s.t.
E(x,yw)∼D [ KL(πw(yw|x)||πθ(yw|x)) ] < ϵ, (5)
Here, ϵ denotes a small positive constant, and KL signifies the Kullback-Leibler divergence. The regulariser is enhanced with an additional SFT term on the preferred data, bolstering the CPO loss as:
LCPO = min θ L(πθ, U)︸ ︷︷ ︸ Lprefer −E(x,yw)∼D[ log πθ(yw|x)]︸ ︷︷ ︸ LNLL
(6)

heading: 3.3 Proposed Evaluation Methodology
text: Prior studies have utilised embedding representations, for assessing the semantics in chemicaldomain models (Jaeger et al., 2018; Edwards et al., 2021; Christofidellis et al., 2023). However, these approaches require domain adaptation for out-ofdistribution data (Edwards et al., 2024) and might lead to opaque and arbitrary outcomes (Steck et al., 2024). We address these limitations by introducing a scalable fine-grained evaluation methodology for assessing the presence of hallucinations1 in generated outputs.
Language Evaluation: For molecule-tolanguage translation, we deploy the QAFactEval (Fabbri et al., 2022) metric to evaluate the factual consistency of generated captions. QAFactEval first selects noun phrases and named entities (NER) from the generated outputs. A question generation (QG) model then formulates associated questions, which a question answering (QA) model addresses based on the reference text. QAFactEval measures the semantic overlap between the QA model’s responses and the selected answers to produce the final metric score. An example is illustrated in Fig. 1. Here, we report the semantic overlap, the f1 accuracy between the QA model and the selected answer, and answerability, which is the probability of the question being answered by the reference caption.
1Hallucination in LLMs refers to a phenomenon where the generated outputs are inaccurate, nonsensical, or contradictory to the provided factual information.
Molecule Evaluation: For language-to-molecule translation, we employ the Chr-F metric, an F-score statistic, to evaluate character n-gram matches between prediction-reference pairs (Popović, 2015). This metric assesses the matches in generated molecules against their references by averaging the scores of unigram, bigram, and trigram matches. A higher Chr-F score indicates better performance.
Bias Evaluation: We also calculate the character and token length bias in generated-reference pairs of molecules and captions, respectively, to investigate potential length bias in the evaluated LLMs.

heading: 4 Experiments

heading: 4.1 Data
text: We conduct experiments on the L+M-24 benchmark dataset, which encompasses both molecule and linguistic modalities (Edwards et al., 2024). It is divided into four categories, each with significant applications in small-molecule domain; biomedical; light and electricity; human interaction and organoleptics; and agriculture and industry. The training and validation subsets consist of approximately 127k and 34k language-molecule pairs, respectively. Here, we utilise 10% of these subsets for training and validation. To operationalise CTO, we recreate a triples dataset consisting of preferred and dis-preferred outputs (see § 2), where the former are the golden references and the latter are generated from MolT5 (Edwards et al., 2022). For evaluation, we randomly selected 3k unseen pairs
from a distinct dataset provided by the research group of L+M-24.2

heading: 4.2 Bechmark Models
text: We compare our results with established languagemolecule models as captured in the literature: • TxtChem-T5 (Christofidellis et al., 2023): A T5
model trained on both linguistic and molecule modalities with a multi-task objective across various datasets, including the CheBI-20 dataset (Edwards et al., 2022), akin to L+M-24. • Chem-LLM (Zhang et al., 2024): An InternLM2Base-7B model, trained on an extensive chemical domain knowledge dataset, with the direct preference optimisation objective (Rafailov et al., 2024), achieves results comparable to GPT-4. • Meditron (Chen et al., 2023): A Meditron-7B model fine-tuned on the entire L+M-24 for unidirectional language-molecule translation. • SFT-Meditron: We fine-tune Meditron-7B on a 10% subset of L+M-24 for bi-directional machine language-molecule translation.

heading: 4.3 Experimental Settings
text: Here, we train Meditron with CTO on a 10% subset of L+M-24. We experiment with both language and molecule weight initialisation obtained from Meditron trained on the entire data (Edwards et al., 2024). We refer to them as CTO-Meditron−−→ Lan. and CTO-Meditron−−→
Mol. , respectively. We train the
models with QLoRA (Dettmers et al., 2024). For evaluation, we adopt established metrics in (Edwards et al., 2022).

heading: 4.4 Experiment Results
text: Table 2 presents a summary of the moleculeto-language results. We observed a significant decrease in performance for benchmark models trained on extensive data with SFT when tested on out-of-distribution data. Among the baseline models, Meditron demonstrated the highest performance, likely due to its training on the entire L+M-24 dataset utilised in our experiments. Training Meditron with SFT for bi-directional languagemolecule translation has demonstrated neither effectiveness (see Table 1) nor efficiency (refer to Appx. B). This suggests that the performance in our experiments is not dependent on memorised patterns from Meditron trained on the entire dataset. In contrast, our models trained with the CTO objective on only 10% of L+M-24 achieved a remarkable
2Sampling is conducted from a distinct subset.
improvement in performance across diverse evaluation metrics, up to 32% compared to Meditron trained on the entire dataset. This improvement is consistent, as our model consistently enhances performance when initialised from agnostic crossmodals, i.e., CTO-Meditron−−→
Lan in Table 1.
We observed similar performance patterns for language-to-molecule translation as reported in Table 2. However, even though our model achieved better performance compared to Meditron when initialised from agnostic cross-modals, it struggled to learn molecular patterns (see CTO-Meditron−−→ Mol. in Table 2). This suggests that language plays a pivotal role in the molecule modality. In the future, we aim to explore more advanced initialised methods to address this challenge.

heading: 4.5 Evaluation Results
text: Fig. 2 illustrates the evaluation results on the factual consistency of generated captions against references for the molecule-to-language task. CTOMeditron−−→
Mol. , trained on 10% of the available data,
exhibited superior factual consistency, achieving a semantic overlap of 2.08, f1 accuracy of 0.34, and answerability of 0.68, compared to 1.34, 0.20, and 0.51, respectively, for Meditron trained on the entire dataset. CTO-Meditron−−→
Lan. also outperformed
Meditron but showed lower performance than CTOMeditron−−→
Mol. . We attribute this to the model being
initialised by agnostic cross-modals.
For the language-to-molecule task, we observed that both Meditron−−→
Lan. and Meditron−−→ Mol. achieved
similar performance in terms of uni-, bi-, and tri-gram overlaps between generated and reference pairs, outperforming Meditron (see Fig. 3). However, when the model was initialized with known cross-modal weights, i.e., Meditron−−→
Lan. , it
achieved a slightly increased performance
For the language-to-molecule task, we observed that Meditron and Meditron−−→
Mol. generated signifi-
cantly shorter and longer outputs, respectively (see Fig. 4). In contrast, Meditron−−→
Lan. did not exhibit
any length bias, producing outputs similar in length to the actual ones. Conversely, for the moleculeto-language task, our models did not show any significant length bias, while Meditron, trained on the entire dataset, generated significantly shorter answers against references.

heading: 5 Conclusion
text: This work address training efficacy and the outof-distribution problem for automatic languagemolecule translation. We train models using only 10% of available data and deploying contrastive preference optimisation which avoids generating translations that are merely adequate but not perfect. We achieve significant improvement in performance when compared with models trained on extensive in and out-of-the-distribution data. Finally, we propose a fine-grained, domain-agnostic evaluation method to assess hallucination in LLMs. Our models show superior factual consistency for caption generation and character-level unigram overlaps for molecule generation.

heading: B Training Effectiveness and Efficiency

heading: A Language-molecule Translation Instructions

heading: Future Work
text: In the future, we aim to explore more advanced initialised methods to address this challenge