NaN-value attributions when using DeepLIFT during fp16 training

Hi! I’m trying to compute attributions for BigBird via DeepLIFT during fp16 training, but I’m getting NaN for many of the attribution values. However, if I do fp32 training, none of the attribution values are NaN. I guess the fp16 precision is truncating some near-zero intermediate values to zero, resulting in division-by-zero. Would appreciate any ideas on how to work around this issue!

For reference, the inputs and baselines parameters for my DeepLIFT attribute function call are constructed as follows:

inputs

• Feed input ids into BigBird word embedding layer to get input embeddings.
• Use input embeddings as inputs for DeepLIFT.

baselines

• Create baseline ids by setting all non-special tokens in the input ids to PAD.
• Feed baseline ids into BigBird word embedding layer to get baseline embeddings.
• Use baseline embeddings as baselines for DeepLIFT.

[aarzchan]

Hi @bilalsal, here is a Colab notebook demonstrating the error. I shared it with your FB email.

[aarzchan]

Hi @bilalsal and @NarineK, thanks for the fix! I don’t have any further remarks/questions, so I’m closing the issue.