Merge pull request #48 from yuxinc17/patch-models

Fixed off by one and other issues (affecting models, log-likelihood computation and predictions).

easy_tpp/config_factory/model_config.py

@@ -202,9 +202,9 @@ def __init__(self, **kwargs):
  self.time_emb_size = kwargs.get('time_emb_size', 16)
  self.num_layers = kwargs.get('num_layers', 2)
  self.num_heads = kwargs.get('num_heads', 2)
- self.mc_num_sample_per_step = kwargs.get('mc_num_sample_per_step', 20)
  self.sharing_param_layer = kwargs.get('sharing_param_layer', False)
- self.loss_integral_num_sample_per_step = kwargs.get('loss_integral_num_sample_per_step', 20)
+ self.use_mc_samples = kwargs.get('use_mc_samples', True) # if using MC samples in computing log-likelihood
+ self.loss_integral_num_sample_per_step = kwargs.get('loss_integral_num_sample_per_step', 20) # mc_num_sample_per_step
  self.dropout_rate = kwargs.get('dropout_rate', 0.0)
  self.use_ln = kwargs.get('use_ln', False)
  self.thinning = ThinningConfig.parse_from_yaml_config(kwargs.get('thinning'))
@@ -227,7 +227,6 @@ def get_yaml_config(self):
  'hidden_size': self.hidden_size,
  'time_emb_size': self.time_emb_size,
  'num_layers': self.num_layers,
- 'mc_num_sample_per_step': self.mc_num_sample_per_step,
  'sharing_param_layer': self.sharing_param_layer,
  'loss_integral_num_sample_per_step': self.loss_integral_num_sample_per_step,
  'dropout_rate': self.dropout_rate,
@@ -265,7 +264,6 @@ def copy(self):
  hidden_size=self.hidden_size,
  time_emb_size=self.time_emb_size,
  num_layers=self.num_layers,
- mc_num_sample_per_step=self.mc_num_sample_per_step,
  sharing_param_layer=self.sharing_param_layer,
  loss_integral_num_sample_per_step=self.loss_integral_num_sample_per_step,
  dropout_rate=self.dropout_rate,

easy_tpp/model/torch_model/torch_attnhp.py

@@ -3,7 +3,7 @@
 import torch
 from torch import nn
 
-from easy_tpp.model.torch_model.torch_baselayer import EncoderLayer, MultiHeadAttention
+from easy_tpp.model.torch_model.torch_baselayer import EncoderLayer, MultiHeadAttention, ScaledSoftplus
 from easy_tpp.model.torch_model.torch_basemodel import TorchBaseModel
 
 
@@ -52,7 +52,7 @@ def __init__(self, model_config):
  if self.use_norm:
  self.norm = nn.LayerNorm(self.d_model)
  self.inten_linear = nn.Linear(self.d_model * self.n_head, self.num_event_types)
- self.softplus = nn.Softplus()
+ self.softplus = ScaledSoftplus(self.num_event_types) # learnable mark-specific beta
  self.layer_event_emb = nn.Linear(self.d_model + self.d_time, self.d_model)
  self.layer_intensity = nn.Sequential(self.inten_linear, self.softplus)
  self.eps = torch.finfo(torch.float32).eps
@@ -151,7 +151,7 @@ def make_layer_mask(self, attention_mask):
  a diagonal matrix, [batch_size, seq_len, seq_len]
  """
  # [batch_size, seq_len, seq_len]
- layer_mask = (torch.eye(attention_mask.size(1)) < 1).unsqueeze(0).expand_as(attention_mask).to(attention_mask.device)
+ layer_mask = (torch.eye(attention_mask.size(1), device=self.device) < 1).unsqueeze(0).expand_as(attention_mask)
  return layer_mask
 
  def make_combined_att_mask(self, attention_mask, layer_mask):
@@ -205,11 +205,11 @@ def loglike_loss(self, batch):
  Returns:
  list: loglike loss, num events.
  """
- time_seqs, time_delta_seqs, type_seqs, batch_non_pad_mask, attention_mask, type_mask = batch
+ time_seqs, time_delta_seqs, type_seqs, batch_non_pad_mask, attention_mask = batch
  # 1. compute event-loglik
  # the prediction of last event has no label, so we proceed to the last but one
  # att mask => diag is False, not mask.
- enc_out = self.forward(time_seqs[:, :-1], type_seqs[:, :-1], attention_mask[:, 1:, :-1], time_seqs[:, 1:])
+ enc_out = self.forward(time_seqs[:, :-1], type_seqs[:, :-1], attention_mask[:, :-1, :-1], time_seqs[:, 1:])
  # [batch_size, seq_len, num_event_types]
  lambda_at_event = self.layer_intensity(enc_out)
 
@@ -227,17 +227,16 @@ def loglike_loss(self, batch):
  time_delta_seqs[:, :-1], # not used
  type_seqs[:, :-1],
  sample_times,
- attention_mask=attention_mask[:, 1:, :-1])
+ attention_mask=attention_mask[:, :-1, :-1])
 
  event_ll, non_event_ll, num_events = self.compute_loglikelihood(lambda_at_event=lambda_at_event,
  lambdas_loss_samples=lambda_t_sample,
  time_delta_seq=time_delta_seqs[:, 1:],
  seq_mask=batch_non_pad_mask[:, 1:],
- lambda_type_mask=type_mask[:, 1:])
+ type_seq=type_seqs[:, 1:])
 
- # return enc_inten to compute accuracy
+ # compute loss to minimize
  loss = - (event_ll - non_event_ll).sum()
-
  return loss, num_events
 
  def compute_states_at_sample_times(self,
@@ -285,7 +284,7 @@ def compute_states_at_sample_times(self,
  encoder_output = encoder_output.permute((1, 2, 0, 3))
  return encoder_output
 
- def compute_intensities_at_sample_times(self, time_seqs, time_delta_seqs, type_seqs, sample_times, **kwargs):
+ def compute_intensities_at_sample_times(self, time_seqs, time_delta_seqs, type_seqs, sample_dtimes, **kwargs):
  """Compute the intensity at sampled times.
 
  Args:
@@ -302,17 +301,17 @@ def compute_intensities_at_sample_times(self, time_seqs, time_delta_seqs, type_s
 
  if attention_mask is None:
  batch_size, seq_len = time_seqs.size()
- attention_mask = torch.triu(torch.ones(seq_len, seq_len), diagonal=1).unsqueeze(0)
- attention_mask = attention_mask.expand(batch_size, -1, -1).to(torch.bool).to(time_seqs.device)
+ attention_mask = torch.triu(torch.ones(seq_len, seq_len), diagonal=1).unsqueeze(0).to(type_seqs.device)
+ attention_mask = attention_mask.expand(batch_size, -1, -1).to(torch.bool)
 
- if sample_times.size()[1] < time_seqs.size()[1]:
+ if sample_dtimes.size()[1] < time_seqs.size()[1]:
  # we pass sample_dtimes for last time step here
  # we do a temp solution
  # [batch_size, seq_len, num_samples]
- sample_times = time_seqs[:, :, None] + torch.tile(sample_times, [1, time_seqs.size()[1], 1])
+ sample_dtimes = time_seqs[:, :, None] + torch.tile(sample_dtimes, [1, time_seqs.size()[1], 1])
 
  # [batch_size, seq_len, num_samples, hidden_size]
- encoder_output = self.compute_states_at_sample_times(time_seqs, type_seqs, attention_mask, sample_times)
+ encoder_output = self.compute_states_at_sample_times(time_seqs, type_seqs, attention_mask, sample_dtimes)
 
  if compute_last_step_only:
  lambdas = self.layer_intensity(encoder_output[:, -1:, :, :])

easy_tpp/model/torch_model/torch_baselayer.py

@@ -16,6 +16,28 @@ def attention(query, key, value, mask=None, dropout=None):
  return torch.matmul(p_attn, value), p_attn
 
 
+class ScaledSoftplus(nn.Module):
+ '''
+ Use different beta for mark-specific intensities
+ '''
+ def __init__(self, num_marks, threshold=20.):
+ super(ScaledSoftplus, self).__init__()
+ self.threshold = threshold
+ self.log_beta = nn.Parameter(torch.zeros(num_marks), requires_grad=True) # [num_marks]
+
+ def forward(self, x):
+ '''
+ :param x: [..., num_marks]
+ '''
+ beta = self.log_beta.exp()
+ beta_x = beta * x
+ return torch.where(
+ beta_x <= self.threshold,
+ torch.log1p(beta_x.clamp(max=math.log(1e5)).exp()) / beta,
+ x, # if above threshold, then the transform is effectively linear
+ )
+
+
 class MultiHeadAttention(nn.Module):
  def __init__(self, n_head, d_input, d_model, dropout=0.1, output_linear=False):
  super(MultiHeadAttention, self).__init__()

easy_tpp/model/torch_model/torch_basemodel.py

@@ -2,6 +2,7 @@
 
 import torch
 from torch import nn
+from torch.nn import functional as F
 
 from easy_tpp.model.torch_model.torch_thinning import EventSampler
 from easy_tpp.utils import set_device
@@ -29,6 +30,7 @@ def __init__(self, model_config):
  self.gen_config = model_config.thinning
  self.event_sampler = None
  self.device = set_device(model_config.gpu)
+ self.use_mc_samples = model_config.use_mc_samples
 
  self.to(self.device)
 
@@ -81,8 +83,7 @@ def get_logits_at_last_step(logits, batch_non_pad_mask, sample_len=None):
  last_logits = torch.gather(logits, dim=1, index=select_index).squeeze(1)
  return last_logits
 
- def compute_loglikelihood(self, time_delta_seq, lambda_at_event, lambdas_loss_samples, seq_mask,
- lambda_type_mask):
+ def compute_loglikelihood(self, time_delta_seq, lambda_at_event, lambdas_loss_samples, seq_mask, type_seq):
  """Compute the loglikelihood of the event sequence based on Equation (8) of NHP paper.
 
  Args:
@@ -92,38 +93,35 @@ def compute_loglikelihood(self, time_delta_seq, lambda_at_event, lambdas_loss_sa
  lambdas_loss_samples (tensor): [batch_size, seq_len, num_sample, num_event_types],
  intensity at sampling times.
  seq_mask (tensor): [batch_size, seq_len], sequence mask vector to mask the padded events.
- lambda_type_mask (tensor): [batch_size, seq_len, num_event_types], type mask matrix to mask the
- padded event types.
+ type_seq (tensor): [batch_size, seq_len], sequence of mark ids, with padded events having a mark of self.pad_token_id
 
  Returns:
  tuple: event loglike, non-event loglike, intensity at event with padding events masked
  """
 
- # Sum of lambda over every type and every event point
- # [batch_size, seq_len]
- event_lambdas = torch.sum(lambda_at_event * lambda_type_mask, dim=-1) + self.eps
-
- # mask the pad event
- event_lambdas = event_lambdas.masked_fill_(~seq_mask, 1.0)
-
- # [batch_size, seq_len)
- event_ll = torch.log(event_lambdas)
+ # First, add an epsilon to every marked intensity for stability
+ lambda_at_event = lambda_at_event + self.eps
+ lambdas_loss_samples = lambdas_loss_samples + self.eps
 
- # Compute the big lambda integral in equation (8) of NHP paper
- # 1 - take num_mc_sample rand points in each event interval
- # 2 - compute its lambda value for every sample point
- # 3 - take average of these sample points
- # 4 - times the interval length
+ log_marked_event_lambdas = lambda_at_event.log()
+ total_sampled_lambdas = lambdas_loss_samples.sum(dim=-1)
 
- # [batch_size, seq_len, n_loss_sample]
- lambdas_total_samples = lambdas_loss_samples.sum(dim=-1)
+ # Compute event LL - [batch_size, seq_len]
+ event_ll = -F.nll_loss(
+ log_marked_event_lambdas.permute(0, 2, 1), # mark dimension needs to come second, not third to match nll_loss specs
+ target=type_seq,
+ ignore_index=self.pad_token_id, # Padded events have a pad_token_id as a value
+ reduction='none', # Does not aggregate, and replaces what would have been the log(marked intensity) with 0.
+ )
 
- # interval_integral - [batch_size, seq_len]
+ # Compute non-event LL [batch_size, seq_len]
  # interval_integral = length_interval * average of sampled lambda(t)
- non_event_ll = lambdas_total_samples.mean(dim=-1) * time_delta_seq * seq_mask
+ if self.use_mc_samples:
+ non_event_ll = total_sampled_lambdas.mean(dim=-1) * time_delta_seq * seq_mask
+ else: # Use trapezoid rule
+ non_event_ll = 0.5 * (total_sampled_lambdas[..., 1:] + total_sampled_lambdas[..., :-1]).mean(dim=-1) * time_delta_seq * seq_mask
 
  num_events = torch.masked_select(event_ll, event_ll.ne(0.0)).size()[0]
-
  return event_ll, non_event_ll, num_events
 
  def make_dtime_loss_samples(self, time_delta_seq):
@@ -160,37 +158,47 @@ def predict_one_step_at_every_event(self, batch):
  Returns:
  tuple: tensors of dtime and type prediction, [batch_size, seq_len].
  """
- time_seq, time_delta_seq, event_seq, batch_non_pad_mask, _, type_mask = batch
+ time_seq, time_delta_seq, event_seq, batch_non_pad_mask, _ = batch
 
  # remove the last event, as the prediction based on the last event has no label
- # time_delta_seq should start from 1, because the first one is zero
- time_seq, time_delta_seq, event_seq = time_seq[:, :-1], time_delta_seq[:, 1:], event_seq[:, :-1]
+ # note: the first dts is 0
+ # [batch_size, seq_len]
+ time_seq, time_delta_seq, event_seq = time_seq[:, :-1], time_delta_seq[:, :-1], event_seq[:, :-1]
 
  # [batch_size, seq_len]
- dtime_boundary = time_delta_seq + self.event_sampler.dtime_max
+ dtime_boundary = torch.max(time_delta_seq * self.event_sampler.dtime_max,
+ time_delta_seq + self.event_sampler.dtime_max)
 
  # [batch_size, seq_len, num_sample]
  accepted_dtimes, weights = self.event_sampler.draw_next_time_one_step(time_seq,
  time_delta_seq,
  event_seq,
  dtime_boundary,
- self.compute_intensities_at_sample_times)
+ self.compute_intensities_at_sample_times,
+ compute_last_step_only=False) # make it explicit
 
- # [batch_size, seq_len]
- dtimes_pred = torch.sum(accepted_dtimes * weights, dim=-1)
-
- # [batch_size, seq_len, 1, event_num]
+ # We should condition on each accepted time to sample event mark, but not conditioned on the expected event time.
+ # 1. Use all accepted_dtimes to get intensity.
+ # [batch_size, seq_len, num_sample, num_marks]
  intensities_at_times = self.compute_intensities_at_sample_times(time_seq,
  time_delta_seq,
  event_seq,
- dtimes_pred[:, :, None],
- max_steps=event_seq.size()[1])
+ accepted_dtimes)
 
- # [batch_size, seq_len, event_num]
- intensities_at_times = intensities_at_times.squeeze(dim=-2)
+ # 2. Normalize the intensity over last dim and then compute the weighted sum over the `num_sample` dimension.
+ # Each of the last dimension is a categorical distribution over all marks.
+ # [batch_size, seq_len, num_sample, num_marks]
+ intensities_normalized = intensities_at_times / intensities_at_times.sum(dim=-1, keepdim=True)
 
- types_pred = torch.argmax(intensities_at_times, dim=-1)
+ # 3. Compute weighted sum of distributions and then take argmax.
+ # [batch_size, seq_len, num_marks]
+ intensities_weighted = torch.einsum('...s,...sm->...m', weights, intensities_normalized)
 
+ # [batch_size, seq_len]
+ types_pred = torch.argmax(intensities_weighted, dim=-1)
+
+ # [batch_size, seq_len]
+ dtimes_pred = torch.sum(accepted_dtimes * weights, dim=-1) # compute the expected next event time
  return dtimes_pred, types_pred
 
  def predict_multi_step_since_last_event(self, batch, forward=False):