Skip to content

Concepts refacto & new generation specific splitter#156

Merged
AntoninPoche merged 125 commits into
mainfrom
concepts-refacto
Jun 22, 2026
Merged

Concepts refacto & new generation specific splitter#156
AntoninPoche merged 125 commits into
mainfrom
concepts-refacto

Conversation

@AntoninPoche

@AntoninPoche AntoninPoche commented Jun 8, 2026

Copy link
Copy Markdown
Collaborator

This PR should absolutely be reviewed commit by commit.

It touches many files due to renaming and some file moves. Also, it has a lot of lines because the notebooks are rerunning. But it is not that complex.

Description

This PR was initiated by the observation that ModelWithSplitPoint was too complicated (with several split points, many granularities, and aggregation strategies). But in practice, either for classification or generation, we always use the same settings. It was initiated for classification in #150.

Hence, the goal of this PR is to remove dead code, simplify API, make computations more efficient

The main modifications are:

  • Limit ModelWithSplitPoints to a single split point and remove the split_point parameter from any other object.
  • get_activations() does not return a dict[torch.Tensor] anymore but a tuple[torch.Tensor, torch.tensor | None]. With a single split point, no need for a dictionary of activations. Only a tuple (activations, predictions) where the predictions can be None. This removes the need for the get_split_activations()
  • Abstract splitting into a BaseSplitter and renaming every model_with_split_points attribute to splitter.
  • Optimize how nnsight is used, also require nnsight>=0.7.0,<0.8.0
  • Introduce a new SplitterForGeneration which forces the granularity to tokens, thus with no aggregations. Though there is still a parameter to include special tokens (but not padding).
  • Adapt documentation to splitters' abstraction.
  • Split the interpreto.model_wrapping module in interpreto.attributions.inference_wrappers, interpreto.concepts.splitters and move llm_interface.py to interpreto.commons. With new splitters and inference wrappers, the module was messy.
  • Add inputs_to_activations to splitters.
  • Rename encode_activations and decode_concepts to activations_to_concepts and concepts_to_activations.
  • Add a device property and a to() method to concept explainers and link them directly to the concept_model device. The splitters are already user exposed, but users had no unified way to move the concept_model device. This also removes device management in interpretations.
  • Adapt generation concept notebooks to modifications and rerun to validate validity.
  • Fix the vendoring script.

Breaking changes

  • ModelWithSplitPoints(split_points=[3]) is deprecated, not breaking. Use ModelWithSplitPoints(split_point=3) instead.
  • AnythingElse.split_point or AnythingElse(split_point=...) will raise an error. The splitters are the only ones managing the split point.
  • get_activations() output is not a dict but a tuple now.
  • splitters do not have a get_split_activations() method anymore. It is not necessary; activations are the first element of the tuple.
# Old and broken
activations_dict = splitter.get_activations(inputs)
activations = splitter.get_split_activationns

# New
activations, _  = splitter.get_activations(inputs)
  • from interpreto.model_wrapping import ModelWithSplitPoints is not the correct path anymore. One should use from interpreto import ModelWithSplitPoints which was already the recommended way.
  • Similarly, from interpreto.model_wrapping.llm_interface import OpenAILLM is not the correct path anymore. One should use from interpreto.commons.llm_interface import OpenAILLM.
  • encode_activations() and decode_concepts() methods from concept explainers have been renamed to activations_to_concepts() and concepts_to_activations(). But these are mainly internal to the library.
  • Interpretation methods do not have a concept_model_device anymore. But it should not change much, since the concept model is usually already on the device.
# Old and broken
interpretation_method = TopKInputs(concept_explainer, concept_model_device="cuda")

# New
concept_explainer.to("cuda")  # only moves the `.concept_model`, not the `.splitter`
interpretation_method = TopKInputs(concept_explainer)

Type of Change

  • 📚 Examples / docs / tutorials / dependencies update
  • 🥂 Improvement (non-breaking change which improves an existing feature)
  • 🚀 New feature (non-breaking change which adds functionality)
  • 💥 Breaking change (fix or feature that would cause existing functionality to change)

Checklist

  • I've read the CODE_OF_CONDUCT.md document.
  • I've read the CONTRIBUTING.md guide.
  • I've successfully run the style checks using make lint.
  • I've written tests for all new methods and classes that I created and successfully ran make test.
  • I've written the docstring in Google format for all the methods and classes that I used.

users should never interact with the inference wrapper, so cpu() and cuda() methods should be moved to the explainer
the device should be managed by the inference wrapper for an easier tracking, preventing memory leakage and hopefully bigger batch_size
…ions

it is the work of the embeddings perturbator to embed, so it was removed from inference wrappers. Furthermore, it is mandatory for embedding perturbators to have a input_embedder. Finally, now gradient wrapping assume they receive input_embeds
leave the target device moving to the inference wrapper. send mask to device for aggregators for them not to care about it. We can do so because with generators and garbage collectors, all these elements are destoyed once consumed.
All dispatch were removed, the only entry point is __call__,
it manages targets, targeted logits, or gradients.
Also the batching algorithm is changed.
Make the difference between Iterator, Iterable, and list. There should not be any Iterable, as we do not know what many times we can iterate on it. Iterator, once. List, indefinitly.
when the mask is required, the decorator is used and it passes the mask to device. No need to put the mask back to cpu afterward as it comes from an Iterator and will be automatically destroyed.
there is a single perturb function doing the embedding, it calls the perturb_embeds on a tensor and not the tensor mapping. Also repeat other keys to match perturbed embeddings.
most changes are minor and concern debugging artifacts.
one for global batching and optimization, the other for the three possible behaviors target, targeted logits, and gradients
Many failure case can come from weird default pad token ids, for Llama it was -1. Therefore, the setup_token_ids harmonize this.
…ation

in models '*ForSequenceClassification' we mostly use the cls token and can easily find the classification head. I wanted to use these features to make a simple and fast wrapper. The goal is ease inputs to concepts attributions.
It is a property which returns a model. If this model is passed to an attribution explainer, it provides inputs-to-concepts attributions.
This new wrapper should support attribution methods to interpret concepts.
Comment thread docs/notebooks/generation_concept_tutorial.ipynb
Comment thread docs/notebooks/generation_concept_tutorial.ipynb
Comment thread docs/notebooks/generation_concept_tutorial.ipynb
Comment thread docs/notebooks/examples/classification_probe_tutorial.ipynb

@fanny-jourdan fanny-jourdan left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Great job, Antonin! The changes were necessary and well done ;)
I've added a few comments to the code (and some questions to make sure I understood it correctly), and most importantly, some things to fix in the notebook that aren't running properly. I'll leave it to you to take care of these small issues, and then you can merge!

Comment thread docs/notebooks/examples/generation_demonstration.ipynb
Comment thread docs/notebooks/examples/classification_demonstration.ipynb
@fanny-jourdan

Copy link
Copy Markdown
Contributor

I reran all the notebooks with your changes, and everything is working fine now. :)

@AntoninPoche AntoninPoche changed the base branch from dev to main June 22, 2026 21:42
@AntoninPoche AntoninPoche merged commit 3aaf4a3 into main Jun 22, 2026
4 checks passed
@AntoninPoche AntoninPoche changed the title Concepts refacto & new generation specific splitter (4) Concepts refacto & new generation specific splitter Jun 22, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants