# Medical SANSformers: Training self-supervised transformers without attention
for Electronic Medical Records

Yogesh Kumar 1, 2<EMAIL_ADDRESS>Alexander Ilin 1
<EMAIL_ADDRESS>Henri Salo 2<EMAIL_ADDRESS>Sangita Kulathinal 2, 3
<EMAIL_ADDRESS>Maarit K. Leinonen 2<EMAIL_ADDRESS>Pekka Marttinen 1, 2<EMAIL_ADDRESS>
1Department of Computer Science, Aalto University, Finland
2Information Services Department, Finnish Institute for Health and Welfare,
Finland
3Department of Mathematics and Statistics, University of Helsinki, Finland

###### Abstract

We leverage deep sequential models to tackle the problem of predicting
healthcare utilization for patients, which could help governments to better
allocate resources for future healthcare use. Specifically, we study the
problem of divergent subgroups, wherein the outcome distribution in a smaller
subset of the population considerably deviates from that of the general
population. The traditional approach for building specialized models for
divergent subgroups could be problematic if the size of the subgroup is very
small (for example, rare diseases). To address this challenge, we first
develop a novel attention-free sequential model, SANSformers, instilled with
inductive biases suited for modeling clinical codes in electronic medical
records. We then design a task-specific self-supervision objective and
demonstrate its effectiveness, particularly in scarce data settings, by pre-
training each model on the entire health registry (with close to one million
patients) before fine-tuning for downstream tasks on the divergent subgroups.
We compare the novel SANSformer architecture with the LSTM and Transformer
models using two data sources and a multi-task learning objective that aids
healthcare utilization prediction. Empirically, the attention-free SANSformer
models perform consistently well across experiments, outperforming the
baselines in most cases by at least $\sim 10$%. Furthermore, the self-
supervised pre-training boosts performance significantly throughout, for
example by over $\sim 50$% (and as high as $800$%) on $R^{2}$ score when
predicting the number of hospital visits.

Keywords: Deep Learning, Electronic Health Records, Transfer Learning, Risk
Adjustment, Healthcare

##  1 Introduction

We develop and rigorously evaluate a new sequential architecture for
Electronic Medical Records (EMR) which is devoid of any recurrence,
convolution or self-attention mechanisms and is adapted from Liu et al. (2021)
and Tolstikhin et al. (2021). Deep learning methods have been successfully
applied in the past to predict the cause for the next visit (Lipton et al.,
2015; Choi et al., 2016a) and time-duration to the next visit (Choi et al.,
2016b; Harutyunyan et al., 2019) using EMR. A common approach in many of these
works has been to treat the patient history from EMR as a discrete
longitudinal sequence and thus naturally, methods from Natural Language
Processing (NLP) were adapted to the EMR domain. Transformer architectures
(Vaswani et al., 2017) that have been pre-trained in an unsupervised fashion
on a larger text corpus (Devlin et al., 2018; Raffel et al., 2019; Radford et
al., 2019) are dominating the state-of-the-art in several of the NLP tasks. Li
et al. (2020) and Rasmy et al. (2021), among others, have demonstrated that we
can obtain similar superior results on EMR by tapping into the benefits of
self-supervised pre-training while using the same architecture and training
regime as the NLP tasks. However, as demonstrated by Choi et al. (2016a) and
Choi et al. (2019), a lot can be gained by making changes to the model design
to tailor it for the EMR.

Unlike natural language, sequences of clinical codes from EMR are multi-
dimensional since there can be multiple codes associated with each visit and
consecutive visits could be even several years apart. For these reasons,
recorded codes within consecutive visits need not be highly correlated. We
thus design a model instilled with inductive biases that suit the EMR data
structure: axial mixing to tackle intra-visit multi-dimensionality and
$\Delta\tau$ embeddings to encode the relative time difference between each
visit. We also adopt a self-supervised regime that utilize the EMR structure,
namely generative pre-training of future given the patient history as input.
We also introduce attention-free SANSformers (a deliberately crafted acronym
for Sans ANy Self-Attention Transformers) which performs surprisingly better
than transformer models, thus questioning the need for the memory and time
complexity intensive self-attention mechanism in Transformer architectures for
EMR.

To evaluate our model, we designed elaborate experiments inspired by real-
world scenarios and applications on the EMR data. Our primary motivating
application was to build a risk adjustment model which can predict the
utilization of healthcare in the future for an individual based on one’s
disease history and other relevant variables. Such models to allocate
resources to healthcare providers according to the expected demand are already
in use in many countries (McGuire and van Kleef, 2018). A major challenge is
to predict accurately for patients that belong to divergent subgroups, for
example patients diagnosed with a certain disease. That is because for severe
or chronic illnesses the trajectory of the patient history may deviate
considerably from the previous pattern of the patient himself/herself and also
from that of other patients. To address this challenge, models for specific
subgroups have been designed in the health economics application field
(Shrestha et al., 2018; Ellis et al., 2018). Indeed, when we categorize
patients into different subgroups based on their first record of certain
diagnosis, we observe that the histogram of the expected number of visits
varies considerably by subgroups. Figure 1 (_top_) compares the histogram of
the number of visits for patients in bipolar disorder subgroup with the rest
of the population.

A trivial solution to this problem would be to train a single model on the
entire population and predict the utilization for each subgroup. However, this
model would not be as accurate as a specialized model trained for each
subgroup. To demonstrate this we train an LSTM model on the general population
(single model) to predict the utilization for patients in a test set, a subset
of the bipolar disorder subgroup, and compare it with a specialized model
trained only with the patients having bipolar disorder. The specialized model
achieves an $R^{2}$ score of 0.16 while the single model scores a mere 0.09 on
the test set. Further, Figure 1 (_down_) shows a box plot of the predicted
number of visits for both the specialized and single LSTM models. The
predictions were uniformly binned into eight categories and their
interquantile range was plotted against the true number of visits. Since the
histogram of visits for the general population is inflated around zero visits,
we can see that the single model mimics this distribution by predicting high
values for low risk patients as compared to the specialized model. This over-
prediction could lead to severely misplaced resource allocation and, thus,
justifying specialized models. However, training neural networks from scratch
usually requires a large amount of data and developing specialized models can
be difficult for rare diseases. This limitation can been seen in Figure 1
(_down_) where the specialized model fails to predict low values and shows
larger variance over the predictions. The data efficiency of model can be
vastly improved by fine-tuning a self-supervised model that was trained on a
larger corpus.

Figure 1: The problem of predicting in a divergent subgroup. _Top_ Histograms
of the number of visits to the hospital for the general population and for
people diagnosed with bipolar disorder. The observed bi-modality in the
histogram of counts is due to the introduction of topcap, where we counter the
long-tail problem by clipping the max count value at 36 visits; _Bottom_ Box
plot comparing the predicted number of visits for patients with bipolar
disorder of an LSTM model trained only on this subgroup (Specialized model)
against an LSTM model trained on the entire population (Single model). The
patients are divided into eight categories based on their predicted values,
where the categories are obtained by dividing the whole range into eight bins
of equal width: $[0,4],[4,8],\ldots,[28,32]$, shown on the x-axis. The
specialized model has a higher $R^{2}$ score of 0.16 compared to 0.09 of the
single model on the test set, but the specialized model has a large variance
and doesn’t predict the low values properly.

In order to strengthen the evaluation rigor and to enable reproducibility of
our proposed methods and results, we also evaluate our model on the publicly
available The Medical Information Mart for Intensive Care (MIMIC)-IV (Johnson
et al., 2020) dataset. With MIMIC, we predict the probability of inpatient
mortality, i.e., the mortality of a patient who has been admitted to the
hospital. We also evaluate on a sequential task of predicting the length-of-
stay of the next visit in an autoregressive fashion.

Our main contributions are summarized as follows

  * •

Introduce an attention-free sequential model for modeling EMR data and
supplement it with architectural changes such as axial decomposition and
$\Delta\tau$ embeddings, to address the specific challenges of the EMR data.

  * •

Rigorously compare the performance of the proposed architecture against strong
baselines on a multi-task learning framework, consisting of classification and
regression tasks, using multiple data sources, and demonstrate it is more data
efficient and accurate as compared to the transformer models.

  * •

Demonstrate that self-supervised pre-training on a larger general population
greatly improves predictions on smaller divergent subgroups irrespective of
the model architecture, highlighting significant potential in the important
healthcare resources allocation application.

##  2 Related Work

Deep Learning on EMR Lipton et al. (2015) was one of the first works to apply
deep learning for EMR phenotyping. They used sequential real-valued
measurements of 13 different vital measurements to predict one of 128
diagnoses using an LSTM. Choi et al. (2016b) was another pivotal work in this
direction, where they used a Gated Recurrent Unit (GRU) model in a multi-task
setting for phenotyping. They also demonstrated the importance of transfer
learning when working with smaller datasets. Other works in this domain
include Choi et al. (2016a), which developed a bidirectional attention based
model focused interpretability; Harutyunyan et al. (2019) created a benchmark
framework for evaluating EMR models; Choi et al. (2019) developed a
convolutional transformer model that has an added inductive bias to model the
interactions between different codes occurring at one visit while Kumar et al.
(2020) used a 1x1 convolutional network to model this interaction; Li et al.
(2020); Rasmy et al. (2021) applied pre-trained BERT transformers to EMR data
and demonstrated superior results compared to traditional recurrent neural
networks (RNN) or convolutional networks.

Self-supervised pre-training Howard and Ruder (2018); Peters et al. (2018)
were one of the earliest works to demonstrate that unsupervised language model
pre-training can be leveraged to vastly improve the model performance for
downstream tasks. However, the greatest improvements were achieved by pre-
training the transformer architectures (Devlin et al., 2018; Radford et al.,
2019; Brown et al., 2020; Raffel et al., 2019).

Attention-less MLP Models There have been some recent works which explored
removing the attention mechanisms from the transformer models, especially for
the computer vision tasks. Tolstikhin et al. (2021) replaced the self-
attention mechanism with a "mixing" mechanism. Liu et al. (2021) simplified
the mixer further by replacing the mixer with spatial gating mechanisms. Other
concurrent works with similar objectives include Touvron et al. (2021); Lee-
Thorp et al. (2021); Melas-Kyriazi (2021). These papers questioned the need
for the complex and expensive self-attention mechanisms, extending this to EMR
domain has been one of the core motivations for our work.

##  3 Setup

###  3.1 Cohorts

We used two data sources for our experiments: a confidential dataset (PUMMEL)
that was sourced from the Care Register for Health Care and Register of
Primary Health Care visits maintained by the Finnish Institute for Health and
Welfare (THL) and a smaller publicly available MIMIC-IV dataset. The datasets
were analysed separately, to demonstrate the generality of the findings in
multiple independent settings. Furthermore, the fact that the MIMIC-IV data is
publicly available allows the readers to run our method and replicate our
findings on that dataset.

#### Pummel

The Pummel dataset consists of the EMR of every Finnish citizen aged 65 or
above, who has interacted with the primary or secondary healthcare services.
The data, which has been pseudonymized in order to protect the privacy of the
patients, consists of both primary and secondary care visits to the healthcare
services between the years 2012 and 2018. These visits can include any form of
contact with the healthcare facilities such as booked appointments, over-the-
phone consultations, nursing home visits, or admission to the hospital.

Table 1: Basic statistics of the two EMR datasets. | Pummel | MIMIC-IV
---|---|---
# of unique patients | 1,050,512 | 256,878
# of visits | 60,896,305 | 523,740
Avg. # of visits per patient | 57.94 | 2.04
Max # of visits | 1,443 | 238
Avg # of codes per visit | 8.46 | 18.84
Max # of codes per visit | 164 | 110
Token Vocabulary Size | 5237 | 5019

The EMR data in its original raw tabular form holds the record of variables
related to each patient visit to the healthcare services. These variables
include the medical diagnosis codes such as the International Classification
of Diseases (ICD-10) and International Classification of Primary Care
(ICPC-2), the codes for surgical procedures performed (Lehtonen et al., ), the
patient demographics like age and gender and the specialty of visit. Since we
are modeling the patient history sequentially, we transformed the tabular data
by grouping the rows based on both PatientID and VisitID (which are unique
numeric codes assigned to each patient and healthcare contact, respectively).
Diagnostic codes from one healthcare contact, i.e. the intra-visit codes,
appeared in the same order in which the health care personnel have reported
them into the register. The resulting dataset consisted of 1,050,512 patient
sequences across the years 2012 – 2018 with an average of $57.94$ visits per
patient.

#### MIMIC

The MIMIC-IV _v1.0_ dataset consists of EMR from hospital and Intensive Care
Unit (ICU) admissions of roughly 250,000 patients admitted to Beth Israel
Deaconess Medical Center (BIDMC). The dataset is extremely rich in its
contents and contains information about the patient vitals, doctor notes,
diagnoses, procedure and medication codes, discharge summaries, and more from
both hospital and ICU admissions. For our study, we only used the hospital
admissions, the primary focus of this investigation. Specifically, we
extracted patient information from the patients, admissions, diagnoses_icd,
procedures_icd, drgcodes and services tables. The information from each visit
was grouped by the admission identifier hadm_id and all the visits were
grouped by the patient identifier subject_id. The resulting dataset consisted
of 256,878 patients with an average of 2.04 visits per patient.

Table 1 lists some basic statistics of both datasets.

###  3.2 Tasks

To perform a rigorous inquiry into the performance of the proposed
architectures, we devised different tasks on both data sources. In general, we
frame our prediction tasks together as a multi-task learning problem where a
single neural network architecture is trained to optimize the losses
corresponding to several tasks. This reflects the real-world scenario where
training specialized models for each task would not be practical due to time
and resource constraints. Moreover, neural nets trained in a multi-task
learning framework have shown to have other benefits such as better
generalization and better quality of learned representation (Ruder, 2017).

On Pummel, our goal was to forecast the burden on healthcare utilization for
each patient belonging to a particular subgroup based on their diagnosed
diseases. We formulate two tasks which together act as a proxy for the actual
healthcare demand in lieu of estimating the monetary demand directly (which
would vary depending on the country and inflation rates). Specifically, we use
the one-year EMR history to predict the following variables for each patient
for the next year:

  * •

Task 1: The number of physical visits to the healthcare centers
($y_{\text{count}}$)

  * •

Task 2: The counts of the physical visits due to six specific disease
categories ($y_{\text{diag}}$)

Both of these tasks are strongly related to healthcare cost. For Task 2, we
handpicked six disease categories that have been identified to consume a lot
of healthcare resources in Finland, namely cancer (ICD-10 codes starting with
C and some D), endocrine and metabolic diseases (E), diseases related to
nervous systems (G), diseases of the circulatory system (I), diseases of the
respiratory system (J) and diseases of the digestive system (K) (Häkkinen et
al., ). For both tasks we model the outcome as a Poisson distribution,
$Pr(X=k)=\lambda^{k}e^{-k}/k!,\quad k=1,2,\dots$ and use the neural network to
estimate the expected event rate $\lambda$ for each patient. More precisely,
if we assume that $f(.)$ represents a neural network, then we estimate
$\log\lambda_{i}=f(x_{i})$ for each patient $i$. During training, we minimize
the negative log-likelihood of the Poisson distribution.

With MIMIC-IV dataset, our goal was to test the versatility of the model
architecture, hence we devised two tasks that are significantly different from
the Pummel tasks: binary classification and sequence tagging. Specifically, we
predict

  * •

Task 3: The probability of inpatient mortality ($y_{\text{death}}$), and

  * •

Task 4: The estimated length of stay (LoS) of the next visit
($y_{\text{los}}$)

For Task 3 we utilize the hospital_expire_flag feature in the MIMIC-IV
admissions table, which is used to denote whether the patient died during the
current hospitalization episode. Since it would be much more beneficial to
predict the probability of death well in advance, we remove the last five
visits from the patient history for each patient. Thus effectively, Task 3
predicts the probability of death after the next 5 visits. Task 4 nicely ties
in with the objective of the Pummel experiments—forecasting healthcare
utilization—because longer LoS implies higher utilization. LoS was calculated
as the difference between discharge date and the admission date for each visit
and our objective is to predict the LoS of the next visit autoregressively. We
used binary cross entropy loss and Huber (or smooth L1) loss for tasks 3 and
4, respectively.

Figure 2: EMR Modeling Pipeline._Left_ We can conceptually modularize the
pipeline into Data Ingestor (A), Visit Summarizer (B), Sequence Transducer (C)
and Task/Output Layer (D) such that the data is being transformed within each
of these modules. This aids in weight reuse since each of the modules can be
pre-trained separately. _Center_ The input data in tabular form is transformed
into various shapes by each module. Here $T$ is the total number of visits,
$V$ is the number of intra-visit codes and $E$ is the embedding size for each
patient.

###  3.3 EMR modeling pipeline

Since the content, structure and size vary significantly between different EMR
databases, we conceptually modularize our data and modeling pipeline into
various self-contained modules (Figure 2). Each module applies some functional
transformations to the input data. The transformed data can then be passed
onto the next module.

#### Data Ingestor

The first module, data ingestor, takes in the raw tabular EMR data and
converts it to a form that machine learning models understand. In our case, we
model the patient history as a sequence and group the visit records $v_{i,t}$
for each patient $i$ at each time-step $t=1,2,\dots T$ where $T$ is the total
number of visits for patient $i$. Since the frequency of visits varies between
patients, using a fixed time frame of reference for all patients would lead to
a highly sparse patient history. Thus we use a relative time axis where $t=1$
denotes the first recorded visit for each patient. Since we use $t$ to denote
the time-step within the sequence, we denote the absolute time using $\tau$.
To capture the relative positions between each visit, we compute the time
difference (in days) between each visit, $\Delta\tau$. By representing EMR
data this way, we aimed to obtain a final representation that looks similar to
natural language wherein each code represents a word, the codes in each visit
represent a sentence and the whole patient history represents a document. This
approach has been widely used in many works involving deep learning models and
EMR data (Choi et al., 2016b, a; Li et al., 2020; Rasmy et al., 2021; Ashfaq
et al., 2019).

The visits $v_{i,t}$ are a multivariate sequence of discrete tokens which need
to be numerically encoded. The tokens can be of different types such as
diagnosis codes, procedure codes or speciality of visit. We create a single
vocabulary $W$ of tokens and use that to numerically encode tokens $c\in W$ as
one-hot-encoded (OHE) vectors $\widehat{c}\in\mathbb{R}^{|W|}$, where
$\widehat{c_{p}}=1$ for the $p_{th}$ code in the vocabulary. As is typical in
every NLP pipeline, we project the OHE vectors into a smaller denser subspace
($X_{emb}\in\mathbb{R}^{E}$) by multiplying each OHE token $\widehat{c}$ with
an embedding matrix, $W_{emb}\in\mathbb{R}^{|W|\times E}$.

After the data ingestor module, the input data is transformed into the shape
$T\times V\times E$ for each patient, where $T$ is the total number of visits,
$V$ is size of the intra-visit dimension, i.e., the number of codes per visit,
and $E$ is the size of the embedding vectors. In general, the number of visits
depends on the patient, and the number of codes per visit depends both on the
patient as well as the visit. Nevertheless, in the implementation the data of
each patient are converted to the same dimensions by zero-padding.

#### Visit Summarizer

Another important design consideration for EMR datasets is how to efficiently
summarize the intra-visit axis containing several codes in a single visit,
since sequential models like RNNs and Transformers expect the input sequence
to be one-dimensional. Typically, the embeddings along the intra-visit axis
are added together (Choi et al., 2016b) or are flattened out along the time
axis (Li et al., 2020; Rasmy et al., 2021). Also, other methods to aggregate
these embeddings have been proposed, resulting in better accuracy and
convergence properties (Choi et al., 2019; Kumar et al., 2020).

#### Sequence Transducer and Task Layers

The sequence transducer module tries to capture some temporal patterns that
are present in the input data. Typically, it consists of some Recurrent,
Convolutional or Transformer Neural Network based models. We discuss these
methods in detail in Section 4. The final module is typically an output layer
in the neural network and is task specific.

In the next section, we propose using an axial decomposition mechanism,
inspired by Ho et al. (2019) which combines the visit summarizer and sequence
transducer modules. The axial decomposition is described in detail in Section
4.1.

##  4 Methods

###  4.1 SANSformer Components

We hypothesize that the transformer architecture without self-attention would
be sufficient to capture the patterns that occur in EMR data. However our
model still heavily relies on the other components of the transformer. The
transformer performs a non-linear operation that transforms an input tensor
$x_{in}\in\mathbb{R}^{B}\times\mathbb{R}^{T}\times\mathbb{R}^{E}$, where $B$
is the batch size, $T$ is the sequence length and $E$ is the embedding
dimension, onto the output tensor of the same dimensions. Its primary
components are the positional encodings, the self-attention mechanism, and the
point-wise feedforward layers. The architecture also consists of skip-
connections and layer normalization (Ba et al., 2016) that help in handling
the vanishing gradient problem with deeper layers and regularization,
respectively. In this section we will discuss the various architectural
choices that were made to adapt the transformer architecture for the EMR
dataset. For a holistic discussion on the Transformer architecture, we refer
to the original paper, Vaswani et al. (2017).

Figure 3: Axial SANSformer. _Left_ The overall architecture of of one layer of
the axial SANSformer model. The $\Delta\tau$ and positional encodings are
added to the embedded input. _Right_ Demonstration of how different parts of
the data are considered when updating a certain token (or code) within one
visit. The model separately considers all previous visits (Causal Time Mixer),
and all other tokens in the same visit (Visit Mixer), where a previous visit
is represented by a sum of the tokens in that visit. The illustration is
inspired by Vaswani et al. (2017); Ho et al. (2019).

#### Encoding positions

Since transformers are order-invariant and cannot utilize temporal information
of a sequence, positional encodings were designed to inject absolute position
information of each token into the model. The original paper proposed using
sinusoidal functions of different wavelengths to encode different positions
along the embedding axis, thus the combination of these values will be unique
for each token position. Specifically, the position encoding takes the form
$PE(t,2i)=\sin(t/10000^{2i/d_{model}})$ and
$PE(t,2i+1)=\cos(t/10000^{2i/d_{model}})$ for the even and odd indices of the
embedding axis, respectively. Alternatively, we could encode each position in
the sequence $t=1,2,\dots$ as a separate token and learn their embedding
weights during training (Devlin et al., 2018). We empirically observed that
the original implementation of positional embedding improved the model
performance and requires fewer trainable parameters.

For EMR data, merely embedding positions as described above would not be
sufficient. This is because the strength of the relationship between tokens in
consecutive visits weakens if the time difference between them increases. For
example, two visits in successive days tend to be more related than two visits
eight months apart. To properly inject this information into the model we
designed $\Delta\tau$ embeddings, where we embed the number of days elapsed
between each pair of consecutive visits into an embedding vector. To
summarize, we add both the positional encoding and the $\Delta\tau$ encoding
to the input of the EMR transformer model.

Figure 4: Detailed schematic of the axial mixer model. We decompose the
original tensor of embedding size (E), intra-visit size (V) and time-steps
(T), axially into tensors of dimension $T\times E$ and $V\times E$. GELU
activation is applied throughout the mixer mechanism.

#### Mixers

The purpose of self-attention in transformers was mainly to capture the
patterns between different tokens across the time (and/or visit) axis. Thus
transformer models without any self-attention would not be effective since the
point-wise feedforward layers do not enforce any cross-time interactions among
the tokens. However, as demonstrated by Tolstikhin et al. (2021); Liu et al.
(2021); Touvron et al. (2021), this can be amended by applying the feedforward
layers across the time axis. We borrow the terminology from Tolstikhin et al.
(2021) and call this mechanism “mixers”, since it allows cross-token
interactions along a given axis using matrix multiplication (or equivalently
$1\times 1$ convolution). Unlike self-attention mechanism that captures up to
third-order interactions (with key, value and query), the mixer mechanism only
captures the second-order (Liu et al., 2021). We hypothesize however, that
this level of interaction is sufficient for most applications in EMR.

Our implementation of cross-token interactions is more in line with the
Spatial Gating Unit (SGU) introduced in Liu et al. (2021), illustrated in
Figure 3. The input $X\in\mathbb{R}^{T\times E}$ of sequence length $T$ and
embedding dimension $E$ is transformed into output $Y$ having the same
dimensions,

$\displaystyle Z$ $\displaystyle=\text{GELU}(XU)$ (1)
$\displaystyle\widehat{Z}$ $\displaystyle=\text{SGU}(Z)$ (2) $\displaystyle Y$
$\displaystyle=\widehat{Z}V$ (3)

Here, $U\in\mathbb{R}^{E\times 2P}$ and $V\in\mathbb{R}^{P\times E}$ are
weight matrices that are optimized during training and $P$ is the projection
dimension which is typically larger than $E$. The GELU stands for Gaussian
Error Linear Units (Hendrycks and Gimpel, 2016), which is used as the
activation function in modern versions of transformers, such as BERT. The SGU
function in Equation 2 applies cross-time mixing, where
$Z\in\mathbb{R}^{T\times 2P}$ is split along the projected dimension into two
chunks, $Z_{1},Z_{2},\in\mathbb{R}^{T\times P}$. Next, an affine
transformation is applied to one of the chunks

$\displaystyle Z_{1},Z_{2}$ $\displaystyle=\text{split}(Z)$ (4)
$\displaystyle\widehat{Z_{2}}$ $\displaystyle=\text{GELU}(WZ_{2}+b)$ (5)
$\displaystyle\widehat{Z}$ $\displaystyle=Z_{1}\odot\widehat{Z_{2}}$ (6)

where $W\in\mathbb{R}^{T\times T}$ is another trainable weight matrix and
$\odot$ denotes element-wise multiplication. For clarity, we have omitted the
normalization operations and skip-connections from the equations above.

As demonstrated by Shazeer (2020), the Gated Liner Unit (GLU) variant of the
GELU activation empirically improves the transformer performance, but has a
higher computational cost. Consequently, we use the GELU activation throughout
the mixer components, but instead use the GLU-variant as the activation
function in the feedforward block. Figure 3 _Left_ shows an overview of one
layer of the mixer.

#### Axial mixing to account for token interactions

The purpose of attention is to capture interactions between tokens, and a
major challenge with the attention mechanism is that both the memory and
computational complexity of the operation is $\mathcal{O}(N^{2})$, where $N$
is the length of the sequence. This severely limits the applicability of self-
attention to longer sequences. When the input sequence is multidimensional, as
in our case, the common practice is to flatten the sequence along the time-
axis before applying attention along this flattened dimension (Chen et al.,
2020). For example, if we have an input sequence of dimension $T\times V$,
where $T$ is the sequence length and $V$ the dimension of intra-visit axis,
the flattened sequence will be of length $TV$ and thus will have complexity
$\mathcal{O}(T^{2}V^{2})$.

Instead of attention, we will account for token interactions by applying
mixing independently along each axis of the input tensor, as illustrated in
Figure 3 _Right_. Specifically, we will apply mixing along the visit and time
axes separately and then add their respective outputs. This means that, when
learning the representation of a specific code in a given visit, all other
codes in the same visit will be considered (visit axis), as well as all the
previous visits (time-axis). As explained in Ho et al. (2019), implementing
this type of interaction mechanism can be done by absorbing the tokens along
the visit axis into the batch dimension while mixing over the time dimension
and vice versa for the visit dimension. As a result we would be reducing the
complexity to $\mathcal{O}(V^{2}T+VT^{2})$.

We still need to modify the timewise axial mixing to better suit the EMR data
since in Ho et al. (2019) it was designed specifically for images. In
particular, if we naively mix the tokens along the time axis (by absorbing the
visit tokens into the batch axis), then the resulting sequences will only
consist of codes that correspond to the specific index in the intra-visit
dimension that we are currently processing. Unlike in images, where this would
correspond to creating a sequence of pixels for each row, this would make
little semantic sense for EMR. For example, index 1 would only consist of the
second diagnosis code assigned by the doctor for each previous visit, but
these codes may not have any meaningful order. Thus, instead, when mixing over
previous visits, each previous visit is represented by a sum of tokens over in
that visit (i.e., a sum over the intra-visit axis). This, effectively, amounts
to considering previous visits as a whole, instead of using just a single code
in the given index.

The SANSformer architecture’s SGU is decomposed along the time and visit axes.
However, since we will be training the model in an autoregressive manner, we
need to employ causal masking on the SGU weight matrices. This can be easily
done by zeroing out all the upper-triangular elements of the matrix before the
multiplication operation. We do not need causal masking with the visit-mixer
and its corresponding SGU because each token within the visit has the same
time-step (from a medical point of view, diagnoses done on an earlier point in
time may be recorded in the visit, but this does not violate the causal
ordering). The axial mixer along with the causal masking scheme for time mixer
is illustrated in Figure 3 Right.

A detailed schematic showing all the operations needed to implement the axial
mixer architecture has been included in Figure 4. The details of the SGU and
the feedforward block are presented in Figure 6 in Appendix A.

###  4.2 Self-supervised pre-training

Since a typical EMR dataset is made up of many features from different
modalities, there are plenty of ways to take advantage of self-supervised pre-
training. In this work we devise a self-supervised regime that is derived from
the task at hand—predicting a patient’s future healthcare utilization. The
Pummel dataset consists of a seven year history for each patient but in
practice the prediction for a specific year is based on the previous year’s
history and this restriction may be, for example, due to the unavailability of
long patient histories, people moving between healthcare providers or changing
their health insurance plans, the annual cycle of resource allocation, etc.
(Ellis et al., 2018). In cases where longer histories are available as in the
Pummel data, their use in the prediction may improve accuracy of all the
models (see e.g., Kumar et al. (2020)). The gain in accuracy, however, needs
to be investigated.

Thus, our task is to predict the healthcare utilization in the year following
the first recorded diagnosis of a certain disease, using the previous one-year
patient history as the input. This leaves potentially plenty of past data
unused for each patient with the specific diagnosis, as well as all data from
those individuals who do not have the specific diagnosis, i.e., the general
population. In order to utilize this information, we devise a self-supervised
method where the model predicts the number of visits for the next year given
the one-year history, one-by-one for all available years, regardless of the
diagnoses that the patient received that year. Nonetheless, in order to
prevent any information leakage, we only pre-train on the set of patients who
do not belong to any of the specific subgroups we are analysing, i.e., on the
general population.

##  5 Experiments

###  5.1 Experiment Setup

With all data sources and subgroups, we use the 80-20 cross validation, thus
using 80% of the available data for training and held out 20% of the data for
test set. We further split 20% of the training data to be used as the
validation set, for tuning the hyperparameters and early stopping. To enforce
realistic constraints when forecasting the healthcare utilization, we ensure
that the patients in the test set occur later in time than patients in the
training set. Thus, the training set contains only patients diagnosed for the
first time with a specific illness that defines the subgroup between years
2012 – 2015 while the test set is for patients diagnosed with the illness for
the first time in the year 2016.

#### Compared methods

For baselines, we consider 5-fold cross-validated linear models - Lasso and
Logistic Regression for the regression and classification tasks, respectively.
We represent the sequence of discrete tokens in bag-of-words form as explained
in Section 3.3. For the neural network baselines we borrow models from NLP
literature. Even though the transformer models have dominated most of the
state-of-the-art leaderboards in NLP (Wang et al., 2018), in general they
require large amount of training data to obtain satisfying results. LSTMs on
the other hand are more data efficient, especially in domains where the time
series is sparse. Thus, throughout all experiments we consider these two
transducers as our neural network baseline. Further, we compare the axial
mixers with the additive variants of the LSTM, transformer and SANSformer
models. In the additive variants, the visit summarizer (Section 3.3) simply
adds the token embeddings along the intra-visit axis.

The linear models are implemented using scikit-learn (Pedregosa et al., 2011).
All neural network models are built using PyTorch (Paszke et al., 2017) and
are trained on a machine with a single NVIDIA Tesla V100 GPU with 16GB VRAM.
The networks are trained with Rectified Adam optimizer (Liu et al., 2019)
where the learning rate is varied with a cyclical schedule with linear warmup
and decay (Smith, 2017).

Table 2: Results from comparison of architectures on Pummel test set. Type 2 diabetes subgroup was trained on $30\%$ of training data ($N$ = 12,528). The Lasso model is trained with 5-fold cross-validation (cv) and the neural networks are trained with 5 random restarts. We report the mean $\pm$ standard deviation of the runs. RI indicates Random Initialization of the neural net weights. ‘Add.’ indicates additive visit summarizer. | Task 1 | Task 2
---|---|---
Model |  $R^{2}$ $y_{\text{count}}$ |  Spearman Corr. $y_{\text{count}}$ |  MAE $y_{count}$ |  Spearman Corr. $y_{\text{diag}}$ |  MAE $y_{\text{diag}}$
Lasso (cv = 5) |  0.13 $\pm$ 0.00 |  0.44 $\pm$ 0.00 |  6.34 $\pm$ 0.01 | 0.1988 |  1.42 $\pm$ 0.01
LSTM (RI) |  0.18 $\pm$ 0.01 |  0.44 $\pm$ 0.00 |  6.15 $\pm$ 0.02 |  0.20 $\pm$ 0.00 |  1.38 $\pm$ 0.01
Add. Transformer (RI) |  0.15 $\pm$ 0.01 |  0.44 $\pm$ 0.00 |  6.29 $\pm$ 0.02 |  0.14 $\pm$ 0.01 |  1.49 $\pm$ 0.01
Add. SANSformer (RI) |  0.20 $\pm$ 0.02 |  0.47 $\pm$ 0.00 |  6.07 $\pm$ 0.04 |  0.19 $\pm$ 0.00 |  1.41 $\pm$ 0.02
Axial SANSformer (RI) |  0.15 $\pm$ 0.02 |  0.44 $\pm$ 0.01 |  6.23 $\pm$ 0.05 |  0.17 $\pm$ 0.01 | $1.42\pm 0.02$
Table 3: Results from comparison of architectures on MIMIC-IV test set. The _Target Offset_ column denotes the number of visits that were removed from the tail-end of the patient history, thus higher offset would mean earlier detection. The Logistic Regression model is trained with 5-fold cross-validation (cv) and the neural networks are trained with 5 random restarts. We report the mean $\pm$ standard deviation of the runs. Ties were resolved based on the max score for each model. RI indicates Random Initialization of the neural net weights.  |  | Task 3 | Task 4
---|---|---|---
Model |  Target Offset | AUC $y_{\text{death}}$ |  Spearman Corr. $y_{\text{los}}$
Logistic Regression (cv = 5) | 2 visits |  0.72 $\pm$ 0.00 | -
Logistic Regression (cv = 5) | 5 visits |  0.68 $\pm$ 0.00 | -
LSTM (RI) | 2 visits |  0.75 $\pm$ 0.00 |  0.41 $\pm$ 0.00
LSTM (RI) | 5 visits |  0.69 $\pm$ 0.00 | -
Transformer (RI) | 2 visits |  0.68 $\pm$ 0.00 |  0.32 $\pm$ 0.01
Transformer (RI) | 5 visits |  0.65 $\pm$ 0.04 | -
SANSformer (RI) | 2 visits |  0.74 $\pm$ 0.00 |  0.39 $\pm$ 0.00
SANSformer (RI) | 5 visits |  0.70 $\pm$ 0.01 | -
Axial SANSformer (RI) | 2 visits |  0.75 $\pm$ 0.00 |  0.41 $\pm$ 0.00
Axial SANSformer (RI) | 5 visits |  0.69 $\pm$ 0.00 | -
Table 4: Results on the divergent subgroups in PUMMEL dataset for Multiple
Sclerosis (MS) and Bipolar disorder (BP) subgroups. We observe that pre-
trained models (PT) consistently outperform the randomly initialized models
(RI). The Lasso model is trained with 5-fold cross-validation (cv) and the
neural networks are trained with 5 random restarts. We report the mean $\pm$
standard deviation of the runs.

|  | Task 1 | Task 2
---|---|---|---
Model |  Subgroup Type |  $R^{2}$ $y_{\text{count}}$ |  Spearman Corr. $y_{\text{count}}$ |  MAE $y_{\text{count}}$ |  Spearman Corr. $y_{\text{diag}}$ |  MAE $y_{\text{diag}}$
Lasso (cv = 5) | MS (N=123) |  0.08 $\pm$ 0.10 |  0.55 $\pm$ 0.02 |  7.73 $\pm$ 0.39 |  0.12 $\pm$ 0.05 |  1.34 $\pm$ 0.12
LSTM (RI) | MS (N=123) |  0.06 $\pm$ 0.12 |  0.47 $\pm$ 0.03 |  7.88 $\pm$ 0.38 |  0.01 $\pm$ 0.07 |  1.45 $\pm$ 0.05
LSTM (PT) | MS (N=123) |  0.30 $\pm$ 0.05 |  0.53 $\pm$ 0.02 |  7.00 $\pm$ 0.23 |  0.25 $\pm$ 0.03 |  1.02 $\pm$ 0.04
Transformer (RI) | MS (N=123) |  -0.86 $\pm$ 0.13 |  0.00 $\pm$ 0.04 |  10.92 $\pm$ 0.46 |  0.00 $\pm$ 0.13 |  1.53 $\pm$ 0.07
Transformer (PT) | MS (N=123) |  0.23 $\pm$ 0.01 |  0.52 $\pm$ 0.00 |  6.88 $\pm$ 0.05 |  0.12 $\pm$ 0.00 |  1.17 $\pm$ 0.03
SANSformer (RI) | MS (N=123) |  -0.24 $\pm$ 0.08 |  0.05 $\pm$ 0.01 |  9.29 $\pm$ 0.19 |  0.00 $\pm$ 0.12 |  1.67 $\pm$ 0.08
SANSformer (PT) | MS (N=123) |  0.19 $\pm$ 0.01 |  0.52 $\pm$ 0.01 |  7.06 $\pm$ 0.04 |  0.23 $\pm$ 0.01 |  1.25 $\pm$ 0.00
Axial SANSformer (RI) | MS (N=123) |  -0.49 $\pm$ 0.17 |  0.00 $\pm$ 0.19 |  9.78 $\pm$ 0.40 |  -0.05 $\pm$ 0.08 |  1.70 $\pm$ 0.15
Axial SANSformer (PT) | MS (N=123) |  0.10 $\pm$ 0.07 |  0.44 $\pm$ 0.02 |  8.14 $\pm$ 0.32 |  0.22 $\pm$ 0.01 |  1.26 $\pm$ 0.10
Lasso (cv = 5) | BP (N=827) |  0.16 $\pm$ 0.04 |  0.44 $\pm$ 0.01 |  8.99 $\pm$ 0.27 |  0.26 $\pm$ 0.02 |  1.49 $\pm$ 0.06
LSTM (RI) | BP (N=827) |  0.19 $\pm$ 0.03 |  0.45 $\pm$ 0.01 |  8.43 $\pm$ 0.10 |  0.16 $\pm$ 0.04 |  1.43 $\pm$ 0.06
LSTM (PT) | BP (N=827) |  0.28 $\pm$ 0.01 |  0.45 $\pm$ 0.01 |  8.09 $\pm$ 0.11 |  0.35 $\pm$ 0.01 |  1.24 $\pm$ 0.05
Transformer (RI) | BP (N=827) |  -0.28 $\pm$ 0.12 |  0.04 $\pm$ 0.00 |  10.00 $\pm$ 0.21 |  0.05 $\pm$ 0.02 |  1.75 $\pm$ 0.08
Transformer (PT) | BP (N=827) |  0.25 $\pm$ 0.01 |  0.53 $\pm$ 0.00 |  7.91 $\pm$ 0.02 |  0.26 $\pm$ 0.00 |  1.41 $\pm$ 0.01
SANSformer (RI) | BP (N=827) |  0.07 $\pm$ 0.04 |  0.41 $\pm$ 0.03 |  8.52 $\pm$ 0.11 |  0.09 $\pm$ 0.03 |  1.56 $\pm$ 0.02
SANSformer (PT) | BP (N=827) |  0.29 $\pm$ 0.02 |  0.54 $\pm$ 0.00 |  7.75 $\pm$ 0.04 |  0.31 $\pm$ 0.01 |  1.34 $\pm$ 0.05
Axial SANSformer (RI) | BP (N=827) |  -0.13 $\pm$ 0.07 |  0.07 $\pm$ 0.01 |  9.82 $\pm$ 0.09 |  -0.01 $\pm$ 0.02 |  1.70 $\pm$ 0.14
Axial SANSformer (PT) | BP (N=827) |  0.24 $\pm$ 0.03 |  0.50 $\pm$ 0.02 |  7.89 $\pm$ 0.10 |  0.36 $\pm$ 0.01 |  1.21 $\pm$ 0.01

###  5.2 Comparison of model architectures without pre-training

First, we perform ablation studies to evaluate each of the proposed models
against the baselines, when no pre-training was applied. We split the study
into two parts where we train the models on the Pummel dataset and the MIMIC-
IV dataset respectively. On each occasion, we tune the hyperparameters for the
baseline models to obtain the optimum values for learning rate, weight decay,
dropout ratio and the number of heads for the transformer models. We use the
hyperparameter optimization framework Optuna (Akiba et al., 2019) to search
through the parameter space. Table 5 in Appendix shows the ranges of the
search queries. The LSTM, transformer and additive SANSformer each contain
four layers, while the axial SANSformer contains two layers. The transformer
and SANSformer models use 16 attention heads each.

As explained in Section 3.2, we predict healthcare utilization in the PUMMEL
dataset as the number of physical visits to a healthcare facility
($y_{\text{count}}$, Task 1). Since this a regression problem, we report the
$R^{2}$ score, Spearman’s rank correlation, and the Mean Absolute Error (MAE)
between the predicted and observed counts. For the task of predicting the
counts of the causes of visits ($y_{\text{diag}}$, Task 2), we report the
Spearman’s rank correlation and the MAE averaged across the six disease
category counts, since $R^{2}$ scores in this case were close to zero due to
the extreme levels of sparsity in the target distribution. In this experiment,
the divergent subgroup of patients for whom we want to make the predictions
(i.e., predict $y_{\text{count}}$ and $y_{\text{diag}}$ for the next year),
consists of patients diagnosed with the type 2 diabetes (T2D) in the previous
year. Because the number of T2D patients is sufficiently large, we also vary
the amount of data used for training in order to understand how the
performance of the methods changes with different training set sizes. Table 2
shows the results when the models were trained on a subset of 30% of patients
diagnosed with T2D, and Appendix C shows results for different training set
sizes.

For the MIMIC experiments we report the Area Under the Receiver Operating
characteristic Curve (AUC) score for the mortality prediction (Task 3) and the
Spearman’s rank correlation for the length of stay prediction (Task 4) in
Table 3. We would like our model to predict the mortality well in advance, so
we report the results at two target offsets: 2 visits and 5 visits before
death. We can see in Table 3 that the AUC score drops when the offset is
larger.

To summarize the results, we can see that when weights are randomly
initialized, i.e., without pre-training, both LSTM and SANSformer models
obtain comparable results. The good performance of the LSTM model, especially
in the low training data settings, is expected since most of the state-of-the-
art results in the domain of EMR with clinical codes are held by some form of
RNN-based sequential model (Choi et al., 2019). We also see that the axial
attention mechanism is helpful in MIMIC but not so much in PUMMEL, and we
hypothesize that this is because the MIMIC tasks (especially Task 4) require
capturing complex interactions between tokens. In the next section we will
compare these models when their weights are initialized using pre-training.

Figure 5: Data efficiency study for pre-trained model architectures (PT) in
the test set. _Row 1_ shows the Mean Absolute Error (MAE) for two Pummel
tasks: forecasting number of visits to the hospital and forecasting the
distribution of visits. _Row 2_ shows the Spearman’s rank correlation for both
these tasks. We have also included the best performing randomly initialized
model (RI) for scale.

###  5.3 Improving predictions in divergent subgroups using pre-training

We hypothesize that pre-training models using data from the general population
(N$\approx$1 million) will improve predictions in mutually exclusive subgroups
of patients whose sample sizes are much smaller. To study this in detail, we
consider three divergent subgroups in the PUMMEL dataset: T2D (N=41,000),
Bipolar disorder (BP, N=800), and Multiple Sclerosis (MS, N=120). For each
subgroup, we predict the utilization in the year following the diagnosis in
terms of the same tasks as before: the number of visits ($y_{\text{count}}$,
Task 1) and the distribution of six disease categories ($y_{\text{diag}}$,
Task 2), as described in Section 3.2. The results in this section are obtained
by fine-tuning models that have been pre-trained on the general population, as
described in Section 4.2.

First, because the T2D sample size is sufficiently large, we artificially vary
the size of the corresponding subgroup during finetuning to get an idea of how
the size affects the performance of the different models. Figure 5 shows the
results of the pre-trained models for the different sizes of the T2D data. We
can see that the LSTM and the SANSformer models perform well even when the
size of the training data is small. Pre-trained additive and axial
SANSformers, however, benefit more than LSTMs when the size of the training
data increases. The SANSformer with axial attention performs the best when
$N>200$, indicating that there is a minimum threshold of training size below
which we should use LSTM.

Next, we compare the performance of the pre-trained models and the randomly
initialized models on the two remaining divergent subgroups: Multiple
Sclerosis (MS) and Bipolar Disorder (BP). Modeling these subgroups has two
main challenges: small sample sizes and target histogram which greatly
deviates from that of the general population (Figure 1). Table 4 shows that
irrespective of the model architecture, pre-training the model on the general
population and fine-tuning on the smaller subgroup greatly improves the
results as compared to random initialization, and this is especially
pronounced when the training data size is small. Furthermore, as expected, we
see that the ordering of the performance of the models depends on the metric
used. When the goal of the analysis is to rank patients according to their
usage of healthcare services (Spearman correlation), we see that the
SANSformer architecture is consistently either the best or close to the best
across all the three subgroups and the two tasks. The limitations of the
randomly initialized neural network models in low datasettings is even more
apparent when we compare their results with the Lasso baseline. This again
makes a strong case for self-supervised pre-training.

###  5.4 Summary of Results

  * •

When randomly initialized, the LSTM and SANSformer are the top-2 models and
they obtain comparable results on the PUMMEL and MIMIC datasets (Tables 2 and
3).

  * •

The benefit of pre-training is far more evident in the SANSformer models as
shown in Figure 5 which shows the data efficiency of each model with varying
training data size (for T2D patients).

  * •

For the other two divergent subgroups (Table 4), we see that the pretrained
SANSformer performs better for the BP subgroup in both tasks. For the MS
subgroup however, multiple methods are competitive depending on the metric and
the task.

  * •

The benefit of the axial attention mechanism can be seen in the autoregressive
length of stay prediction task (Task 4) in MIMIC but not so much in PUMMEL,
and we hypothesize that this is because Task 4 requires capturing complex
interactions between tokens.

##  6 Discussion

Our experiments indicated that the proposed SANSformer model achieves a strong
performance across the different datasets and tasks, in both the randomly
initialized and pre-trained settings, and consistently outperforms the
conventional Transformer with self-attention. Attention-free MLP models
(Tolstikhin et al., 2021; Touvron et al., 2021; Liu et al., 2021) have
previously shown competitive performance in computer vision and NLP tasks,
however we observed that even greater benefits are obtained in the EMR domain.
This is likely because the complexity of the relationships between the tokens
in natural language is much higher than between the tokens of clinical codes,
thus allowing us to do away with the self-attention mechanism. One key
limitation of the mixer models is that the length of the sequence that can be
modeled is limited by the size of the weight matrix $W$ in equation 5, whereas
RNNs and transformers can model, in theory, infinitely long input sequences.
However, this limitation is somewhat alleviated in practice since their
performance deteriorates sharply with longer input sequences, due to the
vanishing gradient problem in RNNs and the memory limitations in transformers.

On the other hand, we found that the benefit of the axial summarizer mechanism
is dependent on the complexity of the prediction tasks and the size of the
data available for training. When randomly initialized, we observed that it is
important for the complex sequence-to-sequence prediction in the MIMIC tasks
but it was not useful in the simpler PUMMEL tasks. When initialized with pre-
trained weights, the axial summarizer clearly improved the data efficiency,
i.e., improved faster when more data are available, as compared to the
commonly used additive visit summarizer. Modeling the intra-visit codes
separately could also be useful for interpreting the predictions as
demonstrated in Choi et al. (2016a, 2019). An interesting extension of our
work would be to compare the features learned by SANSformer models with
features that were engineered by domain experts in the field of risk
adjustment.

Finally and importantly, we observed that for all model architectures, the
prediction accuracy in a divergent subgroup could be greatly improved by self-
supervised pre-training on the general population. Such pre-training has
significant potential to improve the subgroup specific prediction models that
have high practical importance in allocating the invaluable healthcare
resources.

Acknowledgments

This work was funded by the Academy of Finland (grants no. 286607 and 294015)
to PM, and the Academy of Finland Flagship programme: Finnish Center for
Artificial Intelligence (FCAI)). We acknowledge the computational resources
provided by Aalto Science-IT project.

## References

  * Akiba et al. (2019) Takuya Akiba, Shotaro Sano, Toshihiko Yanase, Takeru Ohta, and Masanori Koyama.  Optuna: A next-generation hyperparameter optimization framework.  In _Proceedings of the 25rd ACM SIGKDD International Conference on Knowledge Discovery and Data Mining_ , 2019.
  * Ashfaq et al. (2019) Awais Ashfaq, Anita Sant’Anna, Markus Lingman, and Sławomir Nowaczyk.  Readmission prediction using deep learning on electronic health records.  _Journal of Biomedical Informatics_ , 97:103256, 2019.
  * Ba et al. (2016) Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton.  Layer normalization.  _arXiv preprint arXiv:1607.06450_ , 2016.
  * Brown et al. (2020) Tom B Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al.  Language models are few-shot learners.  _arXiv preprint arXiv:2005.14165_ , 2020.
  * Chen et al. (2020) Mark Chen, Alec Radford, Rewon Child, Jeffrey Wu, Heewoo Jun, David Luan, and Ilya Sutskever.  Generative pretraining from pixels.  In _International Conference on Machine Learning_ , pages 1691–1703. PMLR, 2020.
  * Choi et al. (2016a) Edward Choi, Mohammad Taha Bahadori, Joshua A Kulas, Andy Schuetz, Walter F Stewart, and Jimeng Sun.  RETAIN: An interpretable predictive model for healthcare using reverse time attention mechanism.  _arXiv preprint arXiv:1608.05745_ , 2016a.
  * Choi et al. (2016b) Edward Choi, Mohammad Taha Bahadori, Andy Schuetz, Walter F Stewart, and Jimeng Sun.  Doctor ai: Predicting clinical events via recurrent neural networks.  In _Machine Learning for Healthcare Conference_ , pages 301–318. PMLR, 2016b.
  * Choi et al. (2019) Edward Choi, Zhen Xu, Yujia Li, Michael W Dusenberry, Gerardo Flores, Yuan Xue, and Andrew M Dai.  Graph convolutional transformer: Learning the graphical structure of electronic health records.  _arXiv preprint arXiv:1906.04716_ , 2019.
  * Devlin et al. (2018) Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova.  BERT: Pre-training of deep bidirectional transformers for language understanding.  _arXiv preprint arXiv:1810.04805_ , 2018.
  * Ellis et al. (2018) Randall P Ellis, Bruno Martins, and Sherri Rose.  Risk adjustment for health plan payment.  In _Risk Adjustment, Risk Sharing and Premium Regulation in Health Insurance Markets_ , pages 55–104. Elsevier, 2018.
  * Harutyunyan et al. (2019) Hrayr Harutyunyan, Hrant Khachatrian, David C Kale, Greg Ver Steeg, and Aram Galstyan.  Multitask learning and benchmarking with clinical time series data.  _Scientific Data_ , 6(1):1–18, 2019.
  * Hendrycks and Gimpel (2016) Dan Hendrycks and Kevin Gimpel.  Gaussian error linear units (GELUs).  _arXiv preprint arXiv:1606.08415_ , 2016.
  * Ho et al. (2019) Jonathan Ho, Nal Kalchbrenner, Dirk Weissenborn, and Tim Salimans.  Axial attention in multidimensional transformers.  _arXiv preprint arXiv:1912.12180_ , 2019.
  * Howard and Ruder (2018) Jeremy Howard and Sebastian Ruder.  Universal language model fine-tuning for text classification.  _arXiv preprint arXiv:1801.06146_ , 2018.
  * (15) Unto Häkkinen, Tuukka Holster, Taru Haula, Satu Kapiainen, Petra Kokko, Merja Korajoki, Suvi Mäklin, Lien Nguyen, Tuuli Puroharju, and Mikko Peltola.  Need adjustment for financing health and social services in finland. national institute for health and welfare (THL).  URL https://www.julkari.fi/bitstream/handle/10024/139708/URN_ISBN_978-952-343-483-7.pdf.
  * Johnson et al. (2020) A Johnson, L Bulgarelli, T Pollard, S Horng, LA Celi, and R Mark.  MIMIC-IV (version 1.0), 2020.
  * Kumar et al. (2020) Yogesh Kumar, Henri Salo, Tuomo Nieminen, Kristian Vepsalainen, Sangita Kulathinal, and Pekka Marttinen.  Predicting utilization of healthcare services from individual disease trajectories using rnns with multi-headed attention.  In _Machine Learning for Health Workshop_ , pages 93–111. PMLR, 2020\.
  * Lee-Thorp et al. (2021) James Lee-Thorp, Joshua Ainslie, Ilya Eckstein, and Santiago Ontanon.  FNet: Mixing tokens with fourier transforms.  _arXiv preprint arXiv:2105.03824_ , 2021.
  * (19) Jari Lehtonen, Jukka Lehtovirta, and Päivi Mäkelä-Bengs.  THL-toimenpideluokitus (2013).  URL https://www.julkari.fi/handle/10024/104401.
  * Li et al. (2020) Yikuan Li, Shishir Rao, José Roberto Ayala Solares, Abdelaali Hassaine, Rema Ramakrishnan, Dexter Canoy, Yajie Zhu, Kazem Rahimi, and Gholamreza Salimi-Khorshidi.  BEHRT: transformer for electronic health records.  _Scientific Reports_ , 10(1):1–12, 2020.
  * Lipton et al. (2015) Zachary C Lipton, David C Kale, Charles Elkan, and Randall Wetzel.  Learning to diagnose with LSTM recurrent neural networks.  _arXiv preprint arXiv:1511.03677_ , 2015.
  * Liu et al. (2021) Hanxiao Liu, Zihang Dai, David R So, and Quoc V Le.  Pay attention to MLPs.  _arXiv preprint arXiv:2105.08050_ , 2021.
  * Liu et al. (2019) Liyuan Liu, Haoming Jiang, Pengcheng He, Weizhu Chen, Xiaodong Liu, Jianfeng Gao, and Jiawei Han.  On the variance of the adaptive learning rate and beyond.  _arXiv preprint arXiv:1908.03265_ , 2019.
  * McGuire and van Kleef (2018) Thomas G McGuire and Richard C van Kleef.  Regulated competition in health insurance markets: Paradigms and ongoing issues.  In _Risk Adjustment, Risk Sharing and Premium Regulation in Health Insurance Markets_ , pages 3–20. Elsevier, 2018.
  * Melas-Kyriazi (2021) Luke Melas-Kyriazi.  Do you even need attention? a stack of feed-forward layers does surprisingly well on imagenet.  _arXiv preprint arXiv:2105.02723_ , 2021.
  * Paszke et al. (2017) Adam Paszke, Sam Gross, Soumith Chintala, Gregory Chanan, Edward Yang, Zachary DeVito, Zeming Lin, Alban Desmaison, Luca Antiga, and Adam Lerer.  Automatic differentiation in pytorch.  2017\.
  * Pedregosa et al. (2011) Fabian Pedregosa, Gaël Varoquaux, Alexandre Gramfort, Vincent Michel, Bertrand Thirion, Olivier Grisel, Mathieu Blondel, Peter Prettenhofer, Ron Weiss, Vincent Dubourg, et al.  Scikit-learn: Machine learning in python.  _The Journal of Machine Learning Research_ , 12:2825–2830, 2011.
  * Peters et al. (2018) Matthew E Peters, Mark Neumann, Mohit Iyyer, Matt Gardner, Christopher Clark, Kenton Lee, and Luke Zettlemoyer.  Deep contextualized word representations.  _arXiv preprint arXiv:1802.05365_ , 2018.
  * Radford et al. (2019) Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever.  Language models are unsupervised multitask learners.  _OpenAI blog_ , 1(8):9, 2019.
  * Raffel et al. (2019) Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu.  Exploring the limits of transfer learning with a unified text-to-text transformer.  _arXiv preprint arXiv:1910.10683_ , 2019.
  * Rasmy et al. (2021) Laila Rasmy, Yang Xiang, Ziqian Xie, Cui Tao, and Degui Zhi.  Med-BERT: pretrained contextualized embeddings on large-scale structured electronic health records for disease prediction.  _npj Digital Medicine_ , 4(1):1–13, 2021.
  * Rose and McGuire (2019) Sherri Rose and Thomas G McGuire.  Limitations of p-values and R-squared for stepwise regression building: a fairness demonstration in health policy risk adjustment.  _The American Statistician_ , 73(sup1):152–156, 2019.
  * Ruder (2017) Sebastian Ruder.  An overview of multi-task learning in deep neural networks.  _arXiv preprint arXiv:1706.05098_ , 2017.
  * Shazeer (2020) Noam Shazeer.  GLU variants improve transformer.  _arXiv preprint arXiv:2002.05202_ , 2020.
  * Shrestha et al. (2018) Akritee Shrestha, Savannah Bergquist, Ellen Montz, and Sherri Rose.  Mental health risk adjustment with clinical categories and machine learning.  _Health Services Research_ , 53:3189–3206, 2018.
  * Smith (2017) Leslie N Smith.  Cyclical learning rates for training neural networks.  In _2017 IEEE Winter Conference on Applications of Computer Vision (WACV)_ , pages 464–472. IEEE, 2017.
  * Tolstikhin et al. (2021) Ilya Tolstikhin, Neil Houlsby, Alexander Kolesnikov, Lucas Beyer, Xiaohua Zhai, Thomas Unterthiner, Jessica Yung, Daniel Keysers, Jakob Uszkoreit, Mario Lucic, et al.  MLP-Mixer: An all-MLP architecture for vision.  _arXiv preprint arXiv:2105.01601_ , 2021.
  * Touvron et al. (2021) Hugo Touvron, Piotr Bojanowski, Mathilde Caron, Matthieu Cord, Alaaeldin El-Nouby, Edouard Grave, Armand Joulin, Gabriel Synnaeve, Jakob Verbeek, and Hervé Jégou.  ResMLP: Feedforward networks for image classification with data-efficient training.  _arXiv preprint arXiv:2105.03404_ , 2021.
  * Vaswani et al. (2017) Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin.  Attention is all you need.  _arXiv preprint arXiv:1706.03762_ , 2017.
  * Wang et al. (2018) Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel R Bowman.  GLUE: A multi-task benchmark and analysis platform for natural language understanding.  _arXiv preprint arXiv:1804.07461_ , 2018.

## Appendix A.

Figure 6: Detailed schematic of the Spatial gating unit _Top_ and the
feedforward blocks _Bottom_. GELU activation is applied throughout the mixer
mechanism. The feedforward (FF) block consists of a series of linear layers
with GLU activation and dropout. As explained in Shazeer (2020), the GLU
operation is applied by splitting the input tensor along the embedding
dimension into two chunks, adding GELU activation on one of the chunks and
combining them with a multiplicative gate.

## Appendix B.

Table 5: Search ranges for the hyperparameter search. The best hyperparameters were chosen based on validation loss from 50 trials. Parameters | Range
---|---
Learning Rate | range [1e-5, 1e-3]
Weight Decay | range [1e-5, 1e-1]
Dropout Ratio | pick [0.1, 0.2, 0.3]
# of transformer heads | pick [8, 16]

## Appendix C.

We plot the MAE and Spearman’s rank correlation for each of the randomly
initialized model architectures while varying the size of the training
dataset. We can see that the LSTM models achieve better performance when the
size of the training data is small.

Figure 7: Data efficiency study for randomly initialized model architectures
(RI) in the test set. _Row 1_ shows the Mean Absolute Error (MAE) for two
Pummel tasks: forecasting number of visits to the hospital and forecasting the
distribution of visits. _Row 2_ shows the spearman rank correlation for both
these tasks.

## Appendix D.

The LSTM was found to be competitive in many tasks, especially with a small
training set. Therfore, we performed a more detailed error analysis to further
understand the differences between the models. The box-plot in Figure 1 ranks
the predicted counts of the number of visits for each patient against the true
counts for both the LSTM and SANSformer models. Here, both the models were
fine-tuned with just 80 training examples. We see that the SANSformer model is
clearly better at sorting the high-risk patients into their appropriate risk
groups as indicated by the low-variance of interquantile range of the
predictions, and this finding consistently re-occurred in our experiments.
Thus even though the two models obtained relatively similar $R^{2}$ scores,
SANSformers could be helpful in identifying the high risk-risk patients (Rose
and McGuire, 2019).

Figure 1: Error analysis for pre-trained LSTM (_Left_) and SANSformer
(_Right_) models after fine-tuning with 80 datapoints. Box plot compares the
predicted visits for patients in the type 2 diabetes subgroup. We observe that
both the models are good, but the SANSformer is slightly better at sorting the
patients into their appropriate risk groups and has lower variance.