Why do captum's perturbation and IG treat input & target differently?

rbelew: <p>I’ve been successfull using captum’s <code>LayerIntegratedGradients</code> class,<br>
but none of my attempts trying the same sorts of inputs and targets using<br>
<code>LLMAttribution</code> seem to work.</p>
<p>I’m working with a <code>BertForMultipleChoice</code> model, and the input is a list of<br>
the repeated prompt followed by the choices:</p>
<pre><code>for i,c in enumerate(tst['input_ids'][0]):
    indices = c.detach().tolist()
    sepIdx =  indices.index(SEP_IDX)
    nearSep = indices[sepIdx-prefix:]
    preTokens = tokenizer.convert_ids_to_tokens(indices[sepIdx-prefix:sepIdx-1])
    choiceTokens = tokenizer.convert_ids_to_tokens(indices[sepIdx+1:])
    print(f"{i} {sepIdx} {' '.join(preTokens):&gt;55}\t[SEP] {' '.join(choiceTokens)}")

	0 120           behalf of fake charities . as webster sees it	[SEP] recognizing the guidelines commentary is authoritative [SEP]
	1 109   when he sol ##ici ##ted personal information from the	[SEP] holding that a sentencing guide ##line pre ##va ##ils over its commentary if the two are inconsistent [SEP]
	2 98                                       l ( b ) ( 9 ) ( a	[SEP] holding that sentencing guidelines commentary must be given controlling weight unless it violate ##s the constitution or a federal statute or is plainly inconsistent with the guidelines itself [SEP]
	3 99                                       ( b ) ( 9 ) ( a )	[SEP] holding that commentary is not authoritative if it is inconsistent with or a plainly er ##rone ##ous reading of the guide ##line it interpret ##s or explains [SEP]
	4 119           on behalf of fake charities . as webster sees	[SEP] holding that guidelines commentary is generally authoritative [SEP]
</code></pre>
<p>I’m using  <code>LayerIntegratedGradients</code> with a test example and a target scalar representing the index of the correct (multiple choice) like this:</p>
<pre><code>    tstEGTuple = (tst['input_ids'], 
                  tst['attention_mask'], 
                  tst['token_type_ids'])
    targetIdx = 3 # for this particular test example
        
    lig = LayerIntegratedGradients(custForwardModel, model.bert.embeddings)
    attributions_ig = lig.attribute(tstEGTuple, n_steps=5,target=targetIdx) 
</code></pre>
<p>and that works, eg allowing calculations like <code>summarize_attributions(attributions_ig), viz.VisualizationDataRecord()</code> etc.</p>
<p>For LLMAttribution I am following the <a href="https://captum.ai/tutorials/Llama2_LLM_Attribution#Perturbation-based-Attribution" rel="noopener nofollow ugc">Llama2 tutorial</a> The closest I can get with LLMAttribution seems to require use of <code>TextTokenInput</code> for input, but <em>raw text</em> for the target?</p>
<pre><code>    in0 = tst['input_ids'][0][0]
    in0_tokens = tokenizer.convert_ids_to_tokens(in0)
    in0Txt = ' '.join(in0_tokens)
    in4captum = TextTokenInput(in0Txt, tokenizer,skip_tokens=skip_tokens)

    target = targetList[egIdx]                   
    targetIn = tst['input_ids'][0][target]
    targ_tokens = tokenizer.convert_ids_to_tokens(targetIn)
    targTxt = ' '.join(targ_tokens)
    # targ4captum = TextTokenInput(targTxt, tokenizer,skip_tokens=skip_tokens)
    
	llm_attr = LLMAttribution(fa, tokenizer)
	attributions_fa = llm_attr.attribute(in4captum, target=targTxt) 
</code></pre>
<p>but this raises an exception: <code>prepare_inputs_for_generation</code> isn’t<br>
available for this  <code>BertForMultipleChoice</code> model:</p>
<pre><code>	Traceback (most recent call last):
	File "/Users/rik/Code/eclipse/ai4law/src/run_multChoice.py", line 874, in &lt;module&gt;
	main()
	File "/Users/rik/Code/eclipse/ai4law/src/run_multChoice.py", line 854, in main
	captumPerturb(model,tokenizer,tstEGTensorDict,tstEGtarget,OutDir)
	File "/Users/rik/Code/eclipse/ai4law/src/run_multChoice.py", line 479, in captumPerturb
	attributions_fa = llm_attr.attribute(in4captum, target=targTxt) 
	^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
	File "/Users/rik/data/pkg/miniconda3/envs/ai4law2/lib/python3.11/site-packages/captum/attr/_core/llm_attr.py", line 667, in attribute
	cur_attr = self.attr_method.attribute(
	^^^^^^^^^^^^^^^^^^^^^^^^^^^
	File "/Users/rik/data/pkg/miniconda3/envs/ai4law2/lib/python3.11/site-packages/captum/log/dummy_log.py", line 39, in wrapper
	return func(*args, **kwargs)
	^^^^^^^^^^^^^^^^^^^^^
	File "/Users/rik/data/pkg/miniconda3/envs/ai4law2/lib/python3.11/site-packages/captum/attr/_core/feature_ablation.py", line 288, in attribute
	initial_eval: Union[Tensor, Future[Tensor]] = _run_forward(
	^^^^^^^^^^^^^
	File "/Users/rik/data/pkg/miniconda3/envs/ai4law2/lib/python3.11/site-packages/captum/_utils/common.py", line 588, in _run_forward
	output = forward_func(
	^^^^^^^^^^^^^
	File "/Users/rik/data/pkg/miniconda3/envs/ai4law2/lib/python3.11/site-packages/captum/attr/_core/llm_attr.py", line 567, in _forward_func
	model_inputs = self.model.prepare_inputs_for_generation(
	^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
	File "/Users/rik/data/pkg/miniconda3/envs/ai4law2/lib/python3.11/site-packages/transformers/generation/utils.py", line 376, in prepare_inputs_for_generation
	raise NotImplementedError(
</code></pre>
<p>Thanks for any suggestions!</p>

rbelew: <p>To patch around this unimplemented <code>prepare_inputs_for_generation</code><br>
error, I am now trying to construct the model inputs directly:</p>
<pre><code>            #model_inputs = self.model.prepare_inputs_for_generation(
            #     model_inp, **model_kwargs

            addtl_model_inputs = {'attention_mask': attention_mask}     
            outputs = self.model.forward(model_inp, **addtl_model_inputs)
</code></pre>
<p>The BertForMultipleChoice model’s forward function expects input_ids, but<br>
also attention_mask and token_type_ids.</p>
<p>The FeatureAblation constructor allows specification of a tailored<br>
forward function, and that’s the strategy I’ve used previously with<br>
LayerIntegratedGradients:</p>
<pre><code>MCfwd = partial(multChoice_forward,model)
lig = LayerIntegratedGradients(MCfwd, model.bert.embeddings)
</code></pre>
<p>But providing it for FeatureAblation</p>
<pre><code>MCfwd = partial(multChoice_forward,model)
fa = FeatureAblation(MCfwd) 
</code></pre>
<p>causes  <code>BaseLLMAttribution.__init__</code> to get confused about what’s the<br>
model and what’s the forward_func, with this comment (L#374)</p>
<pre><code>    # alias, we really need a model and don't support wrapper functions
    # coz we need call model.forward, model.generate, etc.  
	self.model: nn.Module = cast(nn.Module, self.forward_func)
</code></pre>
<p>and this confusion immediately causes trouble, eg,  looking for the<br>
device of the function:(</p>
<p>FeatureAblation allows <code>additional_forward_args</code> and that sounds like<br>
what I want: “If the forward function requires additional arguments<br>
other than the inputs for which attributions should not be computed,<br>
this argument can be provided…”</p>
<p>But trying to provide this argument to the FeatureAblation constructor<br>
generates:</p>
<pre><code>	TypeError: captum.attr._core.feature_ablation.FeatureAblation.attribute() got multiple values for keyword argument 'additional_forward_args'
</code></pre>
<p>I note that LLMAttribution._forward_func() helpfully computes an<br>
attention mask and adds it to <code>model_kwargs</code> so I only need to get the<br>
token_type_ids into the model’s forward() method.</p>