All Layers Splitter

AllLayersSplitter captures the residual stream before the first transformer block and after every block. It also applies the wrapped model's native normalization, pooling, and prediction head to one or more residual states.

interpreto.AllLayersSplitter

AllLayersSplitter(model_or_repo_id, *, automodel=AutoModelForCausalLM, tokenizer=None, config=None, device_map=None, layer_path=None, **kwargs)

Bases: LanguageModel

Extract the residual stream before and after every transformer block.

The transformer blocks are inferred from the model configuration or selected explicitly with layer_path. Activations are returned in model order: the input to the first block followed by the output of every block. A model with L transformer blocks therefore returns L + 1 tensors of shape (1, sequence_length, model_width).

This splitter is intended for methods that compare representations across model depths, such as Logit Lens and Tuned Lens. It does not implement the concept-specific BaseSplitter interface because there is no single split point or latent representation.

Parameters:

Name	Type	Description	Default
model_or_repo_id	str | PreTrainedModel	Hugging Face repository ID, local checkpoint path, or preloaded model.	required
automodel	type[AutoModel]	Hugging Face AutoClass used when loading a model from a repository ID or local path.	AutoModelForCausalLM
tokenizer	PreTrainedTokenizerBase | None	Tokenizer associated with a preloaded model.	None
config	PretrainedConfig | None	Optional model configuration passed to the model loader.	None
device_map	device | str | None	Device map passed to the model loader.	None
layer_path	str | None	Path to the transformer block ModuleList, relative to the Hugging Face model. Required when it cannot be identified unambiguously from config.num_hidden_layers.	None
**kwargs	Any	Additional arguments passed to NNsight's LanguageModel.	{}

Raises:

Type	Description
InitializationError	If a preloaded model is provided without a tokenizer.
ValueError	If a tokenizer cannot be inferred for a repository ID.

Example

from transformers import AutoModelForCausalLM from interpreto import AllLayersSplitter splitter = AllLayersSplitter("gpt2") activations = splitter.get_activations("Interpreto is useful.") len(activations) == len(splitter.split_points) + 1 True

Source code in interpreto/concepts/splitters/all_layers_splitter.py

def __init__(
    self,
    model_or_repo_id: str | PreTrainedModel,
    *,
    automodel: type[AutoModel] = AutoModelForCausalLM,
    tokenizer: PreTrainedTokenizerBase | None = None,
    config: PretrainedConfig | None = None,
    device_map: torch.device | str | None = None,
    layer_path: str | None = None,
    **kwargs: Any,
) -> None:
    if isinstance(model_or_repo_id, PreTrainedModel) and tokenizer is None:
        raise InitializationError(
            "Tokenizer is not set. When providing a model instance, the tokenizer must be set."
        )

    super().__init__(
        model_or_repo_id,
        config=config,
        tokenizer=tokenizer,
        automodel=automodel,
        device_map=device_map,
        **kwargs,
    )

    if self.tokenizer is None:
        raise ValueError("`tokenizer` must be provided when it cannot be inferred from the model.")
    if self.tokenizer.pad_token is None:
        self.tokenizer.pad_token = self.tokenizer.eos_token

    layer_name, layers = self._find_layers(self._model, layer_path)
    self.split_points = [f"model.{layer_name}.{index}" for index in range(len(layers))]
    self._block_output_arity: list[int | None] = [None] * len(self.split_points)

activation_names property

activation_names

Names of the residual states returned by :meth:get_activations.

apply_head

apply_head(activations)

Apply the wrapped model's prediction head to residual activations.

The transformer blocks are skipped and activations are used as their output. The wrapped model then executes its own downstream normalization, pooling, and prediction head. This avoids architecture-specific head names and preserves functional operations implemented in model forward methods.

Parameters:

Name	Type	Description	Default
activations	Tensor	Residual activations with shape (n, sequence_length, model_width). The leading dimension may represent several layer boundaries from the same input. Inputs are moved to the model device while preserving gradients.	required

Returns:

Type	Description
Tensor	torch.Tensor: Logits returned by the wrapped model for every activation in the leading dimension.

Source code in interpreto/concepts/splitters/all_layers_splitter.py

def apply_head(self, activations: torch.Tensor) -> torch.Tensor:
    """Apply the wrapped model's prediction head to residual activations.

    The transformer blocks are skipped and ``activations`` are used as
    their output. The wrapped model then executes its own downstream
    normalization, pooling, and prediction head. This avoids
    architecture-specific head names and preserves functional operations
    implemented in model ``forward`` methods.

    Args:
        activations (torch.Tensor): Residual activations with shape
            ``(n, sequence_length, model_width)``. The leading dimension
            may represent several layer boundaries from the same input.
            Inputs are moved to the model device while preserving gradients.

    Returns:
        torch.Tensor: Logits returned by the wrapped model for every
            activation in the leading dimension.
    """
    # NNsight's activation cache can reside on CPU even for a CUDA model.
    # Tensor.to preserves the gradient path back to the caller's activations.
    activations = activations.to(self._model.device)
    embedding_width = self._model.get_input_embeddings().weight.shape[-1]
    # The embedding code still validates its input even though every block is skipped.
    inputs_embeds = activations.new_zeros((*activations.shape[:-1], embedding_width))
    with self.trace(inputs_embeds=inputs_embeds):
        for split_point, output_arity in zip(self.split_points, self._block_output_arity, strict=True):
            replacement = activations if output_arity is None else (activations,) + (None,) * (output_arity - 1)
            self.get(split_point.removeprefix("model.")).skip(replacement)
        logits = self.output.logits.save()
    return logits

get_activations

get_activations(inputs)

Extract the residual stream for one text input.

Parameters:

Name	Type	Description	Default
inputs	str	Text passed to the wrapped model.	required

Returns:

Type	Description
list[Tensor]	list[torch.Tensor]: Input to the first transformer block followed by every transformer block output in split_points order. Each tensor has shape (1, sequence_length, model_width).

Source code in interpreto/concepts/splitters/all_layers_splitter.py

def get_activations(self, inputs: str) -> list[torch.Tensor]:
    """Extract the residual stream for one text input.

    Args:
        inputs (str): Text passed to the wrapped model.

    Returns:
        list[torch.Tensor]: Input to the first transformer block followed by
            every transformer block output in ``split_points`` order.
            Each tensor has shape ``(1, sequence_length, model_width)``.
    """
    return self._trace_activations(inputs)