Why custom forward function is called 5 times in LayerIntegratedGradients?

Captum
1

Hello,
I use captum LayerIntegratedGradients class for attribution method on my bert classification model.
I managed to get the attributions by following the great source.

However, the IG method take some times like 5 times more than inference step.
So I look into it, and find out that It calls the custom forward function 5 times.

def __custom_forward(self, input_ids, mask):
      logits = self.model(input_ids, mask)
      log.info(logits.size())
      log.info(' Custom forward function has been called !!!! ')
      return logits.max(-1).values

Above code is the custom forward function, and Below is the result.

lig2
lig2791×184 2.81 KB

I get that the (10, 3) shape tensor is because of that I choose n_step as 10. So it should be the delta shape. And first (1,3) is for the logits.

self.attr_method = LayerIntegratedGradients(self.__custom_forward, self.model.l1.embeddings)

attributions, delta = self.attr_method.attribute(inputs=ids,
                      baselines=ref_input_ids,
                      additional_forward_args=(mask),
                      n_steps=10,
                      return_convergence_delta=True) # B, input_ids, embed_dim

Then what the other left three of (1, 3) is for?

Thanks in advance!