Could Chemical Language Models benefit from Message Passing (2024)

Abstract: Pretrained language models (LMs) showcase significant capabilities in processing molecular text, while concurrently, message passing neural networks (MPNNs) demonstrate resilience and versatility in the domain of molecular science. Despite these advancements, we find there are limited studies investigating the bidirectional interactions between molecular structures and their corresponding textual representations. Therefore, in this paper, we propose two strategies to evaluate whether an information integration can enhance the performance: contrast learning, which involves utilizing an MPNN to supervise the training of the LM, and fusion, which exploits information from both models. Our empirical analysis reveals that the integration approaches exhibit superior performance compared to baselines when applied to smaller molecular graphs, while these integration approaches do not yield performance enhancements on large scale graphs.

## introduction the success of attention mechanisms on sequential data has introduced a massive family of large language models based on transformer architecture @xcite . it is evident that these large language models are useful for encoding sequential objects such as text @xcite , molecules @xcite , speech @xcite , and forecasting data @xcite . it has been demonstrated that pretrained molecule language models are capable of encoding chemical elements semantically without learning structures @xcite @xcite . especially for proteins which function as natural components of the human body and a representative of molecule family, they could be efficiently encoded by transformer @xcite @xcite which acts as masked language modelers. in contrast to text, molecules contain inherent relationships between their elements, indicating that structural encoding is necessary in addition to word embeddings. message passing neural network (mpnn), emerging as a prominent method for encoding structural information in recent years, has demonstrated its robustness and versatility within the field of molecular sciences. by leveraging the 2-dimensional topological and 3-dimensional geometrical information as augmented features @xcite , it is possible to learn molecular embeddings from structures without sequentially encoding traditional smiles expressions. the advent of mpnns has promoted the exploration of graph-based learning methods for molecular science. graph contrastive learning captures potential different structural distributions to finetune self-learned representations, where both local and global features are enhanced with chemical domain expertise @xcite @xcite . besides, the success of gpt @xcite in traditional natural language processing tasks also motivates the research on graph transformers and graph gpt tailored for the molecule domain @xcite @xcite @xcite . few studies has been investigated to appropriately merge text embeddings and graph embeddings for learning molecule representation better. there has been one study which demonstrated such relationship but with additional prompting with gpt model @xcite , which is out of our scope. in this paper, we aim to explore the interplay between molecular graph embeddings and smile token embeddings. we propose two categories of techniques for integrating information: contrast learning and fusion. in contrast learning-based methods, we incorporate an mpnn as an auxiliary model to supervise the training of the language model, operating at node or graph levels, while we only utilize the language model for downstream tasks. in fusion-based methods, we exploit information from both models to generate outputs for downstream tasks. this is achieved either by merging the output embeddings from both models or by integrating the output embeddings from one model with the input embeddings of the other. our main contributions are to as follows: ## methods we investigate two kinds of information merging methods: contrast learning-based methods and fusion-based methods. in contrast learning-based methods, we use gnn as an auxiliary model to supervise the training of the language model, while we only use the language model for downstream tasks. in fusion-based methods, we make use of information from both models to generate the output for downstream tasks. for each kind of method, we consider different model architectures. in this section, we will first briefly review the two baseline models chemberta @xcite and nnconv @xcite , and then describe the contrast learning-based methods and the fusion-based methods. we ask the chemberta model to produce threelevel information for a molecule. first, node em-beddings are extracted from the final hidden state representations. since the smiles transformation preserves atom orderings, each atom in the original molecule corresponds to a specific output token embedding. second, the graph embedding is extracted from the special token at the beginning of the sequence. third, the property prediction result is the final output. the entire process can be depicted as follows: message passing neural networks. there are different types of graph neural networks, which include graph convolution, graph attention and neural message passing networks (mpnn). edge attributes or edge features are important in message passing mechanisms @xcite . for the baseline model, we follow the model setting in the first paper of mpnn for quantum chemistry dataset qm9 @xcite , which iteratively updates the message m t v and the hidden state h t v for each node v: aggr[•] is the function that aggregates neighbor node u's information as well as the attributes e uv of the shared edge with u. u(•, •) updates hidden states for v. node level contrastive learning in order to compute the constrastive loss, we need a triple ⟨anchor, positive, negative⟩. anchor and positive node embeddings are sampled from language models and message passing networks respectively. negative samples are randomly generated from graphs with a different permutation. for an example node triple @xmath0, p, n⟩, its corresponding contrastive learning loss triplet loss is given by: graph level contrastive learning. apart from establishing negative samples between each pair of nodes at a fine grained level, we estimate contrastive learning at a coarse grained level which aims at computing the difference between language model graph embeddings and mpnn graph embeddings. this could potentially avoid the situation that individual nodes with large difference contribute more to the difference of the molecule property. moreover, the complexity is pretty lower than the complexity of node level comparison, which will be discussed in part 3.4. similar to the node level training loss in ( 6 ), the graph level training loss is defined as: ## results observation 0: protein language models are more preferred. a fundamental observation from experimenting on moleculenet is that purely using message passing neural networks are inferior to language models in molecule property prediction. this phenomenon is also mentioned in the previous research work @xcite . this has indicated some works to include the geometric properties such as 3d information and rotation invariant parameters in message passing observation 1: integration on relatively small graphs are more preferred. using the pretraind chemberta-2, we found that both contrastive learning and fusion methods outperform baseline models in esol, bace, and bbbp where they are relatively small compared with qm9 and hiv datasets. especially, node level contrastive learning performs the best and it seems to be robust among all tasks, followed by late fusion methods and joint fusion methods when injecting llm to mpnns. in large dataset, the tuning strategy might influence the potential performance, where it splits the dataset in a better way therefore we perform one ablation regarding train test split (in section 6) to avoid the difference that brought by dataset itself. observation 2: integration w.r.t both regression and classification are useful. in terms of training convergence, we observe that the accuracy or mean absolute error converges quickly to a high or low score respectively. for small graph datasets bace and bbbp on graph classification problem, an improvement of ≈1% on average accuracy is observed with method mpnn2lm for pretrained chemberta. for version 2, 1.4% improvement is observed with late fusion on bace and 1.3% improvement is observed with node contrastive learning on bbbp . for small graph dataset esol on regression problem, a great improvement is observed where 12.8% improvement on mae with mpnn2lm method with pretrained chemberta, and 7.3% improvement on mae with mpnn2lm method with pretrained chemberta-2. for hiv, we observe a little improvement with node level contrastive learning. using a combination of llm representation and graph representation during the training would make the prediction worse. for qm9, most of the injection / fusion methods would potentially improve the performance except for mpnn2lm fusion. using lm2mpnn would potentially improve 8.6 %. we found that pure mpnn's performance is better than the performance of a chemical llm ( observation 3: pretrained language models are important for downstream predictions. in comparison to chemberta-2, chemberta performs worse when comparing each entry in when it comes with a new pretrained large language model, using our proposed method could tell the similarity between tasks and the model's pretraining strategy. as there is no general conclusion about how a chemical llm and a mpnn could be combined to predict the best, it is still a pioneering area that requires more pretrained models to test its robustness. to select the most appropriate pretrained language model for further training, researchers should first integrate a list of pretrained models, followed by an investigation with different fusion / injection methods. observation 4: joint fusion to some extent helps learn mpnn better but learn original chemical llm worse. we also focus on if such multimodal module (fig. 2 , fig. 3 ) helps learn individual module (fig. 1 ) better. it improves a lot for single mpnn baseline if we consider its language level information as augmented features. for example, for bace dataset, mpnn has an average accuracy of 0.808 with chemberta. with injecting pretrained language information, an improvement of 3.7% is observed (lm2mpnn). however, it might not work very well on the opposite when we inject information from mpnn to llm. for simplicity we just examined with pretrained chemberta-2. for esol dataset, it decreased from 0.4738 to 0.5012 (5.78%). for bbbp dataet, it decreased from 0.9171 to 0.8974 (2.15%). we further suggest that the researchers should not directly use the structural information from graphs as additional input when they want to modify their llm models, but trying to leverage them as auxiliary ground-truth to finetune the token embeddings. effects of dataset split. we want to figure out if different splits of training, validation and test datasets lead to different performance. we run on bbbp (classification) and esol (regression). four ratios are considered: 9:0.5:0.5, 8:1:1, 7:2:1, and 6:2:2. model prediction power is highest at a ratio of 8:1:1 for esol while the prediction power is reducing for bbbp when ratio of training sets is decreasing. ## conclusion in this paper, we delved into various information integration approaches to assess whether the collaborative utilization of chemical large language models (chemical llms) and message passing neural networks (mpnns) surpasses the individual efficacy of these models. we evaluated the integration approaches over different graph scales on both classification and regression tasks. our empirical analysis has demonstrated that the integration approaches outperform the baselines on small-scale graphs but do not yield improvements on datasets of larger scales. furthermore, we have found that differences in dataset splitting strategies, and aggregation choices in fusion have an impact on the overall performance. we wish to extend our proposed methods on large scale benchmark datasets such as pcqm4mv2 @xcite .