AppReferences

# The Medkit-Learn(ing) Environment: Medical Decision Modelling through
Simulation

Alex J. Chan
University of Cambridge
Cambridge, UK
<EMAIL_ADDRESS>
&Ioana Bica
University of Oxford
Oxford, UK
<EMAIL_ADDRESS>
&Alihan Hüyük
University of Cambridge
Cambridge, UK
<EMAIL_ADDRESS>
&Daniel Jarett
University of Cambridge
Cambridge, UK
<EMAIL_ADDRESS>
&Mihaela van der Schaar
University of Cambridge
Cambridge, UK
<EMAIL_ADDRESS>

###### Abstract

Understanding decision-making in clinical environments is of paramount
importance if we are to bring the strengths of machine learning to ultimately
improve patient outcomes. Several factors including the availability of public
data, the intrinsically offline nature of the problem, and the complexity of
human decision making, has meant that the mainstream development of algorithms
is often geared towards optimal performance in tasks that do not necessarily
translate well into the medical regime; often overlooking more niche issues
commonly associated with the area. We therefore present a new benchmarking
suite designed specifically for medical sequential decision making: the
Medkit-Learn(ing) Environment, a publicly available Python package providing
simple and easy access to high-fidelity synthetic medical data. While
providing a standardised way to compare algorithms in a realistic medical
setting we employ a generating process that disentangles the policy and
environment dynamics to allow for a range of customisations, thus enabling
systematic evaluation of algorithms’ robustness against specific challenges
prevalent in healthcare.

##  1 Introduction

Figure 1: _Overview of Medkit._ The central object in Medkit is the _scenario_
, made up of a _domain_ , _environment_ , and _policy_ which fully defines the
synthetic setting. By disentangling the environment and policy dynamics,
Medkit enables us to simulate decision making behaviours with various degrees
of Markovianity, individual consistency, bounded rationality and variation in
practice. An example scenario is highlighted: ICU patient trajectories with
environment dynamics modelled by a sequential VAE and a customised clinical
policy. The output from Medkit will be a batch dataset that can be used for
training and evaluating methods for modelling human decision-making.

Modelling human decision-making behaviour from observed data is a principal
challenge in understanding, explaining, and ultimately improving existing
behaviour. This is the business of decision modelling, which includes such
diverse subfields as reward learning [1, 55, 35, 29], preference elicitation
[32], goal inference [54], interpretable policy learning [28], and policy
explanation [10]. Decision modelling is especially important in medical
environments, where learning interpretable representations of existing
behaviour is the first crucial step towards a more transparent account of
clinical practice.

For research and development in clinical decision modelling, it is important
that such techniques be validated robustly—that is, operating in different
medical domains, guided by different environment dynamics, and controlled by
different behavioural policies. This is difficult due to three reasons. First,
the very nature of healthcare data science is that any learning and testing
must be carried out entirely offline, using batch medical datasets that are
often limited in size, variety, and accessibility [25, 39]. Second, directly
using methods for time-series synthetic data generation is inadequate, as they
simply learn sequential generative models to replicate existing data, making
no distinction between environment and policy dynamics [20, 53, 14]. Because
the environment and policy dynamics are entangled, such models do not allow
for customisation of the decision making policy and thus cannot be used for
evaluating methods for understanding human decision-making. Third, while
various hand-crafted medical simulators have been proposed as stylised proofs-
of-concept for research [43, 24, 16, 22], they often make unrealistic
assumptions and simplifications that are unlikely to transfer well to any more
complicated real-world setting. Moreover, these simulators do not directly
allow obtaining offline data from different types of policy parameterisations.

Desiderata: It is clear that what is desired, therefore, is a tool that
supports: (1) a variety of realistic environment models—learned from actual
data, to reflect real medical settings), thus allowing simulation of (2) a
variety of expressive and customisable policy models that represent complex
human decision-behaviours; as well as (3) ensuring that the environment and
policy components are disentangled—hence independently controllable.

Contributions: We present the Medkit-Learn(ing) Environment (“Medkit”), a
toolbox and benchmarking suite designed precisely for machine learning
research in decision modelling. Fulfilling all of the above key criteria,
Medkit seeks to enable advances in decision modelling to be validated more
easily and robustly—by enabling users to obtain batch datasets with known
ground-truth policy parameterisations that simulate decision making behaviours
with various degrees of Markovianity, bounded rationality, confounding,
individual consistency and variation in practice. Moreover, to facilitate
efficient progress in this area of understanding human decision-making, we
have built Medkit to be freely accessible and transparent in data simulation
and processing to enable reproducibility and fair benchmarking.

##  2 The Medkit-Learn(ing) Environment

Figure 1 gives an overview of the structure of Medkit, demonstrating a modular
design philosophy to enable an ever-growing offering of scenarios as new
algorithms and data become available. Medkit is publicly available on GitHub:
https://github.com/XanderJC/medkit-learn. Written in Python and built using
PyTorch [46] for a unified model framework, Medkit takes advantage of the
OpenAI gym [9] interface for live interaction but has otherwise minimal
dependencies.

###  2.1 Simulating Medical Datasets for Modelling Sequential Decision-Making

Our aim is to build generative models for the decision making process, that
allow for full customisation of: (1) the environment dynamics, that model how
the patient’s state changes; and (2) the policy dynamics through which users
can specify complex decision making behaviours.

Formally, we define a _scenario_ as a tuple ($\Omega,\mathcal{E},\pi$), which
represents the central component of Medkit that fully defines a generative
distribution over synthetic data. A scenario comprises a medical _domain_ ,
$\Omega$ (e.g. the ICU); an _environment_ dynamics model for sequential
observations, $\mathcal{E}$ (e.g. a linear state space model); and a _policy_
mapping the observations to actions, $\pi$ (e.g. a decision tree).

Let $\vec{x}_{T}=x_{s}\cup\\{x_{t}\\}_{t=1}^{T}$ be the individual patient
trajectories and let $\vec{y}_{T}=\\{y_{t}\\}_{t=1}^{T}$ be the clinical
interventions (actions). Here $x_{s}\in\mathcal{X}_{s}$ is a multi-dimensional
vector of _static_ features of the patient, e.g. height, various
comorbidities, or blood type - while $x_{t}\in\mathcal{X}$ a multi-dimensional
vector representing _temporal_ clinical information such as biomarkers, test
results, and acute events. Additionally $y_{t}\in\mathcal{Y}$ is a further
possibly multi-dimensional vector representing the actions and interventions
of medical professionals, for example indicators for ordering tests and
prescribing treatments.

We propose modelling the joint distribution of the patient features and
clinical interventions $p(\vec{x}_{T},\vec{y}_{T})$ using the following
factorisation:

$\displaystyle p(\vec{x}_{T},\vec{y}_{T})$
$\displaystyle=\underbrace{P_{\mathcal{E}}^{\Omega}(x_{s})P_{\mathcal{E}}^{\Omega}(x_{1}|x_{s})\textstyle\prod_{t=2}^{T}P_{\mathcal{E}}^{\Omega}(x_{t}|f_{\mathcal{E}}(\vec{x}_{t-1},\vec{y}_{t-1}))}_{\text{Environment}}$
(1)
$\displaystyle\quad\quad\times\underbrace{Q_{\pi}^{\Omega}(y_{1}|x_{s},x_{1})\textstyle\prod_{t=2}^{T}Q_{\pi}^{\Omega}(y_{t}|g_{\pi}(\vec{x}_{t},\vec{y}_{t-1}))}_{\text{Policy}},$

where the distributions $P^{\Omega}_{\mathcal{E}}(\cdot)$ specify the
transition dynamics for domain $\Omega$ and environment $\mathcal{E}$ and
$Q^{\Omega}_{\pi}$ represents the policy for making clinical interventions in
domain $\Omega$, thus defining the decision making behaviour. Note that the
patient trajectories and interventions depend on the entire history of the
patient. The functions $f$ and $g$ are modelled to be distinct such that the
focus on the past represented in the conditional distributions may be
different for both the policy and the environment. While this factorisation
allows for enough flexibility, we will often make use of a representation
given in the graphical model of Figure 2 which includes a hidden state of the
environment $z_{t}$ to be the underlying driver of the data. Note this is even
more general and we recover equation (1) if $z_{t}$ is simply set as $x_{t}$.

Figure 2: _Graphical model of the generative process we consider._ Usually
there will be some hidden state of the environments that drives the actions
and observations seen.

With the factorisation proposed in Equation 1 we notice a clear separation
between the _environment_ and _policy_ dynamics so that they can be modelled
and learnt separately, with the _domain_ defining the “meta-data” such as the
spaces $\mathcal{X}_{s},\mathcal{X},$ and $\mathcal{Y}$. This disentanglement
between the environment and policy components is not possible in current
synthetic data generation methods (as we explore in section 3). A corollary to
this makes for a useful feature of Medkit - that we can then mix and match
elements of the tuple to create a variety of different scenarios that can be
extended easily in the future when new models or data become available. This
not only satisfies our desiderata, but also enables Medkit users to generate a
variety of batch datasets with customisable policy parameterisations (e.g in
terms of Markovianity, reward, variation in practice) and thus evaluate a
range of methods for understanding decision-making.

###  2.2 User workflow

Medkit was build to facilitate the development of machine learning methods for
clinical decision modelling. Medkit offers users the flexibility to obtain
batch datasets $\mathcal{D}^{\omega}_{syn,\mathcal{E}}$ for any desired type
of parameterisation $\theta$ (e.g. temperature, Markovianity, reward) of the
decision making policy $Q_{\pi_{\theta}}^{\mathcal{E}}$ and thus evaluate a
wide range of methods for modelling sequential decision making. This includes
methods for recovering expert’s reward function [11, 7], subjective dynamics
[28] or interpretable policies in the form of decision trees [6]. For
instance, to evaluate inverse reinforcement learning (IRL) methods, users can
chose among various domains $\Omega$ and environment dynamics $\mathcal{E}$
and define different ground-truth reward functions $R_{\theta}$ with
parameters $\theta$. Then, users can run Q-learning [42] to obtain the optimal
policy $Q^{\omega}_{\pi_{\theta}}$ for reward $R_{\theta}$, and add it to
Medkit, which can then be used to simulate a batch dataset with demonstrations
$\mathcal{D}^{\omega}_{syn,\mathcal{E}}$ for training their IRL algorithm. The
recovered policy parameteriaation $\hat{\theta}$ can then be evaluated against
the ground truth $\theta$.

While, as above, users can specify their own policy to roll-out in the
environments, we also provide as part of Medkit different types of
parameterised policies learnt from the clinicians’ policies in the real
dataset $D_{real}^{\Omega}$. These built-in policies allow users to easily
obtain batch datasets for simulating decision making behaviour with various
(customisable) degrees of Markovianity, rationality, counfounding, individual
consistency and variation in practice. Details can be found in Section 4.2.

Table 1: Summary of related benchmarks key features. Are they focused on the
Medical setting? Are they designed for Offline algorithms? Do they allow
Custom policies? Do they test how Robust algorithms are? Do they incorporate
Non-Markovian environment dynamics?

Benchmark Medical Offline Robust Non-Markovian Simulates Simulated policy RL
envs OpenAI gym [9] ✗ ✗ ✓ ✗ Environment Only N/A ALE [5] ✗ ✗ ✓ ✗ Environment
Only N/A RL and IL benchmarks RL Unplugged [26] ✗ ✓ ✗ ✓ Env. & Policy
(Entangled) Fixed RL Bench [30] ✗ ✓ ✗ ✗ Env. & Policy (Entangled) Fixed
Simitate [41] ✗ ✗ ✗ ✗ Env. & Policy (Entangled) Fixed MAGICAL [50] ✗ ✗ ✓ ✗
Env. & Policy (Entangled) Fixed Synth. gen. TimeGAN [53] ✓ ✓ ✗ ✓ Env. & Policy
(Entangled) Fixed Fourier Flows [2] ✓ ✓ ✗ ✓ Env. & Policy (Entangled) Fixed
Medkit (Ours) ✓ ✓ ✓ ✓ Env. & Policy (Disentangled) Customizable

##  3 Alternative Benchmarks and Simulation

Medkit generates _synthetic_ _batch_ _medical_ datasets for _benchmarking_
algorithms for modelling decision making. There is currently a relative lack
of standardised benchmarks for medical sequential decision making and most of
the few medical simulators used for evaluation are mathematically formulated
as dynamical systems defined by a small set of differential equations (e.g
cancer simulator in Gottesman et al., [24], HIV simulator in Du et al., [16])
or are hand-designed MDPs (e.g sepsis simulator in Oberst and Sontag, [43],
Futoma et al., [22]). Medkit, on the other hand, provides an entire
benchmarking suite and enables users to generate data from various medical
domains, with realistic environment dynamics and with customisable policy
parameterisations. Below, we discuss key differences then with related work,
which are summarised in Table 1.

Most benchmarking work has been done outside of the medical domain, in the
perhaps most similar work to us [50] present a suite specifically designed to
test robustness of imitation learning (IL) algorithms to distributional
shifts. Nevertheless, the properties they consider are specifically designed
for general robotics tasks than for modelling clinical decision making in
healthcare.

Recently offline RL has come more into view and along with it a few
benchmarking datasets [26, 30]. These collect state, action, reward tuples of
agents deployed in various environments, and despite the focus on RL with the
aim to make use of the reward information for some off-policy method like
Q-learning, these datasets can be easily used for simple imitation as well.
However, at their core they are large collections of recorded trajectories
obtained by running trained agents through the live environment. Thus, unlike
in Medkit, the end user is not able to specify properties of the policy that
are unique to describing human decision-making behaviours such as bounded
rationality individual consistency and variation in practice. Indeed this is
an issue with any imitation learning benchmark with its origins in RL: due to
the reward there’s usually only one policy considered the “optimal” one and
methods for these benchmarks are mainly evaluated on their ability to achieve
a high cumulative reward. This neglects the area of decision modelling [10,
32, 28], where we might be more interested in inference over potentially sub-
optimal policies to gain understanding of the human decision-making behaviour.
To address this, Medkit enables users to obtain batch medical datasets for
various different parameterisation $\theta$ (temperature, markovianity,
consistency, bounded rationality, reward) of the policy and the aim is to
evaluate algorithms based on how well they can recover $\theta$. Moreover, RL
benchmarks focus mainly on Markovian environment dynamics, while Medkit
considers the whole history of a patient.

Generative models for decision making. Generative models are a long
established pillar of modern machine learning [34, 23], though notably they
tend to focus on image and text based applications with less focus given to
the static tabular data $p(x_{s})$ and even less for time-series tabular data
$p(\\{x_{t}\\}^{T}_{t=1})$. Medkit presents as a generative model for the
whole process $p(x_{s},\\{x_{t}\\}^{T}_{t=1},\\{y_{t}\\}^{T}_{t=1})$, based on
the factorisation of equation 1. Importantly this allows for control over the
policy, which is very important for the purposes we have in mind, and which
traditional methods for synthetic data generation cannot handle normally.
Typically to apply generative models designed for static data, for example
through normalising flows [15], to this setting it would involve merging all
the static features, series features, and actions into one large feature
vector. This works especially badly for variable length time series requiring
padding and that any relationships between variables cannot be customised.
Methods that are specifically designed to work on time series data have been
proposed based on convolutions [44], deep Markov models [38] and GANs [53]
among others. Generally they model an auto-regressive process - a notable
exception being [2] who use a Fourier transform to model time series within
the frequency domain, making it inapplicable for sequential generation. Once
again though all of these models do not take into account actions (and rarely
static features) meaning they have to be absorbed into the series features and
cannot be customised.

##  4 Medkit Customisable Scenarios

We describe here the the various domains, policies and environment dynamics we
provide in the Medkit package. These can be combined arbitrarily to obtain a
large number of different scenarios for batch data generation. Medkit can also
live simulate the environment but without reward information is inappropriate
for reinforcement learning.

###  4.1 Domains

While Medkit generates _synthetic_ data, the machine learning methods used in
the generation process are trained on _real_ data. This is needed to capture
the complexity of real medical datasets and maximise the realism of the
scenarios and generated synthetic data. Thus, unlike in the toy medical
simulators seen in the literature [43, 16, 24], the batch datasets that can be
simulated from Medkit are high dimensional and governed by complex non-linear
dynamics, providing a much more realistic environment to test policies in
while still maintaining ground-truth information that can be used to evaluate
any learnt policy.

Out-of-the-box Medkit contains two medical domains $\Omega$ for which data can
be generated, capturing different medical settings: (1) Wards: general
hospital ward management at the Ronald Reagan UCLA Medical Center [4] and (2)
ICU: treatment of critically ill patients in various intensive care units [33,
19]. While for each domain, the data has undergone pre-processing to de-
identify and prevent re-identification of individual patients, we add an extra
layer of protection in the form of _differential privacy_ [17] guarantees by
employing differentially private optimisation techniques when training models,
which is readily supported by PyTorch’s Opacus library [21]. By ensuring that
the generated data is synthetic, Medkit enables wider public access without
the risk of sensitive information being inappropriately distributed. Specific
details on the state and action spaces for each domain can be found in the
Appendix along with details of the real data upon which they are based.

###  4.2 Policies

The key advantage of Medkit is that we separate the environment dynamics from
the policy dynamics. This enables us to roll-out customised policies within
the environment, and obtain batch datasets where the ground-truth policy
parameterisation is know. While users can define their own policy
parametrisations, we provide several built-in policies modelling the
distribution:

$\displaystyle p(\vec{y}_{T}|\vec{x}_{T})$
$\displaystyle=\prod^{T}_{t=1}Q_{\pi}^{\Omega}(y_{t}|\vec{x}_{t},\vec{y}_{t-1})$
(2)

By default we might be interested in a policy that seemingly mimics the seen
policy in the data as well as possible and so we include powerful neural-
network based learnt policies. Of course, as we hope to have conveyed already,
the interesting part comes in how the policy seen in the data can be
customised in specific ways that are interesting for imitation learning
algorithms to try and uncover. As such all policies are constructed in a
specific way:

$\displaystyle
Q_{\pi}^{\Omega}(y_{t}|\vec{x}_{t},\vec{y}_{t-1})=\sum_{i}\leavevmode\color[rgb]{1,0,1}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,1}\pgfsys@color@cmyk@stroke{0}{1}{0}{0}\pgfsys@color@cmyk@fill{0}{1}{0}{0}w_{i}\frac{e^{\leavevmode\color[rgb]{1,0.71,0.16}\definecolor[named]{pgfstrokecolor}{rgb}{1,0.71,0.16}\pgfsys@color@cmyk@stroke{0}{0.29}{0.84}{0}\pgfsys@color@cmyk@fill{0}{0.29}{0.84}{0}\beta_{i}\leavevmode\color[rgb]{0.72,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0.72,0,0}\pgfsys@color@cmyk@stroke{0}{0.89}{0.94}{0.28}\pgfsys@color@cmyk@fill{0}{0.89}{0.94}{0.28}q_{i}(y_{t}|\leavevmode\color[rgb]{0,0.44,0.57}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.44,0.57}\pgfsys@color@cmyk@stroke{0.98}{0.13}{0}{0.43}\pgfsys@color@cmyk@fill{0.98}{0.13}{0}{0.43}g_{i}(\vec{x}_{t}\leavevmode\color[rgb]{0,0.88,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.88,0}\pgfsys@color@cmyk@stroke{0.91}{0}{0.88}{0.12}\pgfsys@color@cmyk@fill{0.91}{0}{0.88}{0.12}\langle\mathcal{X}^{\prime}\rangle_{i},\vec{y}_{t-1}\leavevmode\color[rgb]{0,0.44,0.57}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.44,0.57}\pgfsys@color@cmyk@stroke{0.98}{0.13}{0}{0.43}\pgfsys@color@cmyk@fill{0.98}{0.13}{0}{0.43})\leavevmode\color[rgb]{0.72,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0.72,0,0}\pgfsys@color@cmyk@stroke{0}{0.89}{0.94}{0.28}\pgfsys@color@cmyk@fill{0}{0.89}{0.94}{0.28})}}{\sum_{y\in\mathcal{Y}}e^{\leavevmode\color[rgb]{1,0.71,0.16}\definecolor[named]{pgfstrokecolor}{rgb}{1,0.71,0.16}\pgfsys@color@cmyk@stroke{0}{0.29}{0.84}{0}\pgfsys@color@cmyk@fill{0}{0.29}{0.84}{0}\beta_{i}\leavevmode\color[rgb]{0.72,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0.72,0,0}\pgfsys@color@cmyk@stroke{0}{0.89}{0.94}{0.28}\pgfsys@color@cmyk@fill{0}{0.89}{0.94}{0.28}q_{i}(y|\leavevmode\color[rgb]{0,0.44,0.57}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.44,0.57}\pgfsys@color@cmyk@stroke{0.98}{0.13}{0}{0.43}\pgfsys@color@cmyk@fill{0.98}{0.13}{0}{0.43}g_{i}(\vec{x}_{t}\leavevmode\color[rgb]{0,0.88,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.88,0}\pgfsys@color@cmyk@stroke{0.91}{0}{0.88}{0.12}\pgfsys@color@cmyk@fill{0.91}{0}{0.88}{0.12}\langle\mathcal{X}^{\prime}\rangle_{i},\vec{y}_{t-1}\leavevmode\color[rgb]{0,0.44,0.57}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.44,0.57}\pgfsys@color@cmyk@stroke{0.98}{0.13}{0}{0.43}\pgfsys@color@cmyk@fill{0.98}{0.13}{0}{0.43})\leavevmode\color[rgb]{0.72,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0.72,0,0}\pgfsys@color@cmyk@stroke{0}{0.89}{0.94}{0.28}\pgfsys@color@cmyk@fill{0}{0.89}{0.94}{0.28})}}$

that introduces a number of components and properties that Medkit allows us to
model and can be controlled simply through the API, the details of which are
highlighted below:

  1. 1.

Ground-truth Structure \- the policy of a clinician will likely be difficult
if not impossible to describe. Even if they could articulate the policy, the
information will not be available in the data. Alternatively, we might expect
there to be some structure, since for example medical guidelines are often
given in the forms of decision trees [12, 49]. An algorithm that uncovers such
structure on regular medical data cannot be validated, since we do not know if
that inherent structure is in the data or just something the algorithm has
picked out - Medkit allows us to provide this ground truth with which we can
compare against.

  2. 2.

Markovianity \- the common assumption in sequential decision making is usually
that the problem can be modelled as a Markov decision process such that for a
policy that can be expressed
$q(y_{t}|\leavevmode\color[rgb]{0,0.44,0.57}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.44,0.57}\pgfsys@color@cmyk@stroke{0.98}{0.13}{0}{0.43}\pgfsys@color@cmyk@fill{0.98}{0.13}{0}{0.43}g(\vec{x}_{t},\vec{y}_{t-1}\leavevmode\color[rgb]{0,0.44,0.57}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.44,0.57}\pgfsys@color@cmyk@stroke{0.98}{0.13}{0}{0.43}\pgfsys@color@cmyk@fill{0.98}{0.13}{0}{0.43}))$
this is constrained so that $g(x_{t})=g(\vec{x}_{t},\vec{y}_{t-1})$, assuming
that the previous observations contains all of the relevant information. With
Medkit we can simply model more complicated policies that take into account
information much further into the past. We define the Markoviantity of the
policy as the minimum time lag into the past such that the policy is
equivalent to when considering the whole history:
$\inf\\{i\in\mathbb{N}:g(\vec{x}_{t-i:t},\vec{y}_{t-1-i:t-1})=g(\vec{x}_{t},\vec{y}_{t-1})\\}$.

  3. 3.

Bounded Rationality \- clinicians may not always act optimally based on all
the information available to them. In particular they may overlook some
specific variables as though they are not important [36]. We can model this in
Medkit by masking variables going into the policy model so that
$q(y_{t}|g(\vec{x}_{t},\vec{y}_{t-1}))=q(y_{t}|g(\vec{x}_{t}\leavevmode\color[rgb]{0,0.88,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.88,0}\pgfsys@color@cmyk@stroke{0.91}{0}{0.88}{0.12}\pgfsys@color@cmyk@fill{0.91}{0}{0.88}{0.12}\langle\mathcal{X}^{\prime}\rangle,\vec{y}_{t-1}))$,
where $\mathcal{X}^{\prime}$ is a subspace of $\mathcal{X}$ and
$\vec{x}_{T}\langle\mathcal{X}^{\prime}\rangle=x_{s}\cup\\{\operatorname{proj}_{\mathcal{X}^{\prime}}x_{t}\\}_{t=1}^{T}$.
Here, the dimensionality of $\mathcal{X}^{\prime}$ relative to $\mathcal{X}$
given as $\dim\mathcal{X}^{\prime}/\dim\mathcal{X}$ can be used as a measure
of the agent’s rationality.

  4. 4.

Individual Consistency \- some clinicians are very consistent, they will
always take the same action given a specific patient history. Others are more
stochastic, they’ll tend to favour the same actions but might occasionally
choose a different strategy given a “gut feeling” [18]. Medkit can model this
with the temperature of the Boltzmann distribution given in the output of all
of the policies. Formally, for policies of the form
$p(y_{t}|\vec{x}_{t},\vec{y}_{t-1})=\exp{\leavevmode\color[rgb]{1,0.71,0.16}\definecolor[named]{pgfstrokecolor}{rgb}{1,0.71,0.16}\pgfsys@color@cmyk@stroke{0}{0.29}{0.84}{0}\pgfsys@color<EMAIL_ADDRESS>q(y_{t}|g(\cdot))}/\sum_{y\in\mathcal{Y}}\exp{\leavevmode\color[rgb]{1,0.71,0.16}\definecolor[named]{pgfstrokecolor}{rgb}{1,0.71,0.16}\pgfsys@color@cmyk@stroke{0}{0.29}{0.84}{0}\pgfsys@color<EMAIL_ADDRESS>q(y|g(\cdot))}$, the inverse temperature $\beta\in\mathbb{R}_{+}$ measures the
individualised variability of an agent, where $\beta=0$ means that the agent
acts completely at random while $\beta\to\infty$ means that the agent is
perfectly consistent (i.e. their actions are deterministic).

  5. 5.

Variation in Practice \- often (essentially always) medical datasets are not
the recordings of a single clinician’s actions but of a mixture or team that
consult on an individual patient [51]. With Medkit we can model this
effectively using the `Mixture` policy, which takes any number of policies and
a mixing proportion to generate a new mixture policy. Formally, a mixture
policy is given by
$p(y_{t}|\vec{x}_{t},\vec{y}_{t-1})=\sum_{i}\leavevmode\color[rgb]{1,0,1}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,1}\pgfsys@color@cmyk@stroke{0}{1}{0}{0}\pgfsys@color@cmyk@fill{0}{1}{0}{0}w_{i}q_{i}(y_{t}|g(\vec{x}_{t},\vec{y}_{t-1}))$
where $\\{w_{i}\\}$ are the mixing proportions such that $\forall i,w_{i}>0$
and $\sum_{i}w_{i}=1$, and $\\{q_{i}(\cdot)\\}$ are arbitrary base policies.

These different policy parameterizations that are in-built into Medkit are
specific to scenarios that commonly arise in medicine [18, 51, 36], which is
the domain application we consider in this paper. However, note that the main
contribution of Medkit is to provide a framework for obtaining customizable
policies. Thus, users could also incorporate different types of policies if
needed.

###  4.3 Environments

The environment dynamics capture how the patient’s covariates evolve over time
given their history, interventions and the patient’s static features. From the
proposed factorisation in Equation (1), to estimate the environment dynamics,
we model the following conditional distribution in two parts:

$\displaystyle p(\vec{x}_{T}|\vec{y}_{T-1})$
$\displaystyle=\underbrace{P_{\mathcal{E}}^{\Omega}(x_{s},x_{1})}_{\text{Initialisation}}\prod_{t=2}^{T}\underbrace{P_{\mathcal{E}}^{\Omega}(x_{t}|f_{\mathcal{E}}(\vec{x}_{t-1},\vec{y}_{t-1}))}_{\text{Auto-
regression}},$ (3)

allowing for sequential generation of patient trajectories. For all
environments, we model $P_{\mathcal{E}}^{\Omega}(x_{s},x_{1})$ using a
Variational Autoencoder [34],as a powerful generative model that can handle a
mixture of continuous and discrete variables. For the auto-regressive part, to
capture a diverse set of the realistic dynamics of medical datasets, Medkit
contains environments that are (1) directly modelling the patient history
(T-Force and CRN) and (2) building latent variable models (CSS and SVAE). We
describe the models in this section but full details (e.g. on learning) are
given in the Appendix.

Directly modelling the patient history. This relates to attempting to model:

$\displaystyle
p(x_{t}|\vec{x}_{t-1},\vec{y}_{t-1})=P_{\mathcal{E}}^{\Omega}(x_{t}|f_{\mathcal{E}}(\vec{x}_{t-1},\vec{y}_{t-1}))$
(4)

directly, or more specifically that $p(x_{t}|\vec{x}_{t-1},\vec{y}_{t-1})$ is
some $\Theta$ parameterised distribution where
$\Theta=f(\vec{x}_{t-1},\vec{y}_{t-1})$ is a function of the history only. For
the simplest environment model, we use a recurrent neural network trained with
teacher forcing [52] (T-Force) to directly approximate this function. The
network is made up of LSTM units [27] followed by fully connected layers with
ELU activations [13] and is trained to maximise the likelihood of the next
observation given previous observations and interventions. This defines a
factorised Gaussian and Bernoulli distribution over the continuous and binary
covariates respectively with the parameters predicted by the network.

Additionally we extend this method by replacing the LSTM network with the
Counterfactual Recurrent Network (CRN) of Bica et al., [7]. CRN is a causal
inference method that learns balancing representation of the patients’
histories to remove the time-dependent confounding bias present in
observational datasets. This allows the network to more principally be used
for making counterfactual predictions which is what our model for the
environment dynamics needs to do when estimating the next state of a patient
under different possible interventions specified by the policy
$Q^{\Omega}_{\pi}$.

Building latent variable models. We also build environment dynamics where the
observations are driven by a _hidden_ true state of the patient. Formally, we
assume the features $\vec{x}_{T}$ are driven by some evolving latent state
$\vec{z}_{T}=\\{z_{t}\\}_{t=1}^{T},z_{t}\in\mathcal{Z}$ that is not seen in
the data by modelling a factorisation given by:

$\displaystyle
P_{\mathcal{E}}^{\Omega}(x_{t},z_{t}|f_{\mathcal{E}}(\vec{x}_{t-1},\vec{y}_{t-1},\vec{z}_{t-1}))=\underbrace{P_{\mathcal{E}}^{\Omega}(x_{t}|z_{t},x_{s})}_{\text{Emission}}\times\underbrace{P_{\mathcal{E}}^{\Omega}(z_{t}|f_{\mathcal{E}}(\vec{x}_{t-1},\vec{y}_{t-1},\vec{z}_{t-1}))}_{\text{Transition}}.$
(5)

We include as part of Medkit two additional environment dynamics models for
the separate cases when $|\mathcal{Z}|$ is finite or uncountable, as both can
usefully represent patients in the medical context. For $|\mathcal{Z}|$ finite
the latent $z_{t}$ variables then might represent distinct progression
“stages” or various classifications of a disease. Discrete separation like
this is well established in both clinical guidelines and models for a range of
cases including transplantation in patients with CF [8], the diagnosis of
Alzheimer’s disease [45], and cancer screening [47]. Accordingly we use the
Attentive State-Space model of [3] to build an attention-based, customised
state-space (CSS) representation of disease progression. This environment
model accounts for static features and allows Medkit users to customise the
attention mechanism. Given a discrete latent space, the transitions are
parameterised with baseline transition matrices for each action averaged over
attention weights on previous timesteps. The emission distribution allows for
a flexible representation: let $p_{\psi}(x_{t})$ be any distribution with
support over $\mathcal{X}$ and parameter(s) $\psi$ (for example some Gaussian
mixture) then we let:

$\displaystyle
p(x_{t}|z_{t},x_{s})=p_{\psi^{*}}(x_{t}),\quad\text{with}\;\psi^{*}=f_{\gamma}(z_{t},x_{s}).$
(6)

We take $f_{\gamma}$ to be a $\gamma$-parameterised function approximator to
output the parameters of the emission distribution given the current state and
static features of the patient - a standard choice being an MLP that takes in
the concatenation of $z_{t}$ and $x_{s}$. This alleviates a common problem
with state-space models where the observations are ultimately drawn from some
finite mixture of distributions of order $|\mathcal{Z}|$, as now the
dependence on the static features allows for a very flexible output. The CSS
dynamics model allows Medkit users to post-hoc customise the number of states
and the Markovianity of the environment through the attention mechanism (e.g
users can pass a vector that specifies exact weights or an integer
representing the number of states back to look.)

While a discrete representation of hidden states is convenient for
interpretation, it does simplify the problem. It is unlikely that all of the
relevant features of a disease can be adequately captured by a discrete
characterisation - it would seem that in reality diseases evolve gradually and
without step-change. Therefore, to further improve the realism of the
generated trajectories, we also include as part of Medkit’s environments a
deep continuous state space model that extends VAEs in a sequential manner
(SVAE). Principally now we consider a continuous latent state with
$\mathcal{Z}=\mathbb{R}^{d}$. This then allows for more flexibility in the
transition dynamics, in particular by making use of neural architectures. An
encoder network predicts the approximate posterior over the latent variables
and we employ essentially the same method as for teacher forcing in order to
model dynamics in the latent space. With a joint optimisation scheme, we learn
a representation that generates the observations well but also captures the
features relevant for the transitions. This expressiveness allows for a higher
fidelity model than the custom state-space but however comes at the cost of
interpretable structure which we have established may be useful should
algorithms be designed to uncover such things.

Modelling hidden confounding. A common assumption, that is likely not true in
practice, is that there are no hidden confounding variables in the
environment. Medkit allows us to introduce and control these by using a full
set of variables to generate both the actions and the observations but
restrict the visibility of some such that they become hidden. While the
overall generative process $p(\vec{x}_{T},\vec{y}_{T})$ is left unchanged,
only a partially-hidden dataset
$\mathcal{D}=\\{\vec{x}_{T}\langle\bar{\mathcal{X}^{\prime}}\rangle,\vec{y}_{T}\\}$
is provided to the user, where $\mathcal{X}^{\prime}$ is a subspace of
$\mathcal{X}$ and
$\vec{x}_{T}\langle\mathcal{X}^{\prime}\rangle=x_{s}\cup\\{\operatorname{proj}_{\mathcal{X}^{\prime}}x_{t}\\}_{t=1}^{T}$.
Here, the dimensionality of $\mathcal{X}^{\prime}$ relative to $\mathcal{X}$
given as $\dim\mathcal{X}^{\prime}/\dim\mathcal{X}$ can be used as a measure
of the overall confoundedness.

##  5 Practical Demonstrations

In this section we explore some examples of the benefits of using Medkit
compared to existing benchmarks as well as highlight some potential use cases,
in particular how Medkit allows for consistent and systematic evaluation along
with useful ground truth information.

Different reactions to shifting policies. The current literature on imitation
learning focuses on very different environments to those found in the medical
setting and consequently algorithms may not be evaluated against, or designed
to be appropriate for, the quirks of medical data. For example in Figure 3a we
plot the performance of algorithms as the consistency of the policy varies, in
particular we use: Behavioural Cloning (BC) with a deep Q-network; Reward-
regularized Classification for Apprenticeship Learning (RCAL) [48], where the
network is regularised such that the implicit rewards are sparse; ValueDICE
(VDICE) [37], an offline adaptation of the adversarial imitation learning
framework; and Energy-based Distribution Matching (EDM) [31] that uses the
implicit energy-based model to partially correct for the off-policy nature of
BC. What is interesting is not that performance degrades - this is of course
to be expected, but rather that the comparative ranking of algorithms changes
as a function of the consistency. In particular BC performs the worst
(although there is little between them) in the ends up outperform the rest on
average when the variation is highest, suggesting some of the more complicated
algorithms are not robust to these kinds of policies.

Enabling consistent evaluation. Common RL benchmarks like Atari experience
very large variances in the accumulated reward an agent obtains when deployed
in the environment, especially when the reward is sparse. This can make
evaluation and ranking of agents tricky or at least require a large number of
runs in the environment before the variance of the estimator suggests the
results are significant. In Figure 3b we demonstrate this problem in an even
simpler context comparing BC to the AVRIL algorithm of [11], a method for
approximate Bayesian IRL, in the simple Acrobot environment where the aim is
to swing up a pendulum to a correct height. On the right y-axis we plot the
accumulated regret over training of the two agents, and large inconsistencies
in return can be seen so that it is not clear which of the agents is better.
Comparatively on the left y-axis we plot the AUROC on a held out test set as
we train on Medkit data, here evaluation is much more consistent and
statistically significant, demonstrating clearly which algorithm is performing
better.

Figure 3:  Exploring Medkit Practically. Example benefits of Medkit for
exploring and benchmarking imitation learning algorithms.

Ground-truth knowledge comparison. While in the end it only really matters how
an algorithm performs when deployed in the real world, it is challenging to
only use real data to validate them. This is since you run into the key
problem that you will not have any knowledge of the ground truth behind
decisions and so methods that claim to gain insight into such areas cannot
possibly be evaluated appropriately. On the other hand simulating data in
Medkit allows us to do exactly this, and we can compare inferences from an
algorithm to underlying truth in the generating process. A toy example is
shown in Figure 3c where we compare the weights of a linear classifier trained
on Medkit data to those of the true underlying policy, representing the
relative feature importances for the policies.

Figure 4:  t-SNE plots For each policy in the Ward environment we generate
simulated data. We then apply t-SNE and project the real and simulated data
into two components, which is plotted.

Validating realism. It is also of interest to quickly check that we are not
generating completely unrealistic trajectories, rather ones that capture
appropriate properties that will be useful for users. We thus provide
comparisons of the available environment models in Medkit. In particular for
each combination we show in Table 2: the _Predictive Score_ , a classical
“train on synthetic - test on real” evaluation where a network is trained on
the synthetic dataset and applied to a held out test set of the real data,
where the performance is reported; and the _Discriminitive Score_ , where a
classifier is trained to distinguish between the real and synthetic data, and
the AUROC of this task on a held out test set is reported. In aid of
visualisation we also provide in Figure 4 a set of t-SNE plots [40] overlaying
the real and synthetic data. These metrics are standard in the synthetic data
literature [53] and reflect the usefulness of the synthetic data as a
_replacement_ for real data.

Table 2: Predictive and Discriminative Scores. Scores reported on the
different environments for the Wards domain.

| $|\mathcal{Y}|$ | T-Force | CRN | CSS | S-VAE
---|---|---|---|---|---
Pred.$\uparrow$ | 2 | $0.67\pm 0.05$ | $0.94\pm 0.01$ | $0.94\pm 0.01$ | $0.93\pm 0.01$
4 | $0.62\pm 0.02$ | $0.85\pm 0.01$ | $0.86\pm 0.01$ | $0.86\pm 0.02$
8 | $0.61\pm 0.05$ | $0.85\pm 0.03$ | $0.89\pm 0.02$ | $0.87\pm 0.04$
Disc.$\downarrow$ | 2 | $0.41\pm 0.03$ | $0.23\pm 0.02$ | $0.19\pm 0.03$ | $0.22\pm 0.04$
4 | $0.41\pm 0.05$ | $0.24\pm 0.04$ | $0.19\pm 0.04$ | $0.23\pm 0.04$
8 | $0.37\pm 0.07$ | $0.22\pm 0.03$ | $0.20\pm 0.03$ | $0.20\pm 0.02$

Please note though that the highest possible fidelity is not the point of
Medkit: unlike traditional synthetic data, the datasets we produce are _not_
meant to be used as a substitute for real data in training machine learning
algorithms. Rather we would like to produce _realistic_ data that reflects the
difficulties of the medical setting and can be used for development and
benchmarking of algorithms. Additionally, by introducing customisations into
the generative process, we will naturally see departures from real data, but
given our goals this is not a problem. Nevertheless, the high predictive
scores show that Medkit is successfully capturing important trends in the real
data that are useful for prediction, while the discriminative scores and t-SNE
plots confirm that we are not producing trajectories that are
unrepresentative.

##  6 Discussion

Limitations and Societal Impact. As a synthetic data generator, Medkit is
inherently limited by the power of the individual models used and their
ability to accurately model outcomes given specified policies. This is not
such a problem when the focus is on inference over the policy though, as is
the focus in decision modelling. Additionally, Medkit is easily extendable
when new, more powerful, models become available. With Medkit our aim is to
provide a platform allowing for better development of decision modelling
algorithms, the societal impact thus very much depends on the potential use of
such algorithms, for example, they could be used to misrepresent an
individual’s position or identify biases that could be exploited. By focusing
on clinical decision support, we hope to promote a much more beneficial
approach.

Conclusions. We have presented the Medkit-Learn(ing) Environment, a
benchmarking suite for medical sequential decision making. As with many
software libraries, the work is never done and there are always new features
that can be added. Indeed we can, and intend to, always continue to add more
tools and algorithms to be beneficial for the community. One important future
area that Medkit could make an impact in is causality - an area where more
than ever synthetic data is important such that we can actually evaluate the
counterfactuals that are inherently missing from real data, and much can be
done to simulate data for individualised treatment estimation for example.
Overall though our aim with Medkit is to advance the development of algorithms
for _understanding_ , not just imitating, decision making so that we can
better support those high-stakes decisions such as in the clinical setting
without replacing the crucial human aspect needed when the problem is so
important.

## References

  * Abbeel and Ng, [2004] Abbeel, P. and Ng, A. Y. (2004).  Apprenticeship learning via inverse reinforcement learning.  In Proceedings of the twenty-first international conference on Machine learning, page 1.
  * Alaa et al., [2021] Alaa, A., Chan, A. J., and van der Schaar, M. (2021).  Generative time-series modeling with fourier flows.  In International Conference on Learning Representations.
  * Alaa and van der Schaar, [2019] Alaa, A. M. and van der Schaar, M. (2019).  Attentive state-space modeling of disease progression.  In Advances in Neural Information Processing Systems, pages 11338–11348.
  * Alaa et al., [2017] Alaa, A. M., Yoon, J., Hu, S., and Van der Schaar, M. (2017).  Personalized risk scoring for critical care prognosis using mixtures of gaussian processes.  IEEE Transactions on Biomedical Engineering, 65(1):207–218.
  * Bellemare et al., [2013] Bellemare, M. G., Naddaf, Y., Veness, J., and Bowling, M. (2013).  The arcade learning environment: An evaluation platform for general agents.  Journal of Artificial Intelligence Research, 47:253–279.
  * Bewley et al., [2020] Bewley, T., Lawry, J., and Richards, A. (2020).  Modelling agent policies with interpretable imitation learning.  arXiv preprint arXiv:2006.11309.
  * Bica et al., [2020] Bica, I., Alaa, A. M., Jordon, J., and van der Schaar, M. (2020).  Estimating counterfactual treatment outcomes over time through adversarially balanced representations.  In International Conference on Learning Representations.
  * Braun and Merlo, [2011] Braun, A. T. and Merlo, C. A. (2011).  Cystic fibrosis lung transplantation.  Current opinion in pulmonary medicine, 17(6):467–472.
  * Brockman et al., [2016] Brockman, G., Cheung, V., Pettersson, L., Schneider, J., Schulman, J., Tang, J., and Zaremba, W. (2016).  Openai gym.
  * Chakraborti et al., [2018] Chakraborti, T., Fadnis, K. P., Talamadupula, K., Dholakia, M., Srivastava, B., Kephart, J. O., and Bellamy, R. K. (2018).  Visualizations for an explainable planning agent.  In Proceedings of the 27th International Joint Conference on Artificial Intelligence, pages 5820–5822.
  * Chan and van der Schaar, [2021] Chan, A. J. and van der Schaar, M. (2021).  Scalable Bayesian inverse reinforcement learning.  In International Conference on Learning Representations.
  * Chou et al., [2007] Chou, R., Qaseem, A., Snow, V., Casey, D., Cross, J. T., Shekelle, P., and Owens, D. K. (2007).  Diagnosis and treatment of low back pain: a joint clinical practice guideline from the american college of physicians and the american pain society.  Annals of internal medicine, 147(7):478–491.
  * Clevert et al., [2015] Clevert, D.-A., Unterthiner, T., and Hochreiter, S. (2015).  Fast and accurate deep network learning by exponential linear units (elus).  arXiv preprint arXiv:1511.07289.
  * Dash et al., [2020] Dash, S., Yale, A., Guyon, I., and Bennett, K. P. (2020).  Medical time-series data generation using generative adversarial networks.  In International Conference on Artificial Intelligence in Medicine, pages 382–391. Springer.
  * Dinh et al., [2016] Dinh, L., Sohl-Dickstein, J., and Bengio, S. (2016).  Density estimation using real nvp.  arXiv preprint arXiv:1605.08803.
  * Du et al., [2020] Du, J., Futoma, J., and Doshi-Velez, F. (2020).  Model-based reinforcement learning for semi-markov decision processes with neural odes.  arXiv preprint arXiv:2006.16210.
  * Dwork et al., [2014] Dwork, C., Roth, A., et al. (2014).  The algorithmic foundations of differential privacy.  Foundations and Trends in Theoretical Computer Science, 9(3-4):211–407.
  * Eccles et al., [2006] Eccles, M. P., Hrisos, S., Francis, J., Kaner, E. F., Dickinson, H. O., Beyer, F., and Johnston, M. (2006).  Do self-reported intentions predict clinicians’ behaviour: a systematic review.  Implementation Science, 1(1):1–10.
  * Elbers, [2019] Elbers, P. W. G. (2019).  AmsterdamUMCdb v1.0.2 ICU database.
  * Esteban et al., [2017] Esteban, C., Hyland, S. L., and Rätsch, G. (2017).  Real-valued (medical) time series generation with recurrent conditional gans.  arXiv preprint arXiv:1706.02633.
  * Facebook, [2020] Facebook (2020).  Opacus PyTorch library.
  * Futoma et al., [2020] Futoma, J., Hughes, M. C., and Doshi-Velez, F. (2020).  Popcorn: Partially observed prediction constrained reinforcement learning.  International Conference on Artificial Intelligence and Statistics (AISTATS).
  * Goodfellow et al., [2014] Goodfellow, I., Pouget-Abadie, J., Mirza, M., Xu, B., Warde-Farley, D., Ozair, S., Courville, A., and Bengio, Y. (2014).  Generative adversarial nets.  Advances in neural information processing systems, 27:2672–2680.
  * Gottesman et al., [2020] Gottesman, O., Futoma, J., Liu, Y., Parbhoo, S., Brunskill, E., Doshi-Velez, F., et al. (2020).  Interpretable off-policy evaluation in reinforcement learning by highlighting influential transitions.  arXiv preprint arXiv:2002.03478.
  * Gottesman et al., [2019] Gottesman, O., Johansson, F., Komorowski, M., Faisal, A., Sontag, D., Doshi-Velez, F., and Celi, L. A. (2019).  Guidelines for reinforcement learning in healthcare.  Nature medicine, 25(1):16–18.
  * Gulcehre et al., [2020] Gulcehre, C., Wang, Z., Novikov, A., Paine, T., Gómez, S., Zolna, K., Agarwal, R., Merel, J. S., Mankowitz, D. J., Paduraru, C., et al. (2020).  Rl unplugged: A collection of benchmarks for offline reinforcement learning.  Advances in Neural Information Processing Systems, 33.
  * Hochreiter and Schmidhuber, [1997] Hochreiter, S. and Schmidhuber, J. (1997).  Long short-term memory.  Neural computation, 9(8):1735–1780.
  * Hüyük et al., [2021] Hüyük, A., Jarrett, D., Tekin, C., and van der Schaar, M. (2021).  Explaining by imitating: Understanding decisions by interpretable policy learning.  In International Conference on Learning Representations.
  * Jain et al., [2019] Jain, V., Doshi, P., and Banerjee, B. (2019).  Model-free irl using maximum likelihood estimation.  In Proceedings of the AAAI Conference on Artificial Intelligence, volume 33, pages 3951–3958.
  * James et al., [2020] James, S., Ma, Z., Arrojo, D. R., and Davison, A. J. (2020).  Rlbench: The robot learning benchmark & learning environment.  IEEE Robotics and Automation Letters, 5(2):3019–3026.
  * Jarrett et al., [2020] Jarrett, D., Bica, I., and van der Schaar, M. (2020).  Strictly batch imitation learning by energy-based distribution matching.  Advances in Neural Information Processing Systems, 33.
  * Jarrett and van der Schaar, [2020] Jarrett, D. and van der Schaar, M. (2020).  Inverse active sensing: Modeling and understanding timely decision-making.  In International Conference on Machine Learning, pages 4713–4723. PMLR.
  * Johnson et al., [2016] Johnson, A. E., Pollard, T. J., Shen, L., Li-Wei, H. L., Feng, M., Ghassemi, M., Moody, B., Szolovits, P., Celi, L. A., and Mark, R. G. (2016).  Mimic-iii, a freely accessible critical care database.  Scientific data, 3(1):1–9.
  * Kingma and Welling, [2013] Kingma, D. P. and Welling, M. (2013).  Auto-encoding variational bayes.  arXiv preprint arXiv:1312.6114.
  * Klein et al., [2011] Klein, E., Geist, M., and Pietquin, O. (2011).  Batch, off-policy and model-free apprenticeship learning.  In European Workshop on Reinforcement Learning, pages 285–296. Springer.
  * Klerings et al., [2015] Klerings, I., Weinhandl, A. S., and Thaler, K. J. (2015).  Information overload in healthcare: too much of a good thing?  Zeitschrift für Evidenz, Fortbildung und Qualität im Gesundheitswesen, 109(4-5):285–290.
  * Kostrikov et al., [2019] Kostrikov, I., Nachum, O., and Tompson, J. (2019).  Imitation learning via off-policy distribution matching.  In International Conference on Learning Representations.
  * Krishnan et al., [2017] Krishnan, R., Shalit, U., and Sontag, D. (2017).  Structured inference networks for nonlinear state space models.  In Proceedings of the AAAI Conference on Artificial Intelligence, volume 31.
  * Levine et al., [2020] Levine, S., Kumar, A., Tucker, G., and Fu, J. (2020).  Offline reinforcement learning: Tutorial, review, and perspectives on open problems.  arXiv preprint arXiv:2005.01643.
  * Maaten and Hinton, [2008] Maaten, L. v. d. and Hinton, G. (2008).  Visualizing data using t-sne.  Journal of machine learning research, 9(Nov):2579–2605.
  * Memmesheimer et al., [2019] Memmesheimer, R., Mykhalchyshyna, I., Seib, V., and Paulus, D. (2019).  Simitate: A hybrid imitation learning benchmark.  arXiv preprint arXiv:1905.06002.
  * Mnih et al., [2013] Mnih, V., Kavukcuoglu, K., Silver, D., Graves, A., Antonoglou, I., Wierstra, D., and Riedmiller, M. (2013).  Playing atari with deep reinforcement learning.  arXiv preprint arXiv:1312.5602.
  * Oberst and Sontag, [2019] Oberst, M. and Sontag, D. (2019).  Counterfactual off-policy evaluation with gumbel-max structural causal models.  International Conference on Machine Learning.
  * Oord et al., [2016] Oord, A. v. d., Dieleman, S., Zen, H., Simonyan, K., Vinyals, O., Graves, A., Kalchbrenner, N., Senior, A., and Kavukcuoglu, K. (2016).  Wavenet: A generative model for raw audio.  arXiv preprint arXiv:1609.03499.
  * O’Bryant et al., [2008] O’Bryant, S. E., Waring, S. C., Cullum, C. M., Hall, J., Lacritz, L., Massman, P. J., Lupo, P. J., Reisch, J. S., and Doody, R. (2008).  Staging dementia using clinical dementia rating scale sum of boxes scores: a texas alzheimer’s research consortium study.  Archives of neurology, 65(8):1091–1095.
  * Paszke et al., [2019] Paszke, A., Gross, S., Massa, F., Lerer, A., Bradbury, J., Chanan, G., Killeen, T., Lin, Z., Gimelshein, N., Antiga, L., et al. (2019).  Pytorch: An imperative style, high-performance deep learning library.  In Advances in neural information processing systems, pages 8026–8037.
  * Petousis et al., [2019] Petousis, P., Winter, A., Speier, W., Aberle, D. R., Hsu, W., and Bui, A. A. (2019).  Using sequential decision making to improve lung cancer screening performance.  IEEE Access, 7:119403–119419.
  * Piot et al., [2014] Piot, B., Geist, M., and Pietquin, O. (2014).  Boosted and reward-regularized classification for apprenticeship learning.  In Proceedings of the 2014 international conference on Autonomous agents and multi-agent systems, pages 1249–1256. International Foundation for Autonomous Agents and Multiagent Systems.
  * Qaseem et al., [2012] Qaseem, A., Fihn, S. D., Dallas, P., Williams, S., Owens, D. K., and Shekelle, P. (2012).  Management of stable ischemic heart disease: Summary of a clinical practice guideline from the american college of physicians/american college of cardiology foundation/american heart association/american association for thoracic surgery/preventive cardiovascular nurses association/society of thoracic surgeons.  Annals of Internal Medicine, 157(10):735–743.
  * Toyer et al., [2020] Toyer, S., Shah, R., Critch, A., and Russell, S. (2020).  The magical benchmark for robust imitation.  Advances in Neural Information Processing Systems, 33.
  * Undre et al., [2006] Undre, S., Sevdalis, N., Healey, A. N., Darzi, S. A., and Vincent, C. A. (2006).  Teamwork in the operating theatre: cohesion or confusion?  Journal of evaluation in clinical practice, 12(2):182–189.
  * Williams and Zipser, [1989] Williams, R. J. and Zipser, D. (1989).  A learning algorithm for continually running fully recurrent neural networks.  Neural Computation, 1(2):270–280.
  * Yoon et al., [2019] Yoon, J., Jarrett, D., and van der Schaar, M. (2019).  Time-series generative adversarial networks.  In Advances in Neural Information Processing Systems, pages 5508–5518.
  * Zhi-Xuan et al., [2020] Zhi-Xuan, T., Mann, J., Silver, T., Tenenbaum, J., and Mansinghka, V. (2020).  Online bayesian goal inference for boundedly rational planning agents.  Advances in Neural Information Processing Systems, 33.
  * Ziebart et al., [2008] Ziebart, B. D., Maas, A. L., Bagnell, J. A., and Dey, A. K. (2008).  Maximum entropy inverse reinforcement learning.  In Aaai, volume 8, pages 1433–1438. Chicago, IL, USA.

## Checklist

  1. 1.

For all authors…

    1. (a)

Do the main claims made in the abstract and introduction accurately reflect
the paper’s contributions and scope? [Yes]

    2. (b)

Did you describe the limitations of your work? [Yes]

    3. (c)

Did you discuss any potential negative societal impacts of your work? [Yes]

    4. (d)

Have you read the ethics review guidelines and ensured that your paper
conforms to them? [Yes]

  2. 2.

If you are including theoretical results…

    1. (a)

Did you state the full set of assumptions of all theoretical results? [N/A]

    2. (b)

Did you include complete proofs of all theoretical results? [N/A]

  3. 3.

If you ran experiments (e.g. for benchmarks)…

    1. (a)

Did you include the code, data, and instructions needed to reproduce the main
experimental results (either in the supplemental material or as a URL)? [Yes]

    2. (b)

Did you specify all the training details (e.g., data splits, hyperparameters,
how they were chosen)? [Yes]

    3. (c)

Did you report error bars (e.g., with respect to the random seed after running
experiments multiple times)? [Yes]

    4. (d)

Did you include the total amount of compute and the type of resources used
(e.g., type of GPUs, internal cluster, or cloud provider)? [Yes]

  4. 4.

If you are using existing assets (e.g., code, data, models) or
curating/releasing new assets…

    1. (a)

If your work uses existing assets, did you cite the creators? [Yes]

    2. (b)

Did you mention the license of the assets? [No]

    3. (c)

Did you include any new assets either in the supplemental material or as a
URL? [Yes]

    4. (d)

Did you discuss whether and how consent was obtained from people whose data
you’re using/curating? [Yes]

    5. (e)

Did you discuss whether the data you are using/curating contains personally
identifiable information or offensive content? [Yes]

  5. 5.

If you used crowdsourcing or conducted research with human subjects…

    1. (a)

Did you include the full text of instructions given to participants and
screenshots, if applicable? [N/A]

    2. (b)

Did you describe any potential participant risks, with links to Institutional
Review Board (IRB) approvals, if applicable? [N/A]

    3. (c)

Did you include the estimated hourly wage paid to participants and the total
amount spent on participant compensation? [N/A]

Python Repo: https://github.com/XanderJC/medkit-learn

##  Appendix A Domains and Real Data Information

###  A.1 Hospital Wards.

The _Ward_ domain is based on the care of 6,321 patients at the Ronald Reagan
UCLA Medical Center in California who were treated on the general medicine
floor between 2013-2016 alaa2017personalized. These patients were treated for
a variety of conditions including pneumonia, sepsis, and fevers were in
general stable and deterioration that required ICU care was rare. Measurements
were taken roughly every 4 hours, with average stays lasting 9 days, and
include common vital signs such as pulse and blood pressure alongside lab
tests and results and include 8 static (although categorical features are one-
hot encoded, extending the space) and 35 temporal features for which we model
the dynamics. The action space is taken as a choice of one to three binary
actions (encoded for a maximum size of eight) marking the treatment of various
oxygen therapy devices.

###  A.2 Intensive Care Unit.

The _ICU_ domain simulates the treatment of 23,106 of patients in the
intensive care unit from Amsterdam UMC elbers2019amsterdam. These patients are
in a more critical state than those on the general wards while suffering
similarly from a variety of conditions and consequently are monitored more
frequently, with the database containing around 1 billion clinical
observations at varying timesteps down to minute by minute recordings. We
aggregate data into one hour timesteps and model the treatment of mechanical
ventilation alongside the prescription of antibiotics and oxygentherapy. We
include 36 static features including height, initial weight, and
commorbidities with 24 series features focusing on vital signs including heart
rate, blood pressure, and various chemical blood concentration levels.

##  Appendix B Further Environment Model Details

###  B.1 The Customised State Space Model

The CSS is a variety of a deep non-Markovian hidden state-space model. The
_non-Markovianity_ comes from the property in the model that transitions are
parameterised given discrete states simply as follows:

$\displaystyle p(z_{t}|\vec{x}_{t-1},\vec{y}_{t-1},\vec{z}_{t-1})$
$\displaystyle=p(z_{t}|\vec{\alpha}_{t},\vec{z}_{t-1})$
$\displaystyle=\sum_{t^{\prime}=1}^{t-1}\alpha_{t^{\prime}}^{t-1}\mathbf{P}_{y_{t^{\prime}}}(z_{t^{\prime}},z_{t}),$
(7)

with $\mathbf{P}_{y_{t}}$ a baseline transition matrix given intervention
$y_{t}$ is made and $\alpha_{t^{\prime}}^{t-1}$ some attention weight on the
previous timesteps such that
$\sum_{t^{\prime}=1}^{t-1}\alpha_{t^{\prime}}^{t-1}=1$. This departs from the
traditional IOHMM bengio1995input by the inclusion of the attention weights
that induce time-dependency on points further in the past than the previous
time point and can be learnt during training.

As described in the main text, the _deep_ aspect is manifested in the emission
probability distribution whose parameters are predicted from the output of a
neural network that takes the static features and current state.

#### Learning and inference.

Exact inference over the hidden states in this model is intractable and so we
use an inference network to parameterise an approximate posterior over the
latent hidden state, in particular assuming a factorisation that mirrors the
true structure of the posterior as follows:

$\displaystyle q_{\phi}(\vec{z}_{T}|\vec{x}_{T},\vec{y}_{T})$
$\displaystyle=q_{\phi}(z_{1}|\vec{x}_{T},\vec{y}_{T})\prod_{t=2}^{T}q_{\phi}(z_{t}|z_{t-1},\vec{x}_{t:T},\vec{y}_{t:T}).$
(8)

Thus the posterior depends on the previous hidden state and all future
observations and actions. This can be practically achieved by using a backward
LSTM hochreiter1997long to summarise the future into its hidden state before
passing that and the previous state into a new network to obtain the
approximate posterior.

Our inference network leads to a factorised Evidence Lower BOund (ELBO) given
by:

$\displaystyle\log p_{\theta}(\vec{x}_{T}|\vec{y}_{T-1})$
$\displaystyle\geq\mathbb{E}_{q_{\phi}}\big{[}\log
p_{\theta}(\vec{x}_{T},\vec{z}_{T}|\vec{y}_{T-1})-\log
q_{\phi}(\vec{z}_{T}|\vec{x}_{T},\vec{y}_{T})\big{]},$ (9)

leading to an optimisation task over $\theta$ and $\phi$. This is readily done
using stochastic variational inference whereby we take a Monte Carlo estimate
of the ELBO by sampling from the posterior and taking gradients
krishnan2017structured.

###  B.2 The Sequential VAE Model

The SVAE is again a variety of a deep non-Markovian hidden state-space model,
although now the states are no longer discrete but rather live in some latent
space of $\mathbb{R}^{d}$. This again means that an inference network is
required for doing approximate inference and similarly leads to an ELBO
optimisation scheme.

##  Appendix C Further Experimental Details

###  C.1 Computation

All experiments were performed on a 2016 MacBook Pro, using a 2.9 GHz Dual-
Core Intel Core i5 with 8GB of LPDDR3 RAM and no GPU acceleration.

###  C.2 Predictive Score

The predictive score measures how well the synthetic data can be used as a
direct substitute for real data when it comes to a prediction task.

For our experiments we generate 500 sample trajectories of maximum length 30
to be used as the synthetic data. A network consisting of an LSTM layer
followed by two fully connected layers is then trained on this data for 50
epochs. This is then applied to a real data set and the AUROC is reported.

###  C.3 Discriminative Score

The discriminative score measures how well the synthetic data ‘hides’ amongst
real data.

For our experiments a training dataset is generated using 100 real
trajectories and 100 synthetic trajectories along with their associated
real/synthetic label. A single layer LSTM is then trained on this dataset to
predict the label, before being applied to a test set of 100 real trajectories
and 100 synthetic trajectories. We report the absolute value of the test
accuracy minus 0.5.

###  C.4 Imitation Learning Benchmarks

All methods used are based on neural network and so in the experiments we
maintain the same architecture of 2 hidden layers of 64 units each connected
by exponential linear unit (ELU) activation functions.

Publicly available code was used in the implementations of a number of the
benchmarks, specifically:

  * •

VDICE [37]:
https://github.com/google-research/google-research/tree/master/value_dice

  * •

EDM [31]:
https://github.com/wgrathwohl/JEM

Note that VDICE was originally designed for continuous actions with a Normal
distribution output which we adapt for the experiments by replacing with a
Gumbel-softmax.

references apalike


# New Perspectives on the Evaluation of Link Prediction Algorithms for Dynamic
Graphs

Raphaël Romero
Ghent University
<EMAIL_ADDRESS>&Tijl De Bie
Ghent University
<EMAIL_ADDRESS>&Jefrey Lijffijt
Ghent University
<EMAIL_ADDRESS>
###### Abstract

There is a fast-growing body of research on predicting future links in dynamic
networks, with many new algorithms. Some benchmark data exists, and
performance evaluations commonly rely on comparing the scores of observed
network events (positives) with those of randomly generated ones (negatives).
These evaluation measures depend on both the predictive ability of the model
and, crucially, the type of negative samples used. Besides, as generally the
case with temporal data, prediction quality may vary over time. This creates a
complex evaluation space.

In this work, we catalog the possibilities for negative sampling and introduce
novel visualization methods that can yield insight into prediction performance
and the dynamics of temporal networks. We leverage these visualization tools
to investigate the effect of negative sampling on the predictive performance,
at the node and edge level. We validate empirically, on datasets extracted
from recent benchmarks that the error is typically not evenly distributed
across different data segments. Finally, we argue that such visualization
tools can serve as powerful guides to evaluate dynamic link prediction methods
at different levels.

_Keywords_ Temporal Networks $\cdot$ Visualization $\cdot$ Link Prediction

##  1 Introduction

Many real-world phenomena such as biological ecosystems, networks of
computers, email exchanges, and social interactions can be understood as a set
of nodes interacting through time [1]. This type of data is commonly
represented as a dynamic graph. While initial attempts at modeling the
temporal evolution of networks typically aggregated the interactions into a
sequence of static graphs, recent efforts consider how to directly model time
continuously (i.e., without aggregation), preventing the loss of some fine-
grained temporal information during data preprocessing.

The resulting type of data is commonly framed _Continuous-Time Dynamic Graphs_
(CTDGs). Learning from such CTDGs has recently become an active field of
research, as suggested by recent benchmarks [2]. A crucial task thereof is
Dynamic Link Prediction (DLP), where the goal is to predict future links from
a history of observed ones. Successfully predicting interactions in a nearby
future can have numerous applications, such as recommending items on online
platforms, or predicting the evolution of epidemics through time.

However, the DLP task has faced a reproducibility crisis akin to the one
observed in static link prediction [3]. Indeed, while typical DLP evaluation
pipelines merely compared the scores of query events to the scores of
competing negative events selected at random, recent work has emphasized that
the way these negative events are selected can have a dramatic impact on DLP
performance measures, to the point that some sophisticated methods can get
outperformed by trivial, parameter free heuristics [3]. Thus, the awareness is
growing that dynamic link prediction performance measures not only depend on
the model quality, but crucially also on the strategy for sampling negative
events.

Furthermore, unlike conventional machine learning tasks targeting independent
and identically distributed data, prediction errors in DLP are not distributed
uniformly over the data. Hence, the use of visualization emerges as a vital
tool for effectively charting a model’s performance across different data
segments. However, in contrast to static graphs, Dynamic Graphs are
notoriously hard to visualize [4]. Longitudinal visualization, while
inherently limited [5], has not yet been applied to visualize predictive
performance of DLP algorithms. However, as seen in Figure 1, visualizations
such as Temporal Node Activity plots _can_ , to some extent, help to reveal
some aspects of the temporal structure of CTDGs. In this paper, we aim to
explore the opportunities visualization techniques offer in a predictive
context. Specifically, we propose the following contributions.

Contributions:

  * •

_Negative sampling_ (Section 3): We explore the diversity of negative sampling
strategies. To this aim, we propose a taxonomy separating negative sampling
strategies into ten categories, and show their use in evaluation leads to
different insights on the capabilities of a DLP algorithm. Through careful
evaluation of a state-of-the-art algorithm, we validate empirically the
sensitivity of the performance (measured by AUC) to negative sampling.

  * •

_Visualization_ (Section 4): We propose simple and effective lossless
visualization techniques called the Temporal Node Activity and Temporal Edge
Activity plot, enabling the visualization the activity of nodes and edges over
time without the need for aggregation. Finally, we illustrate two ways in
which the proposed visualization tools can yield more detailed insights into
the performance of DLP algorithms.

An open source implementation of the method and experiments is made available
at https://github.com/aida-ugent/dlp_viz.

Figure 1: Example of a Temporal Node Activity plot on the High School dataset.
The nodes are colored by class that the student is taking. The edges are
colored green if they connect two nodes from the same class and red otherwise.
Some structural information can be observed, such as the fact that the inter-
class edges are more likely to occur during break times (10 am, 12–2 pm, 4
pm).

##  2 Link Prediction on Dynamic Graphs

A Continuous-Time Dynamic Graph (CTDG) is defined as a time-ordered stream of
interactions $\mathcal{H}=\\{(u_{m},v_{m},t_{m})\ |\ m=1,\ldots,M\\}$, where
for the $m$-th event, $u_{m}$ and $v_{m}$ are source and destination nodes
belonging to a set of nodes $\mathcal{U}$, and $0\leq t_{1}\leq
t_{2}\leq\ldots\leq t_{M}$ are the timestamps of interactions. A DLP algorithm
is a function that computes a score $s(u,v,t\ |\ \mathcal{H}_{t})$ for any
event $(u,v,t)$ _conditioned_ on the past events denoted
$\mathcal{H}_{t}\overset{\Delta}{=}\\{(u_{m},v_{m},t_{m})\in\mathcal{H}\ |\
t_{m}<t\\}$. DLP algorithms are typically _trained_ on a training history
$\mathcal{H}_{\text{train}}\overset{\Delta}{=}\mathcal{H}_{t<t_{\text{train}}}$
and evaluated on a test history
$\mathcal{H}_{\text{test}}\overset{\Delta}{=}\mathcal{H}_{t\geq
t_{\text{train}}}$. The training set can be further split into train and
validation, with typical proportions of 70%, 15%, and 15% for the train,
validation, and test sets, respectively.

DLP algorithms are typically _evaluated_ by comparing the scores of _positive_
events $s(u,v,t\ |\mathcal{H}_{t})$ with the scores
$s(u^{\prime},v^{\prime},t\ |\mathcal{H}_{t})$ of competing events
$(u^{\prime},v^{\prime},t)$ occurring at time $t$ and involving an edge
$(u^{\prime},v^{\prime})$ other than the positive one. In theory, a perfect
model should be able to rank the positive event $(u,v,t)$ first in the list of
_all possible events_ at time $t$. However, generating this list of all
possible events, let alone calculating a score for each of them, is often
intractable. Instead, most evaluation methods select $K$ negative events at
random and compare the scores of the positive events with those of the
negative events. This setup is summarized in Figure 2. Crucially, the
_procedure_ used to generate the negative samples can have a dramatic impact
on the resulting Link Prediction measure.

Figure 2: A typical evaluation pipeline for Dynamic Link Prediction: the score
of each positive event is calculated along with the scores of $K$ negative
events. The resulting scores are then used to compute the AUC, MRR and AP
measures. The way the negative edges $(u_{i}^{(k)},v_{i}^{(k)})$ are selected
is a crucial component of the evaluation, as it can have a dramatic impact on
the performance measures.

DLP algorithms can be split into several main categories:

  * •

_Heuristic approaches_ consist in calculating a statistic for each edge
$(u,v)$ at time $t$, summarizing the historical information available for
predicting that edge, up to that time $t$. In this paper, we will consider two
simple instances of these heuristic approaches. EdgeBank (introduced in [3])
predicts $s(u,v,t\ |{\mathcal{H}}_{t})=1$ if the edge $(u,v)$ has an
interaction previous to t. Preferential Attachment (PA) (first described in
[6]) scores events $(u,v,t)$ as the product of their temporal degrees
$deg_{u}(t)\times deg_{v}(t)$, where $\deg_{u}(t)$ is the number of events
$(u^{\prime},v^{\prime},t)$ such that $u^{\prime}=u$ or $v^{\prime}=u$. While
simple, these baselines provide good reference points to assess the
performance of more sophisticated methods, as they provide a _controlled
source_ of false negatives and false positive. Indeed, Edgebank produces false
negatives whenever the edge $u,v$ involved in the event $(u,v,t)$, and false
positives whenever the edge involved in the negative event is historical.
Similarly, PA produces false negatives whenever any of the nodes involved in
the positive event is inductive, and false positives whenever both the nodes
involved in the negative event are historical.

  * •

Recent efforts, however, focused on _Dynamic Node Representation Learning_. As
surveyed in [7], these methods can be described by an encoder-decoder
architecture. The encoder is a function that maps each node $u$ to a time-
varying representation $h_{u}(t|\mathcal{H}_{t})$, while the decoder is a
function that maps a pair of nodes $(u,v)$ and a time $t$ to a score
$s(u,v,t|\mathcal{H}_{t})$. This approach extends concepts from the extensive
literature on static graph embedding, where the objective is to learn, for
each node, a representation that summarizes all relevant structural
information at the node level. A similar concept is applied in the context of
dynamic graphs, where the key distinction lies in allowing the representation
to evolve over time as more events are observed. In the remaining of this
paper we will consider the TGN-attn method, an instance of the Temporal Graph
Network [8] that combines graph attention and gated recurrent units into a
single temporal graph embedding method.

##  3 The Impact of Negative Sampling on DLP Evaluation

In this section we explore in detail the impact of the negative sampling
strategy on the DLP task.

###  3.1 Node and Edge-level Categories for Time-based Partitioning

In Section 2, we emphasized the importance of comparing scores for positive
and associated negative edges when evaluating the predictive abilities of a
DLP algorithm. The selection of these negative edges significantly impacts
AUC/MRR/AP scores. To address the challenge of negative edge selection, a
naive random approach may uniformly choose negative nodes $u^{(k)}$ and
$v^{(k)}$ from ${\mathcal{U}}$. However, this often results in overly easy
negative edges due to the low likelihood of most node pairs interacting at
time $t$[3]. A more targeted sampling strategy can avoid this issue of easy
negative samples. Poursafaei et al. [3] propose to distinguish edges that
appear only in the training set (Historical Edges; HE), only in the test set
(Inductive Edges; IE), or both. This approach generates more realistic and
challenging negative events.

Building further on this, we note that predicting for a node never seen during
training (Inductive Node; IN) is potentially more difficult than predicting
for a node that is observed during training (Historical Node; HN). Indeed, for
any HN, the model has direct access to relational behavior during training,
for instance enabling it to effectively position the node in a latent space.
In contrast, for an IN, the predictions will typically be solely based on  1.
connectivity information between ${\mathcal{H}}([t_{cut},t])$ and  2. eventual
node features,  making it a distinct challenge.

To summarize, a given edge $(u,v)$ sampled at random can be categorized in
different ways, as bundled below in Table 1.

Category Type |  Edge category |  Description of edge $(u,v)$
---|---|---
General |  Never Observed |  $(u,v)$ never appears in ${\mathcal{H}}$
Edge-Level |  Historical Edges |  $(u,v)$ appears in ${\mathcal{H}}_{train}$ but never in ${\mathcal{H}}_{test}$
Inductive Edges |  $(u,v)$ appears in ${\mathcal{H}}_{test}$ but never in ${\mathcal{H}}_{train}$
Overlap Edges |  $(u,v)$ appears in both ${\mathcal{H}}_{train}$ and ${\mathcal{H}}_{test}$
Node-Level |  Historical-to-Historical |  $u$ and $v$ both appear in ${\mathcal{H}}_{train}$ but never in ${\mathcal{H}}_{test}$
Historical-to-Inductive |  $u$ appears in ${\mathcal{H}}_{train}$ but never in ${\mathcal{H}}_{test}$, and $v$ appears in ${\mathcal{H}}_{test}$ but never in ${\mathcal{H}}_{train}$
Inductive-to-Inductive |  $u$ and $v$ both appear in ${\mathcal{H}}_{test}$ but never in ${\mathcal{H}}_{train}$
Overlap-to-Overlap |  $u$ and $v$ both appear in ${\mathcal{H}}_{train}$ and ${\mathcal{H}}_{test}$
Overlap-to-Historical |  $u$ appears in ${\mathcal{H}}_{train}$ and ${\mathcal{H}}_{test}$, and $v$ appears in ${\mathcal{H}}_{train}$ but never in ${\mathcal{H}}_{test}$
Overlap-to-Inductive |  $u$ appears in ${\mathcal{H}}_{train}$ and ${\mathcal{H}}_{test}$, and $v$ appears in ${\mathcal{H}}_{test}$ but never in ${\mathcal{H}}_{train}$
Table 1: A taxonomy of the possible negative edges $(u,v)$. Comparing the score of a positive event $s(u,v,t|\mathcal{H}_{t})$ against the score $s(u^{\prime},v^{\prime},t|\mathcal{H}_{t})$ of an edge belonging to either of these categories conveys different cues about the skills of the DLP algorithm. |  | TGN | EdgeBank | Preferential
---|---|---|---|---
Dataset | Negative Sampler |  |  | Attachment
uci | Destination | 0.882 | 0.547 | 0.503
Historical Edge | 0.800 | 0.054 | 0.161
Inductive Edge | 0.626 | 0.554 | 0.544
Historical Destination | 0.883 | 0.545 | 0.437
Inductive Destination | 0.875 | 0.554 | 0.720
wikipedia | Destination | 0.972 | 0.803 | 0.833
Historical Edge | 0.790 | 0.303 | 0.477
Inductive Edge | 0.663 | 0.803 | 0.693
Historical Destination | 0.894 | 0.799 | 0.440
Inductive Destination | 0.908 | 0.803 | 0.879
enron | Destination | 0.812 | 0.707 | 0.527
Historical Edge | 0.825 | 0.258 | 0.338
Inductive Edge | 0.671 | 0.758 | 0.688
Historical Destination | 0.813 | 0.706 | 0.502
Inductive Destination | 0.826 | 0.758 | 0.951
Table 2:  AUC Scores obtained by the 3 methods on 3 different datasets, using
5 different negative sampling strategies. In this example we see that while
the method performs well in general compared to heuristic baselines, it can
get outperformed by a simple heuristic when evaluated on a specific type of
negative samples.

###  3.2 Impact of Negative Sampling: An Empirical Evaluation

In order to illustrate concretely the impact of selecting different types of
negative edges on the prediction performance, we conduct a follow-up numerical
experiment using TGN, a state-of-the-art Dynamic Link Prediction method. We
consider the TGN-attn method, an instance of the Temporal Graph Network [8]
that combines graph attention and gated recurrent units into a single temporal
graph embedding method. We consider 3 datasets extracted from the recent
benchmark [3], and compare TGN-Attn with the EdgeBank and PA baselines,
against 5 types of negative samples.

The resulting comparison is shown in Table 2. As expected, in certain
instances, negative samples are easily identifiable since even basic
heuristics yield very high scores. For example, this is the case with PA and
Inductive Destination negative samples. In this scenario, these negative
samples consistently receive a score of 0 because, by definition, the
destination node linked to the negative event does not appear during the
testing phase. As a result, the score for positive events is always equal to
or greater than that of negative events. On the other hand, TGN attempts to
infer information about inductive destination nodes by employing a combination
of message passing and temporal attention. While this Node Inductivity feature
is beneficial, in the Enron Dataset, it may lead to an overestimation of event
scores involving these negative destination nodes.

##  4 Localizing DLP Errors using Longitudinal Plots

As we have seen, the performance of DLP algorithms can vary drastically
depending on the type of negative samples considered. Moreover, the prevalence
of each type of negative samples varies over time. In this section we propose
a way to visualize this temporal variation.

###  4.1 Temporal Edge and Node Activity plots

A Temporal Node Activity Plot for a CTDG
${\mathcal{H}}=\\{(u_{m},v_{m},t_{m}|m=1,\ldots,M)\\}$ is 2-dimensional
scatter plot representing for each node the timestamps where it is active.
Each node’s activity gets represented by a horizontal line with a specific
coordinate $y(u)$. The y-coordinates $\\{y(u|u\in{\mathcal{U}})\\}$ are
determined by ranking the nodes based on their first arrival time, breaking
ties with last arrival time. For each event $(u,v,t)$, points are plotted at
coordinates $(t,y(u))$ and $(t,y(u))$, representing the temporal activity of
each node as a _horizontal sequence_ , or _row_ of points.

Similarly, for each edge $u,v$ in the network, an _arrival rank_ can be
derived by sorting the edges $u,v$ in terms of their first interaction times,
and breaking the ties using their last interaction time. A Temporal Edge
Activity Plot can then be obtained by plotting each event $(u,v,t)$ associated
with the edge $u,v$ at the coordinate $(t,y(u,v))$. The activity of each
unique edge $(u,v)$ can thus be visualized as a horizontal sequence of points
at the y-coordinate $y(u,v)$.

###  4.2 Competing Events View

From our previous discussion on the impact of negative sampling on aggregated
measures such as the AUC, we understand the need to explore the measure and
visualize the performance across different data segments. More specifically,
it seems sensible to try understand first how the positive event ranks against
a carefully selected set of negative samples (e.g., Historical Destination
Node), depending on which node or edge it involves.

To further explore this question, we propose to view the task as a competition
between various event types, including the positive event. Each event is then
labeled and colored based on which candidate attains the best score. As can be
observed in Figure 3, when we sample negative events
$(u^{\prime},v^{\prime},t)$ such that $(u^{\prime},v^{\prime})$ appears
_exclusively_ in the test set, the model readily distinguishes these negative
events as negative during the training period, but exhibits more errors during
the testing period.

(a) Map of the Predictive error of TGN on a Temporal Node Activity Plot

(b) Map of the Predictive error of TGN on a Temporal Edge Activity Plot

Figure 3:  Visualizing the event with the highest score among the set of
(Positive, Historical Edge/Destination, and Negative Edge/Destination) over
time. Interestingly, on Figure 3(a), it can be seen that, during the training
period, the model scores the event link with an Inductive Destination node
higher than the positive event. This finding supports the notion that the
model requires a ‘warm-up’ set of interactions for each node before accurately
predicting its interactions.

(a) Prediction Rank of events visualized on a Temporal Node Activity plot

(b) Prediction Rank of events visualized on a Temporal Edge Activity plot

Figure 4: Longitudinal maps of the prediction rank of each _true_ event
$(u,v,t)$ among the list of $K$ negative events $(u^{(k)},v^{(k)},t)$. Only
the events having a prediction rank higher than 200 are shown.

(a) Prediction rank vs Node Degree

(b) Prediction rank vs. Edge Degree

Figure 5:  Plot of the prediction rank of the positive event $(u,v,t)$ as a
function of (a) the degree of the destination node $v$ or (b) the edge degree
$(u,v)$.

###  4.3 Visualizing the Prediction Rank of Each Event

Treating DLP as a binary classification leads to now well-understood flaws,
resulting often in unfaithful reporting of model performance, recent efforts
have proposed to instead treat the task as a ranking problem. In this context,
for instance, new benchmark datasets such as the ones proposed by Huang et al.
[2] for edge-level prediction tasks, include a carefully selected set of
$K=1000$ negative events for each positive event. Mean-Reciprocal Rank is then
used as a performance measure. Following the philosophy of the previous sub-
section, we argue that splitting this aggregated measure at the event level
and visualize it exactly on Longitudinal plots can lead to various insights.

Thus, we consider the TGBL-Wiki Dataset [2], splitted using the 70-15-15
proportions, and plot the rank of each event on the TNA and TEA plots on
Figure 4. On both the node-level and edge-level visualization, we observe that
events corresponding to previously unobserved nodes/edges typically have a
higher rank than events that are further in the list of events involving a
given node/edge. We further validate this observation by plotting the
prediction rank of the positive event $(u,v,t)$ as a function of the degree of
the destination node $v$ (Figure 5(a)) or of the edge degree $(u,v)$ (Figure
5(b)). On average, the rank of the positive event is higher for nodes/edges
that have been previously unobserved (i.e., possess a low degree at the time
of prediction). This observation aligns with the intuition that, for newly
encountered entities, the model has limited information about them, making it
more prone to misinterpreting a negative event as a positive one.

##  5 Related Work

Link Prediction is a field with a consistent body of research. Seminal work
includes [6, 9, 10, 11]. However, link prediction in CTDGs come with specific
datasets and challenges, that are still underexplored, as the recent TGB
challenge [2] indicates. Methods for predicting links in (discrete or
continuous) dynamic graphs have attracted substantial interest in recent
years, as surveyed by Kazemi et al. [7]. Competitive methods for DLP include
TGN-Attn [8] that we consider in this paper, but also [12, 13]. In this paper,
our aim is not to propose a new method, but rather to introduce a new
performance visualization method that can be used to improve existing methods
or design new methods.

In terms of the visualization methods considered here, we note that the
Temporal Edge Activity plots proposed in Section 4 are closely related to the
Temporal Edge Traffic plots from [3]. However, while in the latter the authors
considered exclusively edge-level (‘1-st order’) visualizations of the
network, we complement this by introducing node-level (‘0-th order’)
visualizations. Node Activity Maps [14] are also notably very similar to
Temporal Node Activity plots. However, they rely on a time aggregation of the
events, and have not been yet employed to visualize prediction performances of
DLP algorithms.

##  6 Conclusion

Recent academic efforts have been dedicated to standardizing Dynamic Link
Prediction (DLP) as a machine learning task. The goal is to equip the task
with its own evaluation pipelines, baseline methods, and benchmark datasets.
However, as a consequence of the high-dimensionality of the data and its non-
independent, non-identically distributed nature, deriving a single consistent
model validation has encountered numerous challenges. Central to these
challenges is negative sampling. In this paper, we assert that the uneven
error distribution across nodes/edges and time, might be at the source of
these challenges and provide concrete visual cues to support this claim.

Future work in this direction would involve comparing the performance
visualizations obtained through other methods. Moreover, plotting the activity
map of higher-order entities (e.g., cliques or triangles) within the network
could also aid in assessing how performance changes in different parts of the
network.

## Acknowledgements

The research leading to these results has received funding from the European
Research Council under the European Union’s Seventh Framework Programme
(FP7/2007-2013) (ERC Grant Agreement no. 615517), and under the European
Union’s Horizon 2020 research and innovation programme (ERC Grant Agreement
no. 963924), from the Special Research Fund (BOF) of Ghent University
(BOF20/IBF/117), from the Flemish Government under the “Onderzoeksprogramma
Artificiële Intelligentie (AI) Vlaanderen” programme, and from the FWO
(project no. G0F9816N, 3G042220).

## References

  * Holme and Saramäki [2021] Petter Holme and Jari Saramäki.  Temporal networks as a modeling framework, March 2021.
  * Huang et al. [2023] Shenyang Huang, Farimah Poursafaei, Jacob Danovitch, Matthias Fey, Weihua Hu, Emanuele Rossi, Jure Leskovec, Michael Bronstein, Guillaume Rabusseau, and Reihaneh Rabbany.  Temporal Graph Benchmark for Machine Learning on Temporal Graphs, July 2023.
  * Poursafaei et al. [2023] Farimah Poursafaei, Andy Huang, Kellin Pelrine, and Reihaneh Rabbany.  Towards Better Evaluation for Dynamic Link Prediction.  In _Thirty-Sixth Conference on Neural Information Processing Systems Datasets and Benchmarks Track_ , January 2023.
  * Holme and Saramäki [2019] Petter Holme and Jari Saramäki, editors.  _Temporal Network Theory_.  Computational Social Sciences. Springer International Publishing, Cham, 2019.
  * Linhares et al. [2019] Claudio D. G. Linhares, Jean R. Ponciano, Jose Gustavo S. Paiva, Bruno A. N. Travençolo, and Luis E. C. Rocha.  Visualisation of Structure and Processes on Temporal Networks.  In Petter Holme and Jari Saramäki, editors, _Temporal Network Theory_ , pages 83–105. Springer International Publishing, Cham, 2019.
  * Liben-Nowell and Kleinberg [2007] David Liben-Nowell and Jon Kleinberg.  The link-prediction problem for social networks.  _Journal of the American Society for Information Science and Technology_ , 58(7):1019–1031, 2007.
  * Kazemi et al. [2020] Seyed Mehran Kazemi, Rishab Goel, Kshitij Jain, Ivan Kobyzev, Akshay Sethi, Peter Forsyth, and Pascal Poupart.  Representation learning for dynamic graphs: A survey.  _The Journal of Machine Learning Research_ , 21(1):70:2648–70:2720, 2020.  ISSN 1532-4435.
  * Rossi et al. [2020] Emanuele Rossi, Ben Chamberlain, Fabrizio Frasca, Davide Eynard, Federico Monti, and Michael Bronstein.  Temporal graph networks for deep learning on dynamic graphs.  _arXiv_ , page 1–16, 2020.
  * Lichtenwalter et al. [2010] Ryan Lichtenwalter, Jake Lussier, and Nitesh Chawla.  New perspectives and methods in link prediction.  In _Proceedings of the ACM SIGKDD International Conference on Knowledge Discovery and Data Mining_ , pages 243–252, 2010.
  * Lu and Zhou [2011] Linyuan Lu and Tao Zhou.  Link Prediction in Complex Networks: A Survey.  _Physica A: Statistical Mechanics and its Applications_ , 390(6):1150–1170, 2011.  ISSN 03784371.
  * Yang et al. [2015] Yang Yang, Ryan N. Lichtenwalter, and Nitesh V. Chawla.  Evaluating Link Prediction Methods.  _Knowledge and Information Systems_ , 45(3):751–782, 2015.  ISSN 0219-1377, 0219-3116.
  * Trivedi et al. [2019] Rakshit Trivedi, Mehrdad Farajtabar, Prasenjeet Biswal, and Hongyuan Zha.  DYREP: Learning representations over dynamic graphs.  In _ICLR_ , 2019.
  * Wang et al. [2020] Yanbang Wang, Yen-Yu Chang, Yunyu Liu, Jure Leskovec, and Pan Li.  Inductive Representation Learning in Temporal Networks via Causal Anonymous Walks.  In _International Conference on Learning Representations_ , October 2020.
  * Linhares et al. [2017] Claudio D. G. Linhares, Bruno A. N. Travençolo, Jose Gustavo S. Paiva, and Luis E. C. Rocha.  DyNetVis: A system for visualization of dynamic networks.  In _Proceedings of the Symposium on Applied Computing_ , SAC ’17, pages 187–194, New York, NY, USA, April 2017. Association for Computing Machinery.