§ INTRODUCTION

Survival analysis and the comparison of event rates has been greatly
influenced over the last 50 years by the partial likelihood approach of
the Cox proportional hazard model <cit.>. This
approach provides a flexible way of assessing the influence of
covariates on the hazard function, without the need to specify a
parametric survival model. This flexibility comes at the cost of
decoupling the baseline hazard from the effect of the covariates. To
recover the whole survival curve—or the cumulative incidence function
(CIF)—we then need to separately estimate the baseline hazard
<cit.>. This in turn leads to stepwise estimates
of the survival function that can be difficult to interpret.

From the perspective of clinicians and their patients, the most relevant
quantity is often the 5- or 10-year risk of experiencing a certain event
given the patient's particular circumstances, and not the hazard ratio
between a treatment and control group. Therefore, to make sound clinical
decisions, it is important to accurately estimate the full hazard
function, which can subsequently be used to estimate the cumulative
incidence function (CIF). Using a parametric estimator of the hazard
function leads to a smooth function of time; as a consequence, the CIF
and the survival function estimates also vary smoothly over time.

With the goal of fitting smooth-in-time hazard functions, Hanley &
Miettinen  proposed a general framework
for estimating fully parametric hazard models via logistic regression.
Their main idea is simple: comparing person-moments when the event of
interest occurred with moments when patients were at risk. Their
approach handles censored data effortlessly and provides users familiar
with generalized linear models a natural way of fitting parametric
survival models. Moreover, their framework is very flexible: general
functions of time can be estimated (e.g. using splines or general
additive models), and hence these models retain some of the flexibility
of Cox's partial likelihood approach. And since the unit of analysis is
a person moment, time-varying covariates can also easily be included in
this framework.

Logistic regression had already been used in the context of
discrete-time survival modeling <cit.>. But in the
context of continuous-time survival modeling, using the framework of
logistic regression opens the door to an extensive array of powerful
modeling tools. Indeed, lasso and elastic-net regression can be used to
select variables that are associated with the hazard function. Logistic
regression can also be replaced by multinomial regression to analyse
data with multiple competing events.

In this article, we present the casebase package
<cit.> implemented in R <cit.>, for
parametric survival analysis that combines the ideas of Hanley &
Miettinen  into a simple interface. The
purpose of the casebase package is to provide practitioners with
an easy-to-use software tool to compute a patient's risk (or cumulative
incidence) of an event, conditional on a particular patient's covariate
profile. Our package retains the flexibility of case-base sampling and
the familiar interface of the glm function. It also provides
tools for variable selection and competing-risk analysis. In addition,
we provide extensive visualization tools.

In what follows, we first recall some theoretical details on case-base
sampling and its use for estimating parametric hazard functions. We then
give a short review of existing R packages that implement
comparable features as casebase. Next, we provide some details
about the implementation of case-base sampling in our package, and we
give a brief survey of its main functions. This is followed by four case
studies that illustrate the flexibility and capabilities of
casebase. We show how the same framework can be used for competing
risk analyses, penalized estimation, and for studies with time-dependent
exposures. Finally, we end the article with a discussion of the results
and of future directions.

§ THEORETICAL DETAILS

As discussed in Hanley & Miettinen , the
key idea behind case-base sampling is to sample from the study base a
finite amount of person moments. These person moments are indexed
by both an individual in the study and a time point, and therefore each
person moment has a covariate profile, an exposure status and an outcome
status attached to it. We note that there is only a finite number of
person moments associated with the event of interest (what Hanley &
Miettinen call the case series). Case-base sampling refers to the
sampling from the base of a representative finite sample called the
base series.

§.§ Likelihood and estimating

To describe the theoretical foundations of case-base sampling, we use
the framework of counting processes. In what follows, we abuse notation
slightly and omit any mention of \(\sigma\)-algebras. Instead, following
Aalen et al , we use the
placeholder “past” to denote the past history of the corresponding
process. The reader interested in more details can refer to Saarela &
Arjas  and Saarela
. First, let \(N_{i}(t) \in \{0, 1\}\) be
counting processes corresponding to the event of interest for individual
\(i=1, \ldots,n\). For simplicity, we will consider Type I censoring due
to the end of follow-up at time \(\tau\) (the general case of
non-informative censoring is treated in Saarela
). We assume a continuous time model, which
implies that the counting process jumps are less than or equal to one.
We are interested in modeling the hazard functions \(\lambda_{i}(t)\) of
the processes \(N_i(t)\), and which satisfy

\[\lambda_{i}(t) dt = E[dN_{i}(t)\mid \mathrm{past}].\]

Next, we model the base series sampling mechanism using non-homogeneous
Poisson processes \(R_i(t) \in \{0, 1, 2, \ldots\}\), with the
person-moments where \(dR_i(t) = 1\) constituting the base series. The
process \(Q_{i}(t) = R_i(t) + N_{i}(t)\) then counts both the case and
base series person-moments contributed by individual \(i\). This process
is typically defined by the user via its intensity function
\(\rho_i(t)\). The process \(Q_{i}(t)\) is characterized by
\(E[dQ_{i}(t)\mid\mathrm{past}] = \lambda_{i}(t)dt + \rho_i(t)dt\).

If the hazard function \(\lambda_{i}(t; \theta)\) is parametrized in
terms of \(\theta\), we could define an estimator \(\hat{\theta}\) by
maximization of the likelihood expression

\[L_0(\theta) = \prod_{i=1}^n \exp\left\{ -\int_0^{\min(t_i,\tau)} \lambda_i(t; \theta) dt \right\} \prod_{i=1}^{n} \prod_{t\in[0,\tau)} \lambda_{i}(t;\theta)^{dN_{i}(t)},\]

where \(\prod_{t\in[0,u)}\) represents a product integral from \(0\) to
\(u\), and where \(t_i\) is the event time for individual \(i\).
However, the integral over time makes the computation and maximization
of \(L_0(\theta)\) challenging.

Case-base sampling allows us to avoid this integral. By conditioning on
a sampled person-moment, we get individual likelihood contributions of
the form

\[P(dN_{i}(t) \mid dQ_{i}(t) = 1,\mathrm{past}) \stackrel{\theta}{\propto} \frac{\lambda_{i}(t; \theta)^{dN_{i}(t)}}{\rho_i(t) + \lambda_{i}(t;\theta)}.\]

Therefore, we can define an estimating function for \(\theta\) as
\begin{equation}
L(\theta) = \prod_{i=1}^{n} \prod_{t\in[0,\tau)} \left(\frac{\lambda_{i}(t; \theta)^{dN_{i}(t)}}{\rho_i(t) + \lambda_{i}(t;\theta)}\right)^{dQ_i(t)}. \label{eq:lik-function}
\end{equation}
When a logarithmic link function is used for modeling the
hazard function, the above expression is of a logistic regression form
with an offset term \(\log(1/\rho_i(t))\). Note that the sampling units
selected in the case-base sampling mechanism are person-moments, rather
than individuals, and the parameters to be estimated are hazards or
hazard ratios rather than odds or odds ratios. Generally, an individual
can contribute more than one person-moment, and thus the terms in the
product integral are not independent. Nonetheless, Saarela
showed that the logarithm of this
estimating function has mean zero at the true value \(\theta=\theta_0\),
and that the resulting estimator \(\hat{\theta}\) is asymptotically
normally distributed.

In Hanley & Miettinen , the authors
suggest sampling the base series uniformly from the study base.
In terms of Poisson processes, their sampling strategy corresponds
essentially to a time-homogeneous Poisson process with hazard equal to
\(\rho_i(t) = b/B\), where \(b\) is the number of sampled observations
in the base series, and \(B\) is the total population-time for the study
base (e.g. the sum of all individual follow-up times). More complex
examples are also possible; see for example Saarela & Arjas
, where the intensity functions for the
sampling mechanism are proportional to the cardiovascular disease event
rate given by the Framingham score. Non-uniform sampling mechanisms can
increase the efficiency of the case-base estimators.

§.§ Common parametric models

Let \(g(t; X)\) be the linear predictor such that
\(\log(\lambda(t;X)) = g(t; X)\). Different functions of \(t\) lead to
different parametric hazard models. The simplest of these models is the
one-parameter exponential distribution which is obtained by taking the
hazard function to be constant over the range of \(t\):

\begin{equation}
\log(\lambda(t; X)) = \beta_0 + \beta_1 X. \label{eq:exp}
\end{equation}

In this model, the instantaneous failure rate is independent of
\(t\).[The conditional chance of failure in a time interval of specified length is the same regardless of how long the individual has been in the study. This is also known as the memoryless property <cit.>.]

The Gompertz hazard model is given by including a linear term for time:

\begin{equation}
\log(\lambda(t; X)) = \beta_0 + \beta_1 t + \beta_2 X. \label{eq:gomp}
\end{equation}

Use of \(\log(t)\) yields the Weibull hazard which allows for a power
dependence of the hazard on time <cit.>:

\begin{equation}
\log(\lambda(t; X)) = \beta_0 + \beta_1 \log(t) + \beta_2 X. \label{eq:weibull}
\end{equation}

§.§ Competing-risk analysis

Case-base sampling can also be used in the context of competing-risk
analyses. Assuming there are \(J\) competing events, we can show that
each person-moment's contribution to the likelihood is of the form

\[\frac{\lambda_j(t)^{dN_j(t)}}{\rho(t) + \sum_{j=1}^J\lambda_j(t)},\]

where \(N_j(t)\) is the counting process associated with the event of
type \(j\) and \(\lambda_j(t)\) is the corresponding cause-specific
hazard function. As may be expected, this functional form is similar to
the terms appearing in the likelihood function for multinomial
regression.[Specifically, it corresponds to the following parametrization:
\begin{align*} \log\left(\frac{P(Y=j \mid X)}{P(Y = J \mid X)}\right) = X^T\beta_j, \qquad j = 1,\ldots, J-1.\end{align*}

§.§ Variable selection

To perform variable selection on the regression parameters
\(\theta \in \mathbb{R}^p\) of the hazard function, we can add a penalty
to the likelihood and optimise the following equation:
\begin{equation}
\min _{\theta \in \mathbb{R}^{p}}\,\,-\ell\left(\theta\right)+\sum_{j=1}^p w_j P(\theta_j;\lambda,\alpha) \label{eq:penest}
\end{equation}
where \(\ell\left(\theta\right) = \log L(\theta)\) is the
log of the likelihood function given in (<ref>),
\(P(\theta_j;\lambda,\alpha)\) is a penalty term controlled by the
non-negative regularization parameters \(\lambda\) and \(\alpha\), and
\(w_j\) is the penalty factor for the \(j\)th covariate. These penalty
factors serve as a way of allowing parameters to be penalized
differently. For example, we could set the penalty factor for time to be
0 to ensure it is always included in the selected model.

§ EXISTING PACKAGES

Survival analysis is an important branch of applied statistics and
epidemiology. Accordingly, there is already a vast ecosystem of
R packages implementing different methodologies. In this
section, we describe how the functionalities of casebase compare
to these packages.

At the time of writing, a cursory examination of CRAN's task view on
survival analysis reveals that there are over 250 packages related to
survival analysis <cit.>. For the purposes of this
article, we restricted our review to packages that implement at least
one of the following features: parametric modeling, non-proportional
hazard models, competing risk analysis, penalized estimation, and CIF
estimation. By searching for appropriate keywords in the
DESCRIPTION file of these packages, we found 60 relevant
packages. These 60 packages were then manually examined to determine
which ones are comparable to casebase. In particular, we excluded
packages that were focused on a different set of problems, such as
frailty and multi-state models. The remaining 14 packages appear in
Table <ref>, along with some of the functionalities they

Parametric survival models are implemented in a handful of packages:
CFC , flexsurv
, SmoothHazard ,
rsptm2 , mets
, and survival
. The types of models they allow vary for
each package. For example, SmoothHazard is limited to Weibull
distributions , whereas both flexsurv
and survival allow users to supply any distribution of their
choice. Also, flexsurv, smoothhazard, mets and
rstpm2 also have the ability to model the effect of time using
splines, which allows flexible modeling of the hazard function.
Moreover, flexsurv has the ability to estimate both scale and
shape parameters for a variety of parametric families. As discussed
above, casebase can model any parametric family whose log-hazard
can be expressed as a linear combination of covariates (including time).
Therefore, our package allows the user to model the effect of time using
splines. Also, by including interaction terms between covariates and
time, it also allows users to fit (non-proportional) time-varying
coefficient models. However, we do not explicitly model any shape
parameter, unlike flexsurv.

Several packages implement penalized estimation for the Cox model:
glmnet , glmpath
, penalized ,
riskRegression . Moreover, some
packages also include penalized estimation in the context of Cox models
with time-varying coefficients: elastic-net penalization with
CoxRidge  and rstpm2
, while survival
has an implementation of ridge-penalized
estimation. On the other hand, our package casebase provides
penalized estimation of the hazard function. To our knowledge,
casebase and rsptm2 are the only packages to offer this

Next, several R packages implement methodologies for
competing risk analysis; for a different perspective on this topic, see
Mahani & Sharabiani . The package
cmprsk provides methods for cause-specific subdistribution
hazards, such as in the Fine-Gray model
. On the other hand, the package
CFC estimates cause-specific CIFs from unadjusted, non-parametric
survival functions. Our package casebase also provides
functionalities for competing risk analysis by estimating parametrically
the cause-specific hazards. From these quantities, we can then estimate
the cause-specific CIFs.

Finally, several packages include functions to estimate the CIF. The
corresponding methods generally fall into two categories: transformation
of the estimated hazard function, and semi-parametric estimation of the
baseline hazard. The first category broadly corresponds to parametric
survival models, where the full hazard is explicitly modeled. Using this
estimate, the survival function and the CIF can be obtained using their
functional relationships (see Equations <ref> and <ref>
below). Packages providing this functionality include CFC,
flexsurv, mets, and survival. Our package
casebase also follows this approach for both single-event and
competing-risk analyses. The second category outlined above broadly
corresponds to semi-parametric models. These models do not model the
full hazard function, and therefore the baseline hazard needs to be
estimated separately in order to estimate the survival function. This is
achieved using semi-parametric estimators (e.g. Breslow's estimator) or
parametric estimators (e.g. spline functions). Packages that implement
this approach include riskRegression, rstpm2, and
survival. As mentioned in the introduction, a key distinguishing
factor between these two approaches is that the first category leads to
smooth estimates of the cumulative incidence function, whereas the
second category often produces estimates in the form of stepwise
functions. Providing smooth estimates of the CIF was one of the main
motivations for introducing case-base sampling in survival analysis.

Competing     Allows      Penalized                   Semi     Interval/Left     Absolute

Package            Risks      Non PH     Regression     Splines     Parametric     Parametric     Censoring     Risk

Comparison of various R packages for survival analysis based on several defining features. Competing Risks: whether or not an implementation for competing risks is present. Allows Non PH: permits models for non-proportional hazards. Penalized Regression: allows for a penalty term on the regression coefficients when estimating hazards (e.g. lasso or ridge). Splines: permits a flexible fit on time through the use of splines. Parametric: implementation for parametric models. Semi-parametric: implementation for semi-parametric models. Interval/left censoring: models for interval and left-censoring. If this is not selected, the package only handles right-censoring. Absolute Risk: computation for survival curves, cumulative incidence or cumulative hazard is readily available.

§ IMPLEMENTATION DETAILS

The functions in the casebase package can be divided into two
categories: 1) exploratory data analysis, in the form of population-time
plots; and 2) parametric modeling of the hazard function. We strove for
compatibility with both data.frames and data.tables; this
can be seen in the coding choices we made and the unit tests we wrote.

§.§ Population-time plots

Population-time plots are a descriptive visualization of incidence
density, where aggregate person-time is represented by area and events
as points within the area. The case-base sampling approach described in
Section <ref> can be visualized in the form of a population time
plot. These plots are extremely informative graphical displays of
survival data and should be one of the first steps in an exploratory
data analysis. The popTime function and plot method
facilitate this task:

* The casebase::popTime function takes as input the original
dataset along with the column names corresponding to the timescale,
the event status and an exposure group of interest (optional). This
will create an object of class popTime.

* The corresponding plot method for the object created in Step 1
can be called to create the population time plot with several options
for customizing the aesthetics.

By splitting these tasks, we give flexibility to the user. While the
method call in Step 2 allows further customization by using the
ggplot2 <cit.> family of functions, users may choose the
graphics system of their choice to create population-time plots from the
object created in Step 1.

To illustrate these functions, we will use data from the European
Randomized Study of Prostate Cancer Screening (ERSPC)
<cit.> which was extracted using the approach
described in Liu et al. . This
dataset is available through the casebase package. It contains the
individual observations for 159,893 men from seven European countries,
who were between the ages of 55 and 69 years when recruited for the

We first create the necessary dataset for producing the population time
plot using the popTime function. In this example, we stratify the
plot by treatment group. The resulting object inherits from class
popTime and stores the exposure variable as an attribute:

R> pt_object <- casebase::popTime(ERSPC, time = "Follow.Up.Time",
+                                event = "DeadOfPrCa", exposure = "ScrArm")
R> inherits(pt_object, "popTime")

#> [1] TRUE

R> attr(pt_object, "exposure")

#> [1] "ScrArm"

We then pass this object to the corresponding plot method:

R> plot(pt_object, add.base.series = TRUE)

[Population time plots for both treatment arms in the ERSPC dataset]Population time plots for both treatment arms in the ERSPC dataset. The gray area can be thought of as N=88,232 (control group) and N=71,661 (screening group) rows of infinitely thin rectangles (person-moments). More events are observed at later follow-up times, motivating the use of non-constant hazard models.

Figure <ref> is built sequentially by
first adding a layer for the area representing the population time in
gray, with subjects having the least amount of observation time plotted
at the top of the y-axis. We immediately notice a distinctive
stepwise shape in the population time area. This is due to the
randomization of the Finnish cohorts which were carried out on January 1
of each of year from 1996 to 1999. Coupled with the uniform December 31
2006 censoring date, this led to large numbers of men with exactly 11,
10, 9 or 8 years of follow-up. Tracked backwards in time (i.e. from
right to left), the population-time plot shows the recruitment pattern
from its beginning in 1991, and the January 1 entries in successive
years. Tracked forwards in time (i.e. from left to right), the plot for
the first three years shows attrition due entirely to death (mainly from
other causes). Since the Swedish and Belgian centres were the last to
complete recruitment in December 2003, the minimum potential follow-up
is three years. Tracked further forwards in time (i.e. after year 3) the
attrition is a combination of deaths and staggered entries. As we can
see, population-time plots summarise a wealth of information about the
study into a simple graph.

Next, layers for the case series and base series are added. The y-axis
location of each case moment is sampled at random vertically on the plot
to avoid having all points along the upper edge of the gray area. By
randomly distributing the cases, we can get a sense of the incidence
density. In Figure <ref>, we see that
more events are observed at later follow-up times. Therefore, a constant
hazard model would not be appropriate in this instance as it would
overestimate the cumulative incidence earlier on in time, and
underestimate it later on. Finally, the base series is sampled
horizontally with sampling weight proportional to their follow-up time.
The reader should refer to the package vignettes for more examples and a
detailed description of how to modify the aesthetics of a
population-time plot.

§.§ Parametric modeling

The parametric modeling step was separated into three parts:

* case-base sampling;

* estimation of the smooth hazard function;

* estimation of the CIF.

By separating the sampling and estimation functions, we allow the
possibility of users implementing more complex sampling scheme (as
described in Saarela ), or more complex
study designs (e.g. time-varying exposure).

The sampling scheme selected for sampleCaseBase was described in
Hanley & Miettinen : we first sample
along the “person” axis, proportional to each individual's total
follow-up time, and then we sample a moment uniformly over their
follow-up time. This sampling scheme is equivalent to the following
picture: imagine representing the total follow-up time of all
individuals in the study along a single dimension, where the follow-up
time of the next individual would start exactly when the follow-up time
of the previous individual ends. Then the base series could be sampled
uniformly from this one-dimensional representation of the overall
follow-up time. In any case, the output is a dataset of the same class
as the input, where each row corresponds to a person-moment. The
covariate profile for each such person-moment is retained, and an offset
term is added to the dataset. This output could then be used to fit a
smooth hazard function, or for visualization of the base series.

Next, the fitting function fitSmoothHazard starts by looking at
the class of the dataset: if it was generated from
sampleCaseBase, it automatically inherited the class
cbData. If the dataset supplied to fitSmoothHazard does
not inherit from cbData, then the fitting function starts by
calling sampleCaseBase to generate the base series. In other
words, users can bypass sampleCaseBase altogether and only worry
about the fitting function fitSmoothHazard.

The fitting function retains the familiar formula interface of
glm. The left-hand side of the formula should be the name of the
column corresponding to the event type. The right-hand side can be any
combination of the covariates, along with an explicit functional form
for the time variable. Note that non-proportional hazard models can be
achieved at this stage by adding an interaction term involving time
(cf. Case Study 4 below). The offset term does not need to be specified
by the user, as it is automatically added to the formula before calling

To fit the hazard function, we provide several approaches that are
available via the family parameter. These approaches are:

* glm: This is the familiar logistic regression.

* glmnet: This option allows for variable selection using the
elastic-net <cit.> penalty (cf. Case Study 3).
This functionality is provided through the glmnet package

* gam: This option provides support for Generalized
Additive Models via the mgcv package

* gbm: This option provides support for Gradient Boosted
Trees via the gbm package. This feature is still experimental.

In the case of multiple competing events, the hazard is fitted via
multinomial regression as performed by the VGAM package. We
selected this package for its ability to fit multinomial regression
models with an offset.

Once a model-fit object has been returned by fitSmoothHazard, all
the familiar summary and diagnostic functions are available:
print, summary, predict, plot, etc. Our
package provides one more functionality: it computes risk functions from
the model fit. For the case of a single event, it uses the familiar
\begin{equation}\label{eqn:surv}
S(t) = \exp\left(-\int_0^t \lambda(u;X) du\right).
\end{equation}
The integral is computed using either the numerical or
Monte-Carlo integration. The risk function (or cumulative incidence
function) is then defined as
\begin{equation}\label{eqn:CI}
CI(t) = 1 - S(t).
\end{equation}

For the case of a competing-event analysis, the event-specific risk is
computed using the following procedure: first, we compute the overall
survival function (i.e. for all event types):

\[ S(t) = \exp\left(-\int_0^t \lambda(u;X) du\right),\qquad \lambda(t;X) = \sum_{j=1}^J \lambda_j(t;X).\]

From this, we can derive the event-specific subdensities:

\[ f_j(t) = \lambda_j(t)S(t).\]

Finally, by integrating these subdensities, we obtain the event-specific
cumulative incidence functions:

\[ CI_j(t) = \int_0^t f_j(u)du.\]
Again, he integrals are computed using
either numerical integration (via the trapezoidal rule) or Monte Carlo
integration. This option is controlled by the argument method of
the absoluteRisk function.

In the following sections, we illustrate these functionalities in the
context of four case studies.

§ CASE STUDY 1—EUROPEAN RANDOMIZED STUDY OF PROSTATE CANCER

For our first case study, we return to the ERSPC study and investigate
the differences in risk between the control and screening arms.

§.§ Different functional forms of

We fit four models that differ in which functional form of time is used:
1) excluded from the linear predictor as seen in (<ref>), 2)
linear function as seen in (<ref>), 3) log function as seen in
(<ref>), and 4) a smooth function using cubic B-splines. The
models are fit using fitSmoothHazard with the familiar formula

R> fmla <- list(exponential = as.formula(DeadOfPrCa   ScrArm),
+              gompertz = as.formula(DeadOfPrCa   Follow.Up.Time + ScrArm),
+              weibull = as.formula(DeadOfPrCa   log(Follow.Up.Time) + ScrArm),
+              splines = as.formula(DeadOfPrCa   bs(Follow.Up.Time) + ScrArm))
R> fits <- lapply(fmla, function(i)
+   fitSmoothHazard(i, data = ERSPC, ratio = 100)
+ )

The output object from fitSmoothHazard inherits from the
singleEventCB and glm classes. As such, we can directly
use the summary generic:

R> summary(fits[["splines"]])

#> Coefficients:
#>                       Estimate Std. Error z value Pr(>|z|)
#> (Intercept)           -10.2905     0.3185  -32.31  < 2e-16
#> bs(Follow.Up.Time)1     4.3587     0.8008    5.44  5.2e-08
#> bs(Follow.Up.Time)2     2.0890     0.4721    4.42  9.7e-06
#> bs(Follow.Up.Time)3     4.6266     0.6910    6.70  2.1e-11
#> ScrArmScreening group  -0.2354     0.0886   -2.66   0.0079
#> (Dispersion parameter for binomial family taken to be 1)
#>     Null deviance: 6059.0  on 54539  degrees of freedom
#> Residual deviance: 5785.3  on 54535  degrees of freedom
#> AIC: 5795
#> Number of Fisher Scoring iterations: 9

Next, the absoluteRisk function takes as input the
fitSmoothHazard object and returns a matrix where each column
corresponds to the covariate profiles specified in the newdata
argument, and each row corresponds to a specified time point:

R> new_data <- data.frame(ScrArm = c("Control group", "Screening group"))
R> new_time <- seq(0,14,0.1)
R> risks <- lapply(fits, function(i)
+   absoluteRisk(object = i, time = new_time, newdata = new_data, method = "mont")
+ )

In Figure <ref>, we overlay the estimated CIFs from
casebase on the Cox model CIF. The CIF estimates for the
exponential model in panel (1) overestimate the cumulative incidence
earlier on in time, and underestimate it later on. Based on our earlier
discussion of the population-time plot, this poor fit for the
exponential hazard was expected. We notice a better fit with increasing
complexity of our model for time in Figure <ref>
(panels 2–4). As noted above, the usual asymptotic results hold for
likelihood ratio tests built using case-base sampling models. Therefore,
we can easily test the null hypothesis that the exponential model is
just as good as the larger (in terms of number of parameters) spline

#> Analysis of Deviance Table
#> Model 1: DeadOfPrCa   ScrArm + offset(offset)
#> Model 2: DeadOfPrCa   bs(Follow.Up.Time) + ScrArm + offset(offset)
#>   Resid. Df Resid. Dev Df Deviance Pr(>Chi)
#> 1     54538       6052
#> 2     54535       5785  3      267   <2e-16 ***
#> —
#> Signif. codes:  0 '***' 0.001 '**' 0.01 '*' 0.05 '.' 0.1 ' ' 1

The null hypothesis is rejected in favor of the spline model. Similarly,
the AIC provides further evidence that the flexible function of time
provides the best fit:

#>     Exp. Gompertz  Weibull  Splines
#>     6056     5821     5807     5795

CIFs for control and screening groups in the ERSPC data. In each of the panels, we plot the CIF from the Cox model using survival::survfit (solid line) and the CIF from the case-base sampling scheme (dashed line) with different functional forms of time. (1) The time variable is excluded (exponential). (2) Linear function of time (Gompertz). (3) The natural logarithm (Weibull). (4) Cubic B-spline expansion of time.

In Table <ref>, we present a side-by-side
comparison of the hazard ratios and confidence intervals estimated from
fitSmoothHazard and the corresponding parametric model using
survival::survreg, as well as the Cox model estimate. The hazard
ratio estimates and confidence intervals are similar across all four
models. This reinforces the idea that, under proportional hazards, we do
not need to model the full hazard to obtain reliable estimates of the
hazard ratio. Nevertheless, Figure <ref> shows that
different parametric models can still give rise to qualitatively
different estimates for the CIF.

Comparison of estimated hazard ratios and 95% confidence intervals for ERSPC data.
Model     casebase::fitSmoothHazard     survival::survreg

Exponential     0.79 (0.67, 0.94)     0.81 (0.68, 0.96)

Gompertz     0.80 (0.67, 0.95)     0.80 (0.67, 0.95)

Weibull     0.81 (0.68, 0.96)     0.80 (0.65, 0.96)

Splines     0.79 (0.66, 0.94)     –

3l0pt1emCox model estimate: HR (95% CI) = 0.80 (0.67, 0.95)

§.§ Time-dependent hazard

Previous re-analyses of these data suggested that the overall screening
attributed reduction in death due to prostate cancer of 20% was an
underestimate <cit.>. The estimated 20% (from a
proportional hazards model) did not account for the delay between
screening and the time the effect is expected to be observed. As a
result, the null effects in years 1–7 masked the substantial reductions
that began to appear from year 8 onwards. This motivates the use of a
time-dependent hazard ratio which can easily be fit with the
casebase package by including an interaction term with time in the

R> fit_inter <- fitSmoothHazard(DeadOfPrCa   bs(Follow.Up.Time) * ScrArm,
+                              data = ERSPC)

In Figure <ref>, we have the estimated hazard
ratio and 95% confidence interval for screening vs. control group as a
function of time using the plot method for objects of class

R> plot(fit_inter, type = "hr", newdata = new_data,
+      var = "ScrArm", xvar = "Follow.Up.Time", ci = TRUE)

The plot shows that the cures attributable to the screening only begin
to become statistically apparent by year 7 and later. The 25-60%
reductions seen in years 8-12 of the study suggests a much higher
reduction in prostate cancer due to screening than the single overall
20% reported in the original article.

[Estimated hazard ratio and 95% confidence interval for screening vs]Estimated hazard ratio and 95% confidence interval for screening vs. control group as a function of time in the ERSPC dataset. Hazard ratios are estimated from fitting a parametric hazard model as a function of the interaction between a cubic B-spline basis of follow-up time and treatment arm. 95% confidence intervals are calculated using the delta method. The plot shows that the effect of screening only begins to become statistically apparent by year 7. The 25-60% reductions seen in years 8-12 of the study suggests a much higher reduction in prostate cancer due to screening than the single overall 20% reported in the original article.

With this first case study, we explored how casebase allows us to
fit different parametric survival models with possible time-varying
effects, and how we can compare the fit of each model with tools from

§ CASE STUDY 2—BONE-MARROW

In the next case study, we show how case-base sampling can be used in
the context of a competing-risk analysis. For illustrative purposes, we
use the same data that was used in Scrucca et al
. The data was downloaded from the
first author's website, and it is now available as part of the
casebase package.

The data contains information on 177 patients who received a stem-cell
transplant for acute leukemia. The event of interest is relapse, but
other competing causes (e.g. death, progression, graft failure,
graft-versus-host disease) were also recorded. Several covariates were
captured at baseline: sex, disease type (acute lymphoblastic or
myeloblastic leukemia, abbreviated as ALL and AML, respectively),
disease phase at transplant (Relapse, CR1, CR2, CR3), source of stem
cells (bone marrow and peripheral blood, coded as BM+PB, or only
peripheral blood, coded as PB), and age.

First, we look at a population-time plot to visualize the incidence
density of both relapse and the competing events. In Figure
<ref>, failure times are highlighted on the plot using red
dots for the event of interest (panel A) and blue dots for competing
events (panel B). In both panels, we see evidence of a non-constant
hazard function: the density of points is larger at the beginning of
follow-up than at the end.

[Population-time plot for the stem-cell transplant study with both relapse and competing events]Population-time plot for the stem-cell transplant study with both relapse and competing events.

Our main objective is to compute the cumulative incidence of relapse for
a given set of covariates. We start by fitting a smooth hazard to the
data using a linear term for time:

We want to compare our hazard ratio estimates to that obtained from a
Cox regression. To obtain these estimates, we need to treat the
competing event as censoring:

R> library(survival)
R> # Treat competing event as censoring
R> model_cox <- coxph(Surv(ftime, Status == 1)   Sex + D + Phase + Source + Age,
+                    data = bmtcrr)

Estimates and confidence intervals for the hazard ratios for each coefficient. Both estimates from case-base sampling and Cox regression are presented.
1c      2cCase-Base     2cCox

(l3ptr3pt)2-3 (l3ptr3pt)4-5
Covariates     HR     95% CI     HR     95% CI

Sex     0.74     (0.43, 1.3)     0.68     (0.39, 1.19)

Disease     0.52     (0.29, 0.94)     0.52     (0.29, 0.93)

Phase (CR2 vs. CR1)     1.20     (0.48, 3)     1.21     (0.48, 3.02)

Phase (CR3 vs. CR1)     1.59     (0.41, 6.17)     1.67     (0.43, 6.5)

Phase (Relapse vs. CR1)     4.09     (1.9, 8.81)     4.55     (2.09, 9.9)

Source     1.70     (0.55, 5.23)     1.46     (0.47, 4.54)

Age     0.99     (0.97, 1.02)     0.99     (0.97, 1.02)

From the fit object, we can extract both the hazard ratios and their
corresponding confidence intervals. These quantities appear in Table
<ref>. As we can see, the only significant hazard ratio
identified by case-base sampling is the one associated with the phase of
the disease at transplant. More precisely, being in relapse at
transplant is associated with a hazard ratio of 3.89 when compared to

Given the estimate of the hazard function obtained using case-base
sampling, we can compute the absolute risk curve for a fixed covariate
profile. We perform this computation for a 35 year old woman who
received a stem-cell transplant from peripheral blood at relapse. We
compare the absolute risk curve for such a woman with ALL with that for
a similar woman with AML.

Next, we compare our estimates to that obtained from a corresponding
Fine-Gray model . The Fine-Gray model
is a semiparametric model for the cause-specific subdistribution
hazard, i.e. the function \(d_k(t)\) such that

\[CI_k(t) =1 - \exp\left( - \int_0^t d_k(u) \textrm{d}u \right),\]
\(CI_k(t)\) is the cause-specific cumulative incidence. The Fine-Gray
model allows to directly assess the effect of a covariate on the
subdistribution hazard, as opposed to the cause-specific hazard. For the
computation, we use the timereg package <cit.>:

R> library(timereg)
R> model_fg <- comp.risk(Event(ftime, Status)   const(Sex) + const(D) +
+                         const(Phase) + const(Source) + const(Age),
+                       data = bmtcrr, cause = 1, model = "fg")
R> # Estimate absolute risk curve
R> risk_fg <- predict(model_fg, newdata, times = time_points)

Absolute risk curve for a fixed covariate profile and the two disease groups. The estimate obtained from case-base sampling is compared to the Kaplan-Meier estimate.

Figure <ref> shows the absolute risk curves for both
case-base sampling and the Fine-Gray model. As we can see, the two
approaches agree quite well for AML; however, there seems to be a
difference of about 5% between the two curves for ALL. This difference
does not appear to be significant: the curve from case-base sampling is
contained within a 95% confidence band around the Fine-Gray absolute
risk curve (figure not shown).

§ CASE STUDY 3—SUPPORT DATA

In the first two case studies, we described the basic functionalities of
the casebase package: creating population-time plots, fitting
parametric models for hazard functions, and estimating the corresponding
cumulative incidence curves. For the third case study, we show how
casebase can also be used for variable selection through
regularized estimation of the hazard function as given by

To illustrate this functionality, we use the dataset from the Study to
Understand Prognoses Preferences Outcomes and Risks of Treatment
(SUPPORT) <cit.>. The SUPPORT dataset tracks death in
five American hospitals within individuals who are considered seriously
ill. The original data is available online from the Department of
Biostatistics at Vanderbilt University <cit.>. The cleaned
and imputed data consists of 9104 observations and 30 variables, and it
is available as part of the casebase package. In the comparisons
below, all covariates except sps and aps were used. These
two variables correspond to scores for predicting the outcome that were
developed as part of the original study. For more information about this
dataset, the reader is encouraged to look at the documentation in our

For our penalized case-base model, we opt for the natural log of time
which corresponds to a Weibull distribution. For fitting the penalized
hazard, we use fitSmoothHazard.fit, which is a matrix interface
to the fitSmoothHazard function. We supply both a matrix y
containing the time and event variables, and a matrix x
containing all other covariates. We apply the lasso penalty by setting
alpha = 1 and assign a penalty.factor (\(w_j\);
cf. Equation <ref>) of 0 to the time variable to ensure it is
in the selected model. We compare our approach to both Cox regression,
and lasso penalized Cox regression (fitted via the glmnet

To compare the performance of our models, we split the data into 95%
training and 5% test sets. To assess both discrimination and
calibration, we use a time-dependent version of the classical Brier
score that is adjusted for censoring <cit.>. The Brier
score can be used with both parametric and semi-parametric models. We
use the riskRegression package to compute these scores for all

R> # Create matrices for inputs
R> x <- model.matrix(death   . - d.time - aps - sps,
+                   data = train)[, -c(1)] # Remove intercept
R> y <- data.matrix(subset(train, select = c(d.time, death)))
R> # Regularized logistic regression to estimate hazard
R> pen_cb <- casebase::fitSmoothHazard.fit(x, y,
+   family = "glmnet",
+   time = "d.time", event = "death",
+   formula_time =   log(d.time), alpha = 1,
+   ratio = 10, standardize = TRUE,
+   penalty.factor = c(0, rep(1, ncol(x)))
+ )

In Figure <ref>, we show the coefficient estimates for
covariates that we selected by both penalized Cox and penalized
case-base. We note that both penalized approaches produce similar
results. We can also clearly see the shrinkage effect owing to the
\(\ell_1\) penalty.

Coefficient estimates from the Cox model (Cox), penalized Cox model using the glmnet package (Pen. Cox), and our approach using penalized case-base sampling (Pen. CB). Only the covariates that were selected by both penalized approaches are shown. The shrinkage of the coefficient estimates for Pen. Cox and Pen. CB occurs due to the $\ell_1$ penalty.

We then compare the cumulative incidence estimation over the test set.
The probabilities over time for each observation are averaged, resulting
in the absolute risk curves shown in Figure <ref> (A).
We can see some minimal differences between the three models, with the
Kaplan-Meier giving the lowest estimates across follow-up-time. Note
that the apparent smoothness of the Cox and penalized Cox curves is due
to the large number of observations in the training set, which is used
to derive the Breslow estimate of the baseline hazard. As described
above, we compare the performance between the models by computing the
Brier scores over time. In Figure <ref> (B), we can
see that the Brier score is larger for the Kaplan-Meier estimate than
for the other three models. On the other hand, the differences between
these three models are minimal.

Comparison of Cox regression (Cox), penalized Cox regression (Pen. Cox), penalized case-base sampling estimation (Pen. CB), and Kaplan-Meier (K-M). (A)  Probability of death as a function of follow-up time. (B) Brier score as a function of follow-up time, where a lower score corresponds to better performace.

In this third case study, we showed how case-base sampling can be used
in conjunction with penalized logistic regression to perform variable
selection in survival models.

§ CASE STUDY 4—STANFORD HEART TRANSPLANT

In the previous case studies, we only considered covariates that were
fixed at baseline. In this next case study, we use the Stanford Heart
Transplant data <cit.>
to show how case-base sampling can also be used in the context of
time-dependent exposure. This feature of case-base sampling has been
explored in the literature, in the context of vaccination safety
<cit.>. In this study, the exposure period was defined
as the week following vaccination. Hence, the main covariate of
interest, i.e. exposure to the vaccine, was changing over time. In this
context, case-base sampling offers an efficient alternative to nested
case-control designs or self-matching.

Recall the setting of Stanford Heart Transplant study: patients were
admitted to the Stanford program after meeting with their physician and
determining that they were unlikely to respond to other forms of
treatment. After enrollment, the program searched for a suitable donor
for the patient, which could take anywhere between a few days to almost
a year. We are interested in the effect of a heart transplant on
survival; therefore, the patient is considered exposed only after the
transplant has occurred.

As above, we can look at the population-time plot for a graphical
summary of the event incidence (see Figure <ref>. Here,
we colour the exposed person-time (i.e. after transplant) in a darker
shade of gray. As we can see, most events occur early during the
follow-up period, and therefore we do not expect the hazard to be

Population-time for the Stanford Transplant study. The darker shade of gray corresponds to the exposed person-time.

Since the exposure is time-dependent, we need to manually define the
exposure variable after case-base sampling and before
fitting the hazard function. For this reason, we use the
function directly.

R> cb_data <- sampleCaseBase(jasa,
+   time = "futime",
+   event = "fustat", ratio = 100
+ )

Next, we compute the number of days from acceptance into the program to
transplant, and we use this variable to determine whether each
population-moment is exposed or not.

R> # Define exposure variable
R> cb_data <- mutate(cb_data,
+   txtime = time_length(accept.dt
+                        unit = "days"
+   ),
+   exposure = case_when(
+     is.na(txtime)   0L, # No transplant
+     txtime > futime   0L,
+     txtime <= futime   1L
+   )
+ )

Finally, we can fit the hazard using various linear predictors.

R> # Fit several models
R> fit1 <- fitSmoothHazard(fustat   exposure,
+   data = cb_data, time = "futime"
+ )
R> fit2 <- fitSmoothHazard(fustat   exposure + futime,
+   data = cb_data, time = "futime"
+ )
R> fit3 <- fitSmoothHazard(fustat   exposure * futime,
+   data = cb_data, time = "futime"
+ )

Note that the third model includes an interaction term between exposure
and follow-up time. In other words, this model no longer exhibits
proportional hazards. The evidence of non-proportionality of hazards in
the Stanford Heart Transplant data has been widely discussed

We can then compare the goodness of fit of these three models using the
Akaike Information Criterion (AIC).

#> Model1 Model2 Model3
#>    827    791    790

As we can see, the best fit is the third model. By visualizing the
hazard functions for both exposed and unexposed individuals, we can more
clearly see how the hazards are no longer proportional. We can easily
obtain a plot of the hazards by using the

R> plot(fit3, hazard.params = list(xvar = "futime",
+                                 by = "exposure",
+                                 alpha = 0.05,
+                                 ylab = "Hazard",
+                                 data = cb_data))

Plot of the hazard function for exposed and unexposed individuals.

Looking at Figure <ref>, the non-proportionality seems
to be more pronounced at the beginning of follow-up than the end.
Finally, in Figure <ref>, we turn these estimates of the
hazard function into estimates of the cumulative incidence functions.

R> # Compute absolute risk curves
R> newdata <- data.frame(exposure = c(0, 1))
R> absrisk <- absoluteRisk(fit3,
+   newdata = newdata,
+   time = seq(0, 425, length.out = 100)
+ )
R> class(absrisk)

#> [1] "absRiskCB" "matrix"    "array"

R> plot(absrisk,
+      id.names = c("No Tx", "Tx")) +
+   ylab("Cumulative Incidence") +
+   xlab("Follow-up time (in days)") +
+   paper_gg_theme

Plot of the cumulative incidence function for exposed and unexposed individuals.

As we can see in the above case-study, the casebase package can
also be used to model time-varying exposures and non-proportional hazard

§ DISCUSSION

In this article, we presented the R package casebase,
which provides functions for fitting smooth parametric hazards and
estimating CIFs using case-base sampling. Our package also provide
several functions to produce graphical summaries of the data and the
results. We outlined the theoretical underpinnings of the approach, we
provided details about our implementation, and we illustrated the merits
of the approach and the package through four case studies.

As a methodological framework, case-base sampling is very flexible. Some
of this flexibility has been explored before in the literature: for
example, Saarela and Hanley  used case-base
sampling to model a time-dependent exposure variable in a vaccine safety
study. As another example, Saarela and Arjas
combined case-base sampling and a Bayesian
non-parametric framework to compute individualized risk assessments for
chronic diseases. In the case studies above, we further explored this
flexibility along two fronts. On the one hand, we showed how splines
could be used as part of the linear predictor to model the effect of
time on the hazard. This strategy yielded estimates of the survival
function that were qualitatively similar to semiparametric estimates
derived from Cox regression; however, case-base sampling led to
estimates of the survival function that vary smoothly in time and
are thus easier to interpret. On the other hand, we also displayed the
flexibility of case-base sampling by showing how it could be combined
with penalized logistic regression to perform variable selection.
Furthermore, the second and fourth case studies showed how case-base
sampling can respectively be applied to competing risks and time-varying
exposure settings. Even though we did not illustrate it in this article,
case-base sampling can also be combined with the framework of
generalized additive models. This functionality has already been
implemented in the package. Similarly, case-base sampling can be
combined with quasi-likelihood estimation to fit survival models that
can account for the presence of over-dispersion. All of these examples
illustrate how the case-base sampling framework in general, and the
package casebase in particular, allows the user to fit a broad and
flexible family of survival functions.

As presented in Hanley & Miettinen ,
case-base sampling is comprised of three steps: 1) sampling a case
series and a base series from the study; 2) fit the log-hazard as a
linear function of predictors (including time); and 3) use the fitted
hazard to estimate the CIF. Accordingly, our package provides functions
for each step. Moreover, the simple interface of the
fittingSmoothHazard function resembles the glm interface.
This interface should look familiar to new users. Our modular approach
also provides a convenient way to extend our package for new sampling or
fitting strategies.

In the case studies above, we compared the performance of case-base
sampling with that of Cox regression and Fine-Gray models. In terms of
function interface, casebase uses a formula interface that is
closer to that of glm, in that the event variable is the only
variable appearing on the left-hand side of the formula. By contrast,
both survival::coxph and timereg::comp.risk use arrays
that capture both the event type and time. Both approaches to modeling
yield user-friendly code. However, in terms of output, both approaches
differ significantly. Case-base sampling produces smooth hazards and
smooth cumulative incidence curves, whereas Cox regression and Fine-Gray
models produce stepwise CIFs and never explicitly model the hazard
function. Qualitatively, we showed that by using splines in the linear
predictor, all three models yielded similar curves. However, the smooth
nature of the output of casebase provides a more intuitive
interpretation for consumers of these predictions. In Table
<ref>, we provide a side-by-side comparison between the
Cox model and case-base sampling.

Comparison between the Cox model and case-base sampling
Feature     Cox model     Case-base sampling

Model type     Semi-parametric     Fully parametric

Time     Left hand side of the formula     Right hand side (allows flexible modeling of time)

Cumulative incidence     Step function     Smooth-in-time curve

Non-proportional hazards     Interaction of covariates with time     Interaction of covariates with time

Model testing          Use GLM framework (e.g. LRT, AIC, BIC)

Competing risks     Difficult     Cause-specific CIFs

Our choice of modeling the log-hazard as a linear function of covariates
allows us to develop a simple computational scheme for estimation.
However, as a downside, it does not allow us to model location and scale
parameters separately like the package flexsurv. For example, if
we look at the Weibull distribution as parametrised in
stats::pweibull, the log-hazard function is given by

\[ \log \lambda(t; \alpha, \beta) = \left[\log(\alpha/\beta) - (\alpha - 1)\log(\beta)\right] + (\alpha - 1)\log t,\]

where \(\alpha,\beta\) are shape and scale parameters, respectively.
Unlike casebase, the approach taken by flexsurv also allows
the user to model the scale parameter as a function of covariates. Of
course, this added flexibility comes at the cost of interpretability: by
modeling the log-hazard directly, the parameter estimates from
casebase can be interpreted as estimates of log-hazard ratios. To
improve the flexibility of casebase at capturing the scale of a
parametric family, we could replace the logistic regression with its
quasi-likelihood counterpart and therefore model over- and
under-dispersion with respect to the logistic likelihood. We defer the
study of the properties and performance of such a model to a future

Future work will look at some of the methodological extensions of
case-base sampling. First, to assess the quality of the model fit, we
would like to study the properties of the residuals (e.g. Cox-Snell,
martingale). More work needs to be done to understand these residuals in
the context of the partial likelihood underlying case-base sampling. The
resulting diagnostic tools could then be integrated in this package.
Also, we are interested in extending case-base sampling to account for
interval censoring. This type of censoring is very common in
longitudinal studies, and many packages (e.g. SmoothHazard,
survival and rstpm2) provide functions to account for it.
Again, we hope to include any resulting methodology as part of this

In future versions of the package, we also want to increase the
complement of diagnostic and inferential tools that are currently
available. For example, we would like to include the ability to compute
confidence intervals for the cumulative incidence curve. The delta
method or parametric bootstrap are two different strategies we can use
to construct approximate confidence intervals. Furthermore, we would
like to include more functions to compute calibration and discrimination
statistics (e.g. AUC) for our models. Saarela and Arjas
also describe how to obtain a posterior
distribution for the AUC from their model. Their approach could
potentially be included in casebase. Finally, we want to provide
more flexibility in how the case-base sampling is performed. This could
be achieved by adding a hazard argument to the function
sampleCaseBase. In this way, users could specify their own
sampling mechanism. For example, they could provide a hazard that gives
sampling probabilities that are proportional to the cardiovascular
disease event rate given by the Framingham score <cit.>.

In conclusion, we presented the R package casebase
which implements case-base sampling for fitting parametric survival
models and for estimating smooth cumulative incidence functions using
the framework of generalized linear models. We strongly believe that its
flexibility and its foundation on the familiar logistic regression model
will make it appealing to new and established practitioners. The
casebase package is freely available from the Comprehensive
R Archive Network at
<https://cran.r-project.org/package=casebase>. Interested users can
visit <http://sahirbhatnagar.com/casebase/> for detailed package
documentation and vignettes.

§ ACKNOWLEDGMENTS

We would like to thank Yi Yang for helpful discussions on penalized
regression models and gradient boosting.  gratefully
acknowledges funding via a Discovery Grant from the Natural Sciences and
Engineering Research Council of Canada (NSERC).


* Manton and Merabet (1997) N. S. Manton and H. Merabet, Nonlinearity 10, 3 (1997), eprint hep-th/9605038.
  * Vilenkin and Shellard (2000) A. Vilenkin and E. S. Shellard, _Cosmic Strings and Other Topological Defects_ (Cambridge University Press, 2000), ISBN 978-0-521-65476-0.
  * Maeda and Turok (1988) K.-i. Maeda and N. Turok, Phys. Lett. B 202, 376 (1988).
  * Gregory et al. (1990) R. Gregory, D. Haws, and D. Garfinkle, Phys. Rev. D 42, 343 (1990).
  * Gregory (1991) R. Gregory, Phys. Rev. D 43, 520 (1991).
  * Anderson et al. (1997) M. R. Anderson, F. Bonjour, R. Gregory, and J. Stewart, Phys. Rev. D 56, 8014 (1997), eprint hep-ph/9707324.
  * Witten (1985) E. Witten, Nucl. Phys. B 249, 557 (1985).
  * Davis and Shellard (1989) R. L. Davis and E. P. S. Shellard, Nucl. Phys. B 323, 209 (1989).
  * Rayleigh (1887) L. Rayleigh, The London, Edinburgh, and Dublin Philosophical Magazine and Journal of Science 24, 145 (1887).
  * Kovacic et al. (2018) I. Kovacic, R. Rand, and S. Mohamed Sah, Applied Mechanics Reviews 70 (2018).
  * Mettler (1967) E. Mettler, in _Dynamic Stability of Structures_ , edited by G. Herrmann (Pergamon, 1967), pp. 169–188, ISBN 978-1-4831-9821-7, URL https://www.sciencedirect.com/science/article/pii/B9781483198217500156.
  * Saurabh et al. (2020) A. Saurabh, T. Vachaspati, and L. Pogosian, Phys. Rev. D 101, 083522 (2020), eprint 2001.01030.
  * Mur (1981) G. Mur, IEEE Transactions on Electromagnetic Compatibility EMC-23, 377 (1981).


# RoKEPG: RoBERTa and Knowledge Enhancement for Prescription Generation of
Traditional Chinese Medicine

Hua Pu1, Jiacong Mi1, Shan Lu2, Jieyue He1,* *Corresponding author 1School of
Computer Science and Engineering, Key Lab of Computer Network and Information
Integration, MOE,
Southeast University, Nanjing, 210018, Jiangsu, China
2Nanjing Fenghuo Tiandi Communication Technology Co., Ltd, Nanjing, 211161,
Jiangsu, China<EMAIL_ADDRESS><EMAIL_ADDRESS><EMAIL_ADDRESS><EMAIL_ADDRESS>
###### Abstract

Traditional Chinese medicine (TCM) prescription is the most critical form of
TCM treatment, and uncovering the complex nonlinear relationship between
symptoms and TCM is of great significance for clinical practice and assisting
physicians in diagnosis and treatment. Although there have been some studies
on TCM prescription generation, these studies consider a single factor and
directly model the symptom-prescription generation problem mainly based on
symptom descriptions, lacking guidance from TCM knowledge. To this end, we
propose a RoBERTa and Knowledge Enhancement model for Prescription Generation
of Traditional Chinese Medicine (RoKEPG). RoKEPG is firstly pre-trained by our
constructed TCM corpus, followed by fine-tuning the pre-trained model, and the
model is guided to generate TCM prescriptions by introducing four classes of
knowledge of TCM through the attention mask matrix. Experimental results on
the publicly available TCM prescription dataset show that RoKEPG improves the
$F_{1}$ metric by about 2% over the baseline model with the best results.

###### Index Terms:

Deep learning; Traditional Chinese medicine prescription generation; Pre-
trained language models

##  I Introduction

TCM is increasingly being used clinically worldwide, and its efficacy is
recognized by more and more people and countries [1]. TCM has accumulated a
large body of literature and treatment records. A prescription consisting of
TCM is the most critical form of TCM treatment, and it is widely used to treat
various diseases. As of now, over 100,000 classical TCM prescriptions have
been accumulated [2], offering clinicians references and enabling
computational models for computer-assisted therapy and automatic prescription
generation through learning. In TCM prescription generation, herbs for
treating symptoms are generated from symptom text. Recent deep learning
breakthroughs enable mining complex TCM principles.

Although some studies on TCM prescription generation exist, such as PTM [3],
MC-eLDA [4], AttentiveHerb [5], and TCM Translator [6], most of these studies
only consider symptom factors to directly model the symptom-prescription
generation problem, resulting in a lack of guidance of TCM knowledge for model
generation. Moreover, there is a complex nonlinear relationship between the
composition of TCM prescriptions and patient symptoms, and one herb may treat
multiple symptoms, while one symptom may correspond to multiple herbs. Hence,
it is challenging for traditional machine learning methods to fit the complex
relationship between herbs and symptoms adequately. At the same time,
thousands of herbs are covered in TCM prescriptions, and the unbalanced
distribution of TCM data also poses a challenge regarding the accuracy of TCM
prescription generation. Therefore, we take advantage of RoBERTa [7] and
propose a Knowledge Enhancement model for Prescription Generation of
Traditional Chinese Medicine (RoKEPG) to address the above issues.

Our main contributions are as follows:

  * •

We performed a secondary pre-training of the RoBERTa model based on a TCM
domain corpus which we constructed with 13.5 million words.

  * •

We conduct fine-tuning on the pre-trained RoBERTa model based on the attention
mask matrix to introduce the knowledge of nature, taste, channel tropism and
effect of herbs for guiding the model to generate TCM prescriptions.

  * •

The RoKEPG model is tested on a publicly available TCM prescription dataset by
a series of experiments. The experimental results verify the effectiveness of
the RoKEPG model compared with the state-of-the-art related models.

##  II Related Work

This section reviews the work related to prescription generation for TCM.
Early studies mainly relied on topic models for TCM prescription generation
[3, 4]. However, there are over 1000 commonly used herbs, and some of the
correlations between herbs and symptoms are highly complex and challenging to
compute, making topic models insufficient for the TCM prescription generation
task. Subsequently, some scholars modeled TCM prescription generation as a
Multilabel Classification (MLC) problem, considering each herb as a category
label and predicting probabilities [8, 9, 10, 11], but these methods did not
fully utilize the correlations between labels. Inspired by Neural Machine
Translation (NMT) models, research emerged on using Sequence to Sequence
(Seq2Seq) models for generating TCM prescriptions [5, 6, 12, 13].

In recent years, pre-trained language models have rapidly developed, and
research on TCM prescription generation using pre-trained language models has
also emerged [14, 15]. By harnessing the powerful semantic capabilities
obtained from pre-training on large-scale unsupervised corpora, models can
achieve good quality patient prescriptions with fine-tuning on small sample
data. Additionally, methods based on heterogeneous graph contrastive learning
have been used to generate TCM prescriptions [16].

The above studies on TCM prescription generation mainly focused on directly
modeling the symptom-TCM prescription generation problem based on symptoms.
While some studies introduced the effects of herbs to guide prescription
generation, these factors were relatively single. This inspired the current
study to consider more influencing factors in the TCM prescription generation
task, such as the nature, taste, and channel tropism of herbs, aiming to make
the model generation process more in line with the realistic scenario of
prescription writing by physicians.

## Methods

###  II-A Task Definition

We first define the TCM prescription generation task, which is modelled based
on the symptom text-prescription composition data pair, and finally achieve
the generation of a set of herbs for treating symptom text. Specifically, TCM
prescription generation refers to generating a set of herbs corresponding to
the treatment symptoms $Y=\left[y_{1},y_{2}\ldots,y_{m}\right]$ based on the
symptoms $X=\left[x_{1},x_{2}\ldots,x_{n}\right]$ of a TCM prescription, where
$n$ is the length of the symptom text contained in the prescription and $m$ is
the number of herbs contained in the prescription.

###  II-B Overview

The RoKEPG model consists of two main parts: pre-training and fine-tuning. In
the pre-training stage, RoBERTa [7], pre-trained language model, is utilized
and trained on the TCM corpus that we constructed. Subsequently, in the fine-
tuning stage, the model generated during pre-training is used to generate TCM
prescriptions. During fine-tuning, TCM knowledge is incorporated through the
attention mask matrix to guide the model’s generation process. The overall
framework of the RoKEPG model is shown in Figure 1, where the white squares in
the fine-tuning phase indicate contextual information available when
performing sequence-to-sequence learning, such as for setting symptom texts to
do bidirectional attention when performing training. In contrast, grey squares
indicate information that is not visible through the mask, such as for setting
up the prediction of herbs to do unidirectional attention and for setting up
knowledge that is not a counterpart of that herb to be invisible.

Figure 1: Framework diagram of RoKEPG model

###  II-C Pre-training

To improve the performance on the TCM prescription generation task, we pre-
train the RoBERTa [7] language model on a domain-specific TCM corpus. As
publicly available pre-trained language models are generally based on general
domain corpora, using them directly for domain-specific tasks may not yield
satisfactory results. Therefore, we constructed a TCM corpus containing 13.5
million words by processing relevant TCM books like Basic Theory of
Traditional Chinese Medicine [17] and Pharmacology of Traditional Chinese
Medical Formulae [18]. After constructing the TCM corpus, we proceed to pre-
train the model using this corpus, as shown in Figure 1.

The input of the RoBERTa model contains three parts: Token Embedding, Segment
Embedding and Position Embedding. The three parts are summed and fed into the
architecture containing the L-layer Transformer. For each layer $l$, the Self-
Attention Head is used to aggregate the output from the previous layer and
learn the contextual representation $H$. The formula for the Self-Attention
Head $A_{l}$ is shown below:

$\begin{cases}Q=H^{l-1}W_{l}^{Q}\\\ K=H^{l-1}W_{l}^{K}\\\
V=H^{l-1}W_{l}^{V}&\end{cases}$ (1)
$A_{l}=softmax\left(\frac{QK^{T}}{\sqrt{d_{k}}}\right)V_{l}$ (2)

where $Q$ is the query matrix, $K$ is the key matrix, $V$ is the value matrix,
$H^{l-1}$ is the output vector from layer $l-1$, and $W_{l}^{Q}$, $W_{l}^{K}$,
and $W_{l}^{V}$ are the parameter matrices.

###  II-D Fine-tuning

####  II-D1 Overall Description of the Fine-tuning Phase

UniLM [19] is a unified pre-trained language model with natural language
understanding (NLU) and natural language generation (NLG) capabilities, which
enables the model to implement Seq2Seq through Attention Mask. Inspired by the
UniLM model, we enable the RoBERTa model to implement Seq2Seq learning by
changing the attention matrix masks. Specifically, based on the model weights
obtained in the pre-training phase, mask fine-tuning on RoBERTa to generate
TCM prescriptions. The objective function of the TCM prescription generation
task:

$\arg\max_{y\in V}P(Y|X)$ (3)
$P(Y|X,\theta)=\prod_{t=1}^{m}p(y_{t}|y_{t-1},y_{t-2},\ldots,y_{1};X;\theta)$
(4)

where $\theta$ is the parameter for model training, $V$ is the dictionary
corresponding to the model, and the goal of model training is to maximize the
joint probability of generating TCM prescriptions under the condition that the
symptoms are $X$.

Since the original RoBERTa model predicts individual words by masking, this
masking approach cannot capture the complete semantic relationships of TCMs,
and at the same time, increases the difficulty of the model prediction.
Therefore, we choose to perform the mask prediction for the full name of the
Chinese medicine in the composition of the TCM prescription.

####  II-D2 Input Layer

Firstly, the symptom and prescription composition corresponding to each TCM
prescription in the dataset is used as a sequence to the source sequence $X$
and the target sequence $Y$ in the sequence task, respectively. The symptom
and prescription composition sequences $<X,Y>$ are concatenated with the
special characters ”[CLS]” and ”[SEP]” to perform splicing:

$S=\left(\left[CLS\right],X,\left[SEP\right],Y,\left[SEP\right]\right)$ (5)

where ”[CLS]” indicates the beginning of the sequence and ”[SEP]” is used to
separate the source sequence from the target sequence as well as to indicate
the end of the target sequence. To distinguish the source and target sequences
in the input text, we set the segment embedding of each character in the
symptom sequence $X$ to $A$, and the segment embedding of each character in
the prescribed constituent sequence $Y$ to $B$.

TCM prescription generation is not only related to symptoms but also closely
related to TCM knowledge. Therefore, we introduce some additional TCM
knowledge to guide the model generation. Firstly, a TCM-centric knowledge
graph is constructed, containing five categories of nodes: herbs, nature,
taste, channel tropism and effects, taking Ephedrae Herba as an example, as
shown in Figure 2.

Figure 2: Diagram of TCM Knowledge Graph

During the model training, the input sequence of the model needs to add
knowledge of TCM. Take the TCM prescription with the symptom of ”便血不止(bleeding
in the stool)” and the component of the prescription ”人参、白术、当归(Ginseng,
Largehead Atractylodes Rhizome and Chinese Angelica)” as an example, extract
the triples of each herb in the prescription component ”人参、白术、当归(Ginseng,
Largehead Atractylodes Rhizome and Chinese Angelica)”. The ternary groups of
the corresponding entities in the knowledge graph are added to the input
sequence, as shown in Figure 3.

Figure 3: Prescription composition adding knowledge example diagram

Equation 5 is correspondingly modified to:

$S=\left(\left[CLS\right],X,\left[SEP\right],Y,\left[SEP\right],\ Z\right)$
(6)

where $Z$ denotes the sequence consisting of the relevant knowledge of each
herb in $Y$.

The segment embedding corresponding to the joined TCM knowledge is set to $A$
to distinguish it from the prescription component. The token embedding,
position embedding and segment embedding of the spliced sequence $S$ are
summed and input to the RoBERTa model as the final input vector, and the
embedding of each character is calculated as shown in Figure 4.

Figure 4: Embedding calculation diagram

Map the input text sequence $S$ to the vector sequence $H^{0}$
correspondingly:

$h_{n}^{0}=E^{emb}\left(s_{n}\right)+E^{pos}\left(s_{n}\right)+E^{seg}\left(s_{n}\right),n\in\left[1,N\right]$
(7) $H^{0}=\left[h_{1}^{0},h_{2}^{0},\ldots,h_{N}^{0}\right]$ (8)

####  II-D3 Mask Prediction Layer

The $H^{0}$ in the input layer is fed into the network consisting of the
L-layer Transformer to obtain the contextual representation of the input text
at the lth layer:

$H^{l}\ =\ {\rm Transformer}_{l}\left(H^{l-1}\right),l\in\left[1,L\right]$ (9)

The original RoBERTa model utilizes contextual information from both the left
and right sides of a character when making predictions, which is not directly
applicable to sequence-to-sequence tasks. To enable the RoBERTa model for
sequence-to-sequence learning, we regulate the contextual information that
characters in a sequence can access through attention masks. Specifically, for
the TCM prescription generation task, characters in the symptom text component
need to have visibility to information from both the left and right sides. On
the other hand, for characters in the prescription component, only the
information of the first $j-1$ characters is accessible when predicting the
$j^{th}$ character.

During model training, with the introduction of additional TCM knowledge in
the input layer, the prescription component needs to access the knowledge of a
specific herb when making predictions, while excluding knowledge of other
herbs. This necessitates refining the attention mask matrix to precisely
control the accessible information.Taking the Chinese prescription with the
symptom component ”便血不止(blood in stool)” and the target prescription component
”人参、白术、当归(Ginseng, Largehead Atractylodes Rhizome and Chinese Angelica)” as an
example, the mask matrix for introducing the knowledge of Chinese medicine is
shown in Figure 5.

Figure 5: Mask matrix diagram

Specifically, the symptom part of the attention mask matrix is set to 0,
meaning that all symptom sequence characters are involved in the self-
attention calculation. The right part of the characters of the TCM
prescription component that needs to be predicted is set to $-\infty$, while
the mask of the Chinese medicine corresponding to the Chinese medicine
knowledge that needs to be predicted is set to 0, and the mask of the
knowledge corresponding to other Chinese medicines is set to $-\infty$,
indicating that the Chinese medicine can see its own related knowledge but not
other irrelevant knowledge. With such a setting, all the characters of the
symptom sequence can be made to perform self-attentive computation, and the
left part of the predicted characters of the prescription component and the
corresponding Chinese medicine knowledge part participate in the attention
computation, while its right part does not participate in the computation. The
attention head $A_{l}$ is calculated by the formula:

$\begin{cases}Q=H^{l-1}W_{l}^{Q}\\\ K=H^{l-1}W_{l}^{K}\\\
V=H^{l-1}W_{l}^{V}\end{cases}$ (10)
$M_{ij}=\left\\{\begin{matrix}0,&\text{allow to attend}\\\
-\infty,&\text{prevent from attending}\end{matrix}\right.$ (11)
$A_{l}=\text{softmax}\left(\frac{QK^{T}}{\sqrt{d_{k}}}+M\right)V_{l}$ (12)

####  II-D4 Output Layer

To output a set of herbs, the end of a multilayer Transformer is connected to
a linear classifier (such as a softmax classifier) to generate a probability
distribution over the word list to generate a prescription for the herbs:

$P\left(y\right)=softmax\left(h_{\left[Mask\right]}^{L}W^{T}\right)$ (13)

The predicted herb is appended to the input sequence to replace the
$\left[MASK\right]$ tokens to continue decoding until $\left[SEP\right]$
appears. The cross-entropy loss $L$ of the whole target sequence $Y$ is
defined as:

$L=-\frac{1}{|Y|}\sum_{i=1}^{|Y|}\log p\left(y_{i}|y_{<i},X\right)$ (14)

where $Y$ is the target sequence of the TCM prescription, $y_{i}$ is the ith
herb generated.

We view TCM prescription generation as a text generation problem in which the
choice of decoding method significantly affects the quality of the generated
text. During the model’s test inference phase, we refrain from using the
common Beam Search algorithm due to its drawbacks in causing duplication and
generating shorter texts, which are not ideal for TCM prescriptions. Instead,
we employ a combined approach using Temperature-controlled Stochastic Sampling
and Top-p Sampling as our decoding methods. This combination enhances the
RoKEPG model’s prescription generation by considering the diverse nature of
TCM prescriptions. As TCM prescriptions are composed of a varied set of herbs
without repetition and vary in length, our decoding strategy aims to improve
both the diversity and accuracy of the generated prescriptions, ensuring
better results for practical applications.

##  III Experiment

In this section, the RoKEPG model is systematically evaluated on a publicly
available TCM prescription dataset.

###  III-A Dataset

TCM Corpus：We collected books from more than twenty books related to the field
of TCM, such as Basic Theory of Traditional Chinese Medicine [17] and
Pharmacology of Traditional Chinese Medical Formulae [18]. Finally, after
cleaning and organising the data, we constructed a TCM corpus containing
340,000 lines and 13.5 million words.

TCM Knowledge Graph：In TCM theory, herbs have properties such as nature,
taste, channel tropism and effect. We construct a knowledge graph centred on
TCM, which contains five types of entities (TCM, nature, taste, channel
tropism, effect) with a total number of 1430 entities and four types of
relationships (TCM-has_nature$\rightarrow$nature, TCM-
has_taste$\rightarrow$taste, TCM-has_channel$\rightarrow$channel tropism, TCM-
has_effect$\rightarrow$ effect) with a total number of 9853 relationships.

TCM Dataset：The publicly available dataset contains 82,044 Traditional Chinese
Medicine (TCM) prescriptions, obtained from traditional prescription records
in TCM books, and underwent data cleaning and collation. We randomly divide
the whole TCM dataset into three parts: train set (73,839 entries, about 90%),
valid set (4,102 entries, about 5%) and test set (4,103 entries, about 5%).

###  III-B Evaluation Metrics and Experimental Setup

Evaluation Metrics：Adopt the practice with other papers of TCM prescription
generation by Seq2Seq [12], we used the same Precision, Recall and $F_{1}$
scores as metrics to evaluate the results of TCM prescription generation, with
Precision and Recall calculated as follows:

$Precision=\frac{G\left(X\right)\cap R\left(X\right)\mathrm{\ }}{\mathrm{\
G(X)\ }}$ (15) $Recall=\frac{G\left(X\right)\cap R\left(X\right)\mathrm{\
}}{\mathrm{\ R(X)\ }}$ (16)

Where $G(X)$ is a set of herbs generated by the model and $R(X)$ is the
prescription component in the reference prescription. In the TCM prescription
generation task, Recall is more critical than Precision [13]. The higher the
Recall, the more the generated set of herbs also belong to the real
prescription composition. However, in order to avoid generating too long
prescriptions, it is also necessary to consider Precision, so the prescription
generation results are evaluated comprehensively using the $F_{1}$ metric
because $F_{1}$ takes into account the Precision and Recall are considered
together.

$F_{1}=\frac{2\cdot\mathrm{\ Precision\ }\cdot\mathrm{\ Recall\ }}{\mathrm{\
Precision\ }+\mathrm{\ Recall\ }}$ (17)

Experimental Setup: Since almost all TCM prescription compositions in the TCM
dataset contain 20 or fewer herbs, we set the maximum number of generated
prescriptions containing herbs to 20. All experiments were performed in
PyTorch, and we set the maximum length of the input sequence to 256 and
selected the best experimental results based on the valid set.

In the pre-training phase of RoKEPG, the number of hidden layer neurons is set
to 768, the number of hidden layers is set to 12, the head of multi-head
attention is set to 12, the batch size is 32, the number of epochs is 15, the
maximum sequence length is set to 512, and the secondary pre-training is based
on the RoBERTa-WWM-Chinese model.

In the fine-tuning phase, the number of hidden layer neurons is set to 768,
the number of hidden layers is set to 6, the head of multi-head attention is
set to 12, the batch size is 16, the maximum sequence length is set to 256,
the learning rate is 1e-5, the number of epochs is 40, and the model is fine-
tuned based on the model obtained in the pre-training phase.

In the test decoding stage, the p-value of Top-p Sampling is set to 0.8, while
the temperature is set to 0.2.

###  III-C Comparative Experiments

To validate the RoKEPG model, we compared it with the following baseline
models on the same dataset:

  * •

Multi-Label[12]: The encoder part uses BiGRNN, and the generation part uses a
multi-label classification method to predict the model of Chinese medicine.

  * •

Basic Seq2Seq[12]: The encoder is a bidirectional Gated Recurrent Unit (GRU),
and the decoder is a unidirectional GRU, forming a sequence-to-sequence model.

  * •

S2SA-CS[12]: A Seq2Seq model with attention mechanism, coverage mechanism and
soft cross-entropy.

  * •

TCM-Translator[6]: The encoder is a Transformer without positional embedding,
and the decoder is an LSTM, constituting the Seq2Seq model.

  * •

Herb-Know[13]: A sequence-to-sequence model based on pointer networks and
improved model generation by introducing knowledge of the effects of herbs.

  * •

BERT-UniLM[15]: A single BERT generates TCM prescriptions by attention mask
matrix.

  * •

RoBERTa-UniLM[7]: A single RoBERTa is implemented to generate TCM
prescriptions through an attention mask matrix similar to BERT-UniLM.

  * •

GPT-2[20]: An autoregressive model consisting of the decoder part of a
multilayer one-way Transformer.

  * •

BART[21]: A pre-trained language model, built according to Transformer’s
general framework, implements sequence-to-sequence with a bidirectional
encoder and a left-to-right autoregressive decoder.

  * •

T5[22]: A pre-trained language model that retains most of the original
Transformer framework transforms all tasks in natural language processing into
text-to-text tasks.

Through several experiments, the learning rates of GPT-2, BART and T5 are all
set to 1e-5, the batch sizes are all set to 8, and the number of epochs are
all set to 30. the main parameters of BERT-UniLM and RoBERTa-UniLM are the
same as those of RoKEPG, while Multi-Label, Basic Seq2Seq and S2SA-CS keep the
same parameter settings as paper [12], TCM-Translator keeps the same parameter
settings as the original paper [6], and Herb-Know likewise keeps the same
parameter settings as the original paper [13].

####  III-C1 TCM Prescription Generation Comparison Experiment

The experimental results of TCM prescription generation based on the publicly
available TCM dataset are shown in Table I.

TABLE I: TCM prescription generation comparison experiment Model | Precision | Recall | $F_{1}$
---|---|---|---
Multi-Label | 10.83 | 29.72 | 15.87
Basic Seq2Seq | 26.03 | 13.52 | 17.80
S2SA-CS | 29.57 | 17.30 | 21.83
TCM-Translator | 25.72 | 12.13 | 16.48
Herb-Know | 28.35 | 20.56 | 23.83
GPT-2 | 15.47 | 13.24 | 14.27
BART | 20.56 | 14.84 | 17.24
T5 | 19.34 | 16.98 | 18.08
BERT-UniLM | 28.82 | 18.85 | 22.79
RoBERTa-UniLM | 27.76 | 19.46 | 22.88
RoKEPG | 23.43 | 29.01 | 25.92

As can be seen from the experimental results in Table I, from the overall
experimental results, the RoKEPG model significantly outperforms the other
baseline models in terms of recall and $F_{1}$ scores, and the RoKEPG model
outperforms the most effective baseline Herb-Know model by about 2% in the
$F_{1}$ metric.

The Multi-Label model has a high Recall but low Precision and $F_{1}$ values
due to the long prescription length generated by the Multi-Label model. The
results generated based on BERT are slightly lower than those generated based
on RoBERTa, because the RoBERTa model is a tuned and enhanced version of BERT
and does achieve better results than BERT in several NLP tasks. Although
GPT-2, BART and T5 models are more widely used in text generation tasks, they
do not work well in TCM prescription generation tasks.

####  III-C2 TCM Knowledge Introduction Experiment

For the TCM knowledge introduction module, we explored two different
introduction methods: random introduction and all introduction. The random
introduction refers to finding all the triads associated with a certain TCM
from the TCM knowledge graph and randomly selecting one of the triads as the
TCM knowledge to be added to the input sequence. All introduction refers to
adding all the triads associated with herb into the input sequence. We
validate the effect of these two different knowledge introduction methods
based on three models, BERT-UniLM, RoBERTa-UniLM, and RoKEPG together, and the
experimental results are shown in Table II.

TABLE II: Different ways of introducing knowledge Knowledge Methods | Model | Precision | Recall | $F_{1}$
---|---|---|---|---
Random Introduction | BERT-UniLM | 24.01 | 22.82 | 23.40
RoBERTa-UniLM | 20.35 | 28.01 | 23.57
RoKEPG | 21.35 | 32.25 | 25.69
All Introduction | BERT-UniLM | 22.72 | 24.13 | 23.40
RoBERTa-UniLM | 20.75 | 27.32 | 23.58
RoKEPG | 23.43 | 29.01 | 25.92

The experimental results in Table II show that all three models are better
with the all introduction approach than the random introduction approach,
probably because there is a repeated introduction of the same relevant
knowledge of Chinese medicine if the random introduction is performed in the
training phase. Moreover, the random introduction only introduces one
knowledge at a time, and the model receives less knowledge than the all
introduction approach, leading to decreased model performance.

####  III-C3 Different Decoding Strategies

In the decoding generation stage of the model test, different decoding
strategies will affect the results differently. We combine a mixture of two
decoding strategies, Temperature-controlled Stochastic Sampling and Top-p
Sampling, for decoding. The results of comparing different decoding methods on
the performance of RoKEPG model are shown in Table III, and the contents in
parentheses in the table are the parameter settings when the results of
corresponding decoding methods are optimal.

TABLE III: Different decoding methods Decoding Methods | Precision | Recall | $F_{1}$
---|---|---|---
Beam Search(beam-size=5) | 23.09 | 27.97 | 25.29
Top-K (k=3) | 21.26 | 29.50 | 24.71
Top-P (p=0.9) | 23.64 | 28.65 | 25.90
Temperature(temperature=0.2) | 23.00 | 28.91 | 25.62
RoKEPG (mixed decoding) | 23.43 | 29.01 | 25.92

The experimental results in Table III show that the mixture decoding method of
RoKEPG model outperforms the other decoding methods, and the effect of
Temperature-controlled Stochastic Sampling and Top-P Sampling alone is
slightly inferior to that of the mixture decoding of RoKEPG model. The common
Beam Search decoding method is biased to generate shorter prescriptions, and
there are many duplicates in the results compared with other decoding methods.

####  III-C4 Ablation Experiment

In order to explore the effects of the pre-training phase and the introduction
of TCM knowledge in the RoKEPG model, we designed two sets of ablation
experiments. We first validate the effectiveness of the pre-training phase
based on the RoKEPG model and the RoBERTa-UNILM model together, and the
experimental results are shown in Table IV. w/o denotes Without, and the pre-
training module is denoted by pre-trained.

TABLE IV: Ablation experiments with pre-training module Model | Precision | Recall | $F_{1}$
---|---|---|---
RoBERTa-UNILM | 24.18 | 25.75 | 24.94
RoBERTa-UNILM-w/o pre-trained | 27.76 | 19.46 | 22.88
RoKEPG | 23.43 | 29.01 | 25.92
RoKEPG-w/o pre-trained | 20.75 | 27.32 | 23.58

As shown from the experimental results in Table IV, the pre-trained model
significantly improves effectiveness, with more than 2% improvement in the
$F_{1}$ metric over the non-pre-trained model.

The results of the ablation experiment for the effectiveness of introducing
TCM knowledge into the RoKEPG model are shown in Table V, and the introduction
of TCM knowledge is indicated by herb-know.

TABLE V: Ablation experiments for knowledge introduction Model | Precision | Recall | $F_{1}$
---|---|---|---
RoKEPG | 23.43 | 29.01 | 25.92
RoKEPG-w/o herb-know | 24.18 | 25.75 | 24.94

From the experimental results in Table V, it can be seen that the introduction
of TCM knowledge brings a 0.98% improvement in the $F_{1}$ index compared with
the non-introduction of TCM knowledge, indicating that the introduction of TCM
knowledge does help in TCM prescription generation, and the model enhanced by
TCM knowledge can generate TCM prescriptions better.

####  III-C5 Case Study

Further, we randomly selected two TCM prescriptions from the test set for case
study, as shown in Table VI. The RoKEPG model generates herbs mostly present
in the reference prescription, including some not originally listed.

TABLE VI: Case study Cases |  Case 1 |  Case 2
---|---|---
Symptom |  小儿乳癖、积聚。按之苦痛，肌肤渐瘦，面色青黄；小儿阴阳气不顺，虚痞胀满，呕逆腹痛，成癥瘕痞结。
(Lactation and accumulation in children. Bitter and painful when pressed, skin gradually thinned, face green and yellow; Yin and Yang Qi disorder in children, deficiency and fullness, vomiting and rebellious abdominal pain, forming abdominal masses and abnormalities.) |  产后赤白痢，脐下咻痛。
(Postpartum red and white dysentery with painful swoosh under the umbilicus.)
Reference prescription |  麦芽，灵脂，香附，陈皮，神曲，青皮，莱菔，厚朴，半夏，槟榔，枳实，砂仁，三棱。
(Hordei Fructus Germinatus, Trogopterus Dung, Nutgrass Galingale Rhizome, Citrus reticulata Blanco, Medicated Leaven, Citri Reticulatae Pericarpium Viride, Radish root, Magnoliae Officmalis Cortex, Pinellia ternate, Betelnutpalm Seed, Immature Bitter Orange, Villous Amomum Fruit, Sparganii Rhizoma.) |  人参，白术，麦冬，当归，熟地，天麻，防风，荆芥，陈皮，甘草，生姜。
(Ginseng, Largehead Atractylodes Rhizome, Radix Ophiopogonis, Chinese
Angelica, Rehmanniae Radix Praeparata, Tall Gastrodia Tuber, Divaricate
Saposhniovia Root, Schizonepetae Herba, Citrus reticulata Blanco, Licorice,
Fresh Ginger.)
RoKEPG |  枳实，陈皮，厚朴，香附，山楂，麦芽，神曲，砂仁，木香，槟榔， 青皮，甘草。
(Immature Bitter Orange, Citrus reticulata Blanco, Magnoliae Officmalis Cortex, Nutgrass Galingale Rhizome, Crataegi Fructus, Hordei Fructus Germinatus, Medicated Leaven, Villous Amomum Fruit, Common Vladimiria Root, Betelnutpalm Seed, Citri Reticulatae Pericarpium Viride, Licorice.) |  人参，白术，当归，川芎，白芍，熟地，麦冬，陈皮，甘草，生地黄。
(Ginseng, Largehead Atractylodes Rhizome, Chinese Angelica, Sichuan lovase
rhizome, White Peony Root, Rehmanniae Radix Praeparata, Radix Ophiopogonis,
Citrus reticulata Blanco, Licorice, Rehmanniae Radix.)
Result |  Precision=66.67 Recall=61.54 $F_{1}$=64.0 |  Precision=70.0 Recall=63.64 $F_{1}$=66.67

In the first case, the RoKEPG model generated several herbs that were not in
the reference prescription according to the symptoms, including Crataegi
Fructus for treating gastric and epigastric fullness, diarrhea and abdominal
pain [23]; the generated Common Vladimiria Root is mainly used for symptoms
such as epigastric distension and inappetence [24]; and the generated Citri
Reticulatae Pericarpium Viride could be used for treating lactation, canker
sores, and epigastric distension and pain [25]. The symptoms these herbs can
treat are similar to those of the first case, indicating that the RoKEPG model
is not randomly generated herbs, as all the herbs generated by the model are
highly correlated with the prescribed symptoms.

In the second case, the RoKEPG model generated Sichuan lovase rhizome and
White Peony Root, which is one of the well-known pairs of medicines that can
invigorate blood and nourish it and soften the liver to relieve pain [26]. Red
and white dysentery (赤白痢) refers to the presence of pus and blood in the stool
often accompanied by abdominal pain. The medicine pair composed of Sichuan
lovase rhizome and White Peony Root can relieve the pain. The model generated
Rehmanniae Radix, which also treats blood in the stool and stops bleeding
[27]. Based on the above analysis, the RoKEPG model can learn the combination
and pairs of Chinese medicines.

##  IV Conclusion

TCM is an essential part of traditional medicine in China, and it is crucial
to study the prescription generation of TCM and to uncover the complex
relationship between herbs and symptoms. For the TCM prescription generation
task, we propose a RoBERTa and Knowledge Enhancement model for Prescription
Generation of Traditional Chinese Medicine model (RoKEPG). RoKEPG performes a
secondary pre-training of the RoBERTa model based on a TCM domain corpus which
we constructed, and introduces four types of knowledge of TCM through the
attention mask matrix to help the model to generate prescriptions better. The
experimental results based on the publicly available TCM dataset show that
RoKEPG improves the $F_{1}$ index by about 2% over the best baseline model,
which can effectively assist TCM practitioners in prescribing. Since we model
the TCM prescription generation as a text generation problem, there is the
problem of exposure bias in the text generation task, and we will also
consider adding contrastive learning to alleviate this problem in the follow-
up to improve the quality of the generated TCM prescriptions.

## Acknowledgement

This work was supported by National Key R&D Program of China (2019YFC1711000)
and Collaborative Innovation Center of Novel Software Technology and
Industrialization.

## References

  * [1] A. Helfenstein and P. Tammela, “Analyzing user-generated online content for drug discovery: development and use of medcrawler,” _Bioinformatics_ , vol. 33, no. 8, pp. 1205–1209, 2017.
  * [2] J. Qiu, “Traditional medicine: a culture in the balance,” _Nature_ , vol. 448, no. 7150, pp. 126–129, 2007.
  * [3] L. Yao, Y. Zhang, B. Wei, W. Zhang, and Z. Jin, “A topic modeling approach for traditional chinese medicine prescriptions,” _IEEE Transactions on Knowledge and Data Engineering_ , vol. 30, no. 6, pp. 1007–1021, 2018.
  * [4] Y. Zhang, W. Ji, H. Wang, X. Wang, and J. Chen, “Mc-elda: Towards pathogenesis analysis in traditional chinese medicine by multi-content embedding lda,” in _Advances in Knowledge Discovery and Data Mining: 23rd Pacific-Asia Conference, PAKDD 2019, Macau, China, April 14-17, 2019, Proceedings, Part I 23_. Springer, 2019, pp. 489–500.
  * [5] Z. Liu, Z. Zheng, X. Guo, L. Qi, J. Gui, D. Fu, Q. Yao, and L. Jin, “Attentiveherb: a novel method for traditional medicine prescription generation,” _IEEE Access_ , vol. 7, pp. 139 069–139 085, 2019.
  * [6] Z. Wang, J. Poon, and S. Poon, “Tcm translator: A sequence generation approach for prescribing herbal medicines,” in _2019 IEEE International Conference on Bioinformatics and Biomedicine (BIBM)_. IEEE, 2019, pp. 2474–2480.
  * [7] Y. Liu, M. Ott, N. Goyal, J. Du, M. Joshi, D. Chen, O. Levy, M. Lewis, L. Zettlemoyer, and V. Stoyanov, “Roberta: A robustly optimized bert pretraining approach,” _arXiv preprint arXiv:1907.11692_ , 2019.
  * [8] M. R. Boutell, J. Luo, X. Shen, and C. M. Brown, “Learning multi-label scene classification,” _Pattern recognition_ , vol. 37, no. 9, pp. 1757–1771, 2004\.
  * [9] M.-L. Zhang and Z.-H. Zhou, “Ml-knn: A lazy learning approach to multi-label learning,” _Pattern recognition_ , vol. 40, no. 7, pp. 2038–2048, 2007.
  * [10] A. Elisseeff and J. Weston, “A kernel method for multi-labelled classification,” _Advances in neural information processing systems_ , vol. 14, 2001.
  * [11] J. Read, B. Pfahringer, G. Holmes, and E. Frank, “Classifier chains for multi-label classification,” _Machine learning_ , vol. 85, pp. 333–359, 2011\.
  * [12] W. Li, Z. Yang, and X. Sun, “Exploration on generating traditional chinese medicine prescription from symptoms with an end-to-end method,” _arXiv preprint arXiv:1801.09030_ , 2018.
  * [13] C. Li, D. Liu, K. Yang, X. Huang, and J. Lv, “Herb-know: knowledge enhanced prescription generation for traditional chinese medicine,” in _2020 IEEE International Conference on Bioinformatics and Biomedicine (BIBM)_. IEEE, 2020, pp. 1560–1567.
  * [14] Z. Liu, C. Luo, D. Fu, J. Gui, Z. Zheng, L. Qi, and H. Guo, “A novel transfer learning model for traditional herbal medicine prescription generation from unstructured resources and knowledge,” _Artificial Intelligence in Medicine_ , vol. 124, p. 102232, 2022.
  * [15] N. Wang, “Research on clinical decision support for chinese medicine based on electronic medical records,” Ph.D. dissertation, University of Science and Technology of China, 2021.
  * [16] Z. Yin, Y. Wu, and Y. Zhang, “Hgcl: Heterogeneous graph contrastive learning for traditional chinese medicine prescription generation,” in _International Conference on Health Information Science_. Springer, 2022, pp. 88–99.
  * [17] G. Sun and H. Zheng, _Basic theory of Traditional Chinese Medicine_. China Traditional Chinese Medicine Publishing House, 2017.
  * [18] Z. Deng, _Pharmacology of traditional Chinese medical formulae_. China Traditional Chinese Medicine Publishing House, 2019.
  * [19] L. Dong, N. Yang, W. Wang, F. Wei, X. Liu, Y. Wang, J. Gao, M. Zhou, and H.-W. Hon, “Unified language model pre-training for natural language understanding and generation,” _Advances in neural information processing systems_ , vol. 32, 2019.
  * [20] A. Radford and J. Wu, “Rewon child, david luan, dario amodei, and ilya sutskever. 2019,” _Language models are unsupervised multitask learners. OpenAI blog_ , vol. 1, no. 8, p. 9, 2019.
  * [21] M. Lewis, Y. Liu, N. Goyal, M. Ghazvininejad, A. Mohamed, O. Levy, V. Stoyanov, and L. Zettlemoyer, “Bart: Denoising sequence-to-sequence pre-training for natural language generation, translation, and comprehension,” _arXiv preprint arXiv:1910.13461_ , 2019.
  * [22] C. Raffel, N. Shazeer, A. Roberts, K. Lee, S. Narang, M. Matena, Y. Zhou, W. Li, and P. J. Liu, “Exploring the limits of transfer learning with a unified text-to-text transformer,” _The Journal of Machine Learning Research_ , vol. 21, no. 1, pp. 5485–5551, 2020.
  * [23] J. Dong, J. Chen, S. Gong, J. Xu, X. Xu, and G. Zhang, “Chemical composition, pharmacological effects and quality marker (q-marker) prediction of crataegi fructus,” _Chinese herbal medicine_ , vol. 52, no. 9, pp. 2801–2818, 2021\.
  * [24] J. Zheng, M. Shang, J. Wang, G. Dai, J. Song, and B. Duan, “Chemical composition, pharmacological effects, clinical application research progress and quality marker prediction of aucklandiae radix,” _Chinese herbal medicine_ , vol. 53, no. 13, pp. 4198–4213, 2022.
  * [25] R. Xu, F. Zhong, and D. Wu, “The chinese herbal medicine citrus reticulata blanco,” _Chinese medicinal materials_ , vol. 36, no. 6, pp. 1018–1023, 2013\.
  * [26] C. Wu, Z. Liu, H. Qin, W. Zhu, and B. Wang, “Exploring the dispensing rules and application characteristics of chuanxiong in the chinese pharmacopoeia based on data mining,” _Journal of Southwest University for Nationalities (Natural Science Edition)_ , vol. 48, no. 2, pp. 172–180, 3 2022\.
  * [27] Y. Ye, M. Tang, Q. Zhang, J. Su, F. Chu, S. Zhou, S. Xu, H. Lei, and Q. Li, “Advances in the study of the substance basis of the medicinal effects of raw rehmanniae radix,” in _Proceedings of the Eighth Annual Conference of Chinese Medicine Chemistry Branch, Chinese Society of Traditional Chinese Medicine_ , 2013, pp. 148–155.