Title: UniMoT: Unified Molecule-Text Language Model with Discrete Token Representation
Venue: NeurIPS 2024
Accepted: no
Keywords: Large Language Models

Tokenization

Molecule Generation

Molecule Comprehension

Multi-modal Learning
Abstract: The remarkable success of Large Language Models (LLMs) across diverse tasks has driven the research community to extend their capabilities to molecular applications, leading to the development of molecular LLMs. However, most molecular LLMs employ adapter-based architectures that do not treat molecule and text modalities equally and lack a supervision signal for the molecule modality. To address these issues, we introduce UniMoT, a unified molecule-text LLM adopting a tokenizer-based architecture that expands the vocabulary of LLM with molecule tokens.
Specifically, we introduce a Vector Quantization-driven tokenizer that incorporates a Q-Former to bridge the modality gap between molecule and text. This tokenizer transforms molecules into sequences of molecule tokens with causal dependency, encapsulating high-level molecular and textual information. 
Equipped with this tokenizer, UniMoT can unify molecule and text modalities under a shared token representation and an autoregressive training paradigm, enabling it to interpret molecules as a foreign language and generate them as text.
Following a four-stage training scheme, UniMoT emerges as a multi-modal generalist capable of performing both molecule-to-text and text-to-molecule tasks. Extensive experiments demonstrate that UniMoT achieves state-of-the-art performance across a wide range of molecule comprehension and generation tasks.

Abstract

The remarkable success of Large Language Models (LLMs) across diverse tasks has driven the research community to extend their capabilities to molecular applications, leading to the development of molecular LLMs. However, most molecular LLMs employ adapter-based architectures that do not treat molecule and text modalities equally and lack a supervision signal for the molecule modality. To address these issues, we introduce UniMoT , a unified molecule-text LLM adopting a tokenizer-based architecture that expands the vocabulary of LLM with molecule tokens. Specifically, we introduce a Vector Quantization-driven tokenizer that incorporates a Q-Former to bridge the modality gap between molecule and text. This tokenizer transforms molecules into sequences of molecule tokens with causal dependency, encapsulating high-level molecular and textual information. Equipped with this tokenizer, UniMoT can unify molecule and text modalities under a shared token representation and an autoregressive training paradigm, enabling it to interpret molecules as a foreign language and generate them as text. Following a four-stage training scheme, UniMoT emerges as a multi-modal generalist capable of performing both molecule-to-text and text-to-molecule tasks. Extensive experiments demonstrate that UniMoT achieves state-of-the-art performance across a wide range of molecule comprehension and generation tasks.

1 Introduction

The incredible capabilities of Large Language Models (LLMs) [5, 44] have led to their widespread 20 use as versatile tools for completing diverse real-world tasks. This success has sparked interest in 21 Multi-modal LLMs [59, 52], which aim to enhance LLMs by enabling them to process multi-modal 22 inputs and outputs. Prior research efforts [26, 41, 12, 6, 33, 35, 25] have focused on adapting LLMs 23 to molecular tasks, resulting in the development of molecular LLMs. These molecular LLMs can 24 analyze molecule structures [35, 33, 6], address drug-related inquiries [26, 41], assist in synthesis 25 and retrosynthesis planning [12], support drug design [12], and more. 26

27

Prevalent molecular LLMs commonly employ adapter-based architectures, adopting either a linear

28

29

30

projection [26, 41, 6] or a Q-Former [33, 25] as an adapter to translate molecule features into the semantic space of LLM, as illustrated in Figure 1a and Figure 1b.

Despite demonstrating initial capabilities in molecular comprehension and yielding promising results in molecule-to-text generation

tasks, they still lack molecule generation abilities. The critical issue within these methods is their 31

32

unequal treatment of molecules and text, resulting in a lack of supervision for the molecule modality.

This limitation significantly constrains model capacity and effectiveness. Due to limitations imposed 33

by the training paradigm, they are unable to perform text-to-molecule generation tasks. 34

Discretizing continuous molecule features into discrete molecule tokens offers a promising solution 35

for conducting both molecule-to-text and text-to-molecule generation tasks. By treating tokens from 36

S

Text Tokenizer

Molecule Features

Text

Text

/S

S

Text Tokenizer

Molecule Features

Q-Former

Text

Text

/S

Molecule

Embedding

Text

Token

Molecule

Token

S

Text Tokenizer

Molecule

Molecule Tokenizer

Text

Decode into Molecule

Text

/S

(a) Projection-Based Architecture.

(b) Q-Former-Based Architecture.

(c) Tokenizer-Based Architecture.

different modalities equally, we can predict the next molecule or text token in an autoregressive 37 manner. However, directly discretizing molecule features poses several challenges: (i) This approach 38 results in long sequences, with lengths equivalent to the number of atoms in a batch. LLMs typically 39 experience a quadratic increase in computational complexity with sequence length [46]. (ii) Molecule 40 tokens derived from molecule features lack left-to-right causal dependency, which conflicts with 41 the unidirectional attention mechanism in LLMs. (iii) Molecule features lack textual information, 42

hindering effective molecule-text interactions and alignment.

To this end, we present UniMoT , a unified molecule-text LLM that adopts a tokenizer-based architec44 ture, integrating molecule comprehension and generation, as depicted in Figure 1c. A pivotal aspect 45 of UniMoT's architecture is the molecule tokenizer for transforming molecules into molecule tokens. 46 We introduce a Vector Quantization-driven [45] tokenizer, incorporating a Q-Former [23] to bridge 47 the modality gap between molecule and text. Specifically, we incorporate causal masks for the queries, 48 enabling the Causal Q-Former to generate a causal sequence of query embeddings compatible with 49 the unidirectional attention in LLMs. The sequence of query embeddings is subsequently quantized 50 into a sequence of molecule tokens using a learnable codebook. The molecule tokens encapsulate 51 high-level molecular and textual information, which are then aligned with the latent space of a 52 generative model via an MLP adapter, enabling the generation of desired molecules. 53

54

55

56

57

58

59

60

61

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-token-prediction training paradigm of LLM. This unification of representation and training paradigm enhances LLMs' ability to understand molecule-text interactions and alignment. UniMoT interprets molecules akin to understanding a foreign language, and generates them as if they were text. Following a four-stage training scheme, UniMoT serves as a multi-modal generalist capable of performing both molecule comprehension and generation tasks.

62

63

64

65

66

67

68

69

70

71

72

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 molecule 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.

UniMoT exhibits remarkable capabilities in multi-modal comprehension and generation. Extensive experiments demonstrate that UniMoT achieves state-of-the-art performance across a wide spectrum of molecule comprehension tasks and molecule generation tasks.

73

74

75

76

77

78

79

80

81

82

83

84

85

86

87

88

89

90

91

92

93

94

95

96

97

98

99

100

101

102

103

104

105

106

107

108

109

110

111

112

113

114

115

116

117

118

119

120

121

122

123

2 Related Works

Molecular Large Language Models. The recent emergence of Vision Large Language Models (VLLMs) [24, 23, 28] 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 [47, 8, 17], 2D molecule graphs [15, 56], 3D geometric conformations [56, 32], and textual information from the literature [43, 2, 21]. In the multiple modalities domain, various innovative approaches are being employed. MolT5 [11], a T5-based [38] model, is designed for SMILES-to-text and text-to-SMILES translations. Other works, such as MoMu [39], MoleculeSTM [31], MolFM [34], and GIT-Mol [29], leverage cross-modal contrastive learning to align the representation spaces of molecules and text. Additionally, some studies use multi-modal learning architectures to develop molecular LLMs, which often adopt adapter-based architectures. For instance, InstructMol [6], GraphGPT [41], and DrugChat [26] employ a simple projection layer to map molecule features to LLM's input space. MolCA [33] and 3D-MoLM [25] utilize a Q-Former [23] to bridge the modality gap between molecules and text. 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) [13] is a widely used technique in generative models. VQ-VAE [45] converts an image into a set of discrete codes within a learnable discrete latent space by learning to reconstruct the original image. VQ-GAN [57] enhances the generation quality by leveraging adversarial and perceptual objectives. In the context of molecules, VQ has been effectively applied to quantize molecule representations. For example, DGAE [4] introduces a VQ model specifically for molecular graphs, where molecular graphs are encoded into discrete latent codes. Mole-BERT [54] uses VQ to rethink the pre-training of GNNs for molecular tasks. IMoLD [60] proposes using VQ to enhance invariant molecule representations, and VQSynergy [51] 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 transform molecules 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 subsequently fed into an LLM for autoregressive pretraining (Section 3.2), as illustrated in Figure 3. The LLM vocabulary is expanded with molecule codes 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 both molecular 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 node representations that are context-aware and encompass diverse local neighborhood structural information. By employing a molecule encoder, we obtain molecule features X ∈ R N × F , where each atom representation contains context-aware structural information.

Causal Q-Former. We employ a Q-Former model introduced by BLIP-2 [23] to generate query embeddings Z = { z i } i M =1 ∈ R M × d containing high-level molecular and textual information, where M represents the number of queries and d denotes the dimension of query embeddings. Specifically, we incorporate causal masks into the queries, ensuring that they only interact with preceding queries. This ensures the sequence of query embeddings maintains a causal dependency, aligning with the requirements of LLMs operating on text sequence. Details regarding the Causal Q-Former can be found in Appendix A.

124

125

126

127

128

129

130

131

132

133

134

135

136

137

138

139

140

141

142

143

144

Molecule

Encoder

'The molecule is an

indole phytoalexin

that …'

1

3

7

2

Learnable

Queries

Embeddings

Learnable

Molecule

Codebook

Molecule

Tokens

Aligning

(MSE Loss)

Input

Molecule

Straight-

Through

Gradients

Adapter

SMILES

Decoder

Reconstructed

Molecule

Molecule

Caption

SMILES

Encoder

C1=CC=C2C(=C1)C

(=CN2)C3=NC=CS3

Molecule

SMILES

Feed into LLM

Generated

Text

Trainable

Parameters

Frozen

Parameters

Vector Quantization. The Causal Q-Former converts molecule and text features into a causal sequence of query embeddings. Subsequently, we aim to quantize these query embeddings into molecule tokens using a variant of VQ-VAE [45]. These discrete molecule tokens can then be integrated with text tokens to form a multi-modal sequence suitable for feeding into LLMs. The causal sequence of query embeddings { z i } i M =1 is quantized into a causal sequence of molecule tokens { s i } i M =1 by identifying the closest neighbor in a learnable codebook C = { c i } K i =1 , where K represents the size of the codebook. The codebook is randomly initialized and optimized during pretraining. Specifically, token s i is determined as follows:

Intuitively, the query embedding z i is quantized to the closest neighbor c s i in the codebook. As the vector quantization process is non-differentiable, we adopt the straight-through estimator [3] to train the Causal Q-Former by copying the gradient from the molecule tokens to the query embeddings, as shown in Figure 2. The resulting embeddings of molecule tokens, denoted as C = { c s i } i M =1 , are subsequently utilized for reconstructing molecules.

Molecule Reconstruction. An 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. The embeddings of molecule tokens C can be aligned with the latent space of the generative model via an MLP adapter ψ , represented as X R = ψ ( C ) , where X R denotes the embeddings for reconstruction. Subsequently, we can reconstruct the molecule from X R using the pretrained SMILES decoder To achieve alignment, we minimize the Mean Squared Error (MSE) loss between X R and the SMILES [50] embeddings X S produced by the pretrained SMILES encoder.

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 145 updating the codebook embeddings, and the third term is a commitment loss that encourages the 146 query embedding to stay close to the chosen codebook embedding. sg[ · ] denotes the stop-gradient 147 operator, and the hyperparameter β is set to 0.25. 148

149

150

151

152

153

154

155

156

157

158

159

160

161

162

163

164

165

166

167

168

169

170

171

172

173

174

175

Unified Molecule-Text LLM (UniMoT)

MOL

T+3

/MOL

SMI

24

/SMI

Text Tokenizer

Molecule

C1=CC=C2C(

=C1)C(=CN2)

The molecule is

SMILES

3

an indole phytoalexin that …

Prompt

BOS

37

Molecule Tokenizer

16

54

41

T+7

T+3

/MOL

SMI

24

/SMI

3

BOS

EOS

16

54

41

T+7

37

Adapter

MOL

T+3

BOS

24

43

EOS

16

Unified Molecule-Text LLM (UniMoT)

27

T+7

Text Tokenizer

Generate a

molecule with

the following

properties.

Prompt

The molecule

is an indole

phytoalexin

that …

Caption

Molecule

Decoder

Generated

Molecule

During

Inference

MOL

T+3

24

43

EOS

16

27

/MOL

T+7

(a) Molecule-to-Text Autoregression.

(b) Text-to-Molecule Autoregression.

3.2 Unified Molecule-Text Language Model

Expanding Vocabulary. Employing the molecule tokenizer, a molecule can be tokenized into a molecule sequence { s i } i M =1 with causal dependency. The molecule sequence can be concatenated with the text sequence to form a multi-modal sequence { u i } L i =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 = { v i m } K i =1 , which maintains the order of the molecule codebook C = { c i } K i =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 = { v t i } T i =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 { u i } L i =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 text-to-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 { s i } i M =1 with the SMILES [50] sequence and a prompt to form the multi-modal input sequence { u i } L i =1 , as illustrated in Figure 3a. The text sequence of the corresponding molecule caption is used as the generation target.

Text-to-Molecule Autoregression. For molecule generation, a prompt and the molecule caption 176 are concatenated, with a [MOL] token appended to signify the beginning of the molecule sequence, 177 as illustrated in Figure 3b. The molecule sequence { s i } i M =1 produced by the tokenizer is used as the 178

179

180

181

182

183

184

185

186

187

188

189

190

191

192

193

194

195

196

197

198

199

200

201

202

203

204

205

206

207

208

209

210

211

212

213

214

215

216

217

218

219

220

221

222

223

224

225

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 QFormer 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 C.

Stage-1: Causal Q-Former Pretraining. We connect the molecule encoder and Causal Q-Former, leveraging the pretrained MoleculeSTM molecule encoder [31]. 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 MTC, MTM, and MTG for the pretraining of the Causal Q-Former, as detailed in Appendix A.

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 [17] as the generative model. Specifically, we leverage the SMILES encoder and SMILES decoder provided by ChemFormer. The molecule codebook size is set to K = 2048 . As shown in Figure 2, we keep the molecule encoder, SMILES encoder, and SMILES decoder frozen, while updating the Causal Q-Former, codebook, and adapter.

Stage-3: Unified Molecule-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 text-to-molecule autoregression, aimed at enhancing UniMoT's multi-modal comprehension and generation capabilities. To enhance efficiency, we train the LLM using LoRA tuning [14].

Stage-4: Task-Specific Instruction Tuning. UniMoT is fine-tuned on seven comprehension and generation tasks: molecular property prediction, molecule captioning, molecule-text retrieval, captionguided molecule generation, reagent prediction, forward reaction prediction, and retrosynthesis. We also utilize LoRA tuning to improve efficiency. This stage ensures UniMoT can accurately interpret and respond to human instructions, making it versatile and effective for 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 [53]. Models are tasked with generating a single prediction ('yes' or 'no'). We compare UniMoT with the following baselines: KVPLM [58], AttrMask [16], InfoGraph [40], MolCLR [48], GraphMVP [30], MoleculeSTM [31], and InstructMol [6]. 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 D. Compared to traditional graph learning methods and molecular LLMs like InstructMol, 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 [11], MoMu [39], InstructMol [6], MolCA [33], and 3D-MoLM [25]. BLEU [37], ROUGE [27], and METEOR [1] are adopted as evaluation metrics. UniMoT is evaluated for molecule captioning on the PubChem and

226

227

228

229

230

231

232

233

234

235

236

237

238

239

240

241

242

243

244

245

246

247

248

249

250

251

252

CheBI-20 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 D.

From Table 2, we observe that UniMoT consistently outperforms the baselines by a significant margin. 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), Q-Former-based architecture (such as MolCA and 3D-MoLM), and models trained with contrastive learning strategies (such as MoMu). The results demonstrate that the molecule tokenizer can generate molecule tokens with high-level molecular and textual information, enhancing molecule comprehension abilities.

Molecule-Text Retrieval Task. The molecule-text retrieval task involves using a molecule to retrieve text (M2T) and using text to retrieve a molecule (T2M). We compare UniMoT with several baselines: Sci-BERT [2], KV-PLM [58], MoMu [39], MoleculeSTM [31], MolCA [33], and 3DMoLM[25]. We report the performance of retrieval using a batch of 64 random samples and the entire test set, evaluated with the metrics of Accuracy and Recall@20. We use the checkpoint from Stage-1 of pretraining. UniMoT is evaluated on the datasets of PubChem, PCdes, and MoMu. Performance on the PubChem dataset is shown in Table 3, while performance on the PCdes and MoMu datasets is presented in Appendix D. UniMoT can understand complex molecule-text interactions through the introduction of the Causal Q-Former. From Table 3, UniMoT demonstrates superior performance over the baselines on molecule-text retrieval, particularly in molecule-to-text retrieval. This underscores UniMoT's capability in learning fine-grained alignment between molecules and text.

4.2 Molecule Generation Tasks

We employ molecule generation tasks, which encompass caption-guided molecule generation, reagent prediction, forward reaction prediction, and retrosynthesis. 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

253

254

255

256

257

258

259

260

261

262

263

264

265

266

baselines: LLaMA [44], Vicuna [7], Mol-Instructions [12], and InstructMol [6]. The metrics used to evaluate molecule generation tasks include Exact Match, BLEU [37], Levenshtein Distance [22], RDKit Fingerprint Similarity [20], MACCS Fingerprint Similarity [10], and Morgan Fingerprint Similarity [36]. These metrics evaluate structural similarity between generated and target molecules, along with Validity [19], which assesses the proportion of chemically valid molecules generated. We utilize the Mol-Instructions dataset to evaluate the generation capabilities of UniMoT, and the results are presented in Table 4.

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. Regarding the results in Table 4, UniMoT exhibits the capability to generate valid molecules with a higher degree of similarity to the target molecules compared to the baselines. UniMoT can generate molecules as if they were text, demonstrating strong generation capabilities and providing a new perspective to molecule generation tasks.

267

268

269

270

271

272

273

274

275

276

277

278

279

280

281

282

283

284

285

286

287

288

289

290

291

292

293

294

295

296

297

298

299

300

301

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 5. 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. The results show that query embeddings 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 compare the performance of continuous causal query embeddings and discrete tokens quantized from causal embeddings as inputs to LLMs. As shown in Table 5, 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.

Model Size and Tuning Stategy. We conducted a comparison of molecule captioning performance across various model sizes and tuning strategies, as illustrated in Table 6. Our findings indicate that scaling up the LLM to 13B or adopting a fully tuning strategy yields only marginal improvements in performance compared to using LLaMA2-7B with LoRA tuning. While larger models and fully tuning strategies might offer slight gains in performance, they come at a significant cost in terms of efficiency. Considering the trade-off between achieving high performance and maintaining efficiency, we have chosen to utilize LLaMA2-7B with LoRA tuning in our experiments. This ensures that our model remains both powerful and practical.

5 Conclusion

This work introduces UniMoT, an innovation in the field of molecular-textual understanding and generation, which has successfully unified these two distinct modalities under a single, coherent framework. By integrating a Vector Quantization-driven tokenizer with a Causal Q-Former, UniMoT overcomes previous architectural limitations where molecule and text modalities were not treated equally, lacking a dedicated supervision signal for the molecular domain. This unique tokenizer transforms molecules into sequences of discrete tokens, embedding high-level molecular and textual information cohesively. Moreover, by employing a four-stage training scheme, UniMoT has emerged as a versatile multi-modal LLM, adept at handling molecule-to-text and text-to-molecule tasks. Extensive empirical evaluations demonstrate that UniMoT attains state-of-the-art performance across a diverse array of molecule comprehension and generation tasks.

302

303

304

305

306

307

308

309

310

311

312

313

314

315

316

317

318

319

320

321

322

323

324

325

326

327

328

329

330

331

332

333

334

335

336

337

338

339

340

341

342

343

344

345

346

347

348

349

References

Satanjeev Banerjee and Alon Lavie. Meteor: An automatic metric for mt evaluation with improved correlation with human judgments. In Proceedings of the acl workshop on intrinsic and extrinsic evaluation measures for machine translation and/or summarization , pages 65-72, 2005.

Iz Beltagy, Kyle Lo, and Arman Cohan. Scibert: A pretrained language model for scientific text. arXiv preprint arXiv:1903.10676 , 2019.

Yoshua Bengio, Nicholas Léonard, and Aaron Courville. Estimating or propagating gradients through stochastic neurons for conditional computation. arXiv preprint arXiv:1308.3432 , 2013.

Yoann Boget, Magda Gregorova, and Alexandros Kalousis. Vector-quantized graph auto-encoder. arXiv preprint arXiv:2306.07735 , 2023.

Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems , 33:1877-1901, 2020.

He Cao, Zijing Liu, Xingyu Lu, Yuan Yao, and Yu Li. Instructmol: Multi-modal integration for building a versatile and reliable molecular assistant in drug discovery. arXiv preprint arXiv:2311.16208 , 2023.

Wei-Lin Chiang, Zhuohan Li, Zi Lin, Ying Sheng, Zhanghao Wu, Hao Zhang, Lianmin Zheng, Siyuan Zhuang, Yonghao Zhuang, Joseph E Gonzalez, et al. Vicuna: An open-source chatbot impressing gpt-4 with 90%* chatgpt quality. See https://vicuna. lmsys. org (accessed 14 April 2023) , 2(3):6, 2023.

Seyone Chithrananda, Gabriel Grand, and Bharath Ramsundar. Chemberta: large-scale self-supervised pretraining for molecular property prediction. arXiv preprint arXiv:2010.09885 , 2020.

Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.04805 , 2018.

Joseph L Durant, Burton A Leland, Douglas R Henry, and James G Nourse. Reoptimization of mdl keys for use in drug discovery. Journal of chemical information and computer sciences , 42(6):1273-1280, 2002.

Carl Edwards, Tuan Lai, Kevin Ros, Garrett Honke, Kyunghyun Cho, and Heng Ji. Translation between molecules and natural language. arXiv preprint arXiv:2204.11817 , 2022.

Yin Fang, Xiaozhuan Liang, Ningyu Zhang, Kangwei Liu, Rui Huang, Zhuo Chen, Xiaohui Fan, and Huajun Chen. Mol-instructions: A large-scale biomolecular instruction dataset for large language models. arXiv preprint arXiv:2306.08018 , 2023.

Robert Gray. Vector quantization. IEEE Assp Magazine , 1(2):4-29, 1984.

Edward J Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, and Weizhu Chen. Lora: Low-rank adaptation of large language models. arXiv preprint arXiv:2106.09685 , 2021.

Weihua Hu, Bowen Liu, Joseph Gomes, Marinka Zitnik, Percy Liang, Vijay Pande, and Jure Leskovec. Strategies for pre-training graph neural networks. arXiv preprint arXiv:1905.12265 , 2019.

Weihua Hu, Bowen Liu, Joseph Gomes, Marinka Zitnik, Percy Liang, Vijay Pande, and Jure Leskovec. Strategies for pre-training graph neural networks. arXiv preprint arXiv:1905.12265 , 2019.

Ross Irwin, Spyridon Dimitriadis, Jiazhen He, and Esben Jannik Bjerrum. Chemformer: a pre-trained transformer for computational chemistry. Machine Learning: Science and Technology , 3(1):015022, 2022.

Sunghwan Kim, Jie Chen, Tiejun Cheng, Asta Gindulyte, Jia He, Siqian He, Qingliang Li, Benjamin A Shoemaker, Paul A Thiessen, Bo Yu, et al. Pubchem 2023 update. Nucleic acids research , 51(D1):D1373D1380, 2023.

Matt J Kusner, Brooks Paige, and José Miguel Hernández-Lobato. Grammar variational autoencoder. In International conference on machine learning , pages 1945-1954. PMLR, 2017.

Greg Landrum et al. Rdkit: Open-source cheminformatics, 2006.

Jinhyuk Lee, Wonjin Yoon, Sungdong Kim, Donghyeon Kim, Sunkyu Kim, Chan Ho So, and Jaewoo Kang. Biobert: a pre-trained biomedical language representation model for biomedical text mining. Bioinformatics , 36(4):1234-1240, 2020.

350

351

352

353

354

355

356

357

358

359

360

361

362

363

364

365

366

367

368

369

370

371

372

373

374

375

376

377

378

379

380

381

382

383

384

385

386

387

388

389

390

391

392

393

394

395

396

397

Vladimir I Levenshtein et al. Binary codes capable of correcting deletions, insertions, and reversals. In Soviet physics doklady , volume 10, pages 707-710. Soviet Union, 1966.

Junnan Li, Dongxu Li, Silvio Savarese, and Steven Hoi. Blip-2: Bootstrapping language-image pre-training with frozen image encoders and large language models. In International conference on machine learning , pages 19730-19742. PMLR, 2023.

Junnan Li, Dongxu Li, Caiming Xiong, and Steven Hoi. Blip: Bootstrapping language-image pre-training for unified vision-language understanding and generation. In International conference on machine learning , pages 12888-12900. PMLR, 2022.

Sihang Li, Zhiyuan Liu, Yanchen Luo, Xiang Wang, Xiangnan He, Kenji Kawaguchi, Tat-Seng Chua, and Qi Tian. Towards 3d molecule-text interpretation in language models. arXiv preprint arXiv:2401.13923 , 2024.

Youwei Liang, Ruiyi Zhang, Li Zhang, and Pengtao Xie. Drugchat: towards enabling chatgpt-like capabilities on drug molecule graphs. arXiv preprint arXiv:2309.03907 , 2023.

Chin-Yew Lin. Rouge: A package for automatic evaluation of summaries. In Text summarization branches out , pages 74-81, 2004.

Haotian Liu, Chunyuan Li, Qingyang Wu, and Yong Jae Lee. Visual instruction tuning. Advances in neural information processing systems , 36, 2024.

Pengfei Liu, Yiming Ren, Jun Tao, and Zhixiang Ren. Git-mol: A multi-modal large language model for molecular science with graph, image, and text. Computers in Biology and Medicine , 171:108073, 2024.

Shengchao Liu, Mehmet F Demirel, and Yingyu Liang. N-gram graph: Simple unsupervised representation for graphs, with applications to molecules. Advances in neural information processing systems , 32, 2019.

Shengchao Liu, Weili Nie, Chengpeng Wang, Jiarui Lu, Zhuoran Qiao, Ling Liu, Jian Tang, Chaowei Xiao, and Animashree Anandkumar. Multi-modal molecule structure-text model for text-based retrieval and editing. Nature Machine Intelligence , 5(12):1447-1457, 2023.

Shengchao Liu, Hanchen Wang, Weiyang Liu, Joan Lasenby, Hongyu Guo, and Jian Tang. Pre-training molecular graph representation with 3d geometry. arXiv preprint arXiv:2110.07728 , 2021.

Zhiyuan Liu, Sihang Li, Yanchen Luo, Hao Fei, Yixin Cao, Kenji Kawaguchi, Xiang Wang, and Tat-Seng Chua. Molca: Molecular graph-language modeling with cross-modal projector and uni-modal adapter. arXiv preprint arXiv:2310.12798 , 2023.

Yizhen Luo, Kai Yang, Massimo Hong, Xingyi Liu, and Zaiqing Nie. Molfm: A multimodal molecular foundation model. arXiv preprint arXiv:2307.09484 , 2023.

Yizhen Luo, Jiahuan Zhang, Siqi Fan, Kai Yang, Yushuai Wu, Mu Qiao, and Zaiqing Nie. Biomedgpt: Open multimodal generative pre-trained transformer for biomedicine. arXiv preprint arXiv:2308.09442 , 2023.

Harry L Morgan. The generation of a unique machine description for chemical structures-a technique developed at chemical abstracts service. Journal of chemical documentation , 5(2):107-113, 1965.

Kishore Papineni, Salim Roukos, Todd Ward, and Wei-Jing Zhu. Bleu: a method for automatic evaluation of machine translation. In Proceedings of the 40th annual meeting of the Association for Computational Linguistics , pages 311-318, 2002.

Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. Journal of machine learning research , 21(140):1-67, 2020.

Bing Su, Dazhao Du, Zhao Yang, Yujie Zhou, Jiangmeng Li, Anyi Rao, Hao Sun, Zhiwu Lu, and Ji-Rong Wen. A molecular multimodal foundation model associating molecule graphs with natural language. arXiv preprint arXiv:2209.05481 , 2022.

Fan-Yun Sun, Jordan Hoffmann, Vikas Verma, and Jian Tang. Infograph: Unsupervised and semisupervised graph-level representation learning via mutual information maximization. arXiv preprint arXiv:1908.01000 , 2019.

Jiabin Tang, Yuhao Yang, Wei Wei, Lei Shi, Lixin Su, Suqi Cheng, Dawei Yin, and Chao Huang. Graphgpt: 398 Graph instruction tuning for large language models. arXiv preprint arXiv:2310.13023 , 2023. 399

Rohan Taori, Ishaan Gulrajani, Tianyi Zhang, Yann Dubois, Xuechen Li, Carlos Guestrin, Percy Liang, 400 and Tatsunori B Hashimoto. Stanford alpaca: An instruction-following llama model, 2023. 401

Ross Taylor, Marcin Kardas, Guillem Cucurull, Thomas Scialom, Anthony Hartshorn, Elvis Saravia, 402 Andrew Poulton, Viktor Kerkez, and Robert Stojnic. Galactica: A large language model for science. arXiv 403 preprint arXiv:2211.09085 , 2022. 404

405

406

407

Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models (2023). arXiv preprint arXiv:2302.13971 , 2023.

408

409

410

411

412

413

414

415

416

417

418

419

420

421

422

423

424

425

426

427

428

429

430

431

432

433

434

435

436

437

438

439

440

441

442

443

444

445

446

447

Aaron Van Den Oord, Oriol Vinyals, et al. Neural discrete representation learning. Advances in neural information processing systems , 30, 2017.

Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems , 30, 2017.

Sheng Wang, Yuzhi Guo, Yuhong Wang, Hongmao Sun, and Junzhou Huang. Smiles-bert: large scale unsupervised pre-training for molecular property prediction. In Proceedings of the 10th ACM international conference on bioinformatics, computational biology and health informatics , pages 429-436, 2019.

Y Wang, J Wang, Z Cao, and AB Farimani. Molclr: Molecular contrastive learning of representations via graph neural networks. arxiv 2021. arXiv preprint arXiv:2102.10056 , 2021.

Zichao Wang, Weili Nie, Zhuoran Qiao, Chaowei Xiao, Richard Baraniuk, and Anima Anandkumar. Retrieval-based controllable molecule generation. arXiv preprint arXiv:2208.11126 , 2022.

David Weininger. Smiles, a chemical language and information system. 1. introduction to methodology and encoding rules. Journal of chemical information and computer sciences , 28(1):31-36, 1988.

Jiawei Wu, Mingyuan Yan, and Dianbo Liu. Vqsynery: Robust drug synergy prediction with vector quantization mechanism. arXiv preprint arXiv:2403.03089 , 2024.

Shengqiong Wu, Hao Fei, Leigang Qu, Wei Ji, and Tat-Seng Chua. Next-gpt: Any-to-any multimodal llm. arXiv preprint arXiv:2309.05519 , 2023.

Zhenqin Wu, Bharath Ramsundar, Evan N Feinberg, Joseph Gomes, Caleb Geniesse, Aneesh S Pappu, Karl Leswing, and Vijay Pande. Moleculenet: a benchmark for molecular machine learning. Chemical science , 9(2):513-530, 2018.

Jun Xia, Chengshuai Zhao, Bozhen Hu, Zhangyang Gao, Cheng Tan, Yue Liu, Siyuan Li, and Stan Z Li. Mole-bert: Rethinking pre-training graph neural networks for molecules. In The Eleventh International Conference on Learning Representations , 2022.

Canwen Xu, Daya Guo, Nan Duan, and Julian McAuley. Baize: An open-source chat model with parameter-efficient tuning on self-chat data. arXiv preprint arXiv:2304.01196 , 2023.

Yuning You, Tianlong Chen, Yongduo Sui, Ting Chen, Zhangyang Wang, and Yang Shen. Graph contrastive learning with augmentations. Advances in neural information processing systems , 33:5812-5823, 2020.

Jiahui Yu, Xin Li, Jing Yu Koh, Han Zhang, Ruoming Pang, James Qin, Alexander Ku, Yuanzhong Xu, Jason Baldridge, and Yonghui Wu. Vector-quantized image modeling with improved vqgan. arXiv preprint arXiv:2110.04627 , 2021.

Zheni Zeng, Yuan Yao, Zhiyuan Liu, and Maosong Sun. A deep-learning system bridging molecule structure and biomedical text with comprehension comparable to human professionals. Nature communications , 13(1):862, 2022.

Jun Zhan, Junqi Dai, Jiasheng Ye, Yunhua Zhou, Dong Zhang, Zhigeng Liu, Xin Zhang, Ruibin Yuan, Ge Zhang, Linyang Li, et al. Anygpt: Unified multimodal llm with discrete sequence modeling. arXiv preprint arXiv:2402.12226 , 2024.

Xiang Zhuang, Qiang Zhang, Keyan Ding, Yatao Bian, Xiao Wang, Jingsong Lv, Hongyang Chen, and Huajun Chen. Learning invariant molecular representation in latent discrete space. Advances in Neural Information Processing Systems , 36, 2024.

448

449

450

451

452

453

454

455

456

457

458

459

460

461

462

463

464

465

466

467

468

469

470

471

472

473

474

475

476

A Details of Causal Q-Former

The Q-Former operates as a query-based transformer that utilizes learnable query vectors to interact with molecule features extracted by a frozen encoder. These queries are essential for extracting relevant information from the molecule features. The Q-Former comprises both a molecule transformer and a text transformer, sharing self-attention layers. The text transformer architecture is based on BERT [9], while the molecule transformer incorporates cross-attention layers between self-attention and feed-forward layers. Q-Former employs a cross-attention mechanism where the query vectors selectively attend to different aspects of the molecule features, allowing the model to capture critical details necessary for understanding and generating textual descriptions of molecular properties.

Specifically, we incorporate causal masks into the queries, ensuring that they only interact with preceding queries. This ensures the sequence of query embeddings maintains a causal dependency, aligning with the requirements of LLMs operating on text sequence. The Causal Q-Former is illustrated in Figure 4. We employ the Causal Q-Former to generate causal query embeddings Z = { z i } i M =1 ∈ R M × d containing high-level molecular and textual information, where M represents the number of queries and d denotes the dimension of query embeddings. Next, we introduce three tailored objectives MTC, MTM, and MTG for the pretraining of the Causal Q-Former.

Molecule

Encoder

Input

Molecule

'The molecule is an

indole phytoalexin

that …'

Learnable Queries

Molecule Caption

Self Attention

Cross Attention

Feed Forward

Molecule-Text

Matching

Self Attention

Feed Forward

Molecule-Grounded

Text Generation

Molecule-Text

Contrastive

Learning

Every

Block

Attention Mask

Different Masking

Strategies for

Different Tasks

Molecule-Text Contrastive Learning (MTC) aims to align molecule and text features by maximizing their mutual information. This is achieved by maximizing the molecule-text similarity of positive pairs against that of negative pairs. We utilize the last query embedding z M of the query sequence { z i } i M =1 as the query representation, since the output query sequence is causal and the last embedding contains global information from the queries. For text representation, we use the output embedding of the [CLS] token, denoted as y . The contrastive learning loss is expressed as follows:

where B denotes the batch size, and τ represents the temperature parameter. Here, z i M and y i refer to the i -th query and text representations in a batch, respectively.

Molecule-Text Matching (MTM) focuses on learning fine-grained alignment between molecule and text features. As query embeddings Z = { z i } i M =1 capture both molecular and textual information through cross-attention and self-attention layers respectively, we utilize the last query embedding z M as input to a binary classifier. This classifier predicts whether a given molecule-text pair is matched or unmatched. The corresponding loss function is formulated as follows:

477

478

479

480

481

482

483

484

485

486

487

488

489

490

491

492

493

494

495

496

497

498

499

500

501

502

503

504

505

506

507

508

509

510

511

512

513

514

515

516

517

518

where ϕ represents a binary classifier, and X i and t i denote the i -th input molecule features and input text in a batch, respectively.

Molecule-grounded Text Generation (MTG) focuses on generating textual descriptions given a molecule input. In this task, causal masks for queries are not applied since only textual output is required. However, causal masks are applied for text, allowing each text token to attend to its preceding text tokens and all queries, but not subsequent tokens. The Language Modeling (LM) loss function is applied to model the generation of text t i conditioned on the molecule input X i , formulated as:

where t i j represents the j -th token in the text sequence t i . Here, X i and t i denote the i -th input molecule features and generated text in a batch, respectively.

The total loss for training the Q-Former encompasses the three aforementioned objectives:

B Details of Datasets

This section provides detailed information about the datasets used in evaluating the performance of UniMoT across various tasks. The datasets are utilized for molecular property prediction, molecule captioning, molecule-text retrieval, and molecule generation tasks. Each dataset serves a unique purpose in assessing different capabilities of the model.

We present the details of the Molecular Property Prediction Datasets below:

BBBP [53]: The Blood-Brain Barrier Penetration dataset predicts the ability of molecules to penetrate the blood-brain barrier.

Tox21 [53]: This dataset is part of the Toxicology in the 21st Century initiative, used for toxicity prediction.

ToxCast [53]: Another toxicity prediction dataset with a broader range of biological assays.

Sider [53]: Side Effect Resource database, used for predicting drug side effects.

ClinTox [53]: Clinical Toxicity dataset for predicting clinical trial toxicity outcomes.

MUV [53]: Maximum Unbiased Validation dataset for virtual screening.

HIV [53]: Human Immunodeficiency Virus dataset for predicting anti-HIV activities.

BACE [53]: Beta-Secretase 1 dataset for predicting inhibitors of the BACE-1 enzyme, relevant for Alzheimer's research.

QM9 [12]: The quantum mechanics properties dataset, where the objective is to predict key quantum mechanics properties of a given molecule, such as HUMO, LUMO, and the HUMOLUMO gap.

We present the details of the Molecule Captioning Datasets below:

PubChem [18]: A large dataset of chemical molecules used for generating textual descriptions of molecular structures.

ChEBI-20 [11]: A subset of the Chemical Entities of Biological Interest database, provides structured and detailed descriptions of molecules, enhancing the model's ability to generate accurate captions.

We present the details of the Molecule-Text Retrieval Datasets below:

PubChem [18]: Used for both molecule-to-text (M2T) and text-to-molecule (T2M) retrieval tasks.

PCdes [58]: Another dataset for evaluating M2T and T2M retrieval accuracy.

MoMu [39]: Dataset specifically designed for molecule-text interactions and retrieval tasks.

519

520

521

522

523

524

525

526

527

528

529

530

531

532

533

534

535

536

537

538

539

540

541

542

543

544

545

546

547

548

549

550

551

552

553

554

555

556

557

558

559

560

561

562

563

564

565

566

567

568

569

We present the details of the Molecule Generation Datasets below:

Mol-Instructions [12]: This dataset includes tasks such as caption-guided molecule generation, reagent prediction, forward reaction prediction, and retrosynthesis. It is used to evaluate the model's ability to generate molecular structures based on textual descriptions and other related tasks.

We summarize the datasets used for evaluating UniMoT in Table 7. It encompasses various types of datasets, including those for classification, regression, captioning, retrieval, and generation tasks. Each dataset is described in terms of its type, tasks it supports, a brief description of its content, its URL for access, and the license under which it is distributed. The licenses vary, with some datasets being in the public domain and others under CC-BY 4.0 license.

C Details of Training

Stage-1: Causal Q-Former Pretraining. During Stage-1, we only connect the molecule encoder and the Causal Q-Former, leaving out other blocks. We leverage the pretrained molecule encoder from MoleculeSTM [31], which has undergone extensive contrastive learning with molecule-text pairs. We utilize the PubChem dataset [18] for pretraining, keeping the molecule encoder frozen while updating only the Causal Q-Former. Both queries and text serve as input to the Causal Q-Former, while only queries serve as input in subsequent stages. Inspired by BLIP-2 [23], we employ three tailored objectives - Molecule-Text Contrastive Learning (MTC), Molecule-Text Matching (MTM), and Molecule-grounded Text Generation (MTG) - for the pretraining of the Causal Q-Former, as detailed in Appendix A.

The dimension of molecule features is set to 300. We use 16 queries, each with a dimension of 768. The size of Z ( 16 × 768 ) is much smaller than the size of molecule features X (e.g., 150 × 300 ). The Q-former is pretrained for 50 epochs. We adopt the AdamW optimizer with a weight decay of 0.05, and a cosine decay learning rate scheduler, with a minimal learning rate of 1e-5. The batch size is set to 64. The computational overhead for this pretraining is 20 GPU hours on 4 NVIDIA A100 GPUs.

Stage-2: Molecule Tokenizer Pretraining. We connect the Causal Q-Former with the subsequent blocks and train the molecule tokenizer using the objective defined in Equation (2). Following the approach of RetMol [49], we utilize SMILES strings [50] to represent molecules, and employ the pretrained ChemFormer [17] as the generative model. Specifically, we leverage the SMILES encoder and SMILES decoder components provided by ChemFormer. We utilize PubChem [18] and CheBI-20 [11] datasets, keeping the molecule encoder, SMILES encoder, and SMILES decoder frozen, while updating the Causal Q-Former, codebook, and adapter. Once optimized, the molecule tokenizer remains unchanged throughout the subsequent stages.

The molecule codebook size is set to K = 2048 , and the dimension of codebook embedding is 768. The tokenizer is pretrained for 50 epochs. We adopt the AdamW optimizer with a weight decay of 0.05, and a cosine decay learning rate scheduler, with a minimal learning rate of 1e-5. The batch size is set to 64. The computational overhead for this pretraining is 40 GPU hours on 4 NVIDIA A100 GPUs.

Stage-3: Unified Molecule-Text Pretraining. We connect the molecule tokenizer with the LLM and employ the LM objective defined in Equation (3) to pretrain the LLM. We utilize LLaMA [44] as the default LLM. To construct the unified molecule-text vocabulary, we merge 2048 molecule codes with the original text vocabulary. Pretraining the LLM involves molecule-to-text autoregression and text-to-molecule autoregression, aimed at enhancing UniMoT's multi-modal comprehension and generation capabilities. We utilize datasets PubChem [18], CheBI-20 [11], PCdes [58], and MoMu [39] for this purpose. To enhance efficiency, we train the LLM using LoRA tuning [14].

The multi-modal LLM is pretrained for 10 epochs. We adopt the AdamW optimizer with a weight decay of 0.05, and a cosine decay learning rate scheduler, with a minimal learning rate of 1e-5. The batch size is set to 32. The computational overhead for this pretraining is 50 GPU hours on 4 NVIDIA A100 GPUs. To reduce CUDA memory usage, we integrate LoRA with the parameters set to r = 8 , α = 32 , and dropout = 0.1. This integration is applied to the k_proj , v_proj , q_proj , and o_proj modules.

570

571

572

573

574

575

576

577

578

579

580

581

582

583

584

585

586

587

588

589

590

591

592

Stage-4: Task-Specific Instruction Tuning. We perform instruction tuning to align UniMoT with human instructions through supervised fine-tuning on seven tasks: molecular property prediction, molecule captioning, molecule-text retrieval, caption-guided molecule generation, reagent prediction, forward reaction prediction, and retrosynthesis. For the molecular property prediction task, we utilize the quantum mechanics properties dataset [12] for regression prediction and the MoleculeNet datasets [53] for property classification. For the molecule captioning and molecule-text retrieval tasks, we employ datasets PubChem [18], CheBI-20 [11], PCdes [58], and MoMu [39]. For the remaining tasks, we utilize the Mol-Instructions dataset [12] to conduct instruction tuning. We fine-tune UniMoT for 10 epochs on each task using the same optimizer, learning rate scheduler, and LoRA configurations as in Stage-3 pretraining. Instruction samples for comprehension and generation tasks are shown in Table 8.

We have summarized the detailed training hyperparameters of UniMoT in Table 9.

D Details and More Results of Experiments

Molecular Property Prediction Task. Property prediction aims to anticipate a molecule's intrinsic physical and chemical properties based on its structural or sequential characteristics. In the regression task, we conduct experiments on the quantum mechanics properties dataset QM9 [12], where the objective is to predict key quantum mechanics properties of a given molecule, such as HUMO, LUMO, and the HUMO-LUMO gap. We compare UniMoT against several baselines, including Alpaca [42], Baize [55], LLaMA2-7B [44], Vicuna-13B [7], Mol-Instructions [12], and InstructMol [6]. Mean Absolute Error (MAE) serves as our evaluation metric. The performance of the regression task on the QM9 dataset is presented in Table 10. Compared to previous single-modal instruction-tuned LLMs and molecular LLMs, UniMoT exhibits further improvement on the regression task, showcasing its fundamental comprehension abilities in molecular contexts.

Molecule Captioning Task. The molecule captioning task involves generating a comprehensive 593 description of a molecule. For this task, we compare UniMoT with several baselines: MolT5 [11], 594

595

596

597

598

599

600

601

602

603

604

605

606

607

608

609

610

611

612

613

MoMu [39], InstructMol [6], MolCA [33], and 3D-MoLM [25]. We adopt BLEU [37], ROUGE [27], and METEOR [1] as the evaluation metrics. The performance of UniMoT in the molecule captioning task on the CheBI-20 dataset is presented in Table 11. Some concrete examples of molecule captioning task are presented in Table 12. From the results, it is evident that UniMoT consistently outperforms the baselines by a significant margin. These results underscore the effectiveness of the molecule tokenizer in providing molecule tokens with high-level molecular and textual information, thus enhancing molecule comprehension.

Molecule-Text Retrieval Task. The molecule-text retrieval task involves using a molecule to retrieve text (M2T) and using text to retrieve a molecule (T2M). We compare UniMoT with several baselines: Sci-BERT [2], KV-PLM [58], MoMu [39], MoleculeSTM [31], MolCA [33], and 3DMoLM [25]. We report the performance of retrieval using a batch of 64 random samples and the entire test set, evaluated with the metrics of Accuracy and Recall@20. We use the checkpoint from Stage-1 of pretraining. Performance on the PCdes and MoMu datasets is shown in Table 13. UniMoT demonstrates superior performance over the baselines on molecule-text retrieval, particularly in molecule-to-text retrieval. This demonstrates that UniMoT has learned fine-grained alignment between molecules and text, and it can understand molecule-text interactions through the introduction of the Causal Q-Former.

Molecule Generation Tasks. Molecule generation tasks include caption-guided molecule generation, reagent prediction, forward reaction prediction, and retrosynthesis.

614

615

616

617

618

619

620

621

622

623

624

625

626

627

628

629

630

631

632

633

634

635

636

637

638

639

640

641

642

643

644

645

646

647

Caption-guided molecule generation involves creating molecular structures from textual descriptions, leveraging NLP and cheminformatics to interpret and translate descriptions into chemical structures.

Reagent prediction focuses on identifying suitable reagents for given reactants and desired products, optimizing synthetic routes.

Forward reaction prediction forecasts probable products from specific reactants and reagents, using knowledge of chemical reactivity.

Retrosynthesis deconstructs target molecules into simpler starting materials.

In molecule generation tasks, evaluating the quality of generated molecules involves several metrics that measure different aspects of similarity and validity.

Exact Match checks if the generated molecule is identical to the target molecule, offering a stringent criterion for precise replication but potentially overlooking chemically similar variants.

The BLEU score [37], adapted from machine translation, measures the overlap of n-grams (short sequences of atoms or bonds) between generated and target molecules, thus assessing partial similarities.

Levenshtein Distance [22] evaluates the minimum number of edits needed to transform the generated molecule into the target, providing insight into structural changes required.

RDKit [20], MACCS [10], and Morgan [36] Fingerprint Similarities compare the generated and target molecules based on various molecular fingerprinting methods, which capture different aspects of molecular structure and properties.

The Validity [19] metric assesses the proportion of chemically valid molecules generated, ensuring that the output consists of plausible chemical structures.

Together, these metrics offer a comprehensive evaluation framework, balancing exact matches with structural and chemical validity.

E Limitations

While UniMoT demonstrates considerable advancements in unifying molecule and text modalities for comprehensive understanding and generation tasks, several limitations must be acknowledged. Although UniMoT exhibits strong performance in molecule-to-text and text-to-molecule tasks, it has not been extensively tested on more complex molecule generation tasks such as molecule editing, which require precise modifications to molecular structures. Future work could explore extending UniMoT's capabilities to handle such sophisticated molecular manipulations.

Due to the scarcity of annotated data in the molecular field, the training of UniMoT is less extensive compared to fields like computer vision. This limitation restricts the model's ability to fully learn and generalize from diverse molecular structures and properties. In contrast, the visual domain benefits

Molecule

-NH

HC

OH

OH

HO

Generated Molecule Caption

The molecule is an optically active form of phenylalaninate having Dconfiguration. It is a conjugate base of a D-phenylalanine. It is an enantiomer of a L-phenylalaninate.

The molecule is an ammonium ion that is the conjugate acid of 2phenylpropylamine arising from protonation of the primary amino function; major species at pH 7.3. It has a role as a human metabolite, an Escherichia coli metabolite and a mouse metabolite. It is a conjugate acid of a 2-phenylpropylamine.

The molecule is an enamide obtained by the carboxy group of trans-cinnamic acid with the secondary amino group of (2S,5R)1,2,5-trimethylpiperazine. It has a role as an Aspergillus metabolite. It is an alkaloid, a N-acylpiperazine, an enamide and a tertiary carboxamide. It derives from a trans-cinnamic acid.

The molecule is an (omega-1)hydroxy fatty acid ascaroside obtained by formal condensation of the alcoholic hydroxy group of (10R)10-hydroxylauric acid with ascarylopyranose (the alpha anomer). It is a metabolite of the nematode Caenorhabditis elegans. It has a role as a Caenorhabditis elegans metabolite. It is a monocarboxylic acid and an (omega-1)-hydroxy fatty acid ascaroside. It derives from an (11R)11-hydroxylauric acid. It is a conjugate acid of an ascr18(1-).

The molecule is a 2-oxo monocarboxylic acid that is pyruvic acid in which one of the methyl hydrogens is substituted by a 4-vinylcyclohex2-en-1-yl group. It has a role as a plant metabolite. It derives from a pyruvic acid. It is a conjugate acid of a 4-[(1E)-4-vinylcyclohex-2-en1-yl]pyruvate.

Ground Truth

The molecule is the D-enantiomer of phenylalaninate. It is a conjugate base of a D-phenylalanine. It is an enantiomer of a L-phenylalaninate.

The molecule is the cation obtained by protonation of the amino group of 2-phenylethylamine. It has a role as a human metabolite and an Escherichia coli metabolite. It is a conjugate acid of a 2-phenylethylamine.

The molecule is an enamide obtained by formal condensation of the carboxy group of transcinnamic acid with the secondary amino group of (2R,5R)-1,2,5trimethylpiperazine. It has a role as an Aspergillus metabolite. It is a Nacylpiperazine, a N-alkylpiperazine, an alkaloid, an enamide and a tertiary carboxamide. It derives from a trans-cinnamic acid.

The molecule is an (omega-1)hydroxy fatty acid ascaroside obtained by formal condensation of the alcoholic hydroxy group of (10R)10-hydroxyundecanoic acid with ascarylopyranose (the alpha anomer). It is a metabolite of the nematode Caenorhabditis elegans. It is a monocarboxylic acid and an (omega-1)-hydroxy fatty acid ascaroside. It derives from a (10R)10-hydroxyundecanoic acid. It is a conjugate acid of an ascrblue18(1-).

The molecule is a 2-oxo monocarboxylic acid that is pyruvic acid in which one of the methyl hydrogens has been replaced by a methylenecyclopropyl group. It has a role as a rat metabolite and a xenobiotic metabolite. It is a 2-oxo monocarboxylic acid, a member of cyclopropanes and an olefinic compound. It derives from a pyruvic acid.

(a) Accuracy (%) of molecule-text retrieval task on the PCdes dataset.

(b) Accuracy (%) of molecule-text retrieval task on the MoMu dataset.

from abundant labeled datasets, allowing for more comprehensive training and better performance. 648 Addressing this data scarcity in the molecular domain is crucial for improving UniMoT's training 649 effectiveness and overall capabilities. 650

651

652

653

The current empirical evaluations, though extensive, are primarily conducted on standard datasets and benchmarks; expanding the evaluation to a broader array of datasets and real-world scenarios will provide a more comprehensive understanding of the model's robustness and generalizability.

654

655

656

657

658

659

660

661

662

663

664

665

F Broader Impacts

The development of UniMoT, a unified model for molecule and text modalities, has significant potential to positively impact various fields. UniMoT can streamline the drug discovery process by enabling efficient molecule generation and optimization based on textual descriptions. In material science, it can aid in discovering new materials with desirable properties. Additionally, UniMoT can enhance research collaboration between chemists, biologists, and data scientists by integrating molecular and textual data, leading to comprehensive research insights and innovative solutions.

This paper does not pose any ethical concerns. The study does not involve human subjects and follows proper procedures for data set releases. There are no potentially harmful insights, methodologies, or applications. Additionally, there are no conflicts of interest or sponsorship concerns. Discrimination, bias, and fairness issues are not applicable. Privacy and security matters have been appropriately addressed, legal compliance has been maintained, and research integrity has been upheld.

666

667

668

669

670

671

672

673

674

675

676

677

678

679

680

681

682

683

684

685

686

687

688

689

690

691

692

693

694

695

696

697

698

699

700

701

702

703

704

705

706

707

708

709

710

711

712

713

714

715

716

NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

Justification: We conduct extensive experiments to verify our claims.

Guidelines:

The answer NA means that the abstract and introduction do not include the claims made in the paper.

The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.

The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: We discuss the limitations in Appendix E.

Guidelines:

The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.

The authors are encouraged to create a separate "Limitations" section in their paper.

The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.

The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.

The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.

If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [NA]

Justification: The paper does not include theoretical results.

Guidelines:

The answer NA means that the paper does not include theoretical results.

All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: We disclose all the information to reproduce the experimental results in Section 4, Appendix C, and Appendix D.

Guidelines:

The answer NA means that the paper does not include experiments.

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.

If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [No]

Justification: Once our paper is accepted, we will make the code openly accessible.

Guidelines:

The answer NA means that paper does not include experiments requiring code.

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).

Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: We disclose all the details of our experiments in Appendix C and Appendix D. Guidelines:

The answer NA means that the paper does not include experiments.

The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.

The full details can be provided either with the code, in appendix, or as supplemental material.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [No]

Justification: Given the considerable computational resources required for experiments with LLMs, we adhere to the common practice in the community.

Guidelines:

The answer NA means that the paper does not include experiments.

The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)

The assumptions made should be given (e.g., Normally distributed errors).

It should be clear whether the error bar is the standard deviation or the standard error of the mean.

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.

For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: The information regarding compute resources is provided in Appendix C.

Guidelines:

The answer NA means that the paper does not include experiments.

The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.

The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.

The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: We have carefully reviewed the code of ethics to ensure strict adherence to the guidelines.

Guidelines:

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).

10. Broader Impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [Yes]

Justification: We discuss the broader impacts in Appendix F.

Guidelines:

The answer NA means that there is no societal impact of the work performed.

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.

Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.

The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification: The paper poses no such risks.

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.

Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.

We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [Yes]

Justification: The licenses for existing assets are provided in Appendix B.

Guidelines:

The answer NA means that the paper does not use existing assets.

The authors should cite the original paper that produced the code package or dataset.

The authors should state which version of the asset is used and, if possible, include a URL.

The name of the license (e.g., CC-BY 4.0) should be included for each asset.

For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.

925

926

927

928

929

930

931

932

933

934

935

936

937

938

939

940

941

942

943

944

945

946

947

948

949

950

951

952

953

954

955

956

957

958

959

960

961

962

963

964

965

966

967

968

969

970

971

972

973

974

If this information is not available online, the authors are encouraged to reach out to the asset's creators.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [NA]

Justification: The paper does not release new assets.

Guidelines:

The answer NA means that the paper does not release new assets.

Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent was obtained from people whose asset is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.

14. Crowdsourcing and Research with Human Subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Justification: The paper does not involve crowdsourcing nor research with human subjects.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

Justification: The paper does not involve crowdsourcing nor research with human subjects. Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors propsed Uni-MoT, a unified structure to align molecules with texts with a VQ tokenizer and the Q-Former in BLIP-2. By treating molecules as new word tokens in the codebook, Uni-MoT aligns the discrete token representation for molecules and texts, while also following the autoregressive manner of LLMs. The training of Uni-MoT follows four main stages, Causal Q-Former Pretraining, Molecule Tokenizer Pretraining, Unified Molecule-Text Pretraining, Task-Specific Instruction Tuning. The experiments demonstrate that Uni-MoT can achieve better performance compared to the selected baselines.\n\nSTRENGTHS:\n1. The performance of Uni-MoT is overall good and better than the baseline models.\n2. Uni-MoT provides a alternative solution for aliging molecules with texts.\n\nWEAKNESSES:\n1. Although the authors claim that Uni-MoT follows a different structure as shown in Figure 1 c, it turns out that it still follows the BLIP-2 [1] structure, which has been widely used to align 2D molecule graph [2] and 3D molecule information [3]. Thus, **the technical contribution and novelty of this paper are extremely limited**. Especially when VQ-VAE [4] is also a well-developed structure adopted in computer vision. This paper is more like simply swapping the input of the Q-former in BLIP-2. \n2. The experiments are only conducted on a single serie of LLM, Llama-2, **which is not sufficient to demonstrate the model agnosticsm of Uni-MoT**. In fact, LLMs like Mistral [5] and Meditron [6] might possibly achieve a better performance. Meanwhile, the selection of Llama-2 is also not convincing, as Llama-2 is not specially pre-trained on chemistry or biomedical corpora. \n3. **The seletion of the datasets is also worth discussion**. The result on ChEBI-20 is presented in Appendix, while the main experiments are conducted on PubChem. I am wondering why not also conduct the remaining experiments on ChEBI-20, as the data scale of ChEBI-20 is much larger than PubChem. At the same time, the reverse task, text-based molecule generation, on ChEBI-20 and PubChem is surprisingly not presented.\n4. **The comparison with the baselines is not fair enough**. For example, MolCA adopts Galactica-1.3B as its backbone model, which is much smaller than Llama-2-7B. Thus, the proposed method can not demonstrate its superiority compared to previous methods. Notably, since the authors mentioned MolCA and 3D-MoLM, it is necessary to discuss the possible affects of the modalities. Furthermore, several SOTA baseline models like BioT5 [7] and BioT5+ [8] are not discussed.\n5. **The ablation study is also not sufficient without providing the naive fine-tuning performance of Llama-2**. Besides, as Uni-MoT incorporates the VQ tokenizer, it is also important to discuss the size of the codebook.\n6. During the pre-training stages, it should ensure that the **data leakage** is avoided. Considering the pre-training dataset adopted has overlaps with the fine-tuning dataset [7, 8], the performance gain could possibly come from the data leakage.\n7. **Some claims and definitions are confusing**. e.g. In Line 44, \"a unified molecule-text LLM\", I do not agree with the claim of an \"LLM\". The \"LLM\" is still the Llama-2. It should be like \"structure\" or something.\n\n#### References\n[1 ]Li, J., Li, D., Savarese, S., & Hoi, S. (2023, July). Blip-2: Bootstrapping language-image pre-training with frozen image encoders and large language models. In International conference on machine learning (pp. 19730-19742). PMLR.\n[2] Liu, Z., Li, S., Luo, Y., Fei, H., Cao, Y., Kawaguchi, K., ... & Chua, T. S. (2023). Molca: Molecular graph-language modeling with cross-modal projector and uni-modal adapter. arXiv preprint arXiv:2310.12798.\n[3] Li, S., Liu, Z., Luo, Y., Wang, X., He, X., Kawaguchi, K., ... & Tian, Q. (2024). Towards 3D Molecule-Text Interpretation in Language Models. arXiv preprint arXiv:2401.13923.\n[4] Yu, J., Li, X., Koh, J. Y., Zhang, H., Pang, R., Qin, J., ... & Wu, Y. (2021). Vector-quantized image modeling with improved vqgan. arXiv preprint arXiv:2110.04627.\n[5] Jiang, A. Q., Sablayrolles, A., Mensch, A., Bamford, C., Chaplot, D. S., Casas, D. D. L., ... & Sayed, W. E. (2023). Mistral 7B. arXiv preprint arXiv:2310.06825.\n[6] Chen, Z., Cano, A. H., Romanou, A., Bonnet, A., Matoba, K., Salvi, F., ... & Bosselut, A. (2023). Meditron-70b: Scaling medical pretraining for large language models. arXiv preprint arXiv:2311.16079.\n[7] Pei, Q., Zhang, W., Zhu, J., Wu, K., Gao, K., Wu, L., ... & Yan, R. (2023). Biot5: Enriching cross-modal integration in biology with chemical knowledge and natural language associations. arXiv preprint arXiv:2310.07276.\n[8] Pei, Q., Wu, L., Gao, K., Liang, X., Fang, Y., Zhu, J., ... & Yan, R. (2024). BioT5+: Towards Generalized Biological Understanding with IUPAC Integration and Multi-task Tuning. arXiv preprint arXiv:2402.17810.\n\nQUESTIONS:\n1. Could the authors conduct the experiments on more backbone LLMs? For example, on Galactica-1.3B, is the performance better than MolCA? Or on Llama-2-7B, is the performance still better than MolCA(Llama-2-7B)?\n2. Could the authors justify the selection of the datasets?\n3. Could the authors compare the proposed method with stronger baselines? Let's say, when the authors claim the SOTA performance, it should be noted BioT5 and BioT5+ are not included in discussion.\n4. Could the authors provide more solid ablation study?\n\nMy main questions are listed above, but I do expect the authors could discuss more about their weaknesses.\n\nLIMITATIONS:\nYes. They have discussed the limitations as not enough tasks, data scarcity, and real scenarios.", "TITLE": "Official Review", "RECOMMENDATION": 2, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-06-19"}

comments: SUMMARY:
This work presents a new molecule LLM that uses a pretrained tokenizer to replace the projection layer. The tokenizer consists a Q-Former and a VQ module, which are trained with consistency loss. The model is evaluated on molecular understanding and generation tasks.

STRENGTHS:
- The authors provided a novel framework to align text representation in LLMs and the molecules. 
- The authors conducted experiments on many datasets.

WEAKNESSES:
- The claim that adapter-based methods cannot do text-to-molecule generation tasks in not accurate (Sec.1). They can always adapt text-encoder to a pre-trained SMILES or graph generator. This may make this work not well motivated.
- In the tokenizer (Fig.2), it seems that multiple alignment methods are required to train the model, including: (1)the molecule and text contrastive in Q-Former, (2), the aligning MSE loss, (3) the SMILES decoder reconstruction. It's not clear about the design reasons, and if all of them are required.
- Given the learnable query has a fixed size, it may not perform well for larger molecules.

QUESTIONS:
- In Sec 3.3, what are the data that used in different stages?
- In Fig. 1, the model can output both molecule and text at the same time. But in Fig. 3, seems only text or molecule get loss for back propagation. Is this consistent and can authors share more training details?
- In all experiments, does this model pretrained with same data as all baselines?

LIMITATIONS:
Yes
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 2
DATE: 2024-07-09

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors propose to use a vector-quantized tokenizer that incorporates a Q-Former to connect pre-trained molecule encoder, SMILES encoder, and SMILES decoder so that a language model can integrate molecule and text modalities. Based on the proposed tokenizer, the authors introduce a four-stage training strategy to train UniMoT, a unified molecule-text LLM proposed in this submission. The performance of the UniMoT is evaluated empirically with 7 tasks in the areas of molecule comprehension and molecule generation. Some ablation studies are also conducted.\n\nSTRENGTHS:\n- This paper is well-organized and well-written.\n- Using a vector-quantized tokenizer and a Q-Former to connect different modalities could be somewhat novel.\n- The proposed UniMoT outperforms baselines in most cases and reach comparable performances in others.\n- The authors provide many implement details which increases the reproducibility.\n\nWEAKNESSES:\n- All the components are borrowed from existing works. Besides the pre-trained molecule encoder, SMILES encoder and decoder, both vector quantization and Q-Former are proposed in previous works [1][2]. The Q-Former part of this paper (Appendix A) is very similar to Q-Former's original paper. Even Figure 4 in this paper is very similar to Figure 2 in Q-Former's original paper [2].\n- When generating molecules the proposed method relies heavily on a pre-trained decoder. In the decoder's original paper, the reported validity is 99.9 and no guarantee is provided that the generated SMILES string will be always valid [3]. The 100% validity reported in this paper could be attributed to overfitting.\n- Some hyperparameter choices are not well justified and studied. For instance, the molecule codebook size is set to 2048, but there is no explanation why 2048 is chosen. How the molecule codebook size affects the performance is also not studied.\n- There are some existing works about molecule tokenizers [4][5], the paper lacks the comparison of the performance using different tokenizers.\n- The robustness of the model is not studied. Each molecule has many synonyms, how the proposed method performs given different synonyms of the same molecule is desired to know.\n\n[1] Van Den Oord, Aaron, and Oriol Vinyals. \"Neural discrete representation learning.\" Advances in neural information processing systems 30 (2017).\n\n[2] Li, Junnan, et al. \"Blip-2: Bootstrapping language-image pre-training with frozen image encoders and large language models.\" International conference on machine learning. PMLR, 2023.\n\n[3] Irwin, R., Dimitriadis, S., He, J., Bjerrum, E.J., 2021. Chemformer: A Pre-Trained Transformer for Computational Chemistry. Mach. Learn. Sci. Technol.\n\n[4] Li, Xinhao, and Denis Fourches. \"SMILES pair encoding: a data-driven substructure tokenization algorithm for deep learning.\" Journal of chemical information and modeling 61.4 (2021): 1560-1569.\n\n[5] Schwaller, Philippe, et al. \"Molecular transformer: a model for uncertainty-calibrated chemical reaction prediction.\" ACS central science 5.9 (2019): 1572-1583.\n\nQUESTIONS:\n- How do you guarantee that the output SMILES string will be 100% valid?\n- Why the molecule codebook size is chosen as 2048? How will the performance change when the codebook size is changed?\n- If you use other molecule tokenizers, how will the performance change?\n- When pre-training with PubChem dataset, what is the input text for each molecule since there are many entries and synonyms for each molecule?\n- How do you solve the problem that a molecule has many synonyms? Will your method output different answers if different synonyms of the same molecule are used in the prompt?\n\nLIMITATIONS:\nThe authors have adequately addressed the limitations.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-07-11"}

comments: SUMMARY:
introduce a molecule tokenizer based on Codebooks which gets integrated into UniMoT, a unified molecule-text LLM

STRENGTHS:
strong performance on a variety of benchmarks; outperforms/on pair with 3D-MoLM

WEAKNESSES:
- Table 1: include CLAMP [1] as well as standard molecular fingerprints (incl. in said reference); further KV-PLM linear probing for ToxCast results in 66.3 AUROC compared to 55.03 in your paper

[1] Seidl, P., Vall, A., Hochreiter, S., & Klambauer, G. (2023, July). Enhancing activity prediction models in drug discovery with the ability to understand human language. In International Conference on Machine Learning (pp. 30458-30490). PMLR.

QUESTIONS:
what is the reconstruction rate of the SMILES decoder?

LIMITATIONS:
-
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-12