Hello, I am trying to implement this loss function taken from Section 2.1 of Right for the Right Reasons: Training Differentiable Models by Constraining their Explanations  (Ross, et al., 2017)

The first and third term are the Cross-entropy loss and L2 regularization, respectively and are already implemented in Pytorch. The matrix A is a binary mask with dims (Num of samples, W, H, #Color channel). The new loss can be called like this:
optimizer = optim.SGD(model.parameters(), lr=0.01, weight_decay=0.01)  # L2 is weight decay

for batch in batches:
   x, target = batch
   logits = model(x)
   loss = explanation_loss(logits, target, l_one, A[index])
   loss.backward()
   optimizer.zero_grad()
   optimizer.step()


I am confused on how to implement the 2nd term (right reasons).
Specifically:


if my pytorch model outputs logits, I must first convert them to probabilities with softmax and then into a one-hot encoding to take the log, right?


How do I implement the input gradient here? (partial derivative w.r.t. x in the image)


Below is my pseudo-code that I want to make functional. Here is the original code that I’m having trouble to understand 1 (In the function callled “objective”). Thank you!
def explanation_loss(logits, target, l_one, mask_matrix):
    ce = nn.CrossEntropyLoss()(logits, target) 
    if l_one==0:
      return ce
    pred = F.softmax(logits, dim=1)
    label_one_hot = torch.nn.functional.one_hot(target, num_classes).float().to(device)
    label_one_hot = torch.clamp(label_one_hot, min=1e-4, max=1.0)  # prevent log(0)

    # missing the input gradient
    right_reasons = l_one* (torch.sum(mask_matrix * torch.log(label_one_hot), dim=1)).mean()
    return ce + right_reasons