Training Text-to-Molecule Models with Context-Aware Tokenization (2025)

Abstract: Recently, text-to-molecule models have shown great potential across various chemical applications, e.g., drug-discovery. These models adapt language models to molecular data by representing molecules as sequences of atoms. However, they rely on atom-level tokenizations, which primarily focus on modeling local connectivity, thereby limiting the ability of models to capture the global structural context within molecules. To tackle this issue, we propose a novel text-to-molecule model, coined Context-Aware Molecular T5 (CAMT5). Inspired by the significance of the substructure-level contexts in understanding molecule structures, e.g., ring systems, we introduce substructure-level tokenization for text-to-molecule models. Building on our tokenization scheme, we develop an importance-based training strategy that prioritizes key substructures, enabling CAMT5 to better capture the molecular semantics. Extensive experiments verify the superiority of CAMT5 in various text-to-molecule generation tasks. Intriguingly, we find that CAMT5 outperforms the state-of-the-art methods using only 2% of training tokens. In addition, we propose a simple yet effective ensemble strategy that aggregates the outputs of text-to-molecule models to further boost the generation performance.

## introduction discovering molecules that match desired language descriptions has been a long-standing goal in chemistry since it is an essential ingredient for practical deployments like drug-discovery and material design @xcite @xcite . however, achieving such text-to-molecule generation poses a challenge because of the different structural modalities of language and molecules. to address this challenge, researchers have explored the fine-tuning of pre-trained language models with additional molecular data @xcite , which is inspired by the recent success of language models in leveraging various domain knowledge, including chemical concepts @xcite . specifically, they treat each molecule as a sequence of tokens using string-based molecular representations such as smiles @xcite and selfies @xcite . intriguingly, they show that these molecule-aware language models, i.e., text-to-molecule models, can be obtained by learning the text-conditional molecule distribution based on treating atoms as tokens of language models @xcite . however, it is yet underexplored which tokenization strategy for molecules is more effective for text-to-molecule models. current state-of-the-art approaches @xcite adopt atom-level tokenization, where each atom is represented as a single token within the model's token space @xcite @xcite . even though they show remarkable performance as pioneering efforts, such atom-level tokenizations limit the models' ability to capture crucial global contextual patterns within molecules, focusing only on local connectivities @xcite @xcite . this leads to the question of how to tokenize molecules in a context-preserving manner to train text-to-molecule models more effectively. contribution. in this paper, we introduce a novel text-to-molecule model, coined context-aware molecular t5 (camt5), by proposing a contextenriched motif-level token space. specifically, we draw inspiration from the following chemical prior-the structural context of molecules is more effectively captured through their substructurelevel, i.e., motif-level, characteristics rather than the atom-level attributes @xcite @xcite the molecule is a sulfonamide that is benzene-sulfonamide substituted by a trifluoromethyl… text-to-molecule models translating text description… tokenize atom-wise (baseline) [c] [c] [c] [c] [c] … motif-wise (ours) (a) context-aware tokenization text-to-molecule generation (b) importance-based pre-training input camt5 (ours) importance target [m1] in particular, we carefully design the motif-level tokenization for camt5 to additionally alleviate two drawbacks in the tokenization strategies of previous text-to-molecule models. first, camt5 always generates a valid molecule, while molt5 (edwards et al., 2022) often generates an invalid token sequences that do not correspond to any molecule. second, each of our motif-level tokens has a unique interpretation, while some of the tokens in biot5 @xcite have multiple interpretations, e.g., both an atom and the number of atoms in a ring preceding the token @xcite , introducing semantic-level ambiguities for the models. consequently, to leverage our motif-level tokens effectively, we propose an importance-based training approach that prioritizes key motifs. specifically, each token is assigned to an importance value derived from its constituent atoms, and the training loss is adjusted by weighting it according to this pre-defined importance. this loss design is made possible by our carefully designed motif-level tokens, each representing a unified chemical context @xcite , unlike atom-wise tokenizations @xcite in previous text-to-molecule models. we verify our method's effectiveness on popular benchmarks, e.g., chebi-20 @xcite . on chebi-20, the state-of-the-art results are achieved using only 2% of the training tokens required by the previous best-performing baseline, biot5 (pei et al., 2023). foot_0 specifically, camt5 improves the ratio of molecules that exactly match the description (exact; higher is better) by 0.413 → 0.430, and those similar to the description (rdk; higher is better) by 0.801 → 0.840. we also show that a simple ensemble strategy utilizing camt5 further improves the overall performance, e.g., exact by 0.430 → 0.472. finally, we verify camt5's effectiveness in molecule modification. ## related works text-to-molecule models. inspired by the recent advancements in language models @xcite @xcite , significant efforts have been made to adapt these models for molecular applications, leading to the development of molecule-aware language models, i.e., text-to-molecule models @xcite @xcite @xcite . these approaches typically involve fine-tuning pre-trained language models, such as t5 @xcite , using molecular data by representing molecules as sequences of atom-level tokens. for instance, existing models @xcite @xcite leverage molecular representations like smiles @xcite or selfies @xcite , which encode molecules as atom-level token sequences for textto-molecule frameworks. as a result, they primarily focus on capturing local atom-wise connectivity, while overlooking the crucial global structural con- method token validity non-degeneracy molt5 atom ✗ ✓ biot5 atom ✓ ✗ camt5 (ours) motif ✓ ✓ in addition, text-to-molecule models based on atom-level tokenizations come with additional drawbacks. first, smiles-based models often generate invalid token sequences that violate the grammar @xcite , which do not correspond to valid molecules. second, selfies-based models introduce semantic-level ambiguities, i.e., degeneracy, in token interpretations @xcite , leading to sub-optimal performance in modeling the token distribution. for example, the '[o]' token can be interpreted completely differently: an oxygen atom or an indicator of a ring system comprising six atoms preceding this token. to address the aforementioned limitations, we carefully design our context-enriched motif-level tokens, ensuring validity of the generated token sequences and nondegeneracy in token interpretations (see ## method in section 3.1, we explain an overview of our problem of interest. in section 3.2, we provide the description of our context-aware text-to-molecule model, camt5. in section 3.3, we describe our confidence-based ensemble strategy. ## experiments we verify the effectiveness of camt5 through extensive experiments. in section 4.1, we explain our experimental setups. section 4.2 presents the text-to-molecule generation results on the chebi-20 and pcdes benchmarks. in section 3.3, we present the results of our confidence-based ensemble strategy. in section 4.4, we show the textconditional molecule modification task results. in section 4.5, we provide ablation studies on components of camt5. we provide additional experimental results and analyses in appendix d. ## conclusion we propose camt5, a new text-to-molecule model with chemistry-specialized tokenization. specifically, we adapt pre-trained language models by utilizing motif-level tokenization and importance-based training strategy to better understand the chemical structural context of molecules. in addition, we propose a confidence-based ensemble technique to further improve the quality of the generated molecules from camt5, using other text-to-molecules. we hope that our work further accelerates future research on domain-specific adaptations of pre-trained language models.