# VertiBayes: Learning Bayesian network parameters from vertically partitioned
data with missing values

Florian van Daalen, Lianne Ippel, Andre Dekker, and Inigo Bermejo F. van
Daalen, I. Bermejo, and A. Dekker are with the Department of Radiation
Oncology (MAASTRO) GROW School for Oncology and Reproduction Maastricht
University Medical Centre+ Maastricht the Netherlands
L. Ippel is with Statistics Netherlands Heerlen the Netherlands.

###### Abstract

Federated learning makes it possible to train a machine learning model on
decentralized data. Bayesian networks are widely used probabilistic graphical
models. While some research has been published on the federated learning of
Bayesian networks, publications on Bayesian networks in a vertically
partitioned data setting are limited, with important omissions, such as
handling missing data. We propose a novel method called VertiBayes to train
Bayesian networks (structure and parameters) on vertically partitioned data,
which can handle missing values as well as an arbitrary number of parties. For
structure learning we adapted the K2 algorithm with a privacy-preserving
scalar product protocol. For parameter learning, we use a two-step approach:
first, we learn an intermediate model using maximum likelihood, treating
missing values as a special value, then we train a model on synthetic data
generated by the intermediate model using the EM algorithm. The privacy
guarantees of VertiBayes are equivalent to those provided by the privacy
preserving scalar product protocol used. We experimentally show VertiBayes
produces models comparable to those learnt using traditional algorithms.
Finally, we propose two alternative approaches to estimate the performance of
the model using vertically partitioned data and we show in experiments that
these give accurate estimates.

###### Index Terms:

Federated Learning, Bayesian network, privacy preserving, vertically
partitioned data, missing values, parameter learning, structure learning

##  1 Introduction

†† _The views expressed in this paper are those of the authors and do not
necessarily reflect the policy of Statistics Netherlands._

Federated learning is a field that recently rose to prominence due to the
increased focus on data-hungry techniques, privacy concerns and protection of
the data[1][2]. Using federated learning, it is possible to train a machine
learning model without needing to collect the data centrally [1]. Since it
rose to prominence, various techniques for training a model on centrally
collected data have been adapted to be used on data that is either
horizontally or vertically partitioned[2]. Data is said to be horizontally
partitioned if multiple parties collect the same variables though from
different individuals, e.g., two hospitals who want to build a model to
predict heart failure. It is said to be vertically partitioned when multiple
parties collect different variables about the same individuals, for example,
data from a hospital and from a health insurance company where both parties
have unique variables about the same patients.

A type of model that can benefit from federated learning are Bayesian
networks. Bayesian networks are probabilistic graphical models that have been
widely used in artificial intelligence [3]. They are popular because they can
be built, verified, or improved, by combining data with existing expert
knowledge. For example, medical doctors can manually create the network
structure, ensuring it models already known dependencies correctly, while the
conditional probability distributions are estimated from data. Thanks to its
graphical representation and probabilistic reasoning, it is also a relatively
intuitive model for non-technical personnel. This makes it very useful in
scenarios where non-technical personnel needs to make decisions based on the
model, for example when used as a tool to inform healthcare policies.

While research has been published on federated learning of Bayesian networks
[4] [5] [6] [7], publications on Bayesian networks trained on vertically
partitioned data (also referred to as heterogeneous data in the literature)
are limited. For example, to the best of our knowledge there are no published
approaches to train a Bayesian network from vertically partitioned data with
missing values. Furthermore, the published work on vertically partitioned
data, has focused on scenarios involving only two parties. In this article, we
propose a novel method called VertiBayes to train Bayesian networks on
vertically partitioned data, which can handle missing values as well as an
arbitrary number of parties. We will then discuss the strengths and
limitations of our proposal.

##  2 Methods

In this section, we will start by shortly explaining how a Bayesian network is
generally trained in a central setting. After this we will discuss how the
various methods used to build and validate a Bayesian network can be adapted
to a vertically partitioned federated setting.

###  2.1 Training a Bayesian network

Training a Bayesian network consists of two phases: structure learning and
parameter learning. The first phase, structure learning, consists of
determining the structure of the graph (i.e., the set of links between
variables) and can be done either manually, using expert knowledge, or
automatically, using algorithms such as K2[8]. The second phase is the so-
called parameter learning. In this phase, the conditional probability
distributions (CPDs) for each node in the network are determined. Throughout
this paper, we will focus on CPDs in the form of conditional probability
tables (CPTs) as these are the most common form of CPD. In the next
subsections, we will discuss how this is done in a centrally trained scenario.
After which we will discuss how these methods need to be adapted for the
federated scenario.

####  2.1.1 Structure learning

The structure of a Bayesian network can be either determined manually or
learnt using an algorithm. Here, we focus on the latter, since the former does
not involve data analysis. One of the most popular structure learning
algorithms is K2, which performs a heuristic search for a viable structure by
scoring potential parent nodes for a given node and step-wise adding the
highest scoring parent[8]. The scoring function used in K2 is described in
equation 1 below.

$f(i,\pi_{i})=\prod_{j=1}^{q_{i}}\frac{(r_{i}-1)!}{(N_{ij}+r_{i}-1)!}\prod_{k=1}^{r_{i}}a_{ijk}!\\\
$ (1)

where $pi_{i}$ is the set of parents of node $x_{i}$. $q_{i}$ is the number of
possible instantiations of the parents of $x_{i}$ present in the data. $r_{i}$
is the number of possible values the attribute $x_{i}$ can take. $a_{ijk}$ is
the number of cases in the dataset where $x_{i}$ has it’s $k^{th}$ value and
the parents are initiated with their $j^{th}$ combination.
$N_{ij}=\sum_{k=1}^{r_{i}}a_{ijk}$, is the number of instances where the
parents of $x_{i}$ are initiated with their $j^{th}$ combination.

It is important to note that the resulting structure depends on the order in
which nodes are introduced into the K2 algorithm. As such, it is possible to
construct different structures for the same data.

####  2.1.2 Parameter learning

There are two relevant scenarios to consider when performing parameter
learning: with and without missing data. When there is no missing data, CPDs
can be learned using the maximum likelihood [9]. To calculate the maximum
likelihood for an attribute $X$ with a set of parents $Y$ one simply has to
calculate: $P(X=x_{i}|Y_{i}=y_{i})=N_{(}x_{i},y_{i})/N_{(}y_{i})$ , where
$N_{(}x_{i},y_{i})$ is the number of records where $X=x_{i}$ and $Y_{i}=y_{i}$
and $N_{(}y_{i})$ is the number of records where $Y_{i}=y_{i}$. In the
presence of missing data, the maximum likelihood for training a Bayesian
network is commonly estimated using algorithms such as Expectation
Maximization (EM)[10][11]. The EM algorithm consists of the following two
steps repeated iteratively until convergence is reached:

  1. 1.

Estimate the likelihood of your data using your current estimates of the
probabilities.

  2. 2.

Update your estimates.

It is important to note that since this algorithm is a hill climbing-type
algorithm, it can get stuck in local optima. Therefore, it is good practice to
run the algorithm several times with different random initializations and use
the best result[9].

###  2.2 Federated learning in a vertically split setting

Training a Bayesian network in a vertically partitioned federated setting
introduces additional hurdles and concerns. First, we will discuss how to
perform structure learning. In the second subsection, we will discuss
parameter learning.

####  2.2.1 Structure learning

As mentioned previously, structure learning can be done using the K2
algorithm. In this subsection, we will discuss how to adapt the K2 algorithm
to a vertically partitioned scenario. To solve this equation, the following
information needs to be collected:

  1. 1.

The number of possible values for the attribute $X$.

  2. 2.

The number of instances that fulfill $X=x_{i}$ and $Y=y_{j}$, where $X$ is the
child attribute, $x_{i}$ is a given value for $X$, $Y$ is the set of parent
attributes, and $y_{j}$ is a given set of assigned values to $Y$. This needs
to be calculated for every possible set $j$.

The number of possible values of attribute $X$ can be calculated trivially
without revealing any important information to an external party in a
vertically split federated setting as all relevant information is available
locally at one party.

To calculate the number of instances that fulfill $X=x_{i}$ and $Y=y_{i}$ (the
number of instances for each possible value of $X$ for each possible
configuration of $X$’s parents) one has to calculate the number of instances
that fulfill certain conditions across different datasets. Fortunately, this
problem can be solved using a privacy preserving scalar product protocol and
calculating the scalar product of vectors, one for each site, where each
individual is represented as $1$ or $0$ depending on whether they fulfill the
local conditions (in this case whether the child and parent nodes have the
appropriate values). Earlier research has used this approach to calculate the
information-gain when training a decision tree[12], which at its root, poses
the same problem we face here.

Various variants of the privacy preserving scalar product protocol have been
published [12] [13] [14] [15] [16]. Most of these focus on $2$-party scenarios
but variants do exist for $N$ parties[17]. These methods have different
advantages and disadvantages, such as different privacy guarantees and risks,
different runtime complexities, and different communication cost overheads.
Because of this, the preferred method will differ per scenario. A K2
implementation using one of these protocols will have the same privacy
guarantees and risks but will pose no additional privacy concerns beyond those
posed by the chosen protocol.

####  2.2.2 Parameter learning

During parameter learning, the actual CPDs will be calculated. There are two
scenarios that need to be considered: with and without missing values. EM
works under the assumption that data is missing at random or missing
completely at random.

##### Without missing values

As discussed earlier, parameter learning without missing values can be done by
calculating the maximum likelihood for various attribute values. This means
calculating for each node $i$, $N_{ij}$, the number of samples for each
possible configuration of the parents of node $i$ and $N_{ijk}$ the number of
samples for each possible configuration of the parents where the value is $k$,
for each possible value of the node. $N_{ijk}$ can be calculated by simply
summing the various $N_{ij}$ values. For the sake of performance it is advised
to do this. These can be calculated using the scalar product protocol as
explained earlier when describing the solution for K2. As such, performing
parameter learning in a vertically split federated scenario with no missing
values is not a problem and can be done without any significant additional
privacy risks compared to the central variant beyond the risks involved in the
scalar product protocol implementation used.

##### With missing values

In the presence of missing data, it is not possible to use the maximum
likelihood and one has to use more sophisticated approaches such as the EM
algorithm. As discussed earlier, this algorithm consists of two steps repeated
iteratively. The E-step where the likelihood of the training data given the
current estimates of the CPDs is calculated, and a M-step where the current
estimates are updated.

To estimate the likelihood of the current estimates in the E-step the
following equation needs to be solved:

$E=\prod_{i=1}^{n}P(d_{i})=\prod_{i=1}^{n}\prod_{j=1}^{p}P(x_{ij}|y_{ij})$ (2)

where $n$ is the number of samples in the dataset, $p$ is the number of nodes
in the network, $P(d_{i})$ is the likelihood of the $i-th$ sample, $x_{ij}$ is
the value of the $j-th$ node in the $i-th$ sample, and $y_{ij}$ is the set of
values of the parents of the $j-th$ node in the $i-th$ sample. The appropriate
values $P(x_{ij}|y_{ij})$ need to be selected from the current estimate of the
CPDs based on the attribute values of this particular sample.

However, selecting these in a vertically split federated setting is not
trivial without revealing individual level data. In principle, it is possible
to use the scalar product protocol to select the appropriate probability from
the CPD without revealing the attribute values directly. This can be done by
representing the CPD as one of the vectors used in the scalar product
protocol, while the other vectors consists of $0$’s and $1$’s representing if
this particular cell in the CPD is applicable given the local data each party
has. Unfortunately, it should be noted that the selected probability from the
CPD may not be revealed, as the probability might be unique in the table,
which would allow a party to determine the values of the parents and the node
simply by looking up the probability in the CPD. This means that the practical
utility of this alternative is limited to scenarios in which the final model
is never published.

Revealing the selected probability could be prevented by using fully
homomorphic encryption and only decrypting once $P(d_{i})$ has been
calculated. However, this is impractical due to the high runtime complexity of
current fully homomorphic encryption schemes, the runtime complexity of a
privacy preserving scalar product protocol, the large number of scalar product
protocols which would need to be executed for each $P(d_{i})$, and the need to
repeat this process for each E-step.

Therefore, we conclude that the EM algorithm cannot be easily applied in a
vertically split federated scenario without severe limitations. Instead, we
propose the following three-step solution, which we have dubbed VertiBayes.

  1. 1.

Treat “missing” as a valid value and train an intermediate Bayesian network
using maximum likelihood on the training data.

  2. 2.

Generate synthetic data (including “missing” values) using this intermediate
Bayesian network

  3. 3.

Train the final model on this synthetic data using the EM algorithm

As discussed earlier, parameter learning in a vertically split federated
setting without missing values is possible with the privacy guarantees
provided by the privacy preserving scalar product protocol used. Generating
synthetic data by using this intermediate model also does not add any
additional privacy concerns compared to a centrally trained model, as this is
a basic functionality of any Bayesian network. On the contrary, the final
model has a reduced risk of data leak because it is trained on synthetic
data[18][2].

Figure 1: Training process for VertiBayes

The proposed process, which is illustrated in Figure 1, allows us to train a
Bayesian network in a vertically split federated setting with missing values
without any additional privacy concerns compared to a centrally trained model.
However, it should be noted that it is possible that a loss of signal may
occur due to the three-step approach. In the experiment section, we will test
if our proposed method avoids this potential pitfall.

###  2.3 Time complexity when training a model in a federated setting

A major downside to federated learning is that the time complexity is usually
considerably worse compared to the centralized setting. This is unavoidable
due to the extra overhead created by communication as well as the increased
complexity introduced by the privacy preserving mechanisms. In this
subsection, we will discuss the time complexity of VertiBayes.

In our implementation, there are two important factors to take into account.
The first is the number of parties $n$. This is a major bottleneck as the
$n$-party scalar product protocol implementation we have used scales
combinatorically in the number of parties.

The second important factor is the size of the CPDs that need to be
calculated, as each unique probability that needs to be calculated requires a
separate $n$-party scalar product protocol to be solved. As such, our
implementation scales linearly in the number of probabilities that need to be
calculated. The time complexity of various aspects for our implementation can
be found in table I.

TABLE I: Time complexity Number of Scalar product protocols |  $O(m)$, where $m$ is the number of unique parent-child value combinations for which a probability needs to be calculate
---|---
Number of scalar product subprotocols per protocol |  $\frac{n!}{(x!(n-x)!))}$, for each $x,2<=x<=n$, where $n$ is the number of parties involved in the protocol
Number of multiplications per subprotocol |  $O(p*n*(n-1))$, where $p$ is the population size, and $n$ is the number of parties involved in the protocol

Important to note is that the population size is not a main driver of the
runtime. A relatively simple network trained on a small dataset, but with a
high number of unique attribute values will have a significantly longer
runtime than a more complex network with few unique attribute values. This is
because the overall time complexity is dominated by the number of scalar
product protocols and subprotocols, which is independent of the population
size, but dependent on the number of probabilities that need to be calculated.
Finally, it is important to note that there is ample room for parallelization
to improve the running time as each scalar product protocol that is needed for
VertiBayes is fully independent and can easily be run in parallel.

###  2.4 Federated classification and model validation

The process of using the model to classify new instances in a federated
setting is itself a complex problem that depends strongly on the type of model
used. In this subsection, we will discuss the methods that are available to
classify an individual in a vertically partitioned setting using a Bayesian
Network and the implication this has for the validation of the model.

Classification of new samples using a Bayesian Network in a vertically
partitioned federated setting suffers from the same issues as the expectation
step in the EM algorithm. To classify an instance from vertically partitioned
data, we need to select the appropriate probabilities from the CPD. As
discussed before, doing this without revealing data to other parties would
require a combination of intricate data representation, the scalar product
protocol, and fully homomorphic encryption. While this might be feasible when
classifying one instance at a time, it is not practical when classifying a
large number of instances in batch or when the Bayesian network is
particularly complex. This has major consequences for the validation of a new
model in a federated setting. As such, whenever possible the validation should
be done using a publicly available dataset which avoids the need for privacy
preserving measure during validation. If such a dataset is not available, we
propose two different approaches, “Synthetic Cross-fold Validation” (SCV) and
“Synthetic Validation Data Generation” (SVDG), to validate the model in a
privacy preserving manner.

SCV uses the synthetic data generated by the intermediate Bayesian network as
both training and validation data by executing the EM training using k-fold
cross validation. However, it is possible that this results in overfitting on
the synthetic data and therefore the performance estimate may be biased by the
intermediate Bayesian network.

SVDG splits the private dataset into training and validation sets. It will
then train a Bayesian network on the training set in a federated manner as
normal. On the validation set, it will train a federated network using only
the federated maximum likelihood approach. We can then use the Bayesian
network trained on the validation set to generate a synthetic validation
dataset. This approach reduces the risk of overfitting the previous approach
suffered from but may lead to biased estimates if the synthetic validation set
is not representative of the original validation set. These approaches avoid
leaking real data, but as mentioned, they may not be viable in practice. An
illustration of the two new approaches can be seen in Figure 2

Figure 2: Flow diagrams for proposed validation procedures SCV (left) and SVDG
(right)

In the next section, we will perform experiments to validate that VertiBayes
results in networks with similar performance as a centrally trained model. We
will also verify if the two new approaches to validation give the same results
as the validation on the public data, and if there are scenarios where they
are inappropriate due to the aforementioned risks.

###  2.5 Experimental setup

In order to validate our proposed approach, we have implementedit in a
combination of Java and Python111 Our code can be found in the following two
git repositories  • Main algorithm code:
https://github.com/MaastrichtU-CDS/VertiBayes • Vantage6 wrapper code:
https://github.com/MaastrichtU-CDS/VertiBayes_vantage6 using Vantage6[19] and
ran a number of experiments. Vantage6 is an open-source infrastructure for
privacy preserving federated learning which utilizes Docker. We compare the
performance of our algorithm with a centrally trained Bayesian network using
WEKA[20], a machine learning library written in Java.

####  2.5.1 Structure learning

To validate our federated implementation of the K2 algorithm we ran
experiments using the Iris[21], Asia[22], Alarm[23], and Diabetes[24]
datasets. As K2 is deterministic and dependent on the order in which the nodes
are put into the algorithm we ensured this was the same for both the
federatedly and centrally trained model and then compared the resulting
structures.

####  2.5.2 Parameter learning

In our experiments regarding parameter learning, we have used the Iris[21],
Asia[22], Alarm[23], and Diabetes[24] datasets. In the case of the Iris
dataset, we predict the “label” attribute, for Asia we predict “lung”, for
Alarm we will be predicting “BP” and for Diabetes we will predict “outcome”.
The Iris dataset contains $150$ samples, the Diabetes dataset contains $768$
samples, while the other two datasets contain data of $10.000$ samples. The
Asia and Alarm datasets come with a predefined structure. The Iris dataset
uses a naïve Bayes structure. The Diabetes dataset also uses a predefined
structure.

Both the Iris and Diabetes dataset contain continuous variables. For the sake
of a fair comparison between the central and federated models, these were
discretized into bins before starting our experiments, where each bin contains
at least $10\%$ of the total population as well as a minimum of $10$
individuals. If the last bin cannot be made large enough to fulfill these
criteria it is simply added to the previous bin. In the case of a dataset of
less than $10$ individuals, the bin will simply contain all possible values.
For our experiments the bins were predefined alongside the predefined network
structure, but the bins can also be generated on the fly during the training
of the federated model using the same discretization strategy. The simplicity
of this strategy allows it to be executed without needing additional privacy
preserving mechanics. However, it should be noted that this is not the best
possible discretizing strategy. For example, a model using the Minimum
Description Length method (MDL)[25] for discretization, or utilizing expert
knowledge, might outperform this setup.

Slight variations of this discritezation strategy were used in preliminary
experiments. However, we will not list the results of those variants here as
they produced similar results.

To test the effect of missing values we have done experiments where we
randomly set $0\%$, $5\%$, $10\%$ and $30\%$ of the values to missing. The
performance of the models is measured by calculating the area under receiver
operating characteristics curve (AUC). The centrally trained model is
internally validated using $10$-fold cross validation. The federated model is
validated using the two different validation approaches described in the last
section; it is also validated against a “public” central dataset, which is
simply the left-out fold from the original dataset. All of these approaches
use $10$-fold cross validation. We also compare the Akaike information
criterion (AIC)[26] values of the network for their original (private)
training data. This is done to determine if there is any difference in the
CPDs used by the two models.

For all of these experiments the data is randomly partitioned over two parties
by dividing the original dataset into two equally large sets of attributes.

##  3 Results

The results of our experiments can be found in this section. First we will
discuss the results of the structure learning experiments. After that we will
cover the experiments regarding parameter learning.

###  3.1 Structure learning

The federated implementation of the K2 algorithm consistently resulted in the
same structure as the centrally trained model for all datasets.

###  3.2 Parameter learning

Table II shows the results of our parameter learning experiments. It lists the
AUC for the centrally trained model, as well as the AUC’s for the various
(federated) validation schemes. The AUC’s of the federated model are listed in
bold if the difference with the central model is larger than $0.05$. The AIC
is shown for the centrally trained model, this is compared to the AIC for the
federated model trained. The difference between the central and federated AIC
is listed in bold if the difference is at least $5\%$. An AIC closer to $0$ is
better. A negative percentage in the AIC difference column indicates the
federated model performed better.

TABLE II: Results of the experiments.
The experiments are grouped per dataset. The AUC is represented for the centrally trained model, as well as the two federated validation schemes; “Synthetic Cross-fold Validation” (SCV) and “Synthetic Validation Data Generation” (SVDG). AUC values are listed in bold if they differ more than $0.05$ from the central model. An AIC closer to $0$ is better, a negative percentage in the AIC difference collumn indicates the federated model was better. This value is listed in bold if the difference between the federated AIC and the central AIC is at least $5\%$ | AUC |  | AIC
---|---|---|---
| Training method |  | Training Method |
| Centrally trained | Federatedly trained |  |  Centrally trained |  Federatedly trained |  AIC difference
Dataset |  Missing data level | Public validation | Public validation | SCV validation | SVDG validation |
Alarm population size: 10000 | 0 | 0,91 | 0,91 | 0,91 | 0,91 |  | -340571 | -318469 |  -6,49%
0,05 | 0,88 | 0,89 | 0,88 | 0,89 |  | -315856 | -313612 | -0,71%
0,01 | 0,85 | 0,86 | 0,86 | 0,86 |  | -340823 | -350576 | 2,86%
0,3 | 0,76 | 0,76 | 0,76 | 0,76 |  | -444297 | -444866 | 0,13%
Asia population size: 10000 | 0 | 1,00 | 1,00 | 1,00 | 1,00 |  | -22555 | -22559 | 0,02%
0,05 | 0,76 | 0,76 | 0,76 | 0,76 |  | -23430 | -23395 | -0,15%
0,01 | 0,69 | 0,7 | 0,7 | 0,7 |  | -24105 | -24090 | -0,06%
0,3 | 0,58 | 0,59 | 0,58 | 0,59 |  | -25517 | -25613 | 0,37%
Diabetes population size: 768 | 0 | 0,8 | 0,77 | 0,95 | 0,79 |  | -13593 | -14407 | 5,99%
0,05 | 0,79 | 0,74 | 0,92 | 0,76 |  | -13699 | -14556 | 6,25%
0,01 | 0,75 | 0,72 | 0,90 | 0,73 |  | -13761 | -14408 | 4,70%
0,3 | 0,61 | 0,57 | 0,80 | 0,6 |  | -14736 | -15136 | 2,71%
Iris population size: 150 | 0 | 0,99 | 0,98 | 1,00 | 1,00 |  | -1036 | -1022 | -1,33%
0,05 | 0,97 | 0,96 | 0,99 | 0,97 |  | -1243 | -1176 |  -5,37%
0,01 | 0,9 | 0,89 | 0,95 | 0,92 |  | -1491 | -1022 |  -31,49%
0,3 | 0,98 | 0,99 | 1,00 | 0,99 |  | -1099 | -1381 | 25,67%

##  4 Discussion

In this paper, we have proposed a novel method to train Bayesian networks in a
federated setting using vertically partitioned data with missing values. We
have shown that it is possible to perform both structure and parameter
learning of Bayesian networks in such a setting with reasonable accuracy.
Structure learning can be performed by adapting any of the existing structure
learning algorithms to use a secure multiparty computation algorithm. In this
study, we used a protocol within the K2 algorithm, but the same approach could
be applied to other score-based algorithms or even constraint-based algorithms
such as the PC algorithm[27]. Parameter learning requires one of two
approaches. When there is no missing data present, the scalar product protocol
can be used directly to compute the maximum likelihood, or when missing data
is present the three-step solution described in section 2.2 using the EM
algorithm can be applied. We will now discuss the performance of VertiBayes
compared to a centrally trained model as well as the limitations in terms of
scalability. Lastly, we will discuss the sensitive information that may be
leaked by any Bayesian network and the limitations this brings in a federated
setting.

###  4.1 Model performance and validation

Our experiments show that the resulting models produced by VertiBayes are
comparable to the centrally trained models. As such, there is generally no
meaningfull difference in terms of AUC or AIC. Furthermore, the experiments
show that it is possible to validate the model in a privacy-preserving manner
despite it being impossible to efficiently classify an individual in a
privacy-preserving manner.

It is however important to note that in certain edge-cases some validation
approaches can show unreliable results. For example, the SVDG approach can
cause problems when the test-folds are too small and the bins are generated on
the fly while training, as opposed to working with pre-defined bins.If there
are not enough individuals in the test-set to create multiple accurate bins on
the fly this strategy will result in a loss of information. This was most
notable when running the preliminary experiments with the Iris dataset, which
is quite small.

Similarly, the model ran into problems using the SCV approach whenever the
CPDs become too large because a node has multiple parents with a significant
amount of bins. This is notable in the results of the Diabetes model. Certain
nodes with multiple parents would end up with CPDs that contained more cells
than there are individuals in the dataset. This led to an overestimation of
the AUC if the SCV approach was used. However, this was not an issue when
utilizing the SVDG approach, due to the stronger separation between training
and validation data. Using larger bins can alleviate this problem to some
extent. However, the bins cannot be made arbitrarily large as this will
eventually cause a loss of information. Using expert-knowledge based
discretization strategies tailored to each dataset, or a better automatic
discretization strategy such as the MDL method mentioned earlier, would help
avoid these problems.

These problems of overfitting and loss of information, show that it is
extremely important to have an appropriate discretization strategy. So long as
the potential pitfalls surrounding discretization are addressed, VertiBayes
can be used to train and validate a Bayesian network in a vertical federated
setting.

###  4.2 Scalability

Any algorithm that is adapted to a federated setting will be slower than the
central counterpart due to the overhead caused by the protocols used to
protect privacy. Our experiments confirmed the potential issues we brought up
in section 2.3 when we discussed the theoretical time complexity.

The effects of population-size proved to be negligible. The effects of the
complexity of the network structure, that is to say the number of nodes and
links within the network, is only relevant in so much that it creates more
probabilities to calculate. As such, it runs in comparable time as the
experiment that simply adds a few additional values to reach the same number
of probabilities. As expected, the size of the CPDs that had to be calculated
had the greatest effect on the total runtime.

This does mean that there are practical limitations to using VertiBayes, as it
may take too long to train a large network. However, it should be noted that
in certain settings a longer runtime might still be acceptable. For example,
it is perfectly acceptable that training a model for use in a clinical setting
takes an extended amount of time.

###  4.3 Sensitive information in published Bayesian networks

Publishing a Bayesian network, or any machine learning model, will reveal
certain information about the training data, regardless of how the network is
trained. When publishing a Bayesian network two important aspects will be
revealed: the network structure and the CPDs. The network structure will only
reveal which conditional dependencies exist amongst attributes, which is not
sensitive data in most scenarios. The CPDs on the other hand, can potentially
be used to reconstruct individual level data from the trainingset, when the
probabilities in the CPD’s are based on one, or a few individuals. An
effective countermeasure is using k-anonymity[28] to ensure that each
probability in the CPDs represents a minimum amount of samples and that no
probabilities of $0$ or $1$ are present in the CPDs. Such probabilities make
it considerably easier to deduce individual level data, and they can also lead
to artifacts when using the network. Lastly, a public Bayesian network can be
used by one of the parties that participated in the training to guess
(although not reconstruct) the data of the other parties based on their own
data. Similarly, any third party with partial data can use the final model to
estimate the missing values in his dataset. This is unavoidable and it should
be taken into consideration when decisions are made about which models are to
be made public.

These concerns imply that there are practical limits to what privacy
preserving techniques should aim for. Trying to prevent any and all privacy
issues using privacy preserving techniques during the training phase is futile
when models are made public as the models themselves will always reveal some
information.

##  5 Conclusion

In this paper, we have proposed a novel approach to train Bayesian network
parameters from a vertically partitioned data with and without missing values.
This method can deal with an arbitrary number of parties, only limited by the
runtime. We have shown that there are no additional privacy risks compared to
a centrally trained model beyond the ones presented by the specific privacy
preserving scalar product protocol implementation used. Our experiments show
there are no meaningful differences in performance between models trained with
VertiBayes and models trained centrally, provided continuous variables are
adequately discretized. They also show it is possible to estimate the
performance of a model with vertically partitioned data with a reasonable
accuracy. As such, VertiBayes is a useful tool for training Bayesian networks
in a vertically partitioned setting.

##  6 Future Work

When using the model in a federated setting with a vertical split, it is
currently not possible to efficiently classify or predict a new instance in a
privacy preserving manner using a Bayesian network. It would be beneficial if
a solution for this was found and implemented. In addition to this, VertiBayes
could be improved by implementing more advanced discretization methods, such
as MDL, in a vertically partitioned setting as our current implementation
relies either on a very basic automatic discretization approach or the use of
experts, which may not be the best discretization approach possible. Lastly,
the impact of different missing data mechanisms on our proposed approach
should be investigated. In this article we used data that missed completely at
random, however, we did not look at other missing mechanisms, such as missing
at random. It would be worthwhile to investigate whether this significantly
influences the performance of the resulting model.

## References

  * [1] L. Li, Y. Fan, M. Tse, and K.-Y. Lin, “A review of applications in federated learning,” _Computers & Industrial Engineering_, vol. 149, p. 106854, Nov. 2020. [Online]. Available: https://www.sciencedirect.com/science/article/pii/S0360835220305532
  * [2] P. Kairouz, H. B. McMahan, B. Avent, A. Bellet, M. Bennis, A. N. Bhagoji, K. Bonawitz, Z. Charles, G. Cormode, R. Cummings, R. G. L. D’Oliveira, S. E. Rouayheb, D. Evans, J. Gardner, Z. Garrett, A. Gascón, B. Ghazi, P. B. Gibbons, M. Gruteser, Z. Harchaoui, C. He, L. He, Z. Huo, B. Hutchinson, J. Hsu, M. Jaggi, T. Javidi, G. Joshi, M. Khodak, J. Konečný, A. Korolova, F. Koushanfar, S. Koyejo, T. Lepoint, Y. Liu, P. Mittal, M. Mohri, R. Nock, A. Özgür, R. Pagh, M. Raykova, H. Qi, D. Ramage, R. Raskar, D. Song, W. Song, S. U. Stich, Z. Sun, A. T. Suresh, F. Tramèr, P. Vepakomma, J. Wang, L. Xiong, Z. Xu, Q. Yang, F. X. Yu, H. Yu, and S. Zhao, “Advances and Open Problems in Federated Learning,” _arXiv:1912.04977 [cs, stat]_ , Mar. 2021, arXiv: 1912.04977. [Online]. Available: http://arxiv.org/abs/1912.04977
  * [3] J. Pearl, _Probabilistic Reasoning in Intelligent Systems: Networks of Plausible Inference_. Elsevier, Jun. 2014, google-Books-ID: mn2jBQAAQBAJ.
  * [4] Z. Yang and R. Wright, “Privacy-Preserving Computation of Bayesian Networks on Vertically Partitioned Data,” _IEEE Transactions on Knowledge and Data Engineering_ , vol. 18, no. 9, pp. 1253–1264, Sep. 2006. [Online]. Available: http://ieeexplore.ieee.org/document/1661515/
  * [5] R. Wright and Z. Yang, “Privacy-preserving Bayesian network structure computation on distributed heterogeneous data,” in _Proceedings of the tenth ACM SIGKDD international conference on Knowledge discovery and data mining_ , ser. KDD ’04. New York, NY, USA: Association for Computing Machinery, Aug. 2004, pp. 713–718. [Online]. Available: https://doi.org/10.1145/1014052.1014145
  * [6] Zhiqiang Yang and R. Wright, “Improved Privacy-Preserving Bayesian Network Parameter Learning on Vertically Partitioned Data,” in _21st International Conference on Data Engineering Workshops (ICDEW’05)_. Tokyo, Japan: IEEE, 2005, pp. 1196–1196. [Online]. Available: http://ieeexplore.ieee.org/document/1647809/
  * [7] I. Ng and K. Zhang, “Towards Federated Bayesian Network Structure Learning with Continuous Optimization,” _arXiv:2110.09356 [cs, stat]_ , Oct. 2021, arXiv: 2110.09356. [Online]. Available: http://arxiv.org/abs/2110.09356
  * [8] G. F. Cooper and E. Herskovits, “A Bayesian method for the induction of probabilistic networks from data,” _Machine Learning_ , vol. 9, no. 4, pp. 309–347, Oct. 1992. [Online]. Available: https://doi.org/10.1007/BF00994110
  * [9] D. Koller and N. Friedman, _Probabilistic Graphical Models: Principles and Techniques_. MIT Press, Jul. 2009, google-Books-ID: 7dzpHCHzNQ4C.
  * [10] A. P. Dempster, N. M. Laird, and D. B. Rubin, “Maximum Likelihood from Incomplete Data via the EM Algorithm,” _Journal of the Royal Statistical Society. Series B (Methodological)_ , vol. 39, no. 1, pp. 1–38, 1977, publisher: [Royal Statistical Society, Wiley]. [Online]. Available: https://www.jstor.org/stable/2984875
  * [11] S. L. Lauritzen, “The EM algorithm for graphical association models with missing data,” _Computational Statistics & Data Analysis_, vol. 19, no. 2, pp. 191–201, Feb. 1995. [Online]. Available: https://www.sciencedirect.com/science/article/pii/0167947393E0056A
  * [12] W. Du and Z. Zhan, “Building decision tree classifier on private data,” in _Proceedings of the IEEE international conference on Privacy, security and data mining - Volume 14_ , ser. CRPIT ’14. AUS: Australian Computer Society, Inc., Dec. 2002, pp. 1–8.
  * [13] W. Du and M. Atallah, “Privacy-preserving cooperative statistical analysis,” in _Seventeenth Annual Computer Security Applications Conference_. New Orleans, LA, USA: IEEE Comput. Soc, 2001, pp. 102–110. [Online]. Available: http://ieeexplore.ieee.org/document/991526/
  * [14] M. J. Atallah and W. Du, “Secure Multi-party Computational Geometry,” in _Algorithms and Data Structures_ , G. Goos, J. Hartmanis, J. van Leeuwen, F. Dehne, J.-R. Sack, and R. Tamassia, Eds. Berlin, Heidelberg: Springer Berlin Heidelberg, 2001, vol. 2125, pp. 165–179, series Title: Lecture Notes in Computer Science. [Online]. Available: http://link.springer.com/10.1007/3-540-44634-6_16
  * [15] B. Goethals, S. Laur, H. Lipmaa, and T. Mielikäinen, “On Private Scalar Product Computation for Privacy-Preserving Data Mining,” in _Information Security and Cryptology – ICISC 2004_ , D. Hutchison, T. Kanade, J. Kittler, J. M. Kleinberg, F. Mattern, J. C. Mitchell, M. Naor, O. Nierstrasz, C. Pandu Rangan, B. Steffen, M. Sudan, D. Terzopoulos, D. Tygar, M. Y. Vardi, G. Weikum, C.-s. Park, and S. Chee, Eds. Berlin, Heidelberg: Springer Berlin Heidelberg, 2005, vol. 3506, pp. 104–120, series Title: Lecture Notes in Computer Science. [Online]. Available: http://link.springer.com/10.1007/11496618_9
  * [16] J. Vaidya and C. Clifton, “Privacy preserving association rule mining in vertically partitioned data,” in _Proceedings of the eighth ACM SIGKDD international conference on Knowledge discovery and data mining_ , ser. KDD ’02. New York, NY, USA: Association for Computing Machinery, Jul. 2002, pp. 639–644. [Online]. Available: https://doi.org/10.1145/775047.775142
  * [17] F. van Daalen, I. Bermejo, L. Ippel, and A. Dekkers, “Privacy preserving n-party scalar product protocol.” [Online]. Available: http://arxiv.org/abs/2112.09436
  * [18] N. C. Abay, Y. Zhou, M. Kantarcioglu, B. Thuraisingham, and L. Sweeney, “Privacy preserving synthetic data release using deep learning,” in _Machine Learning and Knowledge Discovery in Databases_ , M. Berlingerio, F. Bonchi, T. Gärtner, N. Hurley, and G. Ifrim, Eds. Springer International Publishing, vol. 11051, pp. 510–526, series Title: Lecture Notes in Computer Science. [Online]. Available: http://link.springer.com/10.1007/978-3-030-10925-7_31
  * [19] A. Moncada-Torres, F. Martin, M. Sieswerda, J. Van Soest, and G. Geleijnse, “VANTAGE6: an open source priVAcy preserviNg federaTed leArninG infrastructurE for Secure Insight eXchange,” _AMIA … Annual Symposium proceedings. AMIA Symposium_ , vol. 2020, pp. 870–877, 2020.
  * [20] E. Frank, I. H. Witten, and M. A. Hall, _Data Mining, Fourth Edition:Practical Machine Learning Tools and Techniques | Guide books_. Morgan Kaufmann Publishers Inc, 2016.
  * [21] M. De Marsico, M. Nappi, D. Riccio, and H. Wechsler, “Mobile Iris Challenge Evaluation (MICHE)-I, biometric iris dataset and protocols,” _Pattern Recognition Letters_ , vol. 57, pp. 17–23, May 2015. [Online]. Available: https://www.sciencedirect.com/science/article/pii/S0167865515000574
  * [22] S. L. Lauritzen and D. J. Spiegelhalter, “Local computations with probabilities on graphical structures and their application to expert systems,” vol. 50, no. 2, pp. 157–194, _eprint: https://onlinelibrary.wiley.com/doi/pdf/10.1111/j.2517-6161.1988.tb01721.x. [Online]. Available: https://onlinelibrary.wiley.com/doi/abs/10.1111/j.2517-6161.1988.tb01721.x
  * [23] I. A. Beinlich, H. J. Suermondt, R. M. Chavez, and G. F. Cooper, “The ALARM Monitoring System: A Case Study with two Probabilistic Inference Techniques for Belief Networks,” _AIME 89_ , pp. 247–256, 1989, publisher: Springer, Berlin, Heidelberg. [Online]. Available: https://link.springer.com/chapter/10.1007/978-3-642-93437-7_28
  * [24] J. W. Smith, J. Everhart, W. Dickson, W. Knowler, and R. Johannes, “Using the ADAP Learning Algorithm to Forecast the Onset of Diabetes Mellitus,” _Proceedings of the Annual Symposium on Computer Application in Medical Care_ , pp. 261–265, Nov. 1988. [Online]. Available: https://www.ncbi.nlm.nih.gov/pmc/articles/PMC2245318/
  * [25] U. M. Fayyad and K. B. Irani, “Multi-Interval Discretization of Continuous-Valued Attributes for Classification Learning,” in _Ijcai._ , 1993, pp. 1022–1029.
  * [26] H. Akaike, “A new look at the statistical model identification,” _IEEE Transactions on Automatic Control_ , vol. 19, no. 6, pp. 716–723, Dec. 1974, conference Name: IEEE Transactions on Automatic Control.
  * [27] P. Spirtes, C. N. Glymour, P. Spirtes, and C. Glymour, “An algorithm for fast recovery of sparse causal graphs,” _Social Science Computer Review_ , pp. 62–72, 1991.
  * [28] L. Sweeney, “k-ANONYMITY: A MODEL FOR PROTECTING PRIVACY,” _International Journal of Uncertainty, Fuzziness and Knowledge-Based Systems_ , vol. 10, no. 05, pp. 557–570, Oct. 2002, publisher: World Scientific Publishing Co. [Online]. Available: https://www.worldscientific.com/doi/abs/10.1142/S0218488502001648

|  Florian van Daalen Florian van Daalen received his BSc degree in Knowledge
Engineering from University Maastricht in 2012 and his MSc degree in
Artificial Intelligence in 2014. He is currently working toward the PhD degree
in Clinical Data Science within the Clinical Data Science group, University
Maastricht, Netherlands. His research interests include privacy preserving
techniques, federated learning, and ensemble based learning.
---|---
|  Lianne Ippel Lianne Ippel received her PhD in Statistics from Tilburg
University on analyzing data streams with dependent observations, for which
she won the dissertation award from General Online Research conference (2018).
After a Postdoc at Maastricht University, she now works at Statistics
Netherlands where she works at the methodology department on international
collaborations and innovative methods for primary data collection.
---|---
|  Andre Dekker Prof. Andre Dekker, PhD (1974) is a medical physicist and
professor of Clinical Data Science at Maastricht University Medical Center and
Maastro Clinic in The Netherlands. His Clinical Data Science research group
(50 staff) focuses on 1) federated FAIR data infrastructures, 2) AI for health
outcome prediction models and 3) applying AI to improve health. Prof. Dekker
has authored over 200 publications, mentored more than 30 PhD students and
holds multiple awards and patents on the topic of federated data and AI. He
has held visiting scientist appointments at universities and companies in the
UK, Australia, Italy, USA and Canada.
---|---
|  Inigo Bermejo Inigo Bermejo received the BSc degree on Computer Engineering
from the University of the Basque Country, Spain, in 2006 and the PhD in
Intelligent Systems from UNED, Spain, in 2015. He is currently a postdoctoral
researcher at the Clinical Data Science group, Maastricht University. His
research interests include privacy preserving techniques, prediction modelling
and causal inference.
---|---


# Differentiable Good Arm Identification

Yun-Da Tsai, Tzu-Hsien Tsai, Shou-De Lin

###### Abstract

This paper focuses on a variant of the stochastic multi-armed bandit problem
known as good arm identification (GAI). GAI is a pure-exploration bandit
problem that aims to identify and output as many good arms as possible using
the fewest number of samples. A good arm is defined as an arm whose expected
reward is greater than a given threshold. We present our study in the context
of a structured bandit setting and introduce DGAI, a novel, differentiable
good arm identification algorithm. By leveraging a data-driven approach, DGAI
significantly enhances the state-of-the-art HDoC algorithm empirically.
Additionally, we demonstrate that DGAI can improve the cumulative reward
maximization problem when a threshold is provided as prior knowledge for the
arm set. Extensive experiments have been conducted to validate our algorithm’s
performance. The results demonstrate that our algorithm significantly
outperforms the competitors in both synthetic and real-world datasets for both
the GAI and MAB tasks.

##  1 Introduction

Bandit problems represent a category of sequential decision-making problems
where stochastic rewards are observed as feedback. Within this broad context,
the classic multi-armed bandit (MAB) problem aims to maximize the expected
cumulative reward over a series of trials by navigating the exploration-
exploitation dilemma. Another essential category within the bandit problem is
the best arm identification (BAI), a pure-exploration problem in which the
agent aims to identify the best arm with minimum sample complexity
(kalyanakrishnan2012pac; audibert2010best; kano2019good). This paper’s
specific focus is on a derivative of the pure-exploration problem called the
good arm identification (GAI) problem (kano2019good).

In the GAI problem, a ”good arm” is an arm whose expected reward meets or
exceeds a given threshold. The agent must identify such an arm within a
specified error probability during repeated trials, stopping only when it
believes that no good arms remain. The goal of GAI is to identify as many arms
as possible with as few samples as possible. Key differences between GAI and
BAI include: 1\. GAI identifies arms through absolute comparisons (i.e., above
a threshold) rather than relative (best) comparisons, requiring distinct
optimization strategies. 2\. Unlike BAI, which examines the entire arm set,
GAI can be solved by an anytime algorithm. This flexibility allows GAI to
provide reasonable solutions even if interrupted, making it suitable for time-
sensitive applications like online advertising and high-speed trading
(kano2019good). 3\. GAI handles the exploration-exploitation dilemma with
confidence. Exploitation entails pulling the arm most likely to be good, while
exploration means pulling other arms to build confidence in identifying them
as either good or bad. The overall sample complexity is bounded by union
bounds $\Delta_{i}=|\mu_{i}-\epsilon|$ (i.e., the gap between thresholds for
identification) and $\Delta_{i,j}=\mu_{i}-\mu_{j}$ (i.e., the gap between arms
for sampling) where
$\Delta=\min\\{\min_{i\in[K]}\Delta_{i},\min_{\lambda\in[K-1]}\Delta_{\lambda,\lambda+1}/2\\}$
(see Table 1 for notation explanations).

Common GAI algorithms like HDoC, LUCB-G, and APT-G (kano2019good) encompass a
sampling strategy and an identification criterion (see Algorithm
LABEL:alg:hdoc). The former determines which arm to draw, while the latter
decides whether to accept an arm as good. Once an arm is deemed good, it is
removed from the pool, as there is no need for further sampling. These
strategies formulate confidence bounds through concentration inequalities
based on assumptions about the reward distribution like sub-Gaussianity.
Though essential for theoretical analysis, these bounds may not always align
with practical applications. Constructed confidence bounds tend to be
conservative in real-world situations (osband2015bootstrapped;
kveton2019garbage), as they’re based on assumed reward distributions, leading
to broad, non-adaptive confidence bounds and sub-optimal performance. Our
research highlights the superiority of adaptive, data-driven techniques.

This work introduces the GAI problem within a structured bandit setting at the
intersection of multi-armed bandits and learning halfspaces. We aim for a
data-driven approach to learn the confidence bound, ensuring adaptability to
the actual problem structure. Furthermore, we propose that confidence bounds
for both sampling and identification in GAI should be learned separately in a
data-driven manner to reflect different distributions of gaps between arms and
thresholds (the union bounded sample complexity). We introduce an algorithm
called Differentiable GAI (or DGAI) that employs a differentiable UCB index to
accomplish these objectives. This index offers more information than the
classical UCB index and can be integrated with a smoothed indicator function
to learn the confidence bound for both the sampling strategy and
identification criteria. By using the newly designed differentiable algorithm
and learning objective function, the adaptive confidence bound significantly
enhances empirical performance over existing GAI algorithms in both online and
offline settings, where the latter can learn GAI hyperparameters over multiple
training trajectories, revealing the converged optimal confidence bound.

Next, we demonstrate that DGAI can be extended to enhance cumulative reward
maximization problems with a given threshold as prior knowledge. Generally,
the criteria for eliminating sub-optimal arms in MAB and bad arms in GAI are
the same when the conventional UCB algorithm is chosen for sampling, as shown
in Sec. 4.5. By employing the DGAI algorithm, we facilitate dynamic and
adaptive criteria to discard sub-optimal arms in a data-driven manner, with
the threshold as prior information guiding the learning and adjustment of the
confidence bound. This reformulation enables us to achieve both the
maximization of cumulative reward and the elimination of sub-optimal arms,
contributing to further empirical performance improvements.

In summary, this paper makes the following contributions:

  1. 1.

We introduce a novel structured bandit problem, linear & non-linear GAI,
closing a gap in the literature on structured bandits.

  2. 2.

We propose DGAI, a novel differentiable algorithm designed to learn adaptive
confidence bounds based on actual reward structures rather than assumptions.

  3. 3.

We provide correctness guarantees establishing that linear DGAI is
$\delta$-PAC.

  4. 4.

We show that DGAI outperforms state-of-the-art baselines on both synthetic and
real-world datasets for GAI problems.

  5. 5.

We demonstrate that incorporating DGAI can improve cumulative reward
maximization problems.

##  2 Related works

###  2.1 Good arm identification

The work by Kano et al. (kano2019good) was the first to formulate the Good Arm
Identification (GAI) problem as a pure-exploration problem in the fixed
confidence setting. In the fixed confidence setting, an acceptance error rate
(confidence) is fixed, and the goal is to minimize the number of arm pulls
required to identify good arms (sample complexity). This work addresses a new
kind of dilemma: the exploration-exploitation dilemma of confidence. Here,
exploration involves the agent pulling arms other than the currently best one
to increase the confidence of whether that arm is good or bad. On the other
hand, exploitation entails the agent pulling the currently best arm to
increase confidence in its goodness. To tackle this problem, the authors
propose three algorithms: a hybrid algorithm for the Dilemma of Confidence
(HDoC), the Lower and Upper Confidence Bounds algorithm for GAI (LUCB-G),
which builds upon the LUCB algorithm for best arm identification
(kalyanakrishnan2012pac), and the Anytime Parameter-free Thresholding
algorithm for GAI (APT-G), based on the APT algorithm for the thresholding
bandit problem (locatelli2016optimal). The lower bound on the sample
complexity for GAI is proven to be $\Omega(\lambda\log\frac{1}{\delta})$, and
HDoC can identify $\lambda$ good arms within
$O(\lambda\log\frac{1}{\delta}+(K-\lambda)\log\log\frac{1}{\delta})$ samples.

###  2.2 Differentiable bandit algorithm

SoftUCB (yang2020differentiable) solves the policy-gradient optimization of
bandit policies via a differentiable bandit algorithm. The authors present a
differentiable UCB-typed linear bandit algorithm that combines EXP3
(auer1995gambling) and Phased Elimination (lattimore2020bandit) to allow the
confidence bound to be learned purely in a data-driven fashion and avoid the
need for reliance on concentration inequalities and assumptions about the
reward distribution’s form. By incorporating a differentiable UCB index, the
learned confidence bound becomes adaptable to the problem’s actual reward
structure. Additionally, the authors proved that the differentiable UCB index
has a regret bound that scales with $d$ and $T$ as
$\mathcal{O}(\beta\sqrt{dT})$, compared to classical UCB-typed algorithms,
where $d$ is the dimension of the input to the linear bandit model, and $T$ is
the total number of sampling rounds.

##  3 Preliminary

In this section, we first formulate the GAI problem in the fixed confidence
setting and show lower bound on sample complexity for GAI. Next we will
introduce the basic framework for GAI algorithm that follows the HDoC
algorithm. We give the notations in Table 1.

$K$ |  Number of arms.
---|---
$m$ |  number of good arms (unknown).
$A$ |  set of arms, where $|A|=K$.
$\xi$ |  Threshold determining whether a arm is good or bad.
$\delta$ |  acceptance error rate
$\mu_{i}$ |  The true mean of $i^{th}$ arm.
$\hat{\mu}_{i,t}$ |  The empirical mean of $i^{th}$ arm at time $t$.
$\tau_{\lambda}$ |  The number of rounds to identify whether $\lambda^{th}$ arm is good or arm.
$\tau_{stop}$ |  Round that agent confirms no good arm left.
$N_{i}(t)$ |  The number of samples of arm $i$ by the end of sampling round $t$.
$U_{i,t}$ |  confidence bound at time $t$.
$\overline{\mu}_{i,t}$ |  upper confidence bound.
$\underline{\mu}_{i,t}$ |  lower confidence bound.
$\|x\|_{M}$ |  $\sqrt{x^{T}Mx}$, $x\in\mathbb{R}^{d}$.
$\Delta_{i}$ |  $|\mu_{i}-\xi|$ gap between thresholds for identification.
$\Delta_{i,j}$ |  $\mu_{i}-\mu_{j}$ gap between arms for sampling.
Table 1: Notations

###  3.1 Good arm identification

Let $K$ be the number of arms, $\xi$ be the threshold, and $\delta$ be the
acceptance error rate. The reward of arm $i$, where $i\in[1,\dots,K]$, follows
the mean $\mu_{i}$, which is initially unknown. We define ”good” arms as those
arms whose means are larger than the given threshold $\xi$. Without loss of
generality, we assume the following order for the means:

$\mu_{1}\geq\mu_{2}\geq\dots\geq\mu_{m}\geq\xi\geq\mu_{m+1}\geq\dots\geq\mu_{K}$
(1)

Here, $m$ represents the total number of good arms, which is also unknown to
the agent.

In each round $t$, the agent selects arm $a(t)$ to pull and receives a reward.
Based on the previous rewards this arm has received, the agent classifies it
as either a good or bad arm. The agent continues this process until all arms
are identified as either good or bad arms, and this stopping time is denoted
as $\tau_{stop}$. The agent’s outputs are denoted as
$\hat{a}_{1},\hat{a}_{2},\dots$, representing the identification of good arms
at rounds $\tau_{1},\tau_{2},\dots$, respectively.

###  3.2 Sample Complexity for GAI

The lower bound on the sample complexity for GAI is given in (kano2019good) in
terms of top-$\lambda$ expectations $\\{\mu_{i}\\}_{i=1}^{\lambda}$ and is
tight up to the logarithmic factor $O(\log\frac{1}{\delta})$.

######  Definition 1 ($(\lambda,\delta)$-PAC).

An algorithm satisfying the following conditions is called ($\lambda$,
$\delta$)-PAC: if there are at least $\lambda$ good arms then
$\mathbb{P}[\\{\hat{m}<\lambda\\}\,\cup\,\bigcup_{i\in\\{\hat{a}_{1},\hat{a}_{2},\dots,\hat{a}_{\lambda}\\}}\\{\mu_{i}<\xi\\}]\leq\delta$
and if there are less than $\lambda$ good arms then
$\mathbb{P}[\hat{m}\geq\lambda]\leq\delta$.

######  Theorem 1.

Under any $(\lambda,\delta)$-PAC algorithm, if there are $m\geq\lambda$ good
arms, then

$\displaystyle\mathbb{E}[\tau_{\lambda}]$
$\displaystyle\geq\left(\sum_{i=1}^{\lambda}\frac{1}{d(\mu_{i},\xi)}\log\frac{1}{2\delta}\right)-\frac{m}{d(\mu_{\lambda},\xi)},$
(2)

where $d(x,y)=x\log(x/y)+(1-x)\log((1-x)/(1-y))$ is the binary relative
entropy, with convention that $d(0,0)=d(1,1)=0$.

#### Hybrid algorithm for the Dilemma of Confidence (HDoC)

HDoC’s sampling strategy is based on the upper confidence bound (UCB)
algorithm for the cumulative regret minimization (auer2002finite), and its
identification criterion (i.e., to output an arm as a good one) is based on
the lower confidence bound (LCB) identification (kalyanakrishnan2012pac). The
upper bound on the sample complexity of HDoC is given in the following
corollary based on Theorem 2 in (kano2019good):

######  Corollary 1.

Let
$\Delta$=$\min\\{\min_{\lambda\in[K-1]}\Delta_{\lambda,\lambda+1}/2,\min_{i\in[K]}\Delta_{i}\\}$.
Then, for any $\lambda\leq m$,

$\displaystyle\mathbb{E}[\tau_{\lambda}]=\mathcal{O}\left(\frac{\lambda\log\frac{1}{\delta}+(K-\lambda)\log\log\frac{1}{\delta}+K\log\frac{K}{\Delta}}{\Delta^{2}}\right)$
(3)
$\displaystyle\mathbb{E}[\tau_{\mathrm{stop}}]=\mathcal{O}\left(\frac{K\log(1/\delta)+K\log(K/\Delta)}{\Delta^{2}}\right)$
(4)

Note that evaluating the error probability in HDoC relies on the union bound
over all rounds $t\in\mathbb{N}$, and its use does not worsen the asymptotic
analysis for $\delta\to 0$. However, the empirical performance can be
considerably improved by avoiding the union bound and using an identification
criterion based on such a bound. The complexity measure for the sampling
strategy is based on the gap between arms $\Delta_{i,j}$, and the complexity
measure for the identification criteria is based on the gap between the
threshold $\Delta_{i}$.

##  4 Algorithm

In this section, we proposed the differential GAI algorithm, DGAI, and proved
DGAI is $\delta$-PAC. Next, we show DGAI in linear form and later extend to
non-linear cases. Last we introduced the differential UCB index and followed
by the training objectives. Pseudo codes of all the algorithms are described
in Algorithm LABEL:alg:ours.

###  4.1 Linear Good Arm Identification

First, we introduce the DGAI in linear bandit matrix forms. Each arm
$i\in\mathcal{A}$ is associated with a known feature vector
$\mathbf{x}_{i}\in\mathbb{R}^{d}$. Linear bandit is equivalent to discrete and
independent arms using one-hot feature vectors. The expected reward of each
arm $\mu_{i}=\mathbf{x}_{i}^{T}\bm{\theta}$ follows a linear relationship over
$\mathbf{x}_{i}$ and an unknown parameter vector $\bm{\theta}$. At each round
$t$, the confidence bound is defined as
$|\hat{\mu}_{i,t}-\mu_{i}|\leq\beta||\mathbf{x}_{i}||_{\mathbf{V}^{-1}_{t}},\
\ \forall i\in\mathcal{A}$ and
$\mathbf{V}_{t}=\sum_{t=1}^{T}\mathbf{x}_{t}\mathbf{x}_{t}^{T}$ is the Gram
matrix up to round $t$. $\beta$ in the confidence bound can be derived based
on different concentration inequalities under the assumption of the stochastic
reward structure. e.g., Hoeffding inequality (hoeffding1994probability), self-
normalized (abbasi2011improved), Azuma Inequality (lattimore2018bandit),
Bernstein inequality (mnih2008empirical). The correctness of the output arms
can be verified by the following theorem, whose proof is given in Appendix A.

######  Theorem 2.

The DGAI algorithm is $\delta$-PAC which outputs a bad arm with probability at
most $\delta$.

###  4.2 Differentiable UCB index

In DGAI, our objective is to learn the magnitude of the confidence bound in a
data-driven manner without relying on prior assumptions about the unknown
reward distribution. The design of our UCB index and associated lemmas
primarily follows (yang2020differentiable). Let the arm with the largest lower
confidence bound at round $t$ as
$i_{*}=\arg\max_{i\in\mathcal{A}}\hat{\mu}_{i,t}-U_{i,t}$. We then define:

$\displaystyle\phi_{i,t}=U_{i,t}+U_{i_{*},t}\text{\ \ and \ \
}\hat{\Delta}_{i,t}=\hat{\mu}_{i_{*},t}-\hat{\mu}_{i,t}$ (5)

and $\hat{\Delta}_{i,t}$ is the estimated reward gap between $i_{*}$ and $i$.
The UCB index $S_{i,t}$ is defined as

$S_{i,t}=\beta\phi_{i,t}-\hat{\Delta}_{i,t}\,$ (6)

######  Lemma 1.

The UCB index $S$ provides the following information: $S_{i,t}<0$, i.e.,
$\mu_{*}-\mu_{i}>0$, implies arm $i$ is a sub-optimal arm. $S_{i,t}\geq
S_{j,t}\geq 0$, i.e., $\hat{\mu}_{i,t}+U_{i,t}\geq\hat{\mu}_{j,t}+U_{j,t}$,
implies arm i has larger upper confidence bound.

For each round $t\in[T]$, the likelihood of selecting arm $i$ is:

$p_{i,t}=\frac{\exp(\gamma_{t}S_{i,t})}{\sum_{j=1}^{K}\exp(\gamma_{t}S_{j,t})}~{},~{}\gamma_{t}=\frac{\log\left(\frac{\delta|\mathcal{L}_{t}|}{1-\delta}\right)}{\tilde{S}_{\text{max},t}}$
(7)

where $\gamma_{t}>0$ is the coldness parameter modulating the concentration of
the policy distribution, as elucidated in Lemma 2. The set $\mathcal{L}_{t}$
consists of suboptimal arms (i.e., $S_{t}<0$), while $\mathcal{U}_{t}$ holds
non-suboptimal arms (i.e., $S_{t}\geq 0$). We have
$\mathcal{U}_{t}\cup\mathcal{L}_{t}=\mathcal{A}$ and
$\mathcal{U}_{t}\cap\mathcal{L}_{t}=\emptyset$.
$\tilde{S}_{\text{max},t}=\max_{i\in\mathcal{U}_{t}}S_{i,t}$,
$|\mathcal{L}_{t}|$ denotes the cardinality of $\mathcal{L}_{t}$ and $\delta$
is a probability hyper-parameter.

######  Lemma 2.

At any round $t\in[T]$, for any $\delta\in(0,1)$, setting
$\gamma_{t}\geq\log(\frac{\delta|\mathcal{L}_{t}|}{1-\delta})/\tilde{S}_{\text{max},t}$
guarantees that
$p_{\mathcal{U}_{t}}=\sum_{i\in\mathcal{U}_{t}}p_{i,t}\geq\delta$ and
$p_{\mathcal{L}_{t}}=\sum_{i\in\mathcal{L}_{t}}p_{i,t}<1-\delta$ such that
sub-optimal arms ($i\in L_{t}$) will be selected with near zero probability
when $\delta\approx 1$.

It is worth noting that DGAI can readily be extended to non-linear versions.
Appendix B provides details on these adaptations. Remark LABEL:app:rmk:kernel
describes the adaptation to the kernelized DGAI variant by incorporating
KernelUCB (valko2013finite), and Remark LABEL:app:rmk:neural showcases the
transition to a neural network version of DGAI by incorporating NeuralUCB
(zhou2020neural). A more comprehensive experiment and theoretical analysis are
reserved for future studies.

###  4.3 Training Objectives

The optimization objective for sampling strategy and identification criteria
are independent and learned separately. The optimization objective of the
sampling strategy is to maximize the cumulative reward with constraints shown
as the following differentiable function over $\beta$:

$\begin{split}\max_{\beta}\sum_{t=1}^{T}\mathbb{E}[y_{t}]=\max_{\beta}\sum_{t=1}^{T}\sum_{i=1}^{K}p_{i,t}\mu_{i}~{},\
\ \\\ s.t.\
|\mu_{i}-\hat{\mu}_{i,t}|\lessapprox\beta||\mathbf{x}_{i}||_{\mathbf{V}^{-1}_{t}},\
\ \forall i\in\mathcal{A},t\in[T]\end{split}$ (8)

The addition constraint force to tighten the upper confidence bound as low as
possible and ensures that $\beta$ is minimal but
$\beta||\mathbf{x}_{i}||_{\mathbf{V}^{-1}_{t}}$ is still indeed an actual
upper confidence bound at any round $t\in[T]$ for any arm $i\in\mathcal{A}$.
Applying the Lagrange multipliers gives the objective with $\eta$ as the
hyperparameter for balancing the constraints:

$\begin{split}\max_{\beta}\sum_{t=1}^{T}\sum_{i=1}^{K}p_{i,t}\mu_{i}-\eta_{1}C-\eta_{2}|C|~{},\
\ s.t.\ \ \eta>0\\\ \text{\ \ where \ \
}C=|\mu_{i}-\hat{\mu}_{i,t}|-\beta||\mathbf{x}_{i}||_{\mathbf{V}^{-1}_{t}}\end{split}$
(9)

The objective for the identification criteria is to maximize the Exploit score
which is defined as follow:

$\sum_{i=1}^{K}(T-t_{i})*(\mu_{i}-\xi)$ (10)

where $t_{i}$ is the round that arm $i$ is identified. The exploit score gives
a positive reward when the arm is identified correctly and increases if the
arm is identified earlier. This conditional reward is formed as an indicator
function with the lower confidence bound in the objective over $\alpha$:

$\begin{split}\max_{\alpha}\sum_{t=1}^{T}\sum_{i=1}^{K}\mathbbm{I}[\mu_{i,t}-\alpha\|x_{i}\|_{V_{t}^{-1}}>\xi]*(\mu_{i,t}-\xi)\\\
s.t.|\mu_{i,t}-\xi|\gtrapprox\alpha\|x_{i}\|_{V_{t}^{-1}},\ \ \forall
i\in\mathcal{A},t\in[T]\end{split}$ (11)

The indicator function is relaxed and smoothen with sigmoid function and again
the Lagrange multipliers yields the following objective:

$\max_{\beta}\sum_{t=1}^{T}\sum_{i=1}^{K}\sigma((\mu_{i,t}-\alpha\|x_{i}\|_{V_{t}^{-1}}-\xi)*M)*(\mu_{i,t}-\xi)\\\
-\eta_{1}D-\eta_{2}|D|~{},\ \ s.t.\ \ \eta>0~{},\ \ \text{where~{}}\\\
D=\alpha||\mathbf{x}_{i}||_{\mathbf{V}^{-1}_{t}}-|\mu_{i}-\hat{\mu}_{i,t}|~{},\
\ \sigma(x)=\frac{1}{1+e^{-x}}$ (12)

$\eta$ is for balancing the constraints and $M$ is for tuning the sharpness of
the indicator function.

###  4.4 Training Settings

We trained our differentiable algorithm in both offline and online settings.
The online setting is more suitable for bandit problems, while the offline
setting shows the converged optimal confidence bound and would be the upper
bound of the online setting.

#### Offline setting

In the offline setting, we will run multiple epochs of $T$-rounds training
trajectories with the identical arm set $\mathcal{A}$ to train $\alpha$ and
$\beta$. First, $\alpha_{0}$ and $\beta_{0}$ are initialized. Both parameters
are used to run one entire $T$-rounds training trajectories. After the
trajectory, $\alpha_{0}$ and $\beta_{0}$ is updated with Equation 9 and
Equation 12 respectively. After $N$ epochs of trajectories, we have our
finalized confidence bound with learned $\alpha_{N}$ and $\beta_{N}$. DGAI for
learning the identification criterion in the offline setting is shown in
Algorithm LABEL:alg:ours. The time complexity for training in the offline
setting is $\mathcal{O}(NKT)$.

#### Online setting

In this setting, $\alpha$ and $\beta$ are updated online in one single
$T$-rounds trajectory. First, $\alpha_{0}$ and $\beta_{0}$ are initialized and
for every $b\in[T]$ round, $\alpha$ and $\beta$ will perform batch update
where $b$ is the batch size.

We adopted policy gradient methods, the same as non-episodic reinforcement
learning problems (sutton2018reinforcement) to update the parameters in an
online fashion. Instead of maximizing the cumulative reward until the end of
the trajectory, we update the observed cumulative reward up to round $t$ and
bootstrapped future reward under the current policy
$\bm{\pi}_{t}=[p_{1,t},p_{2,t},...,p_{K,t}]$. The online objective for
sampling and identification are as follows, where $R$ is the reward feedback
we obtain in the original objective function Equation 8 and Equation 10:

$\begin{split}\left(\sum_{s=1}^{t}\sum_{i=1}^{K}R_{i,s}+(T-t)\sum_{i=1}^{K}R_{i,t}\right)/T\end{split}$
(13)

The time complexity for training in the online setting is ${O}(KT)$.

###  4.5 GAI with Cumulative Reward Maximization

Classical bandit problems focus on the cumulative reward maximization problem,
where the objective is to maximize the sum of the samples collected by the
algorithm up to time $T$ (auer1995gambling). In this present section, we delve
into transforming the ”exploration-exploitation dilemma of confidence” that
GAI addresses back into a cumulative reward maximization problem.
Specifically, we tackle the cumulative reward maximization by incorporating
the additional information of a threshold as prior knowledge of the arm set
$A$ allowing us to utilize this threshold to optimize and adjust the
confidence bound to maximize cumulative reward. In identical settings,
(pmlr-v48-locatelli16) utilized the APT algorithm and provided a regret bound
with optimal threshold $\xi=\mu_{*}$ as illustrated in theorem 3. The bandit
policy performs more efficiently in terms of regret if the threshold $\xi$ is
closer to $\mu_{*}$ and achieves optimality when $\xi=\mu_{*}$. However, since
$\mu_{*}$ is always unknown in reality, the proper threshold must often be
estimated beforehand, relying on prior experience.

######  Theorem 3.

Let $K>0,R>0$ and $T\geq 2K$ and consider a problem where the distribution of
the arms is R-sub-Gaussian. Run with parameters $\xi=\mu^{*}$ such that

$\displaystyle Regret(T)=T\mu^{*}-\mathbb{E}\sum_{t\leq T}\mu_{t}$
$\displaystyle\leq\inf_{\delta\geq
1}\Big{[}\sum_{k\not=k^{*}}\frac{4R^{2}\log(T)\delta}{\mu^{*}-\mu_{i}}$
$\displaystyle+(\mu^{*}-\mu_{i})(1+\frac{K}{T^{2\delta-2}})\Big{]}.$

In what follows, we incorporate DGAI to learn a dynamic and adaptive threshold
rather than making random assumptions about the threshold. The upper
confidence bound used for sampling serves as an identification criterion; for
instance, $|\hat{\mu}_{i,t}-\mu_{i}|\leq U$ implies arm $i$ is good
$\mu_{i}\geq\xi$ when $\hat{\mu}_{i,T_{i}(t)}-U\geq\xi$. We have implemented
two distinct confidence bounds to aid in solving the cumulative reward
maximization problem. The first confidence bound, drawn from Equation 8,
directs the sampling strategy to select arms based on criteria that softly
eliminate sub-optimal options. In contrast, the other confidence bound,
derived from Equation 11, selects arms according to criteria that reject those
falling below the threshold. The objective function to solve the multi-arm
bandit with threshold re-formulates both Equation 8 and Equation 11 into a
single, unified objective function:

$\begin{split}\max_{\alpha,\beta}\sum_{t=1}^{T}\mathbb{E}[y_{t}]=\max_{\alpha,\beta}\sum_{t=1}^{T}\sum_{i=1}^{K}p_{i,t}\mu_{i}~{},\\\
p_{i,t}=\frac{exp(S_{i,t}*I^{\xi}_{i,t})}{\sum
exp(S_{j,t}*I^{\xi}_{i,t})}~{},~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}\\\
I_{i,t}=\sigma((\mu_{i,t}-\alpha\|x_{i}\|_{V_{t}^{-1}}-\xi)*M)\end{split}$
(14)

##  5 Experiment

Here we would like to evaluate DGAI in (1) improvements in GAI problem (2)
improvements in cumulative reward maximization with given threshold.

###  5.1 Experiment Setup

#### Datasets

We use two synthetic and two real-world datasets whose details are listed in
Appendix Table 3. We generated one small and one large synthetic dataset with
a number of arms = 50 and 1000, respectively. The reward of each arm is
generated from the uniform distribution $\in[0.5005,0.49975]$, and the
threshold $\xi$ is set to 0.5. For real world datasets, we used two public
recommendation datasets where each item can be treated as an item: OpenBandit
(saito2020large) and MovieLens (harper2015movielens). We use the method
similar to (da2022fast; tsai2024lil), to transform multi–class dataset to
bandit dataset In the MovieLens dataset, the rating of each movie is treated
as a means to sample the stochastic reward of each arm. The threshold for a
good arm is set to the 95 percentile reward of all arms in the entire dataset.

| synth | synth | MovieLens | Openbandit
---|---|---|---|---
| small | large |  |
HDoC | 1.2e6 | 6.5e4 | 1.8e6 | 4e5
LUCB-G | 1.4e6 | 1.1e5 | 1.7e6 | 4e5
APT-G | 1.8e6 | 8.8e4 | 6e6 | 4e5
TT-TS | 1.7e6 | 1.1e5 | 4.1e7 | 4e5
SoftUCB-G | 2.7e6 | 1.2e5 | 3.8e7 | 1e6
DGAI-online | 3.1e6 | 1.7e5 | 4.2e7 | 1.3e6
DGAI-offline | 4.1e6 | 2.5e5 | 6.1e7 | 1.5e6
Table 2: Exploit score of DGAI and the baselines. The result shows that DGAI
outperforms the baselines in all cases. Results exhibit a standard deviation
below 10%, averaged over 10 repetitions per run.

#### Parameter setting

For the two optimized parameters $\alpha,\beta$, we set initial values to 0.
We used stochastic gradient descent as the optimizer and 1e-1 as the learning
rate. $\eta_{1}$ and $eta_{2}$ in Equation 9 and Equation 12 are both set to
1e-3. It is worth mentioning that fine-tuning their parameters can improve the
convergence of the training, and a bad hyperparameter may not converge to
optimal in one single training trajectory.

###  5.2 Good arm identification

#### Evaluation Metric

The metric for evaluating the improvement of good arm identification problem
is by the Exploit score $\sum_{i=1}^{K}(T-t_{i})*(\mu_{i}-\xi)$ which is also
used in Equation 10. The increase in the cumulative score from this metric
indicates that the agent outputs good arms faster.It not only evaluates the
overall performance (when the agent finishes identifying all arms) but also
considers each individual arm (when each arm is identified).

#### Baseline

We compare with three SOTA algorithms: HDoC, APT-G, LUCB-G introduced in
(kano2019good) for solving the good arm identification problem. We also
compare with top-two Thompson sampling (TT-TS) (russo2016simple), which is a
Bayesian method designed for arm identification. We also added SoftUCB-G,
which runs the exact HDoC algorithm, except that we replaced the conventional
upper confidence bound with SoftUCB (yang2020differentiable) such that it also
learns the upper bound with differential data-driven methods. The results of
all baselines are averaged with ten different random seeds.

#### Sampling Bound

  1. 1.

HDoc : Pull arm $\hat{a}^{*}=\operatorname*{arg\,max}_{i\in
A}\overline{\mu}_{i,t}$ , where
$\overline{\mu}_{i,t}=\hat{\mu}_{i,t}+\sqrt{\frac{\log{t}}{2N_{i}(t)}}$

  2. 2.

LUCB-G : Pull arm $\hat{a}^{*}=\operatorname*{arg\,max}_{i\in
A}\overline{\mu}_{i,t}$ , where
$\overline{\mu}_{i,t}=\hat{\mu}_{i,t}+\sqrt{\frac{\log{4KN_{i}^{2}{t}/\delta}}{2N_{i}(t)}}$

  3. 3.

APT-G : Pull arm $\hat{a}^{*}=\operatorname*{arg\,max}_{i\in A}\beta_{i,t}$ ,
where
$\beta_{i,t}=\hat{\mu}_{i,t}+\sqrt{N_{i}(t)}|\xi-\hat{\mu}_{i,t}|$

  4. 4.

TT-TS : Pull arm $\hat{a}^{*}=\operatorname*{arg\,max}_{a\not\in
I_{m}}P_{i,t}$, where $I=\operatorname*{arg\,max}_{a\in A}P_{i,t}$ and $P$ is
the posterior probability of arm $i$ is optimal.

  5. 5.

SoftUCB-G : Pull arm $\hat{a}^{*}=\operatorname*{arg\,max}_{i\in A}P_{i,t}$

  6. 6.

DGAI : Pull arm $\hat{a}^{*}=\operatorname*{arg\,max}_{i\in A}P_{i,t}$

#### Identification Bound

All baseline have identical identification bound:
$\sqrt{\frac{\log{4KN_{i}^{2}(t)/\delta}}{2N_{i}(t)}}$ and DGAI is
$\alpha||\mathbf{x}_{i}||_{\mathbf{V}^{-1}_{t}}$.

###  5.3 Cumulative Reward Maximization

#### Evaluation Metric

We use the cumulative reward (same as the classic multi-arm bandit problem) to
evaluate the results. Baseline We compare our methods to classical UCB,
Thompson sampling and SoftUCB.

###  5.4 Results

The experiment results for DGAI are shown in Table. 2. The complete
performance through the training epochs for GAI is shown in Fig.
LABEL:fig:exp1. The cumulative exploit score for DGAI outperforms other
baselines eventually for both online and offline learning settings. Fig.
LABEL:fig:learning-curve depicts the learning curves of $\alpha$ and $\beta$.
It shows that the confidence bound in both online and offline settings
converges in a similar trend, but the parameters converge must faster in the
offline setting since it updates the parameters with multiple complete
training trajectories. We can observe that the confidence bound w.r.t
$\alpha,\beta$ continually adjusts during the single round training trajectory
and, at last, converge to a scale similar to in the offline setting. In Fig.
LABEL:fig:bound-compare, we plot the confidence bound for the best arm in the
arm set $\mathcal{A}$ during the training trajectory and compare the
identification bounds between DGAI and the baselines. We can observe that as
the training epochs go on, the learned confidence bound w.r.t $\alpha,\beta$
continually converges to optimal, and thus the cumulative exploit score
continues to increase. The results suggested that the learned confidence bound
with DGAI is significantly tighter than the original union confidence bound
consistently in all cases. The experiment results solving the cumulative
reward maximization problem with threshold are shown in Fig. LABEL:fig:exp2.
We can observe that DGAI outperforms other baselines on all datasets with the
help of the additional identification criteria based on the threshold. Note
that the curves in all figures are smoothed to ignore the vibration.

##  6 Conclusion

In this research, we have presented a novel differentiable algorithm for the
good arm identification (GAI) problem within the framework of a structured
bandit setting. Our algorithm, DGAI, has the unique capability to adapt the
confidence bound, learning through gradient ascent, which sets it apart from
existing methods. We showed that DGAI can be applied in both online/offline
settings and could be further extended to non-linear settings. Furthermore, we
also show that DGAI improves cumulative reward maximization problems when a
threshold is provided as prior knowledge. The experiments on all datasets have
shown order-of-magnitude improvements compared to other baselines in terms of
rewards obtained. A more comprehensive experiment for the generalized linear
case and theoretical analysis for sample complexity is reserved for future
studies.

##  Appendix A Proof of Theorem 2

###### Proof of Theorem 2.

We show that DGAI is $(\lambda,\delta)$-PAC for arbitrary $\lambda\in[K]$.

First we consider the case that there are more than or equal to $\lambda$ good
arms and show

$\displaystyle\mathbb{P}\left[\\{\hat{m}<\lambda\\}\,\cup\,\bigcup_{i\in\\{\hat{a}_{1},\hat{a}_{2},\ldots,\hat{a}_{\lambda}\\}}\\{\mu_{i}<\xi\\}\right]\leq\delta.$
(15)

Since we are now considering the case $m\geq\lambda$, the event
$\\{\hat{m}<\lambda\\}$ implies that at least one good arm $j\in[m]$ is
regarded as a bad arm, that is, $\\{\underline{\mu}_{j,n}\leq\xi\\}$ occurs
for some $j\in[m]$ and $n\in\mathbb{N}$. Thus we have

$\displaystyle\mathbb{P}[\hat{m}<\lambda]$
$\displaystyle\leq\sum_{j\in[m]}\mathbb{P}\left[\bigcup_{n\in\mathbb{N}}\\{\overline{\mu}_{j,n}<\xi\\}\right]$
$\displaystyle\leq\sum_{j\in[m]}p_{i,t}$ $\displaystyle\leq\delta\quad\mbox{by
Lemma \ref{lemma:gamma_lemma}}$ (16)

On the other hand, since the event
$\bigcup_{i\in\\{\hat{a}_{1},\hat{a}_{2},\ldots,\hat{a}_{\lambda}\\}}\\{\mu_{i}<\xi\\}$
implies that $j\in\\{\hat{a}_{i}\\}_{i=1}^{\lambda}$ for some bad arm
$j\in[K]\setminus[m]$, we have

$\displaystyle\mathbb{P}\left[\bigcup_{i\in\\{\hat{a}_{1},\hat{a}_{2},\ldots,\hat{a}_{\lambda}\\}}\\{\mu_{i}<\xi\\}\right]$
$\displaystyle\leq\sum_{j\in[K]\setminus[m]}\mathbb{P}[j\in\\{\hat{a}_{i}\\}_{i=1}^{\lambda}]$
$\displaystyle\leq\sum_{j\in[K]\setminus[m]}\mathbb{P}\left[\bigcup_{n\in\mathbb{N}}\\{\underline{\mu}_{j,n}\geq\xi\\}\right]$
$\displaystyle\leq\delta$ (17)

in the same way as (16). We obtain (15) by putting (16) and (17) together.

Next we consider the case that the number of good arms $m$ is less than
$\lambda$ and show

$\displaystyle\mathbb{P}[\hat{m}\geq\lambda]\leq\delta\,.$ (18)

Since there are at most $m<\lambda$ good arms, the event
$\\{\hat{m}\geq\lambda\\}$ implies that
$j\in\\{\hat{a}_{i}\\}_{i=1}^{\lambda}$ for some $j\in[K]\setminus[\lambda]$.
Thus, in the same way as (17) we have

$\displaystyle\mathbb{P}\left[\hat{m}\geq\lambda\right]$
$\displaystyle\leq\sum_{j\in[K]\setminus[m]}\mathbb{P}[j\in\\{\hat{a}_{i}\\}_{i=1}^{\lambda}]$
$\displaystyle\leq\delta\,,$

which proves (18). ∎

##  Appendix B Extension to non-linear DGAI

In this section, we will show how DGAI can be easily extended to non-linear
cases, specifically, kernelised and neural network versions. Note that further
experiments and theoretical analysis will be left for future work.

######  Remark 1.

We can easily extend DGAI in linear form to non-linear kernelised version. In
the following, we applied the kernel trick (cristianini2004kernel) and the
kernelised version of the Mahalanobis (haasdonk2010classification) similar to
the settings in KernelUCB (valko2013finite). The derivation is straightforward
and we provide it for convenience and to introduce the notation. Kernel
methods assume that there exists a mapping $\phi:\mathbb{R}^{d}\rightarrow H$
that maps the data to a (possibly infinite dimensional) Hilbert space in which
a linear relationship can be observed. We call $\mathbb{R}^{d}$ the primal
space and $H$ the associated reproducing kernel Hilbert space (RKHS). We use
matrix notation to denote the inner product of two elements $h,h^{\prime}\in
H$, i.e. $h^{T}h^{\prime}:=\langle h,h^{\prime}\rangle_{H}$ and
$\|h\|=\sqrt{\langle h,h\rangle_{H}}$ to denote the RKHS norm. From the
mapping $\phi$ we have the kernel function, defined by:
$k(x,x^{\prime}):=\phi(x)^{T}\phi(x^{\prime})$, $\forall
x,x^{\prime}\in\mathbb{R}^{d}$, and the kernel matrix of a data set
$\\{x_{1},...,x_{t}\\}\subset\mathbb{R}^{d}$ given by
$K_{t}:=\\{k(x_{i},x_{j})\\}_{i,j\leq t}$. Following the above we obtain:

$\hat{\mu}_{i,t}=k_{x,t}^{T}(K_{t}+\gamma I)^{-1}y_{t}.$ (19)

$\hat{U}_{i,t}:=\gamma^{-1/2}\sqrt{k(x_{i,t},x_{i,t})-k_{x,t}^{T}(K_{t}+\gamma
I)^{-1}k_{x,t}}.$ (20)

where $k_{x,t}:=\Phi_{t}\phi(x)=[k(x,x_{1}),...,k(x,x_{t-1})]^{T}$ and for
some regularization parameters $\gamma>0$. Plugging this into equation 5 and
6, we can construct the kernelised version of differentiable UCB index.

######  Remark 2.

We can easily extend DGAI in linear form to non-linear neural network version.
In the following, we use a neural network $f(x;\theta)$ to predict the reward
of context, and upper confidence bounds computed from the network to guide
exploration similar to NeuralUCB (zhou2020neural). A simplified version of
NeuralUCB where only the first-order Taylor approximation of the neural
network around the initialized parameter is updated through online ridge
regression can be seen as KernelUCB specialized to the Neural Tangent Kernel
(jacot2018neural), or LinUCB (li2010contextual) with Neural Tangent Random
Features (cao2019generalization). Following the above we obtain: Denote the
gradient of the neural network function by
$g(x;\theta)=\nabla_{\theta}f(x;\theta)\in\mathbb{R}^{p}$,
$Z=\gamma\textbf{I}$ and network width $M$. We obtain:

$\mu_{i}=f(x_{i,t};\theta_{t-1})$ (21)
$\hat{U}_{i,t}=\beta_{t-1}\sqrt{g(x_{t,a};\theta_{t-1})^{\top}Z_{t-1}^{-1}g(x_{t,a};\theta_{t-1})/M}$
(22)

where $g(x;\theta)=\nabla_{\theta}f(x;\theta)\in\mathbb{R}^{p}$ is the
gradient of the neural network, $Z=\gamma\textbf{I}$ and $M$ is the network
width. Plugging this into equation 5 and 6, we can construct the neural
network version of differentiable UCB index.

|  $N$ arms |  $T$ trajectory |  $\xi$ threshold
---|---|---|---
synth small | 50 | 1e6 | 0.5
synth large | 1000 | 1e6 | 0.5
MovieLens | 9527 | 1e5 | 0.071
Openbandit | 80 | 107 | 0.005
Table 3: Details of the datasets.