Predicting small molecules solubilities on endpoint devices using deep
  ensemble neural networks
Aqueous solubility is a valuable yet challenging property to predict.
Computing solubility using first-principles methods requires accounting for the
competing effects of entropy and enthalpy, resulting in long computations for
relatively poor accuracy. Data-driven approaches, such as deep learning, offer
improved accuracy and computational efficiency but typically lack uncertainty
quantification. Additionally, ease of use remains a concern for any
computational technique, resulting in the sustained popularity of group-based
contribution methods. In this work, we addressed these problems with a deep
learning model with predictive uncertainty that runs on a static website
(without a server). This approach moves computing needs onto the website
visitor without requiring installation, removing the need to pay for and
maintain servers. Our model achieves satisfactory results in solubility
prediction. Furthermore, we demonstrate how to create molecular property
prediction models that balance uncertainty and ease of use. The code is
available at https://github.com/ur-whitelab/mol.dev, and the model is usable at
https://mol.dev.
Mayk Caldas Ramos, Andrew D. White
# Predicting small molecules solubilities on endpoint devices using deep ensemble neural networks

###### Abstract

Aqueous solubility is a valuable yet challenging property to predict. Computing solubility using first-principles methods requires accounting for the competing effects of entropy and enthalpy, resulting in long computations for relatively poor accuracy. Data-driven approaches, such as deep learning, offer improved accuracy and computational efficiency but typically lack uncertainty quantification. Additionally, ease of use remains a concern for any computational technique, resulting in the sustained popularity of group-based contribution methods. In this work, we addressed these problems with a deep learning model with predictive uncertainty that runs on a static website (without a server). This approach moves computing needs onto the website visitor without requiring installation, removing the need to pay for and maintain servers. Our model achieves satisfactory results in solubility prediction. Furthermore, we demonstrate how to create molecular property prediction models that balance uncertainty and ease of use. The code is available at [https://github.com/ur-whitelab/mol.dev](https://github.com/ur-whitelab/mol.dev), and the model is usable at [https://mol.dev](https://mol.dev).

Solubility, Small Molecule, Deep Ensemble, Recurrent Neural Network

## 1 Introduction

Aqueous solubility measures the maximum quantity of matter that can be dissolved in a given volume of water. It depends on several conditions, such as temperature, pressure, pH, and the physicochemical properties of the compound being solvated.[1] The solubility of molecules is essential in many chemistry-related fields, including drug development[2, 3, 4, 5], protein design[6], chemical[7, 8] and separation[9] processes. In drug development, for instance, compounds with biological activity may not have enough bioavailability due to inadequate aqueous solubility.

Solubility prediction is critical and has driven the development of several methods, including first principles[10, 11], semi-empirical equations[12, 13, 14], molecular dynamics (MD) methods[15, 16, 17, 18], quantum computations[19], and quantitative structure-property relationship (QSPR)[20, 21, 22, 23] methods. Despite significant progress, the development of accurate and reliable models for solubility remains a major concern.[24]

To address the persistent issues of systematic bias and non-reproducibility in aqueous solubility datasets, Llinas et al.[25, 26] introduced two solubility challenges featuring consistent data. The first challenge evaluated participants based on the root mean square error (RMSE) obtained and the percentage of correct values within a \(\pm 0.5\) logS error range. Unfortunately, the authors did not report the methods used by the participants.[27] In contrast, the second challenge showed that regardless participants were open to choosing which approach to use, all submitted responses used an implementation of QSPR or machine learning (ML).[28] Neural networks (NN), multiple linear regression (MLR), and decision trees were the most commonly applied methods in these challenges. Tree-based and MLR models presented the best results. Surprisingly, new state-of-art methods did not yield a significant improvement in predictions compared to the results of the first challenge.[27] The challenges' findings showed that data quality is more critical for
accurate predictions than model selection.[28] The results from the solubility challenges will be discussed in detail in Section 5.

Ideally, solubility models should be accurate and accessible, having clear or minimal instructions on how to use a model. Thus a common idea is to use web servers to provide easier public access. However, maintaining a web server requires an ongoing investment of time and money. There are examples of servers that eventually disappear, even with institutional or government support[29]. For example, eight of the 89 web server tools from the 2020 _Nucleic Acid Research_ special web server issue are already offline[30]1 after just a few years. Additionally, some tasks may require a long computation time[31]. For instance, tools like RoseTTAFold[32] and ATB[33] can take hours to days to complete a job, resulting in long queues and waiting times.

Footnote 1: Tested December 30, 2022

An alternative approach is to perform the computation directly on the user's device, removing the need for the server's maintenance and cost. In this approach, the website is simply a static file that can be hosted on sites like GitHub and be completely archived in the Internet Archive2. We explored this approach in Ansari and White [34] for bioinformatics. The main drawback is that the application runs directly from the browser on a user's device (a personal computer or even a cellphone). This would be infeasible for first-principle methods, like those that rely on molecular dynamics. Nevertheless, it is feasible for deep learning models, especially with the increasing integration of deep learning chips and compiler optimizations.

Footnote 2: [https://archive.org/](https://archive.org/)

In this work, we developed a front-end application using a JavaScript (JS) implementation of TensorFlow framework[35]. Our application can be used to predict the solubility of small molecules with uncertainty. To calibrate the confidence of the prediction, our model implements a deep ensemble approach[36] which allows reporting model uncertainty when reporting the prediction. Our model runs locally on the user's device and can be accessed at [https://mol.dev/](https://mol.dev/).

## 2 Related works

Physic-based models have been developed in the past for aqueous solubilization predictions. Those models may become complex, limiting their use to advanced users only.[24] Despite physic-based models being derived from first principles, they are no more accurate than empirical methods.[37] Data-driven models can outperform physics-based models with the benefit of being less time-consuming. Historically, common approaches computed aqueous solubilities based on QSPR[21, 22, 23] and MLR[38, 39] methods.[24]

Huuskonen [38] used a dataset consisting of 1297 organic molecules to develop two models based on MLR and NN. The author reported a good correlation between predicted properties and labels for training (\(r^{2}=0.94\)) and test (\(r^{2}=0.92\)) data. Delaney [39] used another approach based on MLR called Estimated SOLubility (ESOL) adjusted on a 2874 small organic molecules dataset. The final model presented a \(r^{2}=0.55\) and an average absolute error (AAE) of 0.83. GPSol[40], a Gaussian Process-based model, was trained to predict the aqueous solubilities of electrolytes in addition to non-electrolyte molecules. It used 1664 descriptors computed by Dragon software as input features to train the model on a dataset of \(\sim\)4000 molecules. Depending on the dataset, it presented an RMSE of \(0.77\) or \(0.61\). Lusci _et al.[41]_ trained several Undirected graph recurrent neural networks (UG-RNN) architectures using different sets of node feature vectors. The authors report RMSE from 0.90 to 1.41 for different models on the first Solubility Challenge dataset[25]. McDonagh _et al.[37]_ calculated solubilization free energies using first principle theoretical calculations and cheminformatic methods. Their results have shown that cheminformatic methods have better accuracy than theoretical methods. The authors also point to the promising results of using Random Forest models (RMSE of 0.93 on Llinas first dataset[25]).

Those models use descriptors to represent molecules. Descriptors are a straightforward way to convey physical-chemical information in your input. However, descriptor selection is not an easy task. It requires a good understanding of the problem settings, usually only held by specialists. Some automated methods have been proposed to select descriptors[42]. Nevertheless, computing several descriptors can increase the time needed for inference. Additionally, these descriptors can be valid only for a specific region of the chemical space[43].

More recently, transformers models have been used to compute the solubility of small molecules. Francoeur _et al.[44]_ developed the SolTranNet, a transformers model trained on AqSolDB[1] solubility data. Notably, this architecture results in an RMSE of only \(0.278\) when trained and evaluated on the original ESOL[39] dataset using random split. Nevertheless, it shows an RMSE of \(2.99\) when trained using the AqSolDB[1] and evaluated using ESOL. It suggests that the molecules present in ESOL may have low variability, meaning that samples in the test set are similar to samples in the training set. Hence, models trained on the ESOL training set performed excellently when evaluated on the ESOL test set. Regression Transformer (RT)[45] is a multipurpose transformer model trained using an infilling mask
approach[46]. Their results are comparable to those achieved by the SMILES-BERT[47] and Mol-BERT[48] models. While the RMSE values for SMILES-BERT and Mol-BERT are \(0.47\) and \(0.53\), respectively. Whereas RT presented an RMSE of \(0.73\). All three models were fine-tuned using ESOL. MolFormer[49] is an encoder-only transformers model with a modified embedding. It was pre-trained in a large corpus and fine-tuned for numerous downstream tasks. Specifically for the solubility regression fine-tuning, they reported an RMSE of \(0.278\) on the ESOL dataset. Noticeably, the same value was reported by Francoeur _et al.[44]_ when they trained their model on ESOL.

Comparing the performance of different models is a complex task, as performance metrics cannot be directly compared across models evaluated on distinct datasets. To address this issue, Panapitiya et al. [50] curated a large and diverse dataset to train models with various architectures and molecular representations. They also compared the performance of these models on datasets from the literature[25, 26, 38, 39, 51, 52, 53, 54, 55, 56, 57, 58]. Although their models achieved an RMSE of \(\sim 1.1\) on their test set, using descriptors as molecular representations resulted in RMSE values ranging from \(0.55\) to \(\sim 1.35\) when applied to other datasets from the literature. These findings suggest that some datasets used to train models in the literature may be inherently easier to predict, leading to smaller RMSE values. According to their study, the Solubility Challenge datasets by Llinas _et al.[25, 26]_ were found to be particularly challenging due to their more significant reproducibility error.

## 3 Methods

### Dataset

The data used for training the models were obtained from AgSolDB[1]. This database combined and curated data from 9 different aqueous solubility datasets. The main concern in using a large, curated database is to avoid problems with the generalizability of the model[59] and with the fidelity of the data[60]. AgSolDB consists of aqueous solubility (LogS) values for 9982 unique molecules extended with 17 topological and physicochemical 2D descriptors calculated by RDKit[61].

We augmented AgSolDB to 96,625 molecules. Each entry of AgSolDB was used to generate at most ten new unique randomized SMILES strings. Training the model on multiple representations of the same molecule improves its ability to learn the chemical space constraints of the training set, as demonstrated in previous studies [62, 63]. Duplicates were removed.

After shuffling, the augmented dataset was split into 80%/20% for the training and test datasets, respectively. The curated datasets for the solubility challenges[25, 28] were used as withheld validation data to evaluate the model's ability to predict solubility for unseen compounds. To refer to the validation datasets, we labeled the first solubility challenge dataset as "solubility challenge 1" and the two sets from the second solubility challenge as "solubility challenge 2_1" and "solubility challenge 2_2", respectively. Molecules in these three datasets were not found in train and test datasets.

### Model architecture

Our model uses a deep ensemble approach as described by Lakshminarayanan et al. [36]. This technique was selected due to its ability to estimate prediction uncertainty, thus enhancing the predictive capability of our model. The uncertainty of a model can be divided into two sources: aleatoric uncertainty (AU) and epistemic uncertainty (EU).[64, 65] These uncertainties quantify the intrinsic uncertainty inherent in data observations and the disagreement among model estimations, respectively.[66]

Given a model that outputs two values - \(\hat{\mu}_{m}\) and \(\hat{\sigma}_{m}\) - that characterize a normal distribution \(\mathcal{N}(\hat{\mu}_{m},\hat{\sigma}_{m})\), a deep ensemble creates an ensemble of \(m\) models that can estimate prediction uncertainty. For a given data point \(\vec{x}\), the estimates for the ensemble predictions are computed as follows:

\[\hat{\mu}(\vec{x})=\frac{1}{N}\sum_{m}\hat{\mu}_{m}(\vec{x}) \tag{1}\]

\[\hat{\sigma}_{ale}^{2}(\vec{x})=\frac{1}{N}\sum_{m}\hat{\sigma}_{m}^{2}(\vec{ x})\,\ \ \hat{\sigma}_{epi}^{2}(\vec{x})=\frac{1}{N}\sum_{m}\left(\hat{\mu}(\vec{x})-\hat{ \mu}_{m}(\vec{x})\right)^{2} \tag{2}\]

where \(\hat{\sigma}_{ale}^{2}\) is AU, \(\hat{\sigma}_{epi}^{2}\) is EU, N is the ensemble size, and m indexes the models in the ensemble.

We used a deep neural network (DNN) implemented using Keras[67] and TensorFlow[68] to build the deep ensemble. Our DNN model uses Self-referencing embedded strings (SELFIES)[69] tokens as input. Simplified molecular-input
line-entry system (SMILES)[70] or SELFIES[69] molecule representations are converted to tokens based on a pre-defined vocabulary generated from our training data, resulting in 273 available tokens. Figure 1 illustrates the model architecture. The network can be divided into three sections: (\(i\)) Embedding, (\(ii\)) bi-RNN, and (\(iii\)) fully connected NN.

The embedding layer converts a list of discrete tokens into a fixed-length vector space. Working on a continuous vector space has two main advantages: it uses a more compact representation, and semantically similar symbols can be described closely in vector space. Our embedding layer has an input dimension of 273 (vocabulary size) and an output dimension of 64.

Following the embedding layer, the data are fed into the bidirectional Recurrent Neural Network (RNN) layer. We used two RNN layers, each containing 64 units. The effects of using Gated Recurrent Unit (GRU) or Long Short-Term Memory (LSTM)[71] layers as the RNN layers were investigated (refer to Section 4.1). Using bi-RNN was motivated based on our previous work[34] in which LSTM helped improve the model's performance for predicting peptide properties using its sequences. More details regarding RNN, LSTM, and GRU layers can be found in Ref. 72.

The output from the bi-LSTM stack undergoes normalization via Layer Normalization[73]. There is no agreement on why Layer Normalization improves the model's performance.[74, 75, 76, 77] The absence of a comprehensive theoretical understanding of normalization effects hinders the evolution of novel regularization schemes.[78] Despite the limited understanding, Layer Normalization is employed due to its demonstrated effectiveness.[77]

After normalization, data is processed through three dense layers containing 32, 16, and 1 units, respectively. The 16-unit layer's output goes to two different 1-unit layers. One layer uses a linear function and the other uses a softplus function, producing \(\hat{\mu}_{m}\) and \(\hat{\sigma}_{m}\), respectively.

Negative log-likelihood loss \(l\) was used to train the model. It is defined as the probability of observing the label \(y\) given the input \(\vec{x}\):

\[l(\vec{x},y)=\frac{log(\hat{\sigma}_{m}^{2}(\vec{x}))}{2}+\frac{ \left(y-\hat{\mu}_{m}(\vec{x})\right)^{2}}{2\hat{\sigma}_{m}^{2}(\vec{x})} \tag{3}\]

Figure 1: Scheme of the deep learning DNN. The molecule is input using the SMILES or SELFIES representation. This representation is converted to a tokenized input based on a vocabulary obtained using the training dataset. A set of models represents the deep ensemble model. Each model consists of an embed layer, two bidirectional RNN (bi-RNN) layers, a normalization layer, and three fully connected layers being down-sized in three steps. Dropout layers are present after the embed and after each fully connected layer during training, but they were not represented in this scheme. Predictions of the models in the ensemble are then aggregated.
During the training phase, dropout layers with 0.35 dropout rate were incorporated after the embedding and each dense layer to mitigate over-fitting.[79] Models were trained using the Adam[80] optimizer with a fixed learning rate of 0.0001 and default values for \(\beta_{1}\) and \(\beta_{2}\) (0.9 and 0.999, respectively).

Our model employs adversarial training, following the approach proposed by Lakshminarayanan et al. [36] to improve the robustness of our model predictions. Because the input for our model is a discrete sequence, we generate adversarial examples by modifying the embedded representation of the input data. Each iteration in the training phase consists of first computing the loss using Equation 3 and a second step with a new input \(\vec{x}^{\prime}\) to smooth the model's prediction:

\[\vec{x}^{\prime}=\vec{x}+\epsilon\text{sign}(\nabla_{x}l(\vec{x},y)) \tag{4}\]

where \(\epsilon\) is the strength of the adversarial perturbation.

Details of the model performance, limitations, training data, ethical considerations, and caveats are available as model cards[81] at [http://mol.dev/](http://mol.dev/).

## 4 Results

In order to evaluate the performance of our model using deep ensembles, two baseline models were created: (\(i\)) an XGBoost Random Forest (RF) model using the 17 descriptors available on AqSolDB plus 1809 molecular descriptors calculated by PaDELPy, a python wrapper for the PaDEL-Descriptor[82] software, and (\(ii\)) a model with the same architecture used on our deep ensemble using RMSE as the loss function and no ensemble (referred to as DNN). In addition, we evaluate the effects of (\(i\)) the bi-RNN layer (either GRU or LSTM), (\(ii\)) using an augmented dataset to train, (\(iii\)) the adversarial training, and (\(iv\)) the ensemble size in the model's performance. Table 1 shows the performance of each one of our trained models.

### Gated layer

The most common RNN layers are the GRU and the LSTM. GRU layers use two gates, reset and update, to control the cell's internal state. On the other hand, LSTM layers use three gates: forget, input, and output, with the same objective. Available studies compare GRU and LSTM performances in RNNs for different applications, for instance: forecasting[83], cryptocurrency[84, 85], wind speed[86, 87], condition of a paper press[88], motive classification in

\begin{table}
\begin{tabular}{l|c c c|c c c|c c c}  & \multicolumn{2}{c|}{Solubility Challenge 1} & \multicolumn{2}{c|}{Solubility Challenge 2\_1} & \multicolumn{2}{c}{Solubility Challenge 2\_2} \\ Model & RMSE & MAE & r & RMSE & MAE & r & RMSE & MAE & r \\ \hline RF & 1.121 & 0.914 & 0.547 & **0.950** & **0.727** & **0.725** & **1.205** & **1.002** & **0.840** \\ DNN & 1.540 & 1.214 & 0.433 & 1.315 & 1.035 & 0.651 & 1.879 & 1.381 & 0.736 \\ DNN\({}_{Aug}\) & 1.261 & 1.007 & 0.453 & 1.371 & 1.085 & 0.453 & 2.189 & 1.710 & 0.386 \\ kde4\({}^{GRU}\) & 1.610 & 1.145 & 0.462 & 1.413 & 1.114 & 0.604 & 1.488 & 1.220 & 0.704 \\ kde4\({}^{LSTM}\) & 1.554 & 1.191 & 0.507 & 1.469 & 1.188 & 0.650 & 1.523 & 1.161 & 0.706 \\ kde4\({}^{GRU}\)-NoAdv & 1.729 & 1.348 & 0.525 & 1.483 & 1.235 & 0.622 & 1.954 & 1.599 & 0.517 \\ kde4\({}^{LSTM}\)-NoAdv & 1.425 & 1.114 & 0.505 & 1.258 & 0.972 & 0.610 & 1.719 & 1.439 & 0.609 \\ kde4\({}^{GRU}_{Aug}\) & 1.329 & 1.148 & 0.426 & 1.354 & 1.157 & 0.674 & 1.626 & 1.340 & 0.623 \\ kde4\({}^{LSTM}_{Aug}\) & 1.273 & 0.984 & 0.473 & 1.137 & 0.932 & 0.639 & 1.511 & 1.128 & 0.717 \\ kde4\({}^{LSTM}_{Aug}\) & 1.247 & 0.984 & 0.542 & 1.044 & 0.846 & 0.701 & 1.418 & 1.118 & 0.729 \\ kde10\({}^{LSTM}_{Aug}\) & 1.689 & 1.437 & 0.471 & 1.451 & 1.238 & 0.676 & 1.599 & 1.405 & 0.699 \\ kde10\({}^{LSTM}_{Aug}\) & **1.095** & **0.843** & **0.559** & 0.983 & 0.793 & 0.724 & 1.263 & 1.051 & 0.792 \\ \end{tabular}
\end{table}
Table 1: Summary of the metrics for each trained model. We used the Root Mean Squared Error(RMSE(\(\downarrow\))), Mean Absolute Error (MAE(\(\downarrow\))), and Pearson correlation coefficient (r(\(\uparrow\))) to evaluate our models. The arrows indicate the direction of improvement. Deep ensemble models are referred to as “kde\(N\)”, where \(N\) is the ensemble size. Baseline models using random forest (RF) and the DNN model employed for deep ensemble (DNN) are also displayed. DNN model was trained as described in Section 3. The models in which data augmentation was used were subscribed with the flag \(Aug\). A superscript indicates if the bidirectional layer implements a \(GRU\) or a \(LSTM\) layer. In addition, models trained not using adversarial perturbation are flagged with “-NoAdv”. The columns show the results of each model evaluated on each solubility challenge dataset. 2\_1 represents the tight dataset (set-1), while 2\_2 represents the loose dataset (set-2) as described in the original paper (See Ref. 26). \(r\) stands for the Pearson correlation coefficient. The best-performing model in each dataset is displayed in bold.
thematic apperception tests[89] and music and raw speech[90]. Nevertheless, it is not clear which of those layers would perform better at a given task.

We trained models with four elements in the deep ensemble using GRU or LSTM. Metrics can be found in Table 1; for an explanation of the naming syntax used in this work, refer to Table 1 caption. Using LSTM resulted in a decrease in RMSE and MAE and an increase in the correlation coefficient, indicating better performance. For Solubility Challenges 1, 2_1, and 2_2, the kde4\({}^{GRU}_{Aug}\) model yielded RMSE values of 1.329, 1.354, and 1.626, respectively, while the kde4\({}^{LSTM}_{Aug}\) model achieved 1.049, 1.054, and 1.340, respectively. This trend was also observed for the models trained without data augmentation, but in a smaller proportion (See Table 1). Considering that LSTM performs better regarding this model and data, we will consider only bi-LSTM layers for further discussion. Those results are in accordance with our previous work[34] in which using LSTM helped improve the model's performance.

### Data augmentation

Our model is not intrinsically invariant with respect to the SELFIES representation input. For instance, both "C(C(C1C(=C(C(=O)O1)O)O)O" and "O=C1OC(C(O)CO)C(O)=C1O" are valid SMILES representations for the ascorbic acid (See Figure 1) that will be encoded for different SELFIES tokens. Hence, the model should learn to be invariant concerning changes in the string representation during training. It can be achieved by augmenting the dataset with SMILES randomization and training the model using different representations with the same label. Therefore, the model can learn relations in the chemical space instead of correlating the label with a specific representation.[62] With this aim, we evaluated the effects of augmenting the dataset by generating new randomized SMILES representations for each sample.

Augmenting the dataset had a significant impact on the metrics. It could be seen improvements of \(\sim 0.5\) in the RMSE when evaluating on challenge datasets 1 and 2_1, and a gain of \(\sim 0.2\) on 2_2 (See Table 1). Concerning the first two datasets, augmenting data improved every model used in this study. However, surprisingly, data augmentation led to a deprecation of the DNN model on the solubility challenge 2_2 dataset. This behavior was not further investigated.

### Adversarial training

Using adversarial training improved performance in Lakshminarayanan _et al.[36]_ studies. Hence, they suggested that it should be used in future applications of their deep learning algorithm. Thus, we tested the effects of adversarial perturbation on training models with ensemble sizes of 4 and 10.

Comparing kde4\({}^{LSTM}\)-NoAdv and kde4\({}^{LSTM}\), using adversarial training decreases model performance. It can be seen in Table 1 that using adversarial perturbation increased the RMSE from \(1.425\) to \(1.554\) and \(1.258\) to \(1.469\) in solubility challenges dataset 1 and 2_1, respectively. However, the RMSE decreased from \(1.719\) to \(1.523\) in dataset 2_2. Using adversarial perturbation affected our kde4\({}^{LSTM}\)'s performance by a change in RMSE of \(\pm 0.2\).

The inconsistent performance improvement observed when using adversarial training was further investigated with models in which the dataset was augmented. Due to the lack of multiple string representations in the training dataset, it is known that kde4\({}^{LSTM}\) may have generalization problems. A generalization issue could direct the adversarial perturbation in a non-physical direction because the model does not have complete knowledge about the chemical representation space. This hypothesis is reinforced when we compare kde10\({}^{LSTM}_{Aug}\)-NoAdv and kde10\({}^{LSTM}_{Aug}\). When using adversarial training on a model trained with an augmented dataset, the performance improvement is more evident (\(\sim 0.5\)) and consistent for all the test datasets.

### Deep ensemble size

To investigate the effects of increasing the ensemble size, we trained models with an ensemble of 4, 8, and 10 models. Given the previous results, these models used LSTM as the bi-RNN layer and were trained on the augmented dataset. Specifically for the solubility challenge 2_2, the most complex set to predict, these models presented an RMSE of \(1.511\), \(1.418\), and \(1.263\), respectively. Therefore, increasing the ensemble size consistently improved performance. We also observed this improvement on the other datasets (See Table 1).

Besides the immediate improvement in RMSE, increasing the ensemble size also improves the uncertainty of the model. Figure 2 shows the density distribution of the aleatoric variance and the epistemic variance (respectively related to AU and EU) for kde4\({}^{LSTM}_{Aug}\) (top 6 panels) and kde10\({}^{LSTM}_{Aug}\) (bottom six panels).

The increase in ensemble size led to a decrease in both uncertainties. AU distributions for the kde4\({}^{LSTM}_{Aug}\) are centered around 4 logS\({}^{2}\), displaying a long tail that extends to values as high as 20 logS\({}^{2}\) in the worst case (solubility challenge
2_2). A similar trend is observed in EU distributions. On the other hand, the kde10\({}^{LSTM}_{Aug}\) model results in narrower distributions. The mean of these distributions remains relatively unchanged, but a noticeable reduction in the extent of their tails can be observed. AU distribution ends in values around 10 logS\({}^{2}\).

## 5 Discussion

After extensively investigating the hyperparameter selection, we compared our model with available state-of-art models from the literature. Performance metrics on the solubility challenge datasets can be found in Table 2. Parity plots for our chosen models are presented in Figure 3.

Focusing on the solubility challenge 1 dataset[25], kde10\({}^{LSTM}_{Aug}\) is only \(\sim 0.2\) RMSE units worse than the best model available in the literature[41]. The RMSE of the participants of the challenge was not reported.[27] The primary metric used to evaluate models was the percentage of predictions within an error of 0.5 LogS units (called \(\pm 0.5\)log\(\%\)). Computing the same metric, kde10\({}^{LSTM}_{Aug}\) has a percentage of correct prediction of 44.4%. This result would place our model among the 35% best participants. The participant with the best performance presented a \(\pm 0.5\)log\(\%\) of 60.7%.

The architecture of the models was not published in the findings of the first challenge.[27] Nevertheless, the findings for the second challenge[28] investigated the participants more thoroughly. Participants were asked to identify their models' architecture and descriptors used. The challenge is divided into two datasets. Set-1 contains LogS values with an average interlaboratory reproducibility of 0.17 LogS. Our kde10\({}^{LSTM}_{Aug}\) achieve an RMSE of 0.983 and a \(\pm 0.5\)log\(\%\) of 40.0% in this dataset. Therefore, our model performs better than 26% of the published RMSE values and 50% of the \(\pm 0.5\)log\(\%\). In addition, the model with the best performance is an artificial neural network (ANN) that correctly predicted 61% (\(\pm 0.5\)log\(\%\)) of the molecule's LogS using a combination of molecule descriptors and fingerprints. The

Figure 2: Density distribution of the aleatoric (AU) and epistemic variances (EU) for the: (\(i\)) kde4\({}^{LSTM}_{Aug}\) (top six panels) and (\(ii\)) kde10\({}^{LSTM}_{Aug}\) (bottom six panels). Increasing ensemble size reduces the extent of the distribution’s tail, decreasing uncertainty about predictions. However, the ensemble size does not noticeably affect the distribution center.
second dataset (set-2) contains molecules whose solubility measurements are more challenging, reporting an average error in reproducibility of 0.62 LogS. The kde10\({}^{LSTM}_{Aug}\) achieves an RMSE of 1.263 and a \(\pm 0.5\)log\(\%\) of 23.3%. It performs better than 82% of the candidates when considering the RMSE. Surprisingly, \(\pm 0.5\)log\(\%\) does not follow this outstanding performance, which is more significant than only 32% Regarding the literature, kde10\({}^{LSTM}_{Aug}\) has an RMSE only \(\sim 0.1\) higher than a GNN that used an extensive set of numeric and one-hot descriptors in their feature vector.[50] Our model performs better than a transformer model that uses SMILES-string and an adjacency matrix and inputs.[44] The performance of those models is available in Table 2.

Notably, all participants in the solubility challenge 2 submitted a kind of QSPR or descriptor-based ML model. Using descriptors provides an easy way to ensure model invariance concerning molecule representation and is more informative since they can be physical quantities. However, selecting appropriate descriptors is crucial for developing descriptor-based ML models. It often requires specialists with a strong intuition about the relevant physical and chemical properties for predicting the target quantity. Our approach, on the other hand, is based on extracting information from simple string representations, a more straightforward raw data. Furthermore, we could achieve state-of-the-art performance while balancing the model size and complexity and using a raw input (a simple string). This simplified usage enables running the model on devices with limited computing power.

Lastly, transformers models have been used to address the issue of accurately predicting the solubility of small compounds. The typical workflow for transformers involves pre-training the model using a large dataset and subsequently fine-tuning it for a specific downstream task using a smaller dataset. Most existing models were either pre-trained on the ESOL[39] dataset or pre-trained on a larger dataset and fine-tuned using ESOL. Hence, the generalizability of those models cannot be verified. In a study by Francoeur and Koes [44], they considered two versions of their model, SolTransNet. The first version of SolTransNet was trained with the ESOL dataset using random splits. This approach achieved an RMSE of 0.278. Subsequently, the deployed version of SolTransNet was trained with the AqSolDB[1]. When ESOL was used to evaluate their deployed version, the model presented an RMSE of 2.99. While our model achieved an RMSE of 1.316 on ESOL, outperforming the SolTransNet deployed version, it cannot be compared with other models trained on ESOL.

\begin{table}
\begin{tabular}{l|c c|c c|c c|c c}  & \multicolumn{2}{c|}{SolChal1} & \multicolumn{2}{c|}{SolChal2\_1} & \multicolumn{2}{c|}{SolChal2\_2} & \multicolumn{2}{c}{ESOL} \\ Model & RMSE & MAE & RMSE & MAE & RMSE & MAE & RMSE & MAE \\ \hline RF & 1.121 & 0.914 & **0.950** & **0.727** & **1.205** & **1.002** & & \\ DNN & 1.540 & 1.214 & 1.315 & 1.035 & 1.879 & 1.381 & & \\ DNN\({}_{Aug}\) & 1.261 & 1.007 & 1.371 & 1.085 & 2.189 & 1.710 & & \\ kde\({}^{LSTM}_{Aug}\) & 1.273 & 0.984 & 1.137 & 0.932 & 1.511 & 1.128 & 1.397 & 1.131 \\ kde\({}^{LSTM}_{Aug}\) & 1.247 & 0.984 & 1.044 & 0.846 & 1.418 & 1.118 & 1.676 & 1.339 \\ kde10\({}^{LSTM}_{Aug}\) & **1.095** & **0.843** & 0.983 & 0.793 & 1.263 & 1.051 & **1.316** & **1.089** \\ \hline Linear regression[39] & & & & & & & **0.75** & \\ UG-RNN[41] & 0.90 & **0.74** & & & & & & \\ RF[37] & 0.93 & & & & & & & \\ Consensus[55] & **0.91** & & & & & & & \\ GNN[50] & \(\sim 1.10\) & **0.91** & & **1.17** & & & \\ SolyBert[91] & 0.925 & & & & & & \\ SolTransNet\({}^{a}\)[44] & & & 1.004 & 1.295 & 2.99 & & \\ SMILES-BERT\({}^{b}\)[92] & & & & & & & 0.47 \\ MolBERT\({}^{b}\)[48] & & & & & & & 0.531 \\ RT\({}^{b}\)[45] & & & & & & & 0.73 \\ MolFormer\({}^{b}\)[49] & & & & & & & **0.278** \\ \end{tabular}
\end{table}
Table 2: Metrics for the best models found in the current study (upper section) and for other state-of-art models available in the literature (lower section). Values were taken from the cited references. Missing values stand for entries that the cited authors did not study. SolChal columns stand for the Solubility Challenges. 2_1 represents the tight dataset (set-1), while 2_2 represents the loose dataset (set-2) as described in the original paper (See Ref. 26). The best-performing model in each dataset has its RMSE value in bold. \({}^{a}\) Has overlap between training and test sets. \({}^{b}\) Pre-trained model was fine-tuned on ESOL.
## 6 Conclusions

Our model was able to predict LogS values directly from SMILES or SELFIES string representations. Hence, there is no need for descriptors selection and construction. Using only raw data, our model could match state-of-art performance in datasets that are challenging to predict accurately.

In addition, carefully compromising between performance and complexity, we implemented a web application using TensorFlow JS. This application can satisfactorily run on any device with limited computational resources, such as laptops and smartphones. This excludes the need to rely on a server to run the application, improving usability and flexibility and decreasing implementation costs.

## 7 Data and code availability

All code needed to reproduce those results is publicly available on the following GitHub repository: [https://github.com/ur-whitelab/mol.dev](https://github.com/ur-whitelab/mol.dev). The model is also publicly accessible at the following address: [https://mol.mol.dev/](https://mol.mol.dev/).

Figure 3: Parity plots for two selected models being evaluated on the solubility challenge datasets: (\(i\)) kde\({}^{LSTM}_{Aug}\) (top row), and (\(ii\)) kde\({}^{LSTM}_{Aug}\) (bottom row). The left, middle, and right columns show the parity plots for solubility challenge 1[25], 2-set1, and 2-set2[26], respectively. Pearson correlation coefficient is displayed together with RMSE and MAE. “acc-0.5” stands for the \(\pm 0.5\mathrm{log}\%\) metric. Red dashed lines show the limits for molecules considered a correct prediction when computing the \(\pm 0.5\mathrm{log}\%\). The correlation between predicted values and labels increases when more models are added to the ensemble. RMSE and MAE also follow this pattern. However, the \(\pm 0.5\mathrm{log}\%\) decreases in set-2 of the second solubility challenge dataset (SolChal2-set2). While kde10\({}^{LSTM}_{Aug}\) improved the prediction of molecules that were being poorly predicted by kde\({}^{LSTM}_{Aug}\), the prediction of molecules with smaller errors was not greatly improved.