I am trying to replicate something that the authors in the paper, A Closer Look at Memorization in Deep Networks, have proposed in the Section 3.2: Loss-Sensitivity in Real vs. Random Data which says:

mem-loss-sen-ques501×570 80.4 KB

It’s not clear from the discussion in Maclaurin et al. (2015)  or their github code as to what I want to do which is essentially to backprop through training procedure.
A Digression (?): Although I am not quite familiar with Recurrent Neural Networks (RNNs) and haven’t used them in my work, I think it looks a bit similar to (but not exactly the same as) Truncated Backpropagation Through Time (TBPTT). I went through the discussion here but I am still not sure (in case it is indeed similar to TBPTT) how to implement what I need and have mentioned above. For instance, what changes would I need to make here:
# perform k1-step unrolling through the training procedure

def train(data_loader, optimizer, model):
   model.train()
   cnt = 0
   # store the loss sensitivity values in a Tensor
   loss_sensitivity = torch.zeros((len(data_loader.dataset), data_loader.tensors[0].shape[1], num_classes))

   idx_last_k1 = [] # keeps track of indices of (K-k1)th batch where K is the current grad. step
   for batch_id, (x,y, idx) in enumerate(data_loader):

      idx_last_k1 = '''store (K-k1)th batch indices where K is the current step'''
      x.requires_grad = True

      if batch_id % k1 == 0:
         # store the corresponding loss sensitivity values
         loss_sensitivity[idx,:] = data_loader.tensors[0][idx_last_k1,:].grad.clone().detach()

         # We compute gradients (for sensitivity) 
         # only w.r.t. the (M-k1)th batch
         data_loader.tensors[0][idx_last_k1, :].detach_()
         idx_last_k1 = []
      ...
      # some more usual training operations
      ...
      output = model(x)
      loss = cce_loss(output, y)
      optimizer.zero_grad()
 
      loss.backward(retain_graph = True)
      optimizer.step()
      ...

@albanD Any idea about this?