TeSAN Model for EHRs

• TeSAN Model for EHRs is a neural architecture that embeds timestamped medical events by combining contextual co-occurrence with detailed temporal relationships.
• It utilizes a feature-wise temporal self-attention mechanism, temporal-interval embeddings, and a skip-gram objective to generate robust, interpretable representations.
• Empirical evaluations on datasets like MIMIC-III and CMS demonstrate state-of-the-art performance in clustering, nearest-neighbor search, and mortality prediction tasks.

The Temporal Self-Attention Network (TeSAN) is a neural architecture designed for medical concept embedding in the context of longitudinal electronic health records (EHRs). By explicitly incorporating both contextual co-occurrence and fine-grained temporal relationships between medical events, TeSAN advances over prior embedding approaches for representing timestamped medical concepts as fixed-length vectors suitable for downstream predictive tasks. The model is characterized by a feature-wise temporal self-attention mechanism, temporal-interval embeddings, and a skip-gram–style learning objective. Empirical results across clustering, nearest-neighbor search, and inpatient mortality prediction on diverse EHR datasets demonstrate that TeSAN achieves state-of-the-art performance relative to representative baselines, offering robustness and interpretability in the derived embeddings (Peng et al., 2019).

1. Model Architecture and Components

TeSAN comprises four principal modules: input embedding, temporal-interval embedding, a temporal self-attention block (TeSA), and an attention pooling/output projection. The model processes an input sequence of medical concepts, each associated with a timestamp, and outputs context-rich embeddings that encode both semantic and temporal relationships.

Module	Description	Dimensions/Hyperparameters
Input Embedding	One-hot codes mapped via learned weights $W^{(e)}$	$|\mathcal{C}| \approx 7,000$, $d=100$
Temporal-Interval Embed.	Learnable table $E \in \mathbb{R}^{n_{\text{days}} \times d}$	$n_{\text{days}}$ varies per dataset, $d=100$
Temporal Self-Attention	Single “multi-dimensional” head, gated fusion	$$W^{(1)}, W^{(2)}, W^{(3)}, W^{(f1)}, W^{(f2)} \in \mathbb{R}^{d \times d}$$
Attention Pooling	Multi-dimensional attention pooling	Produces $h \in \mathbb{R}^d$

TeSAN uses a vocabulary size of $|\mathcal{C}| = 6,985$ (MIMIC-III) or $7,873$ (CMS) and an embedding dimension $|\mathcal{C}| \approx 7,000$0, with $|\mathcal{C}| \approx 7,000$1 or $|\mathcal{C}| \approx 7,000$2 as the skip window, and $|\mathcal{C}| \approx 7,000$3 negative samples as in word2vec.

2. Temporal Self-Attention Mechanism

A distinguishing feature of TeSAN is its feature-wise temporal self-attention block, which enables the model to capture pairwise temporal dependencies between medical concepts. Given input embeddings $|\mathcal{C}| \approx 7,000$4 and all pairwise day-differences $|\mathcal{C}| \approx 7,000$5, the compatibility function for a concept pair $|\mathcal{C}| \approx 7,000$6 is

$|\mathcal{C}| \approx 7,000$7

where $|\mathcal{C}| \approx 7,000$8 denotes an activation function such as ReLU. This yields a $|\mathcal{C}| \approx 7,000$9-dimensional vector score for each $d=100$0, allowing for feature-wise (rather than scalar) attention.

For each query position $d=100$1, feature-wise softmax normalization is applied across all “keys” $d=100$2:

$d=100$3

This enables each embedding dimension to attend differently to temporal/contextual neighbors. The attended representation is:

$d=100$4

A gated fusion mechanism—using parameters $d=100$5—combines $d=100$6 and the original $d=100$7 to yield $d=100$8, the output embedding at position $d=100$9:

$E \in \mathbb{R}^{n_{\text{days}} \times d}$0

Stacking $E \in \mathbb{R}^{n_{\text{days}} \times d}$1 across all $E \in \mathbb{R}^{n_{\text{days}} \times d}$2 yields $E \in \mathbb{R}^{n_{\text{days}} \times d}$3.

3. Learning Objective and Optimization

TeSAN’s training objective adapts skip-gram with negative sampling to the EHR context. For each target position $E \in \mathbb{R}^{n_{\text{days}} \times d}$4, an attention-pooled context vector $E \in \mathbb{R}^{n_{\text{days}} \times d}$5 is computed from $E \in \mathbb{R}^{n_{\text{days}} \times d}$6, constituting the “context” for the target concept $E \in \mathbb{R}^{n_{\text{days}} \times d}$7. The loss is given by

$E \in \mathbb{R}^{n_{\text{days}} \times d}$8

where $E \in \mathbb{R}^{n_{\text{days}} \times d}$9 denotes the sigmoid, and $n_{\text{days}}$0 is a unigram noise distribution with $n_{\text{days}}$1 negative samples. No additional regularization or auxiliary losses are used. Training is performed for 30 epochs on MIMIC-III and 20 epochs on CMS, with batch size only specified for downstream GRU (128). Preprocessing includes the exclusion of codes with frequency $n_{\text{days}}$2 and, for CMS, patients with fewer than four visits; no imputation or time-binning is performed.

4. Empirical Evaluation

TeSAN is benchmarked against five state-of-the-art embedding methods: CBOW, Skip-gram, GloVe, med2vec, and MCE. The model is evaluated on clustering (using k-means, NMI), nearest-neighbor search (P@1), and inpatient mortality prediction (PR-AUC/ROC-AUC via GRU using visit-level aggregated embeddings).

Dataset/Task	Metric	Baseline Best	TeSAN Best
MIMIC-III, Clustering	NMI (ICD/CCS) %	26.02 / 53.09 (Skip-gram)	32.84 / 58.33
CMS, Clustering	NMI (ICD/CCS) %	8.52 / 41.70 (CBOW)	14.69 / 45.63
MIMIC-III, NN Search	P@1 (ICD/CCS) %	54.3 / 35.2 (CBOW)	66.1 / 43.8
CMS, NN Search	P@1 (ICD/CCS) %	34.3 / 16.4 (CBOW)	47.8 / 24.9
MIMIC-III, Mortality	PR-AUC / ROC-AUC (GRU)	0.5276 / 0.7785 (Sg+)	0.5544 / 0.8064

TeSAN displays consistently superior performance across all reported metrics, particularly notable in clustering and nearest-neighbor search. TeSAN’s clustering and retrieval metrics are both higher and more robust when varying the skip window length ($n_{\text{days}}$3). In mortality prediction, TeSAN+ (with GRU) achieves the highest PR-AUC and ROC-AUC on MIMIC-III. No tests for statistical significance are reported.

5. Analysis, Discussion, and Limitations

The explicit modeling of temporal intervals as $n_{\text{days}}$4-dimensional continuous embeddings, and their integration via multi-dimensional gated self-attention, enables TeSAN to encode both the co-occurrence of medical concepts and the temporal gaps between events. This dual encoding provides a richer representation space than methods focused solely on context (CBOW, Skip-gram), global co-occurrence (GloVe), hierarchical structure (med2vec), or prior forms of local attention (MCE).

Ablation studies demonstrate that both the temporal interval embedding and the multi-dimensional attention mechanism are essential; removing either degrades performance on clustering and retrieval. However, several limitations remain: only a single-head attention block is used, restricting the expressiveness for long-range dependencies beyond the local context window; training hyper-parameters (optimizer, learning-rate) are omitted; GPU specifications and statistical significance tests are not provided.

Potential extensions include development of multi-head attention variants, adaptive or hierarchical skip-window mechanisms, and integration with structured clinical ontologies for further semantic enrichment. These directions could enhance modeling of complex, heterogeneous temporal patterns in longitudinal EHRs (Peng et al., 2019).