Hey, thanks for the update. Unfortunately this is wrong. Sorry, I had a thought about this and the problem here is that multiple inputs (given one output) vary. So for instance, let's have a look at the tensorflow implementation for multiplicative attribution and deduct it from the definition of shapley values:

$\phi_{i} = \sum_{S \subseteq N \setminus {i}} \frac{|S|!(|N|-|S|-1)!}{|N|!} \left( f(S \cup {i}) - f(S) \right)$

Now let's assume we have just $x, y$ (and therefore $b_{x}, b_{y}$ as baselines).
We are looking for a multiplication (or any function that takes two arguments):
$f(x, y) = x \cdot y$
Then this formula becomes:
$\phi_{x} = \frac{1}{2} (f(x, y) - f(b_{x}, y)) + \frac{1}{2} (f(x, b_{y}) - f(b_{x}, b_{y}))$
We would need to do something similar for the layer norm BUT this is computationally expensive, so we might be able to use the approach that the tensorflow implementation takes for softmax (yes, pytorch impl is wrong here too!). I don't fully understand this yet, but doesn't seem too complicated.

Sorry, I just figured this out too late. Really appreciate your effort. Let me know how you want to proceed with this.

Ah! It was too easy to be true. So, if I understand this correctly, in the original DeepLIFT paper they have the rescale rule which actually only applies to a single input. That's what's implemented in the nonlinear_1d function? And that's why this isn't valid for LayerNorm which takes several inputs. So we need to use the full equation to derive the proper function. Is this right? If so I think I understand and can try to see how to translate the way it's done in tensorflow to pytorch.

P.S. Does this mean that the softmax implementation for pytorch also needs to be fixed?

@JulioHC00 yes, that's right. And also we would need to fix the implementation for pytorch's softmax attribution. I believe that once we understand how the tensorflow softmax is implemented, we can apply this to pytorch LayerNorm and softmax as well.

@CloseChoice I guess the only part I'm confused about is how does this fit with the op_handlers? Don't these handle the gradients in the backward propagation? How does this relate to implementing the shap values
equation?

@CloseChoice When I integrate LayerNorm into a testcase and use the op_handler['LayerNorm'] = softmax, the testcase passes. Do we need to create a testcase that does show how Softmax/LayerNorm are broken? What would be in such a testcase?