SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Haiteng Zhao $^{1}$ , Shengchao Liu $^{2}$ , Chang Ma $^{3}$ , Hannan Xu $^{4}$ , Jie Fu $^{5}$ , Zhi-Hong Deng $^{1}$ , Lingpeng Kong $^{3}$ , Qi Liu $^{3}$
<sup>1</sup> Peking University <sup>2</sup> Mila <sup>3</sup> The University of Hong Kong
<sup>4</sup> University of Oxford <sup>5</sup> Hong Kong University of Science and Technology
Molecule property prediction has gained significant attention in recent years. The main bottleneck is the label insufficiency caused by expensive lab experiments. In order to alleviate this issue and to better leverage textual knowledge for tasks, this study investigates the feasibility of employing natural language instructions to accomplish molecule-related tasks in a zero-shot setting. We discover that existing molecule-text models perform poorly in this setting due to inadequate treatment of instructions and limited capacity for graphs. To overcome these issues, we propose GIMLET, which unifies language models for both graph and text data. By adopting generalized position embedding, our model is extended to encode both graph structures and instruction text without additional graph encoding modules. GIMLET also decouples encoding of the graph from tasks instructions in the attention mechanism, enhancing the generalization of graph features across novel tasks. We construct a dataset consisting of more than two thousand molecule tasks with corresponding instructions derived from task descriptions. We pretrain GIMLET on the molecule tasks along with instructions, enabling the model to transfer effectively to a broad range of tasks. Experimental results demonstrate that GIMLET significantly outperforms molecule-text baselines in instruction-based zero-shot learning, even achieving closed results to supervised GNN models on tasks such as toxcast and muv.
Molecule machine learning has gained significant attention in recent years [10, 52, 85, 77], including tasks like property prediction [14, 21], molecule design [81, 35, 28, 27], and others, which have broad applications in biomedical research. The primary method of molecule tasks is graph machine learning [21], where graphs are employed to represent molecule topology. Presently, graph machine learning approaches for molecules mainly follow the pretraining and finetuning paradigm [26, 60, 40, 84]. After the pretraining on large molecule corpus, models are able to encode informative molecule representations and generalize to downstream tasks by supervised finetuning.
One limitation of the supervised finetuning approach is the requirement on labeled data to acquire task-specific features for prediction on downstream tasks. This requirement poses a challenge in scenarios where obtaining labeled data is arduous or costly, especially in the field of molecules, and the training cost also restricts the efficiency of downstream tasks. Additionally, the supervised method exclusively relies on labeled data for acquiring task information, and is therefore unable to incorporate other additional information about the task. For instance, there exists a wealth of information regarding molecule assay tasks often provided in textual descriptions. Unfortunately, the supervised method fails to leverage such information, thereby restricting its flexibility.
<sup>1</sup>The code, model, and data are available at https://github.com/zhao-ht/GIMLET.
In this work, we propose to investigate the feasibility of employing instructions to accomplish molecule-related tasks in a zero-shot setting. The instructions, also referred to as prompts [50, 6], constitute natural language descriptions of the task to be executed. This is inspired by recent progress in natural language processing, where the large language models possess strong generalization performance to unseen tasks by following instructions [71, 47]. This approach eliminates the need for labeled data and leverages the textual knowledge available for downstream tasks.
Several studies have explored the molecule-text language models, employing either autoregressive pretraining using language models on both the text and molecule strings (SMILES) [16, 85, 64], or contrastive pretraining of molecule and text where molecules are encoded by Graph Neural Networks (GNN) [17, 59, 39, 56]. However, these works lack the investigation into instruct-based zero-shot for downstream tasks, and our experiments have empirically demonstrated their limited performance.
We conjecture that the constraints of current molecule-text language models are mainly imposed by their inadequate treatment of instructions during pretraining and their limited capacity to represent molecule structures. First, the pretraining corpora of these models lack task description about descriptions for abundant molecule tasks as well as the supervised signal. They mainly address the capacity of general molecule-related text understanding, which may be insufficient to capture details of complex instructions to molecule properties. Second, these models' capacity is limited in representing graph structures of molecules. Molecular tasks heavily rely on understanding the structure information, while existing multimodal methods either encode SMILES sequence representation of molecules by language model or encode molecular graphs into dense vector representation using GNN, which are not conducive for language models to gain a deep understanding of graph features.
To facilitate natural language instruction-based graph learning for zero-shot molecule tasks, we propose an innovative structured language model called GIMLET (Graph Instruction based Molecule zEro-shoT learning). First, GIMLET extends large language models for both graph and text data by applying transformer mechanisms with generalized position embedding, ensuring strong capacity for both the learning of graph structure representations and the executing of instruction texts without introducing additional graph encoding modules. We further enhance the model by decoupling the encoding of the graph from specific tasks via employing attention masks. This allows for improved generalization of graph features across novel tasks.
In addition to the advanced model architecture, our approach incorporates instruction-based pretraining on an extensive dataset of molecule tasks with textual instructions. We construct a dataset comprises of 2K tasks accompanied by corresponding instructions tailored for instruction-based molecule zero-shot learning framework. Throughout the pretraining process, molecule graphs are paired with natural language instructions for molecule-related tasks, and the supervision signal is provided to train the model in executing various molecule tasks. This methodology empowers GIMLET to comprehend specific instructions for downstream tasks expressed in natural language and transfer acquired knowledge to a broad range of tasks effectively.
Our experimental results show that GIMLET outperforms molecule-text baselines by a substantial margin in instruction-based zero-shot learning, which is closed to supervised GNN on tasks like muv and toxcast. Additionally, GIMLET also exhibits impressive performance in few-shot learning and demonstrates robustness to instruction. In summary, our contributions are as follows:
- We present comprehensive investigations of natural language instruction-based graph zero-shot learning for molecule tasks, to reduce reliance on labeled data and leverage task textual knowledge. To accomplish this framework, we construct a molecule dataset consisting of two thousand tasks with instructions derived from task descriptions, which is open-sourced.
- We propose GIMLET, which extends language models to handle graph and text data. By applying the transformer mechanism with generalized position embedding and decoupled attention, our model learns graph structure representations and executes instruction texts without additional graph encoding modules. Instruction-based pretraining is applied for GIMLET to comprehend specific instructions expressed in natural language and transfer to a broad range of zero-shot tasks.
- Our experiment results outperform other molecule-text models by a substantial margin and demonstrate promising results in instruction-based zero-shot molecule tasks, demonstrating the viability of this framework. Additionally, GIMLET exhibits impressive performance in few-shot learning and demonstrates robustness to instruction.
Figure 1: Our framework handles molecule tasks in the zero-shot fashion by natural language instruction. Within GIMLET, we employ distance-based joint position embedding to encode graphs and instruction texts. Additionally, we utilize attention masks to decouple the graph encoding process.
Molecule Representation Learning One approach for molecular representation learning is utilizing language models to acquire molecular representations based on Simplified Molecular Input Line Entry System (SMILES) strings [68, 10]. Toward stronger capability to incorporate pertinent substructure information, Graph Neural Networks (GNNs) are proposed to model molecules as graphs [21, 83, 26, 36]. Existing GNNs follow the message-passing paradigm and suffer from problems like long-range dependency vanishing and over-smoothing. Recently, graph transformer [52, 80] has been proposed to better encode structures of graphs, illustrating effectiveness in molecule tasks [42, 31, 8, 88].
Molecule Pretraining To fully explore the inherent structural information of molecules on a large scale, significant efforts have been made to molecular pretraining. Supervised pretraining is commonly used for learning useful representations [26, 80, 62]. As for unsupervised pretraining, one approach involved the generative strategy on molecular SMILES strings [68, 25, 10, 3, 53] and graph [26, 37, 38, 52, 87], which was followed by recent works adopting the contrastive paradigm to distinguish augmented views of the same graph and other graphs [67, 60, 24, 83, 82, 63, 78, 19, 61, 70, 76, 69, 40].
Besides the structure-only pretraining, a few recent works incorporate natural language into molecule pretraining. One class of method is the SMILES-based language model, including KVPLM [85] and MolT5 [16], which use SMILES strings and text for joint representation and translation. Another work Galactica [64] explored the multi-task molecule task learning with instruction. Some other works acquire advanced representations for molecules by GNN, such as Text2Mol [17], MoMu [59], MoleculeSTM [39], and CLAMP [56], trained by contrastive learning between molecule graph and text description for molecule retrieval and caption tasks. MoleculeSTM and CLAMP explored molecule editing and property prediction with text. However, previous works lack the investigation into instruct-based zero-shot scenarios for complex molecule tasks like property prediction.
Instruction-based Zero-Shot Learning Instruction-based zero-shot learning is an innovative approach that leverages natural language instructions to enable neural models to solve a variety of tasks [50, 6, 55, 18, 89, 44, 45, 49]. To enhance the model's ability to follow instructions, some researchers have employed instruction-based pretraining techniques [54, 71, 12, 47], which explicitly train language models to solve tasks with instructions. Besides natural language processing, instruction-based zero-shot learning is also studied in multimodal domains like images [4, 9, 1, 32].
3.1 Problem Formulation and Framework Overview
The molecule task is denoted as  $\tau$  and consists of a set of graph data and their corresponding labels  $(G_{i},y_{i}^{\tau})$ . A molecule graph  $G = (N,E,v,e)$  includes nodes  $N = \{1,\dots ,n\}$  and edges  $E\subset N\times N$  while  $v$  and  $e$  represent node features and edge features, respectively. The label  $y^{\tau}$  can take various
forms, such as class labels or numerical values, depending on the type of task. The instruction of  $\tau$  is denoted as  $T^{\tau}$ , which is a sentence  $[o_1,\dots,o_m]$  describing the task.
In the supervised approach, a model  $F_{\theta}$  predicts the label  $\hat{y}$  given a graph  $G$ , i.e.,  $\hat{y} = F_{\theta^{\tau}}(G)$ , by supervised finetuning individually for each downstream task  $\tau$ . The limitation is that it relies on labeled data to learn the corresponding parameters and output modules, and finetuning does not enable the model to effectively utilize extra task-related knowledge.
To overcome these limitations, our framework leverages natural language instructions  $T^{\tau}$  to provide the model with task information, as illustrated in Figure 1. Our zero-shot graph learning framework incorporates molecule graphs  $G$  and task instructions  $T^{\tau}$  into a graph-text language model GIMLET and decodes the output as text uniformly for different tasks, i.e.  $\hat{y}_{\mathrm{str}} = \mathrm{GIMLET}(G, T^{\tau})$ , where  $\hat{y}_{\mathrm{str}}$  is the label string. This description-based instruction framework empowers the model to handle a wide range of molecule tasks as long as they can be described in natural language, and the uniform text decoding approach accommodates various types of tasks including classification or regression.
Previous pretrained molecule-text models perform poorly in our framework due to the inadequate treatment of instructions and limited capacity for graphs. Our method addresses these from two aspects: First, we propose a unified language model GIMLET for both graph and text, towards the stronger capacity to represent molecules and instructions; Second, we adopt instruction-based pretraining to GIMLET, enabling generalization to new tasks based only on the instructions.
3.2 Unified Graph-Text Transformer
The common method for multimodal language models is obtaining the feature of the other modality by applying an additional encoding module, then integrating it into the language model, as in other molecule language models with GNN encoder [17, 59, 56]. The individual encoding module benefits for decoupling the graph feature encoding from instructions, i.e., the features of graph data can be independent of task instructions in the early encoding stage, helping for generalization when the distribution of tasks changes.
However, for the molecule-text model, individual pre-encoding modules present problems. First, graph learning relies on structure information, but the dense vectors encoded by GNN have a limited capacity to carry structure information. Furthermore, training the additional module is difficult due to the increased layers, since deep transformers have vanishing gradients in early layers [34, 2]. Lastly, the additional modules increase parameters and training costs.
To overcome these issues, we propose a novel approach GIMLET which not only directly unifies the standard language model for graph and text without introducing additional graph encoder module, but also remains the decoupled graph encoding for better generalization.
Given graph  $G$  and text input  $T$ , to utilize GIMLET, we represent the graph nodes and text tokens as tokens. The resulting hidden states are denoted as  $H = [h_1,\dots,h_n,h_{n + 1},\dots,h_{n + m}]^T\in \mathbb{R}^{(n + m)\times d_h}$  for corresponding  $n$  graph nodes and  $m$  text tokens.
In this study, we choose T5 [50] as the backbone language model, due to the encoder-decoder architecture suitable for non-sequential encoding and text output. It utilizes the relative position embedding method [57] to represent sequential structure. In attention layer Attn with parameter  $W^{V}, W^{Q}, W^{K} \in \mathbb{R}^{d_{h} \times d_{k}}$  and  $W^{O} \in \mathbb{R}^{d_{k} \times d_{h}}$ , relative position embedding for i-th and j-th token is formalized as
$$
\hat {A} _ {i j} = \frac {\left(h _ {i} W ^ {Q}\right) \left(h _ {j} W ^ {K}\right) ^ {T}}{\sqrt {d _ {k}}} + b (i, j), A = \operatorname {s o f t m a x} (\hat {A}), \quad \operatorname {A t t n} (H) = A H W ^ {V} W ^ {O}, \tag {1}
$$
where  $b(i,j)$  is embedding of the relative distance between  $i$  and  $j$ , i.e.  $i - j$  for sequence. For graph-text joint data, we construct the position embedding  $b(i,j)$  in E.q. 1 by the conjunction of different types of distances. For the relative position of graph nodes, we adopt the graph shortest distance, which has been widely used in the literature of graph transformer [80, 11, 48]. We also use unidirectional constraint in attention to decouple the graph encoding from the instruction. The overall form of position embedding for GIMLET is:
$$
b (i, j) = b _ {\mathrm {P O S} (i, j)} ^ {D} + b _ {i, j} ^ {M} + \underset {k \in \mathrm {S P} (i, j)} {\text {M e a n}} b _ {e _ {k}} ^ {E}, \tag {2}
$$
where  $b^{D}, b^{M}$ , and  $b^{E}$  are position bias, masking, and path embedding, individually. The relative position POS in  $b_{\mathrm{POS}(i,j)}^{D}$  is defined as the conjunction of different types of distances between tokens,
which allows for the effective encoding of both graph and text data, as well as their interaction:
$$
\operatorname {P O S} (i, j) = \left\{ \begin{array}{l l} i - j & \text {i f} n + 1 \leq i, j \leq n + m \\ \operatorname {G R A P H S H O R T E S T D I S T A N C E} (i, j) & \text {i f} 1 \leq i, j \leq n \\ <   \mathrm {C r o s s} > & \text {o t h e r w i s e} \end{array} , \right. \tag {3}
$$
where  $<\mathrm{CROSS}>$  is a special distance token held out for cross distance between graph and text tokens.  $b_{i,j}^{M}$  aims to represent the cross mask used to decompose graph encoding from text instructions. It imposes a unidirectional constraint from the graph to the text:
$$
b _ {i, j} ^ {M} = - \infty \text {i f} i \leq n \text {a n d} j > n \quad \text {o t h e r w i s e} 0 \tag {4}
$$
With the unidirectional constraint, graph tokens are limited to attending only to other graph tokens. On the other hand, instruction tokens have the ability to receive information from both instructions and graphs. This approach allows us to separate the encoding of graphs from instructions, enabling instructions to selectively utilize graph features for various downstream tasks.
Finally,  $\mathrm{Mean}_{k\in \mathrm{SP}(i,j)}b_{e_k}^E$  is the mean pooling of the edge features  $b_{e_k}^E$  in the shortest path  $\mathrm{SP}(i,j)$  between node  $i$  and  $j$ , which is only defined between graph node tokens, as used along with graph shortest distance in [80].
The generalized position embedding is applied to the encoder transformer. During decoding, the decoder generates the answer based on the text features outputted by the encoder.
GIMLET unifies graph and text data by a single language model, which has the following merits: (i) In comparison to additional encoding module methods, GIMLET not only avoids the challenges of training additional front layers but also provides a stronger capacity for handling graph data by introducing graph inductive bias to the whole transformer. (ii) Our method leverages both the existing knowledge within the language model and facilitates learning to execute instruction-based graph tasks through standard instruction-based learning on the language model. This approach eliminates the need for additional modeling costs. (iii) The decomposed encoding of graph data retains the advantages of the individual encoding module, reducing task disturbance and ensuring that data features remain independent of task instructions. This enhances generalization for novel instructions. We validate these claims in experiments Subsection 4.3.
3.3 Pretraining and Datasets
Figure 2: (Left) Illustration of datasets. Circle size corresponds to task number. Tasks are organized by category. Tasks on the top are more related to biological assay, on the bottom need more chemical and physical properties. GIMLET is trained on pretraining tasks, then tested on downstream tasks in the zero-shot setting. (Right) Our task instructions contain task explanations and questions.
**Pretraining** Our approach leverages the generalization capabilities acquired through learning from the instructions provided during pretraining, where comprehensive linguistic information and knowledge related to molecular tasks are learned from the provided instructions. The pretraining process is
conducted in a supervised manner using task labels. The loss function for supervised pretraining can be formalized as follows:
$$
L = \frac {1}{\sum_ {\tau} | \tau |} \sum_ {\tau} \sum_ {(G _ {i}, y _ {s t r, i} ^ {\tau}) \in \tau} - \frac {1}{| y _ {s t r , i} ^ {\tau} |} \log P _ {\text {G I M L E T}} \left(y _ {s t r, i} ^ {\tau} \mid G _ {i}, T ^ {\tau}\right), \tag {5}
$$
where  $P_{\mathrm{GIMLET}}$  is the model likelihood for the label string,  $|y_{str,i}^{\tau}|$  represents the label string token number aiming to normalize the loss of long sequences, and  $|\tau|$  is the task sample number. The details of pretraining and downstream zero-shot testing are in Appendix.
Pretraining Dataset To effectively pretrain GIMLET for downstream tasks, it is crucial to include a large number of tasks to provide the model with an extensive task corpus. To this end, we select Chembl [20] as the pretraining dataset, which is widely used for supervised graph pretraining [26, 62]. It consists of 1,310 prediction target labels from biological assays for drug discovery. We divided  $80\%$  of the tasks and  $80\%$  of the molecules in random for pretraining, while the remaining non-overlapping tasks were reserved for zero-shot testing. Additionally, we constructed the Chembl-property dataset, which encompasses various physical and chemical properties available in the Chembl database [20] like molecule weights, log p values, hydrogen bound acceptor numbers, etc. Full details are in Appendix. We validate the effect of Chembl biological tasks and Chembl-property physico-chemical tasks in pretraining in Subsection 4.3.
Downstream Dataset We target a large range of molecule tasks, as shown in Figure 2. First, we include large-scale datasets PCBA [72], which contains 128 biological assay tasks with abundant molecules. We also include the hold-out zero-shot set of Chembl, noted as Chembl Zero-Shot. These two datasets form a large-scale dataset. We also target tasks from MoleculeNet [75], a popular benchmark for molecule properties prediction. We adopt the dataset categorization method in [58] which classifies MoleculeNet datasets into four categories: Physico-chemical tasks, Bio-activity tasks, Toxicity tasks, and Pharmacokinetic tasks. Additional dataset CYP450 [33] is also included in the Pharmacokinetic tasks. These tasks cover diverse aspects of molecule properties, posing a significant challenge for a unified model to simultaneously handle them in a zero-shot fashion.
Instructions To provide essential background information and context for each task, we include task explanations and descriptions in our instructions. The description covers a wide range of aspects, including the family, function, and mechanism of the assay target, the assay experiment setting, the approximation method used for determining the property, and others. See examples in Figure 2. Instructions for all tasks are available in the code file. The task explanation is primarily sourced from websites and databases that introduce and compile the respective datasets, or relevant papers. Details are in Appendix. The descriptions are then concatenated with relevant questions as instructions. These instructions are subsequently reviewed and validated by a professional biology Ph.D. student.
In the experiments, we investigate the following inquiries: (i) Can GIMLET effectively handle zero-shot molecule property tasks by instructions? (ii) Can GIMLET performs better by few-shot learning? (iii) What impact does model architecture have on the performance of GIMLET? (iv) How does pretraining affect the performance of GIMLET? (v) How does the form of instruction influence GIMLET for molecule zero-shot learning?
4.1 Instruction-Based Zero-Shot Learning