\section{Introduction}
\label{sec:intro}

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 \cite{liu2021pay} and \cite{tolstikhin2021mlp}. Deep learning methods have been successfully applied in the past to predict the cause for the next visit \citep{lipton2015learning,choi2016retain} and time-duration to the next visit \citep{choi2016doctor,harutyunyan2019multitask} 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 \citep{vaswani2017attention} that have been pre-trained in an unsupervised fashion on a larger text corpus \citep{devlin2018bert,raffel2019exploring,radford2019language} are dominating the state-of-the-art in several of the NLP tasks. \cite{li2020behrt} and \cite{rasmy2021med}, 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 \cite{choi2016retain} and \cite{choi2019graph}, 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 \textbf{S}ans \textbf{AN}y \textbf{S}elf-Attention Trans\textbf{formers}) 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 \citep{mcguire2018regulated}. 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 \citep{shrestha2018mental,ellis2018risk}. 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 \ref{fig:div-subgroups} (\emph{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 (\textit{single model}) to predict the utilization for patients in a test set, a subset of the bipolar disorder subgroup, and compare it with a \textit{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 \ref{fig:div-subgroups} (\emph{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 \ref{fig:div-subgroups} (\emph{down}) where the \textit{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.

\begin{figure}[!ht]
  \centering
    \includegraphics[width=0.8\linewidth]{images/divergent_density2.png}

    \label{fig:div-subgroups-sub1}

    \bigbreak
    \bigbreak

    \includegraphics[width=0.8\linewidth]{images/divergent_boxplot.png}
    \label{fig:div-subgroups-sub2}

  \caption{The problem of predicting in a divergent subgroup. \emph{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 \textit{topcap}, where we counter the long-tail problem by clipping the max count value at 36 visits; \emph{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.}
  \label{fig:div-subgroups}
\end{figure}

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 \citep{johnson2020mimic} 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
\begin{itemize}
  \item 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.
  \item 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.
  \item 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.
\end{itemize}

\section{Related Work}

\textbf{Deep Learning on EMR} \cite{lipton2015learning} 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. \cite{choi2016doctor} 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 \cite{choi2016retain}, which developed a bidirectional attention based model focused interpretability; \cite{harutyunyan2019multitask} created a benchmark framework for evaluating EMR models; \cite{choi2019graph} developed a convolutional transformer model that has an added inductive bias to model the interactions between different codes occurring at one visit while \cite{kumar2020predicting} used a 1x1 convolutional network to model this interaction; \cite{li2020behrt,rasmy2021med} applied pre-trained BERT transformers to EMR data and demonstrated superior results compared to traditional recurrent neural networks (RNN) or convolutional networks.

\textbf{Self-supervised pre-training} \cite{howard2018universal,peters2018deep} 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 \citep{devlin2018bert,radford2019language,brown2020language,raffel2019exploring}.

\textbf{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. \cite{tolstikhin2021mlp} replaced the self-attention mechanism with a "mixing" mechanism. \cite{liu2021pay} simplified the mixer further by replacing the mixer with spatial gating mechanisms. Other concurrent works with similar objectives include \cite{touvron2021resmlp,lee2021fnet,melas2021you}. 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.

\section{Setup}

\subsection{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.

\subsubsection*{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.

\begin{table}[]
  \caption{\textbf{Basic statistics of the two EMR datasets.}}
  \label{tab:basic-ehr-stats}
  \centering
  \footnotesize
  \begin{tabular}{@{}lrr@{}}
  \toprule
                                & \multicolumn{1}{c}{\textbf{Pummel}} & \multicolumn{1}{l}{\textbf{MIMIC-IV}} \\ \midrule
  \# 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 \\ \bottomrule

  \end{tabular}
\end{table}

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 \citep{thlproc}, 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.

\subsubsection*{MIMIC}

The MIMIC-IV {\footnotesize \emph{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 {\fontfamily{qcr}\selectfont patients, admissions, diagnoses\_icd, procedures\_icd, drgcodes} and {\fontfamily{qcr}\selectfont services} tables. The information from each visit was grouped by the admission identifier {\fontfamily{qcr}\selectfont hadm\_id} and all the visits were grouped by the patient identifier {\fontfamily{qcr}\selectfont subject\_id}. The resulting dataset consisted of 256,878 patients with an average of 2.04 visits per patient.

Table \ref{tab:basic-ehr-stats} lists some basic statistics of both datasets.

\subsection{Tasks}
\label{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 \citep{ruder2017overview}.

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:

\begin{itemize}
  \item Task 1: The number of physical visits to the healthcare centers ($y_{\text{count}}$)
  \item Task 2: The counts of the physical visits due to six specific disease categories ($y_{\text{diag}}$)
\end{itemize}

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) \citep{thlneedfor}. 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
\begin{itemize}
  \item Task 3: The probability of inpatient mortality ($y_{\text{death}}$), and
  \item Task 4: The estimated length of stay (LoS) of the next visit ($y_{\text{los}}$)
\end{itemize}

For Task 3 we utilize the {\fontfamily{qcr}\selectfont 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.

\begin{figure}[!ht]
  \centering
  \includegraphics[width=0.8\linewidth]{images/main_figure.png}

  \caption[]{\textbf{EMR Modeling Pipeline.}\emph{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. \emph{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. }
  \label{fig:main-figure}
\end{figure}

\subsection{EMR modeling pipeline}
\label{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 \ref{fig:main-figure}). Each module applies some functional transformations to the input data. The transformed data can then be passed onto the next module.

\subsubsection*{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 \citep{choi2016doctor,choi2016retain,li2020behrt,rasmy2021med,ashfaq2019readmission}.

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.

\subsubsection*{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 \citep{choi2016doctor} or are flattened out along the time axis \citep{li2020behrt, rasmy2021med}. Also, other methods to aggregate these embeddings have been proposed, resulting in better accuracy and convergence properties \citep{choi2019graph, kumar2020predicting}.

\subsubsection*{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 \ref{methods}. 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 \cite{ho2019axial} which combines the visit summarizer and sequence transducer modules. The axial decomposition is described in detail in Section \ref{axial-mixer}.

\section{Methods}
\label{methods}

\subsection{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 \citep{ba2016layer} 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, \cite{vaswani2017attention}.

\begin{figure}[!ht]
\centering
\begin{subfigure}{.5\textwidth}
  \centering
  \includegraphics[width=.85\linewidth]{images/axial_attention_viz.png}
\end{subfigure}%
\begin{subfigure}{.5\textwidth}
  \includegraphics[width=.9\linewidth]{images/axial_mask.png}
\end{subfigure}
    \caption{\textbf{Axial SANSformer}. \emph{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. \emph{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 (\textit{Causal Time Mixer}), and all other tokens in the same visit (\textit{Visit Mixer}), where a previous visit is represented by a sum of the tokens in that visit. The illustration is inspired by \cite{vaswani2017attention,ho2019axial}.}
\label{fig:sansformer-fig}
\end{figure}

\subsubsection*{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 \citep{devlin2018bert}. 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.

\begin{figure}[]
  \centering
  \includegraphics[width=\linewidth]{images/detailed_mixer.png}

  \caption[]{\textbf{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.}
  \label{fig:detailed-mixer}
\end{figure}

\subsubsection*{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 \cite{tolstikhin2021mlp,liu2021pay,touvron2021resmlp}, this can be amended by applying the feedforward layers across the time axis. We borrow the terminology from \cite{tolstikhin2021mlp} 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 \citep{liu2021pay}. 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 \cite{liu2021pay}, illustrated in Figure \ref{fig:sansformer-fig}. 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,

\begin{align}
  Z         &= \text{GELU}(X U) \\
  \widehat{Z}   &= \text{SGU}(Z) \\
  Y         &= \widehat{Z}V
\end{align}
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 \citep{hendrycks2016gaussian}, 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

\begin{align}
  Z_1, Z_2        &=  \text{split}(Z) \\
  \label{eq:sgu} \widehat{Z_{2}} &=  \text{GELU}(W Z_2 + b) \\
  \widehat{Z}     &= Z_1 \odot \widehat{Z_{2}}
\end{align}
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 \cite{shazeer2020glu}, 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 \ref{fig:sansformer-fig} \emph{Left} shows an overview of one layer of the mixer.

\subsubsection*{Axial mixing to account for token interactions}
\label{axial-mixer}
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 \citep{chen2020generative}. 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^2V^2)$.

Instead of attention, we will account for token interactions by applying \textit{mixing} independently along each axis of the input tensor, as illustrated in Figure \ref{fig:sansformer-fig} \emph{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 \cite{ho2019axial}, 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^2T + VT^2)$.

We still need to modify the timewise axial mixing to better suit the EMR data since in \cite{ho2019axial} 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 \ref{fig:sansformer-fig} \textit{Right}.

A detailed schematic showing all the operations needed to implement the axial mixer architecture has been included in Figure \ref{fig:detailed-mixer}. The details of the SGU and the feedforward block are presented in Figure \ref{fig:ff-sgu-fig} in Appendix A.

\subsection{Self-supervised pre-training}
\label{sec-pre-train}

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. \citep{ellis2018risk}. 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., \cite{kumar2020predicting}). 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.

\section{Experiments}

\subsection{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.

\subsubsection*{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 \ref{pipeline}. 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 \citep{wang2018glue}, 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 \ref{pipeline}) simply adds the token embeddings along the intra-visit axis.

The linear models are implemented using scikit-learn \citep{pedregosa2011scikit}. All neural network models are built using PyTorch \citep{paszke2017automatic} and are trained on a machine with a single NVIDIA Tesla V100 GPU with 16GB VRAM. The networks are trained with Rectified Adam optimizer \citep{liu2019variance} where the learning rate is varied with a cyclical schedule with linear warmup and decay \citep{smith2017cyclical}.

\begin{table}[bt]
  \caption{\textbf{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.}
  \label{pummel-regression-results}
  \centering
  \footnotesize
\begin{tabular}{@{}lccc|cc@{}}
\toprule
\multicolumn{1}{c}{\textbf{}} & \multicolumn{3}{c|}{\textbf{Task 1}}                                                                                & \multicolumn{2}{c}{\textbf{Task 2}}                                                                                \\ \cmidrule(l){2-6}
\multicolumn{1}{c}{\textbf{Model}}                & \textbf{\makecell{$R^2$ \\ $y_{\text{count}}$}} & \textbf{\makecell{Spearman \\ Corr. $y_{\text{count}}$}} & \textbf{\makecell{MAE \\ $y_{count}$}} & \textbf{\makecell{Spearman \\ Corr. $y_{\text{diag}}$}} & \textbf{\makecell{MAE \\ $y_{\text{diag}}$}} \\ \midrule
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                                                                                & \textbf{0.20} $\pm$ \textbf{0.00}                                                                                & \textbf{1.38} $\pm$ \textbf{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)          & \textbf{0.20} $\pm$ \textbf{0.02}                                                                                & \textbf{0.47} $\pm$ \textbf{0.00}                                                                                & \textbf{6.07} $\pm$ \textbf{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$                                                                                \\ \bottomrule
\end{tabular}
\end{table}

\begin{table}[bt]
  \caption{\textbf{Results from comparison of architectures on MIMIC-IV test set.} The \emph{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. }
\label{mimic-regression-results}
\centering
\footnotesize
\begin{tabular}{@{}lc|cc@{}}
\toprule
                           & \multicolumn{1}{l|}{}                       & \textbf{Task 3}                                     & \textbf{Task 4}                                              \\ \cmidrule(l){2-4}
\textbf{Model}               & \textbf{\makecell{Target \\ Offset}} & \multicolumn{1}{l}{\textbf{AUC $y_{\text{death}}$}} & \textbf{\makecell{Spearman \\ Corr. $y_{\text{los}}$}} \\ \midrule

Logistic Regression (cv = 5) & 2 visits                                    & 0.72 $\pm$ 0.00                                              & -                                                            \\
Logistic Regression (cv = 5) & 5 visits                                    & 0.68 $\pm$ 0.00                                              & -                                                            \\ \midrule

LSTM (RI)                    & 2 visits                                    & 0.75 $\pm$ 0.00                                              & 0.41 $\pm$ 0.00                                                     \\

LSTM (RI)                    & 5 visits                                    & 0.69 $\pm$ 0.00                                              & -                                                   \\\midrule

Transformer (RI)             & 2 visits                                    & 0.68 $\pm$ 0.00                                              & 0.32 $\pm$ 0.01                                                      \\
Transformer (RI)             & 5 visits                                    & 0.65 $\pm$ 0.04                                              & -                                                       \\\midrule

SANSformer (RI)              & 2 visits                                    & 0.74
$\pm$ 0.00                                              & 0.39 $\pm$ 0.00                                                       \\
SANSformer (RI)              & 5 visits                                    & 0.70 $\pm$ 0.01                                              & -                                                       \\\midrule

Axial SANSformer (RI)        & 2 visits                                    & \textbf{0.75} $\pm$ \textbf{0.00}                                              & \textbf{0.41} $\pm$ \textbf{0.00}                                                       \\

Axial SANSformer (RI)        & 5 visits                                    & \textbf{0.69} $\pm$ \textbf{0.00}                                              & -                                                        \\ \bottomrule
\end{tabular}
\end{table}

\begin{table}[bt]
  \caption{\textbf{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.}
  \label{divergent-results}
  \centering
  \footnotesize
  \scalebox{0.85}{
  \begin{tabular}{@{}lc|ccc|cc@{}}
\toprule
                                 & \multicolumn{1}{l|}{}   & \multicolumn{3}{c|}{\textbf{Task 1}}                                                                                & \multicolumn{2}{c}{\textbf{Task 2}}                                 \\ \midrule
\multicolumn{1}{c}{\textbf{Model}} & \textbf{\makecell{Subgroup \\ Type}} & \textbf{\makecell{$R^2$ \\ $y_{\text{count}}$}} & \textbf{\makecell{Spearman \\ Corr. $y_{\text{count}}$}} & \textbf{\makecell{MAE \\ $y_{\text{count}}$}} & \textbf{\makecell{Spearman \\ Corr. $y_{\text{diag}}$}} & \textbf{\makecell{MAE \\ $y_{\text{diag}}$}} \\ \midrule
Lasso (cv = 5)                         & MS (N=123)             & 0.08 $\pm$ 0.10               & \textbf{0.55} $\pm$ \textbf{0.02}                                         & \multicolumn{1}{c|}{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                                         & \multicolumn{1}{c|}{7.88 $\pm$ 0.38}                          & 0.01 $\pm$ 0.07                             & 1.45 $\pm$ 0.05                         \\
LSTM (PT)                          & MS (N=123)             & \textbf{0.30} $\pm$ \textbf{0.05}      & 0.53 $\pm$ 0.02                               & \multicolumn{1}{c|}{7.00 $\pm$ 0.23}                          & 0.25 $\pm$ 0.03                             & \textbf{1.02} $\pm$ \textbf{0.04}                \\
Transformer (RI)                   & MS (N=123)             & -0.86 $\pm$ 0.13              & 0.00 $\pm$ 0.04                                        & \multicolumn{1}{c|}{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                                         & \multicolumn{1}{c|}{\textbf{6.88} $\pm$ \textbf{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                                          & \multicolumn{1}{c|}{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                                         & \multicolumn{1}{c|}{7.06 $\pm$ 0.04}                          & \textbf{0.23} $\pm$ \textbf{0.01}                    & 1.25 $\pm$ 0.00 \\
Axial SANSformer (RI)              & MS (N=123)             & -0.49 $\pm$ 0.17              & 0.00 $\pm$ 0.19                                        & \multicolumn{1}{c|}{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                                         & \multicolumn{1}{c|}{8.14 $\pm$ 0.32}                          & 0.22 $\pm$ 0.01                             & 1.26 $\pm$ 0.10                         \\ \midrule

Lasso (cv = 5)                         & BP (N=827)             & 0.16 $\pm$ 0.04               & 0.44 $\pm$ 0.01                                         & \multicolumn{1}{c|}{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                                         & \multicolumn{1}{c|}{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                                         & \multicolumn{1}{c|}{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                                         & \multicolumn{1}{c|}{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                                         & \multicolumn{1}{c|}{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                                         & \multicolumn{1}{c|}{8.52 $\pm$ 0.11}                          & 0.09 $\pm$ 0.03                             & 1.56 $\pm$ 0.02                        \\
SANSformer (PT)                    & BP (N=827)             & \textbf{0.29} $\pm$ \textbf{0.02}      & \textbf{0.54} $\pm$ \textbf{0.00}                                & \multicolumn{1}{c|}{\textbf{7.75} $\pm$ \textbf{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                                         & \multicolumn{1}{c|}{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                                               & \textbf{0.36} $\pm$ \textbf{0.01}                             & \textbf{1.21} $\pm$ \textbf{0.01}                \\ \bottomrule
\end{tabular}
}
\end{table}

\subsection{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 \citep{optuna_2019} to search through the parameter space. Table \ref{hyper-search-range} 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 \ref{tasks}, 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 \ref{pummel-regression-results} 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 \ref{mimic-regression-results}. 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 \ref{mimic-regression-results} 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 \citep{choi2019graph}. 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.

\begin{figure}[!ht]
  \centering
  \includegraphics[width=\linewidth]{images/data_efficiency_PT_waxial.png}
  \caption[]{\textbf{Data efficiency study for pre-trained model architectures (PT) in the test set.} \emph{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. \emph{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.}
  \label{fig:data-ablation-pt}
\end{figure}

\subsection{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 \ref{tasks}. The results in this section are obtained by fine-tuning models that have been pre-trained on the general population, as described in Section \ref{sec-pre-train}.

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 \ref{fig:data-ablation-pt} 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 \ref{fig:div-subgroups}). Table \ref{divergent-results} 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.

\subsection{Summary of Results}

\begin{itemize}
    \item When randomly initialized, the LSTM and SANSformer are the top-2 models and they obtain comparable results on the PUMMEL and MIMIC datasets (Tables \ref{pummel-regression-results} and \ref{mimic-regression-results}).

    \item The benefit of pre-training is far more evident in the SANSformer models as shown in Figure \ref{fig:data-ablation-pt} which shows the data efficiency of each model with varying training data size (for T2D patients).

    \item For the other two divergent subgroups (Table \ref{divergent-results}), 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.

    \item 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.
\end{itemize}

\section{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 \citep{tolstikhin2021mlp,touvron2021resmlp,liu2021pay} 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 \ref{eq:sgu}, 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 \cite{choi2016retain,choi2019graph}. 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.

\acks{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.}

\vskip 0.2in