Directional Ablation

steerability.algorithms.state_control.directional_ablation

args

Directional Ablation argument validation.

control

Directional Ablation control: projects a learned direction out of the residual stream.

DirectionalAblation

Bases: InterventionControl

Directional Ablation (feature removal via projection).

Removes a learned feature direction from the residual stream at one or more layers during generation, h' = h - alpha * (d_hat^T h) d_hat at masked positions. This is the abliteration technique of Arditi et al., which learns a direction as the difference in means over contrastive data and projects it out.

The method operates in two phases:

1. Training (offline). Extract residual activations for contrastive pairs and take the mean difference, or the PCA of paired differences, as the feature direction. A precomputed direction (or an orthonormal subspace, K > 1) may be supplied directly.

2. Inference (online). At each target layer's output, project the direction out of the residual stream at masked positions. alpha = 1.0 fully removes the component (h'.d_hat == 0); alpha < 1.0 gives graded partial suppression.

Ablation is a projection (idempotent at alpha=1, norm-reducing). It can compose with the alignment-adaptive gate (AlignmentAdaptiveTransform) to ablate only where the feature is present.

The control is declarative: _configure maps the validated args onto one Intervention whose behavior layers are the target layers intersected with the artifact's coverage.

Reference:

• "Refusal in Language Models Is Mediated by a Single Direction" Andy Arditi, Oscar Obeso, Aaquib Syed, Daniel Paleka, Nina Panickssery, Wes Gurnee, Neel Nanda https://arxiv.org/abs/2406.11717

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

class DirectionalAblation(InterventionControl):
    """Directional Ablation (feature removal via projection).

    Removes a learned feature direction from the residual stream at one or more layers during
    generation, `h' = h - alpha * (d_hat^T h) d_hat` at masked positions. This is the abliteration
    technique of Arditi et al., which learns a direction as the difference in means over
    contrastive data and projects it out.

    The method operates in two phases:

    1. Training (offline). Extract residual activations for contrastive pairs and take the mean
       difference, or the PCA of paired differences, as the feature direction. A precomputed
       direction (or an orthonormal subspace, `K > 1`) may be supplied directly.

    2. Inference (online). At each target layer's output, project the direction out of the
       residual stream at masked positions. `alpha = 1.0` fully removes the component
       (`h'.d_hat == 0`); `alpha < 1.0` gives graded partial suppression.

    Ablation is a projection (idempotent at `alpha=1`, norm-reducing). It can compose with the
    alignment-adaptive gate (`AlignmentAdaptiveTransform`) to ablate only where the feature is
    present.

    The control is declarative: `_configure` maps the validated args onto one `Intervention`
    whose behavior layers are the target layers intersected with the artifact's coverage.

    Reference:

    - "Refusal in Language Models Is Mediated by a Single Direction"
    Andy Arditi, Oscar Obeso, Aaquib Syed, Daniel Paleka, Nina Panickssery, Wes Gurnee, Neel Nanda
    [https://arxiv.org/abs/2406.11717](https://arxiv.org/abs/2406.11717)
    """

    Args = DirectionalAblationArgs
    supports_batching = True

    def _configure(self):
        if self.steering_vector is not None:
            inner = _Precomputed(self.steering_vector.clone())
        else:
            estimator = (
                ContrastiveDirectionEstimator()
                if self.train_spec.method == "pca_pairwise"
                else MeanDifferenceEstimator()
            )
            inner = ContrastiveFit(
                data=self.data, estimator=estimator, estimator_kwargs={"spec": self.train_spec},
            )
        source = LayerFilteredFit(inner, layer_range=self.layer_range)

        transform = ProjectionTransform(source, alpha=self.alpha)
        if self.use_norm_preservation:
            transform = NormPreservingTransform(transform)

        self._template = (Intervention(
            # heuristic default: single layer at ~40% depth (matches CAA)
            layers=CoveredLayers(
                within=tuple(sorted(set(self.layer_ids))) if self.layer_ids is not None
                       else FractionalDepthSelector(fraction=0.4)
            ),
            transform=transform,
            scope=TokenScope(self.token_scope, last_k=self.last_k, from_position=self.from_position),
        ),)

    @property
    def hook_only_hint(self) -> str:
        if self.alpha != 1.0:
            return "graded ablation (alpha < 1) has no intervention-spec form; run on the huggingface backend"
        return "subspace ablation has no intervention-spec form; run on the huggingface backend"

    @property
    def _layer_ids(self) -> list[int]:
        """The resolved target layers (empty before `steer()`)."""
        return list(self.interventions[0].layers) if self.interventions else []

    @property
    def _steering_vector(self) -> SteeringVector | None:
        """The bound steering artifact as a `SteeringVector` view (None before `steer()`)."""
        if not self.interventions:
            return None
        core, _ = unwrap_modifiers(self.interventions[0].transform)
        if getattr(core, "directions", None) is None:
            return None
        return SteeringVector(
            model_type="unknown",
            directions=core.directions,
            meta=core.artifact_meta or {},
        )

RUNTIME_KWARGS_SCHEMA = [] class-attribute instance-attribute

_gate property

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

_layer_ids property

The resolved target layers (empty before steer()).

_steering_vector property

The bound steering artifact as a SteeringVector view (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 property

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/directional_ablation/control.py

def _configure(self):
    if self.steering_vector is not None:
        inner = _Precomputed(self.steering_vector.clone())
    else:
        estimator = (
            ContrastiveDirectionEstimator()
            if self.train_spec.method == "pca_pairwise"
            else MeanDifferenceEstimator()
        )
        inner = ContrastiveFit(
            data=self.data, estimator=estimator, estimator_kwargs={"spec": self.train_spec},
        )
    source = LayerFilteredFit(inner, layer_range=self.layer_range)

    transform = ProjectionTransform(source, alpha=self.alpha)
    if self.use_norm_preservation:
        transform = NormPreservingTransform(transform)

    self._template = (Intervention(
        # heuristic default: single layer at ~40% depth (matches CAA)
        layers=CoveredLayers(
            within=tuple(sorted(set(self.layer_ids))) if self.layer_ids is not None
                   else FractionalDepthSelector(fraction=0.4)
        ),
        transform=transform,
        scope=TokenScope(self.token_scope, last_k=self.last_k, from_position=self.from_position),
    ),)

_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()

Bound steering artifacts, keyed by intervention position.

For intervention i, the core transform's artifact exports as "intervention_{i}/transform" (a SteeringVector view), each wrapper's own artifact as "intervention_{i}/modifier_{j}", and the resolved gate as "intervention_{i}/gate" when the template's gate slot held a source. Must be called after steer().

Returns:

Type	Description
dict[str, Any]	Mapping from logical name to artifact value.

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

def export_state(self) -> dict[str, Any]:
    """Bound steering artifacts, keyed by intervention position.

    For intervention `i`, the core transform's artifact exports as
    `"intervention_{i}/transform"` (a `SteeringVector` view), each wrapper's own artifact
    as `"intervention_{i}/modifier_{j}"`, and the resolved gate as
    `"intervention_{i}/gate"` when the template's gate slot held a source. Must be called
    after `steer()`.

    Returns:
        Mapping from logical name to artifact value.
    """
    from steerability.algorithms.state_control.common.gating import Gate
    from steerability.algorithms.state_control.common.transforms.base import BaseTransform, unwrap_modifiers

    state: dict[str, Any] = {}
    for i, intervention in enumerate(self.interventions):
        if isinstance(intervention.transform, BaseTransform):
            core, wrappers = unwrap_modifiers(intervention.transform)
            view = _core_artifact_view(core)
            if view is not None:
                state[f"intervention_{i}/transform"] = view
            for j, wrapper in enumerate(wrappers):
                own = getattr(wrapper, "steering_vector", None)
                if own is not None:
                    state[f"intervention_{i}/modifier_{j}"] = own
        template_gate = self._template[i].gate if i < len(self._template) else None
        if isinstance(intervention.gate, Gate) and template_gate is not None \
                and not _is_concrete_gate(template_gate):
            state[f"intervention_{i}/gate"] = intervention.gate
    return state

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)

One activation_adapter entry per bound intervention, in intervention order.

Each entry's args mirror the intervention: the bound transform, the resolved behavior layers, the hook point, the resolved gate, and the token scope. Must be called after steer().

Returns:

Type	Description
list[tuple[str, dict[str, Any]]]	List of ("state_control/activation_adapter", kwargs) pairs.

Raises:

Type	Description
NotFreezableError	If the control is unsteered, an intervention hooks the norm_input site (no activation_adapter form), or an intervention follows an externally driven shared gate (an in-memory relationship that does not serialize).

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

def frozen_form(self, state: dict[str, Any]) -> list[tuple[str, dict[str, Any]]]:
    """One `activation_adapter` entry per bound intervention, in intervention order.

    Each entry's args mirror the intervention: the bound transform, the resolved behavior
    layers, the hook point, the resolved gate, and the token scope. Must be called after
    `steer()`.

    Returns:
        List of `("state_control/activation_adapter", kwargs)` pairs.

    Raises:
        NotFreezableError: If the control is unsteered, an intervention hooks the
            `norm_input` site (no `activation_adapter` form), or an intervention follows
            an externally driven shared gate (an in-memory relationship that does not
            serialize).
    """
    if not self.interventions:
        raise NotFreezableError(
            f"{type(self).__name__} has no bound interventions; call steer() before freezing."
        )
    entries: list[tuple[str, dict[str, Any]]] = []
    for intervention in self.interventions:
        if intervention.resolved_site() == "norm_input":
            raise NotFreezableError(
                f"{type(self).__name__} hooks the norm_input site, which has no "
                "activation_adapter form; configure the layer_output intervention point "
                "to freeze this control."
            )
        if intervention.gate_driven_externally:
            raise NotFreezableError(
                f"{type(self).__name__} follows an externally driven shared gate, an "
                "in-memory relationship that does not serialize; freeze the driving "
                "control's pipeline without the follower, or gate this control directly."
            )
        entries.append(("state_control/activation_adapter", {
            "transform": intervention.transform,
            "layer_ids": [int(lid) for lid in intervention.layers],
            "hook_point": intervention.boundary,
            "gate": intervention.gate,
            "token_scope": intervention.scope.kind,
            "last_k": intervention.scope.last_k,
            "from_position": intervention.scope.from_position,
        }))
    return entries

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()

The combined wire kinds of the bound interventions (or the template before steer()), or None when any intervention is hook-only.

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

def wire_kinds(self):
    """The combined wire kinds of the bound interventions (or the template before
    `steer()`), or None when any intervention is hook-only."""
    from steerability.algorithms.state_control.common.specs import combine_kinds

    source = self.interventions or self._template
    return combine_kinds(intervention.wire_kinds() for intervention in source)