Problem with inputs using Integrated Gradients and model with embedding layers

[karppmik]

I'm wondering how to use the Integrated Gradients with the following model as it has the embedding layers?

class TabularModel(nn.Module):
    def __init__(self, embedding_sizes, n_cont):
        super().__init__()
        self.embeddings = nn.ModuleList([nn.Embedding(categories, size) for categories, size in embedding_sizes])
        n_emb = sum(e.embedding_dim for e in self.embeddings)
        self.n_emb, self.n_cont = n_emb, n_cont
        self.lin1 = nn.Linear(self.n_emb + self.n_cont, 100)
        self.lin2 = nn.Linear(100, 50)
        self.lin3 = nn.Linear(50, 5)
        self.bn1 = nn.BatchNorm1d(self.n_cont)
        self.bn2 = nn.BatchNorm1d(100)
        self.bn3 = nn.BatchNorm1d(50)
        self.emb_drop = nn.Dropout(0.2)
        self.drops = nn.Dropout(0.1)

    def forward(self, x_cat, x_cont):
        x = [e(x_cat[:, i]) for i, e in enumerate(self.embeddings)]
        x = torch.cat(x, 1)
        x = self.emb_drop(x)
        x2 = self.bn1(x_cont)
        x = torch.cat([x, x2], 1)
        x = F.relu(self.lin1(x))
        x = self.drops(x)
        x = self.bn2(x)
        x = F.relu(self.lin2(x))
        x = self.drops(x)
        x = self.bn3(x)
        x = self.lin3(x)
        return x

The main issue is in passing the inputs to

ig = IntegratedGradients(model)
ig.attribute(inputs)

after which I get the error:

AssertionError: Baseline can be provided as a tensor for just one input and broadcasted to the batch or input and baseline must have the same shape or the baseline corresponding to each input tensor must be a scalar. Found baseline: tensor([[8.0000e+00, 1.0138e+03, 8.2027e+01,  ..., 1.4000e+01, 0.0000e+00,
         0.0000e+00],
        [8.0000e+00, 1.0161e+03, 8.7000e+01,  ..., 6.6700e+01, 0.0000e+00,
         0.0000e+00],
        [1.0000e+00, 1.0226e+03, 4.8000e+01,  ..., 2.4700e+01, 0.0000e+00,
         0.0000e+00],
        ...,
        [0.0000e+00, 1.0208e+03, 8.2000e+01,  ..., 1.2400e+01, 0.0000e+00,
         0.0000e+00],
        [7.0000e+00, 1.0142e+03, 9.8000e+01,  ..., 1.1000e+00, 1.0000e+00,
         0.0000e+00],
        [0.0000e+00, 1.0230e+03, 7.6000e+01,  ..., 3.6900e+01, 0.0000e+00,
         0.0000e+00]]) and input: tensor([[ 4,  8,  0,  3],
        [ 5,  8,  1, 15],
        [ 2, 13,  0, 29],
        ...,
        [ 5,  1,  0, 21],
        [ 0, 23,  0,  5],
        [ 5,  5,  4, 11]])

[karppmik]

Ok, I realised the inputs needs to be passed as a tuple:

ig.attribute(inputs=(test_set.features_cat, test_set.features_cont))

but now I get the error message:

RuntimeError: Expected tensor for argument #1 'indices' to have scalar type Long; but got torch.FloatTensor instead (while checking arguments for embedding)

The model's forward function works however with

outputs = model(test_set.features_cat, test_set.features_cont)

So, how should I pass the inputs to .attribute function?

[vivekmig]

Hi @karppmik, since gradients cannot be computed with respect to embedding inputs (indices), there are a few workarounds described here in our FAQ that would work. You can either replace your embedding with an InterpretableEmbedding module which takes the embedding directly as input, or use Layer attribution (LayerIntegratedGradients) to compute IG with respect to the embedding output. Take a look at the linked tutorials in the FAQ answer for examples of these two approaches. Hope this helps!

[karppmik]

Thanks for the reply @vivekmig, I will check those out!

[karppmik]

@vivekmig I went through the tutorials but I wasn't able to figure out how they would be applicable in this context. Would you be able to explain how replacing the model's embedding with InterpretableEmbedding works in this case?

[p16i]

Instead of using model for IntegratedGradients's first argument, maybe you can try do the following?

def wrapper_forward(input_tuple): 
    return model(input_tuple[0], input_tuple[1])

ig = IntegratedGradients(wrapper_forward)

[ovshake]

Hi @karppmik , I am facing the same issue, did you find a solution?

[vivekmig]

Hi @karppmik , for the model you shared, something like this might work:

# Compute embedding output for given inputs
embedding_inp = [e(test_set.features_cat[:, i]) for i, e in enumerate(model.embeddings)]
embeddings_cat = torch.cat(embedding_inp, 1)

# Replace embedding layers with IntepretableEmbeddings, 
# which just directly return the input (assuming embeddings are given as input)
interp_embeddings = []
for i in range(len(model.embeddings):
    interp_embeddings.append(configure_interpretable_embedding_layer(model, 'embeddings.' + str(i)))

# Compute attributions, providing input as embedding output directly
ig = IntegratedGradients(model)
ig.attribute(inputs=(embeddings_cat, test_set.features_cont))

# After attribution, return model to original state
for single_interp_embed in interp_embeddings:
    remove_interpretable_embedding_layer(model, single_interp_embed)

The essential idea is to change your model from taking input indices to taking the input embedding directly as input. The Interpretable Embedding simply makes the embedding layer just an identity transform returning the given input, so the model can directly take the embedding.
This approach doesn't necessarily work in all cases, particularly when the model performs operations on the embedding input / indices prior to passing through the embedding layer. From a quick look at your code, it doesn't seem like this is the case, and the idea of the above code should work. If a model forward function does perform operations on the embedding inputs, the Layer Attribution method would be cleaner.

[karppmik]

Hi @vivekmig, thanks a lot for the reply! I was able to do the attributions with that approach. How would I go about visualising the results?

Referring to this tutorial, https://captum.ai/tutorials/House_Prices_Regression_Interpret, I don't quite understand what the following lines do:

ig_attr_test_sum = ig_attr_test.detach().numpy().sum(0)
ig_attr_test_norm_sum = ig_attr_test_sum / np.linalg.norm(ig_attr_test_sum, ord=1)

In my case, ig_attr_test = ig.attribute(inputs=(embeddings_cat, test_set.features_cont), target=0) returns a tuple, and I'm not sure how to deal with that.

[vivekmig]

Hi @karppmik , that's great, I'm glad it worked! In the tutorial, those lines are summing over all examples and normalizing to get an averaged feature importance across examples.

In your case, the tuple in ig_attr_test of 2 tensors corresponds to the importance of each of the input tensors, embeddings_cat and test_set.features_cont. You can interpret each element of the output tuple of tensors as the importance of the corresponding input element towards the output target.

The exact post-processing steps depend on what information you want to extract, if you want to look at the importance of continuous features for a particular input example, you may want to look at a row of the second output tensor, corresponding to test_set.features_cont. For embedding features, the output corresponds to an importance per embedding dimension. Often, these dimensions don't have intuitive meaning, so it may make sense to sum over this dimension to obtain an importance for the overall embedding.

If you want to consider the average feature importance over multiple examples, you can sum or average over dimension 0, like the example in the regression tutorial. Note that if the feature has strong positive contribution for some examples and negative contribution for others, these may cancel when averaging and not reflect the true importance, so looking at importance for individual examples may be helpful as well.

[karppmik]

Okay, got it! Thanks a lot for help @vivekmig!

[hal-314]

@karppmik How did you solve it? Did you change model forward?

I replaced embedding layers with IntepretableEmbeddings and concatenate all embeddings outputs as said by vivekmig. However, model forward pass doesn't work because concatenating embeddings outputs and the pass to IntepretableEmbeddings isn't the same as passing categorical features to embeddings and concatenate them:

    ....
    def forward(self, x_cat, x_cont):
        # Note that x won't be correct when using IntepretableEmbeddings:
        # Normal case:  x_cat shape = (N, # categorical features) -> x = list of tensors (N, # embedding output for embedding i)
        # With captum: x_cat shape = (N, # embedding outputs) -> x  = list of tensors (# embedding outputs)

        x = [e(x_cat[:, i]) for i, e in enumerate(self.embeddings)] 
        x = torch.cat(x, 1) <- here is the error: "IndexError: Dimension out of range (expected to be in range of [-1, 0], but got 1)"
    ....

If I don't concatenate embeddings outputs, captum IntegratedGradients doesn't work and it'll fail in the forward pass :(

[karppmik]

@hal-314 what worked for me was just creating a separate forward function for the attribution. In my case as follows.

def forward(x_cat, x_cont):
        # x = [e(x_cat[:, i]) for i, e in enumerate(self.embeddings)]
        # x = torch.cat(x, 1)
        x = x_cat
        x = model.emb_drop(x)
        x2 = model.bn1(x_cont)
        x = torch.cat([x, x2], 1)
        x = F.relu(model.lin1(x))
        x = model.drops(x)
        x = model.bn2(x)
        x = F.relu(model.lin2(x))
        x = model.drops(x)
        x = model.bn3(x)
        x = model.lin3(x)
        return x

Then, use that forward function in

ig = IntegratedGradients(forward)
ig.attribute(inputs=(embeddings_cat, test_set.features_cont))

Remember to initialize your model and load the trained parameters first.

[hal-314]

many thanks!

[NarineK]

@karppmik, can we close this issue ? Have we resolved the issue ?