3.2 Architecture of STraTS
The architecture of STraTS is illustrated in Figure 3. Unlike most of the existing approaches which
take a time-series matrix as input, STraTS defines its input as a set of observation triplets. Each
observation triplet in the input is embedded using the Initial Triplet Embedding module. The
initial triplet embeddings are then passed through a Contextual Triplet Embedding module which
utilizes the Transfomer architecture to encode the context for each triplet. The Fusion Self-attention
module then combines these contextual embeddings via self-attention mechanism to generate an
embedding for the input time-series which is concatenated with demographics embedding and
passed through a feed-forward network to make the final prediction. The notations used in the
paper are summarized in Table 1.
3.2.1 Initial Triplet Embedding. Given an input time-series T = {(𝑡𝑖
, 𝑓𝑖
, 𝑣𝑖)}𝑛
𝑖=1
, the initial embedding
for the 𝑖
𝑡ℎ triplet ei ∈ R
𝑑
is computed by summing the following component embeddings: (i) Feature
embedding e
f
i
∈ R
𝑑
, (ii) Value embedding e
v
i
∈ R
𝑑
, and (iii) Time embedding e
t
i
∈ R
𝑑
. In other words,
ei = e
f
i
+ e
v
i
+ e
t
i
∈ R
𝑑
. Feature embeddings e
f
(·) are obtained from a simple lookup table similar to
word embeddings. Since feature values and times are continuous unlike feature names which are
categorical objects, we cannot use a lookup table to embed these continuous values unless they
are categorized. Some researchers [30, 32] have used sinusoidal encodings to embed continuous
ACM Trans. Knowl. Discov. Data., Vol. 1, No. 1, Article . Publication date: February 2022.
Self-Supervised Transformer for Sparse and Irregularly Sampled Multivariate Clinical Time-Series 7
Table 1. Notations used in this paper.
Notation Definition
𝑁 # Time-series for target task
𝑁
′
# Time-series for forecasting task
d ∈ R
𝐷 Demographics vector
F Set of clinical variables
𝑡𝑖 ∈ R≥0 Time of 𝑖
𝑡ℎ observation
𝑓𝑖 ∈ F Variable of 𝑖
𝑡ℎ observation
𝑣𝑖 ∈ R Value of 𝑖
𝑡ℎ observation
(𝑡, 𝑓 , 𝑣) Observation triplet
T = {(𝑡𝑖
, 𝑓𝑖
, 𝑣𝑖)}𝑛
𝑖=1 Multivariate time-series
𝑦,𝑦˜ ∈ {0, 1} True and predicted outputs for target task
z, z˜ ∈ R
|F | True and predicted outputs for forecasting task
m ∈ {0, 1}
|F | Forecast mask
e
𝑡
𝑖
, e
𝑣
𝑖
∈ R
𝑑 CVE for time and value
e
𝑓
𝑖
∈ R
𝑑 Variable embedding
e𝑖 ∈ R
𝑑
Initial triplet embedding
e
𝑇 ∈ R
𝑑 Time-series embedding
e
𝑑 ∈ R
𝑑 Demographics embedding
values. We propose a novel continuous value embedding (CVE) technique using a one-to-many
Feed-Forward Network (FFN) with learnable parameters i.e., e
v
i
= 𝐹 𝐹𝑁𝑣
(𝑣𝑖), and e
t
i
= 𝐹 𝐹𝑁𝑡
(𝑡𝑖).
Both FFNs have one input neuron and 𝑑 output neurons and a single hidden layer with ⌊
√
𝑑⌋
neurons and 𝑡𝑎𝑛ℎ(·) activation. They are of the form 𝐹 𝐹𝑁 (𝑥) = 𝑈 𝑡𝑎𝑛ℎ(𝑊 𝑥 + 𝑏) where the
dimensions of weights {𝑊 , 𝑏,𝑈 } can be inferred from the size of hidden and output layers of the
FFN. Unlike sinusoidal encodings with fixed frequencies, this technique offers more flexibility
by allowing end-to-end learning of continuous value and time embeddings without the need to
categorize them.
3.2.2 Contextual Triplet Embedding. The initial triplet embeddings {e1, ..., e𝑛} are then passed
through a Transformer architecture [30] with 𝑀 blocks, each containing a Multi-Head Attention
(MHA) layer with ℎ attention heads and a FFN with one hidden layer. Each block takes 𝑛 input
embeddings E ∈ R
𝑛×𝑑
and outputs the corresponding 𝑛 output embeddings C ∈ R
𝑛×𝑑
that capture
the contextual information. MHA layers use multiple attention heads to attend to information
contained in different embedding projections in parallel. The computations of the MHA layer are
given by
H𝑗 = 𝑠𝑜 𝑓 𝑡𝑚𝑎𝑥 
(EW𝑞
𝑗
) (EW𝑘
𝑗
)
𝑇
√︁
𝑑/ℎ

(EW𝑣
𝑗
) 𝑗 = 1, ..., ℎ (1)
𝑀𝐻𝐴(E) = (H1 ◦ ... ◦ Hℎ) W𝑐 (2)
Each head projects the input embeddings into query, key, and value subspaces using matrices
{W
𝑞
𝑗
, W𝑘
𝑗
, W𝑣
𝑗
} ⊂ R
𝑑×𝑑ℎ . The queries and keys are then used to compute the attention weights
which are used to compute weighted averages of value (different from value in observation triplet)
vectors. Finally, the outputs of all heads are concatenated and projected to original dimension with
ACM Trans. Knowl. Discov. Data., Vol. 1, No. 1, Article . Publication date: February 2022.
8 Sindhu Tipirneni and Chandan K Reddy
Fig. 3. The overall architecture of the proposed STraTS model. The Input Triplet Embedding module embeds
each observation triplet, the Contextual Triplet Embedding module encodes contextual information for the
triplets, the Fusion Self-Attention module computes time-series embedding which is concatenated with
demographics embedding and passed through a dense layer to generate predictions for target and selfsupervision (forecasting) tasks.
W𝑐 ∈ R
ℎ𝑑ℎ×𝑑
. The FFN layer takes the form
F(X) = 𝑅𝑒𝐿𝑈 (XW𝑓
1
+ b
𝑓
1
) W
𝑓
2
+ b
𝑓
2
(3)
with weights Wf
1
∈ R
𝑑×2𝑑
, b
f
1
∈ R
2𝑑
, Wf
2
∈ R
2𝑑×𝑑
, b
f
2
∈ R
𝑑
. Dropout, residual connections, and
layer normalization are added for every MHA and FFN layer. Also, attention dropout randomly
masks out some positions in the attention matrix before the softmax computation during training.
The output of each block is fed as input to the succeeding one, and the output of the last block
gives the contextual triplet embeddings {c1, ..., c𝑛}.
3.2.3 Fusion Self-attention. After computing contextual embeddings using a Transformer, we
fuse them using a self-attention layer to compute time-series embedding e
𝑇 ∈ R
𝑑
. This layer first
computes attention weights {𝛼1, ..., 𝛼𝑛} by passing each contextual embedding through a FFN and
computing a softmax over all the FFN outputs.
𝑎𝑖 = u
𝑇
𝑎
𝑡𝑎𝑛ℎ(W𝑎c𝑖 + b𝑎) (4)
𝛼𝑖 =
𝑒𝑥𝑝(𝑎𝑖)
Í𝑛
𝑗=1
𝑒𝑥𝑝(𝑎𝑗)
∀𝑖 = 1, ..., 𝑛 (5)
W𝑎 ∈ R
𝑑𝑎×𝑑
, b𝑎 ∈ R
𝑑𝑎
, ua ∈ R
𝑑𝑎 are the weights of this attention network which has 𝑑𝑎 neurons in
the hidden layer. The time-series embedding is then computed as
e
𝑇 =
∑︁𝑛
𝑖=1
𝛼𝑖ci
(6)
ACM Trans. Knowl. Discov. Data., Vol. 1, No. 1, Article . Publication date: February 2022.
Self-Supervised Transformer for Sparse and Irregularly Sampled Multivariate Clinical Time-Series 9
3.2.4 Demographics Embedding. We realize that demographics can be encoded as triplets with
a default value for time. However, we found that the prediction models performed better in our
experiments when demographics are processed separately by passing d through a FFN as shown
below. The demographics embedding is thus obtained as
e
𝑑 = 𝑡𝑎𝑛ℎ(W𝑑
2
𝑡𝑎𝑛ℎ(W𝑑
1
d + b
𝑑
1
) + b
𝑑
2
) ∈ R
𝑑
(7)
where the hidden layer has a dimension of 2𝑑.
3.2.5 Prediction Head. The final prediction for target task is obtained by passing the concatenation
of demographics and time-series embeddings through a dense layer with weights w𝑇
𝑜 ∈ R
𝑑
, 𝑏𝑜 ∈ R
and sigmoid activation.
𝑦˜ = 𝑠𝑖𝑔𝑚𝑜𝑖𝑑 (w
𝑇
𝑜
[e
𝑑
◦ e
𝑇
] + 𝑏𝑜 ) (8)
The model is trained on the target task using cross-entropy loss.
3.2.6 Self-supervision. We experimented with both masking and forecasting as pretext tasks for
providing self-supervision and found that forecasting improved the results on target tasks. The
forecasting task uses the same architecture as the target task except for the prediction layer i.e.,
z˜ = W𝑠 [e
𝑑
◦ e
𝑇
] + b𝑠 ∈ R
|F | (9)
A masked MSE loss is used for training on the forecasting task to account for missing values in the
forecast outputs. Thus, the loss for self-supervision is given by
L𝑠𝑠 =
1
|𝑁 ′
|
𝑁
∑︁′
𝑘=1
∑︁
|F |
𝑗=1
𝑚
𝑘
𝑗
(z˜
𝑘
𝑗 − z
𝑘
𝑗
)
2
(10)
where m𝑘
𝑗
= 1 (or m𝑘
𝑗
= 0) if the ground truth forecast z
𝑘
𝑗
is available (or unavailable) for 𝑗
𝑡ℎ variable
in 𝑘
𝑡ℎ sample. The model is first pretrained on the self-supervision task and is then fine-tuned on
the target task.

from argparse import Namespace
from models import TimeSeriesModel
import torch.nn as nn
from transformers import BertModel, BertConfig
import torch
import torch.nn.functional as F
import numpy as np



class CVE(nn.Module):
    def __init__(self, args):
        super().__init__()
        int_dim = int(np.sqrt(args.hid_dim))
        self.W1 = nn.Parameter(torch.empty(1,int_dim), requires_grad=True)
        self.b1 = nn.Parameter(torch.zeros(int_dim), requires_grad=True)
        self.W2 = nn.Parameter(torch.empty(int_dim, args.hid_dim), requires_grad=True)
        nn.init.xavier_uniform_(self.W1)
        nn.init.xavier_uniform_(self.W2)
        self.activation = torch.tanh

    def forward(self, x):
        # x: bsz, max_len
        x = torch.unsqueeze(x, -1)
        x = torch.matmul(x,self.W1) + self.b1[None,None,:] # bsz,max_len,int_dim
        x = self.activation(x)
        x = torch.matmul(x,self.W2) #bsz,max_len,hid_dim
        return x
        

class FusionAtt(nn.Module):
    def __init__(self, args):
        super().__init__()
        int_dim = args.hid_dim
        self.W = nn.Parameter(torch.empty(args.hid_dim,int_dim), requires_grad=True)
        self.b = nn.Parameter(torch.zeros(int_dim), requires_grad=True)
        self.u = nn.Parameter(torch.empty(int_dim,1), requires_grad=True)
        nn.init.xavier_uniform_(self.W)
        nn.init.xavier_uniform_(self.u)
        self.activation = torch.tanh

    def forward(self, x, mask):
        # x: bsz, max_len, hid_dim
        att = torch.matmul(x,self.W) + self.b[None,None,:] # bsz,max_len,int_dim
        att = self.activation(att)
        att = torch.matmul(att,self.u)[:,:,0] # bsz,max_len
        att = att + (1-mask)*torch.finfo(att.dtype).min  
        att = torch.softmax(att,dim=-1) # bsz,max_len
        return att 
    

class Transformer(nn.Module):
    def __init__(self, args):
        super().__init__()
        self.N = args.num_layers
        self.d = args.hid_dim
        self.dff = self.d*2
        self.attention_dropout = args.attention_dropout
        self.dropout = args.dropout
        self.h = args.num_heads
        self.dk = self.d//self.h
        self.all_head_size = self.dk*self.h

        self.Wq = nn.Parameter(self.init_proj((self.N, self.h, self.d, self.dk)), requires_grad=True)
        self.Wk = nn.Parameter(self.init_proj((self.N, self.h, self.d, self.dk)), requires_grad=True)
        self.Wv = nn.Parameter(self.init_proj((self.N, self.h, self.d, self.dk)), requires_grad=True)
        self.Wo = nn.Parameter(self.init_proj((self.N, self.all_head_size, self.d)), requires_grad=True)
        self.W1 = nn.Parameter(self.init_proj((self.N, self.d, self.dff)), requires_grad=True)
        self.b1 = nn.Parameter(torch.zeros((self.N, 1, 1, self.dff)), requires_grad=True)
        self.W2 = nn.Parameter(self.init_proj((self.N, self.dff, self.d)), requires_grad=True)
        self.b2 = nn.Parameter(torch.zeros((self.N, 1, 1, self.d)), requires_grad=True)
        # self.layer_norm1 = nn.ModuleList([nn.LayerNorm(self.d) for i in range(self.N)])
        # self.layer_norm2 = nn.ModuleList([nn.LayerNorm(self.d) for i in range(self.N)])

    def init_proj(self, shape, gain=1):
        x = torch.rand(shape)
        fan_in_out = shape[-1]+shape[-2]
        scale = gain*np.sqrt(6/fan_in_out)
        x = x*2*scale - scale
        return x
    
    def forward(self, x, mask):
        # x: bsz, max_len, d
        # mask: bsz, max_len
        bsz, max_len, _ = x.size()
        mask = mask[:,:,None]*mask[:,None,:]
        mask = (1-mask)[:,None,:,:]*torch.finfo(x.dtype).min
        layer_mask = mask
        for i in range(self.N):
            # MHA
            q = torch.einsum('bld,hde->bhle', x, self.Wq[i]) # bsz,h,max_len,dk
            k = torch.einsum('bld,hde->bhle', x, self.Wk[i]) # bsz,h,max_len,dk
            v = torch.einsum('bld,hde->bhle', x, self.Wv[i]) # bsz,h,max_len,dk
            # x_for_qkv = x[:,None,:,:,None]
            # q = (x_for_qkv*self.Wq[i][None,:,None,:,:]).sum(dim=-2) # bsz,h,max_len,dk
            # k = (x_for_qkv*self.Wk[i][None,:,None,:,:]).sum(dim=-2) # bsz,h,max_len,dk
            # v = (x_for_qkv*self.Wv[i][None,:,None,:,:]).sum(dim=-2) # bsz,h,max_len,dk
            A = torch.einsum('bhle,bhke->bhlk', q, k) # bsz,h,max_len,max_len
            # A = (q[:,:,:,None,:]*k[:,:,None,:,:]).sum(dim=-1) # bsz,h,max_len,max_len
            if self.training:
                dropout_mask = (torch.rand_like(A)<self.attention_dropout
                                ).float()*torch.finfo(x.dtype).min
                layer_mask = mask+dropout_mask
            A = A+layer_mask
            A = torch.softmax(A, dim=-1)
            v = torch.einsum('bhkl,bhle->bkhe',A,v) # bsz,max_len,h,dk
            # v = (A[:,:,:,:,None]*v[:,:,None,:,:]).sum(dim=-2).transpose(1,2) # bsz,max_len,h,dk
            all_head_op = v.reshape((bsz, max_len, -1))
            all_head_op = torch.matmul(all_head_op, self.Wo[i])
            all_head_op = F.dropout(all_head_op, self.dropout, self.training)
            # Add+layernorm
            # x = self.layer_norm1[i](all_head_op+x) # bsz,max_len,d
            x = (all_head_op+x)/2
            # FFN
            ffn_op = torch.matmul(x,self.W1[i])+self.b1[i]
            ffn_op = F.gelu(ffn_op)
            ffn_op = torch.matmul(ffn_op,self.W2[i])+self.b2[i]
            ffn_op = F.dropout(ffn_op, self.dropout, self.training)
            # Add+layernorm
            # x = self.layer_norm2[i](ffn_op+x)
            x = (ffn_op+x)/2
        return x




class Strats(TimeSeriesModel):
    def __init__(self, args: Namespace):
        super().__init__(args)
        self.cve_time = CVE(args)
        self.cve_value = CVE(args)
        self.variable_emb = nn.Embedding(args.V, args.hid_dim)
        self.transformer = Transformer(args)
        self.fusion_att = FusionAtt(args)
        self.dropout = args.dropout
        self.V = args.V
        

    def forward(self, values, times, varis, obs_mask, demo,
                labels=None, forecast_values=None, forecast_mask=None):
        bsz, max_obs = values.size()
        device = values.device
        if self.training:
            with torch.no_grad():
                var_mask = (torch.rand((bsz,self.V), device=device)<=self.dropout).int()
                for v in range(self.V):
                    mask_pos = (varis==v).int()*var_mask[:,v:v+1]
                    obs_mask = obs_mask*(1-mask_pos)

        # demographics embedding
        demo_emb = self.demo_emb(demo) if self.args.model_type=='strats' \
                    else demo
        # initial triplet embedding
        time_emb = self.cve_time(times)
        value_emb = self.cve_value(values)
        # value_emb = 0
        # for i in range(self.args.V):
        #     value_emb = value_emb + self.cve_value[i](values) * (varis==i)
        vari_emb = self.variable_emb(varis)
        triplet_emb = time_emb+value_emb+vari_emb
        triplet_emb = F.dropout(triplet_emb, self.dropout, self.training)
        # contextual triplet emb
        contextual_emb = self.transformer(triplet_emb, obs_mask) 
        # fusion attention
        attention_weights = self.fusion_att(contextual_emb, obs_mask)[:,:,None]
        if self.args.model_type=='istrats':
            ts_emb = (triplet_emb*attention_weights).sum(dim=1)
        else:
            ts_emb = (contextual_emb*attention_weights).sum(dim=1)
        # concat demo and ts_emb
        ts_demo_emb = torch.cat((ts_emb, demo_emb), dim=-1)
        # prediction/loss
        if self.pretrain:
            return self.forecast_final(ts_demo_emb, forecast_values, forecast_mask)
        logits = self.binary_head(self.forecast_head(ts_demo_emb))[:,0] \
                    if self.finetune else self.binary_head(ts_demo_emb)[:,0]
        return self.binary_cls_final(logits, labels)

understand and explain what is going on to me, especially with the triplet and demo part.