MolReFlect: Towards In-Context Fine-grained Alignments between Molecules and Texts

Previously, the molecule-caption translation task treats a molecule as a general SMILES string $m$ and tries to let LLMs learn the direct mappings $m\leftrightarrow c$ between the molecule SMILES string $m$ and the textual caption $c$. To refine the alignments, several multi-modal methods like MolCA (Liu et al., 2023) have been proposed to incorporate molecule graph information $g_{m}$ and learn the direct mapping $(m,g_{m})\rightarrow c$ for the Mol2Cap task. Nevertheless, these methods still treat the molecule as a general SMILES sequence or molecular graph while ignoring the significance of detailed molecular sub-structures.

Instead of directly learning the mappings from molecules to captions, MolReFlect aims to extract fine-grained alignments $K$ between the molecule SMILES strings and molecule captions, thereby learning the mapping chains $m\rightarrow K\rightarrow c$ and $c\rightarrow K\rightarrow m$. Typically, the fine-grained alignments should be labeled by professional chemists, which is not only challenging but also financially prohibitive. As a result, LLMs have emerged as a viable alternative due to their advanced reasoning capabilities and a certain degree of chemical knowledge. Within MolReFlect, we have developed a zero-shot prompting strategy to empower the teacher LLM to engage in chain-of-thought (CoT) reasoning (Wei et al., 2022). This allows the teacher LLM to distill critical fragments from the molecule SMILES representations or captions, offering implications to their corresponding properties or sub-structure patterns. Formally, we have:

$\displaystyle K^{0}_{c}=p_{\theta_{T}}(.|c,I),\ K^{0}_{m}=p_{\theta_{T}}(.|m,I),$

(4)

where $\theta_{T}$ represents the parameters of the larger teacher LLM, $I$ is the CoT instruction, and $K^{0}_{c}$ and $K^{0}_{m}$ signify the alignments extracted in a zero-shot manner from the molecule caption and SMILES string, respectively.

Despite the powerful capabilities of LLMs, they can still generate answers with hallucinations (Yao et al., 2023). Their knowledge of chemistry is also limited due to the absence of domain pre-training on chemical corpora, which can introduce noises into the zero-shot alignments. To mitigate these potential noises and enhance the quality of zero-shot alignments, we propose a strategy that allows the larger teacher LLM to self-reflect on the zero-shot extraction results through in-context few-shot learning, where the previous zero-shot alignments are retrieved by similarity and serve as context examples for reflection. From the perspective of the molecular similarity principle, we do not calculate the similarity among the fine-grained alignments but follow the retrieval strategy adopted in Li et al. (2024a). For caption-based retrieval, we calculate the caption similarities based on the BM25 algorithm (Robertson et al., 2009) and retrieve top $n$ similar captions $\{c_{1},c_{2},...,c_{n}\}$ ranked by the BM25 scores and their corresponding zero-shot alignments $\{K^{0}_{c_{1}},K^{0}_{c_{2}},...,K^{0}_{c_{n}}\}$ for the input caption $c$ to form the context examples $C_{c}$:

$\displaystyle C_{c}=\{(c_{1},K^{0}_{c_{1}}),(c_{2},K^{0}_{c_{2}}),...,(c_{n},K^{0}_{c_{n}})\}$

(5)

Similarly, for molecule retrieval, we employ a pre-trained Mole-BERT (Xia et al., 2022) as the graph encoder and calculate the cosine similarities between the molecule graph embeddings. Top $n$ similar molecules $\{m_{1},m_{2},...,m_{n}\}$ and their corresponding zero-shot alignments $\{K^{0}_{m_{1}},K^{0}_{m_{2}},...,K^{0}_{m_{n}}\}$ are retrieved for the input molecule $m$ as the context examples $C_{m}$:

$\displaystyle C_{m}=\{(m_{1},K^{0}_{m_{1}}),(m_{2},K^{0}_{m_{2}}),...,(m_{n},K^{0}_{m_{n}})\}$

(6)

Based on the input $c$ or $m$, context examples $C_{c}$ or $C_{m}$, and instruction $I$, we could obtain the in-context reflected alignments $K^{1}_{c}$ or $K^{1}_{m}$ through the teacher LLM. Notably, the zero-shot alignments of the current input are not wrapped into the context to prevent the LLM from directly repeating it, and maintain consistent prompt formats across all instances:

$\displaystyle K^{1}_{c}=p_{\theta_{T}}(.|c,C_{c},I),\ K^{1}_{m}=p_{\theta_{T}}(.|m,C_{m},I),$

(7)

However, the context examples might also introduce noises that could misguide the reflection process, potentially leading to a decline in the quality of the reflected alignments $K^{1}$ compared to the zero-shot alignments $K^{0}$. Furthermore, the alignments generated by the teacher LLM can sometimes be too complex for the smaller student LLM to comprehend. Therefore, choosing the superior one between $K^{0}$ and $K^{1}$ is essential. To avoid possible information leaks, an unsupervised metric is required for selection. Specifically, we adopt the perplexity $\rm{ppl}$ as the metric from the information theory perspective:

$\displaystyle{\rm ppl}(K_{x},x)=\log\left[-p_{\theta_{S}}(K_{x}|x)\right],$

(8)

where $x$ is the input, $K_{x}$ denotes the corresponding alignments, and $\theta_{S}$ is the original parameters of the smaller student LLM. Higher perplexity scores suggest the presence of information that conflicts with the existing knowledge of LLMs. Therefore, the student LLM used for perplexity calculation is better to have some chemical knowledge like Galactiva-125M (Taylor et al., 2022) and can be different from the student LLM used for CoT-ICMT. Between the zero-shot alignment and the in-context reflected alignment, the one with lower perplexity will be selected:

	$\displaystyle K_{c}$	$\displaystyle=\left\{\begin{array}[]{ll}K^{0}_{c}&\ \text{if ${\rm ppl}(K^{0}_{c},c)<{\rm ppl}(K^{1}_{c},c)$},\\ K^{1}_{c}&\ \text{elsewise},\end{array}\right.$		(11)
	$\displaystyle K_{m}$	$\displaystyle=\left\{\begin{array}[]{ll}K^{0}_{m}&\text{if ${\rm ppl}(K^{0}_{m},m)<{\rm ppl}(K^{1}_{m},m)$},\\ K^{1}_{m}&\text{elsewise},\end{array}\right.$		(14)

While it is technically possible to leverage fine-grained alignments as contexts to allow the larger teacher LLM to generate final predictions directly in a CoT manner, the teacher LLM still lacks specialized pre-training on chemical corpora and is unfamiliar with the specific output distribution of the dataset. Consequently, directly querying the larger teacher LLM for final generations usually leads to unsatisfactory results. Furthermore, the cost of directly fine-tuning the larger teacher LLM is prohibitively high, making it unaffordable for most institutions. Instead, we fine-tune the smaller student LLM to learn from the fine-grained alignments provided by the larger teacher LLM. Notably, in this phase, we prioritize the reasoning capabilities of the student LLM over their knowledge of chemistry, so it can differ from the student LLM used to calculate perplexity.

In contrast to In-Context Molecule Tuning (Li et al., 2024a), CoT-ICMT organizes the fine-grained alignments of both the input $x$ and the context examples $C_{x}$ into the CoT format. This CoT format empowers LLMs to learn from the fine-grained alignments and the reasoning processes behind the context examples, thereby facilitating more explainable training. During the process of CoT-ICMT, top-$n$ similar examples are retrieved via the same retrieval strategies mentioned in Section 3.2 and then organized into the context with the CoT format to fine-tune the parameters of the smaller student LLM. Formally, similar to Eq. 3, the loss function can be represented as follows:

	$\displaystyle L^{c\!o\!t\!-\!i\!c\!m\!t}$	$\displaystyle(\theta)=$
		$\displaystyle\underset{(x,y)\in D}{\sum}\left[-\log p_{\theta}\!(\!y|x,\!K_{x}\!,\![C_{x\!\rightarrow\!K_{x}\!\rightarrow\!y}\!,\!F_{x\!\rightarrow\!K_{x}\rightarrow\!y}],\!I\!)\right],$		(15)

where $K_{x}$ denotes the fine-grained alignments of input $x$, $C_{x\rightarrow K_{x}\rightarrow y}$ $=\{\mathcal{P}(x_{i},K_{x_{i}},y_{i})\}_{i=1}^{n}$ represents the text content of context examples organized by the CoT format prompt $\mathcal{P}$, and $F_{x\rightarrow K_{x}\rightarrow y}=\{f_{i}:=x_{i}\rightarrow K_{x_{i}}\rightarrow y_{i}\}_{i=1}^{n}$ denotes the mapping chains behind the context examples, which map the original inputs to the fine-grained alignments and then further map the fine-grained alignments to the final targets.

We compare our method with the baseline models across the two sub-tasks of the ChEBI-20 dataset. Specifically, we select MolT5-large (Edwards et al., 2022), MolReGPT (Li et al., 2023a), MolCA (for the Mol2Cap task only) (Liu et al., 2023), BioT5 (Pei et al., 2023), and ICMA (Li et al., 2024a) as the baseline models. Notably, we adopt Mistral-7B as the smaller student LLM in the CoT-ICMT stage of MolReFlect. The overall results are presented in Table 1 for the Mol2Cap task and in Table 2 for the Cap2Mol task. We will proceed to discuss the outcomes for each sub-task individually.

Therefore, across both the Mol2Cap and Cap2Mol tasks, MolReFlect consistently demonstrates state-of-the-art or comparable performance, affirming the effectiveness of our approach.

To enable a better understanding of MolReFlect, we conduct a series of ablation studies to resolve the research questions that have been raised for discussion.

For this question, we conduct an ablation study on MolReFlect by removing context examples and fine-grained alignments, downgrading MolReFlect to Instruction Tuning and ICMT, respectively. Meanwhile, we also provide the naive-SFT performance of Mistral-7B. The results are presented in Table 3 and Table 4. It is evident that the naive-SFT results are actually unsatisfying as Mistral-7B lacks specific pre-training on chemical corpora. Meanwhile, when only the context examples are removed, the performances drop slightly but attain a BLEU-4 score of 0.539 on the Mol2Cap task and a BLEU score of 0.886 on the Cap2Mol task, demonstrating a significant performance improvement compared to naive-SFT. Notably, in the Cap2Mol task, the exact match score nearly doubles compared to naive-SFT, indicating that the fine-grained alignments indeed convey much molecular structure information to the student LLM. Furthermore, when fine-grained alignments are removed during the fine-tuning phase, the performances drop in both Mol2Cap and Cap2Mol tasks. This suggests that the LLMs are able to learn molecule-text alignments more effectively from the fine-grained alignments in the context examples, leading to better final generations.

We also include several cases in Appendix C and conduct a series of extensive experiments in Appendix B for better explanation. As depicted in Figure 5, the larger teacher LLM can generate preliminary indications towards the final target and even directly figure out the molecular structure in fine-grained alignments. However, some of these indications might be inaccurate. With CoT-ICMT, the smaller student LLM could learn from the input distribution, identify these errors, and correct them in the final generation process. In this case, as illustrated in Figure 4, the output distribution generated by MolReFlect aligns better with the ground truth. Conversely, MolT5 and ICMA fail to achieve this owing to the lack of fine-grained alignments.

To resolve this question, we ablate MolReFlect by removing the in-context reflection and the selection processes, which is equivalent to replacing the fine-grained alignments with zero-shot alignments and in-context reflected alignments, respectively. The details are shown in the last two rows of Table 3 (for the Mol2Cap task) and Table 4 (for the Cap2Mol task). From Table 3, we can observe that the results without in-context reflection lead to sub-optimal performance as the teacher LLM could make mistakes or yield hallucinations, underscoring the necessity of in-context reflection. However, the in-context reflected alignments are not necessarily better than zero-shot alignments, as evidenced by Table 4. Sometimes, the zero-shot alignments of similar molecules/captions may contain more noises, like hallucinations and factual errors, than helpful information and inadvertently become part of the context in the in-context reflection phase. The inaccuracies could then carry over to the in-context reflected alignments, potentially harming the final performance. In this case, the zero-shot alignments can be more helpful as the context examples do not pollute them. Therefore, choosing between zero-shot alignments and in-context reflection alignments is imperative to ensure the quality of fine-grained alignments.

From the information theory perspective, our objective is to provide LLMs with more helpful information and less noise while rigorously preventing any disclosure of information about the target. Therefore, perplexity, an unsupervised metric, is an ideal criterion for the selection process. Higher perplexity scores suggest the presence of information that conflicts with the existing knowledge of LLMs, making it a reliable indicator for discerning the quality of the generated alignments. In this work, we utilize the Galactica-125M as the student model to calculate perplexity, which is particularly adept at chemical tasks and offers rapid computation. The alignments with the lower perplexity scores are selected as the fine-grained alignments. According to Table 3 and 4, across both the Cap2Mol and Mol2Cap tasks, MolReFlect consistently demonstrates superior performance compared to those without in-context reflection or selection, thereby substantiating the effectiveness of In-Context Selective Reflection and Selection.

In this part, we address the last research question by removing the student model and completing the tasks using only the teacher LLM (i.e., Llama-3-70B). Since the cost of fine-tuning the teacher LLM is unaffordable for most institutions, we only test the performance of teacher LLM with prompt engineering to avoid modifications of their parameters. Various prompting strategies are implemented to enable the teacher LLM to undertake the molecule-caption translation tasks independently, including direct prompting, chain-of-thought prompting, few-shot prompting, and few-shot chain-of-thought prompting. Notably, in the chain-of-thought and few-shot chain-of-thought prompting, we utilize the fine-grained alignments produced by the teacher LLM itself as context information. The results of these experiments are detailed in Table 5 and 6.

It can be observed that while Llama-3-70B is a powerful LLM, its performance under direct prompting is notably weak, as it is not trained on the ChEBI-20 or a lot of chemical corpora, ensuring that the information of the ChEBI-20 dataset is not leaked in its pre-training stage. In the Mol2Cap task, the chain-of-thought strategy enhances the performance by introducing fine-grained alignments. However, in the Cap2Mol task, the performance declines by 1.05%, indicating that the teacher LLM struggles to filter out the noise inherent in the fine-grained alignments without explicit supervisory signals. Similarly, in the few-shot setting, the fine-grained alignments also do not contribute to a significant performance boost for the teacher LLM. In contrast, the student LLM proves to be indispensable and could benefit from the CoT-ICMT process by enabling a better understanding of molecule-text alignments and identifying the noises behind fine-grained alignments. As shown in Table 3 and 4, the Instruction Tuning (i.e., w/o Context Examples) performance increases by 9.94% and 14.22% in the Mol2Cap and Cap2Mol tasks, respectively, compared to the naive-SFT. This further underscores the necessity of discerning and mitigating noise within the fine-grained alignments, suggesting that LLMs must engage in fine-tuning to learn from the fine-grained alignments effectively. Thus, the teacher-student framework proves to be indispensable. It enables the smaller student LLM to learn from the input distribution, discern noise in the content generated by the teacher, and absorb valuable information to inform the final generation process.