Unified Molecule-Text Language Model with Discrete Token Representation

1 Introduction

The incredible capabilities of Large Language Models (LLMs) [Brown et al., 2020; Touvron et al., 2023] have led to their widespread use as versatile tools for completing diverse real-world tasks. This success has sparked interest in Multi-modal LLMs [Zhan et al., 2024], which aim to enhance LLMs by enabling them to process multi-modal inputs and outputs. In fields like molecular science, Multimodal LLMs present new opportunities by seamlessly integrating molecular data with textual information, opening up fresh possibilities for more efficient and accurate research and development.Prior research efforts [Liang et al., 2023; Fang et al., 2023; Cao et al., 2023; Liu et al., 2023b; Li et al., 2024] have focused on adapting LLMs to molecular tasks, resulting in the development of molecular LLMs. These molecular LLMs can analyze molecule structures [Liu et al., 2023b; Cao et al., 2023], address drug-related inquiries [Liang et al., 2023], assist in synthesis and retrosynthesis planning [Fang et al., 2023], support drug design [Fang et al., 2023], and more.

∗Corresponding author

Prevalent molecular LLMs often use adapter-based architectures, such as linear projection [Liang et al., 2023; Cao et al., 2023] or Q-Former [Liu et al., 2023b; Li et al., 2024], to map molecule features into the LLM's semantic space (Figure 1a, Figure 1b). While effective in molecular comprehension and molecule-to-text generation, these models struggle with text-to-molecule generation. This is due to the reliance on adapters that require LLMs to directly generate SMILES strings [Weininger, 1988], a text-based representation of molecular structures. These architectures depend on strong alignment between SMILES and text, but as shown in Figure 1a and Figure 1b, molecule and text modalities are not treated equally, with insufficient supervision for the molecular side, making alignment difficult.

Discretizing continuous molecule features into discrete molecule tokens offers a promising solution for conducting both molecule-to-text and text-to-molecule generation tasks. By treating tokens from different modalities equally, we can predict the next molecule or text token in an autoregressive manner. However, directly discretizing molecule features poses several challenges: (i) This approach results in long sequences, with lengths equivalent to the number of atoms in a batch; (ii) Molecule tokens derived from molecule features lack left-to-right causal dependency, which conflicts with the unidirectional attention mechanism in LLMs; (iii) Molecule features lack textual information, hindering effective moleculetext interactions and alignment.

To this end, we present UniMoT, a Unified Molecule-Text LLM that adopts a tokenizer-based architecture, integrating molecule comprehension and generation, as depicted in Figure 1c. A pivotal aspect of UniMoT's architecture is the molecule tokenizer for transforming molecules into molecule tokens. We introduce a Vector Quantization-driven [Van Den Oord et al., 2017] tokenizer, which incorporates a QFormer [Li et al., 2023] to bridge the modality gap between molecules and text. Specifically, we incorporate causal masks for the queries, enabling the Q-Former to generate a causal sequence of queries compatible with the unidirectional attention in LLMs. The sequence of queries is subsequently quantized into a sequence of molecule tokens using a learnable codebook. The molecule tokens encapsulate high-level molecular and textual information, which are then aligned with the latent space of a pretrained generative model via an MLP adapter.

Pretrained LLMs can integrate the molecule tokenizer by treating molecule tokens as new words and constructing a molecule vocabulary through mapping the learned codebook. We adopt the unified discrete token representation for molecules and text, coupled with the unified next-tokenprediction training paradigm of LLM. This unification of representation and training paradigm enables effective moleculetext interactions and alignment through molecule-to-text and text-to-molecule autoregressive pretraining. For molecule generation tasks, UniMoT generates molecule tokens in an autoregressive manner rather than producing SMILES strings, and these molecule tokens can then be decoded into molecules using the generative model.

Our contributions can be summarized as follows:

• We introduce a molecule tokenizer specifically designed for LLMs, enabling the tokenization of molecules into short sequences of tokens with causal dependency. These tokens encapsulate high-level molecular and textual information and can be decoded into desired molecules during inference.
• We present UniMoT, a unified molecule-text LLM that adopts a tokenizer-based architecture instead of traditional adapter-based architectures. UniMoT unifies the modalities of molecule and text under a shared token representation and an autoregressive training paradigm. Following a four-stage training scheme, UniMoT effectively achieves molecule-text alignment.
• UniMoT exhibits remarkable capabilities in multi-modal comprehension and generation. Extensive experiments show that UniMoT achieves state-of-the-art performance across a wide range of comprehension and generation tasks, while also offering a new perspective on molecule generation.

2 RelatedWorks

Multi-modal Large Language Models. Multi-modal Large Language Models (LLMs): Current multi-modal LLMs are typically built on a pre-trained LLM backbone and can understand multiple modalities. LLaVA [Liu et al., 2024a] connects the image encoder to the LLM using a simple linear projection, while BLIP-2 [Li et al., 2023] extracts high-level features from images with CLIP [Radford et al., 2021] and uses QFormer to reduce image token counts. While these models excel at multi-modal comprehension, they often lack focus on multi-modal generation. To address this, recent work unifies multi-modal comprehension and generation, such as SEEDLLaMA [Ge et al., 2023] and AnyGPT [Zhan et al., 2024], which unify processing across different modalities. Inspired by these advances, we introduce a tokenizer-based architecture in the molecule-text domain, converting molecular features into tokens compatible with LLMs.

Molecular Large Language Models. The recent emergence of Vision Large Language Models (VLLMs) [Li et al., 2023] has catalyzed advancements in molecular LLMs, which encompass both single modality and multi-modality approaches. In the single modality domain, researchers are exploring diverse molecule representations, such as 1D sequences like SMILES strings [Irwin et al., 2022], 2D molecule graphs [You et al., 2020], 3D geometric conformations [You et al., 2020], and textual information from the literature [Taylor et al., 2022]. In the multiple modalities domain, various innovative approaches are being employed. MolT5 [Edwards et al., 2022], a T5-based [Raffel et al., 2020] model, is designed for SMILES-to-text and text-to-SMILES translations. Other works, such as MoMu [Su et al., 2022], MoleculeSTM [Liu et al., 2023a], and GIT-Mol [Liu et al., 2024b], leverage cross-modal contrastive learning to align the representation spaces of molecules and text. Additionally, some studies [Cao et al., 2023; Liang et al., 2023; Liu et al., 2023b; Li et al., 2024] use multi-modal learning architectures to develop molecular LLMs, which often adopt adapter-based architectures. However, these methods do not treat molecule and text modalities equally and lack a supervision signal for the molecule modality, limiting model capacity and effectiveness.

Vector Quantization. Vector Quantization (VQ) [Gray, 1984] is a widely used technique in generative models. VQVAE [Van Den Oord et al., 2017] converts an image into a set of discrete codes within a learnable discrete latent space by learning to reconstruct the original image. VQ-GAN [Yu et al., 2021] enhances the generation quality by leveraging adversarial and perceptual objectives. In the context of molecules, VQ has been effectively applied to quantize molecule features. For example, DGAE [Boget et al., 2023] introduces a VQ model specifically for molecules, where molecules are encoded into discrete latent codes. Mole-BERT [Xia et al., 2022] uses VQ to rethink the pre-training of GNNs for molecular tasks. IMoLD [Zhuang et al., 2024] proposes using VQ to enhance invariant molecule representations, and VQSynergy [Wu et al., 2024] demonstrates the use of VQ for drug discovery.

3 Method

Our objective is to leverage the reasoning and generation capabilities of LLMs to enhance the comprehension and generation of molecule and text data. To achieve this, we focus on representing these modalities uniformly within the token representation, utilizing the next-token-prediction training paradigm of LLMs. As illustrated in Figure 2, we introduce a molecule tokenizer (Section 3.1) designed to transformmolecules into molecule tokens by learning to reconstruct the input molecule. The molecule sequence can then be concatenated with the text sequence to form a multi-modal sequence, which is fed into an LLM for molecule-to-text and text-to-molecule autoregressive pretraining (Section 3.2), as illustrated in Figure 3. The LLM vocabulary is expanded with molecule tokens mapped from the learned codebook. We introduce a four-stage training scheme for UniMoT (Section 3.3) comprising Causal Q-Former pretraining, molecule tokenizer pretraining, unified molecule-text pretraining, and task-specific instruction tuning. UniMoT is capable of performing molecule comprehension and generation tasks following the training scheme.

3.1 Molecule Tokenizer for LLMs

Molecule Encoder. We represent the structural information of a molecule as a graph, denoted by G = (V, E), where V is the set of atoms and |V| = N is the number of atoms. The task of the molecule encoder is to extract molecule features that are context-aware and encompass diverse local neighborhood structural information. By employing a molecule encoder, we obtain molecule features X ∈ RN×F , where F denotes the dimensionality of the feature vector for each atom.

Causal Q-Former. We employ a Q-Former model introduced by BLIP-2 [Li et al., 2023] to generate queries Z = {zi}Mi=1 ∈ RM×d containing high-level molecular and textual information, where M represents the number of queries and d denotes the dimension of queries. The Q-Former operates as a query-based transformer that utilizes learnable queries {zi}Mi=1 to interact with molecule features X extracted by the molecule encoder. Specifically, we incorporate causal masks into the queries, ensuring that they only interact with preceding queries. This ensures the sequence of queries maintains a causal dependency, aligning with the unidirectional requirements of LLMs operating on text sequence. Details regarding the Causal Q-Former can be found in Appendix ??.

Vector Quantization. The Causal Q-Former converts molecules and text into a causal sequence of queries. Subsequently, the causal sequence of queries {zi}Mi=1 is quantized into a causal sequence of molecule tokens {si}Mi=1 by identifying the closest neighbor in a learnable codebook C = {ci}Ki=1 , where K represents the size of the codebook. The codebook is randomly initialized and optimized during pretraining. Specifically, token si is determined as follows:

Figure 3: Illustration of the multi-modal autoregressive pretraining on molecule-text datasets. UniMoT excels in multi-modal comprehension and generation tasks, enabled by the unified LM objective. T represents the size of the text vocabulary.

Intuitively, the query zi is quantized to the closest neighbor csi in the codebook. As the vector quantization process is nondifferentiable, we adopt the straight-through estimator [Bengio et al., 2013] to train the Causal Q-Former by copying the gradient from the molecule tokens to the queries, as shown in Figure 2. The resulting embeddings of molecule tokens {si}Mi=1 , denoted as C = {csi }Mi=1 , are subsequently utilized for reconstructing molecules.

Molecule Reconstruction. An MLP adapter ψ needs to be trained to align the discrete latent space of molecule tokens with the continuous latent space of a molecular generative model for molecule reconstruction. This can be represented as XR = ψ(C), where XR denotes the embeddings for reconstruction. To achieve alignment, we minimize the Mean Squared Error (MSE) loss between XR and the SMILES [Weininger, 1988] embeddings XS produced by the pretrained SMILES encoder. Subsequently, we can reconstruct the molecule from XR using the pretrained SMILES decoder. The training loss of the tokenizer is expressed as follows:

Here, the first term represents the alignment loss, the second term is a codebook loss aimed at updating the codebook embeddings, and the third term is a commitment loss that encourages the query to stay close to the chosen codebook embedding. sg[·] denotes the stop-gradient operator, and the hyperparameter β is set to 0.25.

3.2 UnifiedMolecule-Text Language Model

Expanding Vocabulary. Employing the molecule tokenizer, a molecule can be tokenized into a molecule sequence {si}Mi=1 with causal dependency. The molecule sequence can be concatenated with the text sequence to form a multi-modal sequence {ui}Li=1 , where L is the length of the multi-modal sequence. To facilitate the representation of the multi-modal sequence, we construct the molecule vocabulary V m = {vmi }Ki=1 , which maintains the order of the molecule codebook C = {ci}Ki=1 . Additionally, V m includes several special tokens such as boundary indicators, e.g., [MOL] and [/MOL], to mark the beginning and end of the molecule sequence. Next, we merge the original text vocabulary V t = {vti }Ti=1 with the molecule vocabulary V m. The unified molecule-text vocabulary V = {V m, V t } facilitates joint learning from molecules and text under a unified next-token-prediction objective. As the vocabulary is expanded, the corresponding embeddings and prediction layers also need to be extended, with the newly introduced parameters initialized randomly.

Unified Molecule-text Modeling. The multi-modal sequence {ui}Li=1 is fed into the pretrained LLM for performing multi-modal autoregression. UniMoT adopts the general Language Modeling (LM) objective to directly maximize the log-likelihood of the data distribution:

where D represents the dataset, I represents the set of indices of the generation target, and θ denotes the parameters of the LLM. The unification of representation and training paradigm for molecules and text enhances the abilities of LLMs to understand molecule-text interactions and alignment. UniMoT can interpret molecules similar to understanding a foreign language, and generate them as if they were text. We conduct autoregressive pretraining on molecule-to-text and textto-molecule tasks to enhance the molecule comprehension and generation capabilities.

Molecule-to-Text Autoregression. While structural information is embedded in molecule features and captured by the molecule tokens through the tokenizer, we also aim to incorporate sequential information of molecules for better comprehension. Therefore, we concatenate the molecule sequence {si}Mi=1 with the SMILES [Weininger, 1988] sequence and a prompt to form the multi-modal input sequence {ui}Li=1 , as illustrated in Figure 3a. The corresponding molecule caption is used as the generation target.

Text-to-Molecule Autoregression. For molecule generation, a prompt and the molecule caption are concatenated, with a [MOL] token appended to signify the beginning of the molecule sequence, as illustrated in Figure 3b. The molecule sequence {si}Mi=1 produced by the tokenizer is used as the generation target. During inference, given a prompt and the molecule caption, the output molecule sequence can be decoded into the desired molecule by the pretrained adapter and SMILES decoder.

3.3 Training Strategy

The training strategy for UniMoT is structured across four stages. Stage-1 focuses on Causal Q-Former pretraining with tailored objectives. In Stage-2, the molecule tokenizer is optimized using the frozen encoders and decoder. Stage-3 integrates the tokenizer with a language model for multi-modal comprehension and generation. Finally, Stage-4 fine-tunes UniMoT for specific tasks, aligning it with human instructions and optimizing performance for various molecular applications. More details regarding the training process can be found in Appendix ??.

Stage-1: Causal Q-Former Pretraining. We connect the molecule encoder and Causal Q-Former, leveraging the pretrained MoleculeSTM molecule encoder [Liu et al., 2023a]. The molecule encoder remains frozen while only the Causal Q-Former is updated. Both queries and text inputs are used, while only queries serve as input in subsequent stages. In our experiments, we utilize 16 queries. We employ three tailored objectives for the pretraining of the Causal Q-Former: Molecule-Text Contrastive Learning (MTC), Molecule-Text Matching (MTM), and Molecule-grounded Text Generation (MTG). The details of these objectives can be found in Appendix ??.

Stage-2: Molecule Tokenizer Pretraining. We connect the Causal Q-Former with subsequent blocks and use the objective defined in Equation (2). We employ the pretrained ChemFormer [Irwin et al., 2022] as the generative model. Specifically, we leverage the SMILES encoder and the SMILES

decoder provided by ChemFormer. The molecule codebook size is set to K = 2048. As shown in Figure 2, we keep the molecule encoder, the SMILES encoder, and the SMILES decoder frozen, while updating the Causal Q-Former, the learnable codebook, and the adapter.

Stage-3: UnifiedMolecule-Text Pretraining. We integrate the molecule tokenizer with the LLM using the unified vocabulary of molecule tokens and text tokens. We employ the LM objective defined in Equation (3) to pretrain the LLM. Pretraining involves molecule-to-text autoregression and textto-molecule autoregression, aimed at enhancing UniMoT's multi-modal comprehension and generation capabilities. To enhance efficiency, we train the LLM using low-rank adaptation (LoRA) [Hu et al., 2021].

Stage-4: Task-Specific Instruction Tuning. UniMoT is fine-tuned on seven comprehension and generation tasks: molecular property prediction, molecule captioning, moleculetext retrieval, caption-guided molecule generation, reagent prediction, forward reaction prediction, and retrosynthesis. We also utilize LoRA to improve efficiency. This stage ensures UniMoT can accurately interpret and respond to human instructions, making it versatile and effective for diverse molecular tasks.

4 Experiments

4.1 Molecule Comprehension Tasks

Molecular Property Prediction Task. The goal of molecular property prediction is to forecast a molecule's intrinsic physical and chemical properties. For the classification task, we incorporate eight binary classification datasets from MoleculeNet [Wu et al., 2018]. Models are tasked with generating a single prediction ("yes" or "no"). We compare UniMoT with the following baselines: KV-PLM [Zeng et al., 2022], AttrMask [Hu et al., 2019], InfoGraph [Sun et al., 2019], MolCLR [Wang et al., 2021], GraphMVP [Liu et al., 2019], MoleculeSTM [Liu et al., 2023a], and InstructMol [Cao et al., 2023]. The ROC-AUC (%) results on the MoleculeNet datasets are shown in Table 1. The performance of the regression task of molecular property prediction is provided in Appendix ??. Compared to traditional graph learning methods and molecular LLMs like InstructMol [Cao et al., 2023], UniMoT demonstrates consistent improvements across the eight datasets, indicating its robust molecule comprehension abilities.

Molecule Captioning Task. The molecule captioning task involves generating a comprehensive description of a molecule. We compare UniMoT with several baselines: MolT5 [Edwards et al., 2022], MoMu [Su et al., 2022], InstructMol [Cao et al., 2023], MolCA [Liu et al., 2023b], and 3D-MoLM [Li et al., 2024]. BLEU [Papineni et al., 2002], ROUGE [Lin, 2004], and METEOR [Banerjee and Lavie, 2005] are adopted as evaluation metrics. UniMoT is evaluated for molecule captioning on the PubChem [Kim et al., 2023] and ChEBI-20 [Edwards et al., 2022] datasets. Performance on the PubChem dataset is shown in Table 2, while the performance on the ChEBI-20 dataset and some concrete examples are presented in Appendix ??. The ChEBI-20 dataset replaces molecular names with "the molecule" to focus on properties. However, predicting molecular names reflects the model's structural understanding, so we conducted the main experiments on PubChem.

From Table 2, we observe that UniMoT consistently outperforms the baselines by a significant margin on the PubChem [Kim et al., 2023] dataset. This task is more complex than classification or regression, providing a robust measure of the model's molecule comprehension abilities. Notably, our proposed tokenizer-based architecture surpasses the projection-based architecture (such as InstructMol [Cao et al., 2023]), Q-Former-based architecture (such as MolCA [Liu et al., 2023b] and 3D-MoLM [Li et al., 2024]), and models trained with contrastive learning strategies (such as MoMu [Su et al., 2022]). This demonstrates that the tokenizer-based architecture achieves better molecule-text alignment through autoregressive molecule-to-text and text-to-molecule pretraining compared to other architectures. Details and More Results of Experiments can be found in Appendix ??.

4.2 Molecule Generation Tasks

We employ molecule generation tasks, which encompass caption-guided molecule generation [Fang et al., 2023], reagent prediction [Fang et al., 2023], forward reaction prediction [Fang et al., 2023], and retrosynthesis [Fang et al., 2023]. Caption-guided molecule generation involves generating molecular structures based on textual descriptions. Reagent prediction entails determining suitable reagents given reactants and products. Forward reaction prediction involves predicting probable products given specific reactants and reagents. Retrosynthesis involves deconstructing a target molecule into simpler starting materials. We compare UniMoT with the following baselines: Llama [Touvron et al., 2023], Vicuna [Chiang et al., 2023], Mol-Instructions [Fang et al., 2023], and InstructMol [Cao et al., 2023]. The metrics used to evaluate molecule generation tasks include Exact Match, BLEU [Papineni et al., 2002], Levenshtein Distance [Levenshtein and others, 1966], RDKit Fingerprint Similarity [Landrum and others, 2006], MACCS Fingerprint Similarity [Durant et al., 2002], and Morgan Fingerprint Similarity [Morgan, 1965]. These metrics evaluate structural similarity between generated and target molecules, along with Validity [Kusner et al., 2017], which assesses the proportion of chemically valid molecules generated.

We utilize the Mol-Instructions [Fang et al., 2023] benchmark to evaluate the generation capabilities of UniMoT. The results of caption-guided molecule generation and reagent prediction are presented in Table 3, and the results of other tasks are in Appendix ??. The caption-guided molecule generation task, the reverse of molecule captioning, is conducted using the PubChem [Kim et al., 2023] dataset, while the other tasks utilize the USPTO [Fang et al., 2023] dataset. As the baselines generate SMILES strings and then convert them to molecules, UniMoT directly leverages the generated molecule tokens and obtains their embeddings from the learned codebook. These embeddings can be decoded to desired molecules through the pretrained adapter and SMILES decoder. As shown in Table 3, UniMoT generates valid molecules with a higher degree of similarity to the target molecules compared to the baselines. This is because UniMoT can generate molecules as if they were text, which is fundamentally different from adapterbased architectures. UniMoT demonstrates strong generation capabilities and offers a new perspective on these tasks.

4.3 Ablation Studies

Cross-Modal Projector. We conducted an ablation study on the cross-modal projector, with the results on the molecule captioning task shown in Table 4a. The linear projection demonstrated the worst performance, indicating that the molecule features lack textual information, thus hindering effective molecule-text interactions and alignment. Additionally, we compared the performance of a Q-Former with bidirectional self-attention to a Causal Q-Former with causal self-attention

(a) Ablation study on the projector and representation form.

(b) Ablation study on the codebook size.

Table 4: Ablation studies on the molecule captioning task using the PubChem dataset.

in the second and third rows. The results show that queries with causal dependency outperform those with bidirectional dependency. This demonstrates that input with left-to-right causal dependency aligns with the unidirectional attention mechanism in LLMs, leading to improved performance.

Discrete vs. Continuous Representation. We compared the performance of continuous causal embeddings and discrete tokens, quantized from causal embeddings, as inputs to LLMs in the third and fourth rows of Table 4a. Continuous embeddings demonstrate better performance than discrete tokens in understanding molecules. This result is reasonable since the quantization process causes information loss in discrete tokens. However, we still use discrete token representation to facilitate the autoregressive training paradigm of LLMs, which supports the unification of comprehension and generation tasks. To achieve this unification, we unavoidably sacrifice some performance in comprehension tasks.

Codebook Size. We conducted experiments with different molecule codebook sizes and reported the performance on the molecule captioning task. The performance is shown in Table 4b. The results demonstrate that the codebook size of 2048 consistently provides the best performance. This choice balances model complexity and performance. A larger codebook could capture more subtle interactions between molecules and text. However, there may be some codes that are not often used. A smaller codebook may result in nearby embeddings being assigned the same code, which reduces the granularity of the representation. More ablation studies are presented in Appendix ??.

5 Conclusion

This work introduces UniMoT, a framework that unifies the modalities of molecules and text. By adopting a tokenizerbased architecture, UniMoT addresses previous limitations where the molecule and text modalities are not treated equally. The molecule tokenizer converts molecules into sequences of discrete tokens, embedding high-level molecular and textual information. The LLM vocabulary is expanded with molecule tokens mapped from a learned codebook. Through a four-stage training scheme, UniMoT has become a versatile multi-modal LLM, capable of handling both molecule-to-text and text-tomolecule tasks. Extensive empirical evaluations show that UniMoT achieves state-of-the-art performance across diverse molecule comprehension and generation tasks. Although UniMoT excels in molecule-to-text and text-to-molecule tasks, it has yet to be extensively tested on more complex tasks like molecule editing, which require precise structural modifications. Additionally, limited annotated data in the molecular domain restricts UniMoT's training, hindering its ability to fully learn and generalize molecular structures and properties. To improve its effectiveness, addressing data scarcity is crucial. Furthermore, expanding evaluations to include a broader range of real-world scenarios will offer a more comprehensive understanding of the model's robustness and generalizability.

Acknowledgments

1. Yao is supported by National Natural Science Foundation of China (under Grant No.92270106), Beijing Natural Science Foundation (under Grant No.4242039), and CCF-Zhipu Large Model Innovation Fund (CCF-Zhipu202402)