AssertionError: Target not provided when necessary, cannot take gradient with respect to multiple outputs.

[kristosh]

The packages I am using are the following:

Python 3.10
PyTorch 2.6.0
captum 0.7.0

I am trying to explain the behavior of a BERT model trained for sequence learning. The model is the following, which is loaded using HuggingFace and the transformers library. The input to the model is a list of tokens that varies depending on the input length, and the target is a label per each input token.

I would like to use an Integrated gradient to explain the model's behavior for each of the tokens and how the context influences the model's decision.

My model looks as follows:

To run the IG, I am using the following code:

`
...

     baseline_ids = baseline_ids.unsqueeze(0) #torch.Size([1, 136])
     input_ids = input_ids.unsqueeze(0)       #torch.Size([1, 1, 136])
     input_embeds = self.model.bert.embeddings.word_embeddings(input_ids) #torch.Size([1, 136, 768])
     baseline_embeds = self.model.bert.embeddings.word_embeddings(baseline_ids) #torch.Size([1, 136, 768])
     ig = IntegratedGradients(model_forward)
     if len(input_embeds.shape) == 4:
       input_embeds = input_embeds.squeeze(1)
       baseline_embeds = baseline_embeds.squeeze(1)    

   attributions = ig.attribute(inputs=input_embeds, baselines=baseline_embeds,target=0, n_steps=n_steps)

`

Where target = 0 is because of the label for the first token. The following code throws the following error:

AssertionError: Target not provided when necessary, cannot take gradient concerning multiple outputs.

The model outputs a 19th-dimensional vector per token. Why am I receiving this error?

[kristosh]

I think one of the issues that I have spotted is that the model_forward, which I am not sure what it does it returns incorrect dimensionalities:

 def model_forward(inputs_embeds):
    return model(inputs_embeds=inputs_embeds).logits

So the input_embeds shape is torch.Size([50, 136, 768]) (where 50 is the number of steps, 136 the number of tokens, and 768 the last layer dimensionality) and the model(inputs_embeds=inputs_embeds).logits returns torch.Size([50, 136, 19]) while in my understanding it should have been torch.Size([50, 19]). Why is that happening? Moreover, what is the need of this function in my code?

[craymichael]

Hi, this is because the function is expecting a simpler classifier than outputs a tensor of shape [num_samples, num_classes]. Your model is outputting a tensor of shape [num_samples, num_tokens, num_classes]. I think you might be able to fix it by setting the target to (0, 0) for class 0 and token 0. However, you would need to run this 19 times, once per class (that you are interested in).

Instead, I would recommend checking out the tutorial for BERT question answering which is similar to your use case: https://captum.ai/tutorials/Bert_SQUAD_Interpret (and https://captum.ai/tutorials/Bert_SQUAD_Interpret2). An introduction to LLM attribution is given in https://captum.ai/tutorials/Llama2_LLM_Attribution but Llama (though it may be easier to follow). We have LLM attribution wrappers so you can get attribute to all output tokens in a nicer way.

[kristosh]

Is there an example of an attribution method for sequence labelling? I am wondering what the status of this type of algorithm is. Is there a unique way to approach this?

As for the code, I have modified the gradient.py file and the assertion and grad calculation (in the function compute_gradients) by using sth like this:

def compute_gradients(
    forward_fn: Callable,
    inputs: Union[Tensor, Tuple[Tensor, ...]],
    target_ind: TargetType = None,
    additional_forward_args: Any = None,
) -> Tuple[Tensor, ...]:
    r"""
    Computes gradients of the output concerning the inputs for an
    arbitrary forward function.

    Args:

        forward_fn: forward function. This can be for example model's
                    forward function.
        input:      Input at which gradients are evaluated,
                    will be passed to forward_fn.
        target_ind: Index of the target class for which gradients
                    must be computed (classification only).
        additional_forward_args: Additional input arguments that forward
                    function requires. It takes an empty tuple (no additional
                    arguments) if no additional arguments are required
    """
    with torch.autograd.set_grad_enabled(True):
        # runs forward pass
        outputs = _run_forward(forward_fn, inputs, target_ind, additional_forward_args)

        assert outputs[0, target_ind].numel() == 1, (
            "Target not provided when necessary, cannot"
            " take gradient with respect to multiple outputs."
        )
        # torch.unbind(forward_out) is a list of scalar tensor tuples and
        # contains batch_size * #steps elements

        grads = torch.autograd.grad(torch.unbind(outputs[:,target_ind]), inputs)   
        #grads = torch.autograd.grad(torch.unbind(outputs), inputs)
    return grads

With this update the function works and returns the grads!

[craymichael]

Taking a quick look that would work, but would change the function behavior in an unexpected way for many uses in the library. Here's something that should work for your use case-

import torch
from captum.attr import IntegratedGradients

def my_forward(x: torch.Tensor) -> torch.Tensor:
    x = x.reshape(-1, 10 ,5)
    w = torch.arange(5 * n_classes, device=x.device, dtype=x.dtype).reshape(5, n_classes)
    y = x @ w
    return y.reshape(-1, n_classes)  # [N, S, C] --> [NS, C]

n_classes = 3
# your input
inp_orig = torch.rand([2, 10, 5], requires_grad=True, device='cuda')
# repeat the input for each target class
inp = inp_orig.repeat(n_classes, 1, 1)
# set the target for each class we're interested in for each input example (batch size * seq length)
targets = torch.arange(n_classes, device=inp.device).repeat_interleave(inp_orig.shape[0] * inp_orig.shape[1])
# compute attributions
ig = IntegratedGradients(my_forward)
attr = ig.attribute(inputs=inp.reshape(-1, 5), target=targets)
# reshape attributions
attr = attr.reshape(n_classes, *inp_orig.shape)

The final attr will be of shape [num classes, batch size, seq length, input dim]. I used a dummy function to test the code but you can substitute your model there.

An implementation reusing the autograd graph would be more efficient but this will get what you want (attributions for all classes and elements in the sequence for each sample in the batch)