# Amortized nonmyopic active search
via deep imitation learning

Quan Nguyen Anindya Sarkar Roman Garnett
Washington University in St. Louis
<EMAIL_ADDRESS>
###### Abstract

Active search formalizes a specialized active learning setting where the goal
is to collect members of a rare, valuable class. The state-of-the-art
algorithm approximates the optimal Bayesian policy in a budget-aware manner,
and has been shown to achieve impressive empirical performance in previous
work. However, even this approximate policy has a superlinear computational
complexity with respect to the size of the search problem, rendering its
application impractical in large spaces or in real-time systems where
decisions must be made quickly. We study the amortization of this policy by
training a neural network to learn to search. To circumvent the difficulty of
learning from scratch, we appeal to imitation learning techniques to mimic the
behavior of the expert, expensive-to-compute policy. Our policy network,
trained on synthetic data, learns a beneficial search strategy that yields
nonmyopic decisions carefully balancing exploration and exploitation.
Extensive experiments demonstrate our policy achieves competitive performance
at real-world tasks that closely approximates the expert’s at a fraction of
the cost, while outperforming cheaper baselines.

##  1 Introduction

Many problems in science and engineering share a common theme where an agent
searches for rare and valuable items in a massive database; examples include
fraud detection, product recommendation, and drug and materials discovery. The
bottleneck of this procedure is often the cost of labeling, that is,
determining whether a candidate is one of the targets of the search. For
instance, in product recommendation, labeling can involve presenting a
customer with a product they might enjoy, at the risk of interrupting the
customer’s shopping experience and consequently losing sales; in drug
discovery, it takes time-consuming computer simulations and/or expensive
laboratory experiments to characterize a potential drug. This labeling cost
rules out exhaustive screening and motivates strategic exploration of the
search space. Active search (as) frames this problem in the language of
adaptive experimental design, and aims to develop policies that iteratively
choose data points to label to uncover as many valuable points as possible
under a labeling budget.

as has been thoroughly studied in previous work and sophisticated search
policies have been developed [13, 18, 19, 20]. Of particular interest is the
work of Jiang et al. [18], who derived the Bayesian optimal policy under the
simplifying assumption that future experiments are chosen simultaneously in a
batch. Here, the number of future experiments is set to be the remaining
labeling budget so that this remaining budget is actively accounted for during
policy computation. The authors called the resulting policy _efficient
nonmyopic search_ (ens), which can be viewed as a budget-aware approximation
to the true optimal policy. They showed this budget-awareness induces
nonmyopic decisions that automatically balance between strategic exploration
of the space and timely exploitation of regions that likely yield many
targets, and ultimately achieve state-of-the-art (sota) search performance
across many tasks. Although the aforementioned simplifying assumption,
combined with aggressive pruning, allows ens to be feasibly applied to
problems of considerable size (100 000+ in Jiang et al. [18, 19, 20], for
example), the policy retains a _superlinear_ computational complexity. This
complexity poses a challenge in (1) deploying in real-time applications where
decisions need to be made quickly and (2) scaling to large spaces. For
instance, a guided data discovery task [37] in visual analytics combines a
search algorithm with an interactive visualization to assist a user with their
analytic goals in real time, and the time available for the search algorithm
to run is thus severely constrained. Modern recommender systems such as
YouTube and Amazon must quickly search over millions of items to make
recommendations for a large number of users [10]; similarly, there exist
databases with billions of synthesizable molecules acting as search spaces for
drug discovery [51].

We aim to alleviate the computational cost of budget-aware search by training
a small, relatively shallow feedforward neural network to mimic the behavior
of the sota, expert policy ens; policy computation is thus amortized as we
deploy the trained network as the search policy. We train this policy using
the imitation learning technique DAgger [43], which aids the goal of behavior
cloning by iteratively querying the expert’s actions on states encountered by
the network being trained. This procedure is done with small, synthetic search
problems where ens is cheap to query. We find that the trained policy network
successfully learns a beneficial nonmyopic search behavior and, despite the
synthetic training data, incurs only a minor decrease in search performance at
real-world tasks, in exchange for much faster decision-making. We showcase the
usefulness of this computationally lightweight policy with a wide range of
search problems spanning diverse applications, including drug discovery tasks
of an unprecedented multi-million scale.

##  2 Preliminaries

We first introduce the problem setting, the active search framework, the sota
search policy and its computational complexity to motivate our goal of
amortization.

###  2.1 Active search and the optimal policy

An active search (as) problem is defined by a finite set of data points
$\mathcal{X}\triangleq\\{x_{i}\\}$, among which there exists a rare, valuable
subset $\mathcal{T}\subset\mathcal{X}$. We use the term “targets” to refer to
the members of this valuable subset, which we wish to collect from the entire
space $\mathcal{X}$. We further use membership in $\mathcal{T}$ as the labels
for the points in $\mathcal{X}$:
$y_{i}\triangleq\mathbb{I}\,[x_{i}\in\mathcal{T}],\forall
x_{i}\in\mathcal{X}$. The targets are not known _a priori_ , but whether a
specific data point $x_{i}$ is one can be determined by querying an oracle
returning the requested label $y_{i}$. This oracle models the process of
performing laborious experiments to characterize the data point in question
(e.g., suggesting a given product to a customer and observing their subsequent
clicking behavior, or performing computer simulations and laboratory
experiments on a candidate molecule for drug discovery). We thus assume the
oracle to be expensive to query. Concretely, we assume we may only query the
oracle $T$ times, where $T\ll n\triangleq|\mathcal{X}|$; $T$ can be viewed as
our querying budget. We denote the set of data points and their revealed
labels as $\mathcal{D}=\\{(x_{i},y_{i})\\}$, and the set of those queried up
to iteration $t\leq T$ as $\mathcal{D}_{t}=\\{(x_{i},y_{i})\\}_{i=1}^{t}$. The
goal of as is to design a policy that sequentially selects which data points
to query so as to find as many targets throughout the $T$ iterations of the
search as possible. To express this preference for maximizing the number of
“hits” across different terminal data sets collected at the end of the search
$\mathcal{D}_{T}$, we use the utility function
$u(\mathcal{D}_{T})=\sum_{y_{i}\in\mathcal{D}_{T}}y_{i}$, which simply counts
the number of targets in $\mathcal{D}_{T}$.

Previous works on as have derived the optimal policy under Bayesian decision
theory that finds the query with the highest expected terminal utility [13,
18]. First, we build a probabilistic classifier that outputs the probability
an unlabeled point $x$ has a positive label given the observed data set
$\mathcal{D}$, denoted as $\Pr(y=1\mid x,\mathcal{D})$. Then, at iteration
$t+1$, having collected $\mathcal{D}_{t}$, we query the data point:

$x_{t+1}^{*}=\operatorname*{arg\,max}_{x_{i+1}\in\mathcal{X}\setminus\mathcal{D}_{t}}\mathbb{E}\big{[}u(\mathcal{D}_{T})\mid
x_{i+1},\mathcal{D}_{t}\big{]},$ (1)

where the expectation is with respect to the label $y_{i+1}$, and the future
queries $(x_{t+2},x_{t+3},\ldots,x_{T})$ are similarly computed in this
optimal manner. While this computation can theoretically be achieved via
dynamic programming [4], it has a complexity of
$O\big{(}(2\,n)^{\ell}\big{)}$, where $n$ is again the size of the search
space and $\ell=T-t$ is the length of the decision-making horizon (i.e., the
remaining querying budget). This exponential blowup stems from the fact that
Bayesian decision theory compels the optimal policy to reason about not only
the possible labels of a putative query, but also how each possible label
affects subsequent future queries. This computation therefore cannot be
realized in almost all practical scenarios except for the very last few
iterations when $\ell$ is sufficiently small [13]. A common strategy that we
can adopt here is to limit the depth of the lookahead in this computation,
effectively pretending $\ell$ is indeed small. The simplest version of this is
obtained by setting $\ell=1$: we assume we have only one query remaining, and
the optimal decision becomes maximizing the expected marginal utility gain,
which is equivalent to greedily querying the most likely target:

$\displaystyle\begin{split}x_{T}^{*}&=\operatorname*{arg\,max}_{x_{T}\in\mathcal{X}\setminus\mathcal{D}_{T-1}}\mathbb{E}\Big{[}u\big{(}\mathcal{D}_{T-1}\cup\\{x_{T},y_{T}\\}\big{)}\mid
x_{T},\mathcal{D}_{T-1}\Big{]}\\\
&=\operatorname*{arg\,max}_{x_{T}\in\mathcal{X}\setminus\mathcal{D}_{T-1}}\Pr(y_{T}=1\mid
x_{T},\mathcal{D}_{T-1}).\end{split}$ (2)

We call this greedy policy the _one-step_ policy, as it computes the one-step
optimal decision. Albeit extremely computationally efficient, one-step does
not consider exploratory queries that may lead to future gains. As a result,
it tends to get stuck in small regions of targets, failing to discover larger
target sets due to insufficient exploration. Theoretically, Garnett et al.
[13] even showed that by limiting its lookahead, a myopic policy could perform
worse than a less myopic one _by any arbitrary degree_. Motivated by the
potential of nonmyopia in as, Jiang et al. [18] developed an approximation to
the optimal policy that accounts for the remaining budget $\ell$, which we
examine next.

###  2.2 Nonmyopic search via budget-awareness

To avoid the high cost of reasoning about the dependence among labels of
future queries, Jiang et al. [18] made the simplifying assumption that, after
our next query, all remaining future queries are made at the same time in a
batch. Under this assumption, the future queries in the lookahead in Eq. 1—to
be optimally chosen to maximize expected terminal utility—can be quickly
identified as the set of $(\ell-1)$ most likely targets [18]. Their policy ens
thus estimates the value of each putative query $x_{i}$ with the expected
utility of the union of $x_{i}$ and the top $(\ell-1)$ unlabeled points that
are adaptively selected based on each possible label $y_{i}$. Again, as the
number of future queries in this policy computation is set to exactly match
the true length of the decision-making horizon, ens actively accounts for the
remaining labeling budget when making its queries. The authors demonstrated
the benefits of this budget-awareness by showing that ens exhibits nonmyopic,
exploratory behavior when the budget is large, and automatically transitions
to more exploitative queries as search progresses. This strategic exploration
ultimately allows ens to outperform many search baselines including the one-
step policy.

While the aforementioned batch assumption avoids an exponential blowup in
computational complexity, ens still incurs a considerable cost, especially
under large values of $n$, the size of the search space. A naïve
implementation with a generic classifier has a complexity of
$O\big{(}n^{2}\log n\big{)}$. The official implementation by Jiang et al.
[18], on the other hand, uses a lightweight $k$-nearest neighbor (nn)
classifier that (reasonably) assumes a certain level of locality when
probabilities $\Pr(y\mid x,\mathcal{D})$ are updated in light of new data.
This structure allows for a faster computation of the batch of future queries
in ens’s lookahead, and brings the complexity down to $O\big{(}n\,(\log
n+m\log m+T)\big{)}$, where $m$ is the largest degree of any node within the
nearest neighbor graph corresponding to the $k$-nn [18]. Unfortunately, this
reduced complexity still poses a substantial challenge in two scenarios
commonly encountered in as: large search spaces (e.g., drug discovery) and
settings where queries must be rapidly computed (recommender and other real-
time systems). We address this problem by training an estimator, specifically
a neural network, to learn the mapping from possible candidate queries to the
output of ens, thus amortizing policy computation; the next section details
our approach.

##  3 Amortizing budget-aware active search

Our goal is to amortize search with a neural network, replacing the time-
consuming policy computation of ens with fast forward passes through the
network. Crucially, this network should learn a beneficial strategy that
outperforms the greedy one-step policy, so that the cost of training and
deploying the network outweighs one-step’s speed and ease of use. We now
discuss our approach using reinforcement learning, specifically imitation
learning, to effectively train one such network.

###  3.1 Learning to search with imitation learning

We start with the goal of training a neural network to learn to search using
reinforcement learning (rl), as the utility function in Sect. 2.1 can be
naturally treated as a reward function, and each search run of $T$ iterations
as belonging to a budget-constrained episodic Markov decision process. Given a
search space defined by $\mathcal{X}$, the current state at iteration $t$ is
given by $\mathcal{D}_{t}$, the data that we have collected thus far, while
the unlabeled data points $\mathcal{X}\setminus\mathcal{D}_{t}$ make up the
possible actions that can be taken. Unlike many rl settings, though, the size
of the action space in a typical as problem makes it challenging for common rl
training algorithms.111Not to mention one of our main goals, scaling as to
large search spaces. This is because many of these algorithms rely on
thoroughly exploring the action space to learn about the value of each
specific action in a given state, and as $n=|\mathcal{X}|$ grows larger, this
task becomes increasingly more daunting.

Noting that we have access to ens, an expert policy with demonstrated superior
performance throughout previous works, we forgo learning to search from
scratch and seek to instead rely on ens for guidance. This proves to be more
feasible, as we can leverage imitation learning techniques in which we collect
a data set of state and expert’s action pairs $S=\\{s,\text{{{ens}}}(s)\\}$
and train a neural network to learn this mapping. Here, we wish our neural
network to output the same decision generated by ens (i.e., which unlabeled
point to query) given the current state (the observed data) of a search
problem. This is done by treating the goal of imitating the expert policy as a
classification problem, where a data point is characterized by a given state
$s$ of a search, and the corresponding label is the expert’s decision
$\text{{{ens}}}(s)$. A neural network classifier is then trained to correctly
classify $\text{{{ens}}}(s)$ as the desirable label among all possible
actions, by minimizing the corresponding cross-entropy loss.

Algorithm 1 DAgger for imitation learning

1:inputs number of training iterations $N$, expert policy $\pi_{*}$, problem
generator $G$

2:initialize $S\leftarrow\emptyset$

3:initialize $\hat{\pi}_{0}$ randomly

4:for $i=1$ to $N$ do

5: sample as problems $\mathcal{X}\sim G$

6: roll out $\hat{\pi}_{i-1}$ on $\mathcal{X}$ to obtain states $\\{s\\}$

7: assemble $S_{i}=\big{\\{}\big{(}s,\pi_{*}(s)\big{)}\big{\\}}$

8: aggregate $S\leftarrow S\cup S_{i}$

9: train $\hat{\pi}_{i}$ on $S$ until convergence

10:end for

11:returns best $\hat{\pi}_{i}$ on validation

The training data $S$ for imitation learning can be assembled in several ways.
For example, we could run ens on training problems and record the states
encountered and the decisions computed. However, this leaves the possibility
that as the trained network is deployed, it will arrive at a state very
different from those seen during training, and thus output unreliable
decisions. We use DAgger [43], a well-established imitation learning
technique, to address this problem. DAgger is a meta-learning algorithm that
iteratively rolls out the policy currently being trained (i.e., it uses the
current policy to make decisions), collects the expert’s actions on the
encountered states, and appends this newly collected guidance to the training
set to improve the policy being trained. This iterative procedure allows the
expert policy to be queried more strategically, targeting states the current
policy network is likely to be in. Alg. 1 summarizes this procedure.

###  3.2 Constructing search problems for training

To realize DAgger, we require access to a search problem “generator” that
provides as problems in which we are free to roll out the network being
trained and observe its performance. One may consider directly using one’s own
real-world use case to train the policy network; however, running DAgger on
real-life as problems might prove infeasible. This is because DAgger is an
iterative training loop that requires many training episodes to be played so
that the collected training data $S$ could cover a wide range of behaviors of
the expert to be imitated. We cannot afford to dedicate many real search
campaigns to this task, especially under our assumption of expensive labels.
Instead, we turn to synthetic problems generated in a way that is sufficiently
diverse to present a wide range of scenarios under which we may observe ens’s
behavior. In addition to constructing these problems and running a policy
currently being trained on them at little computational cost, we can limit the
size of the problems so that ens can be queried efficiently.

When called, our data-generating process constructs a randomly generated set
$\mathcal{X}$. We sample from a Gaussian process (gp) [41] at the locations in
$\mathcal{X}$ to obtain a real-valued label for each $x\in\mathcal{X}$, which
is then converted to a binary label by thresholding at a chosen quantile. The
generated search space and labels are returned as a training problem. Although
this procedure is quite simple and, in using a gp, assumes a certain level of
smoothness in the labels, we observe that the generated problems offer
reasonable variety of structures with “clumps” of targets of variable number
and size. This variety successfully facilitates imitation learning, as later
demonstrated by the empirical performance of our trained policy on real-world
problems. We include more details in Appx. A.

###  3.3 Feature engineering & implementation

The effectiveness of any training procedure in rl crucially depends on the
quality of the representation of a given state during a search. To
characterize a state in a way that aids learning, we use the following
features to represent each unlabeled data point
$x\in\mathcal{X}\setminus\mathcal{D}$ remaining in a search:

  * •

the posterior probability that the data point has a positive label
$\Pr(y=1\mid x,\mathcal{D})$,

  * •

the remaining budget $\ell=T-t$,

  * •

the sum of posterior probabilities of the $(\ell-1)$ unlabeled nearest
neighbors of $x$:

$\sum_{x^{\prime}\in\text{{{nn}}}(x,\,\ell-1)}\Pr(y^{\prime}=1\mid
x^{\prime},\mathcal{D}),$ (3)

where $\text{{{nn}}}(x,k)$ denotes the set of $k$ unlabeled nearest neighbors
of a given $x\in\mathcal{X}$, and

  * •

the sum of similarities between $x$ and its $(\ell-1)$ unlabeled nearest
neighbors:

$\sum_{x^{\prime}\in\text{{{nn}}}(x,\,\ell-1)}s(x,x^{\prime}),$ (4)

where $s(x,x^{\prime})\in[0,1]$ denotes the similarity between two given
points $x,x^{\prime}\in\mathcal{X}$.

Which similarity function $s$ to use to compute the nearest neighbors of each
point and the corresponding similarity values depends on the application. We
use the radial basis function kernel
$s(x,x^{\prime})=\exp\big{(}-\tfrac{\|x-x^{\prime}\|^{2}}{2\,\lambda^{2}}\big{)}$
during training and upon deployment for appropriate tasks in Sect. 5, but this
can be replaced with another function more applicable to a given domain.

We specifically design the third and fourth features to relate the value of an
unlabeled point we may query to the characteristics of its nearest neighbors.
Intuitively, a point whose neighbors are likely targets is a promising
candidate, as it indicates a region that could yield many hits. On the other
hand, points that are close to its neighbors (e.g., cluster centers) could
also prove beneficial to query, as they help the policy explore the space
effectively. Further, the number of nearest neighbors to include in these
computations is set to match the length of the horizon, allowing these
features to dynamically adjust to our remaining budget. Overall, the four
features make up the feature vector of each candidate point, and concatenating
all feature vectors gives the state representation of a given search
iteration. We also note that our state representation is task-agnostic and
applicable across as problems of varying structures and sizes, which is
crucial for training our policy on the different problems we generate, as well
as for when we deploy our trained policy on unseen problems.

Here, finding the nearest neighbors of each point may prove challenging under
large spaces. We leverage the state-of-the-art similarity search library Faiss
[21] to perform efficient approximate nearest neighbor search when exact
search is prohibitive.222We set the number of clusters into which the search
space is split when performing approximate neighbor search at $\lfloor
4\sqrt{n}\rfloor$, following
https://github.com/facebookresearch/faiss/issues/112. Faiss allows us to
significantly accelerate this step, completing, for example, the neighbor
search for our largest problem in Sect. 5 of 6.7 million points in roughly one
hour. Once this neighbor search is done before the actual search campaign, the
time complexity of constructing the features above at each search iteration is
$O(n)$.

We run Alg. 1 for $N=50$ iterations, each consisting of 3 training problems.
At the end of each iteration, we train a small policy network with 5 fully
connected hidden layers (with 8, 16, 32, 16, and 8 neurons, respectively) and
relu activation functions using minibatch gradient descent with Adam optimizer
[23]. The trained policy is then evaluated on a fixed set of 3 unseen
validation problems. We set the labeling budget $T=100$ across all generated
problems.

###  3.4 Demonstration of learned search strategy

Figure 1:  Demonstration of our trained policy’s budget-awareness with a toy
example. Left panel: the probability that a point is a target. Remaining
panels: computed logits and the point selected to be the next query under
different labeling budgets. Our policy appropriately balances between
exploitation under a small labeling budget and strategic exploration if the
budget is large.

Before discussing our experiment results, we briefly demonstrate the learned
behavior of our trained policy network using an illustrative toy example
visualized in Fig. 1. The search space $\mathcal{X}$ consists of uniformly
sampled points as well as 3 distinct clusters of different sizes. The left
panel shows $\Pr(y=1\mid x,\mathcal{D})$, the probability that each point is a
target, specifically set so that:

  * •

for each of the 100 uniformly sampled points, $\Pr(y=1\mid
x,\mathcal{D})=0.1$,

  * •

for each point in the small cluster of size 10 at the top, $\Pr(y=1\mid
x,\mathcal{D})=0.9$,

  * •

for each point in the medium cluster of size 30 on the right, $\Pr(y=1\mid
x,\mathcal{D})=0.3$, and

  * •

for each point in the large cluster of size 100 on the bottom left,
$\Pr(y=1\mid x,\mathcal{D})=0.1$.

This problem structure presents an interesting choice between exploiting the
small cluster of very likely targets and exploring larger clusters that
contain less likely targets. A good policy should select exploitation if the
remaining budget is small and choose to further explore otherwise. The
remaining panels in Fig. 1, which visualize the logits computed by our trained
policy under different remaining budgets $\ell\in\\{10,33,100\\}$, show that
this is exactly the case: the policy targets the cluster of likely targets
when the budget is small, and moves to larger clusters as the budget
increases. Further, when exploring, the policy appropriately favors cluster
centers, which offer more information about the space. This balance between
exploitation and strategic exploration our trained policy exhibits indicates
that the policy has learned a meaningful search behavior from ens, which
translates into good empirical performance, as later shown in Sect. 5.

##  4 Related work

Active search. We continue the line of research on active search (as) [13, 9],
which previous works have also referred to as active learning and adaptive
sampling for discovery [53, 54, 57] or active covering [17]. Garnett et al.
[13] studied the Bayesian optimal policy, and Jiang et al. [18] proposed ens
as a budget-aware approximation demonstrating impressive empirical success.
ens has since then been adopted under various settings, including batch [19],
cost-aware [53, 54, 20], multifidelity [40], and diversity-aware as [39]. We
propose to amortize policy computation of ens using imitation learning,
scaling nonmyopic search to large data sets.

Amortization via neural networks. Using neural networks to amortize expensive
computations has seen increasing interests from the machine learning
community. Of note is the work of Foster et al. [12], who tackled amortizing
maximizing expected information gain [32, 33] for Bayesian experimental design
(bed) [27] with a design network trained on a specialized loss function, and
was a major inspiration for our work. Subsequent works [6, 48] have studied
bed under other settings such as those with discrete action spaces. Liu et al.
[30] tackled amortizing Gaussian process (gp) inference by training a
transformer-based network as a regression model to predict optimal
hyperparameters of gps with stationary kernels; Bitzer et al. [5] later
extended the approach to more general kernel structures. Andrychowicz et al.
[2], on the other hand, learned an optimization policy with a recurrent neural
network that predicts the next update to the parameters to be optimized based
on query history; the policy network was shown to outperform many generic
gradient-based optimizers. Also related is the work of Konyushkova et al.
[26], where a regressor was trained to predict the value of querying a given
unlabeled data point for active learning.

Reinforcement learning. Reinforcement learning (rl) has proven a useful tool
for learning effective strategies for planning tasks similar to as. Examples
include active learning policies for named entity recognition [11, 28], neural
machine translation [29], and active learning on graphs [15]. Igoe et al. [16]
were interested in path planning for drones, framed as a specialized as
setting with linear models and many agents. Sarkar et al. [46] and Sarkar et
al. [47] studied the problem of visual as, a realization of as on images for
geospatial exploration. Overall, the methodologies in these works rely on
being able to generate many training episodes to learn an effective rl policy
from scratch, which cannot be realized in our setting. Having access to ens,
we instead leverage imitation learning to learn to search from this expert on
synthetically generated search problems. When deployed, our trained policy can
be applied to a diverse set of use cases, as demonstrated in the next section.

Table 1:  Average number of targets found and standard errors by each search
policy across 10 repeats of all instances of a given task. Settings that are
computationally prohibitive for a given policy are left blank. The best policy
in each setting is highlighted bold; policies not significantly worse than the
best (according to a two-sided paired $t$-test with a significance level of
$\alpha=0.05$) are in _blue italics_.

|  Disease hotspots |  fashion-mnist |  Bulk metal glass |  Drug discovery (small) |  GuacaMol |  Drug discovery (large)
---|---|---|---|---|---|---
etc$(m=10)$ |  50.50 (4.44) |  47.79 (3.89) |  81.70 (3.66) |  75.46 (2.64) |  10.37 (1.80) |  _33.97_ _(4.14)_
etc$(m=20)$ |  52.25 (3.47) |  42.97 (3.47) |  73.40 (3.67) |  68.46 (2.37) |  9.42 (1.61) |  30.42 (3.70)
etc$(m=30)$ |  48.75 (3.12) |  38.16 (3.03) |  65.70 (3.17) |  60.76 (2.11) |  8.56 (1.45) |  26.59 (3.24)
ucb$(\beta=0.1)$ |  48.20 (4.21) |  51.66 (4.24) |  _88.80_ _(4.00)_ |  80.74 (2.86) |  11.21 (1.93) |  _36.94_ _(4.49)_
ucb$(\beta=0.3)$ |  48.15 (4.20) |  51.66 (4.24) |  _88.80_ _(4.00)_ |  80.74 (2.86) |  11.21 (1.93) |  _36.94_ _(4.49)_
ucb$(\beta=1)$ |  45.48 (3.52) |  51.10 (4.18) |  81.80 (3.15) |  75.14 (2.59) |  11.21 (1.93) |  _36.94_ _(4.49)_
one-step |  48.20 (4.21) |  51.66 (4.24) |  _88.80_ _(4.00)_ |  80.74 (2.86) |  11.21 (1.93) |  _36.94_ _(4.49)_
ids |  48.83 (4.08) |  51.66 (4.24) |  — |  — |  — |  —
ens |  57.67 (3.22) |  88.15 (1.52) |  91.10 (3.48) |  86.82 (2.41) |  — |  —
ans (ours) |  _57.25_ _(3.58)_ |  85.72 (1.69) |  _89.90_ _(4.20)_ |  _85.29_ _(2.47)_ |  15.51 (2.36) |  39.83 (3.76)

##  5 Experiments

We tested our policy network, which we call _amortized nonmyopic search_ , or
ans, and a number of baselines on search problems spanning a wide range of
applications. For each problem included, we run each policy 10 times from the
same set of initial data $\mathcal{D}_{0}$ that contains one target and one
non-target, both randomly sampled. Each of these runs has a labeling budget of
$T=100$.

Baselines. We compare our method against the expert policy ens which ours was
trained to mimic, as well as the one-step policy that greedily queries the
most likely target discussed in Sect. 2.1. We also implement a number of
baseline policies from the literature. The first is a family of upper
confidence bound (ucb) policies [3, 8] that rank candidates by the following
score: $p+\beta\sqrt{p(1-p)}$, where $p=\Pr(y=1\mid x,\mathcal{D})$ is the
posterior probability that a given candidate is a target, and $\beta$ is the
tradeoff parameter balancing exploitation (favoring large $p$) and exploration
(favoring large uncertainty in the label, as measured by $\sqrt{p(1-p)}$).
Here, we have $\beta$ take on values from $\\{0.1,0.3,1\\}$. Another baseline
is from Jiang and Rostamizadeh [17], who proposed a simple explore-then-commit
(etc) scheme that uniformly samples the space for $m$ iterations and then
switches to the greedy sampling strategy of one-step for the remaining of the
search. We run this etc policy with $m\in\\{10,20,30\\}$. Finally, we include
the policy by Xu et al. [57], which uses the information-directed sampling
(ids) heuristic [44, 45] that scores each candidate query $x$ by the ratio
between (1) information about the labels of the current $\ell$ most likely
targets ($\ell=T-t$ is the length of the remaining horizon), gained by
querying $x$ and (2) the expected instant regret from querying $x$. We note
that Xu et al. [57] proposed this policy under specialized as settings that
allow information gain to be efficiently computed; they also only considered
problems with fewer than 1000 points. For our experiments, we can only apply
ids to relatively small spaces where the policy is computationally feasible.

Figure 2:  The time taken per iteration by different active search policies in
the small- and medium-scale experiments. Left: average number of seconds per
iteration with respect to the size of the search space. Right: average number
of targets found and standard errors vs. time per iteration.

Search problems. We now discuss the search problems making up our experiments.
The first was posed by Andrade-Pacheco et al. [1], who sought to identify
disease hotspots within a region of interest. The provided data sets
correspond to 4 distinct as problems of finding locations with a high
prevalence of schistosomiasis in Côte d’Ivoire and Malawi and of lymphatic
filariasis in Haiti and the Philippines. Each problem consists of 1500 points,
with targets accounting for 10%–34% of the space. For our second task,
following Nguyen and Garnett [39], we simulate product recommendation problems
using the fashion-mnist data [56], which contains 70 000 images classified
into 10 classes of clothing articles. We first randomly select 3 out of the 10
classes as products a user is interested in (i.e., our search targets). We
further sub-sample these 3 classes uniformly at random to increase the
difficulty of the problem; the resulting prevalence rate of the targets is
roughly 6%. We repeat this process 10 times to generate 10 search problems
with this data.

Borrowing from previous works [18, 19], we use a data set from the materials
science literature [22, 52] containing 106 810 alloys, of which 4275 can form
bulk metallic glasses with high toughness and wear resistance and are our
search targets. Another application comes from drug discovery, where we aim to
identify “active compounds”, chemical compounds that bind with a targeted
protein. Garnett et al. [14] assembled a suite of such drug discovery
problems, each of which consists of the active compounds for a specific
protein from the bindingdb database [31], and 100 000 molecules sampled from
the zinc database [49] that act as the negative pool. Our experiments include
the first 10 problems where on average the active compounds make up 0.5% of
the search space.

Finally, to demonstrate the ability to perform search on large spaces achieved
by our method ans, we consider two large-scale, challenging drug discovery
tasks. The first employs the GuacaMol database of over 1.5 million drug-like
molecules that were specifically curated for drug discovery benchmarking tasks
involving machine learning [7]. In addition to these molecules, the database
offers a family of objective functions to measure the molecules’ quality using
a variety of criteria. We use each objective function provided to define a
search problem as follows. We first randomly sample a set of 1000 molecules
which we fully label using the objective functions. We then define the search
targets as those of the remaining unlabeled molecules whose scores exceed the
99-th percentile of the labeled set; in other words, the goal of our search is
the top 1% molecules. In total, we assemble 9 such as problems with GuacaMol.
For our second task, we follow the procedure by Garnett et al. [14] described
above with the bindingdb and zinc databases, this time expanding the negative
pool to all drug-like molecules in zinc [51]. This results in a search space
of 6.7 million candidates, of which 0.03% are the active compounds we aim to
search for.

Discussions. Tab. 1 reports the performance of the search policies – measured
in the number of targets discovered – in each of these tasks, where settings
that are computationally prohibitive for a given policy are left blank. From
these results, we observe a clear trend: the state-of-the-art policy ens
consistently achieves the best performance under all settings that it could
feasibly run, while our trained policy network ans closely follows ens,
sometimes outperforming the other baselines by a large margin. Our method also
yields the best result in large-scale problems, demonstrating its usefulness
in large search spaces. Among the baselines, we note the difficulty in setting
the number of exploration rounds $m$ for etc, since no value of $m$ performs
the best across all settings. Results from ucb policies, on the other hand,
indicate that prioritizing exploitation (setting the parameter $\beta$ to a
small value) is beneficial, a trend also observed in previous work [18].

To illustrate the tradeoff between performance and speed achieved by ans, the
left panel of Fig. 2 shows the average time taken by ans, ens, and myopic
baselines per iteration as a function of the size of the data, while the right
panel shows the number of discoveries by each policy vs. the same average time
per iteration. These plots do not include the results from the large-scale
problems so that the comparison with ens is fair, or ids which is slower than
ens but does not perform as well. We see that ans finds almost as many targets
as ens but is much more computationally lightweight. We thus have established
a new point on the Pareto frontier of the performance vs. speed tradeoff with
our search policy. In the 6.7 million-point drug discovery problems, ans on
average takes 36.94 $\pm$ 0.15 minutes per iteration, which we deem entirely
acceptable given the boost in performance compared to faster but myopic
baselines, the fact that the time cost of labeling is typically much higher,
and ens, in comparison, is estimated to take roughly 10 hours per iteration on
the same scale.

Figure 3:  Locations in Côte d’Ivoire selected by the one-step policy and by
ours in an illustrative run with the disease hotspot data, where our policy
discovers a larger target cluster.  Figure 4:  The average difference in
cumulative reward and standard errors between our policy and one-step. Our
policy spends its initial budget exploring the space and finds fewer targets
in the beginning but smoothly switches to more exploitative queries and
outperforms one-step at the end.

As a demonstration of our policy’s strategic explorative behavior learned from
ens, Fig. 3 shows the result of an illustrative run by the one-step policy vs.
ans from the problem of finding schistosomiasis hotspots in Côte d’Ivoire [1].
We note that the queries made by one-step are localized within the center
region containing the target in the initial data $\mathcal{D}_{0}$, while ans
is able to discover a larger cluster of targets to the south. Further, Fig. 4
visualizes the cumulative difference in utility between ans and one-step
across all experimental settings. Here, ans initially finds fewer targets than
one-step, as the former tends to dedicate its queries to exploration of the
space when the remaining budget is large; however, as the search progresses,
ans smoothly transitions to more exploitative queries and ultimately
outperforms the greedy policy. The same pattern of behavior has been observed
from ens in previous works [18, 40, 39].

We defer further analyses to Appx. C, which includes an ablation study showing
the importance of each feature we engineer for our policy network in Sect.
3.3, as well as the benefit of the DAgger loop compared to (1) imitation
learning without iteratively generating more training data and (2) learning to
search from scratch without ens. We also examine the stability between
different training runs giving different policy networks, and observe that
these policies yield comparable results. Finally, we investigate the
effectiveness of various schemes to further improve search performance under
settings where repeated searches are conducted within the same space.

##  6 Conclusion

We propose an imitation learning-based method to scale nonmyopic active search
to large search spaces, enabling real-time decision-making and efficient
exploration of massive databases common in product recommendation and drug
discovery beyond myopic/greedy strategies. Extensive experiments showcase the
usefulness of our policy, which mimics the state-of-the-art policy ens while
being significantly cheaper to run. Future directions include deriving a more
effective reinforcement learning strategy to train our policy network,
potentially outperforming ens, as well as extending to other active search
settings such as batch [18] and diversity-aware search [39, 38].

## References

  * Andrade-Pacheco et al. [2020] Ricardo Andrade-Pacheco, Francois Rerolle, Jean Lemoine, Leda Hernandez, Aboulaye Meïté, Lazarus Juziwelo, Aurélien F Bibaut, Mark J van der Laan, Benjamin F Arnold, and Hugh JW Sturrock.  Finding hotspots: development of an adaptive spatial sampling approach.  _Scientific Reports_ , 10, 2020.
  * Andrychowicz et al. [2016] Marcin Andrychowicz, Misha Denil, Sergio Gómez, Matthew W Hoffman, David Pfau, Tom Schaul, Brendan Shillingford, and Nando De Freitas.  Learning to learn by gradient descent by gradient descent.  In _Advances in Neural Information Processing Systems 29_ , 2016.
  * Auer [2000] Peter Auer.  Using Upper Confidence Bounds for Online Learning .  In _Proceedings of the 41st Annual Symposium on Foundations of Computer Science_ , pages 270–279. IEEE, 2000.
  * Bellman [1957] Richard Bellman.  _Dynamic Programming_.  Princeton University Press, 1957.
  * Bitzer et al. [2023] Matthias Bitzer, Mona Meister, and Christoph Zimmer.  Amortized Inference for Gaussian Process Hyperparameters of Structured Kernels.  In _Uncertainty in Artificial Intelligence_ , pages 184–194, 2023.
  * Blau et al. [2022] Tom Blau, Edwin V Bonilla, Iadine Chades, and Amir Dezfouli.  Optimizing Sequential Experimental Design with Deep Reinforcement Learning.  In _Proceedings of the 39th International Conference on Machine Learning_ , pages 2107–2128, 2022.
  * Brown et al. [2019] Nathan Brown, Marco Fiscato, Marwin HS Segler, and Alain C Vaucher.  GuacaMol: Benchmarking Models for de Novo Molecular Design.  _Journal of Chemical Information and Modeling_ , 59(3):1096–1108, 2019.
  * Carpentier et al. [2011] Alexandra Carpentier, Alessandro Lazaric, Mohammad Ghavamzadeh, Rémi Munos, and Peter Auer.  Upper-Confidence-Bound Algorithms for Active Learning in Multi-armed Bandits.  In _International Conference on Algorithmic Learning Theory_ , pages 189–203, 2011.
  * Coleman et al. [2022] Cody Coleman, Edward Chou, Julian Katz-Samuels, Sean Culatana, Peter Bailis, Alexander C Berg, Robert Nowak, Roshan Sumbaly, Matei Zaharia, and I Zeki Yalniz.  Similarity Search for Efficient Active Learning and Search of Rare Concepts.  In _Proceedings of the AAAI Conference on Artificial Intelligence_ , volume 36, 2022.
  * Dulac-Arnold et al. [2015] Gabriel Dulac-Arnold, Richard Evans, Hado van Hasselt, Peter Sunehag, Timothy Lillicrap, Jonathan Hunt, Timothy Mann, Theophane Weber, Thomas Degris, and Ben Coppin.  Deep Reinforcement Learning in Large Discrete Action Spaces.  _arXiv preprint_ , 2015.  arXiv:1512.07679 [cs.AI].
  * Fang et al. [2017] Meng Fang, Yuan Li, and Trevor Cohn.  Learning how to Active Learn: A Deep Reinforcement Learning Approach.  In _Proceedings of the 2017 Conference on Empirical Methods in Natural Language Processing_ , pages 595–605, 2017.
  * Foster et al. [2021] Adam Foster, Desi R Ivanova, Ilyas Malik, and Tom Rainforth.  Deep Adaptive Design: Amortizing Sequential Bayesian Experimental Design.  In _Proceedings of the 38th International Conference on Machine Learning_ , pages 3384–3395, 2021.
  * Garnett et al. [2012] Roman Garnett, Yamuna Krishnamurthy, Xuehan Xiong, Jeff Schneider, and Richard Mann.  Bayesian Optimal Active Search and Surveying.  In _Proceedings of the 29th International Conference on Machine Learning_ , 2012.
  * Garnett et al. [2015] Roman Garnett, Thomas Gärtner, Martin Vogt, and Jürgen Bajorath.  Introducing the ‘active search’ method for iterative virtual screening.  _Journal of Computer-Aided Molecular Design_ , 29:305–314, 2015.
  * Hu et al. [2020] Shengding Hu, Zheng Xiong, Meng Qu, Xingdi Yuan, Marc-Alexandre Côté, Zhiyuan Liu, and Jian Tang.  Graph Policy Network for Transferable Active Learning on Graphs.  In _Advances in Neural Information Processing Systems 33_ , pages 10174–10185, 2020.
  * Igoe et al. [2021] Conor Igoe, Ramina Ghods, and Jeff Schneider.  Multi-Agent Active Search: A Reinforcement Learning Approach.  _IEEE Robotics and Automation Letters_ , 7(2):754–761, 2021.
  * Jiang and Rostamizadeh [2021] Heinrich Jiang and Afshin Rostamizadeh.  Active Covering.  In _Proceedings of the 38th International Conference on Machine Learning_ , pages 5013–5022, 2021.
  * Jiang et al. [2017] Shali Jiang, Gustavo Malkomes, Geoff Converse, Alyssa Shofner, Benjamin Moseley, and Roman Garnett.  Efficient Nonmyopic Active Search.  In _Proceedings of the 34th International Conference on Machine Learning_ , pages 1714–1723, 2017.
  * Jiang et al. [2018] Shali Jiang, Gustavo Malkomes, Matthew Abbott, Benjamin Moseley, and Roman Garnett.  Efficient nonmyopic batch active search.  In _Advances in Neural Information Processing Systems 31_ , pages 1099–1109, 2018.
  * Jiang et al. [2019] Shali Jiang, Roman Garnett, and Benjamin Moseley.  Cost effective active search.  In _Advances in Neural Information Processing Systems 32_ , pages 4880–4889, 2019.
  * Johnson et al. [2019] Jeff Johnson, Matthijs Douze, and Hervé Jégou.  Billion-Scale Similarity Search with GPUs.  _IEEE Transactions on Big Data_ , 7(3):535–547, 2019.
  * Kawazoe et al. [1997] Y Kawazoe, J-Z Yu, A-P Tsai, and T Masumoto, editors.  _Nonequilibrium Phase Diagrams of Ternary Amorphous Alloys_.  Condensed Matters. Springer-Verlag, 1997.
  * Kingma and Ba [2015] Diederik P Kingma and Jimmy Ba.  Adam: A Method for Stochastic Optimization.  In _Proceddings of the 3rd International Conference for Learning Representations_ , 2015.
  * Kingma and Welling [2014] Diederik P Kingma and Max Welling.  Auto-Encoding Variational Bayes.  In _Proceddings of the 2nd International Conference for Learning Representations_ , 2014.
  * Kingma et al. [2014] Diederik P Kingma, Shakir Mohamed, Danilo Jimenez Rezende, and Max Welling.  Semi-supervised Learning with Deep Generative Models.  _Advances in Neural Information Processing Systems 27_ , 2014.
  * Konyushkova et al. [2017] Ksenia Konyushkova, Raphael Sznitman, and Pascal Fua.  Learning Active Learning from Data.  In _Advances in Neural Information Processing Systems 30_ , 2017.
  * Lindley [1956] Dennis V Lindley.  On a Measure of the Information Provided by an Experiment .  _The Annals of Mathematical Statistics_ , 27(4):986–1005, 1956.
  * Liu et al. [2018a] Ming Liu, Wray Buntine, and Gholamreza Haffari.  Learning How to Actively Learn: A Deep Imitation Learning Approach.  In _Proceedings of the 56th Annual Meeting of the Association for Computational Linguistics_ , pages 1874–1883, 2018a.
  * Liu et al. [2018b] Ming Liu, Wray Buntine, and Gholamreza Haffari.  Learning to Actively Learn Neural Machine Translation.  In _Proceedings of the 22nd Conference on Computational Natural Language Learning_ , pages 334–344, 2018b.
  * Liu et al. [2020] Sulin Liu, Xingyuan Sun, Peter J Ramadge, and Ryan P Adams.  Task-Agnostic Amortized Inference of Gaussian Process Hyperparameters.  In _Advances in Neural Information Processing Systems 33_ , pages 21440–21452, 2020.
  * Liu et al. [2007] Tiqing Liu, Yuhmei Lin, Xin Wen, Robert N Jorissen, and Michael K Gilson.  BindingDB: A web-accessible database of experimentally determined protein–ligand binding affinities.  _Nucleic Acids Research_ , 35:D198–D201, 2007.
  * MacKay [1992a] David MacKay.  The Evidence Framework Applied to Classification Networks.  _Neural Computation_ , 1992a.
  * MacKay [1992b] David MacKay.  Information-Based Objective Functions for Active Data Selection.  _Neural Computation_ , 1992b.
  * Maus et al. [2022] Natalie Maus, Haydn Jones, Juston Moore, Matt J Kusner, John Bradshaw, and Jacob Gardner.  Local Latent Space Bayesian Optimization over Structured Inputs.  _Advances in Neural Information Processing Systems 35_ , pages 34505–34518, 2022.
  * McInnes et al. [2018] Leland McInnes, John Healy, and James Melville.  UMAP: Uniform Manifold Approximation and Projection for Dimension Reduction.  _arXiv preprint_ , 2018.  arXiv:1802.03426 [stat.ML].
  * Mnih et al. [2015] Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness, Marc G Bellemare, Alex Graves, Martin Riedmiller, Andreas K Fidjeland, Georg Ostrovski, Stig Petersen, Charles Beattie, Amir Sadik, Ioannis Antonoglou, Helen King, Dharshan Kumaran, Daan Wierstra, Shane Legg, and Demis Hassabis.  Human-level control through deep reinforcement learning.  _Nature_ , 518(7540):529–533, 2015.
  * Monadjemi et al. [2022] Shayan Monadjemi, Sunwoo Ha, Quan Nguyen, Henry Chai, Roman Garnett, and Alvitta Ottley.  Guided Data Discovery in Interactive Visualizations via Active Search.  In _2022 IEEE Visualization and Visual Analytics (VIS)_ , pages 70–74. IEEE, 2022.
  * Nguyen and Dieng [2024] Quan Nguyen and Adji Bousso Dieng.  Quality-Weighted Vendi Scores And Their Application To Diverse Experimental Design.  In _Proceedings of the 41st International Conference on Machine Learning_ , 2024.  To appear.
  * Nguyen and Garnett [2023] Quan Nguyen and Roman Garnett.  Nonmyopic Multiclass Active Search with Diminishing Returns for Diverse Discovery .  In _Proceedings of the 26th International Conference on Artificial Intelligence and Statistics_ , 2023.
  * Nguyen et al. [2021] Quan Nguyen, Arghavan Modiri, and Roman Garnett.  Nonmyopic Multifidelity Acitve Search.  In _Proceedings of the 38th International Conference on Machine Learning_ , pages 8109–8118, 2021.
  * Rasmussen and Williams [2006] Carl Edward Rasmussen and Christopher K. I. Williams.  _Gaussian Processes for Machine Learning_.  The MIT Press, 2006.
  * Rogers and Hahn [2010] David Rogers and Mathew Hahn.  Extended-Connectivity Fingerprints.  _Journal of Chemical Information and Modeling_ , 50(5):742–754, 2010.
  * Ross et al. [2011] Stéphane Ross, Geoffrey Gordon, and Drew Bagnell.  A Reduction of Imitation Learning and Structured Prediction to No-Regret Online Learning .  In _Proceedings of the 14th International Conference on Artificial Intelligence and Statistics_ , pages 627–635, 2011.
  * Russo and Van Roy [2014] Daniel Russo and Benjamin Van Roy.  Learning to Optimize via Information-Directed Sampling.  In _Advances in Neural Information Processing Systems 27_ , 2014.
  * Russo and Van Roy [2018] Daniel Russo and Benjamin Van Roy.  Learning to Optimize via Information-Directed Sampling.  _Operations Research_ , 66:230–252, 2018.
  * Sarkar et al. [2023] Anindya Sarkar, Nathan Jacobs, and Yevgeniy Vorobeychik.  A Partially Supervised Reinforcement Learning Framework for Visual Active Search.  In _Advances in Neural Information Processing Systems 36_ , pages 12245–12270, 2023.
  * Sarkar et al. [2024] Anindya Sarkar, Michael Lanier, Scott Alfeld, Jiarui Feng, Roman Garnett, Nathan Jacobs, and Yevgeniy Vorobeychik.  A Visual Active Search Framework for Geospatial Exploration.  In _Proceedings of the IEEE/CVF Winter Conference on Applications of Computer Vision_ , pages 8316–8325, 2024.
  * Shen and Huan [2023] Wanggang Shen and Xun Huan.  Bayesian sequential optimal experimental design for nonlinear models using policy gradient reinforcement learning.  _Computer Methods in Applied Mechanics and Engineering_ , 2023.
  * Sterling and Irwin [2015] Teague Sterling and John J Irwin.  ZINC 15–Ligand Discovery for Everyone.  _Journal of Chemical Information and Modeling_ , 55(11):2324–2337, 2015.
  * Sutton et al. [1999] Richard S Sutton, David McAllester, Satinder Singh, and Yishay Mansour.  Policy Gradient Methods for Reinforcement Learning with Function Approximation.  _Advances in Neural Information Processing Systems 12_ , 1999.
  * Tingle et al. [2023] Benjamin I Tingle, Khanh G Tang, Mar Castanon, John J Gutierrez, Munkhzul Khurelbaatar, Chinzorig Dandarchuluun, Yurii S Moroz, and John J Irwin.  ZINC-22—A Free Multi-Billion-Scale Database of Tangible Compounds for Ligand Discovery.  _Journal of Chemical Information and Modeling_ , 63(4):1166–1176, 2023.
  * Ward et al. [2016] Logan Ward, Ankit Agrawal, Alok Choudhary, and Christopher Wolverton.  A general-purpose machine learning framework for predicting properties of inorganic materials.  _npj Computational Materials_ , 2(1):1–7, 2016.
  * Warmuth et al. [2002] Manfred K Warmuth, Gunnar Rätsch, Michael Mathieson, Jun Liao, and Christian Lemmen.  Active learning in the drug discovery process.  In _Advances in Neural Information Processing Systems 15_ , pages 1449–1456, 2002.
  * Warmuth et al. [2003] Manfred K Warmuth, Jun Liao, Gunnar Rätsch, Michael Mathieson, Santosh Putta, and Christian Lemmen.  Active Learning with Support Vector Machines in the Drug Discovery Process.  _Journal of Chemical Information and Computer Sciences_ , 43(2):667–673, 2003.
  * Willett et al. [1998] Peter Willett, John M Barnard, and Geoffrey M Downs.  Chemical Similarity Searching.  _Journal of Chemical Information and Computer Sciences_ , 38(6):983–996, 1998.
  * Xiao et al. [2017] Han Xiao, Kashif Rasul, and Roland Vollgraf.  Fashion-MNIST: a Novel Image Dataset for Benchmarking Machine Learning Algorithms.  _arXiv preprint_ , 2017.  arXiv:1708.07747 [cs.LG].
  * Xu et al. [2022] Ziping Xu, Eunjae Shim, Ambuj Tewari, and Paul Zimmerman.  Adaptive Sampling for Discovery.  In _Advances in Neural Information Processing Systems 35_ , pages 1114–1126, 2022.

##  Appendix A Generation of training search problems

Figure 5:  An example two-dimensional problem generated by Alg. 2, where
bright and dark points indicate targets and non-targets, respectively. The
search space includes both clusters and more widely dispersed points.

We now discuss our procedure of generating active search problems to train the
policy network in DAgger, summarized in Alg. 2, where $U\\{m,n\\}$ denotes a
discrete uniform distribution of the integers between $m$ and $n$ (inclusive),
while $U[a,b]$ refers to a continuous uniform distribution between $a$ and
$b$.

The search space $\mathcal{X}$ of each generated problem exists in a
$d$-dimensional space, where $d$ is a random integer between $2$ and $10$.
Once $d$ is determined, we sample $100d$ points uniformly from the
$d$-dimensional unit hypercube. This set of uniform points is combined with
$n_{\text{cluster}}$ clusters, where $n_{\text{cluster}}$ is a random integer
between $10$ and $10d$. To generate each cluster, we first sample another
random integer between $10$ and $10d$, denoted as $m$, to determine the size
of the cluster. We then draw $m$ points from an isotropic Gaussian
distribution with the mean vector $\boldsymbol{\mu}$ randomly sampled within
the unit hypercube and the diagonal covariance matrix $\sigma^{2}I_{d}$, where
$\sigma$ is drawn from $U[0.1,0.1d]$. Here, $\sigma$, which determines the
spread of a given cluster, is constrained to be between $0.1$ and $0.1d$
(relatively small numbers) to ensure that the points within this cluster are
indeed close to one another.

Again, the union of the uniform points and the clusters make up the entire
search space $\mathcal{X}$. We then draw a sample from a Gaussian process (gp)
at the points in $\mathcal{X}$. This gp is equipped with a zero mean function
and a radial basis function kernel whose length scale $\ell$ scales linearly
with the dimensionality of the space $d$ by a factor of $0.05$. This gp sample
yields a vector $\mathbf{f}$ of real-valued numbers. We then sample uniformly
between $0.01$ and $0.2$ for a prevalence rate $p$, which determines the
proportion of $\mathcal{X}$ corresponds to the targets. As such, we compute
the binary labels $\mathbf{y}$ by thresholding $\mathbf{f}$ at the
$100(1-p)$-th quantile of the values in $\mathbf{f}$. We keep the prevalence
rate $p$ below 20% to ensure that the targets are sufficiently rare. The tuple
$(\mathcal{X},\mathbf{y})$ is finally returned. Fig. 5 shows an example of one
such generated problem in two dimensions, showing a search space with a
considerably complex structure with multiple clusters and groups of targets.

##  Appendix B Data sets

We now describe the data sets used in our experiments in Sect. 5. These data
sets are curated from authors of respective publications respecting their
licenses, as detailed below. No identifiable information or offensive content
is included in the data.

  * •

We downloaded the disease hotspot data set from the GitHub repository provided
in Andrade-Pacheco et al. [1]. We computed the nearest neighbors of each data
point (a location within one of the four countries included) using its
coordinates (longitude and latitude).

  * •

The fashion-mnist data set is published by Xiao et al. [56]. We used umap [35]
to produce a two-dimensional embedding of the images and compute the nearest
neighbors on this embedding.

  * •

We obtained the bulk metal glass data from Jiang et al. [18], who, following
Ward et al. [52], represented each data point with various physical attributes
that were found to be informative in predicting glass-forming ability. Each
feature is subsequently scaled to range between $0$ and $1$. The nearest
neighbor search is performed on these features.

  * •

The data for the first set of drug discovery problems were also obtained from
Jiang et al. [18], where the Morgan fingerprints [42] were used as the feature
vectors and the Tanimoto coefficient [55] as the measure of similarity.

  * •

The molecules in the GuacaMol data are included in the publication of Brown et
al. [7], while those in our large drug discovery tasks were downloaded from
the zinc-22 database [51]. For each of these data sets, we used the state-of-
the-art transformer-based molecular variational autoencoder trained in Maus et
al. [34] to generate a 256-dimensional embedding of the molecules. The
approximate nearest neighbor search described in Sect. 3.3 was conducted on
this embedding.

Algorithm 2 Generate synthetic search problems

1:$d\sim U\\{2,10\\}$ $\triangleright$ sample dimensionality of search space

2:$\mathcal{X}\leftarrow\big{\\{}x_{j}:x_{j}\sim U[0,1]\big{\\}}_{j=1}^{100d}$
$\triangleright$ sample uniform points

3:$n_{\text{cluster}}\sim U\\{10,10d\\}$ $\triangleright$ sample number of
clusters

4:for $i=1$ to  $n_{\text{cluster}}$ do

5: $m\sim U\\{10,10d\\}$ $\triangleright$ sample cluster size

6: $\boldsymbol{\mu}=[\mu_{j}]_{j=1}^{d}$, where $\mu_{j}\sim U[0,1]$
$\triangleright$ sample cluster center

7: $\sigma\sim U[0.1,0.1d]$ $\triangleright$ sample spread of cluster

8:
$\mathcal{X}\leftarrow\mathcal{X}\cup\big{\\{}x_{j}:x_{j}\sim\mathcal{N}(\boldsymbol{\mu},\sigma^{2}I_{d})\big{\\}}_{j=1}^{m}$
$\triangleright$ use an isotropic Gaussian distribution

9:end for

10:$\mathbf{f}\sim\mathcal{N}(\mathcal{X};\mathbf{0},\mathbf{\Sigma})$,
where $\mathbf{\Sigma}=K(\mathcal{X},\mathcal{X})$ and
$K(\mathbf{x}_{1},\mathbf{x}_{2})=\exp\left(-\tfrac{\|\mathbf{x}_{1}-\mathbf{x}_{2}\|^{2}}{2\ell^{2}}\right)$
with length scale $\ell=0.05d$

11:$p\sim U[0.01,0.2]$ $\triangleright$ sample target prevalence

12:$\mathbf{y}=\mathbb{I}\big{[}\mathbf{f}>100(1-p)$-th quantile in
$\mathbf{f}\big{]}$ $\triangleright$ threshold to construct binary labels

13:returns $(\mathcal{X},\mathbf{y})$

##  Appendix C Further details on experiments

Experiments were performed on a small cluster built from commodity hardware
comprising approximately 200 Intel Xeon cpu cores, each with approximately 10
gb of ram. All compute amounted to roughly 15 000 cpu hours, including
preliminary experiments not discussed in the paper, training the policy
network, and all experiments for evaluation discussed in Sect. 5 and here.

Table 2:  Average number of targets found and standard errors by each ablated
policy across 100 product recommendation tasks with fashionmnist. The best
policy is highlighted bold.

without target probability |  without remaining budget |  without neighbor probability sum |  without neighbor similarity sum |  imitation learning without DAgger |  reinforce without imitation learning |  ans (ours)
---|---|---|---|---|---|---
66.93 (1.51) |  80.11 (1.80) |  73.35 (1.61) |  63.29 (3.35) |  59.66 (3.80) |  75.15 (1.52) |  85.72 (1.69)

Figure 6:  Distributions of the number of targets found across 100 product
recommendation tasks with fashionmnist by 10 policy networks trained with
different initial random seeds. The distributions are comparable, indicating
that the trained policy networks behave similarly.

Ablation study. We use the 10 product recommendation tasks from the
fashionmnist data to quantify the value of various components of our
framework. First, we trained four additional policy networks, each learning
from ens without one of the four features discussed in Sect. 3.3. We also
trained another network using imitation learning but without DAgger’s
iterative procedure: we ran the expert policy ens on $3\times 50=150$
generated search problems (the same number of problems generated to train the
policy examined in the main text), kept track of the encountered states and
selected actions, and used these data to train the new network until
convergence only once. Finally, we trained a policy network without imitation
learning using the reinforce policy gradient algorithm [50]. The performance
of these policies, along with that of our main policy ans as a reference, is
shown in Tab. 2. We see that by removing any component of our imitation
learning procedure, we incur a considerable decrease in search performance,
which demonstrates the importance of each of these components.

Training stability. We rerun our training procedure with DAgger for 10 times
using different random seeds and evaluate the trained policy networks using
the experiments with the fashionmnist data. Each row of Fig. 6 shows the
distribution of the number of targets found by each of these 10 policy
networks across the 100 search problems. We observe that the variation across
these 10 distributions is quite small, especially compared to the variation
across different search runs by the same policy network. This shows that our
training procedure is stable, resulting in policy networks that behave
similarly under different random seeds.

Refinement under repeated search. In many settings that as targets, multiple
search campaigns may be conducted within the same search space. For example,
as in our drug discovery experiments in Sect. 5, a scientist may explore a
molecular database to identify candidates with different desirable properties.
As the search for a given property concludes, the next search stays within the
same database but now targets a different property. In these situations, we
may reasonably seek to refine our search strategy throughout these episodes
using the results we observe, so that our search policy could improve using
its past experiences. We identify two approaches to such refinement:

  * •

If a neural network is used as the search policy, it can be updated by a
policy gradient algorithm such as reinforce [50] after each episode.

  * •

If a deep autoencoder (dae) is used to produce a representation of the search
candidates (on which the nearest neighbor search described in Sect. 3.3 is
conducted), the autoencoder can be updated with a semisupervised loss [25]
that accounts for the labels it iteratively uncovers throughout the search.

To investigate the effects of each of these approaches on the search
performance of our policy trained with imitation learning and examined in the
main text, we engineer another version of the fashionmnist data set [56] that
simulates a setting of repeated search. We first randomly choose 5 out of 10
classes in the data set to act as possible target sets throughout the repeated
searches. These selected classes are then sub-sampled uniformly at random so
that there are only 1000 data points per class; this yields a data set of 40
000 points in total. We then use a variational autoencoder [24] to learn a
two-dimensional representation of these 40 000 candidates.

We allow 100 search episodes within this database, where in each episode, 1 of
the chosen 5 classes is randomly selected as the target class. To implement
the second approach to search refinement, we train a variational Gaussian
process classifier on the observed data $\mathcal{D}$ and use the
corresponding evidence lower bound (elbo) to make up the supervised component
of the joint loss of the semisupervised model. While we update the search
policy using the reinforce loss at the end of each episode, an update to the
semisupervised vae is performed for every 20 iterations within one episode.

Figure 7:  Average cumulative difference in the number of targets found and
standard errors between (left) updating the policy network using reinforce or
(right) updating the autoencoder producing the representation of the
candidates vs. performing no updates.

Fig. 7 shows the value of each of the two update schemes as the cumulative
difference in the number of targets found between each scheme compared to
performing no updates (both the search policy and the representation of the
data points are kept fixed) throughout 100 search episodes across 10 repeats.
Surprisingly, attempting to further refine the search policy using reinforce
actually hurts performance, resulting in an increasing gap in reward between
the initial policy and the one continually updated. On the other hand, we see
that updating the initial unsupervised vae to account for the observed labels
yields an improvement in performance on average, but this improvement is not
consistent across the 10 repeats. Overall, we show the difficulty in further
updating our trained search policy using real experiences under repeated
searches, and hypothesize that more sophisticated reinforcement learning
procedures such as the double Q-learning algorithm [36] are needed to improve
learning, which we leave as future work.

##  Appendix D Limitations

Sect. 5 demonstrates that our policy ans, trained with imitation learning,
cannot perfectly capture the search strategy of the state-of-the-art ens by
Jiang et al. [18] and is outperformed by the policy. We view improving the
architecture of our policy network as well as the state representation to
enable more effective learning as a promising future direction. For example,
as examined in Liu et al. [30], a transformer-based network with beneficial
input permutation invariance properties can learn from data sets of different
sizes, which could also prove useful in as. This network architecture can be
combined with a better training strategy, as mentioned in Appx. C, to
potentially yield comparable performance as ens or even to outperform it.

##  Appendix E Broader impact

The development of an efficient as algorithm through imitation learning has
significant potential to positively impact various fields where computational
efficiency is paramount. By reducing the superlinear computational complexity
of the state-of-the-art policy ens, our approach enables the application of as
in significantly larger data sets. This scalability is essential for
industries and research fields that deal with vast amounts of data, including
genomics, astronomy, and environmental monitoring, among others. The ability
to achieve competitive performance at a fraction of the cost also has
substantial economic implications. Organizations can deploy high-performance
as solutions without the need for extensive computational resources, making
advanced data analysis more accessible and affordable. Since as aims to
identify as many targets as possible, the collected data set may end up
unbalanced towards the positives. It is important for the user to ensure that
maximizing the number of labeled targets accurately reflects their objective,
and that the collected data are used by downstream tasks that are not
negatively affected by this imbalance.