alg/attentivess/Tutorial_for_Attentive_State-space_Models.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Attentive State-space Models\n",
    "\n",
    "### A complete guide for using attentive state-space models for sequential, time-series data.\n",
    "\n",
    "This notebook is based on the paper *\"Attentive State-Space Modeling of Disease Progression\"* by *Ahmed M. Alaa* and *Mihaela van der Schaar* submitted to **NeurIPS 2019**. In this paper, we develop a deep probabilistic model that learns accurate and interpretable structured representations for disease trajectories. Unlike Markovian state-space models, in which the dynamics are memoryless, our model uses an attention mechanism to create \"memoryful\" dynamics, whereby attention weights determine the dependence of future disease states on past medical history. In what follows, we briefly explain the basic API for the Python (Tensorflow) implementation of the model, and provide examples for how to use."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### A brief explanation of the model\n",
    "\n",
    "An attentive state-space model is a deep probabilistic model that capitalizes on both the interpretable structured representations of probabilistic models and the predictive strength of deep learning methods. The model uses a state-space representation to segment an observable trajectory into hidden states. But unlike conventional state-space models, which are predominantly Markovian, we use RNNs to capture more complex state dynamics as shown in the figure below.\n",
    "\n",
    "![](images/comparison.png)\n",
    "\n",
    "**Sequential data models.** (a) Graphical model for an RNN. Diamonds denote a deterministic representation, (b) Graphical model for an HMM: circles denote probabilistic states, (c) Graphical depiction of an attentive state space model. With a slight abuse of graphical model notation, thickness of arrows reflect attention weights. Because the attention mechanism can be made arbitrarily complex, our model can capture complex dynamics while maintaining its structural interpretability. The attention mechanism is implemented via a sequence-to-sequence RNN architecture."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Setting up the experiment "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In what follows, we use a synthetic dataset to train an attentive state-space model and display the learned model parameters. Before we start, we need to import two libraries:\n",
    "\n",
    "- **data.make_data**: the file containing the synthetic data generation model.\n",
    "- **models.SeqModels**: the file containing the attentive state-space model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from data.make_data import generate_trajectory\n",
    "from models.SeqModels import attentive_state_space_model as attentive_state_space\n",
    "\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We start by drawing random, variable-sized trajectories from the synthetic data model using the function **generate_trajectory** as follows"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "X_observations, true_states = generate_trajectory(num_states=3, \n",
    "                                                  Num_observations=10, \n",
    "                                                  Num_samples=5000, \n",
    "                                                  Max_seq=30, \n",
    "                                                  Min_seq=3, \n",
    "                                                  alpha=100,\n",
    "                                                  reverse_mode=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The synthetic data model uses hidden states to generate observations. The return variables **X_observations** and **true_states** correspond to the observable variables and the hidden state trajectories. Our model is given only the observable data, and should be able to learn the underlying state trajectory that generated the data. The synthetic data generation function takes the following 6 arguments:\n",
    "\n",
    "- **num_states:** The cardinality of the hidden state space.\n",
    "- **Num_observations:** The number of observable variables (i.e., emissions, or features) that are associated with each state.\n",
    "- **Num_samples:** The total number of state trajectories to simulate.\n",
    "- **Max_seq, Min_seq:** The minimum and maximum number of obervations per trajectory. For each trajectory, the sequence length is drawn from a uniform distribution [Min_seq, Max_seq]\n",
    "- **alpha:** This is a real number that determines the extent of \"memory\" in the state trajectory. The contribution of each state realized at time $t$ (for a sequence of length $T$) on the current state transition is proportion to $e^{-\\alpha (T-t)}$. Hence, the effect of alpha can be understood through the following extreme cases:\n",
    "\n",
    "    1. $\\alpha$ = $\\infty$: Then only the current state determines the next state. This corresponds to a Markov model where transitions are completely memorlyess. Here we set $\\alpha=100$ which makes it very close to a Markov chain.\n",
    "    \n",
    "    2. $\\alpha$ = $0$: In this case, all previous states are equally important, in which cases the model needs to attend to very old state realizations in order to predict future state transitions. \n",
    "\n",
    "\n",
    "- **reverse_mode:** This is a boolean that determines whether the importance is higher for more recent or older state realizations. When it is set to *True*, the importance of previous states are set to $e^{-\\alpha (t-T)}$. In this case, setting $\\alpha$ = $\\infty$ makes the initial state the only factor that determines all state transitions, irrespective of any inetrmediate states.\n",
    "\n",
    "The observations are determined by state-dependent Gaussian distributions, the parameters of which can also be changed, but we will not need to change these during this Tutorial."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Visualizing the true data trajectories"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Both the variables **X_observations** and **true_states** are in the forms of lists, the size of which correspond to the number of trajectories that we simulated. Each list contains a variable-length sequence of observations and states. Let us see whatsequence of states is realized in the first trajectory:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trajectory_index = 0\n",
    "\n",
    "true_states[trajectory_index]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Because we set the state-space to contain only three states, the true states can take on values in the set $\\{0,1,2\\}$. Now let us visualize the first trajectory by overlaying the hidden states and the observable variables in one plot."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from matplotlib import pyplot as plt\n",
    "get_ipython().magic('matplotlib inline')\n",
    "\n",
    "fig, ax1 = plt.subplots()\n",
    "\n",
    "t = list(range(len(true_states[trajectory_index])))\n",
    "\n",
    "color = 'tab:red'\n",
    "ax1.set_xlabel('Time step')\n",
    "ax1.set_ylabel('Observations', color=color)\n",
    "ax1.plot(t, X_observations[trajectory_index], color=color, linewidth=5, alpha=0.2)\n",
    "ax1.tick_params(axis='y', labelcolor=color)\n",
    "\n",
    "ax2 = ax1.twinx()  \n",
    "\n",
    "color = 'tab:blue'\n",
    "ax2.set_ylabel('Hidden states', color=color)  \n",
    "ax2.step(t, true_states[trajectory_index], color=color, linewidth=5)\n",
    "ax2.tick_params(axis='y', labelcolor=color)\n",
    "ax2.set_yticks([0, 1, 2])\n",
    "\n",
    "fig.tight_layout()  \n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here, the blue plot reflects the underlying hidden states that generates the observations, and as we can see, the states modulate the distribution of the observable variables (red). The attentive state-space model should be able to do the following:\n",
    "\n",
    "- Learn an underlying representation for the data in terms of a state space and emission distribution.\n",
    "- Predict the state and observations trajectories.\n",
    "- Uncover the memory dynamics underlying the state transitions.\n",
    "- Generate synthetic samples that mimic those of the original data distribution."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training attentive state-space models"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We first instantiate an attentive state space model object as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = attentive_state_space(num_states=3,\n",
    "                              maximum_seq_length=30, \n",
    "                              input_dim=10, \n",
    "                              rnn_type='LSTM',\n",
    "                              latent=True,\n",
    "                              generative=True,\n",
    "                              num_iterations=50, \n",
    "                              num_epochs=3, \n",
    "                              batch_size=100, \n",
    "                              learning_rate=5*1e-4, \n",
    "                              num_rnn_hidden=100, \n",
    "                              num_rnn_layers=1,\n",
    "                              dropout_keep_prob=None,\n",
    "                              num_out_hidden=100, \n",
    "                              num_out_layers=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The attentive state-space model takes on the following arguments:\n",
    "\n",
    "- **num_states**: The number of states in the state space.\n",
    "- **maximum_seq_length**: The maximum allowable length for any trajectory in the training data. \n",
    "- **input_dim**: The number of observations/features in every time step.\n",
    "- **rnn_type**: The type of RNN used in model and inference networks. Select from (LSTM, GRU, RNN)\n",
    "- **latent, generative**: (Redundant) Whether state labels are provided, and whether the model is generative (can sample trajectories).\n",
    "- **num_iterations, num_epochs, batch_size, learning_rate**: Variational inference learning parameters.\n",
    "- **num_rnn_hidden, num_rnn_layers**: Number of hidden states/layers of all RNNs in the model.\n",
    "- **dropout_keep_prob**: Whether dropout is used for regularization.\n",
    "- **num_out_hidden, num_out_layers**: Number of hidden states/layers of output layers."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To train the model, we provide a simple sklearn-like API where the model is fit by simply providing the list of observations to the fit method of the attentative state-space object as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.fit(X_observations)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that the model is trained, we display the parameters of the representation it learned. The initial state and transition probabilities can be recovered by inspecting the attributes of **model** as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.initial_probabilities"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The initial distribution exactly matches that of the synthetic model, which has an initial state distribution of [0.5, 0.3, 0.2]. The baseline state transition matrix of the attentive state-space models is given by:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.transition_matrix"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The emission distributions for all states are assumed to be Gaussian. The mean and covariance parameters of the three states can be recovered as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.state_means.T "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.state_covars"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Predicting state trajectories via attentive state-space models"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To use our model for prediction, we again provide an sklearn-like API where the model issues prediction by simply providing the list of observations to the predict method of the attentative state-space object as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "state_inference, expected_observations, attention = model.predict([X_observations[0]])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The predict method retrieves three different variables: \n",
    "\n",
    "- **state_inference**: This is the posterior distribution over the \"future\" states given previous observations, i.e. $P(X_{t+1}| X_t,.\\,.\\,., X_0)$. For the returned vector, every element corresponds to a time step $t$, and the values is the distribution of the future states (for the next time step $t+1$). \n",
    "- **expected_observations**: The expected values of the observation in the next time step $t+1$.\n",
    "- **attention**: The attention weights assigned to all previous hidden states to issue the current prediction."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The maximum a posteriori (MAP) estimates of the inferred states can be obtained by picking the most likely state per time step as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.argmax(state_inference[trajectory_index], axis=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In what follows, we visually compare the MAP state estimates with the true states:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax1 = plt.subplots()\n",
    "\n",
    "trajectory_index = 0\n",
    "\n",
    "t = list(range(len(true_states[trajectory_index])))\n",
    "\n",
    "color = 'tab:red'\n",
    "ax1.set_xlabel('time (s)')\n",
    "ax1.set_ylabel('MAP state estimate', color=color)\n",
    "ax1.plot(t, np.argmax(state_inference[trajectory_index], axis=1), color=color, linewidth=5)\n",
    "ax1.tick_params(axis='y', labelcolor=color)\n",
    "ax1.set_yticks([0, 1, 2])\n",
    "\n",
    "ax2 = ax1.twinx()  \n",
    "\n",
    "color = 'tab:blue'\n",
    "ax2.set_ylabel('True states', color=color)  \n",
    "ax2.step(t, true_states[trajectory_index], color=color, linewidth=5)\n",
    "ax2.tick_params(axis='y', labelcolor=color)\n",
    "ax2.set_yticks([0, 1, 2])\n",
    "\n",
    "fig.tight_layout()  \n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We also visually compare the expected observation estimates with the true observations:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax1 = plt.subplots()\n",
    "\n",
    "t = list(range(len(true_states[trajectory_index])))\n",
    "\n",
    "color = 'tab:red'\n",
    "ax1.set_xlabel('Time step')\n",
    "ax1.set_ylabel('True observations', color=color)\n",
    "ax1.plot(t, X_observations[trajectory_index], color=color, linewidth=5, alpha=0.2)\n",
    "ax1.tick_params(axis='y', labelcolor=color)\n",
    "\n",
    "ax2 = ax1.twinx()  \n",
    "\n",
    "color = 'tab:blue'\n",
    "ax2.set_ylabel('Predicted average observations', color=color)  \n",
    "ax2.plot(t, expected_observations[trajectory_index], color=color, linewidth=5, alpha=0.1)\n",
    "ax2.tick_params(axis='y', labelcolor=color)\n",
    "ax2.set_yticks([0, 1, 2])\n",
    "\n",
    "fig.tight_layout()  \n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Visualizing the attention weights over time"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, in what follows, we visualize the attention weights materializing at every time-step for the trajectory above. As we can see, because the true underlying model is Markovian, we can see that the biggest attention weight is almost always placed on the most recent time-step."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import seaborn as sns\n",
    "\n",
    "Attention_weights = []\n",
    "\n",
    "for w in range(len(attention[0])):\n",
    "    \n",
    "    Attention_weights.append(np.vstack((attention[0][w], np.zeros((len(attention[0][-1]) - len(attention[0][w]),1)))))\n",
    "\n",
    "\n",
    "Attention_weights = np.array(Attention_weights).reshape((len(attention[0][-1]), len(attention[0][-1])))[:state_inference[0].shape[0], :state_inference[0].shape[0]]\n",
    "\n",
    "mask = np.zeros_like(Attention_weights)\n",
    "\n",
    "mask[np.triu_indices_from(mask)] = True\n",
    "\n",
    "with sns.axes_style(\"white\"):\n",
    "    ax = sns.heatmap(Attention_weights, mask=mask, vmax=.3, square=True)\n",
    "    ax.set_ylabel('Chronological time')\n",
    "    ax.set_xlabel('Previous time steps')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Sampling from the (generative) attentive state-space model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Because our model is generative, we can also sample new (synthetic) trajectories from the model. This means that our model can be used for generating synthetic data. To sample from the trained model, we can use the **sample** API as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sampled_states, sampled_obervations  = model.sample(trajectory_length=5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the above, the **trajectory_length** parameter decides the number of states materializing in the sampled sequence. In what follows, we visualize the sampled states and observations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from matplotlib import pyplot as plt\n",
    "get_ipython().magic('matplotlib inline')\n",
    "\n",
    "fig, ax1 = plt.subplots()\n",
    "\n",
    "t = list(range(len(sampled_states)))\n",
    "\n",
    "color = 'tab:red'\n",
    "ax1.set_xlabel('Time step')\n",
    "ax1.set_ylabel('Sampled observations', color=color)\n",
    "ax1.plot(t, sampled_obervations[0], color=color, linewidth=5, alpha=0.2)\n",
    "ax1.tick_params(axis='y', labelcolor=color)\n",
    "\n",
    "ax2 = ax1.twinx()  \n",
    "\n",
    "color = 'tab:blue'\n",
    "ax2.set_ylabel('Sampled hidden states', color=color)  \n",
    "ax2.step(t, sampled_states, color=color, linewidth=5)\n",
    "ax2.tick_params(axis='y', labelcolor=color)\n",
    "ax2.set_yticks([0, 1, 2])\n",
    "\n",
    "fig.tight_layout()  \n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## What is the effect of memory on the attention mechanism?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we look at a different synthetic datasets where we use the reverse mode of the generative synthetic model to assign bigger importance weight on older states. We take an extreme case where the initial states controls all state transitions along the trajectory (see the earlier explanation for the synthetic data.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "X_observations_new, true_states_new = generate_trajectory(num_states=3, \n",
    "                                                          Num_observations=10, \n",
    "                                                          Num_samples=5000, \n",
    "                                                          Max_seq=30, \n",
    "                                                          Min_seq=3, \n",
    "                                                          alpha=100,\n",
    "                                                          reverse_mode=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We then fit a new model on this new dataset..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_new = attentive_state_space(num_states=3,\n",
    "                                  maximum_seq_length=30, \n",
    "                                  input_dim=10, \n",
    "                                  rnn_type='LSTM',\n",
    "                                  latent=True,\n",
    "                                  generative=True,\n",
    "                                  num_iterations=50, \n",
    "                                  num_epochs=3, \n",
    "                                  batch_size=100, \n",
    "                                  learning_rate=5*1e-4, \n",
    "                                  num_rnn_hidden=100, \n",
    "                                  num_rnn_layers=1,\n",
    "                                  dropout_keep_prob=None,\n",
    "                                  num_out_hidden=100, \n",
    "                                  num_out_layers=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_new.fit(X_observations_new)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "state_inference_new, expected_observations_new, attention_new = model_new.predict([X_observations_new[0]])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let us visualize the attention weights materializing over time for the new model..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import seaborn as sns\n",
    "\n",
    "Attention_weights = []\n",
    "\n",
    "for w in range(len(attention_new[0])):\n",
    "    \n",
    "    Attention_weights.append(np.vstack((attention_new[0][w], np.zeros((len(attention_new[0][-1]) - len(attention_new[0][w]),1)))))\n",
    "\n",
    "\n",
    "Attention_weights = np.array(Attention_weights).reshape((len(attention_new[0][-1]), len(attention_new[0][-1])))[:state_inference[0].shape[0], :state_inference[0].shape[0]]\n",
    "\n",
    "mask = np.zeros_like(Attention_weights)\n",
    "\n",
    "mask[np.triu_indices_from(mask)] = True\n",
    "\n",
    "with sns.axes_style(\"white\"):\n",
    "    ax = sns.heatmap(Attention_weights, mask=mask, vmax=.3, square=True)\n",
    "    ax.set_ylabel('Chronological time')\n",
    "    ax.set_xlabel('Previous time steps')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As we can see, unlike in the previous case, here we have most of the attention weights assigned to the older states! Thus, our attentive state-space model learned the underlying state dynamics in an unsupervised fashion on the basis of the observed data."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}