# Explainability Techniques for Chemical Language Models

Stefan Hödl
Physical Organic Chemistry, Radboud
<EMAIL_ADDRESS>
&William Robinson
Physical Organic Chemistry, Radboud
<EMAIL_ADDRESS>
&Yoram Bachrach
Google DeepMind
<EMAIL_ADDRESS>
Wilhelm Huck
Physical Organic Chemistry, Radboud
<EMAIL_ADDRESS>
&Tal Kachman
Artificial Intelligence, Donders Institute, Radboud
<EMAIL_ADDRESS>

###### Abstract

Explainability techniques are crucial in gaining insights into the reasons
behind the predictions of deep learning models, which have not yet been
applied to chemical language models. We propose an explainable AI technique
that attributes the importance of individual atoms towards the predictions
made by these models. Our method backpropagates the relevance information
towards the chemical input string and visualizes the importance of individual
atoms. We focus on self-attention Transformers operating on molecular string
representations and leverage a pretrained encoder for finetuning. We showcase
the method by predicting and visualizing solubility in water and organic
solvents. We achieve competitive model performance while obtaining
interpretable predictions, which we use to inspect the pretrained model.

Figure 1: Schematic visualization of the explainability technique. The string
representation “OCC(O)=O” of the molecule “glycolic acid” is the input to the
chemical language model, which predicts the property of interest. The
attention scores are extracted from all 6 Transformer layers (shown as stacked
heatmaps). These attention scores are aggregated throughout the architecture
following simple update rules to obtain a relevance weight for each token in
the input string. The molecules’ structure diagram is visualized with green
highlights, where high color intensity corresponds to high attributed impact
of an individual atom towards the predicted property.

##  1 Introduction

Transformers [1, 2] have achieved state-of-the-art results in vision [3],
language [4, 5] and life sciences [6, 7]. However, they are “black boxes” with
millions of parameters: despite their excellent performance, they cannot
provide explanations or human interpretable reasons behind their choices. This
explainability gap severely limits their adoption and usage by practitioners,
with models of high predictive accuracy sometimes breaking down in real world
physical experiments [8, 9].

Explainable AI (XAI) methods [10] are techniques that address this gap by
visualizing, explaining or validating the predictions of machine learning
models. Visualizations of the importance of inputs or features to the model’s
predictions provide insights as to the reasons behind the predictions that the
model makes, and can increase trust in the model’s outputs, especially when
these explanations correspond with human and scientific understanding of the
task [11].

In the domain of chemistry, a set of molecules might have very similar
properties despite diverse structural features, while the exchange of one
functional group might dictate whether a molecule dissolves. Attributing the
importance of individual atoms or functional groups towards the prediction can
help chemists understand model predictions and leverage their intuition.
Various XAI techniques have been applied to molecules, frequently based on
evaluating the importance by perturbation of the input, such as the model-
agnostic SHAP method based on Shapley values [11, 12, 13, 14].

Molecules are typically represented as molecular graphs or with subgraph-based
“fingerprints” [15, 16], significantly constraining the types of machine
learning models and explainability techniques that can be applied. Specialized
XAI techniques have been proposed [17, 18] for graph neural networks (GNN)
[19, 20] which represent molecules as molecular graphs, but these techniques
may be ill-suited to explain models based on other representations. Molecules
can also be represented as strings of letters and symbols, which are easily
machine-readable but difficult to decode for humans. One recent powerful tool
for training chemical prediction models is applying Transformers to very large
databases of such molecular strings, commonly known as “chemical language
models” [21, 22]. These models have shown impressive results on property
prediction tasks [23] and are especially promising when leveraged to generate
novel molecules with specific properties [24, 25, 26]. However, explainability
techniques for models based on molecular graphs or fingerprints are not
applicable to chemical language models.

One technique specific to Transformer architectures is applying heatmap
visualizations of the attention scores, typically extracted from the final
attention layer in the architecture [27]. This approach is limited as only the
attention scores are considered, ignoring the previous layers’ context and
other components of the architecture. Recent approaches overcome this
limitation by considering all Transformer components and aggregating
importance throughout the layers while retaining context [28, 29]. While
attention-specific visualization techniques have been successfully used in
computer vision tasks, they have not yet been adapted to the molecular domain.

Our contribution. Inspired by recent work on Transformer-based image
classification [29, 30], we propose a framework for regression tasks operating
on molecular strings, and develop an explainable AI technique for chemical
language models, using solely the model without external tools or information.
The attention visualization technique we propose for chemical language models
allows identifying the individual atoms whose presence had the biggest impact
on the output of a model predicting properties of the molecule. We leverage a
Transformer pretrained on molecular strings for finetuning on experimental
chemical property measurements and achieve competitive predictive performance.
This task-agnostic approach can easily be applied to all molecular regression
and classification tasks. In our approach, the relative impact that individual
units have on the model’s prediction are attributed to the input string and
visualized.

We demonstrate the technique by predicting solubility of small molecules and
assessing the model’s predictions through visualizing attributed relevance.
Our demonstration is based on a well-studied task where large datasets are
available, and where intuition can play a crucial role in validating the
visualization explainability results. After visualizing relevance for aqueous
solubility, we study solubility in organic solvents. Our technique is directly
applicable to regression and classification tasks using self-attention
Transformer architectures on molecular string representations. It can be used
following model training and only requires saving the attention scores and
calculating their gradients. Similarly, our method can be applied to other
Transformer variants [2], such as encoder-decoder architectures leveraging
cross-attention, by applying modifications analogous to similar extensions in
computer vision [30]. We fully open-source our methodology for the scientific
community.

##  2 Related work

#### Explainable AI.

“Extrinsic” XAI techniques are methods to visualize predictions of an already
trained model, frequently based on attribution of the prediction to the
model’s inputs [11]. One example are salience maps, which produce local
explanations of how changes in the input influence the output, measuring the
squared norm of the model’s gradient with respect to the input [31]. Many
modern approaches are decomposition-based techniques leveraging gradients and
propagation, such as layerwise relevance propagation [32], GradCAM [33],
attention rollout [28], gradient*input and DeepLIFT [34, 35]. All of these
were initially proposed for computer vision classification tasks.

An extensive overview of current explainability techniques for molecules is
provided by Wellawatte et al. [11]. Due to their simplicity, saliency maps are
commonly used in the chemical domain, for example to explain machine learning
models based on molecular fingerprints [36, 18, 25]. Crippen’s $\log(P)$
prediction is an additive approach which visualizes the contribution from each
atom using experimental fits, a good baseline for heatmap approaches [37, 38].

However, many assumptions of these baseline XAI techniques, such as SHAP’s
additive features and fair coalitions, are not valid for chemical language
models because they operate on molecular string representations. Omission of
some tokens frequently leads to invalid molecules, as these tokens directly
represent the molecules’ physical structure. Determining valid substitution
rules is difficult, and many valid substitutions would significantly change
the predicted property. SHAP can be applied to natural language, but its
application to molecular string tokens is analogous to masking out random
letters of a single word and evaluating importance of each letter.

#### Attribution of relevance.

A related work is layerwise relevance propagation, a method designed for
convolutional graph neural networks [17]. The authors use a solubility
regression task, but focus on adapting their approach to operate on molecular
graphs and convolutional filters, where atoms and edges are represented
explicitly and can thus be assigned relevance directly [38].

Theoretical justifications for layerwise relevance propagation are provided in
the work of Montavon et al. [39], which proposes the deep Taylor
decomposition. Heatmaps are defined as “consistent” if they fulfil
conservation of relevance and yield only positive values without negative
relevance. The authors define a training-free relevance model for
architectures using ReLU nonlinearities which have positive activations. This
allows them to propagate relevancies in higher layers in proportion to the
model’s activations and decompose the architecture layer by layer to the input
[39].

#### Molecular string representations.

There are many ways to represent a molecule as an input to a machine learning
model, which has a high impact on the predictive performance and ease of use
by chemists [40]. One prominent such language is SMILES [41], a ubiquitous
machine-readable format representing molecules as strings, which consists of
letters, symbols, digits and brackets, to denote atoms, bonds, rings and
branches. SMILES has been in use since the 1980s, and offers a language that
is flexible yet easy to understand and communicate [42]. A SMILES
representation can be derived algorithmically through a depth-first graph
traversal and iterative breaking and labeling of cycles and branches. Multiple
valid SMILES strings can be obtained through different numbering choices, thus
a “canonical” version was developed to uniquely identify molecules. However,
invalid SMILES strings can easily be constructed, which is addressed by the
SELFIES [43] notation frequently applied in generative settings. Transformers
can take a SMILES string directly as input after the tokenizer splits the
string into its constituent tokens.

#### Solubility prediction.

We demonstrate and analyze our approach using the regression task of
predicting the solubility of small molecules, a difficult but well-studied
problem crucial to many applications [44]. Many approaches have been proposed
to predict solubility, ranging from simple additive contributions using
empirical measurements to elaborate deep learning architectures. The property
of interest is the log-scaled solubility $\log(S)$, the solid-solution
partitioning coefficient, in units of [mol/L] [45]. While datasets of
$\log(S)$ calculations are available, a crucial limiting factor is the
availability of experimental solubility measurement datasets. Aqueous
solubility datasets are aggregated, such as ESOL ($\sim$3K) [46] and AqSolDB
($\sim$10K) [47]. SolProp is the largest available solubility data collection,
which aggregates experimental measurements for solubility in water
(AqueuosSolu, $\sim$12K) as well as solid solubility in organic solvents
(CombiSolu-Exp, $\sim$5K) [48].

The SolProp authors further propose their own predictive approaches based on
GNNs to predict solubility in water at room temperature
$\log(S_{\mathrm{aq,298K}})$. For solubility in an organic solvent $X$ at
temperature $T$, GNN predictions of the free energy $\Delta G_{\mathrm{solv}}$
and enthalpy $\Delta H_{\mathrm{solv}}$ of the thermodynamic cycle as well as
the $\log(S_{\mathrm{aq,298K}})$ prediction are combined to infer
$\log(S_{X,T})$. All GNN models are pretrained on computational datasets of
$\Delta G_{\mathrm{solv}}$, $\Delta H_{\mathrm{solv}}$ where available and
finetuned on experimental datasets. Particularly interesting is the ability of
their approach to incorporate available experimental measurements to
drastically improve the predictive performance [48].

#### Transformers for solubility.

Transformer-based architectures have been explored to directly predict aqueous
solubility, among them SolTranNet [49] which adapts the
MoleculeAttentionTransformer [50] architecture for the AqSolDB dataset.

Another recent work pretrains a SMILES language model from scratch, focusing
on directly predicting the solvation free energy as well as solubility in
organic solvents. [51] Similar to SolProp, most of their models are pretrained
on the computational datasets, but only the structures and not the calculated
properties are used during self-supervised pretraining via masking. Models
trained on multiple combinations of pretraining and finetuning datasets are
presented, but no results for the solubility datasets of the SolProp
collection are available for comparison (AqueousSolu, CombiSolu-Exp).

##  3 Methods

#### Chemical language model.

As the investigated Transformer model, we use the encoder of MegaMolBART [52,
53], a self-attention architecture based on the Chemformer [23]. Formal
algorithms for Transformers are presented in [2]. We consider an encoder
Transformer, and more specifically a bidirectional self-attention Transformer.
Within an attention layer, $\mathbf{Q}$ and $\mathbf{K}$ are the query and key
projections, their element-wise dot product is scaled by $\sqrt{d_{k}}$ for
numerical stability. The attention scores $\mathbf{A}$ are calculated as

$\mathbf{A}=\text{softmax}\left(\frac{\mathbf{Q}\mathbf{K}^{T}}{\sqrt{d}}\right).$
(1)

The Chemformer is pretrained in a self-supervised fashion (BART [54]) using a
reconstruction loss, which is obtained by masking random parts of the input
tokens with the <MASK> token. The model thus learns to reconstruct the
original SMILES string, which requires only strings but no labels as input and
is crucial to achieve scale. The primary difference between the Chemformer and
MegaMolBART is scale, where MegaMolBART was pretrained on approximately 1.45
billion small molecules from the ZINC-15 database [55] while keeping the size
of the model constant.

Only the encoder is used for regression or classification tasks, which can be
viewed as leveraging the pretrained model as a black box to extract expressive
latent features. A small hierarchical regression head with dense layers,
nonlinearities and layer normalization takes the average of these feature
vectors from the encoder and transforms it into a scalar prediction of the
property of interest. The model is subsequently finetuned on the prediction
task of interest using a dataset of experimental measurements (transfer
learning). This enables the model to rapidly converge and achieve competitive
accuracy even for small datasets using very few computational resources
compared to training from scratch. In principle, only the $\sim$40K parameters
of the regression head need to be trained, while the $\sim$20M parameters of
the pretrained encoder with 6 layers can be frozen.

#### <REG> tokenization.

A tokenizer splits the molecules’ SMILES string into a list of tokens of
variable length using regular expressions. These tokens are the input to the
encoder, which outputs a latent representation for every token. A pooling step
is required to obtain a feature vector of static shape for the regression
head, usually the average or sum of all tokens’ latent representations is
used. However, this pooling step hinders attribution of importance through
backpropagation of the prediction to the input string, since the contribution
of each token is lost in the average or sum. Due to this pooling, saliency
maps of the encoders’ output show uniform relevance for all tokens.

We avoid this pooling step by prepending a <REG> token to the tokenized SMILES
string, which aggregates information throughout the encoder. Only the feature
vector of this <REG> token is processed in the regression head, which allows
us to obtain the relevance of all other token to the <REG> token from the
pairwise attention weights and backpropagate gradients from the regression
prediction to the input tokens. The vocabulary of the MegaMolBART tokenizer
thus needs to be extended with a <REG> token corresponding to the <CLS> token
frequently encountered in the domain of computer vision and natural language
processing for classification tasks [56, 54, 3].

#### Relevance aggregation.

Following Chefer et al. [30], for self-attention Transformers the relevancy
matrix $\mathbf{R}$ is initialized as the identity matrix
$\mathbf{R}^{0}=I^{t\times t}$ for a molecule with $t$ tokens, representing
self-contained relevance of each token. The attention scores $\mathbf{A}^{l}$
of each Transformer layer $l\in L$ are saved during the models’ forward call,
while the prediction’s gradients $\nabla\mathbf{A}^{l}:=\frac{\partial
y}{\partial A^{l}}$ are evaluated and saved during the backward call. In
multi-head attention, each attention head captures different aspects of the
task, the importance of each attention head can be quantified from its
gradient $\nabla\mathbf{A}$. To account for these differences, the attention
scores $\mathbf{A}$ are multiplied by their gradient before averaging over all
8 attention heads ($\mathbb{E}_{h}$). Only the positive importance
$(\cdot)^{+}$ is considered to reflect nonlinearities following the deep
Tayler decomposition theory, $\odot$ denotes the element-wise product.

$\bar{\mathbf{A}}^{l}=\mathbb{E}_{h}\big{(}(\nabla\mathbf{A}^{l}\odot\mathbf{A}^{l})^{+}\big{)}$
(2)

The previous layers’ relevance matrix $\mathbf{R}^{l-1}$ contextualizes the
current layers attention mechanism, thus it is matrix multiplied with the
layer’s contribution $\bar{\mathbf{A}}^{l}$ and added element-wise to
$\mathbf{R}^{l-1}$ reflecting skip connections. Relevance is aggregated
throughout all layers $l\in\\{1,2,\dots,6\\}$ with the update rule

$\mathbf{R}^{l}=\mathbf{R}^{l-1}+\bar{\mathbf{A}}^{l}\cdot\mathbf{R}^{l-1}.$
(3)

The feature vector of the <REG> token encapsulates all information
contributing to the prediction, and the relevance of each token to the <REG>
token is captured in the corresponding entry in the contextualized relevancy
matrix $\mathbf{R}$. As a consequence, visualization of the final layers’
averaged attention matrix $\bar{\mathbf{A}}^{6}$ shows zeros everywhere but in
the this column.

Figure 2: 2d structure plot with a color map corresponding to attributed
relevance extracted from the <REG> token of the aggregated relevancy matrix
$\mathbf{R}$. The relevance of each atom token towards the model’s aqueous
solubility prediction is visualized, the intensity of the green highlights
corresponds to the attributed relevance. The color bar shows relative
importance in arbitrary units.

#### Relevance visualization.

All atom token relevance weights are then visualized on the molecules’ 2D
structure diagram to aid human intuition, as understanding SMILES strings
intuitively is difficult even for experienced practitioners. Once the
attributed importance weights are obtained, the molecule is visualized by
plotting a 2d structure diagram with a corresponding heatmap showing the
attributed relevance to the prediction (Fig. 2).

##  4 Experiments

We first present prediction errors from the finetuned MegaMolBART encoder
model for aqueous solubility and compare them to results from other approaches
available in the literature. We then turn to examining solubility in organic
solvents. After establishing the models’ accuracy, we obtain the internal
representations to visualize and inspect the attributed importance for a range
of molecules.

#### Aqueous solubility prediction.

The results below show that when predicting aqueous solubility
$\log(S_{\mathrm{aq,298K}})$, the predictive accuracy of the finetuned
MegaMolBART model is competitive, but slightly worse than the corresponding
SolProp GNN model. The SolProp authors evaluate their model on a randomly
selected as well as a low-uncertainty test set of the AqueousSolu dataset, as
significant uncertainty in experimental measurements is introduced due to
aggregation from diverse literature sources and multiple datasets. The model
errors of the finetuned MegaMolBART model fall behind the specialized SolProp
GNN model on both the random and the accurate test set. The SolProp authors
also evaluate ALOGpS [57] and SolTranNet [49] on the accurate test set (Table
3).

Figure 3: Parity plots showing the predictive accuracy of the MegaMolBART model finetuned on the AqueousSolu dataset. The left plot (blue) shows a (10%) randomly selected test set, while the right plot (red) shows the low-uncertainty test set [48]. The experimental solubility measurements $\log(S_{\mathrm{aq,298K}})$ (X axes) are plotted against the models’ predictions (Y axes), deviations from the diagonal represent the models’ errors. The predictions are grouped into hexagonal bins to highlight the density of predictions, the raw densities are shown as histograms on the sides of both plots. The grey lines show a linear regression fit and its “slope” with $n$ samples in the test set. The run closest to the average of 5 runs with different seeds for the dataset split is shown. Model comparison for aqueous solubility Model | Random | Random | Accurate | Accurate
---|---|---|---|---
| MAE | RMSE | MAE | RMSE
MegaMolBART (AVG) | 0.604 | 0.895 | 0.475 | 0.638
MegaMolBART (SE) | (0.004) | (0.009) | (0.006) | (0.006)
SolProp | 0.49 | 0.75 | 0.34 | 0.49
ALOGpS | - | - | 0.55 | 0.79
SolTranNet | - | - | 0.58 | 0.76
Table 1: Comparison of models trained on the AqueousSolu dataset for random
10% test sets (“Random”) and the low-uncertainty test set (“Accurate”). The
mean (AVG) and standard error (SE) of the mean absolute error (MAE) and root
mean squared error (RMSE) are reported, calculated from 5 runs with random
seeds affecting the split into both train and validation sets and the random
test set.

#### Attention heatmap visualization.

Fig. 4 shows heatmaps of the first two layers’ averaged attention scores
$\bar{\mathbf{A}}$ to illustrate the aggregation rules. The pairwise attention
between each token can be seen, where light colors reflect large attention
scores. Regions of high importance can be seen near the diagonal of the first
heatmap $\bar{\mathbf{A}}^{1}$, which captures the immediate neighbourhood of
each atom. Brackets and digits denote side-branches and rings, which can lead
to close physical proximity of two atoms despite large separation in the
SMILES string. High attention scores can be seen for ring openings and
closures, in particular for the carbon atoms in close physical proximity. The
final aggregated heatmap $\mathbf{R}$ is plotted on the right side after
removal of the identity matrix $I$ to account for differences in scale. The
<REG> token’s column of $\mathbf{R}$ is extracted and visualized.

Figure 4: Visualization of the attention aggregation (Eq. 2 and 3). The
$\bar{\mathbf{A}}$ heatmap is aggregated throughout all 6 attention layers, of
which layer 1, 2 and the aggregated relevancy matrix $\mathbf{R}$ are shown
for the molecule “CC12CC(C)(NC(=O)N1)NC(=O)N2”. The relevancy matrix
$\mathbf{R}$ is obtained by starting with the identity matrix $I$ and applying
the update rules using the averaged attention matrix $\bar{\mathbf{A}}$
extracted from all 6 attention layers.

#### Solubility in organic solvents.

For solubility prediction in different solvents, the interactions between the
solute and the solvent need to be taken into account in the architecture. The
MegaMolBART model is finetuned on the CombiSolu-Exp dataset in the same
fashion as for aqueous solubility. The visualization technique can be applied
without further modifications as long as the task can be formulated within an
architecture based only on self-attention, which is achieved through the usage
of the <SEP> token to separate the SMILES string of the solute and solvent.
This simplifies the task as well as the visualizations of the solute and
solvent, as the model receives a single concatenated input to process, and
thus only a single relevancy matrix $\mathbf{R}$ needs to be constructed. The
tokenized input for glycolic acid ( OCC(O)=O ) shown in Fig. 1 in ethanol (
CCO ) as a solvent would then be

“ <REG>, O, C, C, (, O, ), =, O, <SEP>, C, C, O ”.

#### Cross-attention model variant.

An alternative approach is to model the interaction using cross-attention
layers, where the query and key vectors correspond to the latent
representations of the solute and solvent. Both are extracted from the
MegaMolBART’s encoder and averaged into two vectors of static shape
independent of the number of tokens in the solute and solvent. One or more
cross-attention layers are used to model the interaction between solute and
solvent, which then needs to be aggregated into a single feature vector for
the regression head. This is achieved through the <REG> token in the
aforementioned approach but does not work here due to the pooling operation
before the cross-attention layers. It contradicts the technique’s assumption
that the <REG> token in the aggregated relevance matrix $\mathbf{R}$ reflects
the relevance of each other token to the predicted property, which are no
longer explicitly represented per token after pooling. Separate attention
scores and gradients need to be stored to obtain relevancy matrices
$\mathbf{R}_{\mathrm{solute}}$ and $\mathbf{R}_{\mathrm{solvent}}$, which
requires extensions to the propagation rules following Chefer et al. [30].
Separate heatmap visualizations can be obtained from the cross-attention
layers, which might explicitly highlight interactions.

The predictive accuracy of both models is excellent, though a direct
comparison cannot be made due a big difference in approach and datasets used
to model solubility in organic solvents. The cross-attention architecture
variant achieves very accurate results but cannot be explained through the
<REG> token approach without further modifications to the explainability
technique. Since the SolProp authors use GNN predictions of $\Delta
G_{\mathrm{solv}}$, $\Delta H_{\mathrm{solv}}$, $\log(S_{\mathrm{aq,298K}})$
and thermodynamic cycle relationships to predict $\log(S_{X,T})$, their model
is not trained on the CombiSolu-Exp dataset itself. However, it does use the
GNN models trained on AqueousSolu and corresponding datasets of $\Delta
H_{\mathrm{solv}}$ and $\Delta G_{\mathrm{solv}}$. We refer to the SolProp
paper and supplementary information for results in various contexts [48]. This
implies that visualization techniques would need to consider all SolProp
architecture components. In contrast, both MegaMolBART models are trained
exclusively on the CombiSolu-Exp dataset.

Figure 5:  The left parity plot (green) shows the finetuned MegaMolBART
models’ predictive accuracy for solubility in organic solvents using the <SEP>
token to concatenate the solute and solvent tokens. It is evaluated on 5
randomly selected 10% test sets of the CombiSolu-Exp dataset, the best run’s
parity plot is shown. The average and standard error as well as the
corresponding parity plot for the cross-attention model evaluated on the
$\mathrm{T}=\mathrm{298K}$ test set from [48] can be found in the
supplementary information. The right visualization shows the attributed
relevance and corresponding color bar of the <REG> model for a solute in an
oil (toluene) as a solvent.

#### Solute $\times$ solvent visualization.

The relevancy matrix from the approach leveraging the <REG> and <SEP> tokens
is obtained in the same manner as for aqueous solubility by aggregating
relevancy throughout the layers for the concatenated input. The importance of
the solute and solvent towards the solubility prediction is visualized in Fig.
5. Different functional groups of a molecule contribute towards solubility
when the same solute is dissolved in distinct solvents, such as an oil
(toluene) and an alcohol (ethanol). However, when the models’ $\log(S)$
predictions are close to each other, the importance attributed to both solutes
is nearly indistinguishable despite distinct interactions contributing to
solubility. These interactions are not captured by the <REG> model and might
be the reason why the cross-attention model variant is able to outperform the
<REG> model on the $\mathrm{T}=\mathrm{298K}$ test set of the CombiSolu-Exp
dataset.

#### Discussion of results.

We worked with experimental organic chemists to verify our methodology by
showing them visualizations of the models’ predictions to understand the
nature and quality of the finetuned model. Our observations suggest the
MegaMolBART model learns to predict solubility accurately without fully
learning the relevant physical context. Atoms with high attributed relevance
generally correspond to relevant functional groups for aqueous solubility
(such as OH, =O or NH2), while explanations deteriorate for molecules with
very high or low solubility.

The model is unable to accurately model symmetry and frequently attributes
very different relevance to symmetric functional groups. Interestingly, the
model frequently attributes high relevance to the molecules’ carbon backbone,
which dictates the structure of the molecule but does not by itself influence
solubility. A possible explanation is that the model learns to predict
solubility using the carbon backbone instead of functional groups due to the
pretraining procedure of MegaMolBART.

Our results indicate that different finetuning approaches and regression head
variants produce quantitatively comparable results and qualitatively similar
visualizations due to the major influence of the pretrained MegaMolBART model.
Similarly, the <REG> approach and the average-pooling achieve very similar
accuracy on the AqueousSolu dataset, thus the <REG> token enables attribution
of relevance and visualizations without any loss in accuracy.

##  5 Limitations & future work

The presented explainability technique could be significantly improved by
visualizing relevance attributed to all tokens in a molecular string, instead
of only atom tokens. SMILES strings contain tokens which do not directly
correspond to atoms or bonds, and thus do not have an explicit correspondence
to structure diagrams. This can be achieved by constructing a mapping between
non-atom SMILES tokens and 2d structure plots, mapping relevance of non-atomic
tokens such as ring opening and closures, brackets and bonds. Crucially, no
ground truth relevancy labels exist, which makes validation of one approach
over alternatives difficult.

A naive approach would be to assign relevance attributed to bond tokens to the
bonding pair and distribute bracket and ring token relevancies to all
constituent atoms. Another possible approach to address this issue is to
initialize the relevancy matrix as 1 only for atom tokens, but this prevents
the propagation of relevance between tokens throughout the layers.
Alternatively, a sum over additive contributions for all atom tokens could be
extracted instead of the <REG> token, but this suffers from the same
explainability issues as with the pooling layer. While naively disregarding
all non-atom tokens biases the visualization of the attributed importance,
different variations of the approach similarly introduce biases and
assumptions. Plotting all explicitly represented bonds does not yield a more
accurate importance either but requires adapting SMILES strings and training
from scratch.

Visualizations as well as predictions of molecular properties should also be
symmetric and invariant to different representations of the same molecule.
Non-canonical SMILES string can be obtained by different choices of labeling
of a given molecule. The MegaMolBART model does not fulfil invariance to
randomized SMILES inputs even though SMILES randomization was used as data
augmentation during training [52]. Contrastive training approaches have been
proposed [58] to enforce SMILES invariance in the model, but this is outside
of the scope of this study due to computational efforts required to
pretraining the model from scratch.

Finally, the method is applicable to Transformer architectures beyond self-
attention, where generative architectures in particular would greatly benefit
from visualization techniques. However, further work and additional
modifications are necessary to apply the explainability technique, among them
the extraction and propagation of multiple relevancy matrices.

##  6 Conclusion

With the advent of large language models in chemistry, increasing trust in
their predictions is a crucial step towards incorporating deep learning model
into the laboratory. Visualization of relevance of molecules’ atoms and
functional groups enables practitioners to inspect the learned features of
chemical language models and validate their correspondence to scientific
understanding and intuition.

We propose an attribution technique for Transformers operating on molecular
string representation, which addresses the lack of attention-specific
explainability techniques for chemical language models. Our method propagates
importance throughout the models’ layers and attributes these to the input
SMILES string components using only the attention scores and their gradients.
The <REG> token encapsulates the relevance of each token towards the property
prediction, which is extracted and visualized.

We demonstrate the effectiveness of our technique by predicting and
visualizing solubility in water and organic solvents using a task-agnostic
approach, achieving competitive accuracy and explainable predictions. These
visualizations allow us to inspect what the chemical language model learns.
Our technique is applicable to any self-attention Transformer architectures
and molecular string representations, and can be extended to other Transformer
architectures and tasks.

## References

  * Vaswani et al. [2017] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin.  Attention is all you need.  In I. Guyon, U. Von Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett, editors, _Advances in Neural Information Processing Systems_ , volume 30. Curran Associates, Inc., 2017.  URL https://proceedings.neurips.cc/paper_files/paper/2017/file/3f5ee243547dee91fbd053c1c4a845aa-Paper.pdf.
  * Phuong and Hutter [2022] Mary Phuong and Marcus Hutter.  Formal algorithms for transformers.  7 2022.  URL http://arxiv.org/abs/2207.09238.
  * Dosovitskiy et al. [2020] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al.  An image is worth 16x16 words: Transformers for image recognition at scale.  _arXiv preprint arXiv:2010.11929_ , 2020.
  * Wolf et al. [2020] Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, Rémi Louf, Morgan Funtowicz, et al.  Transformers: State-of-the-art natural language processing.  In _Proceedings of the 2020 conference on empirical methods in natural language processing: system demonstrations_ , pages 38–45, 2020.
  * Otter et al. [2020] Daniel W Otter, Julian R Medina, and Jugal K Kalita.  A survey of the usages of deep learning for natural language processing.  _IEEE transactions on neural networks and learning systems_ , 32(2):604–624, 2020.
  * Alp [2021] Highly accurate protein structure prediction with alphafold.  _Nature_ , 596:583–589, 8 2021.  ISSN 14764687.  doi: 10.1038/s41586-021-03819-2.
  * Ferruz and Höcker [2022] Noelia Ferruz and Birte Höcker.  Controllable protein design with language models.  _arXiv preprint arXiv:2201.07338_ , 2022.
  * [8] Yoav Wald, Amir Feder, Daniel Greenfeld, and Uri Shalit.  On Calibration and Out-of-domain Generalization.  Technical report, arXiv.  URL http://arxiv.org/abs/2102.10395.  arXiv:2102.10395 [cs] type: article.
  * Hogewind et al. [2022] Yannick Hogewind, Thiago D Simao, Tal Kachman, and Nils Jansen.  Safe reinforcement learning from pixels using a stochastic latent representation.  _arXiv preprint arXiv:2210.01801_ , 2022.
  * Došilović et al. [2018] Filip Karlo Došilović, Mario Brčić, and Nikica Hlupić.  Explainable artificial intelligence: A survey.  In _2018 41st International convention on information and communication technology, electronics and microelectronics (MIPRO)_ , pages 0210–0215. IEEE, 2018.
  * Wellawatte et al. [2023] Geemi P. Wellawatte, Heta A. Gandhi, Aditi Seshadri, and Andrew D. White.  A perspective on explanations of molecular prediction models.  _Journal of Chemical Theory and Computation_ , 3 2023.  ISSN 1549-9618.  doi: 10.1021/acs.jctc.2c01235.  URL https://pubs.acs.org/doi/10.1021/acs.jctc.2c01235.
  * Lundberg and Lee [2017] Scott M Lundberg and Su-In Lee.  A unified approach to interpreting model predictions.  In I. Guyon, U. V. Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett, editors, _Advances in Neural Information Processing Systems 30_ , pages 4765–4774. Curran Associates, Inc., 2017.  URL http://papers.nips.cc/paper/7062-a-unified-approach-to-interpreting-model-predictions.pdf.
  * Cornelisse et al. [2022] Daphne Cornelisse, Thomas Rood, Yoram Bachrach, Mateusz Malinowski, and Tal Kachman.  Neural payoff machines: Predicting fair and stable payoff allocations among team members.  In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh, editors, _Advances in Neural Information Processing Systems_ , volume 35, pages 25491–25503. Curran Associates, Inc., 2022.  URL https://proceedings.neurips.cc/paper_files/paper/2022/file/a38df2dd882bf7059a1914dd5547af87-Paper-Conference.pdf.
  * [14] Roma Patel, Marta Garnelo, Ian Gemp, Chris Dyer, and Yoram Bachrach.  Game-theoretic Vocabulary Selection via the Shapley Value and Banzhaf Index.  URL https://aclanthology.org/2021.naacl-main.223.
  * Rogers and Hahn [2010] David Rogers and Mathew Hahn.  Extended-connectivity fingerprints.  _Journal of Chemical Information and Modeling_ , 50:742–754, 5 2010.  ISSN 1549960X.  doi: 10.1021/ci100050t.
  * [16] David Duvenaud, Dougal Maclaurin, Jorge Aguilera-Iparraguirre, Rafael Gómez-Bombarelli, Timothy Hirzel, Alán Aspuru-Guzik, and Ryan P. Adams.  Convolutional networks on graphs for learning molecular fingerprints.  In _Proceedings of the 28th International Conference on Neural Information Processing Systems - Volume 2_ , NIPS’15, pages 2224–2232. MIT Press.
  * [17] Federico Baldassarre and Hossein Azizpour.  Explainability Techniques for Graph Convolutional Networks.  doi: 10.48550/arXiv.1905.13686.  URL http://arxiv.org/abs/1905.13686.  arXiv:1905.13686 [cs, stat] type: article.
  * Weber et al. [2022] Jeffrey K. Weber, Joseph A. Morrone, Sugato Bagchi, Jan D.Estrada Pabon, Seung gu Kang, Leili Zhang, and Wendy D. Cornell.  Simplified, interpretable graph convolutional neural networks for small molecule activity prediction.  _Journal of Computer-Aided Molecular Design_ , 36:391–404, 5 2022.  ISSN 15734951.  doi: 10.1007/s10822-021-00421-6.
  * [19] Haiqi Zhang, Guangquan Lu, Mengmeng Zhan, and Beixian Zhang.  Semi-Supervised Classification of Graph Convolutional Networks with Laplacian Rank Constraints.  54(4):2645–2656.  ISSN 1573-773X.  doi: 10.1007/s11063-020-10404-7.  URL https://doi.org/10.1007/s11063-020-10404-7.
  * Bronstein et al. [2021] Michael M Bronstein, Joan Bruna, Taco Cohen, and Petar Veličković.  Geometric deep learning: Grids, groups, graphs, geodesics, and gauges.  _arXiv preprint arXiv:2104.13478_ , 2021.
  * Flam-Shepherd et al. [2022] Daniel Flam-Shepherd, Kevin Zhu, and Alán Aspuru-Guzik.  Language models can learn complex molecular distributions.  _Nature Communications_ , 13(1):3293, 2022.
  * Du et al. [2022] Yuanqi Du, Tianfan Fu, Jimeng Sun, and Shengchao Liu.  Molgensurvey: A systematic survey in machine learning models for molecule design.  _arXiv preprint arXiv:2203.14500_ , 2022.
  * [23] Ross Irwin, Spyridon Dimitriadis, Jiazhen He, and Esben Jannik Bjerrum.  Chemformer: a pre-trained transformer for computational chemistry.  3(1):015022.  ISSN 2632-2153.  doi: 10.1088/2632-2153/ac3ffb.  URL https://dx.doi.org/10.1088/2632-2153/ac3ffb.
  * Born and Manica [2022] Jannis Born and Matteo Manica.  Regression transformer: Concurrent sequence regression and generation for molecular language modeling.  2 2022.  URL http://arxiv.org/abs/2202.01338.  https://github.com/IBM/regression-transformer.
  * Bagal et al. [2021] Viraj Bagal, Rishal Aggarwal, P. K. Vinod, and U. Deva Priyakumar.  Molgpt: Molecular generation using a transformer-decoder model.  _Journal of Chemical Information and Modeling_ , 2021.  ISSN 15205142.  doi: 10.1021/acs.jcim.1c00600.
  * [26] Cynthia Shen, Mario Krenn, Sagi Eppel, and Alán Aspuru-Guzik.  Deep molecular dreaming: inverse machine learning for de-novo molecular design and interpretability with surjective representations.  2(3).  ISSN 2632-2153.  doi: 10.1088/2632-2153/ac09d6.  URL https://dx.doi.org/10.1088/2632-2153/ac09d6.
  * Schwaller et al. [2020] Philippe Schwaller, Daniel Probst, Alain C. Vaucher, Vishnu H. Nair, David Kreutter, Teodoro Laino, and Jean-Louis Reymond.  Mapping the space of chemical reactions using attention-based neural networks.  12 2020.  URL http://arxiv.org/abs/2012.06051.
  * Abnar and Zuidema [2020] Samira Abnar and Willem Zuidema.  Quantifying attention flow in transformers.  5 2020.  doi: 10.18653/v1/2020.acl-main.385.
  * Chefer et al. [2020] Hila Chefer, Shir Gur, and Lior Wolf.  Transformer interpretability beyond attention visualization.  12 2020.  doi: 10.1109/cvpr46437.2021.00084.
  * Chefer et al. [2021] Hila Chefer, Shir Gur, and Lior Wolf.  Generic attention-model explainability for interpreting bi-modal and encoder-decoder transformers.  3 2021.  doi: 10.1109/iccv48922.2021.00045.
  * Simonyan et al. [2013] Karen Simonyan, Andrea Vedaldi, and Andrew Zisserman.  Deep inside convolutional networks: Visualising image classification models and saliency maps.  12 2013.  URL http://arxiv.org/abs/1312.6034.
  * Bach et al. [2015] Sebastian Bach, Alexander Binder, Grégoire Montavon, Frederick Klauschen, Klaus Robert Müller, and Wojciech Samek.  On pixel-wise explanations for non-linear classifier decisions by layer-wise relevance propagation.  _PLoS ONE_ , 10, 7 2015.  ISSN 19326203.  doi: 10.1371/journal.pone.0130140.
  * [33] Ramprasaath R Selvaraju, Michael Cogswell, Abhishek Das, Ramakrishna Vedantam, Devi Parikh, and Dhruv Batra.  Grad-cam: Visual explanations from deep networks via gradient-based localization.  ISSN 2380-7504.  URL http://gradcam.cloudcv.org.  ISSN: 2380-7504.
  * Shrikumar et al. [2016] Avanti Shrikumar, Peyton Greenside, Anna Shcherbina, and Anshul Kundaje.  Not just a black box: Learning important features through propagating activation differences.  5 2016.  doi: 10.48550/arXiv.1605.01713.  URL http://arxiv.org/abs/1605.01713.
  * [35] Avanti Shrikumar, Peyton Greenside, and Anshul Kundaje.  Learning important features through propagating activation differences.
  * [36] Muriel Gevrey, Ioannis Dimopoulos, and Sovan Lek.  Review and comparison of methods to study the contribution of variables in artificial neural network models.  160(3):249–264.  ISSN 0304-3800.  doi: 10.1016/S0304-3800(02)00257-0.  URL https://www.sciencedirect.com/science/article/pii/S0304380002002570.
  * Wildman and Crippen [1999] Scott A. Wildman and Gordon M. Crippen.  Prediction of physicochemical parameters by atomic contributions.  _Journal of Chemical Information and Computer Sciences_ , 39:868–873, 1999.  ISSN 00952338.  doi: 10.1021/ci990307l.
  * Rasmussen et al. [2022] Maria H Rasmussen, Diana S Christensen, and Jan H Jensen.  Do machines dream of atoms? a quantitative molecular benchmark for explainable ai heatmaps, 2022.
  * Montavon et al. [2017] Grégoire Montavon, Sebastian Lapuschkin, Alexander Binder, Wojciech Samek, and Klaus Robert Müller.  Explaining nonlinear classification decisions with deep taylor decomposition.  _Pattern Recognition_ , 65:211–222, 5 2017.  ISSN 00313203.  doi: 10.1016/j.patcog.2016.11.008.
  * Wigh et al. [2022] Daniel S Wigh, Jonathan M Goodman, and Alexei A Lapkin.  A review of molecular representation in the age of machine learning.  _Wiley Interdisciplinary Reviews: Computational Molecular Science_ , 12(5):e1603, 2022.
  * [41] David Weininger.  SMILES, a chemical language and information system. 1. Introduction to methodology and encoding rules.  28(1):31–36.  ISSN 0095-2338.  doi: 10.1021/ci00057a005.  URL https://doi.org/10.1021/ci00057a005.
  * Weininger [1988] David Weininger.  Smiles, a chemical language and information system. 1. introduction to methodology and encoding rules.  _Journal of chemical information and computer sciences_ , 28(1):31–36, 1988.
  * Lo et al. [2023] Alston Lo, Robert Pollice, AkshatKumar Nigam, Andrew D. White, Mario Krenn, and Alán Aspuru-Guzik.  Recent advances in the self-referencing embedding strings (selfies) library.  2 2023.
  * Lusci et al. [2013] Alessandro Lusci, Gianluca Pollastri, and Pierre Baldi.  Deep architectures and deep learning in chemoinformatics: the prediction of aqueous solubility for drug-like molecules.  _Journal of chemical information and modeling_ , 53(7):1563–1575, 2013.
  * Atkins et al. [2014] Peter Atkins, Peter William Atkins, and Julio de Paula.  _Atkins’ physical chemistry_.  Oxford university press, 2014.
  * Delaney [2004] John S. Delaney.  Esol: Estimating aqueous solubility directly from molecular structure.  _Journal of Chemical Information and Computer Sciences_ , 44:1000–1005, 5 2004.  ISSN 00952338.  doi: 10.1021/ci034243x.
  * Sorkun et al. [2019] Murat Cihan Sorkun, Abhishek Khetan, and Süleyman Er.  Aqsoldb, a curated reference set of aqueous solubility and 2d descriptors for a diverse set of compounds.  _Scientific Data_ , 6, 12 2019.  ISSN 2052-4463.  doi: 10.1038/s41597-019-0151-1.
  * Vermeire et al. [2022] Florence H. Vermeire, Yunsie Chung, and William H. Green.  Predicting solubility limits of organic solutes for a wide range of solvents and temperatures.  _Journal of the American Chemical Society_ , 144:10785–10797, 6 2022.  ISSN 15205126.  doi: 10.1021/jacs.2c01768.
  * [49] Paul G. Francoeur and David R. Koes.  SolTranNet-A Machine Learning Tool for Fast Aqueous Solubility Prediction.  61(6):2530–2536.  ISSN 1549-960X.  doi: 10.1021/acs.jcim.1c00331.
  * [50] Łukasz Maziarka, Tomasz Danel, Sławomir Mucha, Krzysztof Rataj, Jacek Tabor, and Stanisław Jastrzębski.  Molecule Attention Transformer.  Technical report, arXiv.  URL http://arxiv.org/abs/2002.08264.  arXiv:2002.08264 [physics, stat] type: article.
  * Yu et al. [2023] Jiahui Yu, Chengwei Zhang, Yingying Cheng, Yun-Fang Yang, Yuan-Bin She, Fengfan Liu, Weike Su, and An Su.  Solvbert for solvation free energy and solubility prediction: a demonstration of an nlp model for predicting the properties of molecular complexes.  _Digital Discovery_ , 2023.  doi: 10.1039/d2dd00107a.
  * Nvidia [2022] Nvidia.  MegaMolBART.  _https://github.com/NVIDIA/MegaMolBART_ , 2022.
  * Fabian et al. [2020] Benedek Fabian, Thomas Edlich, Héléna Gaspar, Marwin Segler, Joshua Meyers, Marco Fiscato, and Mohamed Ahmed.  Molecular representation learning with language models and domain-relevant auxiliary tasks, 2020.
  * Lewis et al. [2019] Mike Lewis, Yinhan Liu, Naman Goyal, Marjan Ghazvininejad, Abdelrahman Mohamed, Omer Levy, Ves Stoyanov, and Luke Zettlemoyer.  Bart: Denoising sequence-to-sequence pre-training for natural language generation, translation, and comprehension.  10 2019.  URL http://arxiv.org/abs/1910.13461.
  * ZIN [2015] Zinc 15 - ligand discovery for everyone.  _Journal of Chemical Information and Modeling_ , 55:2324–2337, 11 2015.  ISSN 1549960X.  doi: 10.1021/acs.jcim.5b00559.
  * Devlin et al. [2018] Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova.  Bert: Pre-training of deep bidirectional transformers for language understanding.  10 2018.  URL http://arxiv.org/abs/1810.04805.
  * Tetko et al. [2005] Igor V. Tetko, Johann Gasteiger, Roberto Todeschini, Andrea Mauri, David Livingstone, Peter Ertl, Vladimir A. Palyulin, Eugene V. Radchenko, Nikolay S. Zefirov, Alexander S. Makarenko, Vsevolod Yu Tanchuk, and Volodymyr V. Prokopenko.  Virtual computational chemistry laboratory - design and description.  _Journal of Computer-Aided Molecular Design_ , 19:453–463, 2005.  ISSN 15734951.  doi: 10.1007/s10822-005-8694-y.
  * Shrivastava and Kell [2021] Aditya Divyakant Shrivastava and Douglas B. Kell.  Fragnet, a contrastive learning-based transformer model for clustering, interpreting, visualizing, and navigating chemical space.  _Molecules_ , 26, 4 2021.  ISSN 14203049.  doi: 10.3390/molecules26072065.

##  7 Supplementary Information

#### Hyperparameters.

Relatively few hyperparameters need to be selected due to the pretrained
MegaMolBART model. Most variations of the regression head produce very similar
results, thus a relatively small regression head of dimensionality 512 - 64 -
64 - 1 is chosen. This aligns with the MegaMolBART output of dimensionality
512 (8 heads of dimensionality 64), and produces similar accuracy to bigger
variants. Interleaved are ReLU nonlinearities as well as a LayerNorm layer
before the first projection. The regression head has 38.1K parameters, of
which the first dense layer (512 $\rightarrow$ 64) has 32.8K, the second (64
$\rightarrow$ 64) 4.2K, and the final (64 $\rightarrow$ 1) 65 parameters. The
LayerNorm adds another 1024 parameters. For comparison, the MegaMolBART
encoder has 19.8M parameters. For the CombiSolu-Exp dataset, the scalar
temperature is appended to the <REG> token feature vector for the regression
head. Solvent density is not used as is detrimental to the model’s accuracy.

The HuberLoss loss criterion is chosen to balance the mean absolute error
(MAE) and root mean squared error (RMSE). The AdamW optimizer with betas 0.9
and 0.999 is used, the learning rate is dynamically adjusted by pytorch
lightning. Finetuning only the regression head produces comparable results to
finetuning all parameters. A finetuning-scheduler to incrementally unfreeze an
additional layer is also explored but produces similar results while
introducing training issues related to extraction of the attention scores and
gradient hook. SMILES randomization as data augmentation during finetuning is
explored but consistently yields significantly worse results.

#### Datasets.

Two datasets are used, respectively the AqueousSolu dataset for solubility in
water, and the CombiSolu-Exp dataset for solubility in organic solvents. Both
are from the SolProp data collection (https://zenodo.org/record/5970538). For
AqueousSolu the “accurate” setting proposed by the SolProp authors is obtained
by selecting all molecules with more than 1 measurement, where the standard
deviation does not exceed 0.2. This yields 578 test set molecules, the source
of the mismatch of 1 compared to SolProp (579) is unknown. One extreme outlier
with $\log(S)$ = 6.4 is also removed from the dataset.

For the CombiSolu-Exp dataset, the “$\mathrm{T}=\mathrm{298K}$” test set is
obtained by selecting all molecules with a temperature which rounds to 298K.
Samples which contain missing “temperature, solvent density or experimental
$\log(S)$” values are excluded. The $\log(S)$ measurements of the CombiSolu-
Exp dataset are scaled to the range [0, 1] for training stability. For the
“random” test setting in both datasets, a 10% test set obtained from a random
seed is set aside. The remaining molecules are split into a 90% training and a
10% validation set. The random seed only affects the split into train and
validation set in the “accurate” and “$\mathrm{T}=\mathrm{298K}$” setting.

#### MegaMolBART.

The pretrained MegaMolBART model is obtained through the provided docker
container image (version v0.2:0.2.0) from github.com/NVIDIA/MegaMolBART. The
pretrained pytorch model and weights are accessed directly to avoid the
latency of the InferenceWrapper. The RegExTokenizer vocabulary is extended
with a <REG> token, the tokenize function similarly needs to be adjusted. The
encode function of MegaMolBART takes the SMILES string and produces latent
representations of dimensionality (n_tokens, 512), which includes <PAD> tokens
up to pad up to the maximum token size of the batch.

#### Extracting attention scores & gradients.

To extract the attention scores and gradients, the transformers.py file needs
to be extended with a few lines of code to save and retrieve the attention
scores and attach a hook to evaluate the gradients during inference. These
hooks do not need to present during training but can also be extended
afterwards. It should be noted that the validation, test and predict functions
require enabling the gradient (torch.set_grad_enabled(True)) to obtain the
gradients of the attention scores due to the backwards hook, which is enabled
when training = False.

#### Computation cost.

All models were finetuned on a single Nvidia RTX 2080 GPU, which allows for a
maximum batch size of $\sim 48$ due to memory constraints. A full finetuning
run with 50 epochs takes $\sim 10$ minutes for AqueousSolu and $\sim 3$
minutes for 40 training epochs on CombiSolu-Exp. Obtaining the attention
scores and gradients for the explainability technique is very computationally
efficient, as only a single forward and backward pass of the model is
necessary, while the propagation and attribution of relevance is negligible in
terms of computational cost.

###  7.1 Model results.

In addition to the main AqueousSolu results (Table 1), the following table
presents errors from 5 runs keeping the data seed for the split into train,
validation and test set identical while randomizing the model seed which
affects the initialization of the regression head. Thereafter, results for the
“$\mathrm{T}=\mathrm{298K}$” test set of the CombiSolu-Exp dataset are
presented.

AqueousSolu results with random model seed Model | Random | Random
---|---|---
| MAE (SE) | RMSE (SE)
Average | 0.602 | 0.867
Standard Error | (0.002) | (0.004)
Table 2: Model error of MegaMolBART trained on the AqueousSolu dataset using 5 runs with random model seeds affecting the initialization of the regression head. CombiSolu-Exp results with random data seed Model | Random | Random
---|---|---
| MAE (SE) | RMSE (SE)
Average | 0.378 | 0.558
Standard Error | (0.012) | (0.024)
Table 3: Model error of MegaMolBART trained on the “random” test set of the
CombiSolu-Exp dataset using 5 runs with random data seeds affecting the split
into train, validation and test set.

For the CombiSolu-Exp “$\mathrm{T}=\mathrm{298K}$” test set, the <REG>,<SEP>
model variant performs noticeably worse. An alternative approach is to model
the interaction using cross-attention layers, where the query and key vectors
correspond to the averaged latent representations of the solute and solvent,
which achieves state-of-the-art accuracy on both the random and the
“$\mathrm{T}=\mathrm{298K}$” test set. While this approach models can model
explicit interactions between the solute and solvent, further modifications of
the explainability technique are required to obtain separate relevancy
matrices $\mathbf{R}_{\mathrm{solu}}$ and $\mathbf{R}_{\mathrm{solv}}$ for the
solute and solvent.

Figure 6:  Parity plots showing the predictive accuracy of the finetuned
MegaMolBART model trained on the $\mathrm{T}=\mathrm{298K}$ test set of the
CombiSolu-Exp dataset. The left plot shows the comparatively poor results of
the <REG>,<SEP> model variant next to the state-of-the-art results of the
cross-attention variant shown on the right.

###  7.2 AqueousSolu Visualizations.

Once the relevance weights for each atom token are obtained, the molecules’ 2d
structure diagram is plotted using the MolDraw2DCairo and
DrawMoleculeWithHighlights functions of RDKit. Each atom’s attributed weights
are scaled using the molecules’ minimum and maximum attributed importance to
obtain relative importance in arbitrary units and mapped using a color palette
before plotting. The light color palette “green” of the seaborn library is
chosen as the most intuitive.

Figure 8: A representative set of visualizations of molecules selected from a
random AqueousSolu test set, which showcase what the finetuned MegaMolBART
represents. As elaborated in the discussion of results, high importance is
generally attributed to the correct functional groups. However, for some
molecules relevance is also attributed to the structural carbon backbone,
which is likely a result from the pretraining regiment of MegaMolBART. The
attributed relevance of symmetric molecules can be poor, but addressing this
limitation is not trivial.

###  7.3 CombiSolu-Exp Visualizations.

Figure 9: Visualization of attributed relevance of a solute in four solvents,
specifically an oil (toluene, top left) and three alcohols (isopropanol, top
right; 1-butanol, bottom left; ethanol, bottom right). The differences in
attributed importance of individual atoms or functional groups towards
solubility in different solvents can be seen, which are nearly identical for
the solute. This is contrary to expectation, which suggests different
functional groups should contribute towards solubility in an oil compared to
the alcohols. The visualizations highlight the model does not explicitly model
the interaction between solute and solvent. The color bar is the same as in
Fig. 8.