# Compressing Search with Language Models

Thomas Mulc  GoogleSunnyvaleCaliforniaU.S.A<EMAIL_ADDRESS>and  Jennifer L.
Steele  GoogleSunnyvaleCaliforniaU.S.A<EMAIL_ADDRESS>
###### Abstract.

Millions of people turn to Google Search each day for information on things as
diverse as new cars or flu symptoms. The terms that they enter contain
valuable information on their daily intent and activities, but the information
in these search terms has been difficult to fully leverage. User-defined
categorical filters have been the most common way to shrink the dimensionality
of search data to a tractable size for analysis and modeling. In this paper we
present a new approach to reducing the dimensionality of search data while
retaining much of the information in the individual terms without user-defined
rules. Our contributions are two-fold: 1) we introduce SLaM Compression, a way
to quantify search terms using pre-trained language models and create a
representation of search data that has low dimensionality, is memory
efficient, and effectively acts as a summary of search, and 2) we present
CoSMo, a Constrained Search Model for estimating real world events using only
search data. We demonstrate the efficacy of our contributions by estimating
with high accuracy U.S. automobile sales and U.S. flu rates using only Google
Search data.

##  1\. Introduction

Google Search is the predominant search engine worldwide, and as such there
exists unrivaled information relating the terms111A ”term” is defined as the
string entered into search by the user; it is generally 1-6 words long, but
may include numbers, symbols, etc. users search to real world events such as
consumer purchases, economic activity, or illness rates. There is already a
large corpus of research establishing the value of incorporating Google search
data into forecasting and predictive models (Lampos et al., 2015; Wang et al.,
2023; Varian and Choi, 2009; Ginsberg et al., 2009). These existing approaches
all create machine learning features by summarizing search data from a time
period (e.g., day), and then use these features to predict events (e.g.,
automobile sales). This prior research uses two forms of Google search data:
Google Trends and search logs.

Google Trends groups terms into search categories (such as ”Cold & Flu,” and
”Autos & Vehicles”) and returns an indexed value for the search volume in that
category for a particular day and geographic region. Trends and related
classification methods are a coarse signal of consumer interest, where queries
across a spectrum of intent are lumped under a single trend / category as if
they were identical; the dimensionality of this data is relatively small (due
to the relatively small number of categories) and is easily digestible for
most downstream machine learning applications. In (Varian and Choi, 2009;
Woloszko, 2020), they show the value of using this data to predict economic
activity (such as auto sales and home sales) and GDP. However, since Google
Trends data is coarse, their approaches rely on additional features such as
historical sales or other economic indicators.

Search logs contain pairs of search terms and their frequency (i.e., search
volume) over a given time period in a particular geographic area. Since the
number of unique terms is very large, modeling done using search logs requires
that the data for a given time period be summarized into a digestible format
and dimensionality for machine learning. The primary challenge with using raw
search logs in modeling has been to find a way to transform millions of
distinct textual search terms into useful and tractable features that can be
used by downstream machine learning. In (Ginsberg et al., 2009; Lampos et al.,
2015), they show that by aggressively filtering the search data and one-hot
encoding terms, you can create search features small enough for machine
learning. They demonstrate the efficacy of their approaches by modeling U.S.
flu rates using U.S. search logs.

In our work, we use these large search logs and create a new method to
summarize the search terms and their frequency that relies on language models
(LMs) to quantify each search term. Additionally, we create a custom model
tailored for predicting targets using search data.

In Section 2.1 we outline our framework, SLaM, for compressing search data
into tractable features for modeling. SLaM uses the embedding vectors
generated by LMs to retain the semantics of individual terms. The outputs of
SLaM are features we call “search embeddings.” This approach doesn’t rely on
user-defined filters and can be applied at any time granularity (e.g., daily
or weekly level), yielding an aggregated search embedding for each time period
that is memory efficient while being highly predictive of many events.

Our search embeddings are then incorporated into CoSMo, a constrained search
model (Section 2.2), which outputs a score between zero and one and can be
thought of loosely as the probability of the dependent variable occurring for
an average search term, whether that be the probability of a sale or the
probability of having the flu. In section 3 we outline some of the related
works, and how our approach fits into the existing literature.

Incorporating our search features and novel constrained search model into some
real world applications (Section 4), we are able to estimate the gains from
both parts of our approach. Using the reported U.S. flu rates and the U.S.
auto sector sales as case studies, we present results from nowcasting,
highlighting the model improvement from our language model approach compared
to more traditional Google Trends and classification methods. We find that
using our search embeddings increases predictive power by 30% in auto sales
compared to classification embeddings, and our method is on par or better than
existing autoregressive approaches for flu modeling, despite only using search
data as a model input.

Finally, these embeddings allow us to back out useful insights (Section 5)
into how consumers and patients use search by scoring the individual search
terms and highlighting terms with high scores (i.e. high probability of
purchase / having the flu). This is a new capability in search modeling,
because most classification methods (e.g., Google Trends) treat all terms
within a category as identical, which makes backing out the importance of any
one term impossible, while other classification methods that operate on the
individual term-level (e.g., one-hot encoding each term) cannot handle terms
outside of their very limited set of included terms.

While this paper is focused on Google Search, the approach could be used in
other settings where users are supplying a text input and we have a dependent
variable to estimate.

##  2\. Approach

We view modeling using search data as a two-step problem: 1. compressing /
aggregating search (feature engineering) 2. choosing an appropriate model
given the features to model the downstream target (model selection).

Our approach leverages LMs to collapse the query space to a tractable size
that retains information about the query semantics, without the need for
filters or manual data manipulation. Instead of using a binary classifier to
map a search term to a one-hot vector, we use an LM to map the term to a point
on the $D$-dimensional unit-sphere. We then aggregate the search terms along
these new $D$ dimensions, resulting in a search embedding that has
dimensionality $\mathcal{O}(D)$.

We design a model that takes search embeddings as the primary input and
outputs an estimate for the target variable. The model has inductive biases
and constraints that are specific to search data, the underlying distribution
of search embeddings, and the limited quantity of targets available for model
fitting.

###  2.1. SLaM Compression: Search Language Model Compression

Figure 1. SLaM inputs all searches during a given time period and compressed
them to a fixed-length vector that is effectively a summary of all search
terms. (Left) Each search term is passed through a language model that
produces a fixed-length vector of size $D$. Colors represent unique search
terms while shadings represents different embedding dimensions. (Right) All
the $D$-length vectors are passed to the aggregation step, where they are
reduced to a single vector, the search embedding, of size $\mathcal{O}(D)$,
which is later used as a feature for modeling.

Our approach for compressing search data aims to retain much of the
information from the raw query counts without an explosion in dimensionality.
We do so by leveraging the fixed-length representations learned from language
models, which map search terms that have similar semantic meaning near one-
another in a $D$-dimensional embedding space (Kiros et al., 2015). We name our
framework of using language models to compress search ”SLaM Compression:”
Search Language Model Compression. At a high level, SLaM aims to map
individual terms to a fixed-length embedding using a language model, then
aggregate the embedding statistics to remove the individual terms from the
dimensionality (see Figure 1). Our specific implementation of the compression
is derived by analysing the regression of a linear model fit on top of LM
embeddings, which we leave to Appendix A. Intuitively, our method is a
weighted sum of the LM embeddings for time period $t$, where the weights are
the number of search counts for each term $s$. Mathematically it is simply

(1) $\gamma_{t}=\sum_{s\in S}v_{s,t}\cdot LM(s)$

where $S$ is the set of unique search terms, $LM$ is a language model mapping
that maps each search term $s\in S$ to a fixed length $D$-dimensional vector,
and $v_{s,t}$ is the number of times search term $s$ was queried in period
$t$.

In practice $||LM(s)||_{2}=1$ for LM embeddings, and we assume this to be the
case in our method. While we use this simple summation as the way to aggregate
the embeddings, note that other aggregation techniques that preserve some the
statistics of the embeddings can be used (e.g, binned marginal distributions;
see Appendix G). This is why we say SLaM compresses the space down to
$\mathcal{O}(D)$ instead of $D$.

Note that our compression makes no assumption that terms are filtered, and the
summation happens over all search terms included in the feature set. This is
an important characteristic of our method that allows it to scale to larger
sets of terms while freeing the modeler from the burden of feature
engineering.

We decompose our representation into two parts, the total daily query volume
$V_{t}$ and the normalized search embedding $\gamma^{*}_{t}$ where

(2) $\gamma^{*}_{t}=\gamma_{t}/||\gamma_{t}||_{2}.$

Like $LM(s)$, $\gamma^{*}_{t}$, is also on the unit sphere, and can be viewed
as the weighted average of the embeddings projected back onto the sphere.
Building models using $\gamma_{t}^{*}$ allows us to run inference on the
individual search terms, because both $\gamma_{t}^{*}$ and $LM(s)$ are
distributed on the unit-sphere.

In practice, we use two additional inputs to compute our search embedding:
geography and search category. For search embedding $\gamma_{t,r,c}$, only
searches during time period $t$ that happen in geography region $r$ and belong
to category $c$ are included. We match the search geography restriction with
the target variable geographic granularity, so that search data from a
specific region is used to predict the target in that region. The category
restriction shrinks the search data down to a size that is reasonable to
compute over. For example, in our auto case study Vermont auto category
searches predict Vermont auto sales. Even with these filters, millions of
unique queries are used in each search embedding.

###  2.2. CoSMo: a Constrained Search Model for predicting real-world events.

When predicting real world events at a daily or weekly frequency most models
are prone to overfitting due to the curse of dimensionality, because while the
number of targets is limited by the time period and regions, it is easy to add
dimensionality to the features used for modeling (Theodoridis and Koutroumbas,
2006). In the case of search data, when the search features are represented by
the counts of unique terms, the feature set size grows too large unless it is
capped via a heuristic (e.g., top search terms by volume). Even though the
dimensionality of language models is small ($\sim$512 dimensions) compared to
the number of unique terms, an unregularized model whose goal is to predict
roughly three years of daily targets and that uses SLaM search embeddings as
features suffers from the curse of dimensionality, because the number of
features is roughly equal to the number of data points (Rosa, 2010). Although
there exist approaches like Lasso regression (Tibshirani, 1996) to combat
this, we offer a unique modeling approach that is less dependent on
regularization tricks in the loss function.

We start with a structural model that predicts the probability that the
average search contributes to the target $\hat{y_{t}}$ defined as

(3) $\hat{y}_{t}=V_{t}\cdot P(\gamma_{t}^{*},\theta)$

where $V_{t}$ is the total number of searches that happen during time-period
$t$, $P$ is a function that maps its learned parameters $\theta$ and the
search embedding $\gamma_{t}^{*}$ to a number between zero and one. Note two
characterisitics that make our model self-regularizing:

  1. (1)

$V_{t}$ changes day-to-day, and the model must learn to map the product of
this moving target and the probability to the target value; this is similar to
dropout (Hinton et al., 2012) or jittering (Bishop, 1995; Goodfellow et al.,
2016) where the model must be robust to moving targets.

  2. (2)

Although the model can have many parameters inside $\theta$ the variance of
the model is limited because $0<P(\gamma_{t}^{*},\theta)<1$; this comes at the
cost of a higher bias, but we show that this tradeoff leads to lower test
error.

Figure 2. Model Structure for the CoSMo model used in all models.

Our final CoSMo model is an iteration of (3). We model our target variable
$y_{t,r}$ for time period $t$ and geographical region $r\in R$ as

(4) $\hat{y}_{t,r}=\Psi(V_{t,r})\cdot
P(\gamma_{t,r}^{*},\theta,r)\cdot\prod_{k\in K}M_{k}(t,r)$

where $\Psi$ is a scaling function used to modify the volume222The volume
scaling function is typically set to the identity function. The other common
setup is to define it to map all inputs to $1$, effectively removing the
volume component from our model., $K$ is the set of multiplier
variables333e.g., $K=${g: g $\in$ Geographical Regions} $\cup$ {d: d $\in$
Days of Week}, $R$ is the set of all geographical regions, and

(5) $M_{k}(t,r)=\begin{cases}M_{k}&k=f(t)\text{ or }k=r\\\ 1&\text{else}\\\
\end{cases}$

where $M_{k}\in\mathbb{R}^{+}$ and $f$ represents some function mapping a
time-period (e.g., $f$ could be a function that maps each day to a day-of-
week, such that we have a separate multiplier that can be learned for each
day-of-week).

Note the flexibility in our model: 1. to do national-level modeling (i.e., fit
a model to national-level targets) we can set $R$ to be a single country,
while if we specify $R$ as the set of all U.S. States, our model operates at
the state-level 444If the regions are more granular than a country, we can
recover country-level predictions with a roll-up defined as  (6)
$\hat{y_{t}}=\sum_{r\in R}\hat{y_{t,r}}$ which works in cases like automobile
sales, where national-level statistics are a sum of all state-level
statistics. 555Sub-country geo-level models are popular modeling choice,
because models generally benefit from more granular data when the number of
targets is small (Sun et al., 2017)., and 2. while our model can include the
region $r$ as an input to allow interactions between the search embeddings
$\gamma_{t}^{*}$ and region $r$, we can also mask the regional features to
allow no interactions. Similarly, we can mask regional features to exclude
regional multipliers; see Figure 2. We provide more insights around regional
model variants in Appendix B.

Like in (Gorishniy et al., 2021), the model $P(\gamma_{t,r}^{*},\theta,r)$ is
a modified version of ResNet (He et al., 2016) for tabular data that uses
fully connected layers; in our blocks each layer is added instead of every
other layer. The final layer to our model has a singular unit with a sigmoid
activation function to keep the output bounded between zero and one.

The parameters $M_{k}$ and $\theta$ can be trained by minimizing the typical
(Goodfellow et al., 2016) Mean Square Error loss using the target variables.
We used a modified version of the Mean Aboslute Percent Error (MAPE) as the
loss, which we provide in Appendix F. Unless otherwise stated, we used a
regularized version of this loss to train all models.

##  3\. Related Work

Google Search was launched in 1998, and has become the world’s predominant
search engine. As such billions of people use it each month, typing queries
into the search bar to find information on the internet. As early as 2009
researchers were incorporating data around searches into predictive models
(Varian and Choi, 2009; Ginsberg et al., 2009; Schmidt and Vosen, [n. d.]),
and some of the notable modeling efforts including predicting U.S. flu rates
(Ginsberg et al., 2009) (Lampos et al., 2015) and economic indicators (Varian
and Choi, 2009). Much of economic forecasting relies on surveys as an
intermediate or current estimate of economic activity (Census, 2024; Reserve,
2024). (Schmidt and Vosen, [n. d.]) compares Google Trends to survey-based
metrics in the prediction of consumer confidence indicators, and finds that
Google Trends outperforms surveys.

###  3.1. Search representations

Individual search queries have been represented as embeddings in many
retrieval search tasks like in (Yang et al., 2020), (Zamani and Croft, 2016),
(Zhang et al., 2020), and (Zheng and Callan, 2015). But for predicting
external events using search data, as far as we are aware, all aggregations of
search data to date have been through binary 1/0 classification (Wikipedia,
2024) mappings, where a search query is either a member of the class (1) or
not (0).

####  3.1.1. Classification Embeddings

In (Varian and Choi, 2009), they represent search via the counts of the
individual queries that are grouped according to their search category, which
is determined by some external classification system. They represent each
search term $s$ as a one-hot vector $\phi^{s}$ where

(7) $\phi_{c}^{s}=\begin{cases}1&C(s)=c\\\ 0&\text{else}\\\ \end{cases}$

where $C$ is the classifier. They create the final representation of search
$\Phi$ by summing up these one-hot vectors

(8) $\Phi=\sum_{s}v_{s}\cdot\phi^{s}.$

We note that a classification scheme can be seen as another form of embedding,
where the embedding vector represents the counts for each mutually exclusive
category in each dimension – we call this the classification embedding. This
approach struggles with queries that don’t neatly fit into a single category,
and downstream models built on top of this method are sensitive to the
classification system. Classification embeddings lose much of the nuance in
search queries, where ”best new family SUV” might end up classified the same
as ”New model Lamborghini Urus” although the intent of the searches might be
quite different.

####  3.1.2. Filtered One-Hot Embeddings

(Ginsberg et al., 2009), (Polgreen et al., 2008) and (Lampos et al., 2015) use
individual queries as their representation of search. They use a filtered one-
hot encoding to represent each search term as

(9) $\delta_{s^{\prime}}^{s}=\begin{cases}1&s^{\prime}=s\\\ 0&\text{else}\\\
\end{cases}$

where $s^{\prime}\in A$, and $A$ is the set of accepted terms to include.
Their final search representation is the sum of these term representations

(10) $\Delta=\sum_{s}v_{s}\cdot\delta^{s}.$

They also normalize their vectors by the total search volume during each time
period, and (Ginsberg et al., 2009) also sums the vector components to produce
a univariate search index. We call this approach filtered one-hot embedding.
The filtering technique is almost a requirement to using individual search
terms, otherwise the input size would be too large for modeling due to the
curse of dimensionality (Theodoridis and Koutroumbas, 2006; Rosa, 2010). In
(Ginsberg et al., 2009), they note that model performance suffers for
$|A|>45$. In (Lampos et al., 2015) the number of initial terms was large
(hundreds of thousands), and they filter by only keeping terms whose Pearson
coefficient with the target variable is ¿0.5, which we note is an extra step
that required them to pick a filter threshold.. They also utilize Elastic-Net
(Zou and Hastie, 2005) for further term selection. In (Ginsberg et al., 2009)
the initial terms was even larger, around 5M, and they also filter terms by
using the correlations with the target variable. In (Polgreen et al., 2008) a
manual heuristic was invented to decide what terms to include. While these
approaches generated good enough representations to build flu models, search
terms that have many synonyms or common misspellings are completed ignored.

Our approach with respect to the search representations differs from the
existing literature in that it requires minimal filtering of queries a priori,
the queries are mapped to a continuous space vs a discrete space, and we rely
on a language model to handle misspellings, synonyms, and other types of
related terms.

###  3.2. Modeling with Search Representations

In terms of nowcasting and predictive modeling, to the best of our knowledge,
most of the downstream modeling on top of compressed search takes the form of
linear modeling.

(Varian and Choi, 2009) found that search data added information about
consumer behavior when added to a simple autoregressive. (Moller et al., 2023)
used a linear model with housing-related Google Trends to create a housing
index to predict house prices, and found that the accuracy of the model peaks
with about a 3-8 month lag. (Ginsberg et al., 2009) fitted a linear model on
an index of relevant queries.

The only non-linear modeling that we are aware of is from (Lampos et al.,
2015), where a Gaussian Process was used in conjunction with autoregressive
features.

All these related modeling efforts generally include other features, such as
lagged target variables (autoregressive models), and other economic variables,
and were not built solely on search data.

We believe there is room to improve the predictive capabilities beyond what
can be achieved with existing approaches, while preserving the ability to
interpret the model. For our search embeddings, it is likely that there is
interaction between the embeddings, which can be captured through non-linear
models like neural nets. We address the issue of overfitting, common in models
with many parameters, through regularization and inductive biases, and then
validate that our model generalizes by reporting our metrics over a test set
not included in the train and validation sets.

##  4\. Experiments

We evaluate our compressed search features and our modeling methodology in two
environments: U.S. flu prediction where we attempt to estimate the U.S. flu
case rates from the Center for Disease Control, and U.S. auto sales, where we
estimate the number of weekly vehicle sales; while our method was originally
designed for daily targets, we model both of these targets at the weekly-level
due to data availability. The experiments section is broken down as follows:

  1. (1)

Automotive Sales Predictions. We benchmark our method against existing methods
for the automobile sales prediction task.

  2. (2)

Flu Rate Predictions. We benchmark our method against the most recent methods
for the national ILI prediction task.

  3. (3)

Model Ablations. We run a number of ablations on our model using national and
regional flu models.

  4. (4)

LM Choice. We test the effect of using different language models.

  5. (5)

Zero-Shot. We test the model’s ability to do zero-shot inference using
different geo-level features from what was available during training.

. Finally, we look at the interpretability of the model in Section 5.

###  4.1. Experiment Configurations

All models are trained using three sets of dates
$T_{\text{test}},T_{\text{train}},\text{and }T_{\text{val}}$, which are a
partitioning of the full set of dates $T$. The test data is a subset of the
labels that have a timestamp after the train and validation data. Train and
validation sets are randomly chosen from the dates preceding the timestamp at
a predefined split. We set aside 10% of the days of non-test dataset as the
validation dataset.

Most of our experiments were run using only CPUs; in cases where GPUs were
helpful to speedup the computation, we used V100 GPUs. We implement all our
models using JAX (Bradbury et al., 2018), FLAX (Heek et al., 2023) and Optax
(DeepMind et al., 2020).

All neural networks were trained using Adam (Kingma and Ba, 2014) with a
linear warmup (Goyal et al., 2017), cosine decay (Loshchilov and Hutter,
2016), gradient norm clipping (Mikolov et al., 2012) (Pascanu et al., 2013),
additive noise (Neelakantan et al., 2015), and early stopping (Goodfellow et
al., 2016) using the validation loss. Unlike most modern ML methods which
utilize minibatches, our models were trained using the full gradient, because
the full batch fits in memory; we believe this is why we found additive noise
to be helpful with the optimization. For the flu model we run a hyperparameter
grid search using XManager (DeepMind, 2024). All hyperparameters can be found
in Appendix C. For each experiment, we run five trials with different random
seeds. We select the best model according to the average performance on the
validation dataset. For the auto model we have limited input data points, and
are focused on model interpretability, and search query insights. As such we
run a simple model with two hidden layers, and the 512 search embeddings as
inputs.

Unless otherwise stated, we used the Multilingual Sentence Encoder (MLSE)
(Yang et al., 2019) as the language model to embed search terms into 512
dimensions for the search embeddings.

###  4.2. Automotive Sales Predictions

Frequency | Embedding | Model | Test R2 $\uparrow$ | Test MAPE (%) $\downarrow$
---|---|---|---|---
Weekly | Categorical | Lasso | 0.5869 | 10.90
Weekly | Categorical | CoSMo | 0.5381 | 10.85
Weekly | SLaM | CoSMo | 0.7486 | 7.12
Monthly | SLaM | CoSMo | 0.9065 | 3.03
Table 1. Baseline Regional Auto Models with search and indicator multipliers -
fit metrics reported at the national level.

For automotive sales modeling, (Varian and Choi, 2009) was one of the first to
show the value of using Google Search data in predicting current economic
activity. Our approach further leverages the information present in search
queries to increase the accuracy of the nowcast prediction from accounting for
58% of variance when we use classification metrics to 75% using our search
embeddings, a 30% improvement in model accuracy. Much of the remaining
unexplained variance is due to monthly and quarterly cycles in the data. When
the data is rolled up to monthly blocks as reported in (Varian and Choi, 2009)
our model accounts for 91% of variation in the test set. Our model doesn’t use
historical sales or other external variables in our model, and the fit metrics
reported are $R^{2}$ and MAPE in order to be consistent with the literature.

Table 1 shows the results from modelling U.S. Auto Sales. We used overall US
Auto Sales and trained the model at the weekly level across 16 regions,
rolling our predictions up to national. The search data includes over ten
million distinct queries that are vehicle-related. The model uses both
regional and week-of-the-month features. The regional features are included in
the probability model to account for regional differences in both search
adoption and search behavior across regions. The model is trained across
nearly two years of data and the fit metric is reported over the test set, a
further 6 months of data. The model is trained with a two week lag between
search and sales, an interesting area for future research would be the impact
of varying lags, as (Moller et al., 2023) does for the housing market.

Figure 4 highlights the fit of the search embeddings CoSMo model using a four
week rolling average. The US auto sales data that we use in this paper is
based on registration data, and has large spikes at the end of the month as
well as end of quarter. The large improvement in fit by using four week
rolling average suggests that this monthly cycle is likely a supply-side
effect as opposed to reflective of demand patterns.

At the monthly level the model has an R2 of 0.91, and 3.03 MAPE in the test
period. This fit is remarkable given that the model doesn’t include any annual
seasonality controls, or historical sales. As a point of reference the linear
model in (Varian and Choi, 2009) returns a monthly R2 of 0.79 over the
training data using both lagged sales and Google Trends.

While automotive sales are used in this paper, we expect that our approach can
be used to greatly improve nowcasts across economic indicators. In the next
section we show how the model can accurately predict flu rates, and show the
sensitivity of the model to model specifications.

###  4.3. Flu Rate Predictions

For benchmarking experiments, we model Influenza-Like-Illness (ILI) rates from
the CDC (CDC, 2024) at the national level, like (Lampos et al., 2015). Due to
data availability, we are unable to compare our model on the same time frames
as in previous work. Instead, we use data from 2019 until 2022 for training
and validation data, and we estimate the flu rates for the 2022-2023 flu
season as the test period. In (Lampos et al., 2015) the Pearson correlation
coefficient and the Mean Absolute Percentage Error are provided for multiple
flu seasons from 2008 until 2013; for the methods we implemented, we report
the average values across 5 trials. We provide the best and worst performances
of previous methods in (Lampos et al., 2015) to benchmark our approach. In
previous works, it is unclear how the model’s hyperparameters were selected.
We report the test metrics of our approach using the model whose average
validation MAPE was lowest; for benchmarking purposes, we also report the
model with the best test MAPE.

Additionally, we compare our modeling approach to more typical methods such as
logistic regression and multi-layer perceptron (MLP) neural networks, which
have a history of modeling success but do not have the regularizing structural
components of our approach. For logistic regression, we found the model to
work better without search volume, and only use the normalized search
embeddings. All methods include L1 regularization. We include about two
million cold & flu related terms for our search embeddings.

Figure 3 shows our model’s predicted values for a few years during both
training and testing. Our model, which only uses data from search to estimate
of the flu rates of a given week, is able to closely estimate the actual flu
rates for a new flu season despite not using lagged flu rate data in its
estimates like autoregressive models. Table 2 shows the results from modeling
the U.S. ILI rates at the national level. We can see that CoSMo outperforms
other methods which only use search data. The autoregressive (AR) entries in
Table 2 represent methods that include either a 1-week or 2-week lag of the
most recent ILI rate. Our method is generally on par or better than the best
AR approaches.

| Test MAPE(%) $\downarrow$ | Test $r$ $\uparrow$
---|---|---
Logistic Regression | 24.9 $\pm$ 0.1 | .98
MLP | 7.3 $\pm$ 1.5 | .99
Google Flu Trends (Lampos et al., 2015) | [9.5 - 33.1] | [.66 - .97]
Elastic Net (Lampos et al., 2015) | [9.8 - 15.1] | [.92 - .99]
Guassian Process (Lampos et al., 2015) | [9.4 - 14.6] | [.94 - .99]
AR (Lampos et al., 2015) | [6.7 - 14.3] | [.88 - .98]
AR+Google Flu Trends (Lampos et al., 2015) | [6.2 - 12.5] | [.88 - .99]
AR+Elastic Net (Lampos et al., 2015) | [5.1 - 8.7] | [.93 - $\approx 1$]
AR+Guassian Process (Lampos et al., 2015) | [5.0 - 8.6] | [.93 - $\approx 1$]
CoSMo (Ours) | 5.5 $\pm$ 0.4 | .99
CoSMo (Ours, Test selection) | 3.9 $\pm$ 0.1 |  $\approx$1
Table 2. Benchmarking ILI flu rate prediction at the national level. We show
the standard deviation of MAPE for our experiments; we omit this metric for
the Pearson coefficient because it was close to zero for all experiments.
Figure 3. National U.S. Flu Modeling plot for Training and Test periods. CoSMo
predicted values are the average of 40 trainings with random seeds with the
shaded areas represesnting the 95% confidence interval.

###  4.4. Model Ablations

In Tables 3 and 4 we show the test results from training multiple variants of
flu models at the national and regional levels respectively. We run ablations
on three components of the model: the search volume feature, regional
multipliers, and conditioning the probability model on the region. We show the
effect of including vs excluding the search volume as a feature for both state
and national models.

Interestingly, for the national model, excluding the volume has a large
negative impact (5.46% $\rightarrow$ 12.37% MAPE), while for the regional
models excluding the volume helped for those models without region features in
the probability model (44.94% $\rightarrow$ 31.96% MAPE and 38.05%
$\rightarrow$ 27.27% MAPE), but for the other models there was little effect.
The best performing regional models were those with the region as an input
into the probability model . We hypothesize that for the regional modeling
task, there are important interactions between what users are searching and
where they are located, which is why including the region features is so
beneficial. For the multipliers, we see that including the regional
multipliers helps model performance when the probability model is not
conditioned on the geo (31.96% $\rightarrow$ 27.27% and 44.94% $\rightarrow$
$38.05\%$), and when the geo is present in the probability model, there is
little effect.

Volume | Test MAPE(%) $\downarrow$ | Test $r$ $\uparrow$
---|---|---
$\checkmark$ | $\ 5.46\pm 0.43$ | $.9933\pm.0005$
| $12.37\pm 1.75$ | $.9904\pm.0023$
Table 3. National model with and without the volume feature. Multiplier | $P(|\text{geo})$ | Volume | Test MAPE (%)$\downarrow$ | Test $r\uparrow$
---|---|---|---|---
|  | $\checkmark$ | $44.94\pm 1.04$ | $.7024\pm.0186$
|  |  | $31.96\pm 0.63$ | $.8278\pm.0219$
$\checkmark$ |  | $\checkmark$ | $38.05\pm 0.89$ | $.7589\pm.0164$
$\checkmark$ |  |  | $27.27\pm 1.73$ | $.8880\pm.0110$
| $\checkmark$ | $\checkmark$ | $24.54\pm 0.59$ | $.8966\pm.0112$
| $\checkmark$ |  | $25.45\pm 1.05$ | $.8969\pm.0117$
$\checkmark$ | $\checkmark$ | $\checkmark$ | $24.88\pm 0.63$ | $.8960\pm.0082$
$\checkmark$ | $\checkmark$ |  | $24.41\pm 0.35$ | $.9082\pm.0042$
Table 4. Model ablations for regional flu models. The Multiplier column
indicates whether State multipliers were used, while $P(|$geo$)$ indicates
whether the probability model was conditioned on the State.

###  4.5. Zero-shot inference

We analyze the capability of our model to go from child-geography to parent-
geography predictions and vice versa. Training a model on parent-level (e.g,.
country) data, then evaluating on child-level (e.g., State) is common when
child-level data is either missing or never collect, while training a model at
the child-level and making parent-level predictions is useful when it’s
believed that the increased number of child-geo datapoints will help the model
fit. We use two versions of the best flu models: a no-volume national-level
model and a no-volume state-level model. The national-level model was trained
on national-level targets using national-level search embeddings, but
inference was done using state-level search embeddings and evaluated on state-
level targets; vice versa for the state-level model. The results are shown in
Table 5. The model has a surprising capability to infer with some success (.78
$r$) state-level flu rates, in the test period, without ever being trained on
state-level targets. The zero-shot inference performs better in the opposite
direction, (.99 $r$), perhaps leveraging the greater number of training
examples and taking advantage of the easier task of national modeling.

Training Data | Eval Data | Test MAPE(%)$\downarrow$ | Test $r\uparrow$
---|---|---|---
State | State | $31.96\pm 0.63$ | $.8278\pm.0219$
National | State | $55.08\pm 3.23$ | $.7820\pm.0072$
National | National | $12.37\pm 1.75$ | $.9904\pm.0023$
State | National | $11.22\pm 0.35$ | $.9856\pm.0030$
Table 5. Zero-shot evaluation for Flu ILI rate prediction. The zero-shot
examples are the rows where there is a mismatch between the Training Data
column and the Eval Data column. The rows with alignment serve as comparison
points.  Figure 4. U.S. Automotive Sales Actuals vs. Predictions. A 4-week
rolling average of the model and targets were generated to smooth out spikes
typically caused by end-of-month reporting variability. On the test period the
model has a .9065 R2 and 3.03 MAPE. The vertical line indicates the beginning
of the test period.

###  4.6. LM Choice

In addition to the MLSE embeddings (Yang et al., 2019), we look at variants of
the T5 (Raffel et al., 2020) LLM, the sentence-T5 (sT5) (Ni et al., 2021), a
version of T5 that outputs a fixed-length 768-dimensional vector for every
input sequence 666Our method requires that the LM output a D-dimensional
vector that is not dependent on the input shape. Unfortunately, many LMs have
outputs with shape $L\times D$ where $L$ is the number of input tokens. In
order to study many other LMs using our method, such as mT5, we would need to
first map the LM output to a fixed-length vector. Potential options are using
the output associated with the ¡BOS¿ token, or averaging across the sequence
length dimension. We leave these experiments to future work.. We study the
effect of using these embeddings on the the national Flu ILI prediction tasks.
Table 6 shows the results from using different search embeddings created using
the sT5 Base (110M parameters) and sT5 Large (335M parameters) models.

| Test MAPE(%)$\downarrow$ | Test $r\uparrow$
---|---|---
MLSE (baseline) | $5.46\pm 0.43$ | $.9933\pm.0005$
sT5 Base | $6.51\pm 0.13$ | $.9906\pm.0016$
sT5 Large | $6.51\pm 0.97$ | $.9894\pm.0049$
MLSE (English only) | $9.11\pm 0.99$ | $.9902\pm.0014$
sT5 Base (English only) | $7.69\pm 1.29$ | $.9846\pm.0013$
sT5 Large (English only) | $7.22\pm 0.69$ | $.9878\pm.0040$
Table 6. National ILI rate modeling results from using different embedding
functions from a variety of LMs.

Surprisingly, larger capacity models like sT5 Base and sT5 Large do not
outperform the smaller capacity MLSE model. We believe this has to do with sT5
models being trained on only the English language. The MLSE model being a
multi-lingual model is able to make better use of the multiple languages
present in the search data, where as the sT5 models are unable accurately map
the meanings of these queries. We validate this by generating search
embeddings using only English queries and training models on these English-
only search embeddings. These results are shown in Table 6. We can see that
the sT5 models perform similar to their all-language counter parts, where as
performance for MLSE considerable lowers. We leave further studies to future
work.

##  5\. Model Interpretability

After the model has been fit, we interpret the model by running inference on
the queries and organizing them by their probability score. Table 7 shows some
examples of terms that score high, meaning the model believes queries near
those areas of the embedding space are highly predictive of flu cases, and
examples of terms that score low (zero, or close to zero), which the model
learned to ignore.

Term | Score Percentile
---|---
efluenza | 0.04
flu center | 0.09
uniflu tablets | 0.17
summer flu ohio | 0.26
flu current status | 0.86
flu vactination | 0.95
flu cdc recommendations | 1.30
flu & pneumonia | 1.74
benzonate | 38.51
do 5 year olds get immunizations | 60.76
kansas vaccine locations | 86.79
does putting vicks on your chest help | 86.78
nasal spray toddlers | 95.48
runny nose puffy eyes | 99.90
post nasal drip spitting | 99.90
Table 7. High- and Low-Scoring Flu Terms.

One interesting note is the term ”benzonatate” which is a term known to have
major detrimental behavior when not excluded in previous flu prediction models
(Lampos et al., 2015), but our method learns to ignore the term without any
human intervention or special preprocessing of the data. Our model also seems
to handle misspellings well (see ”efluenza” a misspelling of ”influenza,” and
”flu vactination” a misspelling of ”flu vaccination”) which we attribute to
the LM.

Figure 5. Visualization of auto search terms and estimated impact from model
predictions. Each dot represents a distinct search term, and terms have been
clustered based on embedding vectors, and hand-labeled for exposition.

We visualize the embeddings of auto terms and the corresponding model scores
in Figure 5. The figure was created by clustering the embeddings of all auto-
related search terms through a RAC algorithm (Sumengen et al., 2021),
filtering a subset of the clusters for plotting, projecting the original
embedding to a 2D space using t-SNE, and plotting each point in the clusters.
We use the model scores of each term as proxy for the importance of the term,
and we calculate cluster-level importances by averaging all terms in the
cluster; we represent the relative importance of each cluster through the
darkness of the color (see legend).

We manually labeled clusters according to the their unifying themes. For
example the Craigslist for sale cluster has more than 1000 distinct search
terms such as ”craigslist used cars for sale” and ”for sale cars craigslist”
or some other variation. We see that in general the terms are increasing in
impact as they move upwards and to the right in the projected embedding space.
However the model is able to pick up additional nuances, such as the term
clusters 4runner and awd vs. 4wd which are close in the embedding space, but
differ greatly in predicted impact on auto sales.

In Appendix E we show a similar visualization for the flu model, along with
the impact of seasonal variation on the search embedding.

##  6\. Ethical Use of Data

In the case of modeling the flu, while we report great correlations and error
rates on an unseen flu season with our method, it is not a substitute for
traditional disease reporting; previous methods have shown to be effective in
research environments but fail to be as useful or accurate in practice (Lazer
et al., 2014; Olson et al., 2013; Cook et al., 2011), so we recommend caution
and further testing when using our approach.

The search data in this project was anonymized, with none of the queries
associated with any individuals or accounts. Furthermore, because we aggregate
individual search terms in an embedding space, we believe our method only
increases the privacy of the search data.

##  7\. Conclusion

It has been established in the literature that search data can add efficacy to
predictive models. With billions of Google searches each day, it contains
valuable signals on everything from flu prevalence to auto brand sentiment.
Until now the typical implementation of search data into predictive models has
been through incorporating coarse Google Trends data, similarly aggregated
data using binary classifiers, or through complex filters for including
individual search terms, which leaves the downstream models with either a
diluted signal or signals prone to overfitting.

Using SLaM we propose a method to include search data in a privacy-safe manner
by using the embeddings from language models to create a summary of the search
data. This allows us to retain much of the information about the queries
themselves as well as their relative volumes while greatly reducing the
dimensionality. Although all of our experiments were run using Google Search
data, our search compression can be applied to any scenario where statistics
associated with natural language need to be summarized to a fixed-length
vector.

We also introduce CoSMo, a constrained search model, which has inductive
biases that greatly improve the accuracy of our models built on search data.
For estimating the flu rates, we show our simple approach is on par or better
than the existing complex ensemble methods. For estimating auto sales, we show
large improvements over existing methods that only use categorized versions of
search data. Finally, we demonstrate that our models, despite being highly
non-linear neural networks, offer interpretability that explains what terms
are related to the variables of interest.

###### Acknowledgements.

Thanks you Yi Chao for helping with the embedding visualizations, Saket Kumar
for reviewing and supporting the project, and Karima Zmerli for inspiring this
project and supporting this work.

## References

  * (1)
  * Bishop (1995) Chris M. Bishop. 1995\.  Training with Noise is Equivalent to Tikhonov Regularization.  _Neural Computation_ 7, 1 (01 1995), 108–116.  https://doi.org/10.1162/neco.1995.7.1.108 arXiv:https://direct.mit.edu/neco/article-pdf/7/1/108/812990/neco.1995.7.1.108.pdf
  * Bradbury et al. (2018) James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Milne, and Qiao Zhang. 2018\.  _JAX: composable transformations of Python+NumPy programs_.  http://github.com/google/jax
  * CDC (2024) CDC. 2024.  https://gis.cdc.gov/grasp/fluview/fluportaldashboard.html
  * Census (2024) US Census. 2024\.  Advance Monthly Sales for Retail and Food Services.  https://www.census.gov/retail/sales.html Last accessed 8 February, 2024.
  * Cook et al. (2011) Samantha Cook, Corrie Conrad, Ashley L Fowlkes, and Matthew H Mohebbi. 2011.  Assessing Google flu trends performance in the United States during the 2009 influenza virus A (H1N1) pandemic.  _PloS one_ 6, 8 (2011), e23610.
  * DeepMind et al. (2020) DeepMind, Igor Babuschkin, Kate Baumli, Alison Bell, Surya Bhupatiraju, Jake Bruce, Peter Buchlovsky, David Budden, Trevor Cai, Aidan Clark, Ivo Danihelka, Antoine Dedieu, Claudio Fantacci, Jonathan Godwin, Chris Jones, Ross Hemsley, Tom Hennigan, Matteo Hessel, Shaobo Hou, Steven Kapturowski, Thomas Keck, Iurii Kemaev, Michael King, Markus Kunesch, Lena Martens, Hamza Merzic, Vladimir Mikulik, Tamara Norman, George Papamakarios, John Quan, Roman Ring, Francisco Ruiz, Alvaro Sanchez, Laurent Sartran, Rosalia Schneider, Eren Sezener, Stephen Spencer, Srivatsan Srinivasan, Miloš Stanojević, Wojciech Stokowiec, Luyu Wang, Guangyao Zhou, and Fabio Viola. 2020.  _The DeepMind JAX Ecosystem_.  http://github.com/google-deepmind
  * DeepMind (2024) Google DeepMind. 2024\.  XManager: A platform for managing machine learning experiments.  https://github.com/google-deepmind/xmanager.  Accessed: 2024-02-01.
  * Ginsberg et al. (2009) J. Ginsberg, M. Mohebbi, and R. et al. Patel. 2009.  Detecting influenza epidemics using search engine query data.  _Nature_ 457 (2009), 1012–1014.
  * Goodfellow et al. (2016) Ian J. Goodfellow, Yoshua Bengio, and Aaron Courville. 2016\.  _Deep Learning_.  MIT Press, Cambridge, MA, USA.  http://www.deeplearningbook.org.
  * Gorishniy et al. (2021) Yury Gorishniy, Ivan Rubachev, Valentin Khrulkov, and Artem Babenko. 2021.  Revisiting deep learning models for tabular data.  _Advances in Neural Information Processing Systems_ 34 (2021), 18932–18943.
  * Goyal et al. (2017) Priya Goyal, Piotr Dollár, Ross Girshick, Pieter Noordhuis, Lukasz Wesolowski, Aapo Kyrola, Andrew Tulloch, Yangqing Jia, and Kaiming He. 2017.  Accurate, large minibatch sgd: Training imagenet in 1 hour.  _arXiv preprint arXiv:1706.02677_ (2017).
  * He et al. (2016) Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. 2016\.  Deep residual learning for image recognition. In _Proceedings of the IEEE conference on computer vision and pattern recognition_. 770–778.
  * Heek et al. (2023) Jonathan Heek, Anselm Levskaya, Avital Oliver, Marvin Ritter, Bertrand Rondepierre, Andreas Steiner, and Marc van Zee. 2023.  _Flax: A neural network library and ecosystem for JAX_.  http://github.com/google/flax
  * Hinton et al. (2012) Geoffrey E Hinton, Nitish Srivastava, Alex Krizhevsky, Ilya Sutskever, and Ruslan R Salakhutdinov. 2012\.  Improving neural networks by preventing co-adaptation of feature detectors.  _arXiv preprint arXiv:1207.0580_ (2012).
  * Kingma and Ba (2014) Diederik P Kingma and Jimmy Ba. 2014.  Adam: A method for stochastic optimization.  _arXiv preprint arXiv:1412.6980_ (2014).
  * Kiros et al. (2015) Ryan Kiros, Yukun Zhu, Russ R Salakhutdinov, Richard Zemel, Raquel Urtasun, Antonio Torralba, and Sanja Fidler. 2015.  Skip-thought vectors.  _Advances in neural information processing systems_ 28 (2015).
  * Lampos et al. (2015) Vasileios Lampos, Andrew C Miller, Steve Crossan, and Christian Stefansen. 2015.  Advances in nowcasting influenza-like illness rates using search query logs.  _Scientific reports_ 5, 1 (2015), 12760.
  * Lazer et al. (2014) David Lazer, Ryan Kennedy, Gary King, and Alessandro Vespignani. 2014.  The parable of Google Flu: traps in big data analysis.  _science_ 343, 6176 (2014), 1203–1205.
  * Loshchilov and Hutter (2016) Ilya Loshchilov and Frank Hutter. 2016.  Sgdr: Stochastic gradient descent with warm restarts.  _arXiv preprint arXiv:1608.03983_ (2016).
  * Mikolov et al. (2012) Tomáš Mikolov et al. 2012\.  Statistical language models based on neural networks.  _Presentation at Google, Mountain View, 2nd April_ 80, 26 (2012).
  * Moller et al. (2023) Stig Vinther Moller, Thomas Pedersen, Erik Christian Montes Schutte, and Allan Timmermann. 2023\.  Search and Predictability of Prices in the Housing Market.  _Management Science_ 70, 1 (2023), 415–438.
  * Neelakantan et al. (2015) Arvind Neelakantan, Luke Vilnis, Quoc V Le, Ilya Sutskever, Lukasz Kaiser, Karol Kurach, and James Martens. 2015.  Adding gradient noise improves learning for very deep networks.  _arXiv preprint arXiv:1511.06807_ (2015).
  * Ni et al. (2021) Jianmo Ni, Gustavo Hernández Ábrego, Noah Constant, Ji Ma, Keith B Hall, Daniel Cer, and Yinfei Yang. 2021\.  Sentence-t5: Scalable sentence encoders from pre-trained text-to-text models.  _arXiv preprint arXiv:2108.08877_ (2021).
  * Olson et al. (2013) Donald R Olson, Kevin J Konty, Marc Paladini, Cecile Viboud, and Lone Simonsen. 2013.  Reassessing Google Flu Trends data for detection of seasonal and pandemic influenza: a comparative epidemiological study at three geographic scales.  _PLoS computational biology_ 9, 10 (2013), e1003256.
  * Pascanu et al. (2013) Razvan Pascanu, Tomas Mikolov, and Yoshua Bengio. 2013\.  On the difficulty of training recurrent neural networks. In _International conference on machine learning_. Pmlr, 1310–1318.
  * Polgreen et al. (2008) Philip M Polgreen, Yiling Chen, David M Pennock, Forrest D Nelson, and Robert A Weinstein. 2008.  Using internet searches for influenza surveillance.  _Clinical infectious diseases_ 47, 11 (2008), 1443–1448.
  * Raffel et al. (2020) Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu. 2020.  Exploring the limits of transfer learning with a unified text-to-text transformer.  _The Journal of Machine Learning Research_ 21, 1 (2020), 5485–5551.
  * Reserve (2024) Federal Reserve. 2024\.  Household Spending Survey.  https://www.newyorkfed.org/microeconomics/sce/household-spending Last accessed 8 February, 2024.
  * Rosa (2010) Guilherme JM Rosa. 2010\.  The Elements of Statistical Learning: Data Mining, Inference, and Prediction by HASTIE, T., TIBSHIRANI, R., and FRIEDMAN, J.
  * Schmidt and Vosen ([n. d.]) Torsten Schmidt and Simeon Vosen. [n. d.].  Forecasting Private Consumption: Survey-Based Indicators vs. Google Trends.  _Ruhr Economic Paper_ 155 ([n. d.]).
  * Sumengen et al. (2021) Baris Sumengen, Anand Rajagopalan, Gui Citovsky, David Simcha, Olivier Bachem, Pradipta Mitra, Sam Blasiak, Mason Liang, and Sanjiv Kumar. 2021.  Scaling hierarchical agglomerative clustering to billion-sized datasets.  _arXiv preprint arXiv:2105.11653_ (2021).
  * Sun et al. (2017) Yunting Sun, Yueqing Wang, Yuxue Jin, David Chan, and Jim Koehler. 2017.  Geo-level bayesian hierarchical media mix modeling.  _URL: https://ai. google/research/pubs/pub46000_ (2017).
  * Theodoridis and Koutroumbas (2006) Sergios Theodoridis and Konstantinos Koutroumbas. 2006.  _Pattern recognition_.  Elsevier.
  * Tibshirani (1996) Robert Tibshirani. 1996\.  Regression Shrinkage and Selection via the Lasso.  _Journal of the ROyal Statistical Society_ 58, 1 (1996), 267–288.
  * Varian and Choi (2009) Hal R. Varian and Hyunyoung Choi. 2009.  Predicting the Present with Google Trends.  _ssrn.com_ (2009).  http://dx.doi.org/10.2139/ssrn.1659302
  * Wang et al. (2023) Dawei Wang, Andrea Guerra, Frederick Wittke, John Cameron Lang, Kevin Bakker, Andrew W Lee, Lyn Finelli, and Yao-Hsuan Chen. 2023\.  Real-Time Monitoring of Infectious Disease Outbreaks with a Combination of Google Trends Search Results and the Moving Epidemic Method: A Respiratory Syncytial Virus Case Study.  _Tropical Medicine and Infectious Disease_ 8, 2 (2023), 75\.
  * Wikipedia (2024) Wikipedia. 2024\.  Web Query Classification.  https://en.wikipedia.org/wiki/Web_query_classification Last accessed 8 February, 2024.
  * Woloszko (2020) N. Woloszko. 2020\.  Tracking activity in real time with Google Trends.  _OECD Economics Department Workign Papers_ 1634 (2020).  https://dx.doi.org/10.1787/6b9c7518-en
  * Yang et al. (2020) Ji Yang, Xinyang Yi, Derek Zhiyuan Cheng, Lichan Hong, Yang Li, Simon Xiaoming Wang, Taibai Xu, and Ed H Chi. 2020\.  Mixed negative sampling for learning two-tower neural networks in recommendations. In _Companion proceedings of the web conference 2020_. 441–447.
  * Yang et al. (2019) Yinfei Yang, Daniel Cer, Amin Ahmad, Mandy Guo, Jax Law, Noah Constant, Gustavo Hernandez Abrego, Steve Yuan, Chris Tar, Yun-Hsuan Sung, et al. 2019\.  Multilingual universal sentence encoder for semantic retrieval.  _arXiv preprint arXiv:1907.04307_ (2019).
  * Zamani and Croft (2016) Hamed Zamani and W Bruce Croft. 2016.  Embedding-based query language models. In _Proceedings of the 2016 ACM international conference on the theory of information retrieval_. 147–156.
  * Zhang et al. (2020) Han Zhang, Songlin Wang, Kang Zhang, Zhiling Tang, Yunjiang Jiang, Yun Xiao, Weipeng Yan, and Wen-Yun Yang. 2020\.  Towards personalized and semantic retrieval: An end-to-end solution for e-commerce search via embedding learning. In _Proceedings of the 43rd International ACM SIGIR Conference on Research and Development in Information Retrieval_. 2407–2416.
  * Zheng and Callan (2015) Guoqing Zheng and Jamie Callan. 2015.  Learning to reweight terms with distributed representations. In _Proceedings of the 38th international ACM SIGIR conference on research and development in information retrieval_. 575–584.
  * Zou and Hastie (2005) Hui Zou and Trevor Hastie. 2005.  Regularization and variable selection via the elastic net.  _Journal of the Royal Statistical Society Series B: Statistical Methodology_ 67, 2 (2005), 301–320.

##  Appendix A Linear Model Inspiration

Let’s choose to model the sales on day $t$, $y_{t}$, as

$\hat{y_{t}}=\sum_{s}v_{s,t}\cdot P_{s}$

where $S$ is the set of unique search terms, $v_{s,t}$ is the number of
appearances of term $s$ on day $t$, and $P_{s}$ is the conversion rate of
search term $s$ towards a sale.

Let’s also model $P_{s}$ as

$P_{s}=f(\mathbf{e_{s}})$

where $\mathbf{e_{s}}\in\mathbb{R}^{D}$ is the L2-normed LM $D$-dimensional
embedding of $s$ and $f$ is a constrained linear model of the form

$f(\mathbf{x})=\frac{1}{2}\Big{[}\beta+\sum_{i}\alpha_{i}\cdot x_{i}\Big{]}$
$=\frac{1}{2}\Big{[}\beta+\alpha\cdot\mathbf{x}\Big{]}$

where $\beta,\alpha_{i}\in\mathbb{R}$,

$||\mathbf{\alpha}||_{2}=1,$

and

$\beta=1.$

Claim: $0\leq P_{s}\leq 1$

###### Proof.

By cosine similarity

$\begin{split}\cos(\mathbf{e_{s}},\alpha)&=\frac{\mathbf{e_{s}}\cdot\alpha}{||\mathbf{e_{s}}||_{2}\cdot||\alpha||_{2}}\\\
&=\mathbf{e_{s}}\cdot\alpha,\end{split}$

and

$-1\leq\cos(\mathbf{e_{s}},\alpha)\leq 1.$

Then via substitution

$-1\leq\mathbf{e_{s}}\cdot\mathbf{\alpha}\leq 1$

and

$0\leq\mathbf{e_{s}}\cdot\mathbf{\alpha}+\beta\leq 2.$

Hence,

$0\leq\frac{\mathbf{e_{s}}\cdot\mathbf{\alpha}+\beta}{2}\leq 1$

and

$0\leq f(\mathbf{e_{s}})\leq 1$

∎

Then

$\hat{y_{t}}=\sum_{s}v_{s,t}\cdot f(\mathbf{e_{s}})$

is a constrained linear model that estimates the sales and has the property
where $f(\mathbf{e_{s}})$ represents an estimate of the conversion rate for
term $s$.

At this point, the main issue with the model is in the practical
implementation. Typically, $|S|$ is on the order of at least many millions of
terms. We would like to optimize for the coefficients $\alpha$ by minimizing

$L=\sum_{t\in T}l(\hat{y_{t}},y_{t})$

where $l$ is a differentiable loss function and $T$ is the set of days in our
training data. We would like to use one of the many gradient-based
optimization methods to iteratively update $\alpha$ using $\nabla_{\alpha}L$.
Note that to compute $\nabla_{\alpha}L$ we need, for all days, to compute
$f(\mathbf{e_{s}})$ for all search terms $s$, which requires approximately
$|T|\cdot|S|$ steps if implemented naively.

Let’s examine the compute for training such a model. Assume that each
embedding has 512 dimensions, and that the values are stored as 32-bit floats.
Let’s also assume $|S|\approx 10^{7}$ (roughly 10M unique search terms) and
$|T|\approx 10^{3}$ (roughly 3 years of daily training data, which is common
in practice). Then a naive implementation of minimize our loss would take
$10^{10}$ operations to make a single update to our model. In practice, it’s
common for there to be at least $10^{4}$ parameter updates if training a model
from scratch to convergence. If each operation included a single IO read of
data, then there would be $10^{4}\cdot 10^{7}\cdot 10^{3}\cdot[512\cdot
32\text{ bits}]\approx 200\text{ petabytes}$ of data read. Considering that
model training in many application needs to be done frequently throughout a
calendar year, and we have not yet considered the cost of a hyperparameter
search, this naive approach is infeasible in practice.

If, instead of reading the data from disk naively during each training step
you instead cached the result, you would have to store $10^{7}\cdot
10^{3}\cdot[512\cdot 32\text{ bits}]\approx 20\text{ terabytes}$, which is
able to fit into memory on today’s large distributed systems, but this can be
costly, will not fit into non-distributed system, and forces whoever is
implementing the model to use distributed training techniques.

Rewriting our model using vector notation,

$\hat{y_{t}}=\sum_{s}v_{s,t}\cdot f(\mathbf{e_{s}})$
$=\mathbf{v_{t}}^{T}\cdot\Big{[}\frac{1}{2}\cdot[E\cdot\alpha+\beta]\Big{]}$

where $\mathbf{v_{t}}\in\mathbb{R}^{|S|}$ is the vector of query counts for
day $t$ where each element represent the search count for a term, and
$E\in\mathbb{R}^{|S|\times D}$ is the embedding table where each row contains
the $D$-dimensional embedding of a search term, $\alpha$, previously a
$|S|$-length vector, is now the $|S|\times 1$ matrix, and $\beta$, previously
the scalar equal to one, is now the $|S|$-length vector of ones.

Note, we can rearrange the terms such that

$\hat{y_{t}}=[\mathbf{v_{t}}^{T}\cdot E]\cdot\frac{1}{2}\cdot[\alpha+\beta]$
$=\gamma_{t}\cdot\frac{1}{2}\cdot[\alpha+\beta]$

where we call

$\gamma_{t}=\mathbf{v_{t}}^{T}\cdot E$

the ”un-normalized search embedding” for day $t$ (the ”search embedding”
usually refers to the L2-normed version of this quantity). The key takeaway
here is that $\gamma$ can be pre-computed a priori model training, and it
doesn’t need to be recomputed during training. There are still
$|S|\times|T|\times D$ operations required to compute it, but it happens only
once, and doesn’t happen during model training. Thus, each iteration during
the optimization step only operates on a D-dimensional quantity which make
model selection possible.

Back to our 512-dimensional example trained with $10^{4}$ parameter updates,
there would only be $10^{4}\cdot 10^{3}\cdot[512\cdot 32\text{
bits}]\approx\text{ 20 gigabytes}$ of data read if we didn’t cache the search
embeddings; in practice, we can cache the search embeddings (only
$10^{3}\cdot[512\cdot 32\text{ bits}]\approx\text{ 2 megabytes}$ of data) into
memory and have zero IO reads during training.

In practice, we don’t limit ourselves to this linear model, and instead use a
model of the form

$\hat{y_{t}}=h(\gamma_{t})$

where $h$ is a non-linear neural network. In fact, our models typically take
the form

$\hat{y_{t}}=V_{t}\cdot\sigma\Big{(}h\Big{(}\frac{\gamma_{t}}{||\gamma_{t}||_{2}}\Big{)}\Big{)}\cdot\prod_{k}M_{k}(t)$

where $V_{t}=\sum_{i}\mathbf{v_{t}}_{i}$ is the total search volume for day
$t$, $\sigma$ is the logistic squashing function, and $M_{k}(t)$ is the value
for multiplier $k$ on day $t$.

##  Appendix B Regional Model Formulation

In our methodology, the geographical region can be incorporated into the model
two ways:

  * •

in the probability function $P(\gamma_{t,r},\theta,r)$

  * •

as a multiplier, $M_{r}$.

The general case for CoSMo in (4) can be viewed as a combination of two
specific regional models that differ in how they incorporate the region
information:

(11) $\hat{y_{t,r}}=V_{t,r}\cdot P(\gamma_{t,r},\theta,r)$

and

(12) $\hat{y_{t,r}}=V_{t,r}\cdot P(\gamma_{t,r},\theta)\cdot M_{r}$

where $M_{r}\in\mathbb{R}$ is called the regional multiplier for region $r$,
and $V_{t,r}$ and $\gamma_{t,r}$ are the total volume and search embedding on
day $t$ in region $r$, respectively.

In (11) the model has the flexibility to allow the region to change how the
search embedding is mapped. There may be regional differences in search
behavior across the geographic populations or differences in demographics
across regions. By having the probability model conditioned on this region
information, the model can account for these differences, even if the searches
are similar.

On the other hand, in (12), the model keeps the same probability mapping for
all regions, but learns region-specific multipliers that can nudge the
prediction up or down. With regional multipliers, we can capture regional
variations in search such as search adoption and search frequency. In this
case even when users have the same intent to purchase given a specific term,
if different proportions of the population use search, or the search intensity
differs across regions, the multiplier can pick up these differences.

In practice, it’s worth noting that that running model inference using (11)
requires $|R|$-times more calls to $P$ (where $|R|$ is the number of regions),
which is a relatively computationally expensive neural network. Using (12),
you can avoid these $|R|$-times more calls to $P$ and recover $\hat{y}$ with
just $|R|$-times more multiplications, which a small compute cost; this can be
very useful in practice when running inference over hundreds of millions of
search terms and the cost of inference is dominated by the calls to $P$.

##  Appendix C Hyperparameters

In Table 8 we show the values of hyper-parameters we search over for the flu
models. For each method we perform a grid search over relevant hyperparameter
values and choose the best ones according to the validation set.

Parameter Name | Search Space
---|---
Learning rate |  [$10^{-4}$]
L1 Regularization | [0, 20, 200, 500, 1000, 2000, 5000]
Layer Size | [64, 128]
Number of Layers | [1, 5, 10, 20, 40, 80, 160]
Gradient Noise | [0, .001, .0001]
Training Steps | [10,000]
Patience | [20]
Noise Decay | [.55]
Decay Steps | [5,000]
Warmup Steps | [100]
Initial Learning Rate |  [$10^{-7}$]
Table 8. Hyper parameter spaces for all algorithms used during training.

In Table 9 we show the values of hyper-parameters for the auto model.

Parameter Name | Parameter Value
---|---
Learning rate |  [$10^{-4}$]
Layer Size | [527]
Number of Layers | [2]
Training Steps, | [2,200]
Decay Steps | [20,000]
Warmup Steps | [100]
Table 9.  Hyper parameter specifications for auto model.

##  Appendix D Flu Test Period

In Figure D we can see a zoomed in area of the test period. We notice that the
model is well correlated with the actual flu rates, and most of the error
comes from over predicting the size of the peak, where the model has the most
uncertainty. For the rest of the test period, the model is very accurately
predicting the actual values, which we note is impressive considering that no-
lagged features were included like in autoregressive models.

Figure 6. National U.S. Flu Season Evaluation for 2022-2023.

##  Appendix E Flu Visualizations

Figure 7. A sample of the queries and two search embeddings: summer and
winter. The summer embedding is closest to a ”cough” cluster, while the winter
search embedding is closest to a a ”flu” cluster.

We visualize the flu category term embeddings, the model term scores, and the
summer and winter search embeddings in Figure 7. We show the search embeddings
for one Summer (July 15, 2022) day, and one Winter (November 15, 2022) day in
order to show the difference in search embeddings in trough vs peak flu rates.
We can see that the Winter embedding is closer to flu-related search terms,
indicating that these types of searches increase during peak flu season, while
the Summer embedding is closer to non-flu terms such as those related to
coughs. Table 10 shows the 5-Nearest Neighbor clusters for both the Summer and
Winter search embeddings.

Rank | Summer 5-NN | Winter 5-NN
---|---|---
0 | cough suppressant | flu
1 | benzonate cough | flu pneumonia
2 | people coughing | sore throat flu
3 | cough | flu quarantine
4 | sore throat flu | flu vaccine information statement
Table 10. 5-Nearest neighbor search term clusters for Summer and Winter search
embeddings. Search terms clusters are sorted and ranked according to their
distances in the embedding space.

##  Appendix F Training Loss Function

The loss function that we found to empirically work better is the following
adjusted weighted MAPE loss

(13)
$L(y,\hat{y})=\frac{1}{\sum_{i}w_{i}}\sum_{i=1}^{N}w_{i}\cdot\Big{[}\frac{|y_{i}-\hat{y_{i}}|}{|y_{i}|+\epsilon}+\frac{|y_{i}-\hat{y_{i}}|}{|\hat{y_{i}}|+\epsilon}\Big{]}$

where $w_{i}$ is the weight for example $i$, which we typically set to be
$|y_{i}|$ and $\epsilon$ is a small constant introduced for numerical
stability. We use the following Lasso regularized version of (13)

(14) $L^{*}(y,\hat{y})=L(y,\hat{y})+\lambda\sum_{z\in\theta}|z|$

where $\lambda\in\mathbb{R}^{+}$ is a hyperparameter controlling the degree of
regularization on the parameters inside the probability model. We use (14)
unless otherwise stated.

##  Appendix G Marginal Distribution Aggregate

All the analyses in our paper used aggregated search terms according to (1),
where the embeddings are simply summed and later projected back onto the unit
sphere. This leads to a potential dilution of information, where important
statistics of individual terms are ”averaged out.” To combat this, we explore
an alternative approach that fits under the umbrella of SLaM by keeping the
resulting dimensionality of the search features $\mathcal{O}(D)$ by taking the
marginal distributions of embeddings for each day and geo region. We use all
$D$ marginal distributions as the search embedding. For each day and region,
we take the marginal distribution of each emebdding dimension as the
$K$-binned histogram for that dimension on that day in that region. This
yields a search embedding of size $K\cdot D$, which we L2-normalize and feed
into our model. The comparison of this approach using $K=100$ evenly spaced
bins along the interval of $[-1,1]$ against the typical summation approach is
shown in Table 11.

Aggregation Method | Test MAPE (%) | Test $r$
---|---|---
Summation (1) | $5.46\pm 0.43$ | $.9933\pm.0005$
Marginal Distribution | $6.86\pm 0.35$ | $.9952\pm.0004$
Table 11. Comparing the marginal distribution search embeddings to the
traditional summation search embeddings for the Flu National modeling task.

We can see that the marginal distribution search embeddings preforms similarly
to the summation technique. We believe that the slightly worse performance may
be due to the large increase in model features, which could be ameliorated
with different regularization techniques. Additionally, we note that because
the summation technique performs well compared to a more granular version of
the same search features, this additional granularity might not be needed in
many cases. We’d like to explore other aggregation methods and nuances of our
existing approaches in future research.

##  Appendix H Embedding Comparisons

In Table 12 we highlight the main attributes of the different ways to embed
search data. Our proposed approach maximizes the percentage of search terms
included without sacrificing on memory, while embedding terms in a continuous
space.

| Embedding
---
Method
| Unique Terms
---
Included
| % Terms
---
Included
| Embedding
---
Dimension
| Embedding
---
Space
| Memory
---
Requirement
| Practical Memory
---
Requirement
One-hot | $|S|$ | 100% | $|S|$ | Discrete | $\mathcal{O}(|T|\cdot|S|)$ | Gigabytes
Filtered one-hot | $|A|$ | <1% | $|A|$ | Discrete | $\mathcal{O}(|T|\cdot|A|)$ | Megabytes
Classification | $|S|$ | 100% | $|C|$ | Discrete | $\mathcal{O}(|T|\cdot|C|)$ | Megabytes
SLaM | $|S|$ | 100% | $\mathcal{O}(D)$ | Continuous | $\mathcal{O}(|T|\cdot D)$ | Megabytes
Table 12. A comparison of the different ways to transform raw search data into
usable features. $|C|$ represents the number of classes in the classifier. For
the Practical Memory Requirement column, we assume $|S|\approx 10^{7}$.


# Accelerating the inference of string generation-based
chemical reaction models for industrial applications

Mikhail Andronov  Natalia Andronova  Michael Wand  Jürgen Schmidhuber  Djork-
Arné Clevert

###### Abstract

Template-free SMILES-to-SMILES translation models for reaction prediction and
single-step retrosynthesis are of interest for industrial applications in
computer-aided synthesis planning systems due to their state-of-the-art
accuracy. However, they suffer from slow inference speed. We present a method
to accelerate inference in autoregressive SMILES generators through
speculative decoding by copying query string subsequences into target strings
in the right places. We apply our method to the molecular transformer
implemented in Pytorch Lightning and achieve over 3X faster inference in
reaction prediction and single-step retrosynthesis, with no loss in accuracy.

Reaction models, inference speed, transformers

##  1 Introduction

Automated planning of organic chemical synthesis, first formalized around
fifty years ago (Pensak & Corey, 1977), is one of the core technologies
enabling computer-aided drug discovery. While first computer-aided synthesis
planning (CASP) systems relied on manually encoded rules (Johnson et al.,
1989; Gasteiger et al., 2000), researchers now primarily focus on CASP methods
powered by artificial intelligence techniques. The design principles of the
latter were outlined in the seminal work by Segler et al. (Segler et al.,
2018): a machine learning-based single-step retrosynthesis model combined with
a planning algorithm. The former proposes several candidate retrosynthetic
chemical transformations for a given molecule, and the latter, e.g., Monte-
Carlo Tree Search, uses those candidates to construct a synthesis tree.
Single-step retrosynthesis models are now commonly developed in two paradigms:
template-based models and template-free models. Besides retrosynthesis, one
can also build a model to predict the products of chemical reactions (Figure
1).

Figure 1: Both reaction product prediction and single-step retrosynthesis can
be formulated as SMILES-to-SMILES translation and approached with a model like
an encoder-decoder transformer.

The principle of template-based models is to use a set of SMIRKS templates
extracted from reaction data and a machine learning model for classification
or ranking to select a matching template for a query SMILES that will, upon
application, transform the query into the product SMILES (for product
prediction), or the SMILES of possible reactants (for single-step
retrosynthesis). In contrast, in template-free models, the query transforms
into the result without resorting to SMIRKS templates, e.g., with a sequence
of predicted graph edits (Sacha et al., 2021; Bradshaw et al., 2018) or
through “translation” of the query SMILES into the desired SMILES with a
conditioned text generation model (Schwaller et al., 2019; Tetko et al., 2020;
Irwin et al., 2022). While CASP systems leveraging template-based single-step
models proved to be effective (Genheden et al., 2020), there is an interest in
building CASP with template-free models instead, as they demonstrate state-of-
the-art accuracy in both single-step retrosynthesis and reaction product
prediction. Most accurate template-free models are currently conditional
autoregressive SMILES generators based on the transformer architecture
(Vaswani et al., 2017), which also serves as the backbone for Large Language
Models (LLM) (Brown et al., 2020; Zhao et al., 2023). Unfortunately, the high
accuracy of autoregressive models like Chemformer (Irwin et al., 2022) comes
at the cost of a slow inference speed (Torren-Peraire et al., 2024), which
hinders their successful adoption as part of industrial CASP systems. In our
work, we propose a method to accelerate inference from SMILES-to-SMILES
translation models based on speculative decoding (Leviathan et al., 2023; Xia
et al., 2023), a general technique for LLM inference acceleration, combined
with insights from the chemical essence of the problem. We reimplement the
Molecular Transformer (Schwaller et al., 2019) in Pytorch Lightning and use
our method to demonstrate its inference acceleration in single-step
retrosynthesis and product prediction by up to three times and more without
changing the model architecture, training procedure and without compromising
on accuracy.

##  2 Methods

###  2.1 Algorithm

Autoregressive models, such as Transformer variants (Vaswani et al., 2017;
Brown et al., 2020; Schmidhuber, 1992), generate sequences token by token, and
every prediction of the next token requires a forward pass of the model. Such
a process may be computationally expensive, especially for models with
billions of parameters. Therefore, an intriguing question arises whether could
one generate several tokens in one forward pass of the model, thus completing
the output faster. Speculative decoding (Xia et al., 2023; Leviathan et al.,
2023) answers positively. Recently proposed as a method of inference
acceleration for Large Language Models, it is based on the draft-and-verify
idea. At every generation step, one can append some draft sequence to the
sequence generated by the model so far and see if the model ”accepts” the
draft tokens.

Figure 2: Speculative decoding accelerates product prediction with the
molecular transformer or a similar autoregressive SMILES generator. Before
generating an output sequence, we prepare a list of subsequences of a desired
length, e.g., four, of the tokenized query SMILES of reactants. Then, at every
generation step, the model can copy up to four tokens from one of the draft
sequences to the output, thus generating from one to five tokens in one
forward pass.

If the draft sequence has length $N$, in the best case, the model adds $N+1$
token to the generated sequence in one forward pass, and in the worst case, it
adds one token as in standard autoregressive generation. The acceptance rate
for one generated sequence is the number of accepted draft tokens divided by
the total number of tokens in the generated sequence. One can also test
multiple draft sequences in one forward pass taking advantage of
parallelization, and choose the best one. Speculative decoding does not affect
the content of the predicted sequence compared to the token-by-token decoding
in any way.

One can freely choose a way of generating draft sequences. For LLMs, one would
usually use another smaller language model that performs its forward pass
faster than the main LLM (Leviathan et al., 2023) or additional generation
heads on top of the LLM’s backbone (Cai et al., 2024). However, one can also
construct draft sequences without calling any learned functions. For example,
generate random draft sequences, even though their acceptance rate will be
minimal, or assemble draft sequences out of tokens in the query sequence — a
prompt for decoder-only language models or a source sequence for translation
models. The latter option is perfect for retrosynthesis or reaction prediction
as SMILES-to-SMILES translation. In a chemical reaction, large fragments of
reactants typically remain unchanged, which means that the SMILES of products
and reactants have many common substrings. It is especially true if reactant
and product SMILES are aligned to minimize the edit distance between them
(Zhong et al., 2022). Therefore, we can extract multiple substrings of a
chosen length $N$ from the query SMILES and use them as draft sequences with a
high acceptance rate. Figure 2 demonstrates this method in product prediction.
Before generating the target string, we assemble a list of token subsequences
from the source sequence (reactant tokens) with a sliding window of a desired
length (4 in this case) and stride 1. Then, at every generation step, we can
try all draft sequences in one forward pass of the model to see if the model
can copy up to 4 tokens from one of them. The draft token acceptance rate in
this example reaches 78%.

Speculative decoding does not require any changes to the model architecture or
training of additional models. The cost of generating draft sequences in this
way is negligible compared to that of the forward pass of the reaction model.

###  2.2 Model

We demonstrate the application of our method to the Molecular Transformer
(Schwaller et al., 2019). It is an encoder-decoder transformer model suitable
for SMILES-to-SMILES translation. We conduct our experiments on one H100 GPU
with 80 GB memory.

The original Molecular Transformer (Schwaller et al., 2019) adopts OpenNMT
(Klein et al., 2018), a general framework for neural machine translation, for
SMILES-to-SMILES translation. Since the code in this framework is complex and
intractable to customize, we decided to re-implement the model in PyTorch
Lightning to keep only the necessary code and have more design freedom in the
model’s inference procedure implementation.

Table 1: The top-5 accuracy of the Molecular Transformer (MT) in predicting
reaction products on USPTO MIT (mixed) in the original report and our
reimplementation. Both models use beam search with beam size 5 to generate
five predictions for every query. Our model successfully reproduces the
accuracy of the original model with negligible discrepancy.
Accuracy | Original MT | Our MT | $\Delta$
---|---|---|---
Top-1, % | 88.6 | 88.4 | -0.2
Top-2, % | 92.4 | 92.5 | 0.1
Top-3, % | 93.5 | 93.7 | 0.2
Top-5, % | 94.5 | 94.7 | 0.2

###  2.3 Data

We used the open reaction data from US patents (Lowe, 2012) for training all
models. We trained the model for reaction product prediction as in the
original paper (Schwaller et al., 2019) on the USPTO MIT dataset, a standard
benchmark for product prediction, without reactant-reagent separation. We
trained the model for single-step retrosynthesis on USPTO 50K, a standard
dataset for the task. In this dataset, we augmented every reaction in the
training set 20 times using SMILES augmentation (Tetko et al., 2020) with
root-aligned SMILES (Zhong et al., 2022). We followed the standard atomwise
tokenization procedure (Schwaller et al., 2019) to tokenize SMILES.

##  3 Results and Discussion

Our implementation of the Molecular Transformer (MT) successfully reproduces
the accuracy scores of the original MT (Schwaller et al., 2019) that relies on
OpenNMT. Comparing our MT and the original MT, we observe at most 0.2
percentage points discrepancy of top-1 to top-5 accuracy in product prediction
with beam search (Table 1). Having verified the correctness of our MT’s
implementation in this way, we replace the standard decoding procedures for
the MT with our methods based on speculative decoding and achieve a
significant speed-up in product prediction and single-step retrosynthesis.

###  3.1 Product prediction

We tested our MT for product prediction on USPTO MIT mixed, i.e., without an
explicit separation between reactants and reagents. The test dataset in this
benchmark comprises 40 thousand reactions.

When serving a reaction prediction model as an AI assistant for chemists, one
could use greedy decoding with a batch size of 1 for inference. Table 2
summarizes our experiments with greedy decoding from MT on the test set of
USPTO MIT. The model’s inference with standard greedy decoding with a batch
size of 1 finishes in around 62 minutes. In contrast, if we use greedy
generation enhanced with our speculative decoding, the inference time reduces
to 26 minutes with a draft length of 4 and 17 minutes with a draft length of
10, which corresponds to 137 % and 262 % speedup, respectively. The acceptance
rate in our drafting method averaged over all test reactions is 79%.
Potentially, it can be even higher if one adds more draft sequences to choose
from, for example, subsequences of the source sequence dilated by one token.
Of course, greedy decoding with a large batch size is much faster and
completes in around 4 minutes with 32 reactions in a batch. However,
accelerating inference with a batch size of 1 would be sufficient for an
improved user experience with reaction prediction assistants: chemists would
usually enter one query at a time in a user interface of a reaction model like
IBM RXN. The model’s accuracy is 88.3% with both standard and speculative
greedy decoding.

Table 2: Wall time of the model’s inference on the USPTO MIT test set in
reaction product prediction with standard and speculative greedy decoding. B
stands for batch size, and DL stands for draft length. The time is averaged
over five attempts.
Decoding | Time, min
---|---
Greedy (B=1) |  61.8 $\pm$ 5.88
Greedy speculative (B=1, DL=4) |  26.04 $\pm$ 2.07
Greedy speculative (B=1, DL=10) |  17.06 $\pm$ 0.25
Greedy (B=32) |  4.13 $\pm$ 0.06
Table 3: Wall time of the model’s inference on the USPTO 50k test set (5K
reactions) in single-step retrosynthesis with beam search (BS) and speculative
beam search (SBS). Batch size is 1. DL stands for draft length, and beam width
is denoted as $n$. The time is averaged over five attempts.
Decoding |  $n$=5, min |  $n$=10, min |  $n$=25, min
---|---|---|---
BS |  36.7 $\pm$ 0.3 |  39.9 $\pm$ 0.3 |  46.2 $\pm$ 0.3
SBS, DL=10 |  9.9 $\pm$ 0.1 |  15.4 $\pm$ 0.1 |  28.1 $\pm$ 0.1
SBS, DL=0 |  23.1 $\pm$ 0.3 |  25.7 $\pm$ 0.3 |  34.6 $\pm$ 0.3

###  3.2 Single-step retrosynthesis

We carried out single-step retrosynthesis experiments on USPTO 50K, in which
the training dataset was augmented 20 times. The augmentation procedure is to
construct alternative root-aligned (Zhong et al., 2022) SMILES for every
dataset entry. This augmentation minimizes the edit distance between reactants
and products, which simplifies training, pushes the model’s accuracy higher,
and increases the acceptance rate in our speculative decoding method.
Speculative decoding in single-step retrosynthesis accelerates greedy decoding
as much as in reaction product prediction. However, this has limited utility.
In synthesis planning, one would want a single-step retrosynthesis model to
suggest multiple different reactant sets for every query product so that the
planning algorithm can choose from them. Usually, one would employ beam search
to generate several outputs from the transformer for single-step
retrosynthesis.

We found a way to accelerate beam search with speculative decoding. The main
idea behind it is that at every decoding step, we use a draft sequence to
generate multiple candidate sequences of different lengths in one forward
pass, which we then order by probabilities and keep the best $n$ of them. To
organize such sequences of unequal lengths in a batch for a subsequent forward
pass of the model, we pad them from the left and offset the positional
encodings accordingly for every sequence. We call our algorithm ”speculative
beam search” (SBS). The full description of the algorithm is in the Appendix
B.

Table 4: The top-N accuracy of our single-step retrosynthesis model on USPTO
50K with beam search (BS) and speculative beam search (SBS). DL stands for
draft length.
Accuracy | BS | SBS, DL=10 | SBS, DL=0
---|---|---|---
Top-1, % | 52.07 | 52.07 | 52.07
Top-3, % | 75.16 | 75.16 | 75.16
Top-5, % | 82.07 | 82.07 | 82.07
Top-10, % | 89.08 | 89.08 | 89.08
Top-25, % | 94.07 | 94.05 | 94.09

The wall time our retrosynthesis model takes to process the USPTO 50K test set
with beam search and batch size 1 is around 37 minutes, 40 minutes, and 46
minutes when generating 5, 10, and 25 predictions for every query SMILES,
respectively (Table 3). We keep the beam width and the number of best
sequences equal. When we replace the standard beam search with our SBS with
the draft length of ten, the generation finishes in around 10, 15, and 28
minutes, accelerating the inference by around 3.7, 2.7, and 1.8 times,
respectively.

SBS reduces to the standard beam search when draft tokens are never accepted.
It happens, for example, if we use the start-of-sequence token as a single
draft. The ”effective” batch size also does not increase when using only one
draft per forward pass (see Section 3.3). We added an experiment in this mode,
denoting it as ”SBS, DL=0” (Tables 3 and 4), to verify that the algorithm’s
acceleration comes from the draft sequences. Interestingly, SBS is faster than
the standard beam search even in case of DL=0 (Table 3).

The top-N accuracy when using SBS is practically the same as the accuracy of
the standard beam search (Table 4). Only the difference in the top-25 accuracy
is visible, although tiny, only a couple hundredths of a percent point. Thus,
our SBS accelerates the MT’s generation of multiple reactant sets several
times without having to compromise on accuracy at all. Such a speed-up could
make the transformer a more attractive single-step model for multi-step
synthesis planning.

###  3.3 Limitations

The speed of the model’s forward pass decreases with the increase of the size
of a batch input to the transformer decoder. This effect quickly manifests
itself in our drafting strategy. The latter is ”brute force” in some sense, as
we use various substrings of the source SMILES as drafts in parallel at every
forward pass of the model. It inflates the ”effective” batch size of the input
to the transformer decoder: we copy every sequence in the batch as many times
as there are drafts and concatenate the sequence copy and the corresponding
draft. As a result, the generation of target SMILES may become slower with
speculative decoding than with standard decoding procedures, even when it
requires fewer calls to the model. We put a boundary on the number of drafts
extracted from a query sequence to mitigate this deceleration. However, this
compromises the acceptance rate.

The speed of the model’s forward pass decreases with the increase of the size
of a batch input to the transformer decoder. This effect quickly manifests
itself in our drafting strategy. The latter is ”brute force” in some sense, as
we use various substrings of the source SMILES as drafts in parallel at every
forward pass of the model. It inflates the ”effective” batch size of the input
to the transformer decoder: we copy every sequence in the batch as many times
as there are drafts and concatenate the sequence copy and the corresponding
draft. As a result, the generation of target SMILES may become slower with
speculative decoding than with standard decoding procedures, even when it
requires fewer calls to the model. We put a boundary on the number of drafts
extracted from a query sequence to mitigate this deceleration. However, this
compromises the acceptance rate. In addition, as the size of the input batch
grows, the generation speed with speculative decoding becomes bottlenecked in
terms of the number of calls to the model. The ”least lucky” sequence with the
lowest acceptance rate of speculative tokens would determine the number of
calls. Together, these two effects limit the utility of our speculative
decoding for large batch sizes and beam widths. For example, our SBS is slower
than the standard beam search when the beam size is fifty. Nonetheless, we
consider speculative decoding very promising for synthesis planning
acceleration, as it significantly improves the speed of the standard beam
search without losing accuracy in the practical range of beam widths
($\sim$10-20).

Designing a drafting strategy for SMILES that removes the need for multiple
parallel drafts while retaining a high acceptance rate is an aspect of our
ongoing work.

Our speculative algorithm works well in single-step retrosynthesis because the
MT typically predicts low-entropy next-token distributions in this task. In
most cases, all predicted probability mass concentrates on one token.
Consequently, long sequences often win in the probability competition against
shorter sequences at every speculative beam search iteration. If several top
probabilities from the predicted next-token distribution are approximately the
same, then short sequences would tend to have higher probabilities, and the
model would generate only a small number of tokens per forward pass,
benefitting little from drafts.

##  4 Conclusion

We combine speculative decoding and chemical insights to accelerate inference
in the molecular transformer, a SMILES-to-SMILES translation model. Our method
makes processing the test set more than three times faster in both single-step
retrosynthesis on USPTO 50K and reaction product prediction on USPTO MIT
compared to the standard decoding procedures. Our method aims at making state-
of-the-art template-free SMILES-generation-based models such as the molecular
transformer more suitable for industrial applications such as computer-aided
synthesis planning systems.

## Acknowledgements

This study was partially funded by the European Union’s Horizon 2020 research
and innovation program under the Marie Skłodowska-Curie Innovative Training
Network European Industrial Doctorate grant agreement No. 956832 “Advanced
machine learning for Innovative Drug Discovery, and also by the Horizon Europe
funding programme under the Marie Skłodowska-Curie Actions Doctoral Networks
grant agreement “Explainable AI for Molecules - AiChemist” No. 101120466.

## Code availability

The code and checkpoints for trained models will be available at
https://github.com/Academich/translation-transformer.

## Conflict of interests

The authors have no conflicts of interests.

## References

  * Bradshaw et al. (2018) Bradshaw, J., Kusner, M. J., Paige, B., Segler, M. H., and Hernández-Lobato, J. M.  A generative model for electron paths.  _arXiv preprint arXiv:1805.10970_ , 2018.
  * Brown et al. (2020) Brown, T., Mann, B., Ryder, N., Subbiah, M., Kaplan, J. D., Dhariwal, P., Neelakantan, A., Shyam, P., Sastry, G., Askell, A., et al.  Language models are few-shot learners.  _Advances in Neural Information Processing Systems_ , 33:1877–1901, 2020.
  * Cai et al. (2024) Cai, T., Li, Y., Geng, Z., Peng, H., Lee, J. D., Chen, D., and Dao, T.  Medusa: Simple LLM inference acceleration framework with multiple decoding heads.  _arXiv preprint arXiv:2401.10774_ , 2024.
  * Gasteiger et al. (2000) Gasteiger, J., Pförtner, M., Sitzmann, M., Höllering, R., Sacher, O., Kostka, T., and Karg, N.  Computer-assisted synthesis and reaction planning in combinatorial chemistry.  _Perspectives in Drug Discovery and Design_ , 20:245–264, 2000.
  * Genheden et al. (2020) Genheden, S., Thakkar, A., Chadimová, V., Reymond, J.-L., Engkvist, O., and Bjerrum, E.  Aizynthfinder: a fast, robust and flexible open-source software for retrosynthetic planning.  _Journal of Cheminformatics_ , 12(1):70, 2020.
  * Irwin et al. (2022) Irwin, R., Dimitriadis, S., He, J., and Bjerrum, E. J.  Chemformer: a pre-trained transformer for computational chemistry.  _Machine Learning: Science and Technology_ , 3(1):015022, 2022.
  * Johnson et al. (1989) Johnson, P. Y., Burnstein, I., Crary, J., Evans, M., and Wang, T.  Designing an expert system for organic synthesis: the need for strategic planning.  ACS Publications, 1989.
  * Klein et al. (2018) Klein, G., Kim, Y., Deng, Y., Nguyen, V., Senellart, J., and Rush, A. M.  OpenNMT: Neural Machine Translation Toolkit, 2018.
  * Leviathan et al. (2023) Leviathan, Y., Kalman, M., and Matias, Y.  Fast inference from transformers via speculative decoding.  In _International Conference on Machine Learning_ , pp. 19274–19286. PMLR, 2023.
  * Lowe (2012) Lowe, D. M.  Extraction of chemical structures and reactions from the literature. Ph.D. dissertation, University of Cambridge, Cambridge, UK.  https://doi.org/10.17863/CAM.16293, 2012.
  * Pensak & Corey (1977) Pensak, D. A. and Corey, E. J.  LHASA—logic and heuristics applied to synthetic analysis.  ACS Publications, 1977.
  * Sacha et al. (2021) Sacha, M., Błaz, M., Byrski, P., Dabrowski-Tumanski, P., Chrominski, M., Loska, R., Włodarczyk-Pruszynski, P., and Jastrzebski, S.  Molecule edit graph attention network: modeling chemical reactions as sequences of graph edits.  _Journal of Chemical Information and Modeling_ , 61(7):3273–3284, 2021.
  * Schmidhuber (1992) Schmidhuber, J.  Learning to control fast-weight memories: An alternative to recurrent nets.  _Neural Computation_ , 4(1):131–139, 1992.
  * Schwaller et al. (2019) Schwaller, P., Laino, T., Gaudin, T., Bolgar, P., Hunter, C. A., Bekas, C., and Lee, A. A.  Molecular transformer: a model for uncertainty-calibrated chemical reaction prediction.  _ACS central science_ , 5(9):1572–1583, 2019.
  * Segler et al. (2018) Segler, M. H., Preuss, M., and Waller, M. P.  Planning chemical syntheses with deep neural networks and symbolic AI.  _Nature_ , 555(7698):604–610, 2018.
  * Tetko et al. (2020) Tetko, I. V., Karpov, P., Van Deursen, R., and Godin, G.  State-of-the-art augmented NLP transformer models for direct and single-step retrosynthesis.  _Nature Communications_ , 11(1):5575, 2020.
  * Torren-Peraire et al. (2024) Torren-Peraire, P., Hassen, A. K., Genheden, S., Verhoeven, J., Clevert, D.-A., Preuss, M., and Tetko, I. V.  Models matter: the impact of single-step retrosynthesis on synthesis planning.  _Digital Discovery_ , 3(3):558–572, 2024.
  * Vaswani et al. (2017) Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, Ł., and Polosukhin, I.  Attention is all you need.  _Advances in Neural Information Processing Systems_ , 30, 2017.
  * Xia et al. (2023) Xia, H., Ge, T., Wang, P., Chen, S.-Q., Wei, F., and Sui, Z.  Speculative decoding: Exploiting speculative execution for accelerating seq2seq generation.  In _Findings of the Association for Computational Linguistics: EMNLP 2023_ , pp. 3909–3925, 2023.
  * Zhao et al. (2023) Zhao, W. X., Zhou, K., Li, J., Tang, T., Wang, X., Hou, Y., Min, Y., Zhang, B., Zhang, J., Dong, Z., et al.  A survey of large language models.  _arXiv preprint arXiv:2303.18223_ , 2023.
  * Zhong et al. (2022) Zhong, Z., Song, J., Feng, Z., Liu, T., Jia, L., Yao, S., Wu, M., Hou, T., and Song, M.  Root-aligned SMILES: a tight representation for chemical reaction prediction.  _Chemical Science_ , 13(31):9023–9034, 2022.

##  Appendix A Training details.

For product prediction, we train this model with the same hyperparameters as
in Schwaller et al. with four encoder and decoder layers, eight heads,
embedding dimensionality of 256, and feedforward dimensionality of 2048, which
results in 11,4 million parameters. For single-step retrosynthesis, we set the
hyperparameters as in Zhong et al. (six encoder and decoder layers, eight
heads, embedding dimensionality of 256, and feedforward dimensionality of
2048), which results in 17,4 million parameters. The dictionary is the same
for the encoder and the decoder in both models. We use the Adam optimizer for
both models.

##  Appendix B Speculative Beam Search

The outline of the speculative beam search procedure we implement for our
experiments is as follows:

Algorithm 1 Speculative beam search

Input: Query token sequence (tensor of integers) $s$, number of best sequences
$n$

$memory=encoder(s)$

$drafts=getDrafts(s)$

$res=[BOS]$

repeat

$draftedSequences=concatDraftsToSequences(res,drafts)$

$logits=decoder(draftedSequences,memory)$

$bestDraftedSequences=selectBestDraft(logits)$

$candidates=sample(bestDraftedSequence)$

$bestCandidates=sortAndExtract(candidates,n)$

$res=padLeft(bestCandidates)$

until maximal length or number of iterations is reached

## Functions

getDrafts: Produces a list of subsequences of the query sequence using a
sliding window of a given length with stride one. Only the first $N_{d}$
drafts are kept and the others are discarded to mitigate increasing memory
requirements as the number of drafts grows. $N_{d}$ is typically around 25.

encoder: Forward pass of the encoder of the Molecular Transformer.

decoder: Forward pass of the decoder of the Molecular Transformer.

concatDraftsToSequences: Concatenates every draft to every resulting sequence
decoded so far. This gives a tensor of sequences with the batch size of
length($res$) $\times$ length ($drafts$).

Input: decoded sequences $res$, tensor of drafts $drafts$

$draftedSequences=[]$

for $r$ in $res$ do

for $d$ in $drafts$ do

$draftedSequences.append(r\mathbin{\|}d)$

end for

end for

return $draftedSequences$

selectBestDraft: Uses the logits predicted by the decoder to select the best
draft for every resulting sequence and discard the others. The best draft is
the one with the largest number of accepted tokens.

sample: Proposes many candidate sequences for the decoding step based on the
single accepted draft. The candidate sequences may have different lengths.
Figure 3 demonstrates an example.

Figure 3: An example of two first iterations of the sampling of candidate
sequences in speculative beam search. Here, we select the two best candidates
at each iteration. The first forward pass generates 12 candidate sequences.
The second forward pass generates 24 sequences. The draft length in this
example is 10. The best sequences in the first iteration are c1c[nH] and
c1cn(C(=O)O. The best sequences after the second iteration are
c1c[nH]c2ccc(C(C)= and c1cn(C(=O)OC(C)(C)C)c2.

sortAndExtract: Orders the candidate sequences by probabilities, keeps the $n$
ones with the highest probabilities and discards the others.

padLeft: If the best candidate sequences have different lengths, organizes
them in a tensor by padding them with a special PAD token from the left. In
the next step of the algorithm, the starting positions for the positional
encodings in the transformer decoder get shifted accordingly.