#  BSODA: A Bipartite Scalable Framework for Online Disease Diagnosis

Weijie He  Department of Computer Science and Technology & Institute of
Artificial Intelligence & BNRist, Tsinghua UniversityBeijingChina
<EMAIL_ADDRESS>,  Xiaohao Mao  Department of Computer Science and
Technology & Institute of Artificial Intelligence & BNRist, Tsinghua
UniversityBeijingChina<EMAIL_ADDRESS>,  Chao Ma  Department of
Engineering, University of CambridgeMicrosoft Research
CambridgeCambridgeUnited Kingdom<EMAIL_ADDRESS>,  Yu Huang  Tencent Jarvis
LabShenzhenGuangdongChina<EMAIL_ADDRESS>,  José Miguel Hernández-
Lobato  Department of Engineering, University of CambridgeCambridgeUnited
Kingdom<EMAIL_ADDRESS>and  Ting Chen  Department of Computer Science and
Technology & Institute of Artificial Intelligence & BNRist, Tsinghua
UniversityBeijingChina<EMAIL_ADDRESS>
(2022)

###### Abstract.

A growing number of people are seeking healthcare advice online. Usually, they
diagnose their medical conditions based on the symptoms they are experiencing,
which is also known as self-diagnosis. From the machine learning perspective,
online disease diagnosis is a sequential feature (symptom) selection and
classification problem. Reinforcement learning (RL) methods are the standard
approaches to this type of tasks. Generally, they perform well when the
feature space is small, but frequently become inefficient in tasks with a
large number of features, such as the self-diagnosis. To address the
challenge, we propose a non-RL Bipartite Scalable framework for Online Disease
diAgnosis, called BSODA. BSODA is composed of two cooperative branches that
handle symptom-inquiry and disease-diagnosis, respectively. The inquiry branch
determines which symptom to collect next by an information-theoretic reward.
We employ a Product-of-Experts encoder to significantly improve the handling
of partial observations of a large number of features. Besides, we propose
several approximation methods to substantially reduce the computational cost
of the reward to a level that is acceptable for online services. Additionally,
we leverage the diagnosis model to estimate the reward more precisely. For the
diagnosis branch, we use a knowledge-guided self-attention model to perform
predictions. In particular, BSODA determines when to stop inquiry and output
predictions using both the inquiry and diagnosis models. We demonstrate that
BSODA outperforms the state-of-the-art methods on several public datasets.
Moreover, we propose a novel evaluation method to test the transferability of
symptom checking methods from synthetic to real-world tasks. Compared to
existing RL baselines, BSODA is more effectively scalable to large search
spaces.

online disease diagnosis, self-diagnosis, symptom checking

††journalyear: 2022††copyright: rightsretained††conference: Proceedings of the
ACM Web Conference 2022; April 25–29, 2022; Virtual Event, Lyon,
France.††booktitle: Proceedings of the ACM Web Conference 2022 (WWW ’22),
April 25–29, 2022, Virtual Event, Lyon, France††doi:
10.1145/3485447.3512123††isbn: 978-1-4503-9096-5/22/04††ccs: Applied computing
Health informatics††ccs: Information systems Web applications

##  1\. Introduction

The explosive development of the mobile web has increased the feasibility and
potential of access to online disease diagnosis. A growing number of adults
are now attempting to gain a better understanding of their health conditions
through online services before visiting a doctor. According to a national
survey conducted in the United States (Fox and Duggan, 2013), 35% of American
adults regularly use the internet to self-diagnose. Numerous individuals have
long relied on search engines to access online health information (Fox and
Duggan, 2013). However, current search technology is often ineffective for
self-diagnosis (Zuccon et al., 2015). The fact that the majority of users lack
adequate medical knowledge (Keselman et al., 2008) has an effect on the
quality of health queries (Hu and Haake, 2010; Lopes and Ribeiro, 2013;
Puspitasari, 2017).

In recent years, many symptom checking tools have been developed as an
alternative to online self-diagnosis (Semigran et al., 2015), such as the
well-known symptom checkers from Mayo Clinic (may, 2015) and WebMD (web,
2020). A symptom checker simulates the disease diagnosis process which Ledley
and Lusted (1959) summarized into three steps: 1) A patient presents a self-
report of initial symptoms; 2) A doctor inquires the patient with a series of
questions regarding additional symptoms, history, and other related
information; and 3) The doctor makes a final diagnosis. When it comes to
online self-diagnosis, the questions would concern patients’ symptoms.
Compared to search engines, symptom checkers have the advantage of not
requiring users to formulate appropriate health queries, which they are
incapable of (Zeng et al., 2004; Zuccon et al., 2015). Besides, symptom
checkers are more accurate in disease diagnosis and require less effort and
time.

Although self-diagnosis through a symptom checker is a substantial advance
over patient-directed internet searches, it should be highlighted that symptom
checkers are not a substitute for medical investigation by a doctor. Physical
examinations and diagnostic tests such as X-rays and laboratory tests, which
are critical in the diagnosis of many diseases, could not be performed by
symptom checkers. However, there is indeed a genuine need in the field of
healthcare to develop a symptom checker that is user-friendly, efficient, and
accurate (Cross et al., 2021).

From the healthcare perspective, the primary goal of symptom checking is
diagnosis accuracy. A symptom checker acquires sequentially a number of
symptom values in order to make an accurate diagnosis. Naturally, more
information about possible symptoms could result in a more accurate diagnosis;
for example, most existing questionnaire systems acquire a large number of
symptom values from patients through exhaustive inquiries (Shim et al., 2018;
Lewenberg et al., 2017). Unfortunately, both patients and experts find it
inefficient, time-consuming, and unbearable. The actual goal of symptom
checking is to minimize the number of inquiries while achieving a high
diagnostic accuracy.

From the machine learning perspective, symptom checking can be viewed as a
cost-sensitive sequential feature selection and classification task. Acquiring
a symptom’s value can be viewed as selecting a feature which comes with a
cost. Reinforcement learning (RL) methods have shown good performance in this
type of tasks with a small set of features (Tang et al., 2016; Kachuee et al.,
2018b; Kao et al., 2018; Peng et al., 2018; Xia et al., 2020; Liao et al.,
2020; Lin et al., 2020; Liu et al., 2021). However, they often suffer in tasks
with large feature spaces, because they struggle to train a model. EDDI (Ma et
al., 2019), a framework for instance-based active feature acquisition,
proposes an approach to choose the next feature by maximizing a defined
information reward over all features. It uses a Partial Variational
AutoEncoder (VAE) (Kingma and Welling, 2013) to handle data with missing
values. Such a model can deal with the situation in which an agent only knows
partial information of patients when inquiring. Although promising, EDDI is
limited by its high computational cost dealing with a large number of
symptoms.

To address the aforementioned issues, we propose BSODA, a Bipartite Scalable
framework for Online Disease diAgnosis. Different from prior methods, BSODA is
a novel non-RL approach to automatic symptom-inquiry and disease-diagnosis. RL
methods that perform both symptom-inquiry and disease-diagnosis simultaneously
must deal with a large heterogeneous action space that poses a challenge to
learning with efficiency and effectiveness (Liu et al., 2021). Yet BSODA
divides the task into two cooperative branches: inquiry and diagnosis. The
inquiry branch uses a generative model to produce an inquiry, while the
diagnosis branch uses a supervised classification model to predict the
disease. BSODA leverages the characteristics of these two models so that they
cooperate in decision-making.

Specifically, the inquiry branch determines which symptom to collect next
based on the information-theoretic reward function that was used in EDDI (Ma
et al., 2019). Instead of using the Partial VAE, BSODA employs a Product-of-
Experts encoder to better handle partial observations of symptoms, which has
shown good performance with missing multimodal data (Wu and Goodman, 2018).
Besides, to estimate the reward more precisely, a novel two-step sampling
strategy is developed to approximate the joint posterior by taking advantage
of the predictive probability distributions from the diagnosis model. Special
strategies are designed to accelerate the computation of BSODA.

For the diagnosis branch, BSODA returns predictive disease distributions by a
knowledge-guided self-attention model that has been shown to be efficient at
learning the graphical structure of electronic health record (Choi et al.,
2020). BSODA leverages this model to embed the relationships between diseases
and symptoms to incorporate prior knowledge. In particular, BSODA samples
unobserved features from the generative model to model the uncertainty caused
by partial observations of features, which assists the diagnosis model in
making predictions of the diagnosis results and determining when to
statistically terminate the inquiry process.

Unlike most previous RL baselines that have been proved to be effective
exclusively on synthetic or real-world datasets, BSODA performs well on both
types of evaluations. Synthetic datasets were developed for model evaluation
because obtaining a large amount of real-world medical data for training is
difficult. However, there is a significant gap between existing synthetic and
real-world datasets. Good performance on synthetic datasets is not a guarantee
of transferability to the real world. Hence we propose a novel approach to
evaluating the model transferability using both synthetic and real-world
tasks. Synthetic data consisting of thousands of features are used in model
training to assess a model’s ability to cope with large search spaces, while
real-world data is only used for testing to assess the transferability of the
trained model.

In summary, our main contributions are highlighted as: 1) We propose BSODA, a
non-RL bipartite framework for online disease diagnosis. It is scalable to
large feature spaces at a low computational cost, where RL methods typically
struggle. It demonstrates a high degree of generalizability across a variety
of settings; 2) Technically, we design the models and workflow for BSODA in
order to facilitate collaboration between the symptom-inquiry and disease-
diagnosis branches. We also develop several speedup techniques to make BSODA
practical for online services; and 3) We propose a novel evaluation method and
new datasets to assess the performance of symptom checking methods on large
search spaces and transferability to the real-world tasks.

##  2\. Related Work

There are a significant amount of works that deal with self-diagnosis. One
large family of Bayesian inference and tree-based methods (Kononenko, 1993,
2001; Xu et al., 2013; Kohavi, 1996) uses entropy functions to select symptoms
based on the theory of information gain. For instance, Nan et al. (2015, 2016)
and Zakim et al. (2008) proposed to address the cost of feature acquisition by
using decision tree and random forest methods. Hayashi (1991) attempted to
extract rule-based representations from medical data and human knowledge to
perform diagnosis. Early et al. (2016) and Kachuee et al. (2018a) suggested
performing a sensitivity analysis on trained predictors to measure the
importance of each feature in a given context. Wang et al. (2021) developed a
hierarchical lookup framework based on the symptom embedding where graph
representation learning is tailored towards disease diagnosis. Due to the
intractable nature of global maximization of information gain or global
sensitivity, these methods often employ greedy or approximation algorithms
that result in low accuracy.

Recently, Janisch et al. (2019, 2020) showed that reinforcement learning (RL)
methods outperformed tree-based methods in the task of sequential feature
acquisition with a specific cost function. In particular, Tang et al. (2016)
first formulated the inquiry and diagnosis process as a Markov decision
process and then used RL to perform symptom checking in a simulated
environment. Kao et al. (2018) and Peng et al. (2018) demonstrated that
competitive results could be obtained even in medium search spaces. Kachuee et
al. (2018b) proposed a method based on deep Q-learning with variations of
model uncertainty as the reward function. Xia et al. (2020) implemented an RL
agent for symptom checking using Generative Adversarial Network (GAN) and
policy gradient method, which performed well on two public datasets containing
instances of four and five diseases, respectively. However, such a small
number of diseases is impractical in most real-world situations. Lin et al.
(2020) proposed to combine an inquiry branch of a Q network with a diagnosis
branch that involves a generative sampler. It was only effective on the two
small datasets mentioned previously. Liao et al. (2020) proposed a
hierarchical RL framework with a master for worker appointment, several
workers for symptom inquiry, and a separate worker for screening diseases. Liu
et al. (2021) improved it by using a pretraining strategy to overcome
difficulties in convergence. However, their maximum manageable number of
diseases is 90, and the model did not perform well in current datasets.

##  3\. Methodology

In this section, we first present the technical background of BSODA and then
introduce the inquiry branch, including the PoE encoder, the new sampling
scheme, and the acceleration methods for reward estimation. Next, we develop a
knowledge-guided self-attention model in the diagnosis branch. At the end, we
present the diagnostic method that incorporates these two branches and the
termination criterion of the inquiry and diagnosis process. Figure 1 shows the
BSODA process during a single round of inquiry and diagnosis.

###  3.1. Background

####  3.1.1. Problem Formulation

In this paper, we formulate symptom checking as a sequential feature selection
and classification problem. Let $S$ denote the set of all possible symptoms,
and $D$ denote the set of possible diseases. Suppose $x_{s}$ denotes the
presence or absence of symptom $s\in S$: $x_{s}=1$ if a patient actually
suffers from the symptom $s$, i.e., _positive symptom_ , and $x_{s}=0$
otherwise. Similarly, $\mathbf{x}_{D}$ denotes the one-hot categorical vector
of which each dimension indicates whether a disease is present (= 1) or not (=
0). Then, we are interested in predicting the target variables
$\mathbf{x}_{D}$ given the corresponding observed features $\mathbf{x}_{O}$,
where $O\subset S$ is the set of symptoms that are currently observed, and
$U=S\setminus O$ denotes the set of the unobserved ones. More specifically, we
choose which symptom $x_{s}\in\mathbf{x}_{U}$ to inquire next so that our
belief regarding $\mathbf{x}_{D}$ can be optimally improved. Notations are
summarized in Table 1.

Table 1. Notations. Notation | Description
---|---
$S/D$ | Set of symptoms/diseases.
$S_{c}$ | Set of candidate symptoms by filtering.
$x_{s}$ | Presence or absence of symptom $s$ ($s\in S$).
$\mathbf{x}_{D}$ | Presences or absences of diseases.
$\mathbf{x}_{O}/\mathbf{x}_{U}$ | Observed/Unobserved symptoms.
$\mathbf{x}$ | Complete data ($\mathbf{x}=\mathbf{x}_{S}\cup\mathbf{x}_{D}=\mathbf{x}_{O}\cup\mathbf{x}_{U}\cup\mathbf{x}_{D}$).
$N_{T}$ | Maximum number of inquiries.
$N_{C}$ | Number of candidate symptoms.
$N_{M}$ | Number of Monte Carlo samples.
$\Phi/\theta$ | Parameters of the VAE encoder/decoder.
$p_{D}(\cdot)$ | Predictive distribution from the diagnosis model.
$\mathbf{M}$ | Mask matrix.
$\mathbf{P}$ | Conditional probability matrix.
$\mathbf{A}$ | Attention matrix.
$\mathbf{e}/\mathbf{E}$ | Embedding vector(s) of feature(s).

Figure 1. Overall flowchart of BSODA with one round of inquiry and diagnosis.

At first, observed features $\mathbf{x}_{O}$ are fed into the inquiry model.
For every possible symptom $s$ in $S_{c}$ that is built by the filtering
strategy according to prior knowledge, the inquiry branch estimates the reward
$\hat{R}_{s}$, which applies a two-step sampling strategy and approximation
methods. Next, the inquiry branch inquires the symptom $t$ and updates
$\mathbf{x}_{O}$ and $\mathbf{x}_{U}$, where $t=\text{argmax}_{t}\hat{R}_{t}$.
After that, $\mathbf{x}_{O}$ and $\hat{\mathbf{x}}_{U}\sim
p(\mathbf{x}_{U}|\mathbf{x}_{O})$ from the inquiry branch are fed into the
diagnosis model to make predictions. Finally, the termination criterion
determines whether to continue inquiring symptoms.

####  3.1.2. EDDI

A variational autoencoder (VAE) (Kingma and Welling, 2013) defines a
generative model of the form
$p(\mathbf{x},\mathbf{z})=\prod_{i}p_{\theta}(x_{i}|\mathbf{z})p(\mathbf{z})$
in which the data $\mathbf{x}$ is generated from latent variables
$\mathbf{z}$, $p(\mathbf{z})$ is a prior, e.g., spherical Gaussian, and
$p_{\theta}(\mathbf{x}|\mathbf{z})$ is presented as a neural network decoder
with parameters $\theta$ to specify a simple likelihood, e.g., Bernoulli. A
VAE uses another neural network with parameters $\Phi$ as an encoder to
produce a variational approximation of the posterior, that is,
$q_{\Phi}(\mathbf{z}|\mathbf{x})$. A VAE is trained by maximizing an evidence
lower bound (ELBO):

(1) $\mathbb{E}_{q_{\Phi}(\mathbf{z}|\mathbf{x})}[\log
p_{\theta}(\mathbf{x}|\mathbf{z})]-\beta\cdot
D_{\textrm{KL}}[q_{\Phi}(\mathbf{z}|\mathbf{x})\,\|\,p(\mathbf{z})],$

where $\beta$ is the weight to balance the two terms in the expression and
$D_{\textrm{KL}}$ is the Kullback-Leibler (KL) divergence between two
distributions. The ELBO is usually optimized using stochastic gradient descent
and the reparameterization trick (Kingma and Welling, 2013). EDDI (Ma et al.,
2019) is a recently proposed VAE-based framework for the feature acquisition
problem. EDDI chooses the next feature $x_{s}$ to acquire by maximizing over
the information reward

(2) $R_{s}=\mathbb{E}_{x_{s}\sim
p(x_{s}|\mathbf{x}_{O})}D_{\textrm{KL}}\left[p(\mathbf{x}_{D}|x_{s},\mathbf{x}_{O})\,\|\,p(\mathbf{x}_{D}|\mathbf{x}_{O})\right].$

The VAE is used to express the conditional probabilities
$p(x_{s}|\mathbf{x}_{O})$, $p(\mathbf{x}_{D}|\mathbf{x}_{s},x_{O})$, and
$p(\mathbf{x}_{D}|\mathbf{x}_{O})$. Unfortunately, estimating the values of
these quantities in Eqn. 2 is prohibitively expensive. To address this
problem, Ma et al. (2019) demonstrate how Eqn. 2 can be efficiently
approximated using VAE encoding distributions:

(3) $\displaystyle\hat{R}_{s}=\mathbb{E}_{x_{s}\sim
p_{\theta}(x_{s}|\mathbf{x}_{O})}$ $\displaystyle
D_{\textrm{KL}}[q_{\Phi}(\mathbf{z}|x_{s},\mathbf{x}_{O})\,\|\,q_{\Phi}(\mathbf{z}|\mathbf{x}_{O})]-$
$\displaystyle\mathbb{E}_{\mathbf{x}_{D},x_{s}\sim
p_{\theta}(\mathbf{x}_{D},x_{s}|\mathbf{x}_{O})}D_{\textrm{KL}}$
$\displaystyle[q_{\Phi}(\mathbf{z}|\mathbf{x}_{D},x_{s},\mathbf{x}_{O})\,\|\,q_{\Phi}(\mathbf{z}|\mathbf{x}_{D},\mathbf{x}_{O})].$

The first expectation of Eqn. 3 quantifies how much information that $x_{s}$
provides about $\mathbf{z}$, whereas the second expectation quantifies how
much information that $x_{s}$ provides about $\mathbf{z}$ with an addition of
$\mathbf{x}_{D}$. A feature $x_{s}$ will be penalized by the second term if it
is informative about $\mathbf{z}$ but not about $\mathbf{x}_{D}$. All
quantities in the KL divergences above can be computed analytically using
Gaussian approximations. The expectations in Eqn. 3 can be approximated by a
Monte Carlo process, averaging across samples
$\hat{\mathbf{x}}_{D},\hat{x}_{s}\sim p(\mathbf{x}_{D},x_{s}|\mathbf{x}_{O})$
that can be shared between the two expectations.

###  3.2. Inquiry Branch

We adopt EDDI (Ma et al., 2019) as the backbone of the inquiry branch. The
encoder of the VAE is used to handle partial observations of features during
active feature selection. Let $h(\cdot)$ denote a neural network, $g(\cdot)$
denote the summation or max-pooling operation, $\mathbf{e}_{i}$ denote the
embedding vector of the $i$-th feature, and $\mathbf{y}_{i}$ denote the input
carrying information of the $i$-th observed feature, which is constructed by
concatenation: $\mathbf{y}_{i}=[x_{i},\mathbf{e}_{i}]$. EDDI proposed a
permutation invariant set function as an encoder, given by
$\mathbf{c}(\mathbf{x}_{O}):=g(h(\mathbf{y}_{1}),h(\mathbf{y}_{2}),...,h(\mathbf{y}_{|O|}))$,
where $|O|$ denotes the number of observed features. Finally, the code
$\mathbf{c}(\mathbf{x}_{O})$ is fed into a neural network and transformed into
the statistics of a multivariate Gaussian distribution to approximate
$p(\mathbf{z}|\mathbf{x}_{O})$.

####  3.2.1. Product-of-Experts (PoE) Encoder

A simple $g(\cdot)$ cannot handle a large number of features. In BSODA, we
propose to use a PoE encoder (Wu and Goodman, 2018) to approximate the joint
posterior of the latent variables. It assumes conditional independence among
features. The approximated posterior, including a prior Gaussian expert
$p(\mathbf{z})$ with mean vector $\mu_{0}$ and variance vector $V_{0}$, is
given by

(4) $q_{\Phi}(\mathbf{z}|\mathbf{x}_{O})\propto p(\mathbf{z})\prod
q_{\Phi}(\mathbf{z}|x_{i}),$

where $q_{\Phi}(\mathbf{z}|x_{i})$ is an inference network denoting the expert
associated with the $i$-th observed feature $x_{i}$. Figure 1 illustrates the
architecture of the PoE encoder. $\mathbf{e}_{i}$ is pretrained by the
diagnosis model and fixed during training (please refer to Sec. 3.4.1). Then
we use a Multiple Layer Perceptron (MLP) as $h(\cdot)$ to map the input
$\mathbf{y}_{i}$ to a Gaussian distribution in latent space with mean vector
$\mu_{i}$ and variance vector $V_{i}$. $h(\cdot)$ shares parameters among all
features. Because a product of Gaussian experts is itself Gaussian (Cao and
Fleet, 2014), we can quickly compute variance $V=(V_{0}^{-1}+\sum
V_{i}^{-1})^{-1}$ and mean $\mu=(\mu_{0}V_{0}^{-1}+\sum\mu_{i}V_{i}^{-1})V$
for $q_{\Phi}(\mathbf{z}|\mathbf{x}_{O})$ in Eqn. 4. The decoder,
$p_{\theta}(\mathbf{x}|\mathbf{z})$, is given by a product of Bernoulli
distributions, the probabilities of which are specified by an MLP that accepts
$\mathbf{z}$ as input.

If we train the PoE encoder on all available data for a dataset with no
missing features, its performance would suffer when the input contains missing
entries, since the PoE encoder has never encountered such a case during
training. To address this, we drop a random fraction of the fully observed
features $\mathbf{x}$, which includes both $\mathbf{x}_{S}$ and
$\mathbf{x}_{D}$, for each data point during training.

####  3.2.2. Two-Step Sampling Strategy

To estimate the information reward, we approximate expectations by Monte Carlo
sampling, averaging across the samples $\hat{\mathbf{x}}_{D},\hat{x}_{s}\sim
p(\mathbf{x}_{D},x_{s}|\mathbf{x}_{O})$. This can be accomplished by first
sampling $\hat{\mathbf{z}}\sim q_{\Phi}(\mathbf{z}|\mathbf{x}_{O})$, followed
by $\hat{\mathbf{x}}_{D},\hat{x}_{s}\sim
p_{\theta}(\mathbf{x}_{D},x_{s}|\hat{\mathbf{z}})$. However, we propose a more
accurate method for BSODA, in which $\hat{x}_{s}$ and $\hat{\mathbf{x}}_{D}$
are sampled in two steps. Since

(5)
$p(\mathbf{x}_{D},x_{s}|\mathbf{x}_{O})=p(\mathbf{x}_{D}|x_{s},\mathbf{x}_{O})\cdot
p(x_{s}|\mathbf{x}_{O}),$

we propose to sample $\hat{x}_{s}$ from the VAE by sampling
$\hat{\mathbf{z}}\sim q_{\Phi}(\mathbf{z}|\mathbf{x}_{O})$, and then
$\hat{x}_{s}\sim p_{\theta}(x_{s}|\hat{\mathbf{z}})$. Next we sample from
$\hat{\mathbf{x}}_{D}\sim p_{D}(\hat{x}_{s},\mathbf{x}_{O})$ produced by a
diagnosis model. In this way, we use two networks, the VAE and a diagnosis
model, to approximate the joint posterior. By combining the generative and
classification models, we can improve the performance of reward estimation.

###  3.3. Speedup for Reward Estimation

The computational cost of EDDI (Ma et al., 2019) to estimate the reward is
$O(N_{T}\cdot N_{C}\cdot N_{M})$, where $N_{T}$ denotes the maximum number of
inquiries, $N_{C}$ denotes the total number of candidate symptoms, and $N_{M}$
denotes the number of samples required in the Monte Carlo process. Because
$N_{C}$ can account for thousands of symptoms, it would be too expensive to
support an online disease diagnosis service. This section shows how we may
significantly reduce the cost by taking advantage of feature sparsity and
predictive probability distributions from the diagnosis model.

####  3.3.1. Filtering Candidate Symptoms by Prior Knowledge

In each round, the model chooses a feature $x_{s}$ from $\mathbf{x}_{U}$ in
order to maximize the reward. Under the setting of symptom-inquiry, we can
filter out irrelevant symptoms to reduce the number of candidates that need to
be queried.

For every symptom $s$, we calculate the set $S_{s}$ of additional symptoms
that may co-occur with symptom $s$ with high probability in one patient, based
on prior knowledge or training data statistics. More specifically, we assign a
symptom $j\in S_{s}$ if $P_{data}(x_{j}=1|x_{s}=1)$ is larger than a pre-
defined threshold, which we can simply set to 0 by default. At the beginning
of the inquiry process when symptom $i$ is present, the set of candidate
symptoms, $S_{c}$, is initialized to $S_{i}$. Every time the framework selects
a new positive symptom $j$, we update the set of candidate symptoms to be the
intersection of the current set and $S_{j}$. Thus we reduce the computational
cost $O(N_{T}\cdot N_{C}\cdot N_{M})$ by lowering $N_{C}$.

####  3.3.2. Approximation of Reward

The reward $\hat{R}_{s}$ in Eqn. 3 is estimated by averaging over $N_{M}$
Monte Carlo samples $\hat{\mathbf{x}}_{D},\hat{x}_{s}\sim
p(\mathbf{x}_{D},x_{s}|\mathbf{x}_{O})$. It should be noted that $x_{s}$ is
binary and $\mathbf{x}_{D}$ is discrete among $|D|$ diseases in symptom
checking tasks. Indeed we just need to compute rewards of $2\cdot|D|$ distinct
combinations of $x_{s}$ and $\mathbf{x}_{D}$, each with a weight
$p(\mathbf{x}_{D},x_{s}|\mathbf{x}_{O})$ that is solved using the two-step
sampling strategy described in Sec. 3.2.2. Then the computational cost is
$O(N_{T}\cdot N_{C}\cdot(2|D|))$. In the following, we demonstrate how to
speed up it by reducing $2|D|$.

Firstly, consider that the number of symptoms experienced by a patient with a
specific disease is generally smaller than the number of possible symptoms for
the disease, which is much smaller than the total number of symptoms $|S|$.
With such a sparse feature space, the inquiry process should be prioritized to
acquire positive symptoms (Peng et al., 2018). Motivated by this, we encourage
BSODA to focus on positive symptoms by discarding combinations of $x_{s}$ and
$\mathbf{x}_{D}$ when $\hat{x}_{s}=0$, i.e., setting their rewards to 0. Hence
the number of combinations of $x_{s}$ and $\mathbf{x}_{D}$ to be computed is
reduced to $|D|$.

Secondly, when $|D|$ is large, practically almost all predictive distributions
$p_{D}(\cdot)$ are long-tailed. Hence we rank all possible disease values for
$\mathbf{x}_{D}$ by its probability in $p_{D}(\hat{x}_{s},\mathbf{x}_{O})$,
using the two-step sampling strategy in Sec. 3.2.2, and we discard those that
fall below the 90 percent of the top probability or are less than $1/|D|$.
Then the average number of combinations to be computed can be further reduced
to less than 10 experimentally.

###  3.4. Diagnosis Branch

It should be noted that we may wish to perform diagnosis in any round of the
symptom checking process. Recently, an increasing number of RL methods (Lin et
al., 2020; Liao et al., 2020; Liu et al., 2021) have used a supervised
learning model to handle the diagnosis. In BSODA, we denote $p_{D}(\cdot)$ to
be the disease predictive distributions returned by the diagnosis model.

####  3.4.1. Knowledge-Guided Self-Attention Model

BSODA uses a self-attention mechanism guided by prior knowledge to embed
relationships between all features, consisting of both diseases and symptoms.
We use two matrices, $\mathbf{M}$ and $\mathbf{P}$, to represent prior
knowledge, which are of the same size $(|S|+|D|)\times(|S|+|D|)$. $\mathbf{M}$
is a mask matrix of which each entry $\mathbf{M}_{ij}=0$ denotes that the
relationship between feature $i$ and feature $j$ is considered, and negative
infinity otherwise. Specifically, in $\mathbf{M}$, only disease-disease
relationships are not considered. $\mathbf{P}$ is a matrix of which each entry
$\mathbf{P}_{ij}$ denotes the conditional probability $P(j|i)$, which is
calculated by co-occurrence between $i$ and $j$ according to prior knowledge
or training data statistics, and normalized such that each row sums to 1,
i.e., $P(\cdot|i)=1$. Again, only $P(\text{symptom}|\text{disease})$ and
$P(\text{symptom}|\text{symptom})$ are calculated, and all other probabilities
are set to zeros.

The self-attention model consists of several stacked self-attention blocks.
Each block accepts the outputs from the previous block, including embeddings
of all features and an attention matrix representing the weights of the
relationships between features, and calculates new embeddings and attention
for the next block. The output of an attention function can be described as a
weighted sum of the $\mathbf{V}$alues, where the weights are computed by the
$\mathbf{Q}$uery with the corresponding $\mathbf{K}$ey (Vaswani et al., 2017).
For self-attention, $\mathbf{Q}$, $\mathbf{K}$ and $\mathbf{V}$ are derived
from the same source. Defining $\mathbf{A}^{(j)}$ as the attention matrix and
$\mathbf{E}^{(j)}$ as the embedding vectors calculated by the $j$-th block
($j\geq 1$), the self-attention block is given by:

(6) $\displaystyle\mathbf{A}^{(j)}$
$\displaystyle=\text{Softmax}(\frac{\mathbf{Q}^{(j)}\mathbf{K}^{(j)\top}}{\sqrt{c}}+\mathbf{M}),$
$\displaystyle\mathbf{E}^{(j)}=\text{MLP}^{(j)}(\mathbf{A}^{(j)}\mathbf{V}^{(j)}),$
$\displaystyle\mathbf{Q}^{(j)}=\mathbf{E}^{(j-1)}\mathbf{W}_{Q}^{(j)}$
$\displaystyle,\ \mathbf{K}^{(j)}=\mathbf{E}^{(j-1)}\mathbf{W}_{K}^{(j)},\
\mathbf{V}^{(j)}=\mathbf{E}^{(j-1)}\mathbf{W}_{V}^{(j)},$

where $c$ is the column size of $\mathbf{W}_{K}$. All $\mathbf{W}$s are
trainable parameters. $\mathbf{E}^{(0)}$ is constructed by concatenation: if
$i\in S$, $\mathbf{E}_{i}^{(0)}=[x_{s},\mathbf{e}_{i}]$; if $i\in D$,
$\mathbf{E}_{i}^{(0)}=[0,\mathbf{e}_{i}]$. When calculating the attention
matrix $\mathbf{A}$, the attention weights of the relationships that are not
considered are fixed to 0 by the negative infinities in $\mathbf{M}$. To
incorporate prior knowledge $\mathbf{P}$, we initialize $\mathbf{A}^{(0)}$ by
$\mathbf{P}$ and sequentially penalize the attention of the $j$-th block if it
deviates too much from the previous block by a regularized training loss term,
$\sum_{j}D_{\textrm{KL}}(\mathbf{A}^{(j-1)}||\mathbf{A}^{(j)})$. Finally, the
predictive distribution is obtained from the last embedding vectors of
$\mathbf{x}_{D}$:
$p_{D}(\cdot)=\text{Softmax}(\text{MLP}^{D}(\mathbf{E}_{D}^{(J)}))$. We train
the model on complete data $\mathbf{x}$, which accepts $\mathbf{x}_{S}$ as
input and predicts $\mathbf{x}_{D}$. The learned embeddings $\mathbf{E}$ are
used to initialize the VAE model.

####  3.4.2. Diagnostic Method and Termination Criterion

The uncertainty caused by partial observations of features will make it
difficult to perform an accurate diagnosis. Hence, to model the uncertainty,
we impute the missing variables $\mathbf{x}_{U}$ in the input by drawing
$N_{M}$ samples from $\hat{\mathbf{x}}_{U}\sim
p(\mathbf{x}_{U}|\mathbf{x}_{O})$, for which the sampling process is the same
as $\hat{x}_{s}\sim p(x_{s}|\mathbf{x}_{O})$ according to the two-step
sampling strategy in Sec. 3.2.2. These $N_{M}$ samples
$\\{(\hat{\mathbf{x}}_{U},\mathbf{x}_{O})\\}_{m=1}^{N_{M}}$ are then fed into
the diagnosis model, yielding a set of predictive distributions on
$\mathbf{x}_{D}$, denoted by
$\\{(p^{(1)},...,p^{(|D|)})^{m}\\}_{m=1}^{N_{M}}$. Then, we calculate the
expectations $\mu_{p_{D}}=\mathbb{E}[p^{(1)},...,p^{(|D|)}]$ and standard
deviations $\sigma_{p_{D}}=\sqrt{\mathrm{Var}[p^{(1)},...,p^{(|D|)}]}$. BSODA
would stop inquiring to report the chosen disease if the probability of this
disease is so high that inquiring more symptoms would not overturn the
diagnostic result (Lin et al., 2020). That is, the inquiry process would stop
when the probability of the chosen disease is beyond the upper bound of the
$6\sigma$ interval (Parzen, 1960; Pande and Holpp, 2001) of the probabilities
of other diseases. When a single disease $d$, where $d=\text{argmax}_{d}\
\mu_{p_{D}}^{(d)}$, is chosen, the termination criterion can be formulated as:

(7) $\forall j\neq
d,\mu_{p_{D}}^{(d)}>\mu_{p_{D}}^{(j)}+3\sigma_{p_{D}}^{(j)}.$

##  4\. Experiments

Table 2. Essential characteristics of experimental datasets. Dataset | Type | Data Size | #Diseases | #Symptoms | Usage
---|---|---|---|---|---
SymCAT | Synthetic | - | 801 | 474 | Train & Test
MuZhi | Real-World | 710 | 4 | 66 | Train & Test
Dxy | Real-World | 527 | 5 | 41 | Train & Test
HPO-Synthetic | Synthetic | - | 11,441 | 13,032 | Train & Test
HPO-HMS | Real-World | 83 | 37 | 747 | Test
HPO-MME | Real-World | 43 | 18 | 559 | Test

Table 2 summarizes the essential characteristics of the experimental datasets.
BSODA is evaluated on three categories of experimental datasets including a
synthetic dataset, two real-world common disease datasets, and three rare
disease datasets including one synthetic dataset derived from a knowledge base
and two real-world datasets. They are all publicly accessible and do not
contain personally identifiable information. We use synthetic datasets to
generate several synthetic tasks with different numbers of possible diseases.

###  4.1. Existing Datasets

####  4.1.1. SymCAT

SymCAT is a symptom-disease database (AHEAD Research Inc, 2017). For each
disease, there is information about its symptoms with their marginal
probabilities. It is filtered using the occurrence rates in Centers for
Disease Control and Prevention (CDC) database. Following Peng et al. (2018),
we use SymCAT to build a synthetic dataset. We first sample a disease and its
related symptoms from all diseases uniformly at random. Then, we perform a
Bernoulli trial on each extracted symptom based on its associated probability
to form a symptom set. For instance, if we sample a disease ”abscess of nose”,
we would obtain its associated symptoms: ”cough” and ”fever”, which occur with
probabilities of 73% and 62%, respectively. We then generate one data instance
by sampling Bernoulli random variables according to these probabilities. We
sample $10^{6}$, $10^{5}$, and $10^{4}$ records for training, validation, and
testing, respectively.

####  4.1.2. Two Real-World Datasets

Wei et al. (2018) constructed the MuZhi Medical Dialogue dataset, which was
collected from the dialogue data from the pediatric department of a Chinese
online healthcare website (https://muzhi.baidu.com/). The MuZhi dataset
includes 66 symptoms for 4 diseases: children’s bronchitis, children’s
functional dyspepsia, infantile diarrhea infection, and upper respiratory
infection.

The Dxy Medical Dialogue dataset (Xu et al., 2019) contains data from another
popular Chinese online healthcare website (https://dxy.com/), where users
often seek professional medical advice from specialists. The Dxy dataset
includes 41 symptoms for 5 diseases: allergic rhinitis, upper respiratory
infection, pneumonia, children hand-foot-mouth disease, and pediatric
diarrhea. These two datasets have been structured in the form of disease-
symptoms pairs and split into a training set and a testing set.

###  4.2. Proposed HPO-Based Rare Disease Datasets

Human Phenotype Ontology (HPO) (Köhler et al., 2017) provides a standardized
vocabulary of phenotypic abnormalities associated with human disease. The
initial domain application of the HPO was on rare disorders. There are 123,724
annotations of HPO terms to rare diseases and 132,620 to common diseases.
These were compiled from multiple sources, including the medical literature,
Orphanet (https://www.orpha.net/), OMIM (https://omim.org/) and DECIPHER
(https://decipher.sanger.ac.uk/). We chose diseases and related symptoms from
the public rare disease dictionary, along with their marginal probabilities,
to build new datasets.

The HMS dataset was collected by Knitza et al. (2019), including 93 rare
disease cases from original real-world health records. The Matchmaker Exchange
API (Buske et al., 2015) provides a shared language that databases can use to
query each other to find similar patients. The API is built around a
standardized patient profile including both phenotype and genotype
information. The MME dataset is the standardized rare disease test set for the
API, containing 50 de-identified patients selected from publications. We chose
83 and 43 cases from the HMS and MME datasets, respectively, to form new HPO-
HMS and HPO-MME datasets. Each of them have at least one symptom matching a
particular HPO term.

We propose a novel HPO-based evaluation approach using both synthetic and
real-world datasets. The HPO database is used to create a synthetic dataset as
the simulation process of SymCAT. The real-world datasets, HPO-HMS and HPO-
MME, are used to test the transferability from knowledge-based simulated
environments to real-world scenarios. Specifically, we build a _synthetic
training_ dataset by the aforementioned process for SymCAT, utilizing the HPO
entries that match the symptoms and diseases in HPO-HMS and HPO-MME, which are
_real-world testing_ datasets. This evaluation approach can assess how well a
symptom checking method generalizes when trained solely on synthetic data,
because large amounts of _real-world training_ data are difficult to obtain.

###  4.3. Experimental Protocol

####  4.3.1. Baselines

We have selected the most appropriate RL baseline for every dataset. GAMP (Xia
et al., 2020) uses GAN to implement an RL agent for symptom checking, which
has been proven to be state-of-the-art only on tiny feature spaces (MuZhi and
Dxy). REFUEL (Peng et al., 2018) is a competitive RL method designed for
symptom checking on a medium size search space.

Table 3. Performance of REFUEL (Peng et al., 2018) and BSODA on the synthetic SymCAT dataset and HPO-based rare disease datasets. Dataset |  #-Disease Task | #Symptoms | REFUEL (Peng et al., 2018) | BSODA
---|---|---|---|---
Top1 | Top3 | Top5 | #Rounds | Top1 | Top3 | Top5 | #Rounds
SymCAT- Synthetic | 200 | 328 | 53.76 | 73.12 | 79.53 | 8.24 | $\mathbf{55.65\pm 0.25}$ | $\mathbf{80.71\pm 0.26}$ | $\mathbf{89.32\pm 0.29}$ | $12.02\pm 0.06$
300 | 349 | 47.65 | 66.22 | 71.79 | 8.39 | $\mathbf{48.23\pm 0.25}$ | $\mathbf{73.82\pm 0.32}$ | $\mathbf{84.21\pm 0.17}$ | $13.10\pm 0.05$
400 | 355 | 43.01 | 59.65 | 68.89 | 8.92 | $\mathbf{44.63\pm 0.29}$ | $\mathbf{69.22\pm 0.11}$ | $\mathbf{79.54\pm 0.15}$ | $14.42\pm 0.03$
HPO-Synthetic | 500 | 1901 | 64.33 | 73.14 | 75.34 | 8.09 | $\mathbf{76.23\pm 0.47}$ | $\mathbf{84.17\pm 0.53}$ | $\mathbf{86.99\pm 0.41}$ | $5.34\pm 0.04$
1000 | 3599 | 40.08 | 62.67 | 67.42 | 14.19 | $\mathbf{67.03\pm 0.29}$ | $\mathbf{75.94\pm 0.53}$ | $\mathbf{79.27\pm 0.58}$ | $10.62\pm 0.08$
HPO-HMS | 37 | 747 | 18.22 | 34.07 | 45.31 | 11.26 | $\mathbf{34.22\pm 0.93}$ | $\mathbf{61.39\pm 1.12}$ | $\mathbf{68.43\pm 1.12}$ | $13.40\pm 0.06$
HPO-MME | 18 | 559 | 48.15 | 58.40 | 67.33 | 9.23 | $\mathbf{60.70\pm 1.28}$ | $\mathbf{79.53\pm 1.40}$ | $\mathbf{83.37\pm 0.86}$ | $10.14\pm 0.14$

####  4.3.2. Experimental Design

BSODA first trains the knowledge-guided self-attention model in the diagnosis
branch, and then trains the VAE model with a fraction of the features dropped
at random in the inquiry branch. In comparison to RL, BSODA handles this
training naturally.

For the synthetic datasets and the two rare disease datasets, HPO-HMS and HPO-
MME, a positive symptom is chosen uniformly at random as a patient self-report
at the beginning of the inquiring process. MuZhi and Dxy already have self-
reports. For all real-world datasets, the filtering speedup technique is not
applied because their numbers of candidate symptoms are already small. We show
the results of BSODA in percentage with a 95% confidence interval for 5 and 20
random runs, respectively for the synthetic and real-world datasets.

Figure 2. Accuracy vs the maximum number of inquiries for BSODA and REFUEL
(Peng et al., 2018) on the SymCAT 200-disease task.

Within the beginning several steps, the accuracy curve of REFUEL is far below
the one of BSODA. When the accuracy curve of REFUEL approaches its average
steps (8.24), the slope of the curve is going down, and when it goes over 10,
the slope is nearly zero. On the contrary, the accuracy curve of BSODA
maintains a quick growth rate. For Top3 and Top5, the curves of REFUEL are
always under the ones of BSODA. Although there are intersections in Top1
curves, BSODA would outperform REFUEL at the beginning and the future.

###  4.4. Results

####  4.4.1. SymCAT

We synthesize three diagnostic tasks containing 200, 300, and 400 diseases,
respectively. We set $N_{M}$ to 100 for BSODA and the maximum number of
inquiries to 15, 16, and 17 for these three tasks respectively for both BSODA
and REFUEL. The results are reported in Table 3. We observe that on average,
BSODA conducts more inquiries than REFUEL, and as a result, it has a slightly
higher accuracy for the Top 1 prediction and a significantly higher accuracy
for the Top3 and Top5 predictions.

In Figure 2, we present how accuracy changes as the maximum number of
inquiries increases for the 200-disease task. REFUEL was trained to make a
maximum of 20 inquiries. We chose the predicted disease with the highest
action value under the learned policy. In Figure 2, we observe that the
accuracy curves for REFUEL are significantly lower than those for BSODA in the
initial several inquiries. It reveals that with the partial observed
information of patients, REFUEL may struggle to make accurate predictions,
which is where BSODA’s diagnostic method that leverages the generative model
to handle the uncertainties excels.

When REFUEL’s accuracy curves approach their optimal number of inquiries, the
curves reach a plateau. In comparison, BSODA’s accuracy curves continue to
grow at a relatively faster rate. BSODA always performs better than REFUEL on
the Top3 and Top5 curves. Although the Top1 curves intersect, BSODA
outperforms REFUEL at the beginning and the end. REFUEL stops the inquiry and
performs a prediction when it believes that it has sufficient evidence;
however, BSODA proves that this decision may not be optimal: with a few
additional inquiries, BSODA may provide a significantly better disease
prediction, which is critical in the medical domain. In conclusion, using the
same maximum number of inquiries, an RL agent would converge sooner and make
an inferior prediction. It is mostly likely due to the difficulties to
determine whether the exploration, i.e., the information gathered, is
sufficient to provide a ”mature” diagnosis for RL agents. In comparison, the
heuristic stop criterion of BSODA has an advantage.

Table 4. Disease prediction accuracy (%) of REFUEL (Peng et al., 2018), GAMP (Xia et al., 2020) and BSODA on two medical dialogue datasets. |  REFUEL (Peng et al., 2018) |  GAMP (Xia et al., 2020) | BSODA
---|---|---|---
MuZhi (Wei et al., 2018) | 71.8 | 73 | $73.1\pm 0.5$
Dxy (Xu et al., 2019) | 75.7 | 76.9 | $\mathbf{80.2\pm 0.3}$

####  4.4.2. Real-World Common Disease Datasets

Table 4 shows the results for REFUEL, GAMP and BSODA on the MuZhi and Dxy
datasets. For BSODA, we set $N_{M}$ to 100 and the maximum number of inquiries
to 16 for MuZhi and 20 for Dxy, respectively. For both GAMP and REFUEL, we
limit the maximum number of inquiries to 20. The results in Table 4 show that
among these three approaches, BSODA has a competitive performance on MuZhi and
achieves the highest accuracy on Dxy.

From the Dxy dataset, we conducted a case study to investigate the actual
inquired symptoms of REFUEL, BSODA and a doctor. The patient selected was a
21-month infant, who initially reported _Coughing Up Phlegm_ , and was easily
misdiagnosed as having _Upper Respiratory Infection_. The results are shown in
Table 5 with positive inquired symptoms highlighted in bold. The key implicit
symptoms were _Moist Rales_ and _Difficult Breathing_. We observe that REFUEL
provided a hasty and incorrect diagnosis without inquiring about the key
symptoms. BSODA tended to make a differential diagnosis by first inquiring
_Nose Rubbing_ and _Allergic Symptoms_ in order to exclude _Allergic Rhinitis_
, and then inquiring about the key symptoms to provide a correct diagnosis.

####  4.4.3. Proposed HPO-Based Rare Disease Datasets

We utilize the HPO to generate synthetic tasks with thousands of features to
evaluate the performance of symptom checking methods. Two synthetic tasks were
generated, containing 500 and 1,000 diseases, respectively. Our goal is to
evaluate the models in real-world scenarios with a symptom count six times
larger than that of SymCAT, the previous largest dataset. For BSODA, we set
the maximum number of inquiries to 10 and 18 for the tasks with 500 and 1,000
diseases, respectively, and fix $N_{M}$ to 20. In training REFUEL, we set the
maximum number of inquiries to 14 and 18 for these two tasks, respectively.
The results of HPO-synthetic in Table 3 show that BSODA outperforms REFUEL by
a large margin, with accuracy from 11 to 27 percentage points higher while
making fewer inquiries. When the number of diseases grows from 500 to 1,000,
the Top1 accuracy of REFUEL drops by 24%, but that of BSODA decreases by only
9%. Thus, BSODA shows clearer advantages over RL when dealing with large
feature spaces.

For the HPO-HMS and HPO-MME datasets, we set the maximum number of inquiries
to 15 for both REFUEL and BSODA, and fix $N_{M}$ to 100 for BSODA. As shown in
Table 3, BSODA outperforms REFUEL by from 12 to 21 percentage points. However,
the results on the HPO-HMS dataset are not satisfactory enough since the Top1
accuracy of BSODA is only 34. Indeed, because of the disparity in symptom
distributions between the knowledge base and real-world data, transferability
of symptom checking methods remains to be a challenge.

Table 5. A case study for actual inquired symptoms on the Dxy dataset. The patient was a 21-month infant who initially reported _Coughing Up Phlegm_. Source | Inquired Symptoms | Diagnosis
---|---|---
REFUEL | Runny Nose, Green Stool, Fever, | Upper Respiratory
(Peng et al., 2018) | Excessive Breathing | Infection
BSODA | Fever, Nose Rubbing, | Pneumonia
Allergic Symptoms, Moist Rales,
Difficult Breathing, Runny Nose
Doctor |  Moist Rales, Difficult Breathing | Pneumonia

###  4.5. Ablation Study of Speedup Techniques

We performed an ablation study to show the acceleration provided by various
components of the proposed speedup scheme. We ran BSODA on the synthetic
SymCAT 200-disease task, and calculated the average CPU time per inquiry with
various speedup settings. We set $N_{M}$ to 100 and the maximum number of
inquiries to 15. In Table 6, we present the results produced using an NVIDIA
GeForce GTX 1080Ti GPU. The results demonstrate that while using the symptom
filtering strategy alone achieves about 1.59x speedup, combining the symptom
filtering and approximate sampling strategies result in a much larger speedup
of 70x . Obviously, adopting the approximation strategy alone results in a
significant speedup, and without them, the reward estimation process of EDDI
(Ma et al., 2019) is impracticable for an online service.

##  5\. Limitations

From the view of BSODA, a limitation is that the inquiry branch does not
leverage medical knowledge explicitly. Although we use a knowledge-guided
diagnosis model to assist inquiring, the model-driven symptom-inquiry process
is likely to differ from the way that doctors interact with patients. This is
for our future study.

From the view of symptom checking research, a limitation is that there is lack
of evaluation about co-morbidity, i.e., a patient having two or more diseases,
which is common in the real world. Ideally, we could simply extend BSODA to
handle co-morbidity by assuming that the diseases are independent. All that is
required is to train a multi-label diagnosis model instead of a multi-class
one. Unfortunately, we lack evaluation data in this area since we could not
find co-morbidity datasets in current literature. Another limitation in
symptom checking research is that using synthetic datasets frequently results
in inflated performance. It is because the data generation models make a very
simple assumption of symptom-disease and symptom-symptom relationships, while
the real-world cases are much more complex.

Table 6. Average CPU time per inquiry with various speedup settings for BSODA on the SymCAT 200-disease task. Settings | Time (s)
---|---
BSODA | $\mathbf{0.45}$
BSODA w/o Approximation | 20.47
BSODA w/o Approximation & Filtering | 32.52

##  6\. Conclusion

We offer BSODA, a novel non-RL bipartite framework for online disease
diagnosis that is computationally efficient and scalable to large feature
spaces. We adopt a Product-of-Experts encoder to efficiently handle partial
observations of symptom data and design a two-step sampling strategy to
estimate the reward more precisely. We develop filtering and approximation
strategies to significantly speed up estimations of the information reward to
a level that is practical for online services. As demonstrated by tests, BSODA
outperforms prior methods in their respective settings (Peng et al., 2018;
Kachuee et al., 2018b; Ma et al., 2019; Xia et al., 2020). With a novel
evaluation method, BSODA demonstrates its ability to effectively handle large
search spaces. According to our knowledge, this is the first work that
proposes to test the transferability of symptom checking methods from
knowledge-based simulated environments to real-world data.

###### Acknowledgements.

This work is supported by the National Natural Science Foundation of China
(grants 61872218), National Key R&D Program of China (2019YFB1404804 and
2021YFF1200009), Beijing National Research Center for Information Science and
Technology (BNRist), Tsinghua University Initiative Scientific Research
Program, and Guoqiang Institute, Tsinghua University. The funders had no roles
in study design, data collection and analysis, the decision to publish, and
preparation of the manuscript.

## References

  * (1)
  * may (2015) 2015\.  Mayo Clinic symptom checker.  https://www.mayoclinic.org/symptom-checker/select-symptom/itt-20009075.
  * nhr (2018) 2018\.  National health and nutrition examination survey.  https://www.cdc.gov/nchs/nhanes.
  * web (2020) 2020\.  WebMD symptom checker.  https://symptoms.webmd.com/.
  * AHEAD Research Inc (2017) AHEAD Research Inc. 2017\.  SymCAT: Symptom-based, computer assisted triage.  http://www.symcat.com.
  * Buske et al. (2015) Orion J Buske, François Schiettecatte, Benjamin Hutton, Sergiu Dumitriu, Andriy Misyura, Lijia Huang, Taila Hartley, Marta Girdea, Nara Sobreira, Chris Mungall, et al. 2015\.  The Matchmaker Exchange API: automating patient matching through the exchange of structured phenotypic and genotypic profiles.  _Human mutation_ 36, 10 (2015), 922–927.
  * Cao and Fleet (2014) Yanshuai Cao and David J Fleet. 2014.  Generalized product of experts for automatic and principled fusion of Gaussian process predictions.  _arXiv preprint arXiv:1410.7827_ (2014).
  * Choi et al. (2020) Edward Choi, Zhen Xu, Yujia Li, Michael Dusenberry, Gerardo Flores, Emily Xue, and Andrew Dai. 2020.  Learning the Graphical Structure of Electronic Health Records with Graph Convolutional Transformer. In _Proceedings of the AAAI Conference on Artificial Intelligence_ , Vol. 34. 606–613.
  * Cross et al. (2021) Sebastian Cross, Ahmed Mourad, Guido Zuccon, and Bevan Koopman. 2021.  Search engines vs. symptom checkers: A comparison of their effectiveness for online health advice. In _Proceedings of the Web Conference 2021_. 206–216.
  * Early et al. (2016) Kirstin Early, Stephen E Fienberg, and Jennifer Mankoff. 2016\.  Test time feature ordering with FOCUS: Interactive predictions with minimal user burden. In _Proceedings of the 2016 ACM International Joint Conference on Pervasive and Ubiquitous Computing_. 992–1003.
  * Fox and Duggan (2013) Susannah Fox and Maeve Duggan. 2013.  Health Online 2013.  _Health_ (2013).
  * Hayashi (1991) Yoichi Hayashi. 1991\.  A neural expert system with automated extraction of fuzzy if-then rules and its application to medical diagnosis. In _Advances in neural information processing systems_. 578–584.
  * Hu and Haake (2010) Yifeng Hu and Jessica Haake. 2010.  Search your way to an accurate diagnosis: Predictors of Internet-based diagnosis accuracy.  _Atlantic Journal of Communication_ 18, 2 (2010), 79–88.
  * Janisch et al. (2019) Jaromír Janisch, Tomáš Pevnỳ, and Viliam Lisỳ. 2019.  Classification with costly features using deep reinforcement learning. In _Proceedings of the AAAI Conference on Artificial Intelligence_ , Vol. 33. 3959–3966.
  * Janisch et al. (2020) Jaromír Janisch, Tomáš Pevnỳ, and Viliam Lisỳ. 2020.  Classification with costly features as a sequential decision-making problem.  _Machine Learning_ (2020), 1–29.
  * Kachuee et al. (2018a) Mohammad Kachuee, Sajad Darabi, Babak Moatamed, and Majid Sarrafzadeh. 2018a.  Dynamic feature acquisition using denoising autoencoders.  _IEEE transactions on neural networks and learning systems_ 30, 8 (2018), 2252–2262.
  * Kachuee et al. (2018b) Mohammad Kachuee, Orpaz Goldstein, Kimmo Kärkkäinen, Sajad Darabi, and Majid Sarrafzadeh. 2018b.  Opportunistic Learning: Budgeted Cost-Sensitive Learning from Data Streams. In _International Conference on Learning Representations_.
  * Kao et al. (2018) Hao-Cheng Kao, Kai-Fu Tang, and Edward Y Chang. 2018\.  Context-aware symptom checking for disease diagnosis using hierarchical reinforcement learning. In _Thirty-Second AAAI Conference on Artificial Intelligence_.
  * Keselman et al. (2008) Alla Keselman, Allen C Browne, and David R Kaufman. 2008\.  Consumer health information seeking as hypothesis testing.  _Journal of the American Medical Informatics Association_ 15, 4 (2008), 484–495.
  * Kingma and Welling (2013) Diederik P Kingma and Max Welling. 2013.  Auto-encoding variational bayes.  _arXiv preprint arXiv:1312.6114_ (2013).
  * Knitza et al. (2019) Johannes Knitza, Koray Tascilar, Eva-Maria Messner, Marco Meyer, Diana Vossen, Almut Pulla, Philipp Bosch, Julia Kittler, Arnd Kleyer, Philipp Sewerin, et al. 2019\.  German mobile apps in rheumatology: review and analysis using the mobile application rating scale (MARS).  _JMIR mHealth and uHealth_ 7, 8 (2019), e14991.
  * Kohavi (1996) Ron Kohavi. 1996\.  Scaling up the accuracy of naive-bayes classifiers: A decision-tree hybrid.. In _Kdd_ , Vol. 96. 202–207.
  * Köhler et al. (2017) Sebastian Köhler, Nicole A Vasilevsky, Mark Engelstad, Erin Foster, Julie McMurry, Ségolène Aymé, Gareth Baynam, Susan M Bello, Cornelius F Boerkoel, Kym M Boycott, et al. 2017\.  The human phenotype ontology in 2017.  _Nucleic acids research_ 45, D1 (2017), D865–D876.
  * Kononenko (1993) Igor Kononenko. 1993\.  Inductive and Bayesian learning in medical diagnosis.  _Applied Artificial Intelligence an International Journal_ 7, 4 (1993), 317–337.
  * Kononenko (2001) Igor Kononenko. 2001\.  Machine learning for medical diagnosis: history, state of the art and perspective.  _Artificial Intelligence in medicine_ 23, 1 (2001), 89–109.
  * Ledley and Lusted (1959) Robert S Ledley and Lee B Lusted. 1959.  Reasoning foundations of medical diagnosis.  _Science_ 130, 3366 (1959), 9–21.
  * Lewenberg et al. (2017) Yoad Lewenberg, Yoram Bachrach, Ulrich Paquet, and Jeffrey S Rosenschein. 2017.  Knowing What to Ask: A Bayesian Active Learning Approach to the Surveying Problem.. In _AAAI_. 1396–1402.
  * Liao et al. (2020) Kangenbei Liao, Qianlong Liu, Zhongyu Wei, Baolin Peng, Qin Chen, Weijian Sun, and Xuanjing Huang. 2020.  Task-oriented Dialogue System for Automatic Disease Diagnosis via Hierarchical Reinforcement Learning.  _arXiv preprint arXiv:2004.14254_ (2020).
  * Lin et al. (2020) Junfan Lin, Keze Wang, Ziliang Chen, Xiaodan Liang, and Liang Lin. 2020.  Towards Causality-Aware Inferring: A Sequential Discriminative Approach for Medical Automatic Diagnosis.  _arXiv preprint arXiv:2003.06534_ (2020).
  * Liu et al. (2021) Zhuo Liu, Yanxuan Li, Xingzhi Sun, Fei Wang, Gang Hu, and Guotong Xie. 2021\.  Dialogue Based Disease Screening Through Domain Customized Reinforcement Learning. In _Proceedings of the 27th ACM SIGKDD Conference on Knowledge Discovery & Data Mining_. 1120–1128.
  * Lopes and Ribeiro (2013) Carla Teixeira Lopes and Cristina Ribeiro. 2013.  Query behavior: The impact of health literacy, topic familiarity and terminology. In _International Conference on Human Factors in Computing and Informatics_. Springer, 212–223.
  * Ma et al. (2019) C Ma, S Tschiatschek, K Palla, JM Hernández-Lobato, S Nowozin, and C Zhang. 2019\.  EdDI: Efficient dynamic discovery of high-value information with partial VAE. In _36th International Conference on Machine Learning, ICML 2019_ , Vol. 2019. 7483–7504.
  * Nan et al. (2015) Feng Nan, Joseph Wang, and Venkatesh Saligrama. 2015\.  Feature-budgeted random forest.  _arXiv preprint arXiv:1502.05925_ (2015).
  * Nan et al. (2016) Feng Nan, Joseph Wang, and Venkatesh Saligrama. 2016\.  Pruning random forests for prediction on a budget. In _Advances in neural information processing systems_. 2334–2342.
  * Pande and Holpp (2001) Pete S. Pande and Larry Holpp. 2001.  _What Is Six SIGMA?_ McGraw-Hill Professional.
  * Parzen (1960) Emanuel Parzen. 1960\.  _Modern probability theory and its applications_.  John Wiley & Sons, Incorporated.
  * Peng et al. (2018) Yu-Shao Peng, Kai-Fu Tang, Hsuan-Tien Lin, and Edward Chang. 2018.  Refuel: Exploring sparse features in deep reinforcement learning for fast disease diagnosis. In _Advances in neural information processing systems_. 7322–7331.
  * Puspitasari (2017) Ira Puspitasari. 2017\.  The impacts of consumer’s health topic familiarity in seeking health information online. In _2017 IEEE 15th International Conference on Software Engineering Research, Management and Applications (SERA)_. IEEE, 104–109.
  * Semigran et al. (2015) Hannah L Semigran, Jeffrey A Linder, Courtney Gidengil, and Ateev Mehrotra. 2015.  Evaluation of symptom checkers for self diagnosis and triage: audit study.  _bmj_ 351 (2015), h3480.
  * Shim et al. (2018) Hajin Shim, Sung Ju Hwang, and Eunho Yang. 2018.  Joint active feature acquisition and classification with variable-size set encoding. In _Advances in neural information processing systems_. 1368–1378.
  * Tang et al. (2016) Kai-Fu Tang, Hao-Cheng Kao, Chun-Nan Chou, and Edward Y Chang. 2016.  Inquire and diagnose: Neural symptom checking ensemble using deep reinforcement learning. In _NIPS Workshop on Deep Reinforcement Learning_.
  * Vaswani et al. (2017) Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. 2017.  Attention is all you need. In _Advances in neural information processing systems_. 5998–6008.
  * Wang et al. (2021) Zifeng Wang, Rui Wen, Xi Chen, Shilei Cao, Shao-Lun Huang, Buyue Qian, and Yefeng Zheng. 2021.  Online Disease Diagnosis with Inductive Heterogeneous Graph Convolutional Networks. In _Proceedings of the Web Conference 2021_. 3349–3358.
  * Wei et al. (2018) Zhongyu Wei, Qianlong Liu, Baolin Peng, Huaixiao Tou, Ting Chen, Xuan-Jing Huang, Kam-Fai Wong, and Xiang Dai. 2018\.  Task-oriented dialogue system for automatic diagnosis. In _Proceedings of the 56th Annual Meeting of the Association for Computational Linguistics (Volume 2: Short Papers)_. 201–207.
  * Wu and Goodman (2018) Mike Wu and Noah Goodman. 2018.  Multimodal generative models for scalable weakly-supervised learning. In _Advances in Neural Information Processing Systems_. 5575–5585.
  * Xia et al. (2020) Yuan Xia, Jingbo Zhou, Zhenhui Shi, Chao Lu, and Haifeng Huang. 2020.  Generative Adversarial Regularized Mutual Information Policy Gradient Framework for Automatic Diagnosis. In _Proceedings of the AAAI Conference on Artificial Intelligence_ , Vol. 34. 1062–1069.
  * Xu et al. (2019) Lin Xu, Qixian Zhou, Ke Gong, Xiaodan Liang, Jianheng Tang, and Liang Lin. 2019\.  End-to-end knowledge-routed relational dialogue system for automatic diagnosis. In _Proceedings of the AAAI Conference on Artificial Intelligence_ , Vol. 33. 7346–7353.
  * Xu et al. (2013) Zhixiang Xu, Matt Kusner, Kilian Weinberger, and Minmin Chen. 2013\.  Cost-sensitive tree of classifiers. In _International conference on machine learning_. 133–141.
  * Zakim et al. (2008) David Zakim, Niko Braun, Peter Fritz, and Mark Dominik Alscher. 2008\.  Underutilization of information and knowledge in everyday medical practice: Evaluation of a computer-based solution.  _BMC Medical Informatics and Decision Making_ 8, 1 (2008), 50\.
  * Zeng et al. (2004) Qing T Zeng, Sandra Kogan, Robert M Plovnick, Jonathan Crowell, Eve-Marie Lacroix, and Robert A Greenes. 2004.  Positive attitudes and failed queries: an exploration of the conundrums of consumer health information retrieval.  _International journal of medical informatics_ 73, 1 (2004), 45–55.
  * Zuccon et al. (2015) Guido Zuccon, Bevan Koopman, and Joao Palotti. 2015\.  Diagnose this if you can. In _European Conference on Information Retrieval_. Springer, 562–567.

##  Appendix A Overview of BSODA

Please refer to Algorithm 1.

0: Training dataset $\mathbf{X}$; Test dataset $\mathbf{X}^{*}$ with no
observations; Indices $D$ of diseases.

1: Train knowledge-guided self-attention model on complete $\mathbf{X}$ and
train VAE with PoE encoder on $\mathbf{X}$ with a fraction of features dropped
at random.

2: Actively inquire symptom to estimate $\mathbf{x}_{D}$ by $p_{D}(\cdot)$ for
each test point:

for each test point in $\mathbf{X}^{*}$ do

$\mathbf{x}_{O}\leftarrow\\{x_{i}\\}$ (initial symptom $i$)

$S_{c}\leftarrow S_{i}$

repeat

for every symptom $s$ in $S_{c}$ do

// Two-step sampling

$\hat{\mathbf{z}}\sim q_{\Phi}(\mathbf{z}|\mathbf{x}_{O})$

$p_{VAE}\leftarrow p_{\theta}(x_{s}|\hat{\mathbf{z}})$

$\hat{x}_{s}\leftarrow 1$ // Focusing on positive symptoms

$p_{Diag}\leftarrow p_{D}({\hat{x}_{s},\mathbf{x}_{O}})$

$p(\mathbf{x}_{D},x_{s}|\mathbf{x}_{O})\leftarrow p_{VAE}\cdot p_{Diag}$ (Sec.
3.2.2)

// Approximation of the reward

$\\{\hat{R}_{j}^{k}\\}_{j=0,1;k=1,...,|D|}\leftarrow\\{0\\}$

for disease $d$ with probability in the 90 percent of the top probability in
$p_{Diag}$ and more than $1/|D|$ do

$\hat{R}_{1}^{d}\leftarrow$ Eqn. 3

end for

$\hat{R}_{s}\leftarrow$ a weighted average by
$p(\mathbf{x}_{D},x_{s}|\mathbf{x}_{O})$ and $\\{\hat{R}_{j}^{k}\\}$

end for

$t\leftarrow\text{argmax}_{t}\hat{R}_{t}$, inquire symptom $t$

$\mathbf{x}_{O}\leftarrow\mathbf{x}_{t}\cup\mathbf{x}_{O}$

if $\mathbf{x}_{t}$ is positive then

$S_{c}\leftarrow S_{t}\cap S_{c}$ // Filtering candidate symptoms

end if

$\hat{\mathbf{z}}\sim
q_{\Phi}(\mathbf{z}|\mathbf{x}_{O}),\hat{\mathbf{x}}_{U}\sim
p_{\theta}(\mathbf{x}_{U}|\hat{\mathbf{z}})$

$\mu_{p_{D}}\leftarrow\mathbb{E}[p_{D}(\mathbf{x}_{O}\cup\hat{\mathbf{x}}_{U})],\sigma_{p_{D}}\leftarrow\sqrt{\mathrm{Var}[p_{D}(\mathbf{x}_{O}\cup\hat{\mathbf{x}}_{U})]}$

$d\leftarrow\text{argmax}_{d}\ \mu_{p_{D}}^{(d)}$, report disease $d$

until reaching maximum number of inquiries $N_{T}$ or $\forall j\neq
d,\mu_{p_{D}}^{(d)}>\mu_{p_{D}}^{(j)}+3\sigma_{p_{D}}^{(j)}$

end for

Algorithm 1 Overview of BSODA.

##  Appendix B Study with Hybrid Data Types

Kachuee et al. (2018b) compiled a diabetes dataset from the national health
and nutrition examination survey (NAHNES) data (nhr, 2018). It consists of 45
features, such as demographic information, lab results (total cholesterol,
triglyceride, etc.), examination data (weight, height, etc.), and
questionnaire answers (smoking, alcohol, etc.). In contrast to other symptom
checking tasks, these features are continuous, and each feature is associated
with a cost suggested by an expert in medical studies. Finally, the fasting
glucose values were used to define three classes based on standard threshold
values: normal, pre-diabetes, and diabetes. The dataset consists of 92,062
instances and is randomly split into three subsets, including 15% for testing,
15% for validation, and the rest 70% for training.

Figure 3 shows an evaluation on the diabetes dataset that has hybrid types of
features. We implemented BSODA without using any speedup techniques to handle
these continuous features. We set $N_{M}$ to 50 for BSODA and did not restrict
the maximum number of inquiries. When the maximum number of inquiries is set
to $i$, the $i$-th point on the curve plots the average cost (x-axis) and the
accuracy (y-axis). We observe that BSODA achieves a superior accuracy while
having a lower acquisition cost than OL (Kachuee et al., 2018b), which is an
RL baseline based on deep Q-learning and designed for cost-sensitive feature
selection. Except for the diabetes dataset, OL is not optimal for the other
symptom checking datasets because their feature costs are unknown. Thus, in
addition to symptom checking, BSODA has the ability to handle online diagnosis
problems involving hybrid data types.

Figure 3. Accuracy vs acquisition cost for BSODA, OL (Kachuee et al., 2018b)
and a random method on the diabetes dataset.

BSODA achieves a superior accuracy with a lower acquisition cost than OL.

Figure 4. Partial distributions for information rewards among several symptoms
in first three inquiries on the SymCAT 200-disease task. The symptoms with
maximum reward are marked with star. The initial symptom is ”Nasal
Congestion”. ”SB” means ”Shortness of Breath”. ”CO” means ”Cough”. ”FV” means
”Fever”. ”ST” means ”Sore Throat”. ”EP” means ”Ear Pain”. ”HA” means
”Headache”. ”WZ” means ”Wheezing”. ”PE” means ”Pain in Eye”. ”IE” means
”Itchiness of Eye”.

##  Appendix C Case Study for Speedup Techniques

We conducted a case study to show that the proposed speedup techniques have
little effect on the information reward and the actual symptoms queried by
BSODA. Figure 4 shows the partial distributions of information rewards for
several symptoms in the first three inquiries for a case. The test case comes
from the synthetic SymCAT 200-disease task and is a simulated patient with
_Common Cold_. We limit the maximum number of inquiries to 15 and $N_{M}$ to
100.

Table 7. Accuracy (%) with various model settings for BSODA on the SymCAT 200-disease task. Settings | Top1 | Top3 | Top5 | #Rounds
---|---|---|---|---
BSODA | $\mathbf{55.65\pm 0.25}$ | $\mathbf{80.71\pm 0.26}$ | $\mathbf{89.32\pm 0.29}$ | $12.02\pm 0.06$
BSODA w/o Diag. | $47.37\pm 0.50$ | $72.28\pm 0.57$ | $82.02\pm 0.81$ | $14.31\pm 0.03$
BSODA w/o Diag. & Inq. | $30.36\pm 0.26$ | $56.42\pm 0.36$ | $70.70\pm 0.36$ | $4.21\pm 0.10$

##  Appendix D Ablation Study of Models

We performed an ablation study to investigate the effect of model design on
the synthetic SymCAT 200-disease task. We set $N_{M}$ to 100 and the maximum
number of inquiries to 15. In Table 7, ”Inq.” and ”Diag.” refer to the
symptom-inquiry and disease-diagnosis model designs in BSODA, respectively.
”BSODA w/o Diag.” means we replace the predictive distributions $p_{D}(\cdot)$
from the diagnosis model in the two-step sampling strategy and the diagnostic
method with ones produced by the VAE model. ”BSODA w/o Inq.” means we do not
use the PoE encoder. ”BSODA w/o Diag. & Inq.” is simply the EDDI framework.
The results show that the PoE encoder improves the accuracy significantly by
at least 12% and the average number of inquiries increases. With the same
termination criterion and the maximum number of inquiries, the PoE encoder
provides a better and more patient model, which does not diagnose hastily.
Besides, the diagnosis model is critical in effective diagnosis with fewer
inquiries and providing a narrower confidence interval, indicating a more
stable result.


section]

# Gradient representations in ReLU networks as similarity functions111Accepted
at the 29th European Symposium on Artificial Neural Networks, Computational
Intelligence and Machine Learning (ESANN 2021), 6-8 October 2021. Source code
with additional experiments: https://github.com/danielracz/gradsim

Dániel Rácz1,3 and Bálint Daróczy2,3

1 - Eötvos Loránd University BD was supported by MIS “Learning from Pairwise
Comparisons” of the F.R.S.-FNRS and by MTA Premium Postdoctoral Grant 2018.
Institute of Mathematics
Pázmány Péter sétány 1/C  H-1117  Budapest  Hungary
2 - Université catholique de Louvain  INMA/ICTEAM
Avenue Georges Lemaître 4  B-1348  Louvain-la-Neuve  Belgium
3 - Institute for Computer Science and Control  SZTAKI/ELKH
Kende utca 13-17  H-1111  Budapest  Hungary
e-mail<EMAIL_ADDRESS><EMAIL_ADDRESS>
###### Abstract

Feed-forward networks can be interpreted as mappings with linear decision
surfaces at the level of the last layer. We investigate how the tangent space
of the network can be exploited to refine the decision in case of ReLU
(Rectified Linear Unit) activations. We show that a simple Riemannian metric
parametrized on the parameters of the network forms a similarity function at
least as good as the original network and we suggest a sparse metric to
increase the similarity gap.

##  1 Introduction

We consider feed-forward neural networks with ReLU activations to examine how
the network’s final representation space connected to the gradient structure
of the network. Our motivation is twofold: recently discovered knowledge about
ReLU networks [1, 2, 3] and recent results about higher order optimization
methods [4]. In a way, many of the existing machine learning problems can be
investigated as statistical learning problems, therefore information geometry
[5] plays an important role. It was shown in [6] that over the parameter space
of a neural network we can often determine a Riemannian manifold based on an
error or loss function, moreover the tangent bundle on specific Riemannian
metrics, e.g. Fisher information, has unique invariance properties [7, 8].

Our main hypothesis is that in feed-forward ReLU networks we can utilize the
relation of the parameters and the output due the homogeneity property of the
activation functions therefore we investigate the gradient structure of the
network output w.r.t. the parameters and exploit the space induced by the
partial derivatives together with a metric as a representation of data points.
Albeit the inner product space of the tangent bundle is quadratic, there are
well-defined underlying structures in the tangent space specific to ReLU
networks. In this paper we introduce several similarity functions [9] based on
a block-diagonal, sparse metric and we inspect how they relate to the
similarity induced by the network itself by measuring the similarity gap, the
difference in expected similarity between point pairs with the same label and
point pairs with different labels.

###  1.1 Related work

The geometrical properties of the underlying loss manifold of neural networks
was used as a general framework for optimization in classification [6] and for
generative models [10]. Furthermore Martens and Grosse [4] approximated
Amari’s natural gradient [5] for feed-forward neural networks with block
partitioning the Fisher information matrix by exploiting the structure of the
network. The partial gradients were used as representations in classification
given a known generative probability density function, e.g. through an
approximated Fisher kernel, in [11] where the authors approximated the metric
with the diagonal of the Fisher information matrix, closed formula for
Gaussian mixtures were proposed in [8]. In [12] the authors learned a metric
with a neural network over the partial gradients w.r.t. the parameters. Not
only the partial gradient w.r.t. parameters carry important properties e.g. in
[13] the authors experimented with the norm of the input-output sensitivity,
the Frobenius norm of the Jacobian matrix of the output w.r.t. input in case
of simple architectures.

Recent results show that under simple presumptions the maximal “capacity”
(representational power) of deep ReLU networks is related to the arrangement
of the polytopes in the input space [1] and to the properties of transition
between linear regions [14, 3] instead of the exact number of polytopes with
non-zero volume.

Our last ingredient is Balcan and Blum’s theory of similarity functions [9],
which we will use as our foundation to show that we can define similarity
functions in the gradient space that are at least as good similarity functions
as the output of the network.

###  1.2 Notations

Let $f:\mathbb{R}^{d}\times\mathbb{R}^{|\theta|}\rightarrow\mathbb{R}$ be a
feed-forward ReLU network, as a function of both the input data and the
weights, containing $L$ hidden layers, $N$ neurons trained for solving a
classification problem $P$ for some data distribution $D$. The activation
function is ReLU: $h(x)=\max(0,x),x\in\mathbb{R}$. Let $N^{k}$ denote the
number of neurons in the $k$-th layer and let
$h^{k}_{\theta}(x)\in\mathbb{R}^{N^{k}}$ denote the output of the $k$-th layer
for a given input vector $x$ and a fixed set of weights $\theta$. Denoting the
weight vector after the last hidden layer, the discriminative layer, by
$\theta_{L}$ we have $f_{\theta}(x)=\theta_{L}^{T}h^{L}_{\theta}(x)$. We
denote the set of parameters as $\theta$ and the parameters of the $k$-th
layer with $\theta_{k}$. The actual output of our model is
$sgn(f_{\theta}(x))=sgn(\theta_{L}^{T}h_{\theta}(x))$, i.e the network can be
interpreted as a linear separator acting on
$Im(h_{\theta})\subseteq\mathbb{R}^{N^{L}}$, the image of the nonlinear
mapping $h_{\theta}:\mathbb{R}^{d}\rightarrow\mathbb{R}^{N^{L}}$. Following
[1], we define an activation pattern with assigning a sign to each neuron in
the network, $A=\\{a_{l};l=1,..,N\\}\in\\{-1,1\\}^{N}$. For a particular input
we refer $A(x;\theta)=\\{sign([h_{\theta}(x)]_{l};l=1,..,N\\}$ as the
activation pattern assigned to an input $x$. An activation region with the
corresponding fixed $\theta$ and $A$ is defined as
$R(A;\theta):=\\{x\in\mathbb{R}^{d}|sign([h_{\theta}(x)]_{l})=a_{l}\\}$, the
set of input assigned to the same activation pattern. In comparison, linear
regions are the input regions where the function defines different linear
regions.We define tangent vectors as the change in the output with a
directional derivative of $f_{\theta}(x)$ in the direction of
$g_{\theta}(x_{i})=\frac{\partial
f_{\theta}(x)}{\partial\theta}\big{|}_{x=x_{i}}$:
$(D_{g_{\theta}(x_{i})}f)(\theta)=\frac{d}{dt}[f_{\theta}(x)+tg_{\theta}(x_{i}))]|_{t=0}$.
We will refer $\nabla:\mathbb{R}^{d}\rightarrow\mathbb{R}^{|\theta|}$ as the
tangent mapping of input at $\theta$:
$\nabla_{\theta}f_{\theta}(x):=\frac{\partial{f_{\theta}(x)}}{\partial{\theta}}$.
We denote the partial gradient vectors with
$g_{\theta}(x)\in\mathbb{R}^{|\theta|}$.

##  2 Similarity functions

First, we state the following definition from [9] which will be the measure of
our similarity functions:

######  Definition .

A $K(x,y)$ is a weakly $\gamma$-good similarity function for a learning
problem P if:

$E_{x,y\sim P}[K(x,y)|l(x)=l(y)]\geq E_{x,y\sim P}[K(x,y)|l(x)\neq
l(y)]+\gamma.$ (1)

Note, the strong version of good similarity function says that a $K(x,y)$ is a
good similarity function if $1-\epsilon$ probability mass of the examples are
on average more similar to examples of the same category than to examples of
another category. We consider the weakly good similarity because in this case
we may investigate how the gap changes in case of a finite sample. We refer to
$\gamma$ as the similarity gap.

We define the similarity function of the discriminative layer as
$K_{f^{L}_{\theta}}(x,y)=h_{\theta}^{L}(x)^{T}h_{\theta}^{L}(y)$ and the
similarity function of the network as
$K_{f_{\theta}}(x,y)=f_{\theta}(x)f_{\theta}(y)$. Now, let us consider the
inner product in case of the tangent mapping. Before we argue about the metric
in the tangent space we take a closer look on the structure of the Hessian in
case of feed-forward ReLU networks. By definition the Hessian $H_{\theta}$ of
the function induced by the network for a single output is a
$|\theta|\times|\theta|$ sized matrix with
$[H_{\theta}(x)]_{ij}=\frac{\partial^{2}[f_{\theta}(x)]}{\partial\theta_{i}\partial\theta_{j}},\forall
i,j\in\\{1,..,|\theta|\\}$ which in our case is equal to the sum over the
paths from the input to the output node where both $\theta_{i}$ and
$\theta_{j}$ are present therefore if we split the Hessian into sub-blocks,
where only elements from the same layer are present, these sub-blocks will be
diagonal. This property suggest us to seek for a metric where the parameters
of a layer in these sub-blocks may collide if necessary. Therefore we suggest
a simple metric $M_{\theta}\in\mathbb{R}^{|\theta|^{2}}$ with
$[M_{\theta}]_{ij}=\delta_{l(i),l(j)}\theta_{i}\theta_{j}$ where
$\delta_{i,j}$ is one if $i=j$ and zero otherwise and $l(i)$ is the index of
the layer of the $i$-th parameter. We define our new similarity function, the
block diagonal similarity as

$\displaystyle
K_{M_{\theta}}(x,y)=g_{\theta}(x)^{T}M_{\theta}g_{\theta}(y)=\sum_{i,j\in\theta\times\theta}\delta_{l(i),l(j)}\theta_{i}\theta_{j}[g_{\theta}(x)]_{i}[g_{\theta}(y)]_{j}$

where we utilized that the metric has a Cholesky decomposition as it is
positive semi-definite. One of the most important properties of the block-
diagonal similarity function is that in case of a single output the similarity
function is equal, up to a constant, to the network output as

$\displaystyle\sum_{i,j\in\theta\times\theta}\delta_{l(i),l(j)}\theta_{i}\theta_{j}[g_{\theta}(x)]_{i}[g_{\theta}(y)]_{j}=\sum_{i,j\in\theta\times\theta}\delta_{l(i),l(j)}f_{\theta}(x)f_{\theta}(y)$
$\displaystyle=\sum_{k=1}^{L}K_{f_{\theta}(x,y)}\approx
O(K_{f_{\theta}}(x,y)).$

Since
$K_{f_{\theta}}(x,y)=\sum_{i\in\theta^{L}}\theta^{2}_{i}[h_{\theta}(x)]_{i}[h_{\theta}(y)]_{i}$
the following holds for the block-diagonal similarity:

$\displaystyle\omega_{L}O(K_{f^{L}_{\theta}}(x,y))\leq
K_{M_{\theta}}(x,y)\leq\omega_{L}^{*}O(K_{f^{L}_{\theta}}(x,y))$

where $\omega_{L}=\min_{i\in\theta^{L}}\theta_{i}^{2}$ and
$\omega_{L}^{*}=\max_{i\in\theta^{L}}\theta_{i}^{2}$. Note, the similarity
values are not necessary positive.

###  2.1 Sparse metric

Now, we decompose the gap and suggest a modified metric to increase the gap of
the block-diagonal similarity function. As the norm of the gradient vectors
highly affects the gap, we normalize to avoid linear increases in the gap.
First, for simplicity we estimate the expectation $\mathbf{E}_{x,y\sim
P}[K_{M_{\theta}}(x,x^{\prime})|l(x)=l(y)]$ for a single layer and for a
single label therefore the expectation is

$\displaystyle\sum_{i=1}^{|\theta|}\sum_{j=1}^{|\theta|}\mathbf{E}_{x,x^{\prime}\sim
P}[\frac{[g_{\theta}(x)]_{i}\theta_{i}\theta_{j}[g_{\theta}(x^{\prime})]_{j}}{\|M_{\theta}^{1/2}g_{\theta}(x)\|\|M_{\theta}^{1/2}g_{\theta}(x^{\prime})\|}]$
$\displaystyle\approx\frac{1}{|T^{(+)}|^{2}}\sum_{t_{1},t_{2}\in T^{(+)}\times
T^{(+)}}\frac{1}{\|\hat{g_{\theta}}(x_{t_{1}})\|\|\hat{g}_{\theta}(x_{t_{1}})\|}\sum_{i=1}^{|\theta|}\sum_{j=1}^{|\theta|}[g_{\theta}(x_{t_{1}})]_{i}\theta_{i}\theta_{j}[g_{\theta}(x_{t_{2}})]_{j}$
$\displaystyle=\sum_{i,j\in\theta\times\theta}\frac{1}{|T^{(+)}|^{2}}\sum_{t_{1},t_{2}\in
T^{(+)}\times
T^{(+)}}\frac{[g_{\theta}(x_{t_{1}})]_{i}\theta_{i}\theta_{j}[g_{\theta}(x_{t_{2}})]_{j}}{\|\hat{g_{\theta}}(x_{t_{1}})\|\|\hat{g}_{\theta}(x_{t_{1}})\|}=\sum_{i,j\in\theta\times\theta}[\hat{\psi}^{+}_{M_{\theta}}]_{ij}$

where we denote the set of examples with label “$+$” as $T^{(+)}$, a subset of
known examples $T$ and $M_{\theta}^{1/2}g_{\theta}(x)$ with
$\hat{g_{\theta}}(x)$. Similar calculation can be applied if the labels are
different thus

$\displaystyle\gamma_{M_{\theta}}$ $\displaystyle=\mathbf{E}_{x,y\sim
P}[K_{M_{\theta}}(x,x^{\prime})|l(x)=l(y)]-\mathbf{E}_{x,y\sim
P}[K_{M_{\theta}}(x,x^{\prime})|l(x)\neq l(y)]$
$\displaystyle\approx\sum_{i,j\in\theta\times\theta}[\hat{\psi}^{(+)}_{M_{\theta}}]_{ij}-[\hat{\psi}^{(-)}_{M_{\theta}}]_{ij}=\sum_{i,j\in\theta\times\theta}[\hat{\psi}_{M_{\theta}}]_{ij}.$

Therefore if we consider multiple layers, where $[M_{\theta}]_{ij}=0$ if
$l(i)\neq l(j)$, the similarity gap can be approximated as

$\displaystyle\gamma_{M_{\theta}}\approx\sum_{i,j\in\theta\times\theta}\delta_{l(i),l(j)}[\hat{\psi}_{M_{\theta}}]_{ij}=\sum_{k\in{1,...,L}}\sum_{i,j\in\theta^{k}\times\theta^{k}}[\hat{\psi}_{M_{\theta}}]_{ij}.$

Note, with $|T|\rightarrow\infty$ the error of our approximation converges to
zero with probability one, moreover this convergence is true for every element
in $\hat{\psi}_{M_{\theta}}$. An additional consequence that we can partition
the elements of the similarity gap into disjoint sets and examine each
independently. Observe, the elements of $\hat{\psi}\in\mathbb{R}$ therefore
let us define for every pair their “importance” in the gap as
$\textit{imp}_{i,j}=\max\\{0,[\hat{\psi}_{M_{\theta}}]_{ij}\\}$. By setting
the elements in the metric with low or negative importance we can define a
sparse block-diagonal similarity with an additional step, normalization.
However our new metric can have the same dimensions as the original block-
diagonal. Thus we can argue that according to
$\textit{imp}_{i}=\sum_{j}\textit{imp}_{ij},\forall i\in\theta$ we can select
the most important parameters and delete rows and columns associated with less
important parameters and form the elementwise block-diagonal metric.

So far we assumed arbitrary input and do not take advantage of the gradient
graph of the network. To show a case when the normalized sparse block-diagonal
similarity has better similarity gap than the original normalized block-
diagonal similarity we will assume that our input is a subgaussian random
vector [15] with zero mean and variance one, e.g. the input is element-wise
standard normalized. Additionally, observe that in ReLU networks the partial
gradients can be expressed as $g_{\theta}(x)=S_{\theta}(x)x$ where
$S_{\theta}(x)_{i,j}=\frac{\partial^{2}f_{\theta}(x)}{\partial\theta\partial
x}\rvert_{\theta,x}$ a $d\times|\theta|$ sized matrix and inside an activation
region $S_{\theta}(x)=S_{\theta}(A(x))$ is identical for each $x\in A$ thus
$g_{\theta}(x)=S_{\theta}(A(x))x$. Due the complexity of the proof and the
page limit we only mention that, following Theorem 2.1 in [16], the norm of
the elementwise sparse vector for points in the activation region $A$ is
concentrated as for all $\delta>0$

$\displaystyle
P(\|S^{*}_{\theta}(A(x))x\|_{2}^{2}>Tr[S^{*}_{\theta}(A(x))S^{*}_{\theta}(A(x))^{T}](1+4\delta);x\in
A)\leq e^{-\delta}$

where $S^{*}_{\theta}(A(x))$ is the same as $S_{\theta}(A(x))$ without the
removed rows and columns and therefore the similarity gap in case of the
normalized elementwise sparse block-diagonal similarity is related to the
ratio
$\frac{Tr[S_{\theta}(A(x))S_{\theta}(A(x))^{T}]}{Tr[S^{*}_{\theta}(A(x))S^{*}_{\theta}(A(x))^{T}]}$.

We experimented on the CIFAR-10 dataset [17] with a simple feed-forward
network with five layers. We ranked the parameters per layer according the
elementwise block-diagonal metric and deleted the parameters with low
“importance”. Results in Fig. 1 indicate that the gap can be increased but
further, more detailed experimentation needed to understand how the gap is
actually increasing.

Figure 1: Distribution of the network output, the elementwise sparse block-
diagonal and the block-diagonal similarity values.

##  3 Conclusions

In this paper we defined similarity functions over ReLU networks based on the
gradient structure and investigated the similarity gap. Furthermore, we
introduced a measure to rank the parameter pairs in the network according to
their importance in the similarity gap. In the future we plan to extend our
work to other network structures as our findings were limited to feed-forward
fully connected ReLU networks.

## References

  * [1] Boris Hanin and David Rolnick.  Complexity of linear regions in deep networks.  In International Conference on Machine Learning, pages 2596–2604, 2019.
  * [2] Xiao Zhang and Dongrui Wu.  Empirical studies on the properties of linear regions in deep neural networks.  arXiv preprint arXiv:2001.01072, 2020.
  * [3] Boris Hanin, Ryan Jeong, and David Rolnick.  Deep relu networks preserve expected length.  arXiv preprint arXiv:2102.10492, 2021.
  * [4] James Martens and Roger Grosse.  Optimizing neural networks with kronecker-factored approximate curvature.  In International conference on machine learning, pages 2408–2417, 2015.
  * [5] Shun-ichi Amari.  Neural learning in structured parameter spaces-natural Riemannian gradient.  In NIPS, pages 127–133. Citeseer, 1996.
  * [6] Yann Ollivier.  Riemannian metrics for neural networks i: feedforward networks.  Information and Inference: A Journal of the IMA, 4(2):108–153, 2015\.
  * [7] N. N. C̆encov.  Statistical decision rules and optimal inference.  American Mathematical Society, 53, 1982.
  * [8] Florent Perronnin and Christopher Dance.  Fisher kernels on visual vocabularies for image categorization.  In Computer Vision and Pattern Recognition, 2007. CVPR’07. IEEE Conference on, pages 1–8. IEEE, 2007.
  * [9] Maria-Florina Balcan, Avrim Blum, and Nathan Srebro.  A theory of learning with similarity functions.  Machine Learning, 72(1-2):89–112, 2008.
  * [10] Salah Rifai, Yann Dauphin, Pascal Vincent, Yoshua Bengio, and Xavier Muller.  The manifold tangent classifier.  In NIPS, volume 271, page 523, 2011.
  * [11] Tommi S Jaakkola, David Haussler, et al.  Exploiting generative models in discriminative classifiers.  Advances in neural information processing systems, pages 487–493, 1999.
  * [12] Balint Daroczy, Rita Aleksziev, and András Benczúr.  Tangent space separability in feedforward neural networks.  In proc. of BFO at NeurIPS 2019.
  * [13] Roman Novak, Yasaman Bahri, Daniel A Abolafia, Jeffrey Pennington, and Jascha Sohl-Dickstein.  Sensitivity and generalization in neural networks: an empirical study.  arXiv preprint arXiv:1802.08760, 2018.
  * [14] Maithra Raghu, Ben Poole, Jon Kleinberg, Surya Ganguli, and Jascha Sohl Dickstein.  On the expressive power of deep neural networks.  In Proc. of ICML’17.
  * [15] Tor Lattimore and Csaba Szepesvári.  Bandit algorithms.  Cambridge University Press, 2020.
  * [16] Daniel Hsu, Sham Kakade, Tong Zhang, et al.  A tail inequality for quadratic forms of subgaussian random vectors.  Electronic Communications in Probability, 17, 2012.
  * [17] Alex Krizhevsky, Geoffrey Hinton, et al.  Learning multiple layers of features from tiny images.  Technical Report TR-2009, University of Toronto, Toronto, 2009.