Decoder-only Attribution Models Support

[gsarti]

Description

This PR introduces support to perform attributions on decoder-only models from Hugging Face (e.g. GPT2) using the AutoModelForCausalLM abstraction. The implementation will involve several generalizations of existing classes and methods, notably:

• Support for Batch alongside the current EncoderDecoderBatch usage.
• Support for target-side-only attributions produced by decoder-only models throughout the attribution pipeline and in FeatureAttributionSequenceOutput.
• Handling ForSeq2SeqLM (current) and ForCausalLM (to be added) in the HuggingfaceModel class.

The new loading pipeline will also support using pre-loaded models in inseq.load_model aiming for compatibility with LLM.int8 (see #141).

Related Issue

#128

Type of Change

• 🚀 New feature (non-breaking change which adds functionality)

[gsarti]

In the current state, decoder-only models can be instantiated with inseq.load_model and used to generate text with the model.generate wrapper function. Compatibility with EncoderDecoder models is preserved, with currently existing tests passing. However, the attribution process is currently not functioning for DecoderOnly models.

[gsarti]

Major (non-breaking) changes in progress:

• Disentangled format_attribute_args from attribute_step to enable independent usage and overloading in children classes (see e.g. DiscretizedIntegratedGradients overloading)

• Model architecture classes: Separating architecture and framework responsibility in AttributionModel introducing framework-agnostic architecture classes (EncoderDecoderAttributionModel, DecoderOnlyAttributionModel) to handle input processing and attribution params formatting. Final classes will be children sharing one framework and one architecture (e.g. HuggingfaceEncoderDecoderModel).

• Fix partial target-side attribution: These were added at the beginning but didn't keep up with newest features, resulting in wrong slices when non-default attr_pos_start and attr_pos_end parameters were passed. These should now be fully fixed and functional when attribute_target=True.

Example with full attribution:

Example constraining same attributions between steps [6, 10):

Note how the selected slices attributed in the second case match scores with the same slices in the full attribution.

Todos:

• Add classes responsibility diagrams with attribution flow to Dev docs
• Fix contrastive attribution test
• Add test case for full vs sliced attribution using attribute_target

[gsarti]

Some tests are broken at the moment (need checking) but decoder-only attribution seems to be producing the right output. Visualization still doesn't support decoder-only formats, so .show is not yet functional.

[gsarti]

Minor adjustments to viz functions and aggregators to support an empty source_attribution field for decoder-only models. Tests failing on some issues with the attr_pos_start and attr_pos_end that require further investigation - current results are wrong.

[gsarti]

Fixed the issues with positions, and now both seq2seq and decoder-only attribution are working as expected! 🎉 Also added the first test for decoder-only models to the test suite to ensure continued compatibility. Points still needing discussion/fixing:

• Currently the combination of batched attribution (i.e. more than one input_text provided to attribute) and constrained decoding (i.e. specified generated_text) doesn't work if generated texts do not have the same length. Such functionality would require bidirectional padding, while currently the fact of using padding="left" for decoder-only models forces all batched outputs to have the same length.

EDIT: We opted for enforcing a batch size of one (with an info message to inform the user) to handle this case for the moment. The only downside is a slower attribution process, but it is reasonable in the context.

• Setting start/end positions for decoder-only attributions produce wrong slices (i.e. the only modalities in which decoder-only is working now is single-example generate/constrained or batched generate/constrained with fixed length). This will need fixing before merging.

EDIT: Enforcing a batch size of one for the case of multiple texts with custom start-end positions for a decoder-only model.

[nfelnlp]

This looks amazing! :)

When I tested it with the following setup:

gpt_model = inseq.load_model("gpt2", "saliency")
gpt_out = gpt_model.attribute(
    "The developer argued with the designer because her idea",
    show_progress=True,
    pretty_progress=True,
    generation_args={"max_new_tokens": 10}
)
gpt_out.show()

...I ran into this assertion error:

inseq/inseq/models/attribution_model.py

Lines 196 to 198 in 27d1a77

	assert all(
	generated_texts[idx].startswith(input_texts[idx]) for idx in range(len(input_texts))
	), "Forced generations with decoder-only models must start with the input texts."

It turns out that generated_texts is a str, while input_texts is a list (of length one).

I could produce this after simply commenting out the three lines:

Should I run into this error or should we make the assertion less strict?

[gsarti]

Thanks for spotting this bug @nfelnlp! The format_input_texts call in attribute is supposed to enforce that both input and generated texts are lists of strings, but if the generated text is not forced and is only one, generate wrongfully returns a string instead of a list of length 1. To solve the issue it's sufficient to remove this line

texts = texts[0] if len(texts) == 1 else texts

from HuggingfaceModel.generate.

[gsarti]

@nfelnlp with the last commit both the case you mentioned and the batched decoder-only attribution should work, documented our choice in the comment above!

[gsarti]

@nfelnlp only some robustness tests are left, then we can proceed with the merge. Could you help with that? 🙂

[nfelnlp]

Just tested the latest version and it works flawlessly for the GPT-2 use case I described above.
Looks good to me! ✅