Diversity Beats Size Scaling for Chemical Language Models

Abstract

Chemical language models, such as transformers trained on SMILES strings, are increasingly used in drug design and have seen rapid growth in both model capacity and training dataset size. The impact of this scaling on practical downstream performance remains unclear, however. We systematically evaluate how model size and dataset size affect encoder-decoder transformers trained on paired textual molecular representations. We find that, beyond a minimal threshold, further model scaling yields no gain in hit generation rate, while dataset scaling gives diminishing returns. We further introduce a dataset diversification strategy that substantially increases hit diversity. These results suggest that, for molecular hit discovery, data curation and diversity may be more impactful than continued scaling of model size or dataset volume, and they motivate a shift from scale-first to diversity-first training paradigms.

Introduction

Drug discovery is a multidisciplinary process focused on identifying therapeutics that modulate disease-implicated biological targets¹. Rational drug design leverages detailed knowledge of a target’s structure and function to design novel molecules with optimized pharmacological properties and biological activity, thus providing a systematic framework to accelerate discovery and innovation². The development of artificial intelligence (AI) tools has greatly aided in this endeavor through data-driven approaches, with applications in sub-domains like molecular property prediction^3,4, (retro)synthetic reaction planning^5,6, and molecular design^7,8. Interest in these techniques goes beyond academia, with frameworks such as REINVENT being utilized in the pharmaceutical industry⁹. Despite increasing adaptation, AI-based methods still face several challenges, including model accuracy, generalizability, interpretability, and reproducibility, which limit their trustworthiness¹⁰.

Some of the most successful examples of deep learning models for drug design (e.g., the REINVENT family of models⁹ and GenMol¹¹) belong to the class of chemical language models (CLMs) and make use of techniques originally developed for natural language processing (NLP)¹²: molecules are represented as text, such as Simplified Molecular Input Line Entry System (SMILES) strings, and subsequently used for large-scale pretraining of e.g., transformer models or recurrent neural networks. Through fine-tuning, these models are then adapted to the target applications, like property prediction or hit generation. For AI-guided drug design, previous work has primarily focused on generating individual high-scoring molecules (e.g., high binding affinity)^13,14. This often leads to so-called policy collapse: the model finds a solution that satisfies the target constraints and continues to exploit it, producing a small set of molecules repeatedly¹⁵. This can be undesirable in real drug design scenarios such as hit-finding, where exploration of a broad chemical space is advantageous; a diverse set of candidate hits increases the probability of discovering novel compounds and avoids over-representation of highly related structures^16,17. The ability to sample diverse molecules is thus key to successfully applying generative CLMs in practice.

Despite all the work put into developing CLMs, how to best design and train them remains unclear. In NLP, the development of scaling laws for performance as a function of model and training dataset size has provided workable guidelines for training recipes^18,19. Other work has highlighted the importance of dataset quality. In NLP, scaling behavior can be improved through dataset pruning²⁰, and dataset diversification may improve downstream task performance²¹. In chemistry, training a CLM on synthesizable molecules improves the synthesizability of generated samples²². Scaling laws have also been developed for generative CLMs for de novo design²³, with performance generally increasing with both model size and training data volume. However, these scaling laws rely on performance proxies such as pretraining loss, whose reduction does not necessarily indicate improved downstream performance in NLP^24,25 or chemistry^3,26. It is therefore important that models are evaluated directly on downstream tasks to study the effects of model and dataset design.

To this end, we systematically analyze which factors influence the downstream performance of the Mol2Mol model integrated into REINVENT4^9,27, an encoder-decoder transformer²⁸ trained on molecular pairs that generates SMILES sequences similar to an input SMILES, which may be used for molecular optimization in conjunction with reinforcement learning (RL). Specifically, we investigate the effects of model size, training dataset size, and training dataset distribution by pretraining Mol2Mol variants on datasets stemming from different sub-sampling strategies. Each model is then evaluated both through direct sampling, on measures like the validity of generated compounds, as well as goal-oriented molecular design, where the task is to generate compounds with particular properties.

By exploring these factors, we seek to uncover the drivers of model performance in CLMs for molecular optimization in drug design. Our findings are intended to provide practical guidance for optimizing generative model frameworks in future drug discovery efforts.

Results

Model scaling

All models in Table 3 were pretrained on the 7.8M R dataset, consisting of around 7.8 million molecular pairs. We then performed direct sampling using the test set of 1000 molecules, generating 1000 SMILES for each input. The results are shown in Table 1.

Table 3.

Summary of the varied hyperparameter settings for the models under study

Model	Parameters (millions)	Number of layers	Model dimension	Feed-forward dimension
L1–D1–F1*	17.4	6	256	2048
L1/2–D1/2–F1/2	2.2	3	128	1024
L2–D2–F2	138	12	512	4096
L1–D1–F2	30	6	256	4096
L1–D1–F4	55.3	6	256	8192
L1–D2–F1	44.3	6	512	2048
L1–D4–F1	126	6	1024	2048
L2–D1–F1	34.8	12	256	2048
L4–D1–F1	69.5	24	256	2048

The baseline model is marked with an asterisk. Parameter count refers to trainable parameters. Hyperparameters that differ from those of the baseline model are bolded for clarity.

Table 1.

Direct sampling from models pretrained on the 7.8M R dataset

Model	Validity	Uniqueness	NS uniq.	τ10	r
L1–D1–F1*	0.910	0.950	0.667	0.362 ± 0.24	0.474 ± 0.19
L1/2–D1/2–F1/2	0.710	0.919	0.648	0.380 ± 0.25	0.457 ± 0.17
L2–D2–F2	0.954	0.965	0.702	0.395 ± 0.25	0.482 ± 0.19
L1–D1–F2	0.886	0.946	0.687	0.344 ± 0.24	0.456 ± 0.19
L1–D1–F4	0.892	0.951	0.680	0.316 ± 0.25	0.517 ± 0.19
L1–D2–F1	0.903	0.957	0.691	0.360 ± 0.24	0.492 ± 0.18
L1–D4–F1	0.885	0.976	0.755	0.283 ± 0.24	0.409 ± 0.17
L2–D1–F1	0.945	0.954	0.700	0.354 ± 0.25	0.478 ± 0.18
L4–D1–F1	0.952	0.964	0.702	0.374 ± 0.24	0.472 ± 0.19

The baseline model is marked with an asterisk. Uniqueness is calculated after SMILES canonicalization. Uniqueness after removal of stereochemical information is denoted by NS. τ₁₀ represents the Kendall’s tau score between the Tanimoto similarity between the input molecule and the top 10 generated molecules. r is the Pearson correlation coefficient between the negative log-likelihood (NLL) of the generated molecule and its score. For τ₁₀ and r, we additionally report standard deviations.

Clear trends emerge when all considered hyperparameters are scaled simultaneously. The baseline model, L1–D1–F1, produces valid and unique molecules in the majority of cases: 91% of the generated SMILES were valid, of which 95% were unique. Halving all hyperparameters reduces model performance significantly, with only 71% of the generated molecules being valid. This implies that the L^1/2–D^1/2–F^1/2model is too small to properly learn the SMILES syntax. Doubling all hyperparameters slightly improved performance, with 95% of the model outputs being valid SMILES, of which 97% were unique. The r and τ₁₀ scores do not appear to correlate well with model size; neither L^1/2–D^1/2–F^1/2 nor L2–D2–F2 yielded values that differed substantially from those of the baseline model.

No clear trends emerge when scaling the hyperparameters individually. When scaling up just the feed-forward dimension, model performance is reduced slightly, with both the L1–D1–F2 and L1–D1–F4 models exhibiting lower validity than the baseline model. Samples from the L1–D1–F4 model also exhibited a slightly higher r than those from the baseline model, but a lower τ₁₀. In other words, the model is on average more likely to produce samples similar to the input molecule than the baseline, but this correlation is weaker than the baseline for the 10 most likely molecules. Scaling up only the model dimension similarly leads to slight performance degradation, with fewer valid samples being produced from both the L1–D2–F1 and L1–D4–F1 models than the baseline. The uniqueness of valid samples, both with and without stereochemical information, was slightly higher, however. Additionally, samples from the L1–D4–F1 exhibited lower τ₁₀ and r than the baseline. Finally, scaling the number of layers had an overall positive impact, with both the L2–D1–F1 and L4–D1–F1 models producing more valid and unique samples than the baseline.

The Tanimoto similarities (using 1024-bit Morgan fingerprints of radius 2) between the input molecules and samples drawn from each model are shown in Supplementary Fig. 1. Most of the models produced very similar distributions, with a main peak around 0.7 and a smaller peak at 1, indicating that they encode similar molecular distributions. L^1/2–D^1/2–F^1/2 and L1–D4–F1 produced markedly different results, however, indicating that these models fail to learn the same distribution as the other models.

We additionally performed goal-directed drug design to evaluate the hit generation performance of the models. To this end, we utilized RL and rewarded the generation of molecules with a high predicted binding affinity to a particular target (dopamine receptor type 2 (DRD2), glycogen synthase kinase 3 beta (GSK3β), or c-Jun N-terminal kinase 3 (JNK3)) as well as a high quantitative estimation of drug-likeness (QED). We note that the JNK3 oracle exhibited a far lower hit density than the others on our training datasets, as shown in Supplementary Fig. 4. This likely makes optimizing for JNK3 binding more challenging. Results are shown in Fig. 1, presented as accumulated hits as a function of model parameter count for each input compound.

Fig. 1. Accumulated hits for DRD2, GSK3β, and JNK3 binding affinity after 6500 epochs of RL.

The bar plots are grouped and colored according to which hyperparameters are scaled, with the baseline being dark blue. Each row of plots corresponds to a target, and each column of plots corresponds to a target-specific input compound as indicated by the subscript indices. The compounds are listed in Supplementary Table 1. Colors indicate which hyperparameter was scaled: balanced scaling is shown in blue (B), feed-forward dimension in green (F), model dimension in purple (D), and the number of layers in red (L). The baseline model L1–D1–F1 is shown as a dashed horizontal line. Error bars represent 95% confidence intervals calculated using Eq. (1).

Most models generated a comparable number of hits for DRD2 affinity, with two notable exceptions. The L^1/2–D^1/2–F^1/2 model generated fewer hits than the baseline with DRD2₁ and DRD2₃ as inputs, while the L1–D4–F1 model generated fewer hits than the other models across the board.

We observed similar behavior for GSK3β, with L^1/2–D^1/2–F^1/2 performing worse than the baseline for two input molecules. The performance gap of L1–D4–F1 to the other models was even more pronounced here, resulting in less than half of the baseline’s accumulated hit count. For GSK3β₄, L1–D1–F4 also produced roughly one-fourth of the average hit count of the baseline model; we note that the variance in accumulated hit count of the other models was high here, however, indicating inconsistent performance across runs.

For JNK3, the average number of hits was significantly lower than for the other targets, implying that generating JNK3 binders is more difficult, as reflected by the lower hit density of that oracle on our training data. Here, L1–D1–F4 failed to generate a significant number of hits for any input molecule, while L^1/2–D^1/2–F^1/2 performed significantly worse than the baseline for JNK3₂. The other models performed similarly to the baseline.

To assess whether our model scaling results depend on the choice of pretraining data, we conducted another hit generation experiment using DRD2₁ as the input compound for L^1/2–D^1/2–F^1/2, L1–D1–F1, and L2–D2–F2, pretrained on 161K R as well as 558M R. As shown in Supplementary Fig. 5, trends were consistent with Fig. 1, indicating that our conclusions about model scaling hold across pretraining datasets. Additionally, Supplementary Fig. 6 reports the validation losses of these models on the three randomly subsampled datasets. Comparing models, L^1/2–D^1/2–F^1/2 exhibited higher losses than the other models for each dataset, but L2–D2–F2 had slightly higher losses than L1–D1–F1 for the two smaller datasets. This contrasts with Frey et al., who reported substantial loss reductions with increasing model size even for small datasets²³, and is likely related to the nature of the source-target pretraining of Mol2Mol. In our setting, validation loss is a poor proxy for downstream performance: with 7.8M R pretraining, L1–D1–F1 and L2–D2–F2 perform similarly in both direct sampling and hit generation, and the same holds for 558M R pretraining in direct sampling and the DRD2₁ hit-generation experiment. A simple hit generation performance proxy would nevertheless be useful. Direct sampling metrics could serve this role if they correlate with the number of hits generated during RL. Supplementary Fig. 7 (8) shows the number of hits generated during RL as a function of the sampling validity (uniqueness after canonicalization); there is no clear correlation between them when scaling model size.

Taken together, our hit generation experiments imply that there is little to gain from scaling model size beyond a certain point; while the smaller L^1/2–D^1/2–F^1/2 consistently underperformed, L2–D2–F2 generated a similar number of hits to the baseline L1–D1–F1, while scaling a single hyperparameter proved detrimental in some cases. Similar behavior has been observed for other classes of CLMs. Krüger et al. found that increasing model size beyond 15 million parameters was, on average, detrimental to property prediction performance³, while the results of Chitsaz et al. showed that scaling from 32 to 300 million parameters only gave modest improvements on average on the Practical Molecular Optimization benchmark for goal-directed de novo design^29,30. This contrasts the findings of Frey et al., who trained a transformer decoder on SELFIES (an alternative to SMILES) and found that performance improved with increasing model size even at a fixed dataset size²³. They measured only the pretraining validation loss, however, which does not necessarily correlate well with downstream task performance in general^24,25,29.

Dataset size scaling

To evaluate the effect of dataset size on downstream performance, we pretrained L^1/2–D^1/2–F^1/2, L1–D1–F1, and L2–D2–F2 on three different datasets from Table 4: 558M R (558 million molecular pairs), 7.8M R (7.8 million pairs), and 161K R (161 thousand pairs). As outlined in the “Methods” section, these datasets were constructed using serial sub-sampling followed by pair construction in order to preserve the molecular distributions in the underlying sets of molecules. We additionally included two models from the literature in our evaluations, both using the same architecture as L1–D1–F1: one trained on ChEMBL³¹ (around 7 million pairs) by He et al.³², and one trained on PubChem³³ (around 200 billion pairs) by Tibo et al.²⁷. Note that the training dataset used by Tibo et al. corresponds to applying the pair formation strategy outlined in the “Methods” section of this work to the full PubChem database. For brevity, we will refer to these models as ChEMBL-trained and PubChem-trained in the following.

Table 4.

An overview of the training datasets

Dataset	Type	SMILES count	Pair count
558M R	Random	10,000,000	558,157,922
7.8M R	Random	1,000,000	7,785,326
161K R	Random	100,000	161,181
DivMol	Diverse	1,000,000	4,636,584
DivPairs	Diverse	1,620,000	7,475,100

The random datasets were obtained by serial sub-sampling of PubChem. The DivMol and DivPairs datasets were constructed using the diversification strategies described in Sec. ?.

Direct sampling results from using the 1000 test molecules as input are shown in Table 2. Pretraining on the 161K R dataset resulted in the generation of significantly fewer valid molecules than pretraining on 558M R and 7.8M R for all models. L1–D1–F1 and L2–D2–F2 performed comparably across all three datasets, while L^1/2–D^1/2–F^1/2 produced significantly fewer valid molecules in comparison. The smallest model thus appears to be too small to effectively learn the SMILES syntax, no matter the dataset size. The r and τ₁₀ scores improved for L1–D1–F1 and L2–D2–F2 as the size of the training dataset increased, indicating that larger pretraining datasets increase the correlation between the NLL and the similarity between the input molecules and generated samples. The PubChem-trained model further improves these scores compared to the 558M R-trained baseline. Additionally, the uniqueness after removal of stereochemical information decreased with increasing dataset size, indicating that pretraining on larger datasets may lead to the generation of more stereoisomers.

Table 2.

Direct sampling from models pretrained on different datasets

Model	Dataset	Validity	Uniqueness	NS uniq.	τ10	r
L1/2–D1/2–F1/2	161K R	0.504	0.934	0.806	0.326 ± 0.24	0.303 ± 0.17
L1/2–D1/2–F1/2	7.8M R	0.710	0.919	0.648	0.380 ± 0.25	0.457 ± 0.17
L1/2–D1/2–F1/2	558M R	0.735	0.928	0.594	0.334 ± 0.29	0.457 ± 0.22
L1–D1–F1*	161K R	0.796	0.981	0.876	0.354 ± 0.22	0.224 ± 0.17
L1–D1–F1*	7.8M R	0.910	0.950	0.667	0.362 ± 0.24	0.474 ± 0.19
L1–D1–F1*	558M R	0.932	0.953	0.627	0.427 ± 0.25	0.554 ± 0.20
L2–D2–F2	161K R	0.826	0.982	0.871	0.358 ± 0.21	0.272 ± 0.15
L2–D2–F2	7.8M R	0.954	0.965	0.702	0.395 ± 0.25	0.482 ± 0.19
L2–D2–F2	558M R	0.939	0.960	0.647	0.426 ± 0.23	0.572 ± 0.20
L1–D1–F1*	PubChem	0.930	0.953	0.586	0.565 ± 0.24	0.655 ± 0.17
L1–D1–F1*	DivMol	0.731	0.922	0.628	0.390 ± 0.26	0.564 ± 0.19
L1–D1–F1*	DivPairs	0.852	0.935	0.634	0.422 ± 0.24	0.588 ± 0.16

The baseline model is marked with an asterisk. Uniqueness is calculated after SMILES canonicalization. Uniqueness after removal of stereo information is denoted by NS. τ₁₀ represents the Kendall’s tau score between the Tanimoto similarity between the input molecule and the top 10 generated molecules. r is the Pearson correlation coefficient between the NLL of the generated molecule and its score. For τ₁₀ and r, we additionally report standard deviations.

Tanimoto similarities (using 1024-bit Morgan fingerprints of radius 2) between the input molecules and samples drawn from each model-dataset pairing are shown in Supplementary Fig. 2. The distributions for L1–D1–F1 and L2–D2–F2 pretrained on the 7.8M R and 558M R datasets were similar, while that of L^1/2–D^1/2–F^1/2 was shifted slightly. The distributions after pretraining on the 161K R dataset were markedly different, however, indicating that sub-sampling failed to preserve the underlying molecular distribution in PubChem. Furthermore, the 161K R-trained L^1/2–D^1/2–F^1/2 produced quantitatively different results than the other two models, with the main peak shifted from around 0.2–0.6; the model thus appears to be too small to learn the underlying distribution at all.

We then performed goal-oriented generation using DRD2, GSK3β, and JNK3 binding affinity, along with QED as scoring functions. As we observed above, the architecture appears to play a minor role in RL performance, so we evaluated only L1–D1–F1 here. We additionally included results for the ChEMBL- and PubChem-trained models mentioned above. The number of accumulated hits after 6500 epochs of RL is shown in Fig. 2.

Fig. 2. Accumulated unique hits after 6500 epochs of RL for L1–D1–F1 trained on the 161K R, 7.8M R (7.8 million pairs), and 558M R (558 million pairs) as well as the ChEMBL- (7 million pairs) and PubChem-trained (200 billion pairs) models of He et al.³² and Tibo et al.²⁷, with 7 million and 200 billion pairs, respectively.

Each column corresponds to a target protein, and each bar cluster corresponds to a target-specific input compound as indicated by the subscript indices. Error bars represent 95% confidence intervals calculated using Eq. (1).

For DRD2, performance was similar between the models, with only minor differences in accumulated unique hits. The 161K R dataset consistently yielded a slightly lower number of hits, indicating that dataset size plays at least a minor role in performance. The PubChem-trained model did not consistently outperform models trained on smaller datasets here, however.

With GSK3β as a target, performance was more inconsistent. L1–D1–F1 pretrained on the smallest dataset generated a significantly lower number of hits than the other pretraining datasets for GSK3β₃ and GSK3β₄, while the 7.8M R and 558M R datasets yielded similar results. The PubChem-trained model performed slightly better than the rest; for GSK3β₄ in particular, it generated a significantly greater number of hits than the other models.

The results for the JNK3 scoring function were also inconsistent between models. Pretraining on the 161K R dataset yielded a low number of hits across all inputs, while pretraining on the 7.8M R dataset resulted in a low number of hits for JNK3₁. The 558M R-trained model consistently performed well across inputs, yielding more hits than the 7.8M R-trained model for both JNK3₁ and JNK3₃, while the PubChem-trained model outperformed the 558M R-trained model on all inputs but JNK3₃.

Our results indicate that pretraining dataset size plays a role in the number of generated hits, though the results vary between different targets and input molecules. Pretraining on the 161K R dataset consistently yielded fewer hits than pretraining on 7.8M R and 558M R, indicating that it is too small to train effective models. The downstream performance from pretraining on the 7.8M R and 558M R datasets was largely comparable for DRD2 and GSK3β, but the 558M R-trained model yielded more hits overall for JNK3. Generating molecules that bind to JNK3 appears to be an overall harder task, and evaluating model performance using DRD2 and GSK3β was insufficient to differentiate between models trained on sufficiently large datasets. Interestingly, we observed a linear correlation between direct sampling validity as well as uniqueness and the number of generated hits during RL with increasing dataset size for L1–D1–F1, as shown in Supplementary Figs. 7 and 8. The trend did not persist with the PubChem-trained model, however, which generated more hits than the model pretrained on 558M R, but exhibited lower direct sampling validity and uniqueness. This indicates that these metrics are poor proxies for hit generation performance when scaling the dataset size.

The performance of the 7.8M R dataset and the ChEMBL-trained model was also similar across the board. The latter model generated significantly fewer hits for JNK3₂ and JNK3₄; given the similar pair counts of the ChEMBL and 7.8M R datasets, this may be explained by differences in their molecular distributions. The overall best performer was the PubChem-trained model, which generated the highest number of hits on most tasks. Its performance was, in many cases, similar to that of the 7.8M R- and 558M R-trained models, indicating diminishing returns past a certain size threshold. This performance gain comes at the cost of a substantially higher computational cost, however. Similarity-informed pair formation scales quadratically with the number of molecules; preprocessing the full PubChem database (approximately 100 million molecules) is thus ten thousand times more complex than preprocessing the 1 million molecule PubChem subset used for the 7.8M R dataset. Moreover, assuming an approximately linear relationship between training dataset size and the number of pretraining steps required, the PubChem-trained model required roughly four orders of magnitude more compute than our 7.8M R-trained baseline.

Dataset diversification

To evaluate the effects of dataset distribution, we pretrained L1–D1–F1 on two diversified datasets: DivMol, containing around 4.6 million pairs constructed from 1 million molecules, and DivPairs, with 7.5 million pairs constructed from 1.62 million molecules. The diversification strategy is described in the Methods section. We then performed the same experiments as above.

Results from direct sampling from the models are shown in Table 2. The DivMol- and DivPairs-trained models exhibited a lower validity than the 7.8M R-trained baseline, with slightly lower uniqueness. This indicates that the diversification strategy makes it more difficult for the model to learn the SMILES syntax. The r and τ₁₀ scores for the models pretrained on the diversified datasets were slightly higher than for the 7.8M R-trained baseline; in other words, samples drawn from the DivMol- and DivPairs-trained models are more likely to be similar to the input molecule. The r-scores for the 558M R- and PubChem-trained models were even higher, with the latter also exhibiting a higher τ₁₀, indicating that dataset size is important for these metrics. Supplementary Fig. 3 shows Tanimoto similarities (using 1024-bit Morgan fingerprints of radius 2) between the input molecules and samples generated by L1–D1–F1 pretrained on 7.8M R, 558M R, DivMol, and DivPairs, as well as the PubChem-trained model. Compared to the 7.8M R-trained baseline, models trained on diversified data yielded output distributions shifted towards higher similarity; the PubChem-trained model exhibited an even larger shift, consistent with the direct sampling experiment above.

The number of accumulated unique hits after 6500 epochs of RL is shown in Fig. 3, which additionally includes results from L1–D1–F1 pretrained on the 7.8M R and 558M R datasets as well as the ChEMBL- and PubChem-trained models.

Fig. 3. Accumulated unique hits (top) and diversity of generated hits (bottom) after 6500 epochs of RL for L1–D1–F1 trained on the 7.8M R (7.8 million pairs), DivMol (4.6 million pairs), and DivPairs (7.5 million pairs) datasets.

We additionally include results for the ChEMBL-trained model of He et al.³² and the PubChem-trained model of Tibo et al.²⁷ (200 billion pairs). Each column corresponds to a target protein, and each bar cluster corresponds to a target-specific input compound as indicated by the subscript indices. Error bars represent two-sided 95% confidence intervals of the mean computed using Eq. (1).

For DRD2, all models generated a comparable number of hits. The 7.8M R-trained model yielded a slightly lower number of hits for input compounds DRD2₁ and DRD2₃ compared to the DivMol- and DivPairs-trained models, while the PubChem-trained model generated slightly fewer hits for DRD2₃.

With GSK3β as a target, the models perform comparably for all input molecules except for GSK3β₄, for which pretraining on DivMol resulted in significantly fewer hits.

Optimizing for JNK3 binding paints a more nuanced picture. The DivMol-trained model generated a significantly higher number of hits than the 7.8M R-trained model for JNK3₁, despite the reduction in training pairs. Pretraining on the DivMol dataset resulted in a significantly lower number of hits than the other models for the input compound JNK3₂.

Figure 4 reports normalized hit counts relative to the L1–D1–F1 model trained on 7.8M R, averaged over targets and input compounds. While aggregating hits masks target-specific behavior, it provides a simple metric for comparisons across datasets. On average, both the DivPairs- and DivMol-trained models exceeded the hit rate of the 7.8M R baseline despite the lower sampling validity of the models trained on diverse datasets. Notably, DivPairs pretraining also slightly outperformed the 558M R-trained model, even with substantially fewer training pairs. PubChem pretraining yielded the highest overall hit count; the gain in hit rate was disproportionately small relative to the increase in training set size, however, indicating diminishing returns from dataset size scaling. In line with the aggregate results, DivMol and DivPairs yielded comparable hit rates in most target-specific experiments. GSK3β₄ and JNK3₂ are exceptions here, where the DivMol-pretrained model underperformed. Conversely, for JNK3₁, DivMol generated significantly more hits for JNK3₁ than the 7.8M R-trained baseline, while DivPairs pretraining resulted in a significant hit rate across all tested cases. Overall, dataset composition appears to influence the number of generated hits: diversified datasets often outperformed sub-sampled datasets randomly, although maximizing dataset size is an effective strategy for achieving high hit rates.

Fig. 4. Aggregated hits across all targets and input compounds as a function of the number of training pairs for L1–D1–F1.

Hits are normalized to 7.8M R pretraining.

In practice, hit count is not the only factor of interest in drug design; hit diversity is also important, as efficient exploration of chemical space increases the likelihood of discovering novel ligands. Fig. 3 shows the diversity of unique hits in the form of #Circles for the dataset diversification experiments. The hit diversity for the diversified models exceeded, or at least matched, that of the 7.8M R-trained model while generating more hits on average. Additionally, the overall hit diversities of the diversity-trained models were higher than, or comparable to, that of the PubChem-trained model, despite the diversified datasets being four orders of magnitude smaller in terms of training pairs. Importantly, this dataset size reduction proportionally lowers the computational effort required for data preprocessing and pretraining, enabling more efficient model development. Dataset diversity thus has a greater impact than scale on hit diversification, allowing for broader exploration of the chemical space around each input molecule without a significant reduction in the overall number of hits; this is an advantage in scenarios where hit diversity is more important than hit rate maximization.

Discussion

Assessing the performance of CLMs remains challenging, as there is no single universally accepted metric that can be generalized to every model. For our evaluations, we therefore focused primarily on a realistic drug design scenario: hit generation.

Our results demonstrate that scaling up the model size has a limited, and potentially detrimental, impact on the downstream performance of this class of CLMs once the model is large enough, contributing to the growing evidence that there is little to gain from scaling up the size of CLMs. While we have focused on encoder-decoder models for molecular optimization, our observations are in line with what has previously been found for other classes of CLMs^3,29. Furthermore, direct sampling-based measures like validity or uniqueness correlate poorly with downstream task performance: models with higher validity or uniqueness do not necessarily generate a higher number of hits. This highlights the importance of evaluating models on realistic downstream tasks rather than more easily evaluated proxy measures, in line with previous work on both NLP and CLMs^24,25,29.

We found the size of the pretraining dataset to be more important for downstream performance than model size. The smallest pretraining dataset we evaluated, 161K R, resulted in a model that generated significantly fewer hits than the rest. Pretraining on the 7.8M R and 558M R datasets yielded a comparable number of hits for DRD2 and GSK3β, but the 558M R-trained model generated significantly more hits for JNK3₁ and JNK3₃. The PubChem-trained model performed best overall; the gain in average hit rate was disproportionately small relative to the increase in dataset size, however. Dataset size scaling thus appears to have diminishing returns in terms of downstream task performance, with differences becoming meaningful only for more difficult tasks like JNK3 hit generation. The Therapeutics Data Commons (TDC) DRD2 model we use, a common optimization case study for AI-guided drug design^34,35, may thus be insufficient for comparing the relative performance of generative models in this setting.

Dataset composition plays an important role as well. Pretraining on the DivMol dataset (4.6 million pairs) yielded a similar number of hits as pretraining on the 7.8M R dataset (7.8 million pairs), despite the significant reduction in training pairs. The two datasets led to different failure modes: the former underperformed relative to the other models for GSK3β₃ as well as JNK3₂, while the latter underperformed for JNK3₁. The DivPairs-trained model (7.5 million pairs) performed similarly to the 558M R-trained model and additionally exhibited no failure modes across the tested targets. Similar observations hold when comparing to the PubChem-trained model, despite the latter being trained on 200 billion pairs. In practice, training on the largest dataset available appears to yield good downstream performance. This may be computationally infeasible, however, and the number of hits is not the only factor relevant to drug design; hit diversity is also important. The DivMol- and DivPairs-trained models yielded hits with greater overall diversity than the other models. Dataset diversification thus facilitates greater exploration of chemical space without significantly impacting the overall number of generated hits, allowing for the training of high-performance models at a substantially lower compute cost.

We note that dataset diversification strategies have been explored in NLP²¹, but, to our knowledge, the concept has not previously been applied to generative CLMs. In chemistry, diversity-enhancing techniques are typically applied during RL^36–38, including the scaffold-based penalty used here. Other combinations of these methods with the dataset diversification strategy presented in this work may further increase hit diversity; however, a systematic evaluation of this is left for future studies.

Methods

We explore the factors affecting the performance of chemical language models, specifically focusing on the Mol2Mol encoder-decoder transformer used for molecular optimization in the REINVENT4 framework. Molecular structures are represented as SMILES strings³⁹, which are tokenized using the regular expression-based scheme of Schwaller et al.⁵ The complete set of unique tokens present in the tokenized SMILES forms a fixed vocabulary, which is subsequently used by generative models to construct molecules. Mol2Mol is pretrained on molecular pairs: given an input SMILES string, it generates another SMILES. During pretraining, the model is optimized using a language modeling loss based on next-token likelihoods over target SMILES sequences conditioned on the input, thus learning the conditional distribution of chemically plausible transformations. By constructing training pairs with high Tanimoto similarity (above 0.5 in this work), the model learns to encode prior knowledge about the chemical space around a given input SMILES. While this strategy produces a model (often called a prior in other literature) capable of generating analogs, there is no inherent ordering of source-target pairs. Near neighbors to the input compound in chemical space are thus as likely to be generated as more distant ones. To overcome this, we additionally use the ranking loss of Tibo et al. (see Section 7 of the Supplementary Information) that penalizes mismatches between generation likelihood and source-target similarity, i.e., the model is trained to correlate generative likelihood and source-target similarity²⁷. This makes Mol2Mol capable of probabilistically exhaustive neighborhood enumeration. The resulting pretrained model can then be steered via RL in REINVENT4 to generate property-optimized molecules, enabling goal-oriented design. In this setup, samples are drawn from the model and assigned scores according to how well their predicted properties align with the target objective. The model is then updated using these scores to increase the likelihood of sampling high-scoring molecules. This is repeated until a sufficient number of samples have been produced.

Models

To investigate how variations in model architecture and model size affect performance, we systematically vary three hyperparameters: the hidden dimension, the feed-forward dimension, and the number of layers. The encoder-decoder architecture is visualized in Fig. 5 with these hyperparameters highlighted.

Fig. 5. The encoder–decoder transformer considered in this work.

The hyperparameters under consideration in this work are the number of layers (blue), the model dimension (purple), and the feed-forward dimension (green).

The hidden dimension is defined as the dimensionality of the input and output vectors. This dimension is consistent across all layers and serves as the size of the embeddings and the internal representations within the model. Each transformer layer consists of two main sublayers: a multi-head self-attention mechanism and a feed-forward network. The layout of the encoder and decoder modules is the same, but each decoder layer includes an additional cross-attention sublayer for multi-head attention over the encoder’s output. This allows the decoder output to be conditioned on the input of the encoder. Within each layer, the feed-forward network expands the model’s internal representation from the model dimension to the feed-forward dimension, then down-projects back to the model dimension. The number of attention heads is fixed to 8, with each head projecting the internal representations into an equally sized subspace. Additionally, to mitigate overfitting, a dropout rate of 0.1 is applied to each sublayer before applying the residual connection and layer normalization.

We use the model layout used by Tibo et al.²⁷ as a baseline, with 6 layers (L), a model dimension of 256 (D), and a feed-forward dimension of 2048 (F). Each of these hyperparameters is independently scaled by a factor of 2 or 4 to construct scaled models. We additionally consider the case of scaling them simultaneously by a factor of 0.5 or 2. For brevity, we refer to the baseline model as L1-D1-F1, with the scaled models named according to the applied scaling. Model names and parameters are summarized in Table 3.

Datasets

To study the effect of varying the number of training datapoints, we constructed multiple datasets by sub-sampling molecules from PubChem Compound³³, which contained a total of 102,419,980 entries at the time of download (end of December 2021). Prior to sub-sampling, all PubChem molecules were preprocessed and filtered using the REINVENT4 data pipeline to ensure data consistency. The pipeline first applies a regular expression filtering step to tokenize each molecule and validate basic syntax, excluding those containing unsupported tokens. Additionally, an RDKit-based filtering step is conducted, which includes standard RDKit normalization. Any compounds containing elements outside [I, O, Cl, N, C, F, S, Br] were excluded. Stereochemistry is kept, but isotope information is removed. Molecules are retained if they have between 2 and 90 heavy atoms. Molecules failing RDKit sanitization are removed, and charge neutralization is performed. We retain only the largest organic fragment in multi-fragment cases, discarding all others. Additional filtering criteria include a maximum molecular weight (MW) of 1200 Da, a minimum of 2 carbon atoms per molecule, a maximum of 12 rings, and a maximum ring size of 7.

The preprocessed PubChem database was then used to construct an initial set of 10 million molecules through random sub-sampling without replacement, of which 90% was used for training and 8% for validation. From the remainder, 1000 molecules were set aside as a test set.

To create smaller datasets, the initial training and validation splits were reduced in size by a factor of 10 by random sub-sampling, producing a set of 1 million molecules. The procedure was repeated to produce a third set of 100,000 molecules.

The molecular optimization models considered in this work are trained on molecular pairs. To enable this, all possible unique pairs within the training and validation splits were enumerated, including self-pairs. We then computed the Tanimoto similarity of each pair using count-based Morgan fingerprints of radius 2. Following He et al.^32,40 and Tibo et al.²⁷, all pairs with a Tanimoto similarity above 0.5 were included in the final datasets. The training datasets respectively contain 558 million, 7.8 million, and 161,181 pairs. We will refer to them as 558M R, 7.8M R, and 161K R, respectively. A summary of the training datasets we consider is reported in Table 4.

Diverse datasets

To study the effect of dataset composition on model performance, we explore an alternative to random sub-sampling: diversification. To this end, we constructed the DivMol dataset. The preprocessed PubChem database was first divided into five groups based on MW range: <200 Da, 200–300 Da, 300–400 Da, 400–600 Da, and >600 Da. Each group was then clustered using the BitBIRCH algorithm⁴¹ using a threshold of 0.15 and a branching factor of 1000, yielding a total of 28,633 clusters. Individual molecules were randomly sampled from each cluster until the number of molecules reached 1 million while keeping the MW distribution similar to the filtered PubChem database. We note that the approach is conceptually similar to diversity-based selection techniques for virtual screening^42,43.

The DivMol dataset contains about 4.6 million pairs, considerably fewer than the 7.8M R dataset. To facilitate a fairer comparison between randomly sub-sampled and diversified datasets, we additionally constructed the DivPairs dataset by applying the diversification strategy until we reached 1.62 million molecules, which yielded a total of 7.5 million molecular pairs. We note that the diverse datasets contain fewer targets per source compound than 7.8M R and 558M R. The diversities (in terms of #Circles, defined in Eq. 2 below) of the 161K R, 7.8M R, 558M R, DivMol, and DivPairs datasets are shown in Supplementary Fig. 9, showing that the diversification strategy produced sets of molecules with enhanced diversity in comparison to datasets of similar size. The diversified datasets are also summarized in Table 4.

Training procedures

All models were pretrained using a batch size of 128 per GPU using the Adam optimizer and a learning rate of 1e − 4 using the pretraining framework developed by Tibo et al.^27,44,45. 8 GPUs were used for model training, except for the 161K dataset for which two GPUs were used. Training data was shuffled in each epoch. We use the NLL loss regularized with the ranking term proposed by Tibo et al.²⁷ using λ = 10 and the Tanimoto similarity (using count-based Morgan fingerprints of radius 2) between the training pairs as the kernel function. The loss is described in more detail in Section 7 of the Supplementary Information. The NLL acts as a proxy for precedence, i.e., the probability of a given source compound being transformed into a given target compound; a lower NLL implies a higher precedence. The ranking term effectively penalizes the loss if the similarities between training pairs and their precedence are ordered differently; given two training pairs P₁ and P₂ with NLLs and Tanimoto similarities (NLL₁, S₁) and (NLL₂, S₂), the penalty is positive if NLL₁ < NLL₂ but S₁ < S₂, or vice versa, and zero otherwise. As mentioned above, we apply dropout with a probability of 0.1 during training to avoid overfitting.

Additionally, we apply an early stopping strategy based on validation loss. For models trained on the 161K R dataset, training was stopped after 40 epochs without a reduction in validation loss. For the 7.8M R dataset, as well as the diverse datasets, training was stopped after 10 epochs without improvement. Models trained on the 558M R dataset were trained for two full epochs without early stopping due to resource constraints.

Model evaluation

We employed two complementary strategies for model evaluation: direct sampling from the models and goal-oriented generation using RL. For direct sampling, the reserved test set, consisting of 1000 molecules, was used as input. We employed beam search (beam size 1000) with a softmax temperature of 1, following Tibo et al.^27,44, generating 1,000,000 SMILES samples. For these, we assessed validity (the fraction of samples that RDKit parses successfully) and uniqueness (the fraction of non-repeated samples after canonicalization). Because the training data consisted of isomeric SMILES, we also report stereochemistry-agnostic uniqueness, computed after removing stereochemical information, i.e., considering stereoisomers to be identical. Direct sampling from the models we consider should ideally yield an enumeration of molecules similar to the input, and more similar molecules should have a higher probability of being generated, i.e., a lower NLL²⁷. Following Tibo et al., we therefore computed correlation metrics between the Tanimoto similarity (using count-based Morgan fingerprints of radius 2) and NLL in the form of the Pearson correlation coefficient r (the linear correlation between the two quantities) and Kendall’s tau score for the top 10 beam search samples τ₁₀²⁷. The latter is a measure of ordering between samples, that is, how likely it is for a sample with a lower NLL than another to also have a higher similarity with the input molecule. Both measures lie in [−1, 1], with 1 implying perfect correlation, −1 perfect anti-correlation, and 0 lack of correlation.

We then performed RL to encourage the generation of molecules active against different targets, to evaluate performance on goal-oriented tasks. For this purpose, we used pretrained oracle models from TDC⁴⁶ for DRD2, JNK3, and GSK3β activity as scoring functions. Additionally, we included QED⁴⁷ as a scoring component to bias generation toward drug-like molecules. The final score was computed as the geometric mean of the oracle output and the QED. Finally, we encouraged structural diversity by halving the score of molecules that share Bemis-Murcko scaffolds⁴⁸ with previously generated compounds¹⁵ after a given scaffold had been generated more than 25 times. We employed multinomial sampling with a softmax temperature of 1 following He et al.^32,40. Four input molecules were used for each oracle, which were selected from ChEMBL³¹ compounds active against each target with pIC₅₀ ≥ 5. For DRD2, we selected the same four inputs as He et al.⁴⁰. For brevity, we will use enumerated short-hands for these molecules: DRD2₁, DRD2₂, etc. The SMILES for each of these, along with a depiction of their molecular graphs, is given in Supplementary Table 1. Five independent runs were performed for every RL experiment, each running for 6 500 epochs. One hundred and twenty-eight molecules were generated using multinomial sampling with a temperature of 1 in each epoch, resulting in a total of 832,000 molecules per run. The results are reported as averaged promising hits with 95% confidence intervals computed as

$$\overline{x}\pm t_{0.975,n-1}\frac{s}{\sqrt{n}},$$ 1

where $\overline{x}$ is the sample mean, s is the sample standard deviation, n is the number of observations, and t_0.975,n−1 is the 97.5% quantile of the Student’s t-distribution with n − 1 degrees of freedom.

Here, promising hits are defined as molecules with unique Bemis-Murcko scaffolds with a target activity score greater than 0.6, as well as a QED above 0.6. We additionally analyzed the diversity of generated hits through the #Circles diversity measure proposed by Xie et al.⁴⁹ as implemented by Renz et al.¹⁴, with a distance threshold t of 0.65 using Tanimoto distances and 2048-bit Morgan fingerprints of radius 2, as well as the mean bulk Tanimoto similarity (using 1024-bit Morgan fingerprints of radius 2) against the input molecule for some of the experiments. We note that Renz et al. (in Section S1) argued that #Circles aligns better with the goal of generating structurally diverse molecules than internal diversity, and therefore chose to adapt the metric ourselves¹⁴. The #Circles metric is defined as⁴⁹

$$\begin{array}{c}\#\mathrm{Circles}(\mathcal{S};d,t):=\underset{\mathcal{C}\in \mathcal{S}}{\max }\mid \mathcal{C}\mid \hfill \\ s.t.d(x,y)>t,\forall x\ne y\in \mathcal{C},\hfill \end{array}$$ 2

where $\mathcal{C}$ is a subset of a set of molecules $\mathcal{S}$, t a distance threshold corresponding to the diameter of a circle, and d a distance metric that measures the distance between molecules x and y. It can be thought of as finding the maximal number of molecules with pairwise distances greater than a given threshold¹⁴, or equivalently, finding the maximum number of mutually exclusive molecule-centered n-dimensional circles of diameter t required to cover the chemical space spanned by the n-dimensional fingerprints of a set of molecules⁴⁹. A higher #Circles thus implies higher chemical diversity.

Author contributions

B.M. trained the models and performed all experiments. B.M., J.P.J., and N.Ö. wrote the main paper text. A.T. and J.H. provided technical and scientific support. J.P.J. and N.Ö. supervised the work. All authors contributed to conceptualizing the study and reviewed the paper.

Data availability

Final model checkpoints as well as training, validation, and testing data are available on Zenodo at https://zenodo.org/records/17552157.

Competing interests

The authors declare no competing interests.

Supplementary information

The online version contains supplementary material available at 10.1038/s44386-026-00056-3.