# Bidirectional Generation of Structure and Properties Through a Single
Molecular Foundation Model

Jinho Chang1 and Jong Chul Ye1,†

###### Abstract

The recent success of large foundation models in artificial intelligence has
prompted the emergence of chemical pre-trained models. Despite the growing
interest in large molecular pre-trained models that provide informative
representations for downstream tasks, attempts for multimodal pre-training
approaches on the molecule domain were limited. To address this, we present a
novel multimodal molecular pre-trained model that incorporates the modalities
of structure and biochemical properties, drawing inspiration from recent
advances in multimodal learning techniques. Our proposed model pipeline of
data handling and training objectives aligns the structure/property features
in a common embedding space, which enables the model to regard bidirectional
information between the molecules’ structure and properties. These
contributions emerge synergistic knowledge, allowing us to tackle both
multimodal and unimodal downstream tasks through a single model. Through
extensive experiments, we demonstrate that our model shows remarkable
capabilities in solving various meaningful chemical challenges, including
conditional molecule generation, property prediction, molecule classification,
and reaction prediction.

Graduate School of AI, KAIST, Daejeon, South Korea

†Correspondence should be addressed to J.C.Y<EMAIL_ADDRESS>
## Abstract

## Introduction

Capturing complex relations between chemical objects and their properties is
the essence of numerous chemical challenges. During the last decade,
artificial intelligence has emerged as a promising tool in chemistry research
for estimating many biochemical properties and interactions between molecules,
polymers, and proteins, which are difficult to obtain experimentally[1, 2, 3].
Various deep learning-based approaches in the chemical domain employed deep
neural networks to extract desired characteristics like intrinsic properties,
biochemical activities, and chemical reactions from raw molecule data[4, 5,
6]. Especially, de novo molecule design has been extensively studied using
recurrent networks[7], variational autoencoders[8, 9], graph networks[10],
etc[11, 12, 13]. More recently, unsupervised learning approaches of learning
better representations of the chemical inputs have been suggested [14, 15, 16]
to overcome the limitation of learning separate features for each task in a
supervised manner. These recent approaches are on the same track as the
concept of the foundation models that are trained with large datasets and are
often considered as a new paradigm of deep learning[17, 18].

Specifically, a concept of pre-training a neural network in a self-supervised
manner for a better feature representation has been adapted for various
chemical fields[15, 16, 14]. N-Gram Graph[19] and GROVER[20] used a graph
neural network and a graph transformer network, respectively, to obtain a pre-
trained model from the molecular graph. ChemBERTa-2 [21] trained a roBERTa
model with 77 million molecules to build a molecular foundation model, by
training the model to predict 200 different chemical property values.

Meanwhile, in the computer vision field, multimodal pre-training methods like
Vision-Language Pre-training (VLP)[22] have achieved outstanding performance
in downstream tasks that require an understanding of both image and text. Most
of the modern VLP models utilize Transformer[23] architecture and its cross-
attention mechanism to learn the correlation between different modalities[24,
25]. Moreover, several works introduced contrastive learning, which
assimilates features with the same context and distances semantically
unrelated features, to align image and language features in the common feature
space[26, 27, 28]. VLP enables various tasks such as visual question
answering[29], image-text retrieval[30], text-driven image generation[31],
image-driven text generation[32], etc., which are not possible using single
modality foundation models.

Inspired by the success of multimodal learning, several recent works tried to
obtain a better feature of a molecule by leveraging knowledge from different
data representations. Winter _et al._ trained a translation model between
Simplified Molecular-Input Line-Entry System (SMILES) and International
Chemical Identifier (InChI) key to get a feature vector with meaningful
information that both molecular representations have in common[33]. Zhu _et
al._ used a self-supervised training method of BYOL[34] between different
molecule representations of SMILES and molecular graphs to build a dual-view
model[35]. However, these works introduced multimodality only for the
enhancement of a molecule feature for unimodal tasks, not for the interplay
between those different modalities. Furthermore, since SMILES, InChI, and
graph representations contain almost identical information about the
connection between atoms in a molecule, it is unlikely to expect new emergence
properties by multimodal learning between these different molecule
representations.

In this work, we are interested in the cross-modal comprehension between
molecule structure and the associate properties, which facilitates solving
meaningful tasks in many applications like property predictions, conditional
molecule design[36, 37], etc. Taking a step further from multi-task learning
methods[38] which use the prepared properties as labels to extract general
features[21], our approach regards a set of properties as a stand-alone
modality that represents the input molecule and suggests that multimodal
learning for molecules with this property modality can provide much more
informative features. Specifically, we propose a novel molecule Structure-
Property Multi-Modal foundation model(SPMM) which allows various chemistry
experiments in silico, which is pre-trained with a wide range of molecules’
structures and a vector of its properties. By employing a Transformer
architecture[23], the intramodal feature extraction and intermodal fusion can
be done with self-attention and cross-attention mechanisms, respectively.

Our experimental results show that simultaneous learning of structural
features with information from the associate properties through a single
foundation model gives us a better representation that can be fine-tuned for
various downstream tasks. Specifically, by treating both structure and
property symmetrically, the model can perform bidirectional generation and
prediction with a single pre-trained model, which was not possible before.

Figure 1: (a) Overview of the model architecture and pre-training objectives
of SPMM. The contrastive loss aligns the output feature of two unimodal
encoders into the same embedding space. The fusion encoder learns the
relations between two modalities, trained with Next Word Prediction (NWP),
Next Property Prediction (NPP), and SMILES-Property Matching loss (SPM). (b)
Downstream tasks that require multimodal comprehension: i) PV-to-SMILES
generation, ii) SMILES-to-PV generation. (c) Downstream tasks for single
modality inputs: i) property prediction, ii) forward and retro reaction
prediction.

Fig. 1(a) illustrates the overall model architecture and training objectives
for SPMM. The framework of SPMM extends the structure of the dual-stream VLP
models[28, 27, 39]. Dual-stream VLP models encode the input for each modality
with a unimodal encoder, then use another encoder module to perform cross-
attention by using one modality feature as a query and the other modality
feature as a key/value. When a training molecule is given, SPMM takes the
molecule’s SMILES string and its property vector (PV) as multimodal data
inputs as shown in Fig. 1(a). The SMILES and PV are passed through their
corresponding unimodal encoders, which perform self-attention where embedded
inputs become the key, query, and value. After two unimodal features are
obtained, contrastive learning aligns the SMILES and PV features into the same
embedding space by assimilating the features that contain the same context.
This is known to improve the model performance by making cross-modal encoding
easier and guiding the unimodal encoded features to reflect more semantics of
the input[27]. Then, the encoded SMILES and PV features are passed through the
fusion encoders, which perform cross-attention between SMILES and PV features.
This single fusion encoder can perform cross-attention with an alternation of
its query and key/value input because the contrastive learning aligns the
output of the SMILES encoder and the PV encoder into the same feature
space.[39] The fusion encoder is pre-trained with Next Word Prediction (NWP)
for SMILES, Next Property Prediction (NPP), and SMILES-PV Matching loss (SPM).
Prediction of the next component from the given transformer input is a
commonly used self-supervised learning objective, and our NWP and NPP tasks
make the model learn the contextual relationship between SMILES tokens and
properties with the aid of the other modality’s semantic feature.
Additionally, SPM predicts whether a given pair of SMILES and PV represents
the same molecule or not.

Once trained, SPMM can be used for various bidirectional downstream tasks that
require an understanding of both SMILES and properties like property
prediction (SMILES-to-properties) and property-conditioned molecule generation
(properties-to-SMILES, also referred to as inverse-QSAR[37]) as shown in Fig.
1(b). Furthermore, the pre-training objectives that we’ve used allow the pre-
trained SPMM to be applied for single-modality tasks as well, such as molecule
classification and reaction predictions (see Fig. 1(c)). The pre-trained SPMM
showed comparable performances to state-of-the-art models in these unimodal
tasks, which suggests the model’s generalization ability as a foundation
model.

## Results

###  0.1 The model learns bidirectional comprehension between SMILES and
properties.

Once SPMM was pre-trained, we made the model generate SMILES with given PV
inputs only, which is a crucial challenge for many chemical tasks such as de
novo molecule design. As one of the major approaches for drug discovery,
various methods have been suggested for generating molecules with desired
properties[11, 9, 13, 10]. In the approaches presented so far, the maximum
number of simultaneously controllable properties wasn’t very large. Also, the
length of the input property vector cannot be changed. Whenever the target
properties change, the model needs to be trained again for the new wanted
conditions. In contrast, the pre-trained SPMM can take 53 properties used in
pre-training as input conditions and generate molecules that satisfy all of
them, without separate additional training for each property combination.
Moreover, for the properties that we don’t want to control, we can let the
model ignore those conditions by replacing them with the [UNK] token that we
used in pre-training. This is very useful because controlling all 53 input
properties is not a usual scenario in practice, and is also not easy since the
properties are correlated and entangled (e. g., ‘5 atoms & 30 bonds’ or ‘2
rings & 5 aromatic rings’ is unlikely to be a valid PV input).

To demonstrate the molecule generation capability of SPMM, we prepared a
number of PV-to-SMILES generation scenarios and let the pre-trained SPMM
autoregressively generate SMILES using the input properties. This process of
SPMM is very similar to the sequence-to-sequence translation tasks in terms of
the model pipeline (see Figure S3-(a) for details), from the property sentence
of PV to the molecular structure sentence of SMILES.

The validity, uniqueness, and novelty of the generated molecules are the
quantitative metrics of SPMM’s molecule generation, defined as follows:

$\displaystyle\text{validity}=\frac{\text{\\#SMILES with valid
syntax}}{\text{\\#generated SMILES}}$
$\displaystyle\text{uniqueness}=\frac{\text{\\#non-duplicate valid
SMILES}}{\text{\\#valid SMILES}}$
$\displaystyle\text{novelty}=\frac{\text{\\#unique SMILES not in the pre-
training data}}{\text{\\#unique SMILES}}$

Additionally, as a qualitative metric to see how the generated SMILES match
the property input, we measured the normalized Root Mean Square Error
(normalized RMSE) between the input conditions and the generated molecules’
properties. More specifically, we calculate the average of the RMSE of all
controlled properties, after those values are normalized with the
corresponding property’s mean and standard deviation in the pre-training
dataset. We note that RMSE was calculated on the normalized scale of each
property because the values of the properties span multiple orders of
magnitude.

sampling | input PV | validity | uniqueness | novelty | normalized RMSE
---|---|---|---|---|---
deterministic | 1,000 unseen PubChem SMILES’ PV | 0.982$\pm$0.003 | 0.830$\pm$0.077 | 0.954$\pm$0.004 | 0.194$\pm$0.006
stochastic | full PV of the molecule 1 | 0.882$\pm$0.004 | 0.999$\pm$0.001 | 1.000$\pm$0.000 | 0.189$\pm$0.007
Molecular weight=150 | 0.913$\pm$0.008 | 0.999$\pm$0.000 | 0.932$\pm$0.005 | 0.262$\pm$0.004
#ring=2, #aromatic ring=1, TPSA=30, QED=0.8 | 0.933$\pm$0.007 | 0.999$\pm$0.001 | 0.979$\pm$0.005 | 0.343$\pm$0.009
no property control | 0.751$\pm$0.008 | 1.000$\pm$0.000 | 0.998$\pm$0.002 | -

Table 1: Quantitative and qualitative results on various scenarios of PV-to-
SMILES generation tasks, with the mean value and standard deviations. For
deterministic sampling, we ran the experiment with four different random sets
of 1,000 unseen PVs. In the case of stochastic scenarios, four different
random seeds were used for each experiment.

For the first PV-to-SMILES generation scenario, we prepared 1,000 PVs of
SMILES from PubChem[40] that are not contained in the pre-training dataset and
fed them to the pre-trained SPMM to generate appropriate SMILES. Here, the
sampling process was done in a deterministic manner (greedy sampling):
starting from the SMILES [CLS] token ([CLS]S), the model predicts the
probability distribution of the next token and chooses the option with the
highest probability. The first row of Table 1 shows its results. Among the
output of deterministic PV-to-SMILES generation for 1,000 PVs, 98.2% of the
generated output were valid SMILES. The mean RMSE of the 53 normalized
properties was 0.194, which implies that the properties of the generated
samples agree with the property input.

Figure 2: Property distribution of the generated molecules with different PV
inputs and [UNK] token masking. The red vertical dotted lines are the input
property values, and the grey vertical lines are the mean of that property in
the pre-training dataset. The controlled properties are colored in red, and
uncontrolled properties (=masked with [UNK] token) are colored in blue. Due to
the lack of space, only 12 out of 53 properties are shown for each case. For
each PV-to-SMILES scenario, we included the structure of two of the generated
molecules. (a) All 53 properties are controlled, without using the [UNK]
token. The input PV was obtained from the molecule 1. (b) Molecular Weight to
150, and the other property inputs are masked. (c) #ring, #aromatic ring,
TPSA, and QED are controlled to 2, 1, 30, and 0.8. The other property inputs
are masked. (d) Every property is replaced with [UNK] token.

Application fields like drug discovery often require generating multiple
molecules for a single wanted target property condition. This can be done by
sampling the next token stochastically from the modeled probability
distribution instead of using a token with the highest probability. To verify
our model’s ability to generate multiple molecules from a single PV input, we
generated 1,000 SMILES with stochastic sampling on a fixed PV. Figure 2 shows
the property distributions of 1,000 molecules generated from a single PV
input. The mode of each property distribution lands on the input property
value (Fig. 2-(a)). In the situation when only some of the properties are
given, the model only regards the known properties while the other masked
properties are not restricted (Fig. 2-(b), Fig. 2-(c)). SPMM can generate
molecules even with no property information at all; when all input properties
are replaced with [UNK] token (Fig. 2-(d)), the model performs an
unconditional molecule generation, and the output follows the distribution of
the pre-training dataset. The validity, uniqueness, and novelty of the
generated molecules under conditions in Figure 2 are listed in the
“stochastic” rows of Table 1. The validity fluctuated depending on how
feasible or difficult the property input is, and it was between 0.75 and 0.9
in most cases. The uniqueness, the ratio between the number of unique
molecules against the number of validly generated molecules, was almost 100%
in every condition we have experimented with. More examples of the generated
molecule can be found in Supplementary result S1.

Figure 3: Examples of molecule editing, by changing specific values from the
original PV and performing PV-to-SMILES generation with it. The colored output
values correspond to the changed properties from the original PV. (1) The
output of the same PV of the source molecule. (2) The output when
_#aromatic_ring_ is changed to 0. (3) The output when _#ring_ is changed to 2
and _#aromatic_ring_ is changed to 1. (4) The output when _logP_ is changed to
7. (5) The output when _#rotatable_bond_ is changed to 12. For the generation,
the other 41 property conditions are masked by the [UNK] token.

The aforementioned results demonstrate that SPMM can perform molecule
generation with arbitrary PV inputs, which enables simple molecule designing
and editing. Figure 3 contains the output of the SPMM’s stochastic molecule
generation for five PV inputs, which all originated from the PV of the
molecule 1 but four of them had certain values changed. The generated
molecules follow the input modification while maintaining unmodified
properties similarly. SPMM is even able to generate molecules with the out-of-
domain conditions such as _$\log P=7$_ (note that $\sim 5\%$ of the pre-
training dataset has $\log P>$7).

Regarding the overall molecule generation performance of SPMM, we want to
emphasize that SPMM can generate suitable SMILES for many property conditions
that the model has not seen in its pre-training. When we trained SPMM without
50% of random property masking with [UNK] token, the model only worked when
all 53 properties are given since the model has not seen the partially-given
properties. However, even with the technique of [UNK] token masking, the model
cannot face most of the $2^{53}$ possible property combination during the pre-
training process. The SPMM’s ability to handle arbitrary property conditions
for SMILES generation comes from treating PV as a ‘language with 53 words’ and
focusing on each property separately, not simply considering the entire
property input as a single condition. This innovative approach for conditional
molecule generation has never been demonstrated with the existing methods and
thus can be used for many important chemical fields.

Figure 4: Scatter plots of the 1,000 ZINC15 molecules’ real property value
against the generated output, for 12 selected properties. The $x$-axis is the
real property value, and the $y$-axis is the model output. The grey dotted
line is the $y=x$ line.

With the same approach as SMILES generation, the pre-trained SPMM can also be
used to generate a PV with SMILES input only. This task is equivalent to
performing 53 property predictions of a given SMILES at once. Similar to the
PV-to-SMILES generation, properties are predicted in an autoregressive manner:
the model predicts the first property value using only the property [CLS]
token ([CLS]P), then takes all previous outputs again to get the next
prediction value, and so on (see Figure S3-(b)). Although 53 properties that
we’ve used can be calculated using the Python module, the purpose of this
experiment is to verify that the data-driven way of property estimation
coincides with the analytic approach.

Specifically, we fed 1,000 SMILES from the ZINC15 dataset[41], which are not
contained in the pre-training dataset, to the pre-trained SPMM and generated
their corresponding PV. Figure 4 is the scatter plot of the real property
value against the generated output for 12 selected properties out of 53 that
we used for pre-training. It is clear that SPMM’s predicted property is very
close to the actual value, and most of the data point lies on the $y=x$ line.
Although the model virtually has never seen a full-filled PV in the pre-
training due to the 50% of random property masking, the model could
autoregressively predict all 53 properties as a whole. The mean $r^{2}$ score
of the 53 properties was 0.932. The mean of the normalized RMSE for 53
properties was 0.118. The full scatter plot for all 53 properties with each
$r^{2}$ score and raw RMSE is in the Supplementary Figure S2.

To provide an interpretation of the pre-trained SPMM’s performance presented
so far, we further analyzed the learned cross-modal comprehension between
SMILES and property vectors by visualizing the attention scores from the pre-
trained SPMM. Transformer-based models have the benefit of intuitive attention
visualization that shows how the model considers the relation between the
input queries and keys, by providing cross-attention scores between them.

Figure 5: The mean attention score from the attention heads in the SPMM fusion
encoder’s final cross-attention layer for two sample molecules. A darker green
means a higher attention score. For the attention process, the property
features were used as queries, and the SMILES features are used as keys and
values. The corresponding fragments for each token are indicated with ivory
boxes on the molecular structure, while fragments for duplicated tokens are
color-coded with purple. We have calculated cross-attention scores for all 53
properties and SMILES tokens, but only 12 of those properties are shown.

In Figure 5, we plotted the cross-attention score from the last fusion layer
of our pre-trained SPMM when SMILES and its property vector inputs were given.
Since there are multiple heads for the cross-attention, we took the mean of
their attention scores. It is interesting that the aspect of cross-attention
scores followed the intuitive relations between chemical properties and
molecular fragments. The properties related to hydrogen bonding (_NumHDonors,
NumHAcceptors_) show high attention scores for tokens with oxygen and nitrogen
atoms. The property _RingCount_ focuses on the tokens that are involved with
rings while showing weak attention to side groups, and the property
_NumAromaticRings_ only gives high attention score to the components of
aromatic rings. When different SMILES tokens played a similar role in the
molecule such as ‘ _c1ccccc1)_ ’ and ‘ _c1ccccc1_ ’ in the molecule 15, their
attention patterns were similar as well. This result demonstrated that SPMM
could capture the relations between molecule structures and chemical
properties without explicitly-given supervision between them.

###  0.2 Generalization ability as a molecular foundation model.

So far, we have demonstrated that the pre-trained SPMM can be applied to tasks
that require an understanding of the relationship between SMILES and
properties. However, we can also employ the pre-trained SPMM for challenges
that only use SMILES data, such as molecular property prediction. One
advantage of having a dual-stream VLP model structure is that the SPMM’s
multimodal pre-training process includes adjusting the output of one unimodal
encoder to contain contextual information from the other modality, by aligning
it with the other unimodal encoder’s output. This implies that the SMILES
encoder output is a unimodal representation vector, that not only embeds the
input molecule’s structural information but it’s also enhanced by its property
information.

We have analyzed if our pre-trained model had learned an informative
representation that can be readily used for other tasks, even for a single
modality. So we only utilized the SMILES encoder of pre-trained SPMM (see
Supplementary Figure S3-(c)) and made a benchmark study on nine
MoleculeNet[42] downstream tasks and a Drug-Induced Liver Injury (DILI)
prediction task. Each MoleculeNet task is a regression or classification task
for pharmaceutical/biochemical applications like solubility, toxicity, and
brain penetrability. The DILI classification task was done to overcome the
potential limitation of open databases[43, 44] and verify if SPMM could be
extended to more complex endpoints. The task is to classify whether the given
molecule has a risk of causing liver injury. Since many proposed DILI machine
learning models have built their dataset rather than using common benchmarks,
we took the dataset preparations from a known publication[45] and compared the
performance with it for a fair evaluation.

Table 2 contains the performance of SPMM and other models for MoleculeNet.
Using only 6 BERT encoder layers, SPMM showed comparable performances with
state-of-the-art models for all tasks. It achieved the best performance for
Clearance, BBBP, and Clintox tasks, showing its capability as a foundation
model. We’ve also observed that the score of our model dramatically decreased
without pre-training. SPMM also outperformed the proposed 5-ensemble models on
the DILI classification task under the same data preparation, which was not
the case for the naive BERT layers without SPMM pre-training.

| regression[RMSE, $\downarrow$] | classification[AUROC in $\%$, $\uparrow$]
---|---|---
Dataset | Delaney ESOL | LIPO | Freesolv | BACE | Clearance | BBBP | BACE | Clintox | SIDER
#data | 1128 | 4200 | 642 | 1513 | 837 | 2039 | 1513 | 1478 | 1427
#task | 1 | 1 | 1 | 1 | 1 | 1 | 1 | 2 | 27
D-MPNN[46] | 1.050$\pm$0.008 |  0.683$\pm$0.016 | 2.082$\pm$0.082 | 2.253∗ | 49.754∗ | 71.0$\pm$0.3 | 80.9$\pm$0.6 | 90.6$\pm$0.6 | 57.0$\pm$0.7
N-GramRF[19] | 1.074$\pm$0.107 | 0.812$\pm$0.028 | 2.688$\pm$0.085 | 1.318∗ | 52.077∗ | 69.7$\pm$0.6 | 77.9$\pm$1.5 | 77.5$\pm$4.0 |  66.8$\pm$0.7
N-GramXGB[19] | 1.083$\pm$0.082 | 2.072$\pm$0.030 | 5.061$\pm$0.744 | - | - | 69.1$\pm$0.8 | 79.1$\pm$1.3 | 87.5$\pm$2.7 | 65.5$\pm$0.7
PretrainGNN[47] | 1.100$\pm$0.006 | 0.739$\pm$0.003 | 2.764$\pm$0.002 | - | - | 68.7$\pm$1.3 | 84.5$\pm$0.7 | 72.6$\pm$1.5 | 62.7$\pm$0.8
GROVERlarge[20] | 0.895$\pm$0.017 | 0.823$\pm$0.010 | 2.272$\pm$0.051 | - | - | 69.5$\pm$0.1 | 81.0$\pm$1.4 | 76.2$\pm$3.7 | 65.4$\pm$0.1
ChemRL-GEM[48] |  0.798$\pm$0.029 |  0.660$\pm$0.008 |  1.877$\pm$0.094 | - | - | 72.4$\pm$0.4 |  85.6$\pm$1.1 | 90.1$\pm$1.3 |  67.2$\pm$0.4
ChemBERTa-2${}_{\text{(MTR-77M)}}$[21] | 0.889∗ | 0.798∗ | - | 1.363∗ | 48.515∗ | 72.8∗ | 79.9∗ | 56.3∗ | -
MolFormer†[49] | 0.880$\pm$0.028 | 0.700$\pm$0.012 | 2.342$\pm$0.052 |  1.047$\pm$0.029 |  43.175$\pm$1.537 |  73.6$\pm$0.8 |  86.3$\pm$0.6 |  91.2$\pm$1.4 | 65.5$\pm$0.2
SPMM(w/o pre-train) | 1.272$\pm$0.015 | 1.009$\pm$0.021 | 3.018$\pm$0.179 | 1.675$\pm$0.010 | 53.544$\pm$0.312 | 66.6$\pm$0.3 | 78.7$\pm$2.6 | 76.3$\pm$1.5 | 57.1$\pm$1.6
SPMM |  0.818$\pm$0.008 | 0.692$\pm$0.008 |  1.907$\pm$0.058 |  1.096$\pm$0.011 |  42.841$\pm$1.251 |  74.9$\pm$0.8 | 84.8$\pm$0.3 |  91.8$\pm$0.9 | 65.5$\pm$0.9

Table 2: Benchmark results on MoleculeNet downstream tasks. The best
performance for each task was written in bold, and the second-best performance
was underlined. For each task, we fine-tuned our model in four random seeds
and recorded the mean and the standard deviation of those results. The
benchmark model results were taken from ChemRL-GEM and ChemBERTa-2. ∗The
standard deviation cannot be found in the source of the benchmark results.
†Unofficial results, obtained from the official checkpoint under our data
preparation.

model | Acc in %[$\uparrow$] | Selectivity in %[$\uparrow$] | Specificity in %[$\uparrow$] | AUROC in %[$\uparrow$]
---|---|---|---|---
Ai _et al._[45] (best single model on training set) | 81.1 | 81.0 | 81.5 | 89.6
Ai _et al._ (5-ensemble) | 84.3 | 86.9 | 75.4 | 90.4
SPMM(w/o pre-train) | 72.6 | 70.6 | 79.2 | 82.0
SPMM | 84.7 | 85.7 | 81.6 | 91.2

Table 3: The DILI classification task performance of Ai _et al._[45] and SPMM.
The best performance for each metric was written in bold.

We also trained SPMM for the forward and retro-reaction prediction tasks,
which require the model to predict the product SMILES from the reactant SMILES
and vice versa. Regarding both tasks as sequence-to-sequence generation, the
model pipeline for these reaction prediction tasks is the same as the PV-to-
SMILES generation tasks, except the PV encoder is replaced with the SMILES
encoder (see Supplementary Figure S3-(d)). The detailed task definition and
dataset preparation are described in the Methods section.

forward prediction | molecule modality | top-k accuracy in %[$\uparrow$]
---|---|---
| string-based | graph-based | k=1 | k=2 | k=3 | k=5
Molecular Transformer[50] | O |  | 88.7 | 92.1 | 93.1 | 94.2
Augmented Transformer[51] | O |  | 90.6 | 94.4 | - | 96.1
Chemformerlarge[52] | O |  | 91.3 | - | - | 93.7
Graph2SMILES[53] | O | O | 90.3 | - | 94.0 | 94.8
MEGAN[54] |  | O | 86.3 | 90.3 | 92.4 | 94.0
LocalTransform[6] |  | O | 90.8 | 94.8 | 95.7 | 96.3
SPMM | O |  | 91.5 | 93.4 | 94.6 | 95.3

retro-reaction prediction | top-k accuracy in %[$\uparrow$]
---|---
| k=1 | k=5 | k=10
SCROP[55] | 43.7 | 65.2 | 68.7
Two-way Transformer[56] | 47.1 | 73.1 | 76.3
Augmented Transformer[51] | 48.3 | 73.4 | 77.4
Chemformerlarge[52] | 54.3 | 62.3 | 63.0
SPMM | 53.0 | 67.4 | 70.3

Table 4: The performance of SPMM and other works on the forward and retro-
reaction prediction task. For the retro-reaction prediction task, we only
prepared the benchmark results of string-based models. The highest accuracy is
written in bold, and the performance of the runner-up model is underlined. The
benchmark model results are from the paper of LocalTransform[6] and
Chemformer[52].

Table 4 shows the performances of SPMM and other benchmark models on forward
and retro-reaction prediction tasks. Although the reaction prediction tasks
are not the best scenario for the property-emergence features to play
significant roles, SPMM showed the highest top-1 accuracy in the forward-
reaction task with a relatively small pre-training data size (_i.e._ 20M
molecules, compared to 100M molecules of Chemformer). SPMM also achieved the
second-best top-1 accuracy among the string-based retro-reaction task models.

## Discussion

In this work, we proposed a transformer-based multimodal chemical foundation
model SPMM. The proposed model allows for bidirectional generation/prediction
of molecular structure and properties, as well as unimodal tasks like reaction
prediction. During the process, we introduced a method of treating property
collections as a language so that the model could learn the relationship
between SMILES tokens and each property independently. We demonstrated that
pre-trained SPMM showed remarkable performances in problems for interactions
between SMILES and PV domains. And not only for multimodal challenges but even
its unimodal feature for SMILES, SPMM also provides a useful representation
that can be fine-tuned for many molecular downstream tasks. It is important to
note that all of these results were obtained with a pre-training of 20 million
molecules, which is relatively small compared to other large pre-training
approaches and still has room for better performance with more data and
parameters. We also note that we’ve gathered our 53 properties to let them
cover the widest range possible, rather than paying the best effort to select
the most effective combination of properties. This implies the proposed
structure-property multimodal training can be flexibly adopted with different
property selections, according to the given specified scenarios.

Despite the noticeable performances of SPMM, it has several chances for
improvement. One of those comes from using the SMILES notation. Although
SMILES can contain full details about the 2D structure of the molecule, the
information on how atoms and bonds are connected only exists implicitly. Also,
a slight modification in molecular structure can be a drastic change in
SMILES. Graph format is another widely used modality for molecule
representation that contains the explicit information of the adjacency matrix,
which can be an alternative for SMILES. Another limitation in our current SPMM
is that the 53 properties we used happen to be invariant with the changes in
the stereochemistry of the given molecule. It is known that considering
stereochemistry plays a crucial part in various biochemical tasks. However,
the 53 properties we used cannot provide any knowledge about stereochemical
information since their values are unchanged in different stereoisomers. This
makes the SMILES encoder output of different stereoisomers converge since the
contrastive loss aligns them to the same PV feature. We believe this is the
prominent factor that lowered the performance of SPMM in MoleculeNet tasks,
which could be resolved by using more properties that reflect the molecule’s
stereochemistry. Moreover, validation through wet-lab experiments to verify
the model’s predicted/generated properties is another possible further study.
Overcoming these drawbacks of the current study and making the model more
applicable to other chemical tasks could be the works for the future.

Nevertheless, we believe that our approach can provide a pre-trained model
capable of encompassing each input domain and their multimodal domain
simultaneously, which has a vast potential utility. We expect this approach to
be applied to more various and practical chemical situations by using broader
and richer molecular modalities, and possibly, different biochemical domains
like polymers and proteins.

## Methods

###  0.3 Handling SMILES and property values as a language.

Molecules can be represented with various formats such as fingerprints,
strings like SMILES, InChI, or a molecular graph. Since these different
notations contain almost the same information about complete molecular
structure, we employed SMILES to describe a molecule structure. SMILES is a
sequence of characters that represents the connection structure of the
molecule. Many researchers treat SMILES as a variant of language data and
utilize a concept of language models for chemical tasks on SMILES data[11, 21,
57].

Figure 6: Embedding process for SMILES and the corresponding PV.

Figure 6-(a) illustrates our embedding procedure for the input SMILES. The raw
SMILES string is tokenized by the tokenizer and embedded by the SMILES encoder
with the [CLS]S token and the [SEP] token. Here, [CLS] token is a special
token attached to the beginning of every input sequence[58]. Although the
[CLS] token itself doesn’t contain any meaning, the bidirectional attention
mechanism of the model allows the [CLS] token to contain contextual
information of the entire input. Once the model is pre-trained, the [CLS]
token output of the given sequence can be considered as an input
representation vector and be used for classification/regression downstream
tasks, as in many BERT variations for images[59, 60] and VLP[27].

In the SMILES tokenization, our tokenizer tokenizes a given SMILES into
fragments that are contained in a prepared token dictionary of 300 subwords.
This dictionary was obtained from the pre-training data SMILES corpus by the
BPE algorithm[61], which starts from a set of simple characters and
iteratively appends the most frequent token pairs as a merged subword. Being
widely adopted for various language models[62, 63], the BPE algorithm has
provided a subword dictionary containing common functional groups and
substructures like benzene rings, carbonyl groups, two-letter atoms, and amino
groups. Compared to naive character-wise tokenization which considers each
character as a separate token, the merged subwords help the model’s chemical
inference for chemical groups and reduce the total number of tokens.

Meanwhile, a set of chemical properties does not change its carrying
information by changing the internal order, but they certainly have
correlations between the properties. And it is known that a transformer
architecture also performs well for different modalities like images, by
giving specific order to its components and treating them as a sequence. For
this work, we built a PV for each molecule that contains 53 molecular
properties and considered this as a sentence with a length of 53. These
properties from the RDKit python module[64] cover a wide range from simple
ones, such as the number of rings and molecular weight, to complex properties
like solubility, TPSA, and druggability.

The transformer architecture of our model considers each element of PV as a
token to perform the attention mechanism, which is equivalent to regarding PV
as a semi-sentence of 53 properties. Although the size of the vocabulary is
more limited and their order is fixed compared to natural language, it
provides much more precise and compact information about the 53 properties.
One benefit of regarding PV as a language is that we do not have to collect
all elements to build a valid PV. In contrast to a simple vector input, some
property elements can be removed or masked in our approach.

Figure 6-(b) shows our embedding procedure for the input PV. Each property
element in the PV is a numerical value and normalized with the mean and
standard deviation of that property. The order of these 53 properties is
predetermined. Each value in the PV is encoded to a feature vector using a
linear layer as a value encoding. Then we randomly replace 50% of the property
features into the [UNK] token, which is the special token utilized to simulate
that the property is unknown. This is possible since there is no problem in
describing a molecule using only a part of these properties. Random property
feature masking prevents the model from overly dependent on the specific
property, has the effect of data augmentation, and improves the model’s
generalization ability. Although every property we used in this work can be
easily and thoroughly prepared by the computer, this might not be the case for
other properties in real-world situations. SPMM still can be trained when some
properties for certain training molecules are not known, by replacing those
unknown properties with the [UNK] token. On top of the randomly-masked value
encoding, we added a positional encoding similar to that in BERT. Since a PV
explicitly contains the values only, this positional embedding provides
information about what property each value corresponds to. Also, because of
the pre-defined order of these properties, this position embedding is
equivalent to giving a unique index for each property and adding an embedding
of that corresponding index. Then we pass the final result to the PV encoder
with the [CLS]P token.

###  0.4 Pre-training objectives.

Contrastive learning aims to learn better unimodal representation by aligning
the features from different modalities into the same feature space[26]. When
the encoded features of [CLS] tokens of SMILES $S$ and PV $P$ are given as
$S_{cls}$ and $P_{cls}$, we calculate the similarity function $sim(S,P)$ and
$sim(P,S)$ as:

$\displaystyle sim(S,P)=(h_{S}(S_{cls}))^{\intercal}h_{P}(P_{cls}),\quad
sim(P,S)=(h_{P}(P_{cls}))^{\intercal}h_{S}(S_{cls})$ (1)

where $h_{S}$ and $h_{P}$ are the linear projection + normalization layer for
SMILES and property vector, respectively. Now, for a given pair of $S$ and
$P$, we calculate the SMILES-to-PV and PV-to-SMILES intermodal similarities as
follows[27, 26]:

$\displaystyle
s_{s2p}={\exp{(sim(S,P)/\tau)}\over\Sigma_{n=1}^{N}{\exp{(sim(S,P_{n})/\tau})}},\quad
s_{p2s}={\exp{(sim(P,S)/\tau)}\over\Sigma_{m=1}^{M}{\exp{(sim(P,S_{m})/\tau})}}$
(2)

where $M$ and $N$ are the total numbers of SMILES and PV used in the loss
calculation. Here, $\tau$ is a learnable temperature parameter, which has a
sharpening effect by exaggerating the similarity difference. The intramodal
similarities can be calculated in the same way.

$\displaystyle
s_{s2s}={\exp{(sim(S,S)/\tau)}\over\Sigma_{m=1}^{M}{\exp{(sim(S,S_{m})/\tau})}},\quad
s_{p2p}={\exp{(sim(P,P)/\tau)}\over\Sigma_{n=1}^{N}{\exp{(sim(P,P_{n})/\tau})}}$
(3)

The overall contrastive loss is defined using the cross-entropy loss $H$ and
one-hot similarity $y$, which contains 1 for pairs originating from the same
molecule and contains 0 otherwise.

$\displaystyle L_{contrastive}={1\over
2}(H(y_{s2p},s_{s2p})+H(y_{p2s},s_{p2s})+H(y_{s2s},s_{s2s})+H(y_{p2p},s_{p2p}))$
(4)

Following the recent contrastive loss application in VLP[65], we build the
SMILES and PV queues that store the $k$ most recent SMILES and PV instances
and use them for contrastive loss. We set our queue size $k$ to 32,768.

Next Word Prediction (NWP) trains the model to predict the $(n+1)$-th SMILES
token when $0\sim n$-th tokens and the corresponding PV are given. Predicting
the next token is a common objective for training language models, known for
being utilized in the pre-training of GPT[66]. This can be done with a single
flow for each SMILES by applying a causal mask in the self-attention of the
SMILES encoder and the fusion encoder. Let $S=\\{s_{0},s_{1},\dots,s_{n}\\}$
and $P$ denote the input SMILES and the corresponding PV, and
$p^{NWP}(s_{n}|s_{0:n-1},P)$ denote the model’s predicted probability
distribution of the $n$-th token with given $P$ and $0\sim(n-1)$-th SMILES
tokens. The loss for NWP is defined as follows:

$\displaystyle
L_{NWP}=\sum^{n}_{i=1}H(y_{n}^{NWP},p^{NWP}(s_{n}|s_{0:n-1},P))\ $ (5)

where $y_{n}^{NWP}$ is a one-hot label for the $n$-th SMILES token $s_{n}$.

We applied a similar concept of NWP for the property vector as Next Property
Prediction (NPP). NPP makes the model predict the next property value using
its corresponding SMILES and the previous properties. Since each property
element is a numerical value, we replaced the cross-entropy loss in NWP with
mean-square-error loss. When $S$ and $P=\\{p_{0},p_{1},\dots,p_{n}\\}$ denotes
the input SMILES-PV pair and $\hat{p_{n}}(p_{0:n-1},S)$ denotes the model’s
predicted next property values with causal mask in the PV and the fusion
encoder, the loss for NPP is given as follows:

$\displaystyle L_{NPP}=\sum^{n}_{i=1}(p_{n}-\hat{p_{n}}(p_{0:n-1},S))^{2}$ (6)

In NPP, the model does not predict the property value if it is replaced with
[UNK] token.

SMILES-Property Matching (SPM) learns if a given SMILES-PV pair $(S,P)$ is
matched or not. We concatenate the feature of [CLS]S and [CLS]P token from the
fusion encoder output and pass this through a linear-layer SPM head. When
$p^{SPM}(S,P)$ is the output of the SPM head, the SPM loss can be defined as

$\displaystyle L_{SPM}=H(y^{SPM},p^{SPM}(S,P))$ (7)

where $y^{SPM}$ is a one-hot vector for a binary label of SMILES-PV matching;
the label is 1 if $S$ and $P$ originated from the same molecule and 0
otherwise. To build negative samples for SPM, we randomly select a “negative”
pair for each SMILES and PV instance from the other modality and match them as
negative pairs. This negative pair selection was done by hard-negative mining,
which gives a higher chance of being selected as a negative pair for instances
that has a higher similarity of Eq. (2) but isn’t a positive match. This makes
the training more difficult and forces the model to learn how to distinguish
similar instances.

In contrastive learning, using a one-hot label could be too strict since it
regards all instances that came from other pairs as equally-negative
instances. However, some PVs might agree with many SMILES, not only one SMILES
that they’re paired with. Even SMILES can be matched with different PVs since
there’s a 50% of masking in a PV (e.g., “ _MW_ =[UNK], _logP_ =2.1, _#atom_
=12” and “ _MW_ =78, _logP_ =2.1, _#atom_ =[UNK]” both explain Benzene, even
if they came from different molecules). A similar problem also occurs for NWP.
Sometimes there could be multiple sensible options for being the next token,
but using a one-hot label for ground truth might ignore this.

To resolve this issue, we built the momentum teacher model[65, 27] and
utilized its output for contrastive learning and NWP. The momentum teacher
performs a knowledge distillation by providing a pseudo-label that reflects
how the teacher model comprehends. Specifically, the label for the contrastive
learning and NWP are mixed with the momentum model’s output
$s_{*,momentum}(*\in\\{s2p,p2s,s2s,p2p\\})$ and
$p^{NWP}_{momentum}(s_{n}|s_{0:n-1},P)$, with an adjusting hyperparameter
$\alpha$. The detailed formulas for utilizing the momentum model for
contrastive learning and NWP are described in Eq. (8)$\sim$(9) and Eq.
(10)$\sim$(11).

$\displaystyle\tilde{y}_{*}$ $\displaystyle=(1-\alpha)y_{*}+\alpha
s_{*,momentum}\quad(*\in\\{s2p,p2s,s2s,p2p\\})$ (8)
$\displaystyle\tilde{L}_{contrastive}$ $\displaystyle={1\over
2}(H(\tilde{y}_{s2p},s_{s2p})+H(\tilde{y}_{p2s},s_{p2s})+H(\tilde{y}_{s2s},s_{s2s})+H(\tilde{y}_{p2p},s_{p2p}))$
(9) $\displaystyle\tilde{y}_{n}^{NWP}$
$\displaystyle=(1-\alpha)y_{n}^{NWP}+\alpha
p^{NWP}_{momentum}(s_{n}|s_{0:n-1},P)$ (10) $\displaystyle\tilde{L}_{NWP}$
$\displaystyle=\sum^{n}_{i=1}H(\tilde{y}_{n}^{NWP},p^{NWP}(s_{n}|s_{0:n-1},P))$
(11)

After the student model’s parameters ${w_{model}}$ are updated for each batch,
the parameters of the momentum teacher model ${w_{momentum}}$ are updated by
the exponential moving average (EMA) using ${w_{model}}$ and an EMA
hyperparameter $\lambda$ according to Equation 12.

$\displaystyle w_{momentum}=(1-\lambda)w_{model}+\lambda w_{momentum}$ (12)

The overall pre-training objective is the combined loss of Contrastive, NWP,
NPP, and SPM loss.

$\displaystyle L=\tilde{L}_{contrastive}+\tilde{L}_{NWP}+L_{NPP}+L_{SPM}$ (13)

###  0.5 Training for downstream tasks.

Supplementary Figure S3 describes how we utilized our pre-trained model for
downstream tasks. For PV generation and SMILES generation (Supplementary
Figure S3-(a), (b)), we don’t need additional fine-tuning since their training
objectives are already included in the pre-training (NWP, NPP). For the
inference procedure, the model generates PV or SMILES with autoregressive
sampling. Specifically, starting from the [CLS] token of the modality that we
want to generate, the model predicts the first component and repeats taking
the previous outputs to predict the next component until it’s done or meets a
sign to stop. A causal mask has to be used in the self-attention of the fusion
encoder and the unimodal encoder of the generating modality to enforce the
autoregressive generation.

For MoleculeNet downstream tasks and the DILI classification task that only
provide SMILES data, we utilized only the SMILES encoder part of the model
(Supplementary Figure S3-(c)). After the input molecule is encoded with the
SMILES encoder, we pass the feature of the [CLS]S token through a
classification/regression head to get an output. The classification/regression
head consists of MLP with one hidden layer. We fine-tuned our model with the
given training set and get a checkpoint with the lowest loss on the validation
set, and recorded that checkpoint’s performance on the test set.

The forward reaction prediction task provides a reactant SMILES (including
multiple reagent molecules) and a product SMILES. We encode these two inputs
with the SMILES encoder, then feed them into the fusion encoder + prediction
head. The model is trained to autoregressively generate the original product
SMILES (Supplementary Figure S3-(d)). In the inference stage, starting from
the [CLS]S token, the model predicts the next token until it generates the
[SEP] token. Similar to the SMILES generation, the self-attention of the
fusion encoder and the reactant SMILES encoder uses a causal mask. The retro-
reaction prediction task was done in the same way, but the role of the
reactant and product SMILES are swapped. We fine-tuned SPMM for the forward
reaction prediction task with an approach of ‘mixed task’, meaning that the
information about the major reactant is not given to the model. For both
forward and retro-reaction tasks, we replaced the input reactants and products
with their random non-canonical augmented SMILES[67] with a probability of
0.5.

###  0.6 Data preparation.

We obtained 20,000,000 SMILES of general molecules from PubChem[40] for pre-
training. All 53 properties we used can be calculated with SMILES using the
RDKit Python module[64]. The dataset for the MoleculeNet downstream tasks is
provided by the DeepChem[68] python library. We split every dataset into
train/valid/test sets in a ratio of 8:1:1 using a scaffold splitter from
DeepChem, which is a more harsh condition for the model than random splitting.
For the reaction prediction task, we used the USPTO-480k dataset which
contains 479,035 pairs of reactants and the major product of their reaction.
The retro-reaction prediction task used the USPTO-50k dataset, containing
50,037 product-reactant pairs with corresponding reaction types. Although the
USPTO-50k dataset provides tags of reaction type for each reaction data, we
didn’t use them, following the previous retro-reaction prediction
publications.

###  0.7 Implementation details.

We employed the architecture of 6 BERTbase encoder layers for our PV encoder
and SMILES encoder, and 6 BERTbase encoder layers with cross-attention layers
for our fusion encoder. With given $Q\in\mathbb{R}^{len_{q}\times
d_{k}},K\in\mathbb{R}^{len_{k}\times d_{k}},$ and
$V\in\mathbb{R}^{len_{k}\times d_{v}}$ as query, key, and value inputs, the
self-attention and cross-attention layers in BERT compute the output of the
scaled-dot attention according to the following formula:

$\displaystyle Attention(Q,K,V)=Softmax({QK^{T}\over\sqrt{d_{k}}})V$ (14)

We pre-trained the model until it converges using a batch size of 128 and the
AdamW optimizer with a weight decay of 0.02. The learning rate is warmed up to
${1e-4}$ and decreased to ${1e-5}$ with a cosine scheduler. We used the
momentum-adjusting hyperparameter $\alpha$ of 0.4. Since the pseudo-label from
the momentum teacher is not useful in the early stages of the training, we
linearly increased $\alpha$ from 0 to 0.4 during the first epoch. The EMA
hyperparameter $\lambda$ was fixed to 0.995, and the size of the PV and SMILES
queue $k$ was set to 32,768. The momentum models are not used for downstream
tasks. The full description of training for downstream tasks is in
Supplementary Table S1.

Correspondence and requests for materials should be addressed to Jong Chul Ye.
(email: jong.ye@kaist.ac.kr).

This work was supported by the Institute of Information & communications
Technology Planning & Evaluation (IITP) grant funded by the Korea government
(MSIT) (No.2019-0-00075, Artificial Intelligence Graduate School Program
(KAIST)), National Research Foundation (NRF) of Korea grant
NRF-2020R1A2B5B03001980, and by the KAIST Key Research Institute
(Interdisciplinary Research Group) Project.

J.C. prepared the code, performed all experiments and analyses, collected
data, and wrote the manuscript. J.C.Y. supervised the project in conception
and discussion, and prepared the manuscript.

The authors declare that they have no competing financial interests.

## Data Availability

The pre-training dataset is publicly available in PubChem
https://pubchem.ncbi.nlm.nih.gov/. ZINC15 test molecules for SMILES-to-PV
generation are accessible through the ZINC15 website. Scaffold-split
MoleculeNet datasets are available via DeepChem python module
https://deepchem.io/, and raw databases can be found in the MoleculeNet
website https://moleculenet.org/. The DILI training and test data preparation
can be found in https://pubmed.ncbi.nlm.nih.gov/29788510/. The USPTO-480k and
USPTO-50k dataset is available at https://github.com/wengong-jin/nips17-rexgen
and https://github.com/MolecularAI/Chemformer/tree/main. All of these data
that we’ve used in this work are also uploaded at
https://github.com/jinhojsk515/SPMM/.

## Code Availability

The source code for SPMM, a list of 53 properties for PV, experimental codes,
and datasets are available at https://github.com/jinhojsk515/SPMM/ to allow
reproducing the results.

## References

## References

  * [1] Ryu, S. & Lee, S.  Accurate, reliable and interpretable solubility prediction of druglike molecules with attention pooling and bayesian learning (2022).
  * [2] Kuenneth, C. & Ramprasad, R.  polybert: A chemical language model to enable fully machine-driven ultrafast polymer informatics (2022).
  * [3] Moon, S., Zhung, W., Yang, S., Lim, J. & Kim, W. Y.  Pignet: A physics-informed deep learning model toward generalized drug–target interaction predictions (2022).
  * [4] Xu, C., Wang, Y. & Farimani, A. B.  Transpolymer: A transformer-based language model for polymer property predictions.  _arXiv.org_ (2022).
  * [5] Paul, D. _et al._ Artificial intelligence in drug discovery and development.  _Drug Discovery Today_ (2020).
  * [6] Chen, S. & Jung, Y.  A generalized-template-based graph neural network for accurate organic reactivity prediction (2022).
  * [7] Segler, M. H. S., Kogej, T., Tyrchan, C. & Waller, M. P.  Generating focused molecule libraries for drug discovery with recurrent neural networks.  _ACS Central Science_ 4, 120–131 (2017).
  * [8] Gómez-Bombarelli, R. _et al._ Automatic chemical design using a data-driven continuous representation of molecules.  _ACS Central Science_ 4, 268–276 (2018).
  * [9] Lim, J., Ryu, S., Kim, J. W. & Kim, W. Y.  Molecular generative model based on conditional variational autoencoder for de novo molecular design (2018).
  * [10] Lim, J., Hwang, S.-Y., Moon, S., Kim, S. & Kim, W. Y.  Scaffold-based molecular design with a graph generative model.  _Chem. Sci._ 11, 1153–1164 (2020).
  * [11] Wang, W., Wang, Y., Zhao, H. & Sciabola, S.  A transformer-based generative model for de novo molecular design (2022).
  * [12] Olivecrona, M., Blaschke, T., Engkvist, O. & Chen, H.  Molecular de-novo design through deep reinforcement learning.  _Journal of Cheminformatics_ 9 (2017).
  * [13] Zhou, Z., Kearnes, S., Li, L., Zare, R. N. & Riley, P.  Optimization of molecules via deep reinforcement learning (2019).
  * [14] Chithrananda, S., Grand, G. & Ramsundar, B.  Chemberta: Large-scale self-supervised pretraining for molecular property prediction.  _arXiv preprint arXiv:2010.09885_ (2020).
  * [15] Zhang, Z. _et al._ Can pre-trained models really learn better molecular representations for ai-aided drug discovery?  _arXiv:2209.07423_ (2022).
  * [16] Melnyk, I. _et al._ Reprogramming large pretrained language models for antibody sequence infilling (2022).
  * [17] Bommasani, R. _et al._ On the opportunities and risks of foundation models (2021).
  * [18] Horawalavithana, S. _et al._ Foundation models of scientific knowledge for chemistry: Opportunities, challenges and lessons learned.  In _Proceedings of BigScience Episode #5 – Workshop on Challenges & Perspectives in Creating Large Language Models_, 160–172 (Association for Computational Linguistics, virtual+Dublin, 2022).
  * [19] Liu, S., Demirel, M. F. & Liang, Y.  N-gram graph: Simple unsupervised representation for graphs, with applications to molecules (2018).
  * [20] Rong, Y. _et al._ Self-supervised graph transformer on large-scale molecular data (2020).
  * [21] Ahmad, W., Simon, E., Chithrananda, S., Grand, G. & Ramsundar, B.  Chemberta-2: Towards chemical foundation models (2022).
  * [22] Chen, F. _et al._ Vlp: A survey on vision-language pre-training (2022).
  * [23] Vaswani, A. _et al._ Attention is all you need (2017).
  * [24] Chen, Y.-C. _et al._ Uniter: Universal image-text representation learning (2019).
  * [25] Li, X. _et al._ Oscar: Object-semantics aligned pre-training for vision-language tasks (2020).
  * [26] Radford, A. _et al._ Learning transferable visual models from natural language supervision (2021).
  * [27] Li, J. _et al._ Align before fuse: Vision and language representation learning with momentum distillation (2021).
  * [28] Yu, J. _et al._ Coca: Contrastive captioners are image-text foundation models (2022).
  * [29] Agrawal, A. _et al._ Vqa: Visual question answering.  _arXiv:1505.00468 [cs]_ (2016).
  * [30] Kaur, P., Pannu, H. S. & Malhi, A. K.  Comparative analysis on cross-modal information retrieval: A review.  _Computer Science Review_ 39, 100336 (2021).
  * [31] Zhang, C., Zhang, C., Zhang, M. & Kweon, I. S.  Text-to-image diffusion models in generative ai: A survey.  _arXiv:2303.07909 [cs]_ (2023).
  * [32] Ming, Y. _et al._ Visuals to text: A comprehensive review on automatic image captioning.  _IEEE/CAA Journal of Automatica Sinica_ 9, 1339–1365 (2022).
  * [33] Winter, R., Montanari, F., Noé, F. & Clevert, D.-A.  Learning continuous and data-driven molecular descriptors by translating equivalent chemical representations.  _Chem. Sci._ 10, 1692–1701 (2019).
  * [34] Grill, J.-B. _et al._ Bootstrap your own latent: A new approach to self-supervised learning (2020).
  * [35] Zhu, J. _et al._ Dual-view molecule pre-training (2021).
  * [36] Meyers, J., Fabian, B. & Brown, N.  De novo molecular design and generative models.  _Drug Discovery Today_ (2021).
  * [37] Bort, W. _et al._ Inverse qsar: Reversing descriptor-driven prediction pipeline using attention-based conditional variational autoencoder.  _Journal of Chemical Information and Modeling_ 62, 5471–5484 (2022).
  * [38] Crawshaw, M.  Multi-task learning with deep neural networks: A survey (2020).
  * [39] Park, S., Lee, E. S., Shin, K. S., Lee, J. E. & Ye, J. C.  Self-supervised co-learning of uncurated images and reports enables oversight ai in radiology (2022).
  * [40] Kim, S. _et al._ Pubchem in 2021: new data content and improved web interfaces.  _Nucleic Acids Research_ 49, D1388–D1395 (2021).
  * [41] Sterling, T. & Irwin, J. J.  Zinc 15 – ligand discovery for everyone.  _Journal of Chemical Information and Modeling_ 55, 2324–2337 (2015).
  * [42] Wu, Z. _et al._ Moleculenet: a benchmark for molecular machine learning.  _Chem. Sci._ 9, 513–530 (2018).
  * [43] Gogishvili, D., Nittinger, E., Margreitter, C. & Tyrchan, C.  Nonadditivity in public and inhouse data: implications for drug design 13 (2021).
  * [44] Kwapien, K. _et al._ Implications of additivity and nonadditivity for machine learning and deep learning models in drug design.  _ACS Omega_ 7, 26573–26581 (2022).
  * [45] Ai, H. _et al._ Predicting drug-induced liver injury using ensemble learning methods and molecular fingerprints.  _Toxicological Sciences_ 165, 100–107 (2018).
  * [46] Yang, K. _et al._ Analyzing learned molecular representations for property prediction.  _Journal of Chemical Information and Modeling_ 59, 3370–3388 (2019).  PMID: 31361484.
  * [47] Hu, W. _et al._ Strategies for pre-training graph neural networks (2019).
  * [48] Fang, X. _et al._ Geometry-enhanced molecular representation learning for property prediction.  _Nature Machine Intelligence_ 4, 127–134 (2022).
  * [49] Ross, J. _et al._ Large-scale chemical language representations capture molecular structure and properties.  _Nature Machine Intelligence_ 4, 1256–1264 (2022).
  * [50] Schwaller, P. _et al._ Molecular transformer: A model for uncertainty-calibrated chemical reaction prediction.  _ACS Central Science_ 5, 1572–1583 (2019).
  * [51] Tetko, I. V., Karpov, P., Van Deursen, R. & Godin, G.  State-of-the-art augmented nlp transformer models for direct and single-step retrosynthesis.  _Nature Communications_ 11 (2020).
  * [52] Irwin, R., Dimitriadis, S., He, J. & Bjerrum, E. J.  Chemformer: a pre-trained transformer for computational chemistry.  _Machine Learning: Science and Technology_ 3, 015022 (2022).
  * [53] Tu, Z. & Coley, C. W.  Permutation invariant graph-to-sequence model for template-free retrosynthesis and reaction prediction.  _Journal of Chemical Information and Modeling_ 62, 3503–3513 (2022).
  * [54] Sacha, M. _et al._ Molecule edit graph attention network: Modeling chemical reactions as sequences of graph edits.  _Journal of Chemical Information and Modeling_ 61, 3273–3284 (2021).
  * [55] Zheng, S., Rao, J., Zhang, Z., Xu, J. & Yang, Y.  Predicting retrosynthetic reactions using self-corrected transformer neural networks 60, 47–55 (2020).
  * [56] Kim, E., Lee, D.-S., Kwon, Y., Park, M.-S. & Choi, Y.-S.  Valid, plausible, and diverse retrosynthesis using tied two-way transformers with latent variables 61, 123–133 (2021).
  * [57] Lee, K., Jang, J., Seo, S., Lim, J. & Kim, W. Y.  Drug-likeness scoring based on unsupervised learning (2021).
  * [58] Devlin, J., Chang, M.-W., Lee, K. & Toutanova, K.  Bert: Pre-training of deep bidirectional transformers for language understanding (2018).
  * [59] Caron, M. _et al._ Emerging properties in self-supervised vision transformers.  _arXiv:2104.14294 [cs]_ (2021).
  * [60] Dosovitskiy, A. _et al._ An image is worth 16x16 words: Transformers for image recognition at scale.  _arXiv:2010.11929 [cs]_ (2020).
  * [61] Gage, P.  A new algorithm for data compression (1994).
  * [62] Liu, Y. _et al._ Roberta: A robustly optimized bert pretraining approach (2019).
  * [63] Raffel, C. _et al._ Exploring the limits of transfer learning with a unified text-to-text transformer (2019).
  * [64] Landrum, G.  Rdkit: Open-source cheminformatics software (2016).
  * [65] He, K., Fan, H., Wu, Y., Xie, S. & Girshick, R.  Momentum contrast for unsupervised visual representation learning (2019).
  * [66] Brown, T. B. _et al._ Language models are few-shot learners (2020).
  * [67] Bjerrum, E.  Smiles enumeration as data augmentation for neural network modeling of molecules (2017).
  * [68] Ramsundar, B. _et al._ _Deep Learning for the Life Sciences_ (O’Reilly Media, 2019).

## Supplementary Materials

Figure S1: The examples of the output molecules in the stochastic PV-to-SMILES
generations in the Fig. 2, by pre-trained SPMM with four different given PVs.

Figure S2: The scatter plots of the model’s generated PVs for 1,000 unseen
ZINC15 molecules, against the actual property value for all 53 properties. The
$r^{2}$ score and RMSE for each property are described at the top of each
plot.

Figure S3: Overview of the inference and fine-tuning of SPMM for various
downstream tasks: (a) The inference process of pre-trained SPMM for molecule
generation. (b) The inference process of pre-trained SPMM for PV generation.
(c) The model architecture for MoleculeNet downstream tasks. The SMILES
encoder of pre-trained SPMM is used as a backbone. (d) The model architecture
for the reaction prediction task. We adopted the SMILES encoder and the fusion
encoder of pre-trained SPMM and built a sequence-to-sequence model.

task | Delaney ESOL | LIPO | Freesolv | BACE(reg.) | BACE(cls.) | Clearance | BBBP | Clintox | SIDER | DILI
---|---|---|---|---|---|---|---|---|---|---
optimizer | AdamW, weight decay=0.02
scheduler | cosine + warmup
batch size | 4 | 16 | 4 | 8 | 8 | 8 | 16 | 4 | 4 | 8
learning rate(min, max) | 3e-5, 5e-6 | 3e-5, 5e-6 | 5e-5, 5e-6 | 2e-5, 2e-6 | 5e-5, 1e-5 | 2e-5, 5e-6 | 3e-5, 5e-6 | 3e-5, 5e-6 | 5e-5, 5e-6 | 2e-5, 5e-6
epoch | 25 | 25 | 30 | 30 | 10 | 10 | 10 | 15 | 10 | 20

task | forward reaction prediction | retro-reaction prediction
---|---|---
optimizer | AdamW, weight decay=0.02
scheduler | cosine + warmup
batch size | 16 | 16
learning rate(min, max) | 1e-4, 3e-6 | 1e-4, 5e-6
epoch | 100 | 300

Table S1: Detailed training hyperparameters for fine-tuning SPMM in
MoleculeNet downstream tasks, DILI classification task, forward reaction
prediction task, and the reverse-reaction prediction task.


# Learning Optimal Test Statistics in the Presence of Nuisance Parameters

Lukas Heinrich  Technical University of Munich<EMAIL_ADDRESS>
###### Abstract

The design of optimal test statistics is a key task in frequentist statistics
and for a number of scenarios optimal test statistics such as the profile-
likelihood ratio are known. By turning this argument around we can find the
profile likelihood ratio even in likelihood-free cases, where only samples
from a simulator are available, by optimizing a test statistic within those
scenarios. We propose a likelihood-free training algorithm that produces test
statistics that are equivalent to the profile likelihood ratios in cases where
the latter is known to be optimal.

##  1 Introduction

Statistical data analysis in the natural sciences is most often founded on a
probabilistic modelling of the underlying data-generating process, where
$p(x|\theta)$ denotes the probability of experimentally observing data $x$
given a set of theory parameters $\theta$. Inference aims at assessing the
theory space in light of the observed data in order to estimate points or
intervals in this space that are compatible with the data as well as test
hypotheses for data-driven decision-making. In frequentist statistics the main
tools for these tasks are estimates based on the well-developed methodology of
maximum-likelihood estimation, confidence intervals construction and test
statistics. In a Bayesian context, most inference tasks derive their results
from methods that aim to compute posterior densities of the form
$p(\theta|x)$. A major problem for both approaches, however, are _likelihood-
free_ settings, i.e. experimental situations where samples $x\sim p(x|\theta)$
are available but evaluating the likelihood $p(x|\theta)$ is computationally
intractable. The field of _likelihood-free inference_ thus aims to develop
methods that allow us to still perform the desired inference tasks without
requiring explicit evaluation of the model. High-Energy Physics data analysis
is a prominent example of such a likelihood-free problem, which appears due to
a rich, but unobservable evolution of the original particle collision through
many latent intermediate states $z_{i}$ culminating into a high-dimensional
measurement $x$. While the evolution probability itself is $p(x,z|\theta)$ is
tractable, the model of the observable data
$p(x|\theta)=\int\mathrm{d}z\,p(x,z|\theta)$ is not. Classical approaches to
_likelihood-free inference_ often use simulation and summary statistics $f(x)$
to derive a low-dimensional approximate model $\hat{p}(f(x)|\theta)$ to which
the standard methodology can then be applied. More recently a new breed of
methods are developed that aim to use machine learning to eschew an explicit
approximation of the statistical model, in favor of directly targeting only
the model-derived quantities required for inference. In this work we add to
this program by presenting a method to learn a test statistic with _best
average power_. For models which lie in the asymptotic regime, this is
equivalent to the _profile likelihood ratio_ test statistic, which is a key
quantity in frequentist data analysis for models that incorporate systematic
uncertainties through nuisance parameters.

##  2 Related Work

Likelihood-free and simulation-based inference using machine-learning methods
are growing field of research, where inference approaches for both Bayesian
and frequentist settings are studied [1, 2]. The present work is most closely
connected to the carl method of calibrated binary classifiers [3], where
parametrized networks $t(x|\theta,\theta_{0})$ are trained that are shown to
be equivalent to the likelihood-ratio test statistic
$p(x|\theta)/p(x|\theta_{0})$. It is also closely connected to the ACORE [4]
of likelihood-free hypothesis testing, but in our approach, we target
particular proposal distribution designed to asymptotically recover the
profile-likelihood through optimization. In addition, other approaches aim for
parametrizing just a subset of model parameters to train an optimal observable
but without using the resulting classifier as a direct proxy of a test
statistic [5]. In our approach we target a test statistic that differentiates
between parameters of interest $\mu$ and nuisance parameters $\nu$ and the
resulting network is only parametrized by $\mu$ as a result. In a larger
context this work is part of the program of exploring differentiable
programming for high-energy physics inference [6, 7, 8] as it provides an
approximation of a key inference-level quantity which is differentiable with
respect to upstream analysis parameters.

##  3 Likelihood Ratio Tests for Nested Hypotheses

###  3.1 Frequentist Hypothesis Testing

Frequentist hypothesis testing and interval estimation aims at analyzing the
data $x$ through the lens of a set hypotheses $H_{i}$. In binary testing the
goal is to assess whether to reject a null hypothesis $H_{0}$ in favor of an
alternative $H_{1}$. Tests are based on the comparison of the observed data to
the sampling distribution $p(x|\theta_{i})$ of the data under various theories
$\theta_{i}$. Through an analysis of these distribution a _rejection region_
$\omega$ is defined in the data space and $H_{0}$ is rejected if the observed
data lies within that region. Practically, such regions are implicitly defined
as the regions where a chosen _test statistic_ $t(x)$ exceeds a threshold
value $t_{0}$: $\omega=\\{x|t(x)>t_{0}\\}$ and the choice of $t(x)$ directly
affects the performance characteristics such as the achievable _power_
$\beta=p(x\in\omega|H_{1})$ of the test at a given test size
$\alpha=p(x\in\omega|H_{0})$ . The search for _optimal_ test statistics is
therefore a priority in frequentist statistics.

###  3.2 Optimal Test Statistics

For simple hypotheses, where hypotheses correspond uniquely to a specific
parameter points, the Neyman-Pearson Lemma [9] states that the universally
most powerful (UMP) test is given by the (log-)likelihood ratio

$t(x)=-2\log\frac{p(x|\theta_{0})}{p(x|\theta_{1})},$ (1)

where $\theta_{0}$ and $\theta_{1}$ are the parameters of the null and
alternative hypothesis, respectively. Composite hypotheses, where hypotheses
correspond to sets in parameter space, generally do not admit UMP tests.
However optimal tests can be found in more restricted cases. An important
scenario is that of _nested hypotheses_ , where the null hypothesis is
completely contained in the alternative. A recurring setup is one where the
full $n$-dimensional parameter space $\theta\in\mathbb{R}^{n}$ is partitioned
into $k$ _parameters of interest_ $\mu=\theta_{1:k}$ and $n-k$ _nuisance
parameters_ $\nu=\theta_{k+1:n}$. The null hypothesis is then identified as a
subspace on which the parameters of interest assume a certain value
$H_{0}=\\{(\mu,\nu)|\mu=\mu_{0}\\}$. Wald [10] has shown that in this case the
_profile likelihood ratio_ test statistic given by

$t_{\mu_{0}}(x)=-2\log\frac{p(x|\mu_{0},\hat{\vphantom{\rule{1.0pt}{5.71527pt}}\smash{\hat{\nu}}})}{p(x|\hat{\mu},\hat{\nu})}$
(2)

has a number of optimal properties in the asymptotic regime for any $\mu_{0}$.
Here, $\hat{\mu}$ and $\hat{\nu}$ denote the global maximum-likelihood
estimates, while $\hat{\vphantom{\rule{1.0pt}{5.71527pt}}\smash{\hat{\nu}}}$
denotes the maximum-likelihood estimate when the search space is restricted to
the null hypothesis set. In particular it can be shown to exhibit _best
average power_ under a suitable definition of the term, which we discuss in
the Section 4.1.

###  3.3 Likelihood-Free Learning of Test Statistics

An immediate corollary of the optimality properties of test statistics $t(x)$
such as the likelihood ratio and the profile likelihood ratio is that they can
be found by optimizing the parameters $\phi$ of a high-capacity function
approximation $s_{\phi}(x)$, such as a neural network, with respect to a loss
function which is minimized by the sought-after statistic $t(x)$. If the
solution is unique, the approximation $s(x;\phi)$ will naturally converge to
$t(x)$ with sufficient training. If a family of solutions exist, the
optimization may converge on any member of such a family. In particular,
evaluations of test statistics that rely on tail probabilities such as power
and size, are invariant under a bijective variable transform and thus a
solution $s(x)$ may be related to $t(x)$ through a bijective transform
$t=f(s)$. If an external approximation $t_{\mathrm{ext}}(x)$ of $t(x)$ is
available, for example through e.g. explicit density approximation of
$\hat{p}(x|\theta)$, the bijective calibration function $f$ can be found
between the two through e.g. isotonic regression on
$(s_{i},t_{i,\mathrm{ext.}})$ pairs. The procedure above yields a _likelihood-
free_ approximation of (a bijection of) $t(x)$ if the training process only
requires samples from the probability models $x\sim p(x|\theta)$. For the
simple hypotheses, the described approach is captured in the “likelihood ratio
trick” of binary classifiers trained on the binary cross-entropy as a loss
function [3]. The loss is minimized by a classifier that is monotonically
related to the likelihood-ratio. We will now consider the likelihood-free
learning for complex hypotheses and connect the results to asymptotic theory.

##  4 Learning Optimal Test Statistics for Nested Hypotheses

###  4.1 Best Average Power Statistics

We can extend the approach beyond simple hypotheses to find optimal test
statistics $s(x;\mu_{0})$ indexed by the values of the values $\mu_{0}$
assumed by parameters of interest for the null hypothesis set $H_{0}$ by
training a parametrized neural network $s_{\phi}(x;\mu_{0})$. As no UMP tests
are available a metric through which to assess optimality needs to be defined.
We choose to optimize for a measure of “best average power” as described by
Wald [10], which we briefly recapitulate below. A reasonable approach is to
define “best average power” with respect to “equally distant” alternatives.
For example in the 1-dimensional case with a single parameter of interest
$\mu$ and no nuisance parameters, a test of $H_{0}=\mu_{0}$ is desirable that
has similar power for alternatives $\mu_{\pm}=\mu_{0}\pm c\cdot
I_{\mu\mu}^{-1}$, that is for alternatives at a distance $c$ as measured
within the Fisher metric. Generalizing to the case of nuisance parameters and
higher dimensions we pick a given $\theta_{0}\in H_{0}$ and first consider the
a $k$-dimensional hyper-surface spanned by $n-k$ linear equations

$\displaystyle\gamma_{ij}\theta^{j}=d_{i}(\theta_{0}),\;\text{with}\;i=k+1\dots
n,\;j=1..n,$ (3)

where the constants $\gamma_{ij},d_{i}$ are chosen such that $\theta_{0}$ lies
within that surface. This surface slices through the parameters space and have
dimension $k$, which corresponds to the dimensionality of the parameters of
interest. Within this surface of parameter points, we can now consider
alternatives that are equidistant with respect to the $k$ parameters of
interest to $\theta_{0}$, by computing the distance using the Fisher metric at
$\theta_{0}$:

$\displaystyle(\mu^{i}-\mu_{0}^{i})(\mu^{j}-\mu_{0}^{j})I_{ij}(\theta_{0})=c,\;\text{with}\;i,j=1...k$
(4)

The surfaces $S_{c}(\gamma,d,c)$ are $k-1$ dimensional (hyper-)ellipsoids
lying in the linear hyperplane defined by the coefficients $\gamma$. To
compute the average power within the surface $S_{c}$ a density
$p(\theta_{\mathrm{alt}}|S_{c})$ is needed to define the integral

$\displaystyle\beta_{\mathrm{avg}}(s_{\phi},s_{0},S_{c})=\int_{S_{c}}\mathrm{d}\theta_{\mathrm{alt}}\,\beta(\theta_{\mathrm{alt}},s_{0},s_{\phi})\,p(\theta_{\mathrm{alt}}|S_{c}),$
(5)

where $\beta(\theta_{\mathrm{alt}},s_{0})$ is the power of the test based on
the test statistic $s_{\phi}$ and a rejection region defined by
$\omega=\\{x|s_{\phi}(x)>s_{0}\\}$. The metric Wald uses is one that is
uniform over the ellipsoids $S_{c}$. In general, computing the density on the
surfaces $S_{c}$ of alternatives, or sampling from it, requires knowledge of
the Fisher information. The somewhat complicated construction of $S_{c}$
significantly simplifies for the case of a single parameter of interest: the
surfaces are just the two points that lie at a fixed distance from the null
$(\mu-\mu_{0})^{2}=\mathrm{const.}$ and on a linear subspace where
$\gamma^{T}\theta=\mathrm{const}.$ holds. Thus in this case, detailed
knowledge of the Fisher Information is not necessary.

With Wald’s definition of surfaces of alternatives above we can define a
global measure of “best average power”. Given some densities over the space of
parameter points within a null hypothesis $p(\theta_{0}\in H_{0})$, of
surfaces “equidistant alternatives” $p(S_{c}|\theta_{0})$ and finally of
alternatives on those surfaces $p(\theta_{\mathrm{alt}}|S_{c})$, as

$\displaystyle\beta_{\mathrm{global}}(s_{\phi},s_{0})=\int\mathrm{d}\theta_{0}\,\mathrm{d}S_{c}\;\beta_{\mathrm{avg}}(s_{\phi},s_{0},S_{c})p(S_{c}|\theta_{0})p(\theta_{0}),$
(6)

and aim for finding a test statistic $s_{\phi}$ that has best average power
for any threshold value $s_{0}$. If there is a solution $s_{\phi}$ that is
optimal for any choice of $s_{0}$, $S_{c}$ and $\theta_{0}$ simultaneously, it
will also be optimal under _any choice of densities_ $p(\theta_{0})$,
$p(S_{c}|\theta_{0})$.

###  4.2 Likelihood-Free Optimization

We can transform our optimization goal of “best average power” across a
surface $S_{c}$ onto a likelihood-free optimization by realizing that
optimizing for $\beta_{\mathrm{avg}}$, that is the average power across a
density of alternatives $p(\theta_{\mathrm{alt}})$ is equivalent to optimizing
the power of a hypothesis test with null distribution $p(x|\theta_{0})$ and
simple alternative designed as a mixture model on the surface
$S_{c}(\theta_{0})$ given by

$p_{\mathrm{mix}}(x|S_{c})=\int\mathrm{d}\theta_{\mathrm{alt}}\;p(x|\theta_{\mathrm{alt}})p(\theta_{\mathrm{alt}}|S_{c}).$
(7)

By the Neyman-Pearson Lemma and the Likelihood-Ratio Trick, the test statistic
with the optimal power for a test of $H_{0}:p(x|\theta_{0})$ and
$H_{1}:p_{\mathrm{mix}}(x|S_{c})$ can be found by optimizing the binary cross-
entropy:

$\mathcal{L}_{\mathrm{BXE}}(\phi|S_{c},\theta_{0})=-\mathbb{E}_{(y,x\sim
p(x|\theta_{0}),p_{\mathrm{mix}}(x|S_{c})}[(y\log
s_{\phi}(x)+(1-y)\log(1-s_{\phi}(x))],$ (8)

where samples from $\theta_{0}$ are labeled $y=0$ and samples from the mixture
of alternatives are labelled $y=1$. To optimize the global power from Equation
6 in a likelihood-free way we can now average over $S_{c}$ and $\theta_{0}$ to
define a global loss $\mathcal{L}(\phi)$

$\displaystyle\mathcal{L}(\phi)=\mathbb{E}_{S_{c},\theta_{0}}[\mathcal{L}_{\mathrm{BXE}}(\phi|S_{c},\theta_{0})]$
(9)

given choice of densities
$p(\theta_{\mathrm{alt}}|S_{c}),p(S_{c}|\theta_{0})$, $p(\theta_{0})$. The
final optimization procedure is simple: for each minibatch, sample
$\theta_{0}$ and a random set of alternatives from $\theta_{\mathrm{alt}}$
from a random surface $S_{c}(\theta_{0})$, evaluate the neural network on data
$x\sim p(x|\theta_{i})$ and compute the average of the binary cross-entropy
losses for each alternative, where samples from $\theta_{0}$ are labeled $y=0$
and any samples from any $\theta_{\mathrm{alt}}$ are labelled $y=1$. With such
a procedure and average loss definition, the parameters $\phi$ can be
optimized via standard stochastic gradient descent. The algorithm is
summarized in Algorithm 1.

Algorithm 1 Training a Test Statistic with Best Average Power

1:$\eta$: learning rate

2:$\phi_{0}$: initial parameters

3:$\theta\sim p(\theta),\;\theta\sim p(\theta,S_{c}|\theta_{0})$: sampling
routines

4:while not converged do

5: $\theta_{0}=(\mu_{0},\nu_{0})\sim p(\theta)$$\triangleright$ sample null

6: $\theta_{i}=(\mu_{i},\nu_{i})\sim
p(\theta,S_{c}|\theta_{0})$$\triangleright$ sample alternatives

7: $(x_{i},y_{i})\sim p(x|\theta_{0}),p(x|\theta_{i})$ $\triangleright$ null:
$y_{i}=0$, all alternatives have $y_{i}=1$

8: $p_{i}\leftarrow s_{\phi}(x_{i};\mu_{0})$

9: $L=\sum_{\mathrm{null},\mathrm{alts}}L_{\mathrm{BXE}}(y_{i},p_{i})$

10: $\phi_{i+1}\leftarrow\phi_{i}-\eta\nabla_{\phi}L$

11:end while

12:return $\phi_{N}$

###  4.3 Inference and Asymptotic Alignment with the Profile Likelihood Ratio

In general the test statistic found through the training procedure above
depends on the choice of densities
$p(\theta_{\mathrm{alt}}|S_{c}),p(S_{c}|\theta_{0})$, $p(\theta_{0})$.
Furthermore, the sampling distributions $p(s|\theta)$ under a member of the
null hypothesis set $\theta\in H_{0}$ or its encompassing alternative set
$\theta\in H_{1}$ is unknown. This is no impediment for inference as the
sampling distributions for any $\theta$ can be found easily by evaluating the
learned test statistic $s_{\phi}$. Based on such empirical distributions a
consistent frequentist inference procedure for hypothesis testing as well as
point and interval estimation can be constructed.

Our choice of “best average power” as an optimization target, however, is not
an accident and an interesting connection appears in the limit of asymptotic
behavior. Due to the results to Wald mentioned in Section 3.2, we know that
the _profile likelihood ratio_ is the test statistic that has best average
power for _any_ surface $S_{c}$, $\theta_{0}$ and threshold value $s_{0}$ and
thus will also minimize the global measure power $\beta_{\mathrm{global}}$ as
well as the loss defined in Equation 9. In case of models $p(x|\theta)$ which
satisfy sufficient asymptotic properties, the likelihood-free search for a
best average power test statistic $s_{\phi}(x;\mu_{0})$ is expected to yield a
test statistic that is one-to-one to the profile likelihood ratio. We explore
this connection empirically below.

##  5 Experiments

We demonstrate the approach of implicitly learning a best average power test
statistic two examples that have tractable $p(x|\theta)$ and are known to be
within the asymptotic regime. This allows us to verify that the training
procedure finds the optimal solution as we can independently compute the
profile likelihood ratio. Furthermore, the calibration function relating
$s_{\phi}(x;\mu_{0})$ and $t_{\mu_{0}}(x)$ can be found through isotonic
regression. On the one hand we study a synthetic Gaussian Example with fixed
covariance matrix and on the other hand we study the ‘on-off’ problem
measuring two Poisson processes with shared parameters. In each case we train
neural network with three fully connected hidden layers and with width 100 and
subsequent tanh activation. The final output is passed through a sigmoid
activation and then evaluated through the binary cross-entropy as described
above. The networks are trained with the Adam optimizer [11] with learning
rate $\eta=5\cdot 10^{-5}$ and a mini-batch size of $n=1000$. We implement the
training in PyTorch [12] and provide the code to reproduce the figures on
GitHub 111See the repository at
https://github.com/lukasheinrich/learning_best_average_power. The examples are
partially chosen because closed-form solutions of the profile likelihood ratio
are available, which facilitates the comparison with ground-truth and the
derivation of calibration functions. We follow a simplified sampling scheme
for the null and alternative hypotheses suited for one dimension, where we
choose $\gamma_{11}=0$ and $\gamma_{12}=1$, such that the null and the
alternatives to the left and right of the null have the same nuisance
parameter value $\nu$, which is sampled according to a uniform distribution.
The distance $(\mu-\mu_{0})^{2}=\mathrm{const}.$ is also sampled uniformly for
each minibatch.

###  5.1 Gaussian Example

Figure 1: Results of the described method for a Gaussian example. On the left-
most column the trained test statistic distribution $s_{\phi}(x;\mu_{0})$
(top) and the true profile likelihood ratio (botttom) that closely follows
asymptotic theory. To the right, the ROC curves (top) of the trained
classifier, likelihood ratio, and asymptotic theory are shown to match well.
The joint distribution (bottom) suggests a monotonic relationship between the
two. In the middle, the learned test statistic is shown in data space for two
parameters values (top) and the true likelihood is shown for a example
observation together with the target profile. The right pane compares ground
truth and trained profile likelihood values as a function of $\mu$.

In this example we consider a bivariate Gaussian Model with an arbitrary mean
$\theta=(\mu,\nu)^{T}$ but fixed covariance $\Sigma$

$p(x_{1},x_{2}|\mu,\nu)=\mathcal{N}(\left[\begin{matrix}x_{1}\\\
x_{2}\end{matrix}\right]|\left[\begin{matrix}\mu\\\
\nu\end{matrix}\right],\Sigma)$

As asymptotic theory describes models in which maximum likelihood estimates
achieve unbiased normal distributions and their variance is described by the
Cramér-Rao bound, this example is a good proxy for such models and we expect
the optimal solution found in training to reproduce the profile likelihood
ratio (modulo bijection) to a very high degree of precision. Moreover, there
is a trivial relationship between data and the model parameters. The results
are shown in Figure 1. By comparing the ROC curves of the two test statistics,
we see that they are very compatible, which suggests a bijective relationship
between the learned test statistic and the true profile likelihood ratio. The
bijection is evident in the joint distribution of the (known) profile
likelihood ratio and the learned statistic from which we can extract a
calibration function through isotonic regression as implemented in scikit-
learn [13]. We can recast the learned statistic $s_{\phi}(x;\mu_{0})$ into
$t_{\mu_{0}}$-like units through either the calibration function or through
percentile-matching to the $\chi^{2}$ distribution to compare the trained
classifier to the true profile likelihood ratio. As shown in Figure 1, both
the network is an excellent approximation of the true profile likelihood also
in the uncalibrated case, where no ground-truth information is used.

###  5.2 On-Off Problem

In this example we study the classic "on-off" problem [14] of simultaneous
measurements of two Poisson processes:

$p(x_{1},x_{2}|\mu,\nu)=\mathrm{Pois}(x_{1}|\mu s+\nu
b)\mathrm{Pois}(x_{2}|\nu\tau b),$ (10)

where the hyperparameters $s,b$ indicate the nominally expected signal and
background counts. The hyperparameter $\tau$ describes the relationship in
measurement time between the two processes. The signal strength $\mu$ is the
parameter of interest while the scaling factor for the background $\nu$ is a
nuisance parameter. This model deviates from a pure Gaussian setup and
introduces a more complicated relationship between data and parameters. The
results are shown in Figure 2. With hyperparameters set at $s=15$, $b=70$,
$\tau=1$ the model is comfortably in the asymptotic regime but with $\tau\sim
1$ significant deviation the profiled values of the nuisance parameters is
expected as one moves away of the maximum-likelihood estimate. As in the
Gaussian case, the neural network directly approximates the profile likelihood
ratio to very high degree of precision.

Figure 2: Results of the described method for the on-off problem.

##  6 Conclusion

We have demonstrated that we can find powerful test statistics in a
likelihood-free way through optimizing for average cross-entropy, which was
shown to be equivalent to converege to best average power statistics. In the
limit, where the intractable model $p(x|\theta)$ displays asympptic behavior,
the learned test statistic was shown to be equivalent to the profile-
likelihood ratio. This equivalence was empirically shown to hold to a very
high degree of precision in two example cases that are known to behave
asymptotically. The trained classifier can be directly used for frequentist
inference tasks such as point and interval estimation. During inference, no
actual data-dependent minimization needs to be performed and as such the
network $s_{\phi}(x;\phi)$ is an _amortized inference_ tool, in which a large
one-time cost, i.e. the training, is traded off against fast instance-level
inference. It’s interesting to note that in cases where asymptotic assumptions
do not hold this procedure may produce test statistics that perform better
with respect to the average power metric than the profile likelihood ratio as
the latter may not be the optimal solution anymore.

## Acknowledgements

LH thanks Allen Caldwell, Kyle Cranmer, Nathan Simpson, Alexander Held and
Michael Kagan for fruitful discussions and comments on the manuscript. LH is
supported by the Excellence Cluster ORIGINS, which is funded by the Deutsche
Forschungsgemeinschaft (DFG, German Research Foundation) under Germany’s
Excellence Strategy - EXC-2094-390783311. Our code makes use of PyTorch [12],
Numpy [15], SciPy [16], Scikit-Learn [13], Matplotlib [17] and Jupyter [18].

## References

## References

  * [1] Cranmer K, Brehmer J and Louppe G 2020 Proceedings of the National Academy of Sciences 117 30055–30062 (Preprint https://www.pnas.org/doi/pdf/10.1073/pnas.1912789117) URL https://www.pnas.org/doi/abs/10.1073/pnas.1912789117
  * [2] Tejero-Cantero Á, Boelts J, Deistler M, Lueckmann J, Durkan C, Gonçalves P J, Greenberg D S and Macke J H 2020 CoRR abs/2007.09114 (Preprint 2007.09114) URL https://arxiv.org/abs/2007.09114
  * [3] Cranmer K, Pavez J and Louppe G 2015 (Preprint 1506.02169)
  * [4] Dalmasso N, Izbicki R and Lee A B 2020 URL https://arxiv.org/abs/2002.10399
  * [5] Ghosh A, Nachman B and Whiteson D 2021 Phys. Rev. D 104 056026 (Preprint 2105.08742)
  * [6] Heinrich L and Kagan M 2022 Differentiable Matrix Elements with MadJax (Preprint 2203.00057) URL https://arxiv.org/abs/2203.00057
  * [7] De Castro P and Dorigo T 2019 Comput. Phys. Commun. 244 170–179 (Preprint 1806.04743)
  * [8] Simpson N and Heinrich L 2022 neos: End-to-End-Optimised Summary Statistics for High Energy Physics (Preprint 2203.05570) URL https://arxiv.org/abs/2203.05570
  * [9] Neyman J and Pearson E S 1933 Philosophical Transactions of the Royal Society of London. Series A, Containing Papers of a Mathematical or Physical Character 231 289–337 ISSN 02643952 URL http://www.jstor.org/stable/91247
  * [10] Wald A 1943 Transactions of the American Mathematical Society 54 426–482 ISSN 00029947 URL http://www.jstor.org/stable/1990256
  * [11] Kingma D P and Ba J 2014 arXiv preprint arXiv:1412.6980
  * [12] Paszke A, Gross S, Massa F, Lerer A, Bradbury J, Chanan G, Killeen T, Lin Z, Gimelshein N, Antiga L, Desmaison A, Kopf A, Yang E, DeVito Z, Raison M, Tejani A, Chilamkurthy S, Steiner B, Fang L, Bai J and Chintala S 2019 Advances in Neural Information Processing Systems 32 ed Wallach H, Larochelle H, Beygelzimer A, d'Alché-Buc F, Fox E and Garnett R (Curran Associates, Inc.) pp 8024–8035 URL http://papers.neurips.cc/paper/9015-pytorch-an-imperative-style-high-performance-deep-learning-library.pdf
  * [13] Pedregosa F, Varoquaux G, Gramfort A, Michel V, Thirion B, Grisel O, Blondel M, Prettenhofer P, Weiss R, Dubourg V, Vanderplas J, Passos A, Cournapeau D, Brucher M, Perrot M and Duchesnay E 2011 Journal of Machine Learning Research 12 2825–2830
  * [14] Cowan G, Cranmer K, Gross E and Vitells O 2011 Eur. Phys. J. C 71 1554 [Erratum: Eur.Phys.J.C 73, 2501 (2013)] (Preprint 1007.1727)
  * [15] Harris C R, Millman K J, van der Walt S J, Gommers R, Virtanen P, Cournapeau D, Wieser E, Taylor J, Berg S, Smith N J, Kern R, Picus M, Hoyer S, van Kerkwijk M H, Brett M, Haldane A, del Río J F, Wiebe M, Peterson P, Gérard-Marchant P, Sheppard K, Reddy T, Weckesser W, Abbasi H, Gohlke C and Oliphant T E 2020 Nature 585 357–362 URL https://doi.org/10.1038/s41586-020-2649-2
  * [16] Virtanen P, Gommers R, Oliphant T E, Haberland M, Reddy T, Cournapeau D, Burovski E, Peterson P, Weckesser W, Bright J, van der Walt S J, Brett M, Wilson J, Millman K J, Mayorov N, Nelson A R J, Jones E, Kern R, Larson E, Carey C J, Polat İ, Feng Y, Moore E W, VanderPlas J, Laxalde D, Perktold J, Cimrman R, Henriksen I, Quintero E A, Harris C R, Archibald A M, Ribeiro A H, Pedregosa F, van Mulbregt P and SciPy 10 Contributors 2020 Nature Methods 17 261–272
  * [17] Hunter J D 2007 Computing in Science & Engineering 9 90–95
  * [18] Kluyver T, Ragan-Kelley B, Pérez F, Granger B, Bussonnier M, Frederic J, Kelley K, Hamrick J, Grout J, Corlay S, Ivanov P, Avila D, Abdalla S and Willing C 2016 Positioning and Power in Academic Publishing: Players, Agents and Agendas ed Loizides F and Schmidt B (IOS Press) pp 87 – 90