Title: MGP-ATTTCN: AN INTERPRETABLE MACHINE LEARNING MODEL FOR THE PREDICTION OF SEPSIS

Abstract: With a mortality rate of 5.4 million lives worldwide every year and a healthcare cost of more than 16 billion dollars in the USA alone, sepsis is one of the leading causes of hospital mortality and an increasing concern in the ageing western world. Recently, medical and technological advances have helped re-define the illness criteria of this disease, which is otherwise poorly understood by the medical society. Together with the rise of widely accessible Electronic Health Records, the advances in data mining and complex nonlinear algorithms are a promising avenue for the early detection of sepsis. This work contributes to the research effort in the field of automated sepsis detection with an open-access labelling of the medical MIMIC-III data set. Moreover, we propose MGP-AttTCN: a joint multitask Gaussian Process and attention-based deep learning model to early predict the occurrence of sepsis in an interpretable manner. We show that our model outperforms the current state-of-the-art and present evidence that different labelling heuristics lead to discrepancies in task difficulty.

1. INTRODUCTION:

Every year, it is estimated that 31.5 million people worldwide contract sepsis. With a mortality rate of 17% in its benign state and 26% for its severe state (Fleischmann et al., 2016) , sepsis is one of the leading causes of hospital mortality (Vincent et al., 2014) , costing the healthcare system more than 16 billion dollars in the USA alone (Angus et al., 2001) . Studies demonstrated that early treatment has a significant positive effect on the survival rate (Kumar et al., 2006; Nguyen et al., 2007) . In particular, Castellanos-Ortega et al. (2010) demonstrated that each hour delay in treating a patient results in a 7.6% increase in mortality.

Current methods of screening, such as Modified Early Warning System (MEWS) and Systemic Inflammatory Response Syndrome (SIRS) have been criticised for their lack of specificity, leading to low accuracies and high false alarm rates. In 2015, the Third International Consensus Definitions for Sepsis (Singer et al., 2016; Seymour et al., 2016; Shankar-Hari et al., 2016) committee worked towards incorporating medical and technological advances into an up-to-date definition of sepsis, providing scientists with widely acknowledged illness criteria. Together with the rise of Electronic Health Records (EHR), the scientific community is now armed with both the data and labelling techniques to experiment with novel prediction methods (Islam et al., 2019; Henry et al., 2015; Ghosh et al., 2017; Calvert et al., 2016; Desautels et al., 2016) , which are already proving effective in increasing survival rate (Shimabukuro et al., 2017) and promising in decreasing costs.

New models developed so far either relied on some interpretable yet simple prediction methods, such as logistic regression (Calvert et al., 2016) and decision tree based classifiers (Mao et al., 2018; Delahanty et al., 2019) , or on effective yet black-box methods such as Recurrent Neural Networks (Futoma et al., 2017b) . Moreover, the results achieved by different authors are rarely comparable: although most use the MIMIC-III data set, the disparities in labelling rules result in highly variable data sets (eg. Raghu et al. (2018) have 17,898 septic patients vs. 2,577 for Desautels et al. (2016) ). This work presents an attempt at reconciling interpretability and predictive performance on the sepsis prediction task and makes the following contributions:

• Gold standard for labelling. We provide a gold standard for Sepsis-3 labelling implemented on the MIMIC-III data set. • Novel interpretable model. We present an explainable and end-to-end trainable model based on Multitask Gaussian Processes and Attentive Neural Networks for the early prediction of sepsis.

• Empirical evaluation. We assess our model on real-world medical data and report superior predictive performance and interpretability compared to previous methods.

An overview of our proposed method is shown in Figure 1 .

2. RELATED WORK:

Medical time series diagnosis Multiple researchers have tackled the task of predicting sepsis and septic shock. Works on septic shock include exploration of survival models (Henry et al., 2015) and Hidden Markov Models (Ghosh et al., 2017) . However, these models rely on the assumption that a patient has already developed sepsis and focus on patterns of patients' further deterioration.

Other authors (Calvert et al., 2016; Desautels et al., 2016; Mao et al., 2018; Delahanty et al., 2019) use linear models and decision trees on engineered features to predict sepsis onset, thus failing to capture temporal patterns. More recently, Kam & Kim (2017) and Raghu et al. (2018) employed recurrent neural networks to better capture time dependencies. Crucially, all these models rely on either averaging or forward imputation of data points to create equidistant inputs. This transformation creates data artefacts and discards relevant uncertainty: in the medical field, the absence of data is a conscious decision made by professionals implying an underlying belief of the patient state. Futoma et al. (2017b) and Moor et al. (2019) tackled this issue with Multitask Gaussian Processes (MGPs), however their models lack the interpretability necessary in the medical field.

Irregularly sampled time series The most common solution to missing values is forward imputation (Calvert et al., 2016) . Lipton et al. (2016) utilise forward imputation coupled with a missingness indicator fed into a black-box model. Although this method retains information about data presence, it is not clear how the information is later interpreted by the model and hence does not meet our transparency criteria. Ghassemi et al. (2015) use MGPs to fit sparse medical data, however they optimise their model for the data fit and use the parametrisation as input for a classifier rather than optimising the model for a classification task. Both Futoma et al. (2017a) and Moor et al. (2019) use MGPs with end-to-end training, although their temporal covariance function is shared across all variables. Finally, Futoma et al. (2017b) uses MGPs with multiple time kernels in a similar fashion to our model, although they infer the number of kernels from hyperparameter tuning rather than the data itself.

Attention based neural networks Attention was first introduced on the topic of machine translation (Bahdanau et al., 2014) . Since then, the concept has been used in natural language processing (Yang et al., 2016; Yu et al., 2018) and image analysis (Mnih et al., 2014; Schlemper et al., 2019) 

   (1)

where N p is the patient's observation period length. We also define static features s p = {s p,M +1 , ..s p,M +Q } with features identifiers k ∈ {M + 1, . . . , M + Q}, corresponding to timeindependent quantities, such as age, gender and first admission unit. Finally, we define sepsis labels l p ∈ {0, 1}. Given the sparsity of the data, we can define the compact representation of all observed values:

EQUATION

The goal of the model is, for a given set {t p , y p , s p } to predict the label l p . In order to remove clutter, we will from now on drop the patient-specific subscript p from all notation, and the feature subscript k from time notation, simplifying t p,k,i to t i .

3.2. MULTITASK GAUSSIAN PROCESS (MGP):

Gaussian processes are commonly known for their ability to generate coherent function fits to a set of irregular samples, by modelling the data covariance. As they easily account for uncertainty and do not require homogeneously sampled data, Gaussian processes are the perfect candidate model to deal with relatively small amounts of medical data.

Following Bonilla et al. (2008) , we use a Multitask Gaussian Process (MGP) to capture feature correlation and Li & Marlin (2016)'s end-to-end training framework, in a similar manner to Futoma et al. (2017a) . Given an hourly spaced time series {t i } 0 i=-Np (where 0 is the time of prediction), the MGP layer produces a set of posterior predictions for each feature, which will then be fed into a classification model.

We define a patient-independent prior over the true values of {y:

EQUATION

where {K tt l (t i , t j )} l∈L are time point covariances varying in smoothness, {K k l (k, k )} l∈L are feature covariances at a given smoothness level, independent of time, and L are smoothness clusters. Over all variables and time points, the multivariate model has covariance

EQUATION

where D = diag(σ k ) are the noise terms associated to each feature and ⊗ is the Kronecker product.

The posterior over t = {t i } 0 i=-Np is a multivariate Gaussian with mean and covariance:

µ = l∈L K k l ⊗ K tt l l∈L K k l ⊗ K tt l + D ⊗ I -1 y Σ = l∈L K k l ⊗ K t t l - l∈L K k l ⊗ K tt l l∈L K k l ⊗ K tt l + D ⊗ I -1 l∈L K k l ⊗ K t t l (6)

In order to approximate the posterior distribution, we then take Monte Carlo samples y MC from Y MGP ∼ N (µ, Σ).

Note that there are two main feature clusters: vital signs (vitals) and laboratory results (labs). Vitals are noisier and sampled more often, whereas labs are more monotone and rarely sampled. As opposed to Futoma et al. (2017b) , we do not treat the number of clusters L as hyperparameters but set L = 2 and define

K t l (t i , t j ) = exp -|t i -t j | λ l (7)

as Ornstein-Uhlenbeck (OU) kernels with lengths λ 1 and λ 2 , each representing a cluster smoothness. OU kernels are well suited to capture local variations and do not assume infinite differentiability as Squared Exponential kernels do. In our case, differentiablity implies a level of smoothness which does not apply to medical records and only introduces unnecessary bias. In addition, given the scarce availability of labs, all short lengthscales would be an ill fit to the data. We hence discarded kernels varying over lengthscales such as the Cauchy and the Rational Quadratic kernels. K k l (k, k ) are free-form covariance matrice that are learned by gradient descent.

To feed the MGP samples into the classifier, we fix the model time window to N = 25 by either zero padding or truncating the beginning of the time series. We choose to do so at the beginning of the time series in order to align prediction times as the last step of the temporal classification model. Here, we also integrate the static variables by broadcasting them over each time point1 .

3.3. ATTENTION TIME CONVOLUTIONAL NETWORK (ATTTCN):

The concept of attention was born in machine translation (Bahdanau et al., 2014) : given an input sentence embedding S = {h 1 , . . . h |S| }, the attention mechanism produces weights {α i 1 , . . . α i |S| } such that α i j ∈ [0, 1], j α i j = 1, and a context vector c i = j α i j h j used to predict target word i. The weights α i j can therefore be interpreted as the importance of the input sentence's j th word to produce the i th word of the translation.

More recently, Choi et al. (2016) have applied attention to clinical time series. Given a time series {x 1 , . . . x T } ⊂ R r , the authors first create a time-independent embedding of the data {v 1 , . . . v T } ⊂ R m . They then use inversed recurrent neural networks (RNN) to create weights α ∈ R T and β ∈ R T ×m , where α j ∈ [0, 1] and β ij ∈ [-1, 1], with softmax and tanh activations respectively. The context vectors then take the form c i = j≤i α j β j v j and are fed into a multilayer perceptron with softmax activation to yield a prediction.

The attention model we devised borrows some ideas from Choi et al. (2016) . The interpolated data y MC ∈ R N ×(M +Q) is directly fed into two temporal convolutional networks (TCNs) (Lea et al. (2017) ) and generates embeddings z

= [z 1 , . . . , z N ] ∈ R N ×(M +Q) and z = [z 1 , . . . , z N ] ∈ R N ×(M +Q) .

TCNs are a class of neural networks composed of causal convolutions stacked into Residual Blocks. A causal convolution is a 1D convolutional layer which only takes inputs from the past to generate its output, avoiding any information leakage from the future. Residual Blocks are made of two causal convolutional layers together with ReLU activation functions, dropout and L2 regularisations. The Residual Blocks also include an identity map from the input of the block added to the output. As we only use up to 12 layers, this last step is omitted in our architecture. TCNs have shown to outperform RNNs (Bai et al. (2018)) , are faster at training and do not suffer from vanishing gradients. Given the latter, inverting the time series similarly to Lea et al. (2017) also becomes an unnecessary step which we omit.

We generate the attention weights α and β as

EQUATION

W α,0 , W α,1 ∈ R M +Q b α,0 , b α,1 ∈ R (10) W β,0 , W β,1 ∈ R (M +Q)×(M +Q) b β,0 , b β,0 ∈ R M +Q (11) such that α = [α 0 , α 1 ] ∈ R N ×2 and β = [β 0 , β 1 ] ∈ R N ×(M +Q)×2 .

We then create two context vectors, one for each of negative and positive label predictions

EQUATION

where y MC,j is broadcast to meet the dimensionality of β j,δ . We then predict the labels as

EQUATION

In our case, we are only interested in making predictions with the latest available data. We therefore only use llast to train the model. This of course can be easily modified to suit any specific use-case.

Since the MGP output is directly multiplied by weights c i , the classification model can be interpreted as a scoring mechanism where each past point y MC,ij contributes α i,0 β ij,0 to the time series being classified as positive, and α i,1 β ij,1 to the time series being classified as negative. The positive and negative scores are then normalised to represent probabilities of the positive or negative labelling. As we design both α and β to be non-negative, we can hence directly look at the average α and β over Monte Carlo samples to see which time points and features contribute most strongly to the network's positive or negative decision.

4. DATA:

Sepsis is defined as a life-threatening organ dysfunction caused by a dysregulated host response to infection (Singer et al., 2016) . The latter is usually interpreted as the administration of antibiotics coupled with the culture of blood samples, generating a suspicion of infection window, whereas the former is interpreted as a two point increase in Sequential Organ Failure Assessment (SOFA) within such a suspected infection window. We make use of the MIMIC-III data set (Johnson et al., 2016) and encode the Sepsis (2016) . In order to validate our results, we carry out all experiments using both labelling techniques.

We proceed to extract times series of case and control patients for a set of commonly recorded vitals, labs and static variables and normalise their values. Following Moor et al. (2019) , in order to keep the data set length balanced, we match the time series lengths of control patients to those of case patients using the class balance ratio. In addition, we create up to seven copies of each time series and truncate the last zero to six hours of data, effectively creating early prediction patients and augmenting our data set. We remove excessively noisy or computationally intensive data and train the model over It can be seen that our proposed labels are harder to fit than the ones by Moor et al. (2019) . Moreover, our proposed model outperforms the baselines on both label sets, especially for earlier prediction horizons. different hyperparameters, randomly resampling an equal number of case and control patients to counteract the data set imbalance.

5. EXPERIMENTAL RESULTS:

We compare our model's performance to the performance of the InSight algorithm (Calvert et al., 2016) and to the state-of-the-art MGP-TCN algorithm (Moor et al., 2019) . Figure 3 shows the predictive performance of the models for different time horizons.

5.1. COMPARISON BETWEEN DIFFERENT DATA LABELS:

The first result is the difference in performance of models applied to the different labelling methods. The SOFA contributor assumption Moor et al. make has two main implications. Firstly, it considerably restricts the number of patients. Assuming that sicker patients receive more medical attention, the patients included are likely to be in worse conditions than the septic patients excluded and hence easier to classify. Secondly, it delays sepsis onset. For example, a patient having a severe liver failure with few other recorded vitals, followed by an overall collapse further in time will have septic onset at the time of its liver failure in our records, whereas it will only be considered septic at the time of the overall collapse in Moor et al.'s labels. On the other hand, the labels we produce reflect the incomplete nature of medical data: even if only a part of all the potentially relevant tests are carried out at any given time, a doctor must be able to assess a patient's well-being and foresee potential complications. The difference in labels implies a discrepancy in task difficulty: Moor et al.'s labels present an easier learning problem, however they define a more narrow use-case in real-world scenarios.

Indeed, when assessing the performance of the different models on the two different data labellings, it becomes evident that our proposed labels are harder to fit. This means that predicting sepsis in a realistic setting on the intensive care unit is probably much harder than previous work would suggest.

5.2. MODEL PERFORMANCE:

We find that our MGP-AttTCN model has a better performance when presented with patients further in time from sepsis onset. In the case of Moor et al.'s labels the difference is clearly noticeable, whereas with our labels it is of lower statistical significance. With our labels, both MGP-TCN and MGP-AttTCN have a stronger performance than InSight. The intuition behind this result is the robustness of the models to missing data: both MGP-TCN and MGP-AttTCN account for the data uncertainty and hence have a better performance on lower resolution and more irregular data. On Moor et al.'s labels, the MGP-TCN model does not seem to significantly outperform the InSight model, suggesting that those labels might be easy enough to not require a particularly pronounced robustness to missing data. However, the additional attention of the proposed MGP-AttTCN model does seem to gain a clearer advantage here than on our labels, presumably due to a more complete set of features that can be attended to.

5.3. MGP INTERPRETABILITY:

Inspecting the learned covariances (Figure 4 ), we notice that the two OU lengthscales converged to represent two clusters within the selected variables: a shorter lengthscale (around two hours) represents noisy data, whereas a larger lengthscale (around 64 hours) represents smoother observations. In addition, the feature covariance matrix for the short lengthscale puts more emphasis on vitals, while the one for the long lengthscale puts more emphasis on labs, fitting our initial intuition that vitals vary more rapidly. Graphically, one can observe this by inspecting the diagonals on the covariance heatmaps.

On a more granular level, the two covariance matrices also provide insights about the underlying variables. One can for instance observe that the body temperature (tempc) has a larger variance than the systolic and diastolic blood pressure (sysbp, diabp), following the general clinical intuition. Moreover, we can observe correlations between different features, such as a negative correlation between temperature and heart rate, which also seems to coincide with the general medical expectation. These covariances can then for instance be used by the model to extrapolate a full function from a single INR observation with an inverse correlation to the pulse oximetry observations (Fig. 5 ).

5.4. ATTENTION WEIGHTS:

One important benefit of our model compared to current approaches is its interpretability due to the attention mechanism. Once the samples have been drawn, the weights α and β provide us with more information about the importance of different time points and features for the model's behaviour. The attention weights for an exemplary patient trajectory are depicted in Figure 5 .

Overall, the absolute values of α are small for points further from the prediction time and increasingly larger closer to it. A good example of this behaviour is the fourth row in Figure 5 , where feature importance increases in time. We can also see there, that different features can have opposing effects on the prediction. While the elevated heart rate close to the prediction time increases the likelihood of a sepsis prediction (first column, yellow weights), the lowered prothrombin values reduce this likelihood (third column, blue weights). Interestingly, the low prothrombin values are not actually measured in this example, but predicted by the MGP purely based on the other measured features and the learned covariances. Finally, α × β × y MC gives the individual score contribution of each feature at each time point. These weights are shown in the last row of the figure. It can again be seen that the attention weights are generally larger in magnitude closer to the prediction time. Moreover, about half of the features have significant non-zero attention weights, while the others seem to not be important for the prediction in this example.

These visualizations could be used by doctors to make an informed decision about whether or not to trust the prediction of the model for each given patient, thus facilitating the interpretability and accountability that is crucial in medical applications.

6. CONCLUSION:

We have shown that current data sets for the early prediction of sepsis underestimate the true difficulty of the problem and proposed a new labelling for the MIMIC-III data set that corresponds more closely to a realistic intensive care setting. Moreover, we have proposed a new machine learning model, the MGP-AttTCN, which outperforms the state-of-the-art approaches on the easier labels from the literature as well as on our proposed harder labels. Additionally, our model provides an interpretable attention mechanism that would allow clinicians to make more informed decisions about trusting its predictions on a case-by-case basis.

Potential avenues for future work include a more thorough discussion with clinicians to potentially make our proposed labels even more representative of the real-world task, and architectural improvements, for instance by meta-learning the MGP prior (Fortuin & Rätsch, 2019) , amortizing the latent MGP inference for performance gains (Fortuin et al., 2019) , or discretizing the latent space for improved interpretability (Fortuin et al., 2018) . 76.4 ± 0.9 81.4 ± 0.9 80.5 ± 0.5 where b static is generated by the output of a two layers multilayer perceptron applied to the static data. Although this solution is computationally lighter and provides a more elegant interpretation, it does not allow the attention mechanism to utilise the information about the patient's static state when making a decision about its vitals and labs values. For the scope of this paper, we hence decided for earlier static data integration.

C.3 PARAMETER HYPERSEARCH AND TRAINING:

As the datasets are highly unbalanced, we carry out a case set oversampling: we randomly resample the case set to have the same size as the control set. In addition, at each iteration we sample equally the same number of cases and controls, then feed a shuffled version into the model. In this manner the model will see an equal number of controls and cases and will not become biased towards zero labels. This procedure does not happen for either of the validations and test sets, as the results would not compare to real life settings.

For both our core model MGP-AttTCN and all baselines, in order to select the best possible hyperparameters, we performed a hyperparameter random search, as described in Table 4 .

C.4 NUMERICAL RESULTS:

The results are shown in Tables 5 and 6 .

Table 6 : Area under the ROC curve for our labels Time to onset InSight MGP-TCN MGP-AttTCN 6h 55.5 ± 2.7 60.9 ± 0.9 63.6 ± 0.9 5h 55.5 ± 1.4 62.4 ± 1.0 63.9 ± 1.6 4h 55.5 ± 1.6 62.4 ± 0.9 64.1 ± 1.2 3h 55.9 ± 1.3 61.6 ± 0.7 63.6 ± 0.9 2h 57.3 ± 1.5 62.4 ± 0.9 63.6 ± 0.6 1h 58.8 ± 1.8 64.5 ± 0.9 64.7 ± 0.9 0h 61.4 ± 1.1 67.0 ± 0.6 66.5 ± 0.9

:

see Appendix C.2 for more information on this design choice

Figures and Tables:

Figure 1: Proposed model architecture.

-3 criteria following Johnson & Pollard (2018)'s code available on GitHub, with the help of Moor et al. (2019)'s code that the authors have generously provided. One key difference between our assumptions and Moor et al. (2019)'s is the handling of missing SOFA contributor values: if one or more SOFA contributors are missing, Moor et al. do not calculate the total score. On the other hand, we assume such a contributor to be within a healthy norm, implying a zero contribution. The latter heuristics is in line with the official Sepsis 3 definition in Singer et al.

Figure 2: Interpretation of the different attention weights in our model.

Figure3: Performance of different models. It can be seen that our proposed labels are harder to fit than the ones byMoor et al. (2019). Moreover, our proposed model outperforms the baselines on both label sets, especially for earlier prediction horizons.

Figure 4: Heatmaps of the learned MGP covariance matrices between the data features for the two different smoothness clusters.

Figure 5: Visualization of the journey of an exemplary patient trajectory through our proposed model architecture. The raw features (row 1), measured at irregular time points, are interpolated by an MGP (row 2). Samples from the MGP posterior can then be aggregated into means and variances for each feature on a fixed, regularly-spaced time grid (row 3). These values are then attended to by the TCN (row 4), where positive attention weights are yellow and negative ones blue. Row 5 shows the attention weights separated by features (x-axis) and time point (y-axis).

Hyperparameter search

Table content: <table><tr><td>Hyperparameter</td><td colspan="2">Random Search</td></tr><tr><td/><td>min</td><td>max</td></tr><tr><td>MGP Monte Carlo samples</td><td>4</td><td>20</td></tr><tr><td>TCN kernel size</td><td>2</td><td>6</td></tr><tr><td>TCN number of Residual Blocks</td><td>2</td><td>12</td></tr><tr><td>TCN number of hidden layers</td><td>10</td><td>55</td></tr><tr><td>TCN dropout rate</td><td>0</td><td>0.99</td></tr><tr><td>TCN L2 regularisation</td><td>0</td><td>250</td></tr></table>

Area under the ROC curve forMoor et. al. labels

Table content: <table><tr><td>Time to onset</td><td>InSight</td><td colspan="2">MGP-TCN MGP-AttTCN</td></tr><tr><td>6h</td><td colspan="2">66.4 ± 0.7 65.0 ± 2.2</td><td>67.0 ± 4.3</td></tr><tr><td>5h</td><td colspan="2">65.9 ± 1.5 66.4 ± 2.8</td><td>73.2 ± 2.7</td></tr><tr><td>4h</td><td colspan="2">64.7 ± 0.6 68.3 ± 2.9</td><td>73.4 ± 2.9</td></tr><tr><td>3h</td><td colspan="2">67.0 ± 1.0 66.6 ± 2.4</td><td>72.2 ± 3.6</td></tr><tr><td>2h</td><td colspan="2">67.7 ± 1.2 66.4 ± 2.7</td><td>73.9 ± 2.5</td></tr><tr><td>1h</td><td colspan="2">69.5 ± 0.8 68.2 ± 2.0</td><td>76.8 ± 1.7</td></tr><tr><td>0h</td><td/><td/></tr></table>

Peer review (iclr2020):
I*ve read the rebuttal and I*d like to keep my score as is. My main concern is the questionable role of attention in making the model more interpretable (which is the main contribution of the paper). ########################### The paper proposes a new model for automated sepsis detection using multitask GP and attention-based GP. The sepsis detection problem is of paramount importance in the clinical domain and the authors rightly emphasized that point. Also the paper tries to combine interpretability with prediction accuracy by using attention mechanism. The paper is generally well written and well motivated; however, in terms of technical novelty and empirical evidence the paper can be further improved. The MGP-AttnTCN model is mostly a minor modification of the model proposed by Moor et al. 2019 and has the additional attention element to be more interpretable. Unfortunately, it’s not easy for an ICLR reader without any medical background to evaluate the validity of the interpretability results provided in the paper. Furthermore, the recent works in NLP have agued against the value of attention for interpretability (see for instance “Attention is not Explanation” by Jain & Wallace 2019). That said, I believe the paper is probably a better fit for a machine learning in healthcare venue (such as MLHC). In terms of empirical evidence of the prediction accuracy the paper only compares with Moor et al 2019 (which does not show a significant improvement in the realistic setting) and a much older InSight paper (2016). This would have been typically enough for a paper with major technical novelty; however, for this paper, I believe adding more recent baselines and discussing the advantages of the method over these baselines would be necessary. Minor: Caption in Figure 1 can be more informative and useful for the reader if you add more details on different parts of the model. “Graphically, once can observe” should be “one can observe” .