PASTA

steerability.algorithms.state_control.pasta

args

control

SUPPORTED_ATTN_IMPLEMENTATIONS = ('eager', 'sdpa') module-attribute

logger = logging.getLogger(__name__) module-attribute

PASTA

Bases: HookControl

Implementation of PASTA (Post-hoc Attention STeering Approach) from Zhang et al., 2023.

PASTA performs controlled text generation by dynamically modifying attention patterns during inference to amplify or suppress the influence of specific text spans. This allows for fine-grained steering of model behavior without requiring model retraining or parameter updates.

The algorithm works by:

1. Substring Identification: Locate target substrings within the input prompt using tokenizer offset mapping to determine precise token ranges.

2. Attention Modification: Inject scaling factors into the attention mask of specified layers and heads to increase or decrease attention weights for the identified token ranges.

3. Dynamic Steering: Apply different scaling strategies (include, exclude, or generation-focused) to control how the model attends to relevant spans during text generation.

This approach enables real-time control over model focus and can be used for tasks like concept amplification, bias mitigation, or content filtering without architectural changes.

The head_config argument accepts a dict (layer index to head indices), a list (layer indices, all heads), or a HeadProfile recipe. A HeadProfile moves the paper's one-time head-profiling stage into steer(): each candidate head is steered on its own on a task-agnostic set of profiling prompts, scored by a SampleScorer, and ranked by the paired lift over an unsteered baseline; the selected heads become the dict head map, and the resolution is available as head_profile (a HeadProfileResult). Profiling runs on the live model through the pipeline's session, and the resolved head map freezes into a .spipe as a plain-dict PASTA.

Note

PASTA injects a 4D additive attention mask, which only the "eager" and "sdpa" attention implementations consume. Load the model with attn_implementation="eager" (or "sdpa"); "flash_attention_2" expects a 2D padding mask and silently ignores or errors on the injected scaling. steer() fails fast if the model reports an unsupported implementation.

Reference: - "PASTA: Tell Your Model Where to Attend: Post-hoc Attention Steering for LLMs" Qingru Zhang, Chandan Singh, Liyuan Liu, Xiaodong Liu, Bin Yu, Jianfeng Gao, Tuo Zhao https://arxiv.org/abs/2311.02262

Source code in steerability/algorithms/state_control/pasta/control.py

class PASTA(HookControl):
    """
    Implementation of PASTA (Post-hoc Attention STeering Approach) from Zhang et al., 2023.

    PASTA performs controlled text generation by dynamically modifying attention patterns during inference to amplify or
    suppress the influence of specific text spans. This allows for fine-grained steering of model behavior without
    requiring model retraining or parameter updates.

    The algorithm works by:

    1. **Substring Identification**: Locate target substrings within the input prompt using tokenizer offset mapping to
    determine precise token ranges.

    2. **Attention Modification**: Inject scaling factors into the attention mask of specified layers and heads to
    increase or decrease attention weights for the identified token ranges.

    3. **Dynamic Steering**: Apply different scaling strategies (include, exclude, or generation-focused) to control how
    the model attends to relevant spans during text generation.

    This approach enables real-time control over model focus and can be used for tasks like concept amplification, bias
    mitigation, or content filtering without architectural changes.

    The `head_config` argument accepts a dict (layer index to head indices), a list (layer indices,
    all heads), or a `HeadProfile` recipe. A `HeadProfile` moves the paper's one-time head-profiling
    stage into `steer()`: each candidate head is steered on its own on a task-agnostic set of
    profiling prompts, scored by a `SampleScorer`, and ranked by the paired lift over an unsteered
    baseline; the selected heads become the dict head map, and the resolution is available as
    `head_profile` (a `HeadProfileResult`). Profiling runs on the live model through the pipeline's
    session, and the resolved head map freezes into a `.spipe` as a plain-dict PASTA.

    Note:
        PASTA injects a 4D additive attention mask, which only the `"eager"` and `"sdpa"` attention
        implementations consume. Load the model with `attn_implementation="eager"` (or `"sdpa"`);
        `"flash_attention_2"` expects a 2D padding mask and silently ignores or errors on the
        injected scaling. `steer()` fails fast if the model reports an unsupported implementation.

    Reference:
    - "PASTA: Tell Your Model Where to Attend: Post-hoc Attention Steering for LLMs"
    Qingru Zhang, Chandan Singh, Liyuan Liu, Xiaodong Liu, Bin Yu, Jianfeng Gao, Tuo Zhao
    [https://arxiv.org/abs/2311.02262](https://arxiv.org/abs/2311.02262)
    """

    Args = PASTAArgs
    RUNTIME_KWARGS_SCHEMA = [
        {
            "name": "substrings",
            "type": "list[str]",
            "required": True,
            "scope": "row",
            "help": (
                "Substrings whose attention should be steered. Required at inference time. A `str` "
                "broadcasts to every row; a `list[list[str]]` of batch length carries one group per "
                "row; a flat `list[str]` is accepted only at batch size 1, as that row's group."
            ),
        },
    ]

    supports_batching: bool = True

    # placeholders
    model: PreTrainedModel | None = None
    tokenizer: PreTrainedTokenizerBase | None = None
    device: torch.device | str | None = None
    head_profile: "HeadProfileResult | None" = None

    _head_map: dict[int, list[int]] | None = None
    _layers: list[int] | None = None
    _attn_module_names: dict[int, str] | None = None
    _num_heads_by_layer: dict[int, int] | None = None
    _layout = None

    def requirements(self) -> Requirements:
        """Backend requirements for attention-map editing.

        PASTA writes into attention maps through torch hooks, which fused paged-attention
        kernels never materialize, so the generate phase requires `Capability.IN_PROCESS_TORCH`.
        The spec constraint reports a configured incompatible attention implementation
        (`hf_model_kwargs["attn_implementation"]` outside `"eager"`/`"sdpa"`) at `check()` time,
        before any model loads; the same condition is re-checked against the live model in
        `steer()`.

        Returns:
            The control's phase-keyed requirements.
        """
        return Requirements(
            generate=needs(Capability.IN_PROCESS_TORCH),
            spec_constraints=(
                SpecConstraint(
                    description=(
                        "PASTA requires attn_implementation 'eager' or 'sdpa' to inject a 4D "
                        "attention mask; set attn_implementation=\"eager\" in hf_model_kwargs."
                    ),
                    predicate=_attn_implementation_supported,
                ),
            ),
        )

    def steer_access(self) -> ModelAccess:
        """`ModelAccess.MODULE`; the attention module paths resolve on the live model, which
        is retained for the hook closures (the generate phase is in-process)."""
        return ModelAccess.MODULE

    def steer(
        self,
        model: PreTrainedModel,
        tokenizer: PreTrainedTokenizerBase | None = None,
        session=None,
        **__,
    ) -> PreTrainedModel:
        """Initialize PASTA by configuring attention head mappings and model references.

        Resolves the attention module path and head count of every attention layer of the model,
        resolves a `HeadProfile` head map when `head_config` is a profiling recipe (steering each
        candidate head through the pipeline's session), sets up the resolved head map, and fails
        fast on unsupported layers or attention implementations (rather than deep inside
        generation).

        Args:
            model (PreTrainedModel): The base language model to be steered.
            tokenizer (PreTrainedTokenizerBase | None): Tokenizer for substring identification.
                If None, attempts to retrieve from model attributes.
            session: The `ScopedSession` the pipeline hands to `steer()`, scoped to
                `ModelAccess.MODULE`, through which a `HeadProfile` runs its rollouts.

        Returns:
            PreTrainedModel: The input model (unchanged).

        Raises:
            ValueError: If a configured attention module path is missing on the model, if the
                model's attention implementation is not one of `"eager"` / `"sdpa"`, or if a
                `HeadProfile` selects no heads.
        """
        self.model = model
        self.tokenizer = tokenizer or getattr(model, "tokenizer", None)
        self.device = next(model.parameters()).device
        self._resolve_attention_modules(model)
        self._check_attention_implementation(model)

        if isinstance(self.head_config, HeadProfile):
            self.head_profile = self.head_config.resolve(self, model, self.tokenizer, session=session)
            head_map = self.head_profile.head_config
        else:
            head_map = self.head_config

        self._setup_layers(head_map)
        self._finalize_head_map()
        return model

    def _resolve_attention_modules(self, model: PreTrainedModel) -> None:
        """Resolve every attention layer's module path and head count from the model layout.

        The per-layer attention module paths come from `resolve_model_layout` (`.self_attn` for
        `model.layers.*`, `.attn` for `transformer.h.*`). Every attention layer of the layout is
        resolved (not only the configured layers), so a `HeadProfile` can enumerate candidates
        from the same table `get_hooks` reads. Each layer's head count is read per layer from the
        module tree, so a model whose head count varies across layers is sized per layer.

        Raises:
            ValueError: If the architecture is unrecognized, or an attention module path is
                absent.
        """
        layout = resolve_model_layout(model)
        self._layout = layout
        attn_names = layout.attn_names

        self._attn_module_names = {}
        self._num_heads_by_layer = {}
        for layer in layout.attention_layers:
            path = attn_names[layer]
            try:
                model.get_submodule(path)
            except AttributeError as error:
                raise ValueError(f"PASTA could not resolve attention module {path!r}.") from error
            self._attn_module_names[layer] = path
            self._num_heads_by_layer[layer] = head_geometry(model, layout, layer).num_heads

    @staticmethod
    def _check_attention_implementation(model: PreTrainedModel) -> None:
        """Fail fast unless the model uses an attention implementation PASTA can steer.

        Raises:
            ValueError: If `model.config._attn_implementation` is not `"eager"` or `"sdpa"`.
        """
        impl = getattr(model.config, "_attn_implementation", "eager")
        if impl not in {"eager", "sdpa"}:
            raise ValueError(
                f"PASTA requires attn_implementation 'eager' or 'sdpa' to inject a 4D attention "
                f"mask; got {impl!r}. Load the model with attn_implementation=\"eager\"."
            )

    def get_hooks(
        self,
        input_ids: torch.Tensor,
        runtime_kwargs: dict | None,
        **__,
    ) -> dict[str, list]:
        """Create attention modification hooks for specified substrings.

        Identifies token ranges corresponding to target substrings and prepares hooks that will modify attention weights
        during the forward pass, at the configured head map and the control's `alpha`.

        Args:
            input_ids (torch.Tensor): Input token IDs of shape [batch_size, seq_len].
            runtime_kwargs (dict | None): Must contain "substrings" key with target text spans:

                - str: one substring, broadcast to every batch item
                - list[list[str]]: one substring group per batch item, of batch length
                - list[str]: accepted only at batch size 1, as that row's group

        Returns:
            dict[str, list]: Hook specifications with "pre", "forward", "backward" keys. Only "pre" hooks are populated for attention modification.

        Raises:
            ValueError: If "substrings" is missing from runtime_kwargs, a flat list of strings is
                passed at batch size > 1, a group is a `str` or contains a non-`str` element, or
                the number of groups does not match the batch size.
        """
        if not runtime_kwargs or "substrings" not in runtime_kwargs:
            raise ValueError("PASTA requires 'substrings' inside runtime_kwargs")

        token_ranges, input_len = self.locate_spans(input_ids, runtime_kwargs["substrings"])
        return self.build_hooks_for(token_ranges, input_len, self._head_map, self.alpha)

    def locate_spans(self, input_ids: torch.Tensor, substrings) -> tuple[list[torch.Tensor], int]:
        """Locate the `substrings` spans in `input_ids`, in the tensor's own coordinates.

        Runs the substring-to-token-range resolution once for a prepared batch, so a caller
        steering many head maps over one batch (head profiling) locates the spans once and only
        reassembles the hook dict per candidate. The returned ranges are in the coordinate system
        of `input_ids` (its key axis includes any BOS/template and left-pad tokens).

        Args:
            input_ids: Prompt token ids of shape `[batch_size, seq_len]`.
            substrings: The `substrings` runtime-kwarg value, in any of PASTA's accepted forms.

        Returns:
            A `(token_ranges, input_len)` pair, where `token_ranges` is one `[G, 2]` tensor of
            `(start, end)` spans per row and `input_len` is the sequence length.

        Raises:
            ValueError: If `substrings` is malformed (see `_normalize_substrings`).
        """
        batch_size = input_ids.size(0)
        groups = self._normalize_substrings(substrings, batch_size)

        # decode *with* special tokens so offsets share the attention mask's coordinate system
        # (its key axis includes BOS/template tokens)
        prompts = self.tokenizer.decode(input_ids, skip_special_tokens=False)

        # round-trip the substrings through the tokenizer so they match the decoded text
        # (tokenization can subtly change text, e.g. drop spaces)
        for group_idx, group in enumerate(groups):
            try:
                groups[group_idx] = self.tokenizer.decode(
                    self.tokenizer(group, return_tensors="pt", padding=True)["input_ids"],
                    skip_special_tokens=True,
                )
            except Exception as error:
                raise ValueError(
                    f"PASTA failed to re-tokenize substrings {group!r}: {error}"
                ) from error

        # per item: re-encode the decoded text (no specials) and locate token ranges. A faithful
        # fast tokenizer reproduces the real input ids, so its offsets are already in real
        # coordinates; processing per item also avoids any batch padding / padding_side handling.
        token_ranges: list[torch.Tensor] = []
        for item_idx in range(batch_size):
            row_ids = input_ids[item_idx]
            enc = self.tokenizer(
                prompts[item_idx],
                return_tensors="pt",
                return_offsets_mapping=True,
                add_special_tokens=False,
            )
            enc_ids = enc["input_ids"][0].to(row_ids.device)
            offsets = enc["offset_mapping"][0].tolist()

            if torch.equal(enc_ids, row_ids):
                # id-faithful: offsets are already in the real sequence's coordinates
                ranges = [
                    torch.tensor(self._find_token_range(prompts[item_idx], sub, offsets))
                    for sub in groups[item_idx]
                ]
            else:
                # not id-faithful (e.g. sentencepiece whitespace normalization): compute the range
                # in re-encoded space, then relocate the token-id window inside the real ids
                ranges = [
                    torch.tensor(
                        self._locate_range_by_ids(
                            prompts[item_idx], sub, offsets, enc_ids, row_ids
                        )
                    )
                    for sub in groups[item_idx]
                ]
            token_ranges.append(torch.stack(ranges))

        # input_len is the real (padded) sequence length, matching the attention mask's key axis
        input_len = input_ids.size(1)
        return token_ranges, input_len

    def build_hooks_for(
        self,
        token_ranges: list[torch.Tensor],
        input_len: int,
        head_map: dict[int, list[int]],
        alpha: float,
    ) -> dict[str, list]:
        """Assemble the pre-hook dict for a head map at a strength, from located spans.

        The log-alpha constant is bound into each pre-hook partial, so a profiling strength and
        the control's own `alpha` never cross. Only `"pre"` hooks are populated.

        Args:
            token_ranges: Located spans, one `[G, 2]` tensor per row (from `locate_spans`).
            input_len: The sequence length the ranges were located in.
            head_map: Layer index to head indices to steer.
            alpha: The emphasis strength for this hook set.

        Returns:
            Hook specifications with `"pre"`, `"forward"`, `"backward"` keys.
        """
        scale_constant = torch.tensor([alpha], device=self.device, dtype=torch.float32).log()
        hooks: dict[str, list] = {"pre": [], "forward": [], "backward": []}
        for layer, heads in head_map.items():
            hooks["pre"].append(
                {
                    "module": self._attn_module_names[layer],
                    "hook_func": partial(
                        self._attention_pre_hook,
                        head_idx=heads,
                        token_ranges=token_ranges,
                        input_len=input_len,
                        layer_idx=layer,
                        scale_constant=scale_constant,
                    ),
                }
            )
        return hooks

    def attention_layers(self) -> list[int]:
        """The model's attention layers, ascending (available after `steer()`)."""
        return sorted(self._num_heads_by_layer)

    def num_heads_of_layer(self, layer: int) -> int:
        """The attention head count of `layer` (available after `steer()`)."""
        return self._num_heads_by_layer[layer]

    def steer_fits(self) -> tuple[tuple[str, str], ...]:
        """The fit artifacts the steer step produces, for the steer plan.

        A `HeadProfile` head map produces one `("HeadProfile", "direction")` fit (the lift grid,
        classed as a translation-robust direction). A dict or list head map produces no fit. This
        is a pure function of the args, so `check()` can read it before `steer()`.
        """
        if isinstance(self.head_config, HeadProfile):
            return (("HeadProfile", "direction"),)
        return ()

    def export_state(self) -> dict:
        """The resolved profile's lift grid under `"head_profile"` (after a profiled `steer()`).

        Empty for a dict or list head map, or before `steer()` resolves a `HeadProfile`. The full
        `HeadProfileResult` is not serialized here; it stays on `head_profile` and the caller
        writes it through `HeadProfileResult.save`.
        """
        if self.head_profile is not None:
            return {"head_profile": self.head_profile.lift}
        return {}

    def frozen_form(self, state: dict) -> tuple[str, dict]:
        """A same-class `state_control/pasta` frozen form carrying the resolved dict head map.

        Only used when `head_config` is a `HeadProfile`; a dict or list head map keeps the
        `BaseControl` default (the recipe is the frozen form). The loaded control is a plain-dict
        PASTA whose `steer()` does no rollouts.
        """
        if isinstance(self.head_config, HeadProfile):
            return "state_control/pasta", {
                "substrings": self.substrings,
                "head_config": self.head_profile.head_config,
                "alpha": self.alpha,
                "scale_position": self.scale_position,
            }
        return super().frozen_form(state)

    def fit_identity(self):
        """The fit-relevant recipe inputs for staleness detection, or None.

        For a `HeadProfile` head map, the profile's `fit_ingredients()` together with the
        control's `scale_position` (the profile is scored under it). None for a dict or list head
        map. The control's own `alpha` is an application parameter and is excluded.
        """
        if isinstance(self.head_config, HeadProfile):
            return {"profile": self.head_config.fit_ingredients(), "scale_position": self.scale_position}
        return None

    @staticmethod
    def _normalize_substrings(substrings, batch_size: int) -> list[list[str]]:
        """Normalize the `substrings` runtime kwarg to one group per batch row, in a local copy.

        Accepted forms are a `str` (one substring, broadcast to every row), a `list[list[str]]`
        (one group per row, of batch length), and a flat `list[str]` (accepted only at batch
        size 1, as that row's group). Every group must be a non-`str` sequence of `str`; a `str`
        where a group is expected raises rather than being iterated character by character.

        Args:
            substrings: The runtime-kwarg value to normalize.
            batch_size: Number of prompt rows.

        Returns:
            One list of substrings per batch row.

        Raises:
            ValueError: If a flat `list[str]` is passed at batch size > 1 (the message names the
                accepted forms and the `[[...]] * batch_size` broadcast workaround), the number of
                groups does not match the batch size, or any group is a `str` or contains a
                non-`str` element.
        """
        if isinstance(substrings, str):
            return [[substrings] for _ in range(batch_size)]
        if not isinstance(substrings, Sequence):
            raise ValueError(
                f"PASTA 'substrings' must be a str, a list[list[str]] of batch length, or a flat "
                f"list[str] at batch size 1; got {type(substrings).__name__}."
            )
        substrings = list(substrings)
        if all(isinstance(element, str) for element in substrings):
            if batch_size > 1:
                raise ValueError(
                    f"PASTA received a flat list[str] for 'substrings' with batch size {batch_size}. "
                    "Accepted forms are a str (broadcast to every row) and a list[list[str]] with one "
                    "group per row; to broadcast one group over the batch, pass "
                    "[[...]] * batch_size."
                )
            return [substrings]
        if len(substrings) != batch_size:
            raise ValueError(
                f"Need {batch_size} substring groups (one per prompt); got {len(substrings)}"
            )
        groups: list[list[str]] = []
        for group in substrings:
            if isinstance(group, str) or not isinstance(group, Sequence):
                raise ValueError(
                    f"PASTA substring groups must be non-str sequences of str; got "
                    f"{type(group).__name__} for one row."
                )
            group = list(group)
            invalid = [element for element in group if not isinstance(element, str)]
            if invalid:
                raise ValueError(
                    f"PASTA substring groups must contain only str elements; got "
                    f"{type(invalid[0]).__name__} in one row's group."
                )
            groups.append(group)
        return groups

    def _setup_layers(self, head_config) -> None:
        """Derive the configured layer set (and explicit head lists for the dict form).

        The head map is finalized in `_finalize_head_map` once the per-layer head counts are
        known; the list form defers to that step for its all-heads default.

        Args:
            head_config: Configuration specifying which layers/heads to modify:

                - dict: Maps layer indices to lists of head indices
                - list: Layer indices (applies to all heads in those layers)

        Raises:
            ValueError: If the configuration format is invalid.
        """
        if isinstance(head_config, dict):
            self._head_map = {int(layer): list(heads) for layer, heads in head_config.items()}
            self._layers = sorted(self._head_map.keys())
        elif isinstance(head_config, list):
            self._layers = [int(layer) for layer in head_config]
            self._head_map = None
        else:
            raise ValueError(f"Invalid head configuration: {head_config!r}")

    def _finalize_head_map(self) -> None:
        """Validate the configured layers, fill the all-heads default, and check head indices.

        Every configured layer must be an attention layer of the model (its head count is in
        `_num_heads_by_layer`, recorded in `_resolve_attention_modules`); a layer out of range or
        without an attention module raises. The list form of `head_config` then expands to every
        head of each configured layer, and the dict form's explicit head lists are validated,
        both against that layer's own head count. Models whose head count varies across layers are
        handled per layer.

        Raises:
            ValueError: If a configured layer is out of range or carries no attention module, or a
                head index is out of range for its layer.
        """
        num_layers = self._layout.num_layers
        for layer in self._layers:
            if layer < 0 or layer >= num_layers:
                raise ValueError(
                    f"PASTA layer {layer} out of range for model with {num_layers} layers."
                )
            if layer not in self._num_heads_by_layer:
                raise ValueError(
                    f"PASTA layer {layer} carries no attention module; attention layers of this "
                    f"model are {list(self._layout.attention_layers)}."
                )

        if self._head_map is None:  # list form: all heads of each configured layer
            self._head_map = {
                layer: list(range(self._num_heads_by_layer[layer])) for layer in self._layers
            }

        for layer, heads in self._head_map.items():
            num_heads = self._num_heads_by_layer[layer]
            for head in heads:
                if not 0 <= head < num_heads:
                    raise ValueError(
                        f"Head {head} out of range for layer {layer} (0-{num_heads - 1})"
                    )

    @staticmethod
    def _find_token_range(
        string: str,
        substring: str,
        offset_mapping: Sequence[tuple[int, int]],
        occurrence: int = 0,
    ) -> tuple[int, int]:
        """Map a substring to its token index range using offset mapping.

        Locates the character positions of a substring and converts them to token indices using the tokenizer's offset mapping.

        Args:
            string: Full text to search within.
            substring: Target substring to locate.
            offset_mapping: List of (start_char, end_char) tuples for each token.
            occurrence: Which occurrence to find if substring appears multiple times.
                Defaults to 0 (first occurrence).

        Returns:
            tuple[int, int]: Start (inclusive) and end (exclusive) token indices. If the substring is
                absent from `string`, returns the `(0, 0)` sentinel (skipped downstream) after warning.

        Raises:
            ValueError: If the substring is present in `string` but cannot be aligned to token offsets.
        """
        if substring not in string:
            logger.warning(
                "PASTA: substring %r not found in input (len=%d chars); skipping steering for this range.",
                substring, len(string),
            )
            return 0, 0

        char_index = -1
        for _ in range(occurrence + 1):
            char_index = string.index(substring, char_index + 1)
        char_start = char_index
        char_end = char_start + len(substring)

        token_start = token_end = None
        for token_idx, (start_char, end_char) in enumerate(offset_mapping):
            if token_start is None and start_char <= char_start < end_char:
                token_start = token_idx
            if token_end is None and start_char < char_end <= end_char:
                token_end = token_idx

        if token_start is None or token_end is None:
            raise ValueError("Could not map substring to token range")

        return token_start, token_end + 1

    @staticmethod
    def _locate_range_by_ids(
        text: str,
        substring: str,
        offsets: Sequence[tuple[int, int]],
        enc_ids: torch.Tensor,
        real_ids: torch.Tensor,
    ) -> tuple[int, int]:
        """Fallback range resolution when re-encoding is not id-faithful to the real sequence.

        Computes the token range in the re-encoded (special-free) space, extracts that window of
        token ids, and finds the matching contiguous id window inside the real sequence (choosing
        the occurrence nearest the naive index). Returns the `(0, 0)` skip sentinel if the
        substring is absent or its id window cannot be located in the real ids.

        Args:
            text: The decoded text the offsets were computed against.
            substring: The target substring.
            offsets: Offset mapping for the re-encoded text.
            enc_ids: Token ids of the re-encoded text (same coordinate system as `offsets`).
            real_ids: Token ids of the real (padded) sequence to relocate the window into.

        Returns:
            `(start, end)` token indices in the real sequence, or `(0, 0)` if unresolved.
        """
        naive_start, naive_end = PASTA._find_token_range(text, substring, offsets)
        if naive_start == naive_end:  # absent (already warned) or empty
            return 0, 0

        window = enc_ids[naive_start:naive_end]
        w_len = window.size(0)
        real = real_ids.tolist()
        window_list = window.tolist()

        # every contiguous position in real_ids that matches the window
        matches = [
            i for i in range(len(real) - w_len + 1)
            if real[i:i + w_len] == window_list
        ]
        if not matches:
            logger.warning(
                "PASTA: could not relocate substring %r window into the real sequence; skipping.",
                substring,
            )
            return 0, 0

        # pick the occurrence nearest the naive index
        best = min(matches, key=lambda i: abs(i - naive_start))
        return best, best + w_len

    def _attention_pre_hook(
        self,
        module,
        input_args: tuple,
        input_kwargs: dict,
        head_idx: list[int],
        token_ranges: list[torch.Tensor],
        input_len: int,
        scale_constant: torch.Tensor,
        layer_idx: int = 0,
    ):
        """Modify attention mask to steer focus toward/away from target tokens.

        Pre-forward hook that adjusts attention weights by adding scaling factors to the attention mask for specified token ranges and attention heads.

        In `"include"` mode each prompt column is edited once: the highlighted span columns are
        left untouched and `scale_constant` is subtracted from the complement (the non-span prompt
        columns), the paper's Eq. 3 (non-highlighted prompt columns scaled by the coefficient).
        Overlapping spans therefore net to zero, and no column is touched twice, which removes the
        bfloat16 residue of a double touch. `"exclude"` and `"generation"` add `scale_constant` to
        the non-span columns and the whole prompt respectively.

        Args:
            module: The attention module being hooked.
            input_args: Positional arguments to the forward pass.
            input_kwargs: Keyword arguments to the forward pass.
            head_idx: List of attention head indices to modify.
            token_ranges: Token index ranges to apply scaling to.
            input_len: Length of input sequence (for generation positioning).
            scale_constant: The `log(alpha)` edit magnitude, bound into the hook so a profiling
                strength and the control's own `alpha` cannot cross.
            layer_idx: Decoder layer index of the hooked attention module, used to read this
                layer's cached key length when `cache_position` is absent (transformers v5).

        Returns:
            Tuple of potentially modified (input_args, input_kwargs).

        Raises:
            RuntimeError: If hidden states cannot be located.
            ValueError: If scale_position is invalid.
        """
        hidden_states = (
            input_args[0] if input_args else input_kwargs.get("hidden_states")
        )
        if hidden_states is None:
            raise RuntimeError("PASTA: could not locate hidden states")

        attention_mask = input_kwargs.get("attention_mask")
        if attention_mask is None:  # build it
            batch_size, query_len, _ = hidden_states.size()
            num_heads = self._num_heads_by_layer[layer_idx]

            # during decoding the query attends to the full kv cache, so the mask spans the cached key
            # positions rather than just the current query window. transformers v4 exposes the key axis
            # via cache_position; v5 removed that kwarg from attention calls, so fall back to this
            # layer's cache length (the pre-hook runs before the layer's cache update, so cached length
            # plus query_len is the post-update key length). an undersized mask would rely on sdpa
            # broadcasting, and the cuda mem-efficient kernel rejects the stride-0 last dimension that
            # its mask expansion produces ("(*bias): last dimension must be contiguous")
            cache_position = input_kwargs.get("cache_position")
            past_key_values = input_kwargs.get("past_key_values")
            if cache_position is not None:
                key_len = int(cache_position[-1]) + 1
            elif past_key_values is not None and callable(getattr(past_key_values, "get_seq_length", None)):
                key_len = int(past_key_values.get_seq_length(layer_idx) or 0) + query_len
            else:
                key_len = query_len

            # query row i sits at absolute position (key_len - query_len + i) and attends to keys 0..position
            query_positions = torch.arange(
                key_len - query_len, key_len, device=hidden_states.device
            ).unsqueeze(1)
            key_positions = torch.arange(key_len, device=hidden_states.device).unsqueeze(0)
            causal = torch.where(
                key_positions <= query_positions,
                hidden_states.new_zeros(()),
                hidden_states.new_full((), float("-inf")),
            )
            attention_mask = causal[None, None]  # (1,1,q,k)
            attention_mask = attention_mask.expand(
                batch_size, num_heads, -1, -1
            ).contiguous()
            input_kwargs["attention_mask"] = attention_mask

        if attention_mask.dtype == torch.bool:
            # transformers v5 materializes boolean masks (true = attend) for sdpa; convert to the
            # additive convention before applying the log-alpha scaling edits
            attention_mask = torch.where(
                attention_mask,
                hidden_states.new_zeros(()),
                hidden_states.new_full((), torch.finfo(hidden_states.dtype).min),
            )
        attention_mask = attention_mask.to(hidden_states.dtype).contiguous().clone()
        if attention_mask.size(1) == 1:
            num_heads = self._num_heads_by_layer[layer_idx]
            attention_mask = attention_mask.expand(-1, num_heads, -1, -1).contiguous()

        batch_size = attention_mask.size(0)

        # beam search expands the batch via repeat_interleave ([item0, item0, item1, item1]); index
        # token_ranges the same way (a modulo would misorder the beams)
        if batch_size > len(token_ranges):
            if batch_size % len(token_ranges) != 0:
                raise RuntimeError(
                    f"Hidden batch {batch_size} is not a multiple of the prompt batch "
                    f"{len(token_ranges)}; cannot align PASTA token ranges."
                )
            expand = batch_size // len(token_ranges)
            token_ranges = [token_ranges[i // expand] for i in range(batch_size)]

        if self.scale_position not in ("include", "exclude", "generation"):
            raise ValueError(f"Unknown scale_position '{self.scale_position}'")

        for batch_index in range(batch_size):
            ranges = token_ranges[batch_index].tolist()
            has_valid_range = any(start != end for start, end in ranges)
            if not has_valid_range:
                continue

            if self.scale_position == "include":
                # edit each prompt column once: subtract on the complement of the highlighted
                # span union, leaving the span columns untouched (overlapping spans net to zero).
                # a per-column delta applied over the [:input_len] slice keeps the write a single
                # subscription (advanced head indexing writes back only through one __setitem__)
                delta = attention_mask.new_full((input_len,), 0.0)
                delta[:] = -scale_constant
                for start_idx, end_idx in ranges:
                    if start_idx != end_idx:
                        delta[start_idx:end_idx] = 0.0
                attention_mask[batch_index, head_idx, :, :input_len] += delta
            elif self.scale_position == "exclude":
                for start_idx, end_idx in ranges:
                    if start_idx == end_idx:
                        continue
                    attention_mask[batch_index, head_idx, :, :start_idx] += scale_constant
                    attention_mask[batch_index, head_idx, :, end_idx:input_len] += scale_constant
            else:  # generation
                attention_mask[batch_index, head_idx, :, :input_len] += scale_constant

        input_kwargs["attention_mask"] = attention_mask
        return input_args, input_kwargs

RUNTIME_KWARGS_SCHEMA = [{'name': 'substrings', 'type': 'list[str]', 'required': True, 'scope': 'row', 'help': "Substrings whose attention should be steered. Required at inference time. A `str` broadcasts to every row; a `list[list[str]]` of batch length carries one group per row; a flat `list[str]` is accepted only at batch size 1, as that row's group."}] class-attribute instance-attribute

args = self.Args.validate(*args, **kwargs) instance-attribute

device = None class-attribute instance-attribute

enabled = True class-attribute instance-attribute

head_profile = None class-attribute instance-attribute

model = None class-attribute instance-attribute

supports_batching = True class-attribute instance-attribute

tokenizer = None class-attribute instance-attribute

attention_layers()

The model's attention layers, ascending (available after steer()).

Source code in steerability/algorithms/state_control/pasta/control.py

def attention_layers(self) -> list[int]:
    """The model's attention layers, ascending (available after `steer()`)."""
    return sorted(self._num_heads_by_layer)

build_hooks_for(token_ranges, input_len, head_map, alpha)

Assemble the pre-hook dict for a head map at a strength, from located spans.

The log-alpha constant is bound into each pre-hook partial, so a profiling strength and the control's own alpha never cross. Only "pre" hooks are populated.

Parameters:

Name	Type	Description	Default
token_ranges	list[Tensor]	Located spans, one [G, 2] tensor per row (from locate_spans).	required
input_len	int	The sequence length the ranges were located in.	required
head_map	dict[int, list[int]]	Layer index to head indices to steer.	required
alpha	float	The emphasis strength for this hook set.	required

Returns:

Type	Description
dict[str, list]	Hook specifications with "pre", "forward", "backward" keys.

Source code in steerability/algorithms/state_control/pasta/control.py

def build_hooks_for(
    self,
    token_ranges: list[torch.Tensor],
    input_len: int,
    head_map: dict[int, list[int]],
    alpha: float,
) -> dict[str, list]:
    """Assemble the pre-hook dict for a head map at a strength, from located spans.

    The log-alpha constant is bound into each pre-hook partial, so a profiling strength and
    the control's own `alpha` never cross. Only `"pre"` hooks are populated.

    Args:
        token_ranges: Located spans, one `[G, 2]` tensor per row (from `locate_spans`).
        input_len: The sequence length the ranges were located in.
        head_map: Layer index to head indices to steer.
        alpha: The emphasis strength for this hook set.

    Returns:
        Hook specifications with `"pre"`, `"forward"`, `"backward"` keys.
    """
    scale_constant = torch.tensor([alpha], device=self.device, dtype=torch.float32).log()
    hooks: dict[str, list] = {"pre": [], "forward": [], "backward": []}
    for layer, heads in head_map.items():
        hooks["pre"].append(
            {
                "module": self._attn_module_names[layer],
                "hook_func": partial(
                    self._attention_pre_hook,
                    head_idx=heads,
                    token_ranges=token_ranges,
                    input_len=input_len,
                    layer_idx=layer,
                    scale_constant=scale_constant,
                ),
            }
        )
    return hooks

cleanup()

Release resources allocated during steer().

Override in subclasses that allocate GPU memory or other resources during steering. Default no-op.

Source code in steerability/algorithms/core/base_control.py

def cleanup(self) -> None:
    """Release resources allocated during `steer()`.

    Override in subclasses that allocate GPU memory or other resources during steering. Default
    no-op.
    """
    pass

clone_for_call(seed=None)

A configuration-preserving shallow clone for one generation call.

The clone shares steer-time artifacts (memories, steering vectors, attached tokenizers) with the original but has its own attribute namespace, so per-call attribute mutation on the clone never races another call using the original. When seed is given and the control defines reseed(seed), the clone's client-side RNG is re-seeded.

Parameters:

Name	Type	Description	Default
seed	int | None	Optional seed forwarded to the clone's reseed().	None

Returns:

Type	Description
BaseControl	The clone.

Source code in steerability/algorithms/core/base_control.py

def clone_for_call(self, seed: int | None = None) -> "BaseControl":
    """A configuration-preserving shallow clone for one generation call.

    The clone shares steer-time artifacts (memories, steering vectors, attached tokenizers)
    with the original but has its own attribute namespace, so per-call attribute mutation on
    the clone never races another call using the original. When `seed` is given and the
    control defines `reseed(seed)`, the clone's client-side RNG is re-seeded.

    Args:
        seed: Optional seed forwarded to the clone's `reseed()`.

    Returns:
        The clone.
    """
    clone = copy.copy(self)
    if seed is not None:
        reseed = getattr(clone, "reseed", None)
        if callable(reseed):
            reseed(seed)
    return clone

export_intervention_spec(runtime_kwargs=None)

The control's InterventionSpec for intervention-capable backends, or None.

The spec is the second serialization of the tuple the control's hooks close over, emitted from the same transform, gate, and scope objects. Must be called after steer(). Returns None when the configuration has no wire form (the configuration is then hook-only) or when the control does not implement spec export at all.

Parameters:

Name	Type	Description	Default
runtime_kwargs	dict | None	Per-call parameters, mirroring get_hooks; per-item values (strengths, positions) serialize into the returned spec.	None

Returns:

Type	Description
InterventionSpec | None	The validated InterventionSpec with tensor payloads attached, or None.

Source code in steerability/algorithms/state_control/base.py

def export_intervention_spec(self, runtime_kwargs: dict | None = None) -> "InterventionSpec | None":
    """The control's `InterventionSpec` for intervention-capable backends, or None.

    The spec is the second serialization of the tuple the control's hooks close over,
    emitted from the same transform, gate, and scope objects. Must be called after
    `steer()`. Returns None when the configuration has no wire form (the configuration is
    then hook-only) or when the control does not implement spec export at all.

    Args:
        runtime_kwargs: Per-call parameters, mirroring `get_hooks`; per-item values
            (strengths, positions) serialize into the returned spec.

    Returns:
        The validated `InterventionSpec` with tensor payloads attached, or None.
    """
    return None

export_state()

The resolved profile's lift grid under "head_profile" (after a profiled steer()).

Empty for a dict or list head map, or before steer() resolves a HeadProfile. The full HeadProfileResult is not serialized here; it stays on head_profile and the caller writes it through HeadProfileResult.save.

Source code in steerability/algorithms/state_control/pasta/control.py

def export_state(self) -> dict:
    """The resolved profile's lift grid under `"head_profile"` (after a profiled `steer()`).

    Empty for a dict or list head map, or before `steer()` resolves a `HeadProfile`. The full
    `HeadProfileResult` is not serialized here; it stays on `head_profile` and the caller
    writes it through `HeadProfileResult.save`.
    """
    if self.head_profile is not None:
        return {"head_profile": self.head_profile.lift}
    return {}

fit_identity()

The fit-relevant recipe inputs for staleness detection, or None.

For a HeadProfile head map, the profile's fit_ingredients() together with the control's scale_position (the profile is scored under it). None for a dict or list head map. The control's own alpha is an application parameter and is excluded.

Source code in steerability/algorithms/state_control/pasta/control.py

def fit_identity(self):
    """The fit-relevant recipe inputs for staleness detection, or None.

    For a `HeadProfile` head map, the profile's `fit_ingredients()` together with the
    control's `scale_position` (the profile is scored under it). None for a dict or list head
    map. The control's own `alpha` is an application parameter and is excluded.
    """
    if isinstance(self.head_config, HeadProfile):
        return {"profile": self.head_config.fit_ingredients(), "scale_position": self.scale_position}
    return None

frozen_form(state)

A same-class state_control/pasta frozen form carrying the resolved dict head map.

Only used when head_config is a HeadProfile; a dict or list head map keeps the BaseControl default (the recipe is the frozen form). The loaded control is a plain-dict PASTA whose steer() does no rollouts.

Source code in steerability/algorithms/state_control/pasta/control.py

def frozen_form(self, state: dict) -> tuple[str, dict]:
    """A same-class `state_control/pasta` frozen form carrying the resolved dict head map.

    Only used when `head_config` is a `HeadProfile`; a dict or list head map keeps the
    `BaseControl` default (the recipe is the frozen form). The loaded control is a plain-dict
    PASTA whose `steer()` does no rollouts.
    """
    if isinstance(self.head_config, HeadProfile):
        return "state_control/pasta", {
            "substrings": self.substrings,
            "head_config": self.head_profile.head_config,
            "alpha": self.alpha,
            "scale_position": self.scale_position,
        }
    return super().frozen_form(state)

get_hooks(input_ids, runtime_kwargs, **__)

Create attention modification hooks for specified substrings.

Identifies token ranges corresponding to target substrings and prepares hooks that will modify attention weights during the forward pass, at the configured head map and the control's alpha.

Parameters:

Name	Type	Description	Default
input_ids	Tensor	Input token IDs of shape [batch_size, seq_len].	required
runtime_kwargs	dict | None	Must contain "substrings" key with target text spans: str: one substring, broadcast to every batch item list[list[str]]: one substring group per batch item, of batch length list[str]: accepted only at batch size 1, as that row's group	required

Returns:

Type	Description
dict[str, list]	dict[str, list]: Hook specifications with "pre", "forward", "backward" keys. Only "pre" hooks are populated for attention modification.

Raises:

Type	Description
ValueError	If "substrings" is missing from runtime_kwargs, a flat list of strings is passed at batch size > 1, a group is a str or contains a non-str element, or the number of groups does not match the batch size.

Source code in steerability/algorithms/state_control/pasta/control.py

def get_hooks(
    self,
    input_ids: torch.Tensor,
    runtime_kwargs: dict | None,
    **__,
) -> dict[str, list]:
    """Create attention modification hooks for specified substrings.

    Identifies token ranges corresponding to target substrings and prepares hooks that will modify attention weights
    during the forward pass, at the configured head map and the control's `alpha`.

    Args:
        input_ids (torch.Tensor): Input token IDs of shape [batch_size, seq_len].
        runtime_kwargs (dict | None): Must contain "substrings" key with target text spans:

            - str: one substring, broadcast to every batch item
            - list[list[str]]: one substring group per batch item, of batch length
            - list[str]: accepted only at batch size 1, as that row's group

    Returns:
        dict[str, list]: Hook specifications with "pre", "forward", "backward" keys. Only "pre" hooks are populated for attention modification.

    Raises:
        ValueError: If "substrings" is missing from runtime_kwargs, a flat list of strings is
            passed at batch size > 1, a group is a `str` or contains a non-`str` element, or
            the number of groups does not match the batch size.
    """
    if not runtime_kwargs or "substrings" not in runtime_kwargs:
        raise ValueError("PASTA requires 'substrings' inside runtime_kwargs")

    token_ranges, input_len = self.locate_spans(input_ids, runtime_kwargs["substrings"])
    return self.build_hooks_for(token_ranges, input_len, self._head_map, self.alpha)

locate_spans(input_ids, substrings)

Locate the substrings spans in input_ids, in the tensor's own coordinates.

Runs the substring-to-token-range resolution once for a prepared batch, so a caller steering many head maps over one batch (head profiling) locates the spans once and only reassembles the hook dict per candidate. The returned ranges are in the coordinate system of input_ids (its key axis includes any BOS/template and left-pad tokens).

Parameters:

Name	Type	Description	Default
input_ids	Tensor	Prompt token ids of shape [batch_size, seq_len].	required
substrings		The substrings runtime-kwarg value, in any of PASTA's accepted forms.	required

Returns:

Type	Description
list[Tensor]	A (token_ranges, input_len) pair, where token_ranges is one [G, 2] tensor of
int	(start, end) spans per row and input_len is the sequence length.

Raises:

Type	Description
ValueError	If substrings is malformed (see _normalize_substrings).

Source code in steerability/algorithms/state_control/pasta/control.py

def locate_spans(self, input_ids: torch.Tensor, substrings) -> tuple[list[torch.Tensor], int]:
    """Locate the `substrings` spans in `input_ids`, in the tensor's own coordinates.

    Runs the substring-to-token-range resolution once for a prepared batch, so a caller
    steering many head maps over one batch (head profiling) locates the spans once and only
    reassembles the hook dict per candidate. The returned ranges are in the coordinate system
    of `input_ids` (its key axis includes any BOS/template and left-pad tokens).

    Args:
        input_ids: Prompt token ids of shape `[batch_size, seq_len]`.
        substrings: The `substrings` runtime-kwarg value, in any of PASTA's accepted forms.

    Returns:
        A `(token_ranges, input_len)` pair, where `token_ranges` is one `[G, 2]` tensor of
        `(start, end)` spans per row and `input_len` is the sequence length.

    Raises:
        ValueError: If `substrings` is malformed (see `_normalize_substrings`).
    """
    batch_size = input_ids.size(0)
    groups = self._normalize_substrings(substrings, batch_size)

    # decode *with* special tokens so offsets share the attention mask's coordinate system
    # (its key axis includes BOS/template tokens)
    prompts = self.tokenizer.decode(input_ids, skip_special_tokens=False)

    # round-trip the substrings through the tokenizer so they match the decoded text
    # (tokenization can subtly change text, e.g. drop spaces)
    for group_idx, group in enumerate(groups):
        try:
            groups[group_idx] = self.tokenizer.decode(
                self.tokenizer(group, return_tensors="pt", padding=True)["input_ids"],
                skip_special_tokens=True,
            )
        except Exception as error:
            raise ValueError(
                f"PASTA failed to re-tokenize substrings {group!r}: {error}"
            ) from error

    # per item: re-encode the decoded text (no specials) and locate token ranges. A faithful
    # fast tokenizer reproduces the real input ids, so its offsets are already in real
    # coordinates; processing per item also avoids any batch padding / padding_side handling.
    token_ranges: list[torch.Tensor] = []
    for item_idx in range(batch_size):
        row_ids = input_ids[item_idx]
        enc = self.tokenizer(
            prompts[item_idx],
            return_tensors="pt",
            return_offsets_mapping=True,
            add_special_tokens=False,
        )
        enc_ids = enc["input_ids"][0].to(row_ids.device)
        offsets = enc["offset_mapping"][0].tolist()

        if torch.equal(enc_ids, row_ids):
            # id-faithful: offsets are already in the real sequence's coordinates
            ranges = [
                torch.tensor(self._find_token_range(prompts[item_idx], sub, offsets))
                for sub in groups[item_idx]
            ]
        else:
            # not id-faithful (e.g. sentencepiece whitespace normalization): compute the range
            # in re-encoded space, then relocate the token-id window inside the real ids
            ranges = [
                torch.tensor(
                    self._locate_range_by_ids(
                        prompts[item_idx], sub, offsets, enc_ids, row_ids
                    )
                )
                for sub in groups[item_idx]
            ]
        token_ranges.append(torch.stack(ranges))

    # input_len is the real (padded) sequence length, matching the attention mask's key axis
    input_len = input_ids.size(1)
    return token_ranges, input_len

num_heads_of_layer(layer)

The attention head count of layer (available after steer()).

Source code in steerability/algorithms/state_control/pasta/control.py

def num_heads_of_layer(self, layer: int) -> int:
    """The attention head count of `layer` (available after `steer()`)."""
    return self._num_heads_by_layer[layer]

requirements()

Backend requirements for attention-map editing.

PASTA writes into attention maps through torch hooks, which fused paged-attention kernels never materialize, so the generate phase requires Capability.IN_PROCESS_TORCH. The spec constraint reports a configured incompatible attention implementation (hf_model_kwargs["attn_implementation"] outside "eager"/"sdpa") at check() time, before any model loads; the same condition is re-checked against the live model in steer().

Returns:

Type	Description
Requirements	The control's phase-keyed requirements.

Source code in steerability/algorithms/state_control/pasta/control.py

def requirements(self) -> Requirements:
    """Backend requirements for attention-map editing.

    PASTA writes into attention maps through torch hooks, which fused paged-attention
    kernels never materialize, so the generate phase requires `Capability.IN_PROCESS_TORCH`.
    The spec constraint reports a configured incompatible attention implementation
    (`hf_model_kwargs["attn_implementation"]` outside `"eager"`/`"sdpa"`) at `check()` time,
    before any model loads; the same condition is re-checked against the live model in
    `steer()`.

    Returns:
        The control's phase-keyed requirements.
    """
    return Requirements(
        generate=needs(Capability.IN_PROCESS_TORCH),
        spec_constraints=(
            SpecConstraint(
                description=(
                    "PASTA requires attn_implementation 'eager' or 'sdpa' to inject a 4D "
                    "attention mask; set attn_implementation=\"eager\" in hf_model_kwargs."
                ),
                predicate=_attn_implementation_supported,
            ),
        ),
    )

steer(model, tokenizer=None, session=None, **__)

Initialize PASTA by configuring attention head mappings and model references.

Resolves the attention module path and head count of every attention layer of the model, resolves a HeadProfile head map when head_config is a profiling recipe (steering each candidate head through the pipeline's session), sets up the resolved head map, and fails fast on unsupported layers or attention implementations (rather than deep inside generation).

Parameters:

Name	Type	Description	Default
model	PreTrainedModel	The base language model to be steered.	required
tokenizer	PreTrainedTokenizerBase | None	Tokenizer for substring identification. If None, attempts to retrieve from model attributes.	None
session		The ScopedSession the pipeline hands to steer(), scoped to ModelAccess.MODULE, through which a HeadProfile runs its rollouts.	None

Returns:

Name	Type	Description
PreTrainedModel	PreTrainedModel	The input model (unchanged).

Raises:

Type	Description
ValueError	If a configured attention module path is missing on the model, if the model's attention implementation is not one of "eager" / "sdpa", or if a HeadProfile selects no heads.

Source code in steerability/algorithms/state_control/pasta/control.py

def steer(
    self,
    model: PreTrainedModel,
    tokenizer: PreTrainedTokenizerBase | None = None,
    session=None,
    **__,
) -> PreTrainedModel:
    """Initialize PASTA by configuring attention head mappings and model references.

    Resolves the attention module path and head count of every attention layer of the model,
    resolves a `HeadProfile` head map when `head_config` is a profiling recipe (steering each
    candidate head through the pipeline's session), sets up the resolved head map, and fails
    fast on unsupported layers or attention implementations (rather than deep inside
    generation).

    Args:
        model (PreTrainedModel): The base language model to be steered.
        tokenizer (PreTrainedTokenizerBase | None): Tokenizer for substring identification.
            If None, attempts to retrieve from model attributes.
        session: The `ScopedSession` the pipeline hands to `steer()`, scoped to
            `ModelAccess.MODULE`, through which a `HeadProfile` runs its rollouts.

    Returns:
        PreTrainedModel: The input model (unchanged).

    Raises:
        ValueError: If a configured attention module path is missing on the model, if the
            model's attention implementation is not one of `"eager"` / `"sdpa"`, or if a
            `HeadProfile` selects no heads.
    """
    self.model = model
    self.tokenizer = tokenizer or getattr(model, "tokenizer", None)
    self.device = next(model.parameters()).device
    self._resolve_attention_modules(model)
    self._check_attention_implementation(model)

    if isinstance(self.head_config, HeadProfile):
        self.head_profile = self.head_config.resolve(self, model, self.tokenizer, session=session)
        head_map = self.head_profile.head_config
    else:
        head_map = self.head_config

    self._setup_layers(head_map)
    self._finalize_head_map()
    return model

steer_access()

ModelAccess.MODULE; the attention module paths resolve on the live model, which is retained for the hook closures (the generate phase is in-process).

Source code in steerability/algorithms/state_control/pasta/control.py

def steer_access(self) -> ModelAccess:
    """`ModelAccess.MODULE`; the attention module paths resolve on the live model, which
    is retained for the hook closures (the generate phase is in-process)."""
    return ModelAccess.MODULE

steer_fits()

The fit artifacts the steer step produces, for the steer plan.

A HeadProfile head map produces one ("HeadProfile", "direction") fit (the lift grid, classed as a translation-robust direction). A dict or list head map produces no fit. This is a pure function of the args, so check() can read it before steer().

Source code in steerability/algorithms/state_control/pasta/control.py

def steer_fits(self) -> tuple[tuple[str, str], ...]:
    """The fit artifacts the steer step produces, for the steer plan.

    A `HeadProfile` head map produces one `("HeadProfile", "direction")` fit (the lift grid,
    classed as a translation-robust direction). A dict or list head map produces no fit. This
    is a pure function of the args, so `check()` can read it before `steer()`.
    """
    if isinstance(self.head_config, HeadProfile):
        return (("HeadProfile", "direction"),)
    return ()

profiling

Rollout-scored attention-head profiling for PASTA, resolved at steer time.

HeadProfile is a third accepted form of PASTAArgs.head_config, beside the dict and list forms. It is a declarative recipe that PASTA.steer() resolves through the pipeline's session: each candidate (layer, head) is steered on its own on a task-agnostic set of profiling prompts, scored by a SampleScorer, and ranked by the paired lift of its score over an unsteered baseline. The selected heads become the control's dict-form head map, and the resolution HeadProfileResult records the per-head lift, standard error, and selection so the run can be inspected and rebuilt.

The recipe follows the toolkit's source idiom (ContrastiveFit, ConditionPointSearch): it declares access and artifact_class for the steer plan, memoizes its resolution per model, and its resolved head map freezes into a .spipe as a same-class state_control/pasta entry with the fit digest on the lift grid.

Reference:

- "Tell Your Model Where to Attend: Post-hoc Attention Steering for LLMs"
  Qingru Zhang, Chandan Singh, Liyuan Liu, Xiaodong Liu, Bin Yu, Jianfeng Gao, Tuo Zhao
  [https://arxiv.org/abs/2311.02262](https://arxiv.org/abs/2311.02262)

logger = logging.getLogger(__name__) module-attribute

HeadProfile dataclass

Rollout-scored head profiling, resolved at PASTA.steer().

A recipe passed as PASTA(head_config=HeadProfile(...)). It scores each candidate (layer, head) by steering that head alone at the profiling strength alpha and measuring the mean gain in scorer over an unsteered baseline on rows, ranks the candidates deterministically, filters to those that beat the baseline, and returns the top num_heads as a dict head map. Resolution runs through the session the pipeline hands to steer(), so the rollouts execute on the steering backend.

The paired lift and its standard error are the statistic; ranking by lift and ranking by the raw follow rate order candidates identically (the baseline is constant across candidates), but the pairing gives a standard error the raw rate does not. A two-stage screen bounds cost: stage 1 scores every candidate on a fixed subset of screen_rows rows, and stage 2 rescores only the top screen_keep candidates on every row.

The profile uses the control's scale_position; it is not a recipe field, since a profile is only meaningful for the mechanism it is deployed with.

Attributes:

Name	Type	Description
rows	Sequence[Mapping]	One mapping per profiling prompt. "input" is the user turn (the SampleScorer row convention); "substrings" is the row's PASTA runtime kwarg in its per-row form (list[str]). An optional "group" enables the per-group statistics and the "intersection" selection. Every other key passes through to the scorer. The rows encode into a frozen .spipe inline, so the set must stay under the codec's per-entry inline limit (1 MB, about a few thousand short prompts).
scorer	SampleScorer	(response, row) -> float, higher is better. For a strict pass/fail checker this is 0.0 or 1.0; a loose checker or a reward score are drop-in alternatives with finer granularity.
alpha	float	The profiling strength. Required, with no default. The toolkit parameterizes alpha as the reciprocal of the paper's coefficient, so alpha=100.0 is the paper's operating point (coefficient 0.01). Stating it here keeps every point of an alpha sweep the same fit, so the profile is resolved once per run.
num_heads	int	Number of heads to select.
layers	Sequence[int] | None	Candidate layers. None means every attention layer of the resolved layout; non-attention layers of a hybrid stack are never candidates.
selection	Literal['pooled', 'intersection']	"pooled" selects the first num_heads eligible candidates in ranking order. "intersection" is the paper's rule and requires a "group" on every row.
min_lift	float	A candidate is eligible when its pooled lift exceeds this value. The default 0.0 keeps only heads that beat the baseline.
screen_rows	int | None	Rows scored in stage 1 of the two-stage screen. Set together with screen_keep, or leave both unset for a single-stage profile over every row.
screen_keep	int | None	Candidates rescored on every row in stage 2. At least num_heads; 2 * num_heads or more is recommended, since the stage-1 standard error is large enough to push a true positive out of a stage-2 pool sized exactly to the target.
gen_kwargs	dict	Profiling generation parameters, normalized through GenerationParams.from_gen_kwargs. Greedy by default.
batch_size	int	Number of rows generated per session call. Rows in a batch share one hook entry and generate in one batched pass, so larger values trade memory for wall-clock; excluded from the fit digest.
seed	int	Fixes the stage-1 subsample.
progress_callback	Callable[[dict], None] | None	Optional callback invoked once per scored candidate with a small dict (stage, candidate, completed, total, lift). Opt-in; no-op when None.

Reference:

- "Tell Your Model Where to Attend: Post-hoc Attention Steering for LLMs"
  Qingru Zhang, Chandan Singh, Liyuan Liu, Xiaodong Liu, Bin Yu, Jianfeng Gao, Tuo Zhao
  [https://arxiv.org/abs/2311.02262](https://arxiv.org/abs/2311.02262)

Source code in steerability/algorithms/state_control/pasta/profiling.py

@dataclass
class HeadProfile:
    """Rollout-scored head profiling, resolved at `PASTA.steer()`.

    A recipe passed as `PASTA(head_config=HeadProfile(...))`. It scores each candidate
    `(layer, head)` by steering that head alone at the profiling strength `alpha` and measuring
    the mean gain in `scorer` over an unsteered baseline on `rows`, ranks the candidates
    deterministically, filters to those that beat the baseline, and returns the top `num_heads`
    as a dict head map. Resolution runs through the session the pipeline hands to `steer()`, so
    the rollouts execute on the steering backend.

    The paired lift and its standard error are the statistic; ranking by lift and ranking by the
    raw follow rate order candidates identically (the baseline is constant across candidates),
    but the pairing gives a standard error the raw rate does not. A two-stage screen bounds cost:
    stage 1 scores every candidate on a fixed subset of `screen_rows` rows, and stage 2 rescores
    only the top `screen_keep` candidates on every row.

    The profile uses the control's `scale_position`; it is not a recipe field, since a profile is
    only meaningful for the mechanism it is deployed with.

    Attributes:
        rows: One mapping per profiling prompt. `"input"` is the user turn (the `SampleScorer`
            row convention); `"substrings"` is the row's PASTA runtime kwarg in its per-row form
            (`list[str]`). An optional `"group"` enables the per-group statistics and the
            `"intersection"` selection. Every other key passes through to the scorer. The rows
            encode into a frozen `.spipe` inline, so the set must stay under the codec's
            per-entry inline limit (1 MB, about a few thousand short prompts).
        scorer: `(response, row) -> float`, higher is better. For a strict pass/fail checker this
            is 0.0 or 1.0; a loose checker or a reward score are drop-in alternatives with finer
            granularity.
        alpha: The profiling strength. Required, with no default. The toolkit parameterizes
            `alpha` as the reciprocal of the paper's coefficient, so `alpha=100.0` is the paper's
            operating point (coefficient 0.01). Stating it here keeps every point of an `alpha`
            sweep the same fit, so the profile is resolved once per run.
        num_heads: Number of heads to select.
        layers: Candidate layers. None means every attention layer of the resolved layout;
            non-attention layers of a hybrid stack are never candidates.
        selection: `"pooled"` selects the first `num_heads` eligible candidates in ranking order.
            `"intersection"` is the paper's rule and requires a `"group"` on every row.
        min_lift: A candidate is eligible when its pooled lift exceeds this value. The default
            `0.0` keeps only heads that beat the baseline.
        screen_rows: Rows scored in stage 1 of the two-stage screen. Set together with
            `screen_keep`, or leave both unset for a single-stage profile over every row.
        screen_keep: Candidates rescored on every row in stage 2. At least `num_heads`;
            `2 * num_heads` or more is recommended, since the stage-1 standard error is large
            enough to push a true positive out of a stage-2 pool sized exactly to the target.
        gen_kwargs: Profiling generation parameters, normalized through
            `GenerationParams.from_gen_kwargs`. Greedy by default.
        batch_size: Number of rows generated per session call. Rows in a batch share one hook entry
            and generate in one batched pass, so larger values trade memory for wall-clock; excluded
            from the fit digest.
        seed: Fixes the stage-1 subsample.
        progress_callback: Optional callback invoked once per scored candidate with a small dict
            (`stage`, `candidate`, `completed`, `total`, `lift`). Opt-in; no-op when None.

    Reference:

        - "Tell Your Model Where to Attend: Post-hoc Attention Steering for LLMs"
          Qingru Zhang, Chandan Singh, Liyuan Liu, Xiaodong Liu, Bin Yu, Jianfeng Gao, Tuo Zhao
          [https://arxiv.org/abs/2311.02262](https://arxiv.org/abs/2311.02262)
    """

    access: ClassVar[ModelAccess] = ModelAccess.MODULE
    artifact_class: ClassVar[str] = "direction"

    rows: Sequence[Mapping]
    scorer: SampleScorer
    alpha: float
    num_heads: int = 48
    layers: Sequence[int] | None = None
    selection: Literal["pooled", "intersection"] = "pooled"
    min_lift: float = 0.0
    screen_rows: int | None = None
    screen_keep: int | None = None
    gen_kwargs: dict = field(default_factory=lambda: {"max_new_tokens": 128, "do_sample": False})
    batch_size: int = 8
    seed: int = 0
    progress_callback: Callable[[dict], None] | None = None

    def __post_init__(self) -> None:
        # the memo is a plain instance attribute rather than a dataclass field: identity
        # canonicalization walks every dataclasses.fields() entry, so a fitted memo stored as a
        # field would change a sweep's config_id after steer()
        self._model_ref: "weakref.ref | None" = None
        self._result: "HeadProfileResult | None" = None
        self._result_scale_position: str | None = None

        if not callable(self.scorer):
            raise ValueError("HeadProfile.scorer must be callable (response, row) -> float.")
        if not (isinstance(self.alpha, (int, float)) and self.alpha > 0):
            raise ValueError("HeadProfile.alpha must be a positive number.")
        if int(self.num_heads) < 1:
            raise ValueError("HeadProfile.num_heads must be at least 1.")
        self.num_heads = int(self.num_heads)
        if self.selection not in ("pooled", "intersection"):
            raise ValueError(f"HeadProfile.selection must be 'pooled' or 'intersection'; got {self.selection!r}.")
        if not math.isfinite(float(self.min_lift)):
            raise ValueError("HeadProfile.min_lift must be finite.")
        rows = list(self.rows)
        if not rows:
            raise ValueError("HeadProfile.rows must be non-empty.")
        for index, row in enumerate(rows):
            if "input" not in row or "substrings" not in row:
                raise ValueError(
                    f"HeadProfile.rows[{index}] must carry 'input' and 'substrings'; got keys {sorted(row)}."
                )
        self.rows = rows

        screen_set = self.screen_rows is not None
        keep_set = self.screen_keep is not None
        if screen_set != keep_set:
            raise ValueError("HeadProfile.screen_rows and screen_keep must both be set or both be unset.")
        if screen_set:
            if not 0 < int(self.screen_rows) < len(rows):
                raise ValueError(
                    f"HeadProfile.screen_rows must be in (0, {len(rows)}); got {self.screen_rows}."
                )
            if int(self.screen_keep) < self.num_heads:
                raise ValueError(
                    f"HeadProfile.screen_keep ({self.screen_keep}) must be at least num_heads ({self.num_heads})."
                )
            self.screen_rows = int(self.screen_rows)
            self.screen_keep = int(self.screen_keep)

    def budget(self, num_layers: int, num_heads: int) -> dict[str, int]:
        """Rollout counts for a `(num_layers, num_heads)` grid, without loading a model.

        A rollout is one prompt generated once. The candidate count is restricted to `layers`
        when set. With a screen, stage 1 scores every candidate on `screen_rows` rows and stage 2
        rescores `screen_keep` candidates on every row; without one, every candidate is scored on
        every row in stage 1.

        Args:
            num_layers: Number of attention layers in the model.
            num_heads: Number of attention heads per layer.

        Returns:
            A mapping with keys `candidates`, `baseline`, `stage_1`, `stage_2`, and `total`.
        """
        if self.layers is None:
            candidates = int(num_layers) * int(num_heads)
        else:
            candidates = len([layer for layer in self.layers if 0 <= int(layer) < int(num_layers)]) * int(num_heads)
        num_rows = len(self.rows)
        baseline = num_rows
        if self.screen_rows is not None:
            stage_1 = candidates * self.screen_rows
            stage_2 = min(self.screen_keep, candidates) * num_rows
        else:
            stage_1 = candidates * num_rows
            stage_2 = 0
        return {
            "candidates": candidates,
            "baseline": baseline,
            "stage_1": stage_1,
            "stage_2": stage_2,
            "total": baseline + stage_1 + stage_2,
        }

    def fit_ingredients(self) -> dict:
        """The fit-relevant recipe inputs, digested for staleness detection.

        Covers `rows`, `scorer`, `alpha`, `num_heads`, `layers`, `selection`, `min_lift`,
        `screen_rows`, `screen_keep`, `gen_kwargs`, and `seed`. `batch_size` and
        `progress_callback` are execution details and are excluded; the control's own `alpha` is
        an application parameter and is excluded here (it lives on the control, not the recipe).
        The scorer digests as its qualified name, so renaming the scorer invalidates the frozen
        profile.
        """
        return {
            "rows": list(self.rows),
            "scorer": self.scorer,
            "alpha": float(self.alpha),
            "num_heads": self.num_heads,
            "layers": None if self.layers is None else [int(layer) for layer in self.layers],
            "selection": self.selection,
            "min_lift": float(self.min_lift),
            "screen_rows": self.screen_rows,
            "screen_keep": self.screen_keep,
            "gen_kwargs": dict(self.gen_kwargs),
            "seed": int(self.seed),
        }

    def resolve(
        self,
        control: "PASTA",
        model,
        tokenizer,
        *,
        session: "SteeringSession",
    ) -> "HeadProfileResult":
        """Resolve the profile against `model`, fitting once and memoizing.

        Reruns the profiling loop through `session` and returns the ranked, filtered
        `HeadProfileResult`. The result is memoized per model (a weakref slot) and the control's
        `scale_position`, so a second control sharing this recipe on the same model reuses it and
        a different model refits.

        Args:
            control: The steering PASTA, consulted for `scale_position` and its hook builder.
            model: The live model to profile against.
            tokenizer: The tokenizer for prompt preparation and decoding.
            session: The `ScopedSession` the pipeline hands to `steer()`, scoped to
                `ModelAccess.MODULE`, through which the rollouts run.

        Returns:
            The `HeadProfileResult`.

        Raises:
            ValueError: If no candidate beats the baseline (the message names the fixes).
        """
        scale_position = control.scale_position
        if (
            model is not None
            and self._model_ref is not None
            and self._model_ref() is model
            and self._result is not None
            and self._result_scale_position == scale_position
        ):
            return self._result

        result = self._run(control, model, tokenizer, session=session)
        if model is not None:
            self._model_ref = weakref.ref(model)
            self._result = result
            self._result_scale_position = scale_position
        return result

    def _run(
        self,
        control: "PASTA",
        model,
        tokenizer,
        *,
        session: "SteeringSession",
    ) -> "HeadProfileResult":
        """Score every candidate, rank, filter, and select (no caching)."""
        device = next(model.parameters()).device
        params = GenerationParams.from_gen_kwargs(**self.gen_kwargs)
        rows = self.rows
        num_rows = len(rows)

        groups = [row.get("group") for row in rows]
        has_groups = all(group is not None for group in groups)
        if self.selection == "intersection" and not has_groups:
            raise ValueError("HeadProfile.selection='intersection' requires a 'group' on every row.")
        group_names = sorted({str(group) for group in groups}) if has_groups else None

        # prepare each batch of rows once (left-padded on the model device, spans located in the
        # padded coordinates) and reuse the tensors for the baseline and every candidate
        batches = self._prepare_batches(control, tokenizer, device)

        # baseline: unsteered scores per row, in row order
        s_base = self._score_rows(control, session, batches, params, head_map=None, rows=rows)

        candidates = self._enumerate_candidates(control)
        screen_rows_idx = self._screen_subset(groups) if self.screen_rows is not None else None
        if screen_rows_idx is not None:
            # prepare the screen subset once (spans located once), reused across every stage-1
            # candidate, and take its baseline slice in row order
            screen_batches = self._prepare_batches(control, tokenizer, device, row_indices=screen_rows_idx)
            screen_rows = [rows[i] for i in screen_rows_idx]
            screen_base = torch.tensor([s_base[i].item() for i in screen_rows_idx], dtype=torch.float32)
            screen_groups = [str(groups[i]) for i in screen_rows_idx] if has_groups else None

        num_candidate_layers = max((layer for layer, _ in candidates), default=-1) + 1
        max_heads = max((head for _, head in candidates), default=-1) + 1
        lift = torch.full((num_candidate_layers, max_heads), float("nan"))
        se = torch.full((num_candidate_layers, max_heads), float("nan"))
        n = torch.zeros((num_candidate_layers, max_heads), dtype=torch.long)
        stage = torch.zeros((num_candidate_layers, max_heads), dtype=torch.long)
        group_lift = None
        group_rows = None
        if has_groups:
            group_lift = torch.full((len(group_names), num_candidate_layers, max_heads), float("nan"))
            group_index = {name: i for i, name in enumerate(group_names)}
            group_rows = [sum(1 for group in groups if str(group) == name) for name in group_names]

        total = len(candidates)
        completed = 0

        def score_candidate(layer, head, scored_batches, scored_rows, base, scored_groups, stage_value):
            nonlocal completed
            steered = self._score_rows(
                control, session, scored_batches, params, head_map={layer: [head]}, rows=scored_rows,
            )
            diff = steered - base
            lift[layer, head] = diff.mean()
            se[layer, head] = self._std_err(diff)
            n[layer, head] = diff.numel()
            stage[layer, head] = stage_value
            if has_groups:
                for name in group_names:
                    mask = [i for i, g in enumerate(scored_groups) if g == name]
                    if mask:
                        group_lift[group_index[name], layer, head] = diff[mask].mean()
            completed += 1
            if self.progress_callback is not None:
                self.progress_callback({
                    "stage": stage_value,
                    "candidate": (layer, head),
                    "completed": completed,
                    "total": total,
                    "lift": float(diff.mean()),
                })

        all_groups = [str(group) for group in groups] if has_groups else None
        if screen_rows_idx is None:
            for layer, head in candidates:
                score_candidate(layer, head, batches, rows, s_base, all_groups, 2)
            scored = candidates
        else:
            for layer, head in candidates:
                score_candidate(layer, head, screen_batches, screen_rows, screen_base, screen_groups, 1)
            screened = sorted(
                candidates, key=lambda lh: (-_nan_low(lift[lh[0], lh[1]].item()), lh[0], lh[1]),
            )[: self.screen_keep]
            total = len(candidates) + len(screened)
            for layer, head in screened:
                score_candidate(layer, head, batches, rows, s_base, all_groups, 2)
            scored = screened

        baseline_value = float(s_base.mean())
        selected, tie_at_cutoff = self._select(scored, lift, group_lift, group_names, group_rows)
        head_config = {}
        for layer, head in selected:
            head_config.setdefault(layer, []).append(head)
        head_config = {layer: sorted(heads) for layer, heads in sorted(head_config.items())}

        return HeadProfileResult(
            head_config=head_config,
            selected=selected,
            lift=lift,
            se=se,
            n=n,
            stage=stage,
            group_lift=group_lift,
            groups=group_names,
            group_rows=group_rows,
            baseline=baseline_value,
            num_rows=num_rows,
            tie_at_cutoff=tie_at_cutoff,
            alpha=float(self.alpha),
            scale_position=control.scale_position,
        )

    def _prepare_batches(
        self, control: "PASTA", tokenizer, device, *, row_indices: Sequence[int] | None = None,
    ) -> list["_ProfileBatch"]:
        """Left-padded batches for the selected rows, with PASTA spans located once per batch.

        Reproduces the pipeline's prompt path: rows render as chat messages when the tokenizer
        has a chat template and as text otherwise, then tokenize with an empty input-control
        chain. Batching by `batch_size` matches the session's batched generate path, which
        re-stacks and left-packs the rows before the forward pass, an identity on a batch already
        left-padded to a common length. The `substrings` spans are located in these padded
        coordinates once and reused for every candidate, so a candidate only reassembles the hook
        dict rather than re-decoding and re-tokenizing.
        """
        rows = self.rows if row_indices is None else [self.rows[i] for i in row_indices]
        chat = has_chat_template(tokenizer)
        warnings_state = PromptWarnings()
        batches: list[_ProfileBatch] = []
        for start in range(0, len(rows), self.batch_size):
            chunk = rows[start:start + self.batch_size]
            inputs = [row["input"] for row in chunk]
            if chat:
                input_ids, attention_mask, message_handled, _ = resolve_messages_prompt(
                    [[{"role": "user", "content": text}] for text in inputs],
                    runtime_kwargs={},
                    input_controls=[],
                    tokenizer=tokenizer,
                )
            else:
                input_ids, attention_mask, _ = resolve_text_prompt(
                    inputs, input_controls=[], tokenizer=tokenizer, warnings_state=warnings_state,
                )
                message_handled = frozenset()
            input_ids, attention_mask = prepare_inputs(
                input_ids, attention_mask,
                input_controls=[], tokenizer=tokenizer, device=device,
                runtime_kwargs={}, message_handled=message_handled, warnings_state=warnings_state,
            )
            substrings = [list(row["substrings"]) for row in chunk]
            token_ranges, input_len = control.locate_spans(input_ids, substrings)
            batches.append(_ProfileBatch(input_ids, attention_mask, token_ranges, input_len))
        return batches

    def _score_rows(
        self,
        control: "PASTA",
        session: "SteeringSession",
        batches: Sequence["_ProfileBatch"],
        params: GenerationParams,
        *,
        head_map: dict[int, list[int]] | None,
        rows: Sequence[Mapping],
    ) -> torch.Tensor:
        """Generate every batch through the session and score each response, in row order.

        With `head_map` None the batches generate unsteered (the baseline). Otherwise one
        `HookEntry` per batch is built from the batch's located spans and shared across the
        batch's items, so the session takes its batched generate path (identical entries).
        """
        scores: list[float] = []
        row_iter = iter(rows)
        for batch in batches:
            if head_map is None:
                state_entries: tuple = ()
            else:
                hooks = control.build_hooks_for(batch.token_ranges, batch.input_len, head_map, self.alpha)
                state_entries = (HookEntry(hooks=hooks),)
            items = [
                GenerationItem(
                    prompt=PreparedPrompt.from_token_ids(
                        batch.input_ids[i:i + 1], batch.attention_mask[i:i + 1],
                    ),
                    state_entries=state_entries,
                )
                for i in range(batch.input_ids.size(0))
            ]
            results = session.generate(items, params)
            for result in results:
                row = next(row_iter)
                text = session.tokenizer.decode(result.output.output_ids[0], skip_special_tokens=True)
                scores.append(float(self.scorer(text, row)))
        return torch.tensor(scores, dtype=torch.float32)

    def _enumerate_candidates(self, control: "PASTA") -> list[tuple[int, int]]:
        """Every `(layer, head)` candidate in `(layer, head)` order."""
        layers = control.attention_layers() if self.layers is None else [int(layer) for layer in self.layers]
        candidates: list[tuple[int, int]] = []
        for layer in sorted(set(layers)):
            for head in range(control.num_heads_of_layer(layer)):
                candidates.append((layer, head))
        return candidates

    def _screen_subset(self, groups: Sequence) -> list[int]:
        """A fixed subset of `screen_rows` row indices, stratified by group when present."""
        rng = random.Random(self.seed)
        num_rows = len(self.rows)
        indices = list(range(num_rows))
        if all(group is not None for group in groups):
            by_group: dict[str, list[int]] = {}
            for index in indices:
                by_group.setdefault(str(groups[index]), []).append(index)
            names = sorted(by_group)
            per_group = max(1, self.screen_rows // len(names))
            chosen: list[int] = []
            for name in names:
                pool = by_group[name]
                chosen.extend(rng.sample(pool, min(per_group, len(pool))))
            if len(chosen) < self.screen_rows:
                remaining = [index for index in indices if index not in set(chosen)]
                chosen.extend(rng.sample(remaining, min(self.screen_rows - len(chosen), len(remaining))))
            return sorted(chosen[: self.screen_rows])
        return sorted(rng.sample(indices, self.screen_rows))

    @staticmethod
    def _std_err(diff: torch.Tensor) -> float:
        """Standard error of the paired differences (sample std over sqrt n)."""
        count = diff.numel()
        if count < 2:
            return float("nan")
        return float(diff.std(unbiased=True) / math.sqrt(count))

    def _select(
        self,
        scored: Sequence[tuple[int, int]],
        lift: torch.Tensor,
        group_lift: torch.Tensor | None,
        group_names: Sequence[str] | None,
        group_rows: Sequence[int] | None,
    ) -> tuple[list[tuple[int, int]], int]:
        """Rank the scored candidates, filter by eligibility, and select per the mode."""
        eligible = [(layer, head) for layer, head in scored if lift[layer, head].item() > self.min_lift]
        if not eligible:
            raise ValueError(
                "HeadProfile selected no heads: no candidate beat the baseline "
                f"(min_lift={self.min_lift}). Lower min_lift, add rows, or restrict layers."
            )

        ranked = sorted(eligible, key=lambda lh: (-lift[lh[0], lh[1]].item(), lh[0], lh[1]))

        if self.selection == "pooled":
            selected = ranked[: self.num_heads]
        else:
            selected = self._intersection_select(eligible, group_lift, group_names, group_rows, ranked)

        if len(selected) < self.num_heads:
            warnings.warn(
                f"HeadProfile requested {self.num_heads} heads but only {len(selected)} candidates were "
                f"eligible (lift > {self.min_lift}); selecting the eligible set.",
                UserWarning,
            )

        cutoff_lift = lift[selected[-1][0], selected[-1][1]].item()
        tie_at_cutoff = sum(1 for layer, head in ranked if lift[layer, head].item() == cutoff_lift)
        return selected, tie_at_cutoff

    def _intersection_select(
        self,
        eligible: Sequence[tuple[int, int]],
        group_lift: torch.Tensor,
        group_names: Sequence[str],
        group_rows: Sequence[int],
        pooled_ranked: Sequence[tuple[int, int]],
    ) -> list[tuple[int, int]]:
        """The paper's rule: intersection of per-group top-k sets at the smallest reaching k."""
        smallest = min(group_rows)
        if smallest < 200:
            smallest_name = group_names[group_rows.index(smallest)]
            warnings.warn(
                f"HeadProfile selection='intersection' uses group {smallest_name!r} with only {smallest} rows; "
                "the paper reports the per-group ranking as robust down to 200 samples per group.",
                UserWarning,
            )
        per_group_rank: list[list[tuple[int, int]]] = []
        for group_index in range(len(group_names)):
            ranked = sorted(
                eligible,
                key=lambda lh: (-_nan_low(group_lift[group_index, lh[0], lh[1]].item()), lh[0], lh[1]),
            )
            per_group_rank.append(ranked)

        pooled_order = {lh: i for i, lh in enumerate(pooled_ranked)}
        for k in range(1, len(eligible) + 1):
            top_sets = [set(ranked[:k]) for ranked in per_group_rank]
            intersection = set.intersection(*top_sets) if top_sets else set()
            if len(intersection) >= self.num_heads:
                ordered = sorted(intersection, key=lambda lh: pooled_order[lh])
                return ordered[: self.num_heads]
        # no k reaches num_heads: return the full final intersection, pooled-ordered
        final = set.intersection(*[set(ranked) for ranked in per_group_rank]) if per_group_rank else set()
        return sorted(final, key=lambda lh: pooled_order[lh])

access = ModelAccess.MODULE class-attribute

alpha instance-attribute

artifact_class = 'direction' class-attribute

batch_size = 8 class-attribute instance-attribute

gen_kwargs = field(default_factory=lambda: {'max_new_tokens': 128, 'do_sample': False}) class-attribute instance-attribute

layers = None class-attribute instance-attribute

min_lift = 0.0 class-attribute instance-attribute

num_heads = 48 class-attribute instance-attribute

progress_callback = None class-attribute instance-attribute

rows instance-attribute

scorer instance-attribute

screen_keep = None class-attribute instance-attribute

screen_rows = None class-attribute instance-attribute

seed = 0 class-attribute instance-attribute

selection = 'pooled' class-attribute instance-attribute

budget(num_layers, num_heads)

Rollout counts for a (num_layers, num_heads) grid, without loading a model.

A rollout is one prompt generated once. The candidate count is restricted to layers when set. With a screen, stage 1 scores every candidate on screen_rows rows and stage 2 rescores screen_keep candidates on every row; without one, every candidate is scored on every row in stage 1.

Parameters:

Name	Type	Description	Default
num_layers	int	Number of attention layers in the model.	required
num_heads	int	Number of attention heads per layer.	required

Returns:

Type	Description
dict[str, int]	A mapping with keys candidates, baseline, stage_1, stage_2, and total.

Source code in steerability/algorithms/state_control/pasta/profiling.py

def budget(self, num_layers: int, num_heads: int) -> dict[str, int]:
    """Rollout counts for a `(num_layers, num_heads)` grid, without loading a model.

    A rollout is one prompt generated once. The candidate count is restricted to `layers`
    when set. With a screen, stage 1 scores every candidate on `screen_rows` rows and stage 2
    rescores `screen_keep` candidates on every row; without one, every candidate is scored on
    every row in stage 1.

    Args:
        num_layers: Number of attention layers in the model.
        num_heads: Number of attention heads per layer.

    Returns:
        A mapping with keys `candidates`, `baseline`, `stage_1`, `stage_2`, and `total`.
    """
    if self.layers is None:
        candidates = int(num_layers) * int(num_heads)
    else:
        candidates = len([layer for layer in self.layers if 0 <= int(layer) < int(num_layers)]) * int(num_heads)
    num_rows = len(self.rows)
    baseline = num_rows
    if self.screen_rows is not None:
        stage_1 = candidates * self.screen_rows
        stage_2 = min(self.screen_keep, candidates) * num_rows
    else:
        stage_1 = candidates * num_rows
        stage_2 = 0
    return {
        "candidates": candidates,
        "baseline": baseline,
        "stage_1": stage_1,
        "stage_2": stage_2,
        "total": baseline + stage_1 + stage_2,
    }

fit_ingredients()

The fit-relevant recipe inputs, digested for staleness detection.

Covers rows, scorer, alpha, num_heads, layers, selection, min_lift, screen_rows, screen_keep, gen_kwargs, and seed. batch_size and progress_callback are execution details and are excluded; the control's own alpha is an application parameter and is excluded here (it lives on the control, not the recipe). The scorer digests as its qualified name, so renaming the scorer invalidates the frozen profile.

Source code in steerability/algorithms/state_control/pasta/profiling.py

def fit_ingredients(self) -> dict:
    """The fit-relevant recipe inputs, digested for staleness detection.

    Covers `rows`, `scorer`, `alpha`, `num_heads`, `layers`, `selection`, `min_lift`,
    `screen_rows`, `screen_keep`, `gen_kwargs`, and `seed`. `batch_size` and
    `progress_callback` are execution details and are excluded; the control's own `alpha` is
    an application parameter and is excluded here (it lives on the control, not the recipe).
    The scorer digests as its qualified name, so renaming the scorer invalidates the frozen
    profile.
    """
    return {
        "rows": list(self.rows),
        "scorer": self.scorer,
        "alpha": float(self.alpha),
        "num_heads": self.num_heads,
        "layers": None if self.layers is None else [int(layer) for layer in self.layers],
        "selection": self.selection,
        "min_lift": float(self.min_lift),
        "screen_rows": self.screen_rows,
        "screen_keep": self.screen_keep,
        "gen_kwargs": dict(self.gen_kwargs),
        "seed": int(self.seed),
    }

resolve(control, model, tokenizer, *, session)

Resolve the profile against model, fitting once and memoizing.

Reruns the profiling loop through session and returns the ranked, filtered HeadProfileResult. The result is memoized per model (a weakref slot) and the control's scale_position, so a second control sharing this recipe on the same model reuses it and a different model refits.

Parameters:

Name	Type	Description	Default
control	'PASTA'	The steering PASTA, consulted for scale_position and its hook builder.	required
model		The live model to profile against.	required
tokenizer		The tokenizer for prompt preparation and decoding.	required
session	'SteeringSession'	The ScopedSession the pipeline hands to steer(), scoped to ModelAccess.MODULE, through which the rollouts run.	required

Returns:

Type	Description
'HeadProfileResult'	The HeadProfileResult.

Raises:

Type	Description
ValueError	If no candidate beats the baseline (the message names the fixes).

Source code in steerability/algorithms/state_control/pasta/profiling.py

def resolve(
    self,
    control: "PASTA",
    model,
    tokenizer,
    *,
    session: "SteeringSession",
) -> "HeadProfileResult":
    """Resolve the profile against `model`, fitting once and memoizing.

    Reruns the profiling loop through `session` and returns the ranked, filtered
    `HeadProfileResult`. The result is memoized per model (a weakref slot) and the control's
    `scale_position`, so a second control sharing this recipe on the same model reuses it and
    a different model refits.

    Args:
        control: The steering PASTA, consulted for `scale_position` and its hook builder.
        model: The live model to profile against.
        tokenizer: The tokenizer for prompt preparation and decoding.
        session: The `ScopedSession` the pipeline hands to `steer()`, scoped to
            `ModelAccess.MODULE`, through which the rollouts run.

    Returns:
        The `HeadProfileResult`.

    Raises:
        ValueError: If no candidate beats the baseline (the message names the fixes).
    """
    scale_position = control.scale_position
    if (
        model is not None
        and self._model_ref is not None
        and self._model_ref() is model
        and self._result is not None
        and self._result_scale_position == scale_position
    ):
        return self._result

    result = self._run(control, model, tokenizer, session=session)
    if model is not None:
        self._model_ref = weakref.ref(model)
        self._result = result
        self._result_scale_position = scale_position
    return result

HeadProfileResult dataclass

The resolution of a HeadProfile: the selected head map plus per-candidate statistics.

lift, se, n, and stage are [num_layers, max_heads] tensors with NaN (0 for n and stage) at positions that were not candidates. stage is 1 for a candidate scored only on the screen subset and 2 for a candidate scored on every row. group_lift is [num_groups, num_layers, max_heads] and group_rows the per-group row counts, both None when the rows carry no "group".

Attributes:

Name	Type	Description
head_config	dict[int, list[int]]	The selected head map, layer index to sorted head indices.
selected	list[tuple[int, int]]	The selected candidates, in selection order.
lift	Tensor	Per-candidate mean gain over the baseline.
se	Tensor	Per-candidate standard error of the paired differences.
n	Tensor	Per-candidate number of rows scored.
stage	Tensor	Per-candidate screen stage (1 or 2), 0 for non-candidates.
group_lift	Tensor | None	Per-group lift, or None.
groups	list[str] | None	The group names, or None.
group_rows	list[int] | None	The per-group row counts, or None.
baseline	float	The unsteered mean score over all rows.
num_rows	int	The number of profiling rows.
tie_at_cutoff	int	Candidates whose lift equals the last selected candidate's.
alpha	float	The profiling strength.
scale_position	str	The control's scale position the profile was scored under.

Source code in steerability/algorithms/state_control/pasta/profiling.py

@dataclass
class HeadProfileResult:
    """The resolution of a `HeadProfile`: the selected head map plus per-candidate statistics.

    `lift`, `se`, `n`, and `stage` are `[num_layers, max_heads]` tensors with `NaN` (0 for `n`
    and `stage`) at positions that were not candidates. `stage` is 1 for a candidate scored only
    on the screen subset and 2 for a candidate scored on every row. `group_lift` is
    `[num_groups, num_layers, max_heads]` and `group_rows` the per-group row counts, both None
    when the rows carry no `"group"`.

    Attributes:
        head_config: The selected head map, layer index to sorted head indices.
        selected: The selected candidates, in selection order.
        lift: Per-candidate mean gain over the baseline.
        se: Per-candidate standard error of the paired differences.
        n: Per-candidate number of rows scored.
        stage: Per-candidate screen stage (1 or 2), 0 for non-candidates.
        group_lift: Per-group lift, or None.
        groups: The group names, or None.
        group_rows: The per-group row counts, or None.
        baseline: The unsteered mean score over all rows.
        num_rows: The number of profiling rows.
        tie_at_cutoff: Candidates whose lift equals the last selected candidate's.
        alpha: The profiling strength.
        scale_position: The control's scale position the profile was scored under.
    """

    head_config: dict[int, list[int]]
    selected: list[tuple[int, int]]
    lift: torch.Tensor
    se: torch.Tensor
    n: torch.Tensor
    stage: torch.Tensor
    group_lift: torch.Tensor | None
    groups: list[str] | None
    group_rows: list[int] | None
    baseline: float
    num_rows: int
    tie_at_cutoff: int
    alpha: float
    scale_position: str

    def to_frame(self) -> "pandas.DataFrame":  # noqa: F821
        """One row per candidate, columns `layer`, `head`, `lift`, `se`, `n`, `stage`,
        `selected`, `rank`, plus one `lift_<group>` column per group when groups are present."""
        import pandas as pd

        selected_set = {tuple(pair) for pair in self.selected}
        rank_of = {tuple(pair): i for i, pair in enumerate(self.selected)}
        records: list[dict] = []
        num_layers, max_heads = self.lift.shape
        for layer in range(num_layers):
            for head in range(max_heads):
                if int(self.stage[layer, head]) == 0:
                    continue
                record = {
                    "layer": layer,
                    "head": head,
                    "lift": float(self.lift[layer, head]),
                    "se": float(self.se[layer, head]),
                    "n": int(self.n[layer, head]),
                    "stage": int(self.stage[layer, head]),
                    "selected": (layer, head) in selected_set,
                    "rank": rank_of.get((layer, head)),
                }
                if self.groups is not None:
                    for group_index, name in enumerate(self.groups):
                        record[f"lift_{name}"] = float(self.group_lift[group_index, layer, head])
                records.append(record)
        return pd.DataFrame(records)

    def save(self, path: str | Path) -> None:
        """Write the complete result to one JSON file (tensors as nested lists, NaN as null)."""
        payload = {
            "head_config": {str(layer): list(heads) for layer, heads in self.head_config.items()},
            "selected": [list(pair) for pair in self.selected],
            "lift": _tensor_to_json(self.lift),
            "se": _tensor_to_json(self.se),
            "n": _tensor_to_json(self.n),
            "stage": _tensor_to_json(self.stage),
            "group_lift": None if self.group_lift is None else _tensor_to_json(self.group_lift),
            "groups": self.groups,
            "group_rows": self.group_rows,
            "baseline": self.baseline,
            "num_rows": self.num_rows,
            "tie_at_cutoff": self.tie_at_cutoff,
            "alpha": self.alpha,
            "scale_position": self.scale_position,
        }
        Path(path).write_text(json.dumps(payload, indent=2))

    @classmethod
    def load(cls, path: str | Path) -> "HeadProfileResult":
        """Read a result written by `save`, reproducing every field exactly."""
        payload = json.loads(Path(path).read_text())
        return cls(
            head_config={int(layer): list(heads) for layer, heads in payload["head_config"].items()},
            selected=[tuple(pair) for pair in payload["selected"]],
            lift=_json_to_tensor(payload["lift"], torch.float32),
            se=_json_to_tensor(payload["se"], torch.float32),
            n=_json_to_tensor(payload["n"], torch.long),
            stage=_json_to_tensor(payload["stage"], torch.long),
            group_lift=None if payload["group_lift"] is None else _json_to_tensor(payload["group_lift"], torch.float32),
            groups=payload["groups"],
            group_rows=payload["group_rows"],
            baseline=payload["baseline"],
            num_rows=payload["num_rows"],
            tie_at_cutoff=payload["tie_at_cutoff"],
            alpha=payload["alpha"],
            scale_position=payload["scale_position"],
        )

alpha instance-attribute

baseline instance-attribute

group_lift instance-attribute

group_rows instance-attribute

groups instance-attribute

head_config instance-attribute

lift instance-attribute

n instance-attribute

num_rows instance-attribute

scale_position instance-attribute

se instance-attribute

selected instance-attribute

stage instance-attribute

tie_at_cutoff instance-attribute

load(path) classmethod

Read a result written by save, reproducing every field exactly.

Source code in steerability/algorithms/state_control/pasta/profiling.py

@classmethod
def load(cls, path: str | Path) -> "HeadProfileResult":
    """Read a result written by `save`, reproducing every field exactly."""
    payload = json.loads(Path(path).read_text())
    return cls(
        head_config={int(layer): list(heads) for layer, heads in payload["head_config"].items()},
        selected=[tuple(pair) for pair in payload["selected"]],
        lift=_json_to_tensor(payload["lift"], torch.float32),
        se=_json_to_tensor(payload["se"], torch.float32),
        n=_json_to_tensor(payload["n"], torch.long),
        stage=_json_to_tensor(payload["stage"], torch.long),
        group_lift=None if payload["group_lift"] is None else _json_to_tensor(payload["group_lift"], torch.float32),
        groups=payload["groups"],
        group_rows=payload["group_rows"],
        baseline=payload["baseline"],
        num_rows=payload["num_rows"],
        tie_at_cutoff=payload["tie_at_cutoff"],
        alpha=payload["alpha"],
        scale_position=payload["scale_position"],
    )

save(path)

Write the complete result to one JSON file (tensors as nested lists, NaN as null).

Source code in steerability/algorithms/state_control/pasta/profiling.py

def save(self, path: str | Path) -> None:
    """Write the complete result to one JSON file (tensors as nested lists, NaN as null)."""
    payload = {
        "head_config": {str(layer): list(heads) for layer, heads in self.head_config.items()},
        "selected": [list(pair) for pair in self.selected],
        "lift": _tensor_to_json(self.lift),
        "se": _tensor_to_json(self.se),
        "n": _tensor_to_json(self.n),
        "stage": _tensor_to_json(self.stage),
        "group_lift": None if self.group_lift is None else _tensor_to_json(self.group_lift),
        "groups": self.groups,
        "group_rows": self.group_rows,
        "baseline": self.baseline,
        "num_rows": self.num_rows,
        "tie_at_cutoff": self.tie_at_cutoff,
        "alpha": self.alpha,
        "scale_position": self.scale_position,
    }
    Path(path).write_text(json.dumps(payload, indent=2))

to_frame()

One row per candidate, columns layer, head, lift, se, n, stage, selected, rank, plus one lift_<group> column per group when groups are present.

Source code in steerability/algorithms/state_control/pasta/profiling.py

def to_frame(self) -> "pandas.DataFrame":  # noqa: F821
    """One row per candidate, columns `layer`, `head`, `lift`, `se`, `n`, `stage`,
    `selected`, `rank`, plus one `lift_<group>` column per group when groups are present."""
    import pandas as pd

    selected_set = {tuple(pair) for pair in self.selected}
    rank_of = {tuple(pair): i for i, pair in enumerate(self.selected)}
    records: list[dict] = []
    num_layers, max_heads = self.lift.shape
    for layer in range(num_layers):
        for head in range(max_heads):
            if int(self.stage[layer, head]) == 0:
                continue
            record = {
                "layer": layer,
                "head": head,
                "lift": float(self.lift[layer, head]),
                "se": float(self.se[layer, head]),
                "n": int(self.n[layer, head]),
                "stage": int(self.stage[layer, head]),
                "selected": (layer, head) in selected_set,
                "rank": rank_of.get((layer, head)),
            }
            if self.groups is not None:
                for group_index, name in enumerate(self.groups):
                    record[f"lift_{name}"] = float(self.group_lift[group_index, layer, head])
            records.append(record)
    return pd.DataFrame(records)