ITI

steerability.algorithms.state_control.iti

Inference-Time Intervention (ITI) state control.

args

ITI argument validation.

control

Inference-Time Intervention (ITI) state control.

ITI

Bases: InterventionControl

Inference-Time Intervention (ITI).

Steers model behavior by shifting activations at a sparse set of attention heads during inference. The intervention operates at the residual stream level by adding direction vectors to head-associated slices of the hidden dimension.

ITI operates in two phases:

1. Offline (during steer()): For every attention head across all layers, extract the head's output activations on labeled true/false statements. Train a per-head linear probe on an 80/20 held-out split (over groups when the data carries them) and rank heads by validation accuracy. For the top-K heads, compute the mass mean shift: direction = mean(activations_true)

2. mean(activations_false).

3. Online (during generation): At each generated token, for each selected (layer, head) pair, add alpha * direction to that head's slice of the residual stream. The intervention fires unconditionally on every token in the specified token_scope.

The control is declarative: _configure maps the validated args onto one Intervention at the attention output projection (the site derived from the head_additive transform kind), over the layers hosting selected heads.

Reference:

• "Inference-Time Intervention: Eliciting Truthful Answers from a Language Model" Kenneth Li, Oam Patel, Fernanda Viégas, Hanspeter Pfister, Martin Wattenberg https://arxiv.org/abs/2306.03341

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

class ITI(InterventionControl):
    """Inference-Time Intervention (ITI).

    Steers model behavior by shifting activations at a sparse set of attention heads
    during inference. The intervention operates at the residual stream level by adding
    direction vectors to head-associated slices of the hidden dimension.

    ITI operates in two phases:

    1. **Offline (during steer())**: For every attention head across all layers,
       extract the head's output activations on labeled true/false statements.
       Train a per-head linear probe on an 80/20 held-out split (over groups when
       the data carries them) and rank heads by validation accuracy. For the
       top-K heads, compute the mass mean shift: direction = mean(activations_true)
       - mean(activations_false).

    2. **Online (during generation)**: At each generated token, for each selected
       (layer, head) pair, add alpha * direction to that head's slice of the
       residual stream. The intervention fires unconditionally on every token
       in the specified token_scope.

    The control is declarative: `_configure` maps the validated args onto one `Intervention`
    at the attention output projection (the site derived from the `head_additive` transform
    kind), over the layers hosting selected heads.

    Reference:

    - "Inference-Time Intervention: Eliciting Truthful Answers from a Language Model"
    Kenneth Li, Oam Patel, Fernanda Viégas, Hanspeter Pfister, Martin Wattenberg
    [https://arxiv.org/abs/2306.03341](https://arxiv.org/abs/2306.03341)
    """

    Args = ITIArgs
    supports_batching = True
    hook_only_hint = (
        "norm preservation over per-head streams has no intervention-spec form; "
        "run on the huggingface backend"
    )

    def _configure(self):
        if self.steering_vector is not None:
            if isinstance(self.steering_vector, SteeringVector):
                source = _Precomputed(self.steering_vector.clone())
            else:
                source = self.steering_vector
        else:
            source = _ProbeMassShiftFit(self.data, self.train_spec)

        self._template = (Intervention(
            layers=CoveredLayers(),
            transform=_HeadSelectionBuild(
                source,
                selected_heads=self.selected_heads,
                num_heads=self.num_heads,
                alpha=self.alpha,
                norm_preserving=self.use_norm_preservation,
            ),
            scope=TokenScope(self.token_scope, last_k=self.last_k, from_position=self.from_position),
            boundary="layer_input",
        ),)

    def wire_kinds(self):
        """`head_additive` kinds for the bound configuration; conservative before `steer()`.

        The factory-built transform is unknown before binding, but its kind is definitional
        for this control, so the plan is stated directly: `head_additive` unless norm
        preservation is on (the wire modifier rescales the residual row, not the per-head
        stream).
        """
        from steerability.algorithms.core.execution.contracts import InterventionKinds
        from steerability.algorithms.state_control.common.specs import combine_kinds

        if self.interventions:
            return combine_kinds(intervention.wire_kinds() for intervention in self.interventions)
        if self.use_norm_preservation:
            return None
        return InterventionKinds(
            transforms=frozenset({"head_additive"}),
            scopes=frozenset({self.token_scope}),
        )

    @property
    def _active_layer_ids(self) -> set[int]:
        """Layers hosting selected heads (empty before `steer()`)."""
        return set(self.interventions[0].layers) if self.interventions else set()

    @property
    def _steering_vector(self) -> SteeringVector | None:
        """The bound head-shift artifact (None before `steer()`)."""
        if not self.interventions:
            return None
        core, _ = unwrap_modifiers(self.interventions[0].transform)
        return getattr(core, "steering_vector", None)

    def export_state(self) -> dict:
        """The bound per-head steering vector under the `"steering_vector"` key (after `steer()`)."""
        vector = self._steering_vector
        return {"steering_vector": vector} if vector is not None else {}

    def frozen_form(self, state: dict) -> tuple[str, dict]:
        """A same-class frozen form: the bound per-head vector plus the resolved head selection."""
        core, _ = unwrap_modifiers(self.interventions[0].transform)
        selected = sorted(
            (int(layer_id), int(head_id))
            for layer_id, heads in core.active_heads.items()
            for head_id in heads
        )
        return "state_control/iti", {
            "steering_vector": state["steering_vector"],
            "selected_heads": selected,
            "num_heads": self.num_heads,
            "alpha": self.alpha,
            "token_scope": self.token_scope,
            "last_k": self.last_k,
            "from_position": self.from_position,
            "use_norm_preservation": self.use_norm_preservation,
        }

RUNTIME_KWARGS_SCHEMA = [] class-attribute instance-attribute

_active_layer_ids property

Layers hosting selected heads (empty before steer()).

_gate property

The first intervention's gate (None before steer() or when unconditional).

_steering_vector property

The bound head-shift artifact (None before steer()).

_template = () class-attribute instance-attribute

_transform property writable

The first intervention's transform (None before steer()).

Assignment replaces the transform on the first intervention, so a wrapped or instrumented transform takes effect in subsequently built hooks.

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

enabled = True class-attribute instance-attribute

hook_only_hint = 'norm preservation over per-head streams has no intervention-spec form; run on the huggingface backend' class-attribute instance-attribute

interventions = () class-attribute instance-attribute

supports_batching = True class-attribute instance-attribute

tokenizer = None class-attribute instance-attribute

__init__(*args, **kwargs)

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

def __init__(self, *args, **kwargs) -> None:
    if self.Args is None:  # null / arg-free control
        if args or kwargs:
            raise TypeError(f"{type(self).__name__} accepts no constructor arguments.")
        self._configure()
        return

    self.args: BaseArgs = self.Args.validate(*args, **kwargs)

    # move fields to attributes, skipping any name the control exposes as a property
    # (e.g. CAST.condition_point); the raw value stays reachable via self.args.<name>
    for field in fields(self.args):
        if isinstance(getattr(type(self), field.name, None), property):
            continue
        setattr(self, field.name, getattr(self.args, field.name))

    self._configure()

_configure()

Post-construction hook, called after Args fields are mirrored onto the instance.

Driver presets override this to map their mirrored args onto the fields their generic base reads, so subclasses never bypass a parent __init__. Default no-op.

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

def _configure(self):
    if self.steering_vector is not None:
        if isinstance(self.steering_vector, SteeringVector):
            source = _Precomputed(self.steering_vector.clone())
        else:
            source = self.steering_vector
    else:
        source = _ProbeMassShiftFit(self.data, self.train_spec)

    self._template = (Intervention(
        layers=CoveredLayers(),
        transform=_HeadSelectionBuild(
            source,
            selected_heads=self.selected_heads,
            num_heads=self.num_heads,
            alpha=self.alpha,
            norm_preserving=self.use_norm_preservation,
        ),
        scope=TokenScope(self.token_scope, last_k=self.last_k, from_position=self.from_position),
        boundary="layer_input",
    ),)

_fit_sources()

Yield the template's fit sources: unbound sources carrying an artifact_class whose resolution does model-side work (access above ModelAccess.FACTS).

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

def _fit_sources(self):
    """Yield the template's fit sources: unbound sources carrying an `artifact_class`
    whose resolution does model-side work (`access` above `ModelAccess.FACTS`)."""
    for source in self._unbound_sources():
        if getattr(source, "artifact_class", None) is None:
            continue
        if getattr(source, "access", ModelAccess.MODULE) == ModelAccess.FACTS:
            continue
        yield source

_resolve_module_layout(model=None)

The module-path layout, resolved from the module tree on first use.

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

def _resolve_module_layout(self, model=None):
    """The module-path layout, resolved from the module tree on first use."""
    layout = getattr(self, "_module_layout", None)
    if layout is None:
        from steerability.algorithms.core.internals.model_layout import resolve_model_layout

        if model is None:
            raise RuntimeError(
                f"{type(self).__name__} was steered without a live model, so hook module "
                "names are unresolved; provide the model (the pipeline does) or steer with "
                "a model."
            )
        layout = resolve_model_layout(model)
        self._module_layout = layout
    return layout

_unbound_sources()

Yield each unbound template element's source (or undeclared factory slot).

Yields the transform sources of unbound transform elements, factory transform slots themselves (which declare their own access or default to the live model), and unresolved gate sources, in template order.

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

def _unbound_sources(self):
    """Yield each unbound template element's source (or undeclared factory slot).

    Yields the transform sources of unbound transform elements, factory transform slots
    themselves (which declare their own `access` or default to the live model), and
    unresolved gate sources, in template order.
    """
    from steerability.algorithms.state_control.common.transforms.base import BaseTransform, unwrap_modifiers

    for intervention in self._template:
        transform = intervention.transform
        if isinstance(transform, BaseTransform):
            core, wrappers = unwrap_modifiers(transform)
            for element in (core, *wrappers):
                if not element.is_bound and element.source is not None:
                    yield element.source
        else:
            yield transform
        gate = intervention.gate
        if gate is not None and not _is_concrete_gate(gate):
            yield gate

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 per-call clone whose interventions carry independent gate state.

Gates are deep-copied with one shared memo across the control's interventions, so a gate instance shared by several interventions stays shared inside the clone while being isolated from the original and from sibling clones. Transforms, scorers, and steer-time artifacts stay shared.

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

def clone_for_call(self, seed: int | None = None):
    """A per-call clone whose interventions carry independent gate state.

    Gates are deep-copied with one shared memo across the control's interventions, so a
    gate instance shared by several interventions stays shared inside the clone while
    being isolated from the original and from sibling clones. Transforms, scorers, and
    steer-time artifacts stay shared.
    """
    import dataclasses

    clone = super().clone_for_call(seed)
    if self.interventions:
        memo: dict = {}
        clone.interventions = tuple(
            dataclasses.replace(intervention, gate=copy.deepcopy(intervention.gate, memo))
            for intervention in self.interventions
        )
    return clone

export_intervention_spec(runtime_kwargs=None)

The control's InterventionSpec, lowered from the bound interventions, or None.

Must be called after steer(). Returns None when the configuration has no wire form.

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

def export_intervention_spec(self, runtime_kwargs: dict | None = None):
    """The control's `InterventionSpec`, lowered from the bound interventions, or None.

    Must be called after `steer()`. Returns None when the configuration has no wire form.
    """
    from steerability.algorithms.state_control.common.lowering import lower_interventions

    if not self.interventions or getattr(self, "_num_layers", None) is None:
        return None
    kinds = self.wire_kinds()
    if kinds is None:
        return None
    return lower_interventions(self.interventions, num_layers=self._num_layers)

export_state()

The bound per-head steering vector under the "steering_vector" key (after steer()).

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

def export_state(self) -> dict:
    """The bound per-head steering vector under the `"steering_vector"` key (after `steer()`)."""
    vector = self._steering_vector
    return {"steering_vector": vector} if vector is not None else {}

fit_identity()

The template's fit sources in template order, each in its encodable form, or None when the template declares no fits.

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

def fit_identity(self) -> Any | None:
    """The template's fit sources in template order, each in its encodable form, or None
    when the template declares no fits."""
    sources = list(self._fit_sources())
    if not sources:
        return None
    return tuple(_fit_ingredients(source) for source in sources)

frozen_form(state)

A same-class frozen form: the bound per-head vector plus the resolved head selection.

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

def frozen_form(self, state: dict) -> tuple[str, dict]:
    """A same-class frozen form: the bound per-head vector plus the resolved head selection."""
    core, _ = unwrap_modifiers(self.interventions[0].transform)
    selected = sorted(
        (int(layer_id), int(head_id))
        for layer_id, heads in core.active_heads.items()
        for head_id in heads
    )
    return "state_control/iti", {
        "steering_vector": state["steering_vector"],
        "selected_heads": selected,
        "num_heads": self.num_heads,
        "alpha": self.alpha,
        "token_scope": self.token_scope,
        "last_k": self.last_k,
        "from_position": self.from_position,
        "use_norm_preservation": self.use_norm_preservation,
    }

get_hooks(input_ids, runtime_kwargs=None, attention_mask=None, **kwargs)

Compile the bound interventions to hooks for the current generation.

Delegates to build_hooks: a fresh hook runtime is created, gates reset to the logical batch, and one behavior hook is emitted per (intervention, layer).

Generation-time context arrives through **kwargs; model is consulted to resolve hook module names when steering ran without a live model.

Parameters:

Name	Type	Description	Default
input_ids	Tensor | dict	Prompt token ids of shape [B, T] or [T].	required
runtime_kwargs	dict | None	Unused.	None
attention_mask	Tensor | None	The prompt attention mask matching input_ids, forwarded to gate evidence pooling on the prefill pass. When None and the tokenizer defines a pad token, a mask is inferred from leading and trailing pad runs.	None

Returns:

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

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

def get_hooks(
    self,
    input_ids: torch.Tensor | dict,
    runtime_kwargs: dict | None = None,
    attention_mask: torch.Tensor | None = None,
    **kwargs,
) -> dict[str, list[HookSpec]]:
    """Compile the bound interventions to hooks for the current generation.

    Delegates to `build_hooks`: a fresh hook runtime is created, gates reset to the
    logical batch, and one behavior hook is emitted per (intervention, layer).

    Generation-time context arrives through `**kwargs`; `model` is consulted to resolve
    hook module names when steering ran without a live model.

    Args:
        input_ids: Prompt token ids of shape `[B, T]` or `[T]`.
        runtime_kwargs: Unused.
        attention_mask: The prompt attention mask matching `input_ids`, forwarded to
            gate evidence pooling on the prefill pass. When None and the tokenizer defines a
            pad token, a mask is inferred from leading and trailing pad runs.

    Returns:
        Hook specifications with `"pre"`, `"forward"`, `"backward"` keys.
    """
    from steerability.algorithms.state_control.common.runtime import build_hooks
    from steerability.algorithms.state_control.common.token_scope import compute_prompt_lens
    from steerability.utils.tokenization import infer_attention_mask_from_ids

    ids = input_ids if isinstance(input_ids, torch.Tensor) else input_ids["input_ids"]
    if ids.ndim == 1:
        ids = ids.unsqueeze(0)

    layout = self._resolve_module_layout(kwargs.get("model"))
    pad_token_id = getattr(self.tokenizer, "pad_token_id", None) if self.tokenizer is not None else None
    prompt_lens = compute_prompt_lens(ids, pad_token_id)

    if attention_mask is not None:
        mask = attention_mask if isinstance(attention_mask, torch.Tensor) else torch.as_tensor(attention_mask)
        prompt_mask = (mask.unsqueeze(0) if mask.ndim == 1 else mask).to(torch.bool)
    elif pad_token_id is not None:
        prompt_mask = infer_attention_mask_from_ids(ids, pad_token_id).to(torch.bool)
    else:
        prompt_mask = None

    return build_hooks(self.interventions, layout, prompt_lens, prompt_mask, model=kwargs.get("model"))

requirements()

Backend requirements derived from the declared interventions, per phase.

Generate offers the intervention-spec alternative whenever every component of every intervention has a wire form; hook-only configurations require the in-process backend. Score is in-process: remote prompt-logprob scoring anchors token scopes at the request's prompt end (the end of the prompt-plus-reference concatenation), which would silently unanchor prompt-relative interventions.

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

def requirements(self) -> Requirements:
    """Backend requirements derived from the declared interventions, per phase.

    Generate offers the intervention-spec alternative whenever every component of every
    intervention has a wire form; hook-only configurations require the in-process backend.
    Score is in-process: remote prompt-logprob scoring anchors token scopes at the
    request's prompt end (the end of the prompt-plus-reference concatenation), which would
    silently unanchor prompt-relative interventions.
    """
    from steerability.algorithms.core.execution.contracts import Capability, Requirements, any_of, needs

    kinds = self.wire_kinds()
    in_process = needs(Capability.IN_PROCESS_TORCH)
    score = needs(
        Capability.IN_PROCESS_TORCH,
        hint=(
            "remote prompt-logprob scoring anchors token scopes at the request's prompt "
            "end, so scoped interventions would not cover the reference; score on the "
            "huggingface backend"
        ),
    )
    if kinds is None:
        return Requirements(
            generate=needs(Capability.IN_PROCESS_TORCH, hint=self.hook_only_hint),
            score=score,
        )
    return Requirements(
        generate=any_of(
            in_process,
            needs(
                Capability.INTERVENTION_SPECS,
                kinds=kinds,
                hint="serve this intervention through the vLLM-Hook plugin",
            ),
        ),
        score=score,
    )

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

Bind the intervention template against the model (or the session's layout).

Structural facts come from the steering session's layout when a session is given, so a fully concrete template (precomputed vectors, manual thresholds) binds with model=None. Templates carrying sources (fits, searches) resolve them here.

Parameters:

Name	Type	Description	Default
model	PreTrainedModel | None	The base language model, or None for concrete templates bound against a session layout.	None
tokenizer	PreTrainedTokenizerBase | None	Tokenizer used when fitting sources.	None
session	SteeringSession | None	SteeringSession on the steering backend, provided by the pipeline.	None

Returns:

Type	Description
PreTrainedModel | None	The input model, unchanged.

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

def steer(
    self,
    model: PreTrainedModel | None = None,
    tokenizer: PreTrainedTokenizerBase | None = None,
    session: "SteeringSession | None" = None,
    **kwargs,
) -> PreTrainedModel | None:
    """Bind the intervention template against the model (or the session's layout).

    Structural facts come from the steering session's layout when a session is given, so a
    fully concrete template (precomputed vectors, manual thresholds) binds with
    `model=None`. Templates carrying sources (fits, searches) resolve them here.

    Args:
        model: The base language model, or None for concrete templates bound against a
            session layout.
        tokenizer: Tokenizer used when fitting sources.
        session: `SteeringSession` on the steering backend, provided by the pipeline.

    Returns:
        The input model, unchanged.
    """
    from steerability.algorithms.core.internals.model_layout import resolve_model_layout
    from steerability.algorithms.state_control.common.layout_facts import resolve_layout

    layout = resolve_layout(model, session)
    self._num_layers = layout.num_layers
    self._module_layout = resolve_model_layout(model) if model is not None else None
    if tokenizer is not None:
        self.tokenizer = tokenizer
    self.interventions = tuple(
        intervention.bind(model, tokenizer, layout=layout, session=session)
        for intervention in self._template
    )
    return model

steer_access()

The model access the template's steer step requires, folded over its sources.

A fully bound template requires ModelAccess.FACTS, since pure layer selectors resolve from structural facts available on any session. Otherwise the strongest declared source access wins; a source or factory slot without an access declaration builds against the live model, so it counts as ModelAccess.MODULE.

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

def steer_access(self) -> ModelAccess:
    """The model access the template's steer step requires, folded over its sources.

    A fully bound template requires `ModelAccess.FACTS`, since pure layer selectors
    resolve from structural facts available on any session. Otherwise the strongest
    declared source access wins; a source or factory slot without an `access` declaration
    builds against the live model, so it counts as `ModelAccess.MODULE`.
    """
    access = ModelAccess.FACTS
    for source in self._unbound_sources():
        access = max(access, getattr(source, "access", ModelAccess.MODULE))
    return access

steer_fits()

The template's fit artifacts, i.e. every unbound source carrying an artifact_class, as (artifact, artifact_class) pairs in template order.

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

def steer_fits(self) -> tuple[tuple[str, str], ...]:
    """The template's fit artifacts, i.e. every unbound source carrying an
    `artifact_class`, as `(artifact, artifact_class)` pairs in template order."""
    return tuple(
        (type(source).__name__, source.artifact_class) for source in self._fit_sources()
    )

wire_kinds()

head_additive kinds for the bound configuration; conservative before steer().

The factory-built transform is unknown before binding, but its kind is definitional for this control, so the plan is stated directly: head_additive unless norm preservation is on (the wire modifier rescales the residual row, not the per-head stream).

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

def wire_kinds(self):
    """`head_additive` kinds for the bound configuration; conservative before `steer()`.

    The factory-built transform is unknown before binding, but its kind is definitional
    for this control, so the plan is stated directly: `head_additive` unless norm
    preservation is on (the wire modifier rescales the residual row, not the per-head
    stream).
    """
    from steerability.algorithms.core.execution.contracts import InterventionKinds
    from steerability.algorithms.state_control.common.specs import combine_kinds

    if self.interventions:
        return combine_kinds(intervention.wire_kinds() for intervention in self.interventions)
    if self.use_norm_preservation:
        return None
    return InterventionKinds(
        transforms=frozenset({"head_additive"}),
        scopes=frozenset({self.token_scope}),
    )

_HeadSelectionBuild

Transform factory: resolve the head-shift artifact, select heads, build the transform.

Head selection is a fact of the artifact (top-K heads by probe accuracy), so the transform is constructed at bind time from the resolved SteeringVector. The factory declares its own steer access: a precomputed vector builds model-free, while fitting captures pre-o_proj per-head activations, which no backend serves remotely.

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

class _HeadSelectionBuild:
    """Transform factory: resolve the head-shift artifact, select heads, build the transform.

    Head selection is a fact of the artifact (top-K heads by probe accuracy), so the transform
    is constructed at bind time from the resolved `SteeringVector`. The factory declares its
    own steer access: a precomputed vector builds model-free, while fitting captures
    pre-`o_proj` per-head activations, which no backend serves remotely.
    """

    def __init__(self, source, selected_heads, num_heads: int, alpha: float, norm_preserving: bool):
        self._source = source
        self._selected_heads = selected_heads
        self._num_heads = num_heads
        self._alpha = alpha
        self._norm_preserving = norm_preserving

    @property
    def access(self) -> ModelAccess:
        return getattr(self._source, "access", ModelAccess.MODULE)

    @property
    def artifact_class(self) -> str | None:
        return getattr(self._source, "artifact_class", None)

    def fit_ingredients(self) -> dict:
        """The encodable fit-identity form: the source's ingredients plus the head selection."""
        from steerability.algorithms.state_control.base import _fit_ingredients

        return {
            "kind": "head_selection",
            "source": _fit_ingredients(self._source),
            "selected_heads": self._selected_heads,
            "num_heads": self._num_heads,
        }

    def __call__(self, ctx) -> BaseTransform:
        steering_vector = ctx.resolve(self._source)

        if self._selected_heads is not None:
            selected = self._selected_heads
        else:
            if steering_vector.probe_accuracies is None:
                raise ValueError(
                    "steering_vector has no probe_accuracies. "
                    "Either provide selected_heads explicitly or use data to train a new vector."
                )
            selected = TopKHeadSelector(self._num_heads).select(steering_vector=steering_vector)

        active_heads: dict[int, set[int]] = {}
        for layer_id, head_id in selected:
            active_heads.setdefault(layer_id, set()).add(head_id)

        transform: BaseTransform = HeadAdditiveTransform(
            steering_vector, active_heads=active_heads, strength=self._alpha,
        )
        if self._norm_preserving:
            transform = NormPreservingTransform(transform)
        return transform

_alpha = alpha instance-attribute

_norm_preserving = norm_preserving instance-attribute

_num_heads = num_heads instance-attribute

_selected_heads = selected_heads instance-attribute

_source = source instance-attribute

access property

artifact_class property

__call__(ctx)

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

def __call__(self, ctx) -> BaseTransform:
    steering_vector = ctx.resolve(self._source)

    if self._selected_heads is not None:
        selected = self._selected_heads
    else:
        if steering_vector.probe_accuracies is None:
            raise ValueError(
                "steering_vector has no probe_accuracies. "
                "Either provide selected_heads explicitly or use data to train a new vector."
            )
        selected = TopKHeadSelector(self._num_heads).select(steering_vector=steering_vector)

    active_heads: dict[int, set[int]] = {}
    for layer_id, head_id in selected:
        active_heads.setdefault(layer_id, set()).add(head_id)

    transform: BaseTransform = HeadAdditiveTransform(
        steering_vector, active_heads=active_heads, strength=self._alpha,
    )
    if self._norm_preserving:
        transform = NormPreservingTransform(transform)
    return transform

__init__(source, selected_heads, num_heads, alpha, norm_preserving)

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

def __init__(self, source, selected_heads, num_heads: int, alpha: float, norm_preserving: bool):
    self._source = source
    self._selected_heads = selected_heads
    self._num_heads = num_heads
    self._alpha = alpha
    self._norm_preserving = norm_preserving

fit_ingredients()

The encodable fit-identity form: the source's ingredients plus the head selection.

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

def fit_ingredients(self) -> dict:
    """The encodable fit-identity form: the source's ingredients plus the head selection."""
    from steerability.algorithms.state_control.base import _fit_ingredients

    return {
        "kind": "head_selection",
        "source": _fit_ingredients(self._source),
        "selected_heads": self._selected_heads,
        "num_heads": self._num_heads,
    }

_ProbeMassShiftFit

A fit recipe for ITI's per-head mass-shift vector.

Fitting captures pre-o_proj per-head activations, a capture kind no backend advertises, so the fit requires a live model.

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

class _ProbeMassShiftFit:
    """A fit recipe for ITI's per-head mass-shift vector.

    Fitting captures pre-`o_proj` per-head activations, a capture kind no backend advertises,
    so the fit requires a live model.
    """

    access = ModelAccess.MODULE
    artifact_class = "direction"

    def __init__(self, data, train_spec):
        self._data = data
        self._train_spec = train_spec

    def fit_ingredients(self) -> dict:
        """The encodable fit-identity form: the labeled data and the train spec."""
        return {"kind": "probe_mass_shift", "data": self._data, "train_spec": self._train_spec}

    def resolve(self, model, tokenizer, *, session=None) -> SteeringVector:
        if model is None:
            raise ValueError("Fitting ITI from data requires a live model at steer time.")
        return ProbeMassShiftEstimator().fit(model, tokenizer, data=self._data, spec=self._train_spec)

_data = data instance-attribute

_train_spec = train_spec instance-attribute

access = ModelAccess.MODULE class-attribute instance-attribute

artifact_class = 'direction' class-attribute instance-attribute

__init__(data, train_spec)

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

def __init__(self, data, train_spec):
    self._data = data
    self._train_spec = train_spec

fit_ingredients()

The encodable fit-identity form: the labeled data and the train spec.

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

def fit_ingredients(self) -> dict:
    """The encodable fit-identity form: the labeled data and the train spec."""
    return {"kind": "probe_mass_shift", "data": self._data, "train_spec": self._train_spec}

resolve(model, tokenizer, *, session=None)

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

def resolve(self, model, tokenizer, *, session=None) -> SteeringVector:
    if model is None:
        raise ValueError("Fitting ITI from data requires a live model at steer time.")
    return ProbeMassShiftEstimator().fit(model, tokenizer, data=self._data, spec=self._train_spec)

utils

ITI utility modules.

__all__ = ['ProbeMassShiftEstimator'] module-attribute

ProbeMassShiftEstimator

Bases: BaseEstimator[SteeringVector]

Learns per-head direction vectors using probe-based mass mean shift.

The fit has three parts, run per attention head across every layer:

1. Extract the head's pre-o_proj activation on labeled positive/negative statements, pooled to one vector per statement.
2. Fit a logistic-regression probe on an 80/20 held-out split and record its validation accuracy for later head selection. The split is drawn over groups when the data carries group keys, so no group's statements straddle the train/validation partition.
3. Compute the mass-mean-shift direction as mean(positives) - mean(negatives), L2-normalize it to a unit direction, and scale it by the standard deviation of all statements projected onto that direction.

Returns a SteeringVector with directions shaped [num_heads, head_dim] per layer, num_heads/head_dim metadata, and probe_accuracies populated for every (layer, head) pair. "Probe" here is the fit-time head classifier used to rank heads, not the toolkit's Probe detector.

Per-token activations are pooled inside the capture hook and never retained beyond the chunk that produced them, so host memory scales with the number of statements rather than their length.

Reference:

- "Inference-Time Intervention: Eliciting Truthful Answers from a Language Model"
  Kenneth Li, Oam Patel, Fernanda Viégas, Hanspeter Pfister, Martin Wattenberg
  [https://arxiv.org/abs/2306.03341](https://arxiv.org/abs/2306.03341)

Source code in steerability/algorithms/state_control/iti/utils/estimator.py

class ProbeMassShiftEstimator(BaseEstimator[SteeringVector]):
    """Learns per-head direction vectors using probe-based mass mean shift.

    The fit has three parts, run per attention head across every layer:

    1. Extract the head's pre-`o_proj` activation on labeled positive/negative statements, pooled
       to one vector per statement.
    2. Fit a logistic-regression probe on an 80/20 held-out split and record its validation
       accuracy for later head selection. The split is drawn over groups when the data carries
       group keys, so no group's statements straddle the train/validation partition.
    3. Compute the mass-mean-shift direction as `mean(positives) - mean(negatives)`, L2-normalize
       it to a unit direction, and scale it by the standard deviation of all statements projected
       onto that direction.

    Returns a `SteeringVector` with directions shaped `[num_heads, head_dim]` per layer,
    `num_heads`/`head_dim` metadata, and `probe_accuracies` populated for every `(layer, head)`
    pair. "Probe" here is the fit-time head classifier used to rank heads, not the toolkit's
    `Probe` detector.

    Per-token activations are pooled inside the capture hook and never retained beyond the chunk
    that produced them, so host memory scales with the number of statements rather than their
    length.

    Reference:

        - "Inference-Time Intervention: Eliciting Truthful Answers from a Language Model"
          Kenneth Li, Oam Patel, Fernanda Viégas, Hanspeter Pfister, Martin Wattenberg
          [https://arxiv.org/abs/2306.03341](https://arxiv.org/abs/2306.03341)
    """

    def fit(
        self,
        model: PreTrainedModel,
        tokenizer: PreTrainedTokenizerBase,
        *,
        data: LabeledExamples,
        spec: VectorTrainSpec,
    ) -> SteeringVector:
        """Extract per-head steering directions and probe accuracies from labeled statements.

        Args:
            model: Model to extract attention outputs from.
            tokenizer: Tokenizer for encoding the labeled examples.
            data: Independent positive/negative statements (true/false statements for ITI). Unlike
                `ContrastivePairs`, these do not need to be equal length. When the data carries
                group keys, the probe train/validation split is drawn over groups.
            spec: Training configuration (`accumulate` mode and `batch_size`).

        Returns:
            A `SteeringVector` with directions shaped `[num_heads, head_dim]` per layer,
            `num_heads`/`head_dim` metadata, `probe_accuracies` for every head, and
            `meta["probe_split"]` recording whether the split was drawn over groups
            (`"group"`) or statements (`"statement"`).

        Raises:
            ValueError: If the model is a hybrid attention stack (some decoder layers carry no
                attention module), its attention head geometry is not uniform across layers,
                `spec.accumulate` is unsupported, or the probe partition cannot be drawn.
        """
        model_type = getattr(model.config, "model_type", "unknown")

        # ITI reshapes every layer's o_proj input with one head geometry, so every layer must carry
        # an attention module and the geometry must be uniform. refuse hybrid stacks, then read the
        # geometry per layer from the module tree and fail loudly before any capture on a model with
        # heterogeneous heads (e.g. Gemma 4 alternates sliding and global head dims).
        layout = resolve_model_layout(model)
        if layout.is_hybrid:
            raise ValueError(
                "ITI requires an attention module on every decoder layer, but only layers "
                f"{list(layout.attention_layers)} of {layout.num_layers} carry one; hybrid "
                "attention stacks such as Qwen3.5 are not supported by ITI."
            )
        geometries = {lid: head_geometry(model, layout, lid) for lid in range(layout.num_layers)}
        distinct = {(g.num_heads, g.head_dim) for g in geometries.values()}
        if len(distinct) > 1:
            differing = sorted((lid, g.num_heads, g.head_dim) for lid, g in geometries.items())
            raise ValueError(
                "ITI requires uniform attention head geometry across layers, but the model has "
                f"heterogeneous heads: {differing} as (layer, num_heads, head_dim). Models such as "
                "Gemma 4 that alternate sliding and global head dimensions are not supported by ITI."
            )
        geometry = next(iter(geometries.values()))
        num_heads = geometry.num_heads
        head_dim = geometry.head_dim

        if spec.accumulate not in ("last_token", "all"):
            raise ValueError(f"ProbeMassShiftEstimator does not support accumulate='{spec.accumulate}'.")

        device = next(model.parameters()).device
        pos_texts = list(data.positives)
        neg_texts = list(data.negatives)
        n_pos = len(pos_texts)
        n_neg = len(neg_texts)

        logger.debug("Tokenizing %d positive and %d negative statements", n_pos, n_neg)
        enc_pos = tokenize_texts(tokenizer, pos_texts, device)
        enc_neg = tokenize_texts(tokenizer, neg_texts, device)

        logger.debug("Extracting pooled head features with batch_size=%d", spec.batch_size)
        feats_pos = _extract_pooled_head_features(model, enc_pos, spec.batch_size, spec.accumulate)
        feats_neg = _extract_pooled_head_features(model, enc_neg, spec.batch_size, spec.accumulate)

        labels = np.array([1] * n_pos + [0] * n_neg)
        groups = None
        if data.groups:
            groups = list(data.positive_groups) + list(data.negative_groups)
        train_idx, val_idx = _probe_partition(labels, groups)

        num_layers = len(feats_pos)
        logger.debug("Computing probe-based directions for %d layers x %d heads", num_layers, num_heads)

        directions: dict[int, torch.Tensor] = {}
        probe_accuracies: dict[tuple[int, int], float] = {}

        # sklearn's lbfgs fit warns under recent scipy (iprint) and at the iteration cap; neither affects the fit
        with warnings.catch_warnings():
            warnings.simplefilter("ignore", category=OptimizeWarning)
            warnings.simplefilter("ignore", category=ConvergenceWarning)

            for layer_id in range(num_layers):
                pos_heads = feats_pos[layer_id].view(n_pos, num_heads, head_dim)
                neg_heads = feats_neg[layer_id].view(n_neg, num_heads, head_dim)

                layer_directions = []
                for head_id in range(num_heads):
                    hp = pos_heads[:, head_id, :].float()  # [n_pos, head_dim]
                    hn = neg_heads[:, head_id, :].float()  # [n_neg, head_dim]
                    x = torch.cat([hp, hn], dim=0)  # [N, head_dim]
                    features = x.numpy()

                    probe = LogisticRegression(max_iter=1000, solver="lbfgs")
                    probe.fit(features[train_idx], labels[train_idx])
                    probe_accuracies[(layer_id, head_id)] = float(probe.score(features[val_idx], labels[val_idx]))

                    raw = hp.mean(dim=0) - hn.mean(dim=0)  # [head_dim]
                    norm = raw.norm()
                    theta_hat = raw / norm if norm > 0 else raw
                    sigma = (x @ theta_hat).std()
                    layer_directions.append((sigma * theta_hat).to(dtype=torch.float32))

                directions[layer_id] = torch.stack(layer_directions, dim=0)  # [num_heads, head_dim]

        logger.debug("Finished fitting probe-based head directions")
        meta = {
            **artifact_provenance_meta(model, tokenizer),
            "probe_split": "group" if groups is not None else "statement",
        }
        return SteeringVector(
            model_type=model_type,
            directions=directions,
            num_heads=num_heads,
            head_dim=head_dim,
            probe_accuracies=probe_accuracies,
            meta=meta,
        )

fit(model, tokenizer, *, data, spec)

Extract per-head steering directions and probe accuracies from labeled statements.

Parameters:

Name	Type	Description	Default
model	PreTrainedModel	Model to extract attention outputs from.	required
tokenizer	PreTrainedTokenizerBase	Tokenizer for encoding the labeled examples.	required
data	LabeledExamples	Independent positive/negative statements (true/false statements for ITI). Unlike ContrastivePairs, these do not need to be equal length. When the data carries group keys, the probe train/validation split is drawn over groups.	required
spec	VectorTrainSpec	Training configuration (accumulate mode and batch_size).	required

Returns:

Type	Description
SteeringVector	A SteeringVector with directions shaped [num_heads, head_dim] per layer,
SteeringVector	num_heads/head_dim metadata, probe_accuracies for every head, and
SteeringVector	meta["probe_split"] recording whether the split was drawn over groups
SteeringVector	("group") or statements ("statement").

Raises:

Type	Description
ValueError	If the model is a hybrid attention stack (some decoder layers carry no attention module), its attention head geometry is not uniform across layers, spec.accumulate is unsupported, or the probe partition cannot be drawn.

Source code in steerability/algorithms/state_control/iti/utils/estimator.py

def fit(
    self,
    model: PreTrainedModel,
    tokenizer: PreTrainedTokenizerBase,
    *,
    data: LabeledExamples,
    spec: VectorTrainSpec,
) -> SteeringVector:
    """Extract per-head steering directions and probe accuracies from labeled statements.

    Args:
        model: Model to extract attention outputs from.
        tokenizer: Tokenizer for encoding the labeled examples.
        data: Independent positive/negative statements (true/false statements for ITI). Unlike
            `ContrastivePairs`, these do not need to be equal length. When the data carries
            group keys, the probe train/validation split is drawn over groups.
        spec: Training configuration (`accumulate` mode and `batch_size`).

    Returns:
        A `SteeringVector` with directions shaped `[num_heads, head_dim]` per layer,
        `num_heads`/`head_dim` metadata, `probe_accuracies` for every head, and
        `meta["probe_split"]` recording whether the split was drawn over groups
        (`"group"`) or statements (`"statement"`).

    Raises:
        ValueError: If the model is a hybrid attention stack (some decoder layers carry no
            attention module), its attention head geometry is not uniform across layers,
            `spec.accumulate` is unsupported, or the probe partition cannot be drawn.
    """
    model_type = getattr(model.config, "model_type", "unknown")

    # ITI reshapes every layer's o_proj input with one head geometry, so every layer must carry
    # an attention module and the geometry must be uniform. refuse hybrid stacks, then read the
    # geometry per layer from the module tree and fail loudly before any capture on a model with
    # heterogeneous heads (e.g. Gemma 4 alternates sliding and global head dims).
    layout = resolve_model_layout(model)
    if layout.is_hybrid:
        raise ValueError(
            "ITI requires an attention module on every decoder layer, but only layers "
            f"{list(layout.attention_layers)} of {layout.num_layers} carry one; hybrid "
            "attention stacks such as Qwen3.5 are not supported by ITI."
        )
    geometries = {lid: head_geometry(model, layout, lid) for lid in range(layout.num_layers)}
    distinct = {(g.num_heads, g.head_dim) for g in geometries.values()}
    if len(distinct) > 1:
        differing = sorted((lid, g.num_heads, g.head_dim) for lid, g in geometries.items())
        raise ValueError(
            "ITI requires uniform attention head geometry across layers, but the model has "
            f"heterogeneous heads: {differing} as (layer, num_heads, head_dim). Models such as "
            "Gemma 4 that alternate sliding and global head dimensions are not supported by ITI."
        )
    geometry = next(iter(geometries.values()))
    num_heads = geometry.num_heads
    head_dim = geometry.head_dim

    if spec.accumulate not in ("last_token", "all"):
        raise ValueError(f"ProbeMassShiftEstimator does not support accumulate='{spec.accumulate}'.")

    device = next(model.parameters()).device
    pos_texts = list(data.positives)
    neg_texts = list(data.negatives)
    n_pos = len(pos_texts)
    n_neg = len(neg_texts)

    logger.debug("Tokenizing %d positive and %d negative statements", n_pos, n_neg)
    enc_pos = tokenize_texts(tokenizer, pos_texts, device)
    enc_neg = tokenize_texts(tokenizer, neg_texts, device)

    logger.debug("Extracting pooled head features with batch_size=%d", spec.batch_size)
    feats_pos = _extract_pooled_head_features(model, enc_pos, spec.batch_size, spec.accumulate)
    feats_neg = _extract_pooled_head_features(model, enc_neg, spec.batch_size, spec.accumulate)

    labels = np.array([1] * n_pos + [0] * n_neg)
    groups = None
    if data.groups:
        groups = list(data.positive_groups) + list(data.negative_groups)
    train_idx, val_idx = _probe_partition(labels, groups)

    num_layers = len(feats_pos)
    logger.debug("Computing probe-based directions for %d layers x %d heads", num_layers, num_heads)

    directions: dict[int, torch.Tensor] = {}
    probe_accuracies: dict[tuple[int, int], float] = {}

    # sklearn's lbfgs fit warns under recent scipy (iprint) and at the iteration cap; neither affects the fit
    with warnings.catch_warnings():
        warnings.simplefilter("ignore", category=OptimizeWarning)
        warnings.simplefilter("ignore", category=ConvergenceWarning)

        for layer_id in range(num_layers):
            pos_heads = feats_pos[layer_id].view(n_pos, num_heads, head_dim)
            neg_heads = feats_neg[layer_id].view(n_neg, num_heads, head_dim)

            layer_directions = []
            for head_id in range(num_heads):
                hp = pos_heads[:, head_id, :].float()  # [n_pos, head_dim]
                hn = neg_heads[:, head_id, :].float()  # [n_neg, head_dim]
                x = torch.cat([hp, hn], dim=0)  # [N, head_dim]
                features = x.numpy()

                probe = LogisticRegression(max_iter=1000, solver="lbfgs")
                probe.fit(features[train_idx], labels[train_idx])
                probe_accuracies[(layer_id, head_id)] = float(probe.score(features[val_idx], labels[val_idx]))

                raw = hp.mean(dim=0) - hn.mean(dim=0)  # [head_dim]
                norm = raw.norm()
                theta_hat = raw / norm if norm > 0 else raw
                sigma = (x @ theta_hat).std()
                layer_directions.append((sigma * theta_hat).to(dtype=torch.float32))

            directions[layer_id] = torch.stack(layer_directions, dim=0)  # [num_heads, head_dim]

    logger.debug("Finished fitting probe-based head directions")
    meta = {
        **artifact_provenance_meta(model, tokenizer),
        "probe_split": "group" if groups is not None else "statement",
    }
    return SteeringVector(
        model_type=model_type,
        directions=directions,
        num_heads=num_heads,
        head_dim=head_dim,
        probe_accuracies=probe_accuracies,
        meta=meta,
    )

estimator

Probe-based mass mean shift estimator for ITI.

The per-head classifiers this estimator fits are fit-time selection statistics. Their held-out accuracies rank heads and their weights are discarded. They are ordinary scikit-learn logistic regressions over per-head activation slices, unrelated to the toolkit's Probe detector (a calibrated, model-free readout over pooled residual-stream features consumed by gates and routers). "Probe" here means only ITI's fit-time head classifier.

_SPLIT_SEED = 42 module-attribute

_VAL_FRACTION = 0.2 module-attribute

logger = logging.getLogger(__name__) module-attribute

ProbeMassShiftEstimator

Bases: BaseEstimator[SteeringVector]

Learns per-head direction vectors using probe-based mass mean shift.

The fit has three parts, run per attention head across every layer:

1. Extract the head's pre-o_proj activation on labeled positive/negative statements, pooled to one vector per statement.
2. Fit a logistic-regression probe on an 80/20 held-out split and record its validation accuracy for later head selection. The split is drawn over groups when the data carries group keys, so no group's statements straddle the train/validation partition.
3. Compute the mass-mean-shift direction as mean(positives) - mean(negatives), L2-normalize it to a unit direction, and scale it by the standard deviation of all statements projected onto that direction.

Returns a SteeringVector with directions shaped [num_heads, head_dim] per layer, num_heads/head_dim metadata, and probe_accuracies populated for every (layer, head) pair. "Probe" here is the fit-time head classifier used to rank heads, not the toolkit's Probe detector.

Per-token activations are pooled inside the capture hook and never retained beyond the chunk that produced them, so host memory scales with the number of statements rather than their length.

Reference:

- "Inference-Time Intervention: Eliciting Truthful Answers from a Language Model"
  Kenneth Li, Oam Patel, Fernanda Viégas, Hanspeter Pfister, Martin Wattenberg
  [https://arxiv.org/abs/2306.03341](https://arxiv.org/abs/2306.03341)

Source code in steerability/algorithms/state_control/iti/utils/estimator.py

class ProbeMassShiftEstimator(BaseEstimator[SteeringVector]):
    """Learns per-head direction vectors using probe-based mass mean shift.

    The fit has three parts, run per attention head across every layer:

    1. Extract the head's pre-`o_proj` activation on labeled positive/negative statements, pooled
       to one vector per statement.
    2. Fit a logistic-regression probe on an 80/20 held-out split and record its validation
       accuracy for later head selection. The split is drawn over groups when the data carries
       group keys, so no group's statements straddle the train/validation partition.
    3. Compute the mass-mean-shift direction as `mean(positives) - mean(negatives)`, L2-normalize
       it to a unit direction, and scale it by the standard deviation of all statements projected
       onto that direction.

    Returns a `SteeringVector` with directions shaped `[num_heads, head_dim]` per layer,
    `num_heads`/`head_dim` metadata, and `probe_accuracies` populated for every `(layer, head)`
    pair. "Probe" here is the fit-time head classifier used to rank heads, not the toolkit's
    `Probe` detector.

    Per-token activations are pooled inside the capture hook and never retained beyond the chunk
    that produced them, so host memory scales with the number of statements rather than their
    length.

    Reference:

        - "Inference-Time Intervention: Eliciting Truthful Answers from a Language Model"
          Kenneth Li, Oam Patel, Fernanda Viégas, Hanspeter Pfister, Martin Wattenberg
          [https://arxiv.org/abs/2306.03341](https://arxiv.org/abs/2306.03341)
    """

    def fit(
        self,
        model: PreTrainedModel,
        tokenizer: PreTrainedTokenizerBase,
        *,
        data: LabeledExamples,
        spec: VectorTrainSpec,
    ) -> SteeringVector:
        """Extract per-head steering directions and probe accuracies from labeled statements.

        Args:
            model: Model to extract attention outputs from.
            tokenizer: Tokenizer for encoding the labeled examples.
            data: Independent positive/negative statements (true/false statements for ITI). Unlike
                `ContrastivePairs`, these do not need to be equal length. When the data carries
                group keys, the probe train/validation split is drawn over groups.
            spec: Training configuration (`accumulate` mode and `batch_size`).

        Returns:
            A `SteeringVector` with directions shaped `[num_heads, head_dim]` per layer,
            `num_heads`/`head_dim` metadata, `probe_accuracies` for every head, and
            `meta["probe_split"]` recording whether the split was drawn over groups
            (`"group"`) or statements (`"statement"`).

        Raises:
            ValueError: If the model is a hybrid attention stack (some decoder layers carry no
                attention module), its attention head geometry is not uniform across layers,
                `spec.accumulate` is unsupported, or the probe partition cannot be drawn.
        """
        model_type = getattr(model.config, "model_type", "unknown")

        # ITI reshapes every layer's o_proj input with one head geometry, so every layer must carry
        # an attention module and the geometry must be uniform. refuse hybrid stacks, then read the
        # geometry per layer from the module tree and fail loudly before any capture on a model with
        # heterogeneous heads (e.g. Gemma 4 alternates sliding and global head dims).
        layout = resolve_model_layout(model)
        if layout.is_hybrid:
            raise ValueError(
                "ITI requires an attention module on every decoder layer, but only layers "
                f"{list(layout.attention_layers)} of {layout.num_layers} carry one; hybrid "
                "attention stacks such as Qwen3.5 are not supported by ITI."
            )
        geometries = {lid: head_geometry(model, layout, lid) for lid in range(layout.num_layers)}
        distinct = {(g.num_heads, g.head_dim) for g in geometries.values()}
        if len(distinct) > 1:
            differing = sorted((lid, g.num_heads, g.head_dim) for lid, g in geometries.items())
            raise ValueError(
                "ITI requires uniform attention head geometry across layers, but the model has "
                f"heterogeneous heads: {differing} as (layer, num_heads, head_dim). Models such as "
                "Gemma 4 that alternate sliding and global head dimensions are not supported by ITI."
            )
        geometry = next(iter(geometries.values()))
        num_heads = geometry.num_heads
        head_dim = geometry.head_dim

        if spec.accumulate not in ("last_token", "all"):
            raise ValueError(f"ProbeMassShiftEstimator does not support accumulate='{spec.accumulate}'.")

        device = next(model.parameters()).device
        pos_texts = list(data.positives)
        neg_texts = list(data.negatives)
        n_pos = len(pos_texts)
        n_neg = len(neg_texts)

        logger.debug("Tokenizing %d positive and %d negative statements", n_pos, n_neg)
        enc_pos = tokenize_texts(tokenizer, pos_texts, device)
        enc_neg = tokenize_texts(tokenizer, neg_texts, device)

        logger.debug("Extracting pooled head features with batch_size=%d", spec.batch_size)
        feats_pos = _extract_pooled_head_features(model, enc_pos, spec.batch_size, spec.accumulate)
        feats_neg = _extract_pooled_head_features(model, enc_neg, spec.batch_size, spec.accumulate)

        labels = np.array([1] * n_pos + [0] * n_neg)
        groups = None
        if data.groups:
            groups = list(data.positive_groups) + list(data.negative_groups)
        train_idx, val_idx = _probe_partition(labels, groups)

        num_layers = len(feats_pos)
        logger.debug("Computing probe-based directions for %d layers x %d heads", num_layers, num_heads)

        directions: dict[int, torch.Tensor] = {}
        probe_accuracies: dict[tuple[int, int], float] = {}

        # sklearn's lbfgs fit warns under recent scipy (iprint) and at the iteration cap; neither affects the fit
        with warnings.catch_warnings():
            warnings.simplefilter("ignore", category=OptimizeWarning)
            warnings.simplefilter("ignore", category=ConvergenceWarning)

            for layer_id in range(num_layers):
                pos_heads = feats_pos[layer_id].view(n_pos, num_heads, head_dim)
                neg_heads = feats_neg[layer_id].view(n_neg, num_heads, head_dim)

                layer_directions = []
                for head_id in range(num_heads):
                    hp = pos_heads[:, head_id, :].float()  # [n_pos, head_dim]
                    hn = neg_heads[:, head_id, :].float()  # [n_neg, head_dim]
                    x = torch.cat([hp, hn], dim=0)  # [N, head_dim]
                    features = x.numpy()

                    probe = LogisticRegression(max_iter=1000, solver="lbfgs")
                    probe.fit(features[train_idx], labels[train_idx])
                    probe_accuracies[(layer_id, head_id)] = float(probe.score(features[val_idx], labels[val_idx]))

                    raw = hp.mean(dim=0) - hn.mean(dim=0)  # [head_dim]
                    norm = raw.norm()
                    theta_hat = raw / norm if norm > 0 else raw
                    sigma = (x @ theta_hat).std()
                    layer_directions.append((sigma * theta_hat).to(dtype=torch.float32))

                directions[layer_id] = torch.stack(layer_directions, dim=0)  # [num_heads, head_dim]

        logger.debug("Finished fitting probe-based head directions")
        meta = {
            **artifact_provenance_meta(model, tokenizer),
            "probe_split": "group" if groups is not None else "statement",
        }
        return SteeringVector(
            model_type=model_type,
            directions=directions,
            num_heads=num_heads,
            head_dim=head_dim,
            probe_accuracies=probe_accuracies,
            meta=meta,
        )

fit(model, tokenizer, *, data, spec)

Extract per-head steering directions and probe accuracies from labeled statements.

Parameters:

Name	Type	Description	Default
model	PreTrainedModel	Model to extract attention outputs from.	required
tokenizer	PreTrainedTokenizerBase	Tokenizer for encoding the labeled examples.	required
data	LabeledExamples	Independent positive/negative statements (true/false statements for ITI). Unlike ContrastivePairs, these do not need to be equal length. When the data carries group keys, the probe train/validation split is drawn over groups.	required
spec	VectorTrainSpec	Training configuration (accumulate mode and batch_size).	required

Returns:

Type	Description
SteeringVector	A SteeringVector with directions shaped [num_heads, head_dim] per layer,
SteeringVector	num_heads/head_dim metadata, probe_accuracies for every head, and
SteeringVector	meta["probe_split"] recording whether the split was drawn over groups
SteeringVector	("group") or statements ("statement").

Raises:

Type	Description
ValueError	If the model is a hybrid attention stack (some decoder layers carry no attention module), its attention head geometry is not uniform across layers, spec.accumulate is unsupported, or the probe partition cannot be drawn.

Source code in steerability/algorithms/state_control/iti/utils/estimator.py

def fit(
    self,
    model: PreTrainedModel,
    tokenizer: PreTrainedTokenizerBase,
    *,
    data: LabeledExamples,
    spec: VectorTrainSpec,
) -> SteeringVector:
    """Extract per-head steering directions and probe accuracies from labeled statements.

    Args:
        model: Model to extract attention outputs from.
        tokenizer: Tokenizer for encoding the labeled examples.
        data: Independent positive/negative statements (true/false statements for ITI). Unlike
            `ContrastivePairs`, these do not need to be equal length. When the data carries
            group keys, the probe train/validation split is drawn over groups.
        spec: Training configuration (`accumulate` mode and `batch_size`).

    Returns:
        A `SteeringVector` with directions shaped `[num_heads, head_dim]` per layer,
        `num_heads`/`head_dim` metadata, `probe_accuracies` for every head, and
        `meta["probe_split"]` recording whether the split was drawn over groups
        (`"group"`) or statements (`"statement"`).

    Raises:
        ValueError: If the model is a hybrid attention stack (some decoder layers carry no
            attention module), its attention head geometry is not uniform across layers,
            `spec.accumulate` is unsupported, or the probe partition cannot be drawn.
    """
    model_type = getattr(model.config, "model_type", "unknown")

    # ITI reshapes every layer's o_proj input with one head geometry, so every layer must carry
    # an attention module and the geometry must be uniform. refuse hybrid stacks, then read the
    # geometry per layer from the module tree and fail loudly before any capture on a model with
    # heterogeneous heads (e.g. Gemma 4 alternates sliding and global head dims).
    layout = resolve_model_layout(model)
    if layout.is_hybrid:
        raise ValueError(
            "ITI requires an attention module on every decoder layer, but only layers "
            f"{list(layout.attention_layers)} of {layout.num_layers} carry one; hybrid "
            "attention stacks such as Qwen3.5 are not supported by ITI."
        )
    geometries = {lid: head_geometry(model, layout, lid) for lid in range(layout.num_layers)}
    distinct = {(g.num_heads, g.head_dim) for g in geometries.values()}
    if len(distinct) > 1:
        differing = sorted((lid, g.num_heads, g.head_dim) for lid, g in geometries.items())
        raise ValueError(
            "ITI requires uniform attention head geometry across layers, but the model has "
            f"heterogeneous heads: {differing} as (layer, num_heads, head_dim). Models such as "
            "Gemma 4 that alternate sliding and global head dimensions are not supported by ITI."
        )
    geometry = next(iter(geometries.values()))
    num_heads = geometry.num_heads
    head_dim = geometry.head_dim

    if spec.accumulate not in ("last_token", "all"):
        raise ValueError(f"ProbeMassShiftEstimator does not support accumulate='{spec.accumulate}'.")

    device = next(model.parameters()).device
    pos_texts = list(data.positives)
    neg_texts = list(data.negatives)
    n_pos = len(pos_texts)
    n_neg = len(neg_texts)

    logger.debug("Tokenizing %d positive and %d negative statements", n_pos, n_neg)
    enc_pos = tokenize_texts(tokenizer, pos_texts, device)
    enc_neg = tokenize_texts(tokenizer, neg_texts, device)

    logger.debug("Extracting pooled head features with batch_size=%d", spec.batch_size)
    feats_pos = _extract_pooled_head_features(model, enc_pos, spec.batch_size, spec.accumulate)
    feats_neg = _extract_pooled_head_features(model, enc_neg, spec.batch_size, spec.accumulate)

    labels = np.array([1] * n_pos + [0] * n_neg)
    groups = None
    if data.groups:
        groups = list(data.positive_groups) + list(data.negative_groups)
    train_idx, val_idx = _probe_partition(labels, groups)

    num_layers = len(feats_pos)
    logger.debug("Computing probe-based directions for %d layers x %d heads", num_layers, num_heads)

    directions: dict[int, torch.Tensor] = {}
    probe_accuracies: dict[tuple[int, int], float] = {}

    # sklearn's lbfgs fit warns under recent scipy (iprint) and at the iteration cap; neither affects the fit
    with warnings.catch_warnings():
        warnings.simplefilter("ignore", category=OptimizeWarning)
        warnings.simplefilter("ignore", category=ConvergenceWarning)

        for layer_id in range(num_layers):
            pos_heads = feats_pos[layer_id].view(n_pos, num_heads, head_dim)
            neg_heads = feats_neg[layer_id].view(n_neg, num_heads, head_dim)

            layer_directions = []
            for head_id in range(num_heads):
                hp = pos_heads[:, head_id, :].float()  # [n_pos, head_dim]
                hn = neg_heads[:, head_id, :].float()  # [n_neg, head_dim]
                x = torch.cat([hp, hn], dim=0)  # [N, head_dim]
                features = x.numpy()

                probe = LogisticRegression(max_iter=1000, solver="lbfgs")
                probe.fit(features[train_idx], labels[train_idx])
                probe_accuracies[(layer_id, head_id)] = float(probe.score(features[val_idx], labels[val_idx]))

                raw = hp.mean(dim=0) - hn.mean(dim=0)  # [head_dim]
                norm = raw.norm()
                theta_hat = raw / norm if norm > 0 else raw
                sigma = (x @ theta_hat).std()
                layer_directions.append((sigma * theta_hat).to(dtype=torch.float32))

            directions[layer_id] = torch.stack(layer_directions, dim=0)  # [num_heads, head_dim]

    logger.debug("Finished fitting probe-based head directions")
    meta = {
        **artifact_provenance_meta(model, tokenizer),
        "probe_split": "group" if groups is not None else "statement",
    }
    return SteeringVector(
        model_type=model_type,
        directions=directions,
        num_heads=num_heads,
        head_dim=head_dim,
        probe_accuracies=probe_accuracies,
        meta=meta,
    )

_MaskHolder

Carries the current chunk's attention mask so capture hooks can pool per row.

Source code in steerability/algorithms/state_control/iti/utils/estimator.py

class _MaskHolder:
    """Carries the current chunk's attention mask so capture hooks can pool per row."""

    __slots__ = ("mask",)

    def __init__(self):
        self.mask: torch.Tensor | None = None

__slots__ = ('mask',) class-attribute instance-attribute

mask = None instance-attribute

__init__()

Source code in steerability/algorithms/state_control/iti/utils/estimator.py

def __init__(self):
    self.mask: torch.Tensor | None = None

_extract_pooled_head_features(model, enc, batch_size, accumulate)

Extract pooled pre-o_proj head features from every layer via temporary hooks.

Registers one forward_pre_hook on each layer's output projection (o_proj / c_proj), which captures the concatenated per-head attention output before the projection mixes it. Each chunk's [b, T, D] input is pooled to [b, D] on the model's device inside the hook, then moved to CPU in the captured dtype and appended to the layer's list. Nothing 3-D is ever retained, so host memory scales with the number of statements rather than their length.

Parameters:

Name	Type	Description	Default
model	PreTrainedModel	The model to extract from.	required
enc	dict[str, Tensor]	Tokenized input with input_ids and (optionally) attention_mask.	required
batch_size	int	Chunk size for the forward passes.	required
accumulate	str	"last_token" selects the last non-pad position per row; "all" mean-pools over the non-pad positions.	required

Returns:

Type	Description
dict[int, Tensor]	A mapping from layer_id to a [N, num_heads * head_dim] CPU tensor in the captured dtype.

Source code in steerability/algorithms/state_control/iti/utils/estimator.py

@torch.no_grad()
def _extract_pooled_head_features(
    model: PreTrainedModel,
    enc: dict[str, torch.Tensor],
    batch_size: int,
    accumulate: str,
) -> dict[int, torch.Tensor]:
    """Extract pooled pre-`o_proj` head features from every layer via temporary hooks.

    Registers one `forward_pre_hook` on each layer's output projection (`o_proj` / `c_proj`),
    which captures the concatenated per-head attention output before the projection mixes it. Each
    chunk's `[b, T, D]` input is pooled to `[b, D]` on the model's device inside the hook, then
    moved to CPU in the captured dtype and appended to the layer's list. Nothing 3-D is ever
    retained, so host memory scales with the number of statements rather than their length.

    Args:
        model: The model to extract from.
        enc: Tokenized input with `input_ids` and (optionally) `attention_mask`.
        batch_size: Chunk size for the forward passes.
        accumulate: `"last_token"` selects the last non-pad position per row; `"all"` mean-pools
            over the non-pad positions.

    Returns:
        A mapping from `layer_id` to a `[N, num_heads * head_dim]` CPU tensor in the captured dtype.
    """
    input_ids = enc["input_ids"]
    attention_mask = enc.get("attention_mask")
    num_examples = input_ids.size(0)

    layout = resolve_model_layout(model)
    oproj_names = layout.oproj_names
    num_layers = layout.num_layers

    storage: dict[int, list[torch.Tensor]] = {i: [] for i in range(num_layers)}
    holder = _MaskHolder()
    handles: list[torch.utils.hooks.RemovableHandle] = []

    def make_pre_hook(layer_id: int):
        def hook(_module, args, kwargs):
            x = args[0] if args else kwargs.get("input")  # [b, T, D]
            mask = holder.mask
            batch = x.size(0)
            if accumulate == "last_token":
                positions = get_last_token_positions(mask, x.size(1), batch)
                pooled = select_at_positions(x, positions.to(x.device))  # [b, D]
            else:
                pooled = masked_mean(x, mask)  # [b, D]
            storage[layer_id].append(pooled.detach().to("cpu"))
        return hook

    try:
        for layer_id, oproj_name in enumerate(oproj_names):
            oproj_module = model.get_submodule(oproj_name)
            handles.append(oproj_module.register_forward_pre_hook(make_pre_hook(layer_id), with_kwargs=True))

        for start in range(0, num_examples, batch_size):
            end = min(start + batch_size, num_examples)
            batch_ids = input_ids[start:end]
            batch_mask = attention_mask[start:end] if attention_mask is not None else None
            holder.mask = batch_mask
            model(input_ids=batch_ids, attention_mask=batch_mask, use_cache=False)
    finally:
        for handle in handles:
            handle.remove()

    return {layer_id: torch.cat(tensors, dim=0) for layer_id, tensors in storage.items()}

_probe_partition(labels, groups)

Draw one train/validation partition over the concatenated positives-then-negatives space.

One partition is computed before the head loop and shared by every head. Without groups the split is stratified by label; with groups it is drawn over groups so that no group's statements appear on both sides.

Parameters:

Name	Type	Description	Default
labels	ndarray	Binary labels over the concatenated index space (positives then negatives).	required
groups	Sequence[str | int] | None	Group key per statement in the same index space, or None for a stratified split.	required

Returns:

Type	Description
tuple[ndarray, ndarray]	A (train_idx, val_idx) pair of integer index arrays.

Raises:

Type	Description
ValueError	If a grouped partition has fewer than two distinct groups, or either side of the resulting partition lacks a class.

Source code in steerability/algorithms/state_control/iti/utils/estimator.py

def _probe_partition(
    labels: np.ndarray,
    groups: Sequence[str | int] | None,
) -> tuple[np.ndarray, np.ndarray]:
    """Draw one train/validation partition over the concatenated positives-then-negatives space.

    One partition is computed before the head loop and shared by every head. Without groups the
    split is stratified by label; with groups it is drawn over groups so that no group's statements
    appear on both sides.

    Args:
        labels: Binary labels over the concatenated index space (positives then negatives).
        groups: Group key per statement in the same index space, or None for a stratified split.

    Returns:
        A `(train_idx, val_idx)` pair of integer index arrays.

    Raises:
        ValueError: If a grouped partition has fewer than two distinct groups, or either side of
            the resulting partition lacks a class.
    """
    indices = np.arange(len(labels))
    if groups is None:
        train_idx, val_idx = train_test_split(
            indices, test_size=_VAL_FRACTION, random_state=_SPLIT_SEED, stratify=labels
        )
        return train_idx, val_idx

    groups = np.asarray(groups)
    if len(np.unique(groups)) < 2:
        raise ValueError(
            "Grouped probe split requires at least two distinct groups; add data or regroup."
        )
    splitter = GroupShuffleSplit(n_splits=1, test_size=_VAL_FRACTION, random_state=_SPLIT_SEED)
    train_idx, val_idx = next(splitter.split(indices, labels, groups))
    if len(np.unique(labels[train_idx])) < 2 or len(np.unique(labels[val_idx])) < 2:
        raise ValueError(
            "Grouped probe split left one side without both classes; add data or regroup."
        )
    return train_idx, val_idx