CAA

steerability.algorithms.state_control.caa

args

CAA argument validation.

control

CAA

Bases: InterventionControl

Contrastive Activation Addition (CAA).

Steers model behavior by adding a learned mean-difference direction vector to the residual stream at a single layer during generation.

CAA operates in two phases:

1. Training (offline): Given contrastive prompt pairs where each pair shares the same question but ends with opposite answer tokens, extract residual stream activations at the answer-token position. The steering vector is the mean difference between positive and negative activations.

2. Inference (online): Add multiplier * v_L to the residual stream at a chosen layer L, at all token positions after the user's prompt. A positive multiplier increases the target behavior; negative decreases it.

The control is declarative: _configure maps the validated args onto one Intervention (an additive transform at one layer, optionally norm-preserving), and the base class binds it at steer().

Reference:

• "Steering Llama 2 via Contrastive Activation Addition" Nina Panickssery, Nick Gabrieli, Julian Schulz, Meg Tong, Evan Hubinger, Alexander Matt Turner https://arxiv.org/abs/2312.06681

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

class CAA(InterventionControl):
    """Contrastive Activation Addition (CAA).

    Steers model behavior by adding a learned mean-difference direction
    vector to the residual stream at a single layer during generation.

    CAA operates in two phases:

    1. **Training (offline)**: Given contrastive prompt pairs where each pair
       shares the same question but ends with opposite answer tokens, extract
       residual stream activations at the answer-token position. The steering
       vector is the mean difference between positive and negative activations.

    2. **Inference (online)**: Add `multiplier * v_L` to the residual stream
       at a chosen layer L, at all token positions after the user's prompt.
       A positive multiplier increases the target behavior; negative decreases it.

    The control is declarative: `_configure` maps the validated args onto one `Intervention`
    (an additive transform at one layer, optionally norm-preserving), and the base class binds
    it at `steer()`.

    Reference:

    - "Steering Llama 2 via Contrastive Activation Addition"
    Nina Panickssery, Nick Gabrieli, Julian Schulz, Meg Tong, Evan Hubinger, Alexander Matt Turner
    [https://arxiv.org/abs/2312.06681](https://arxiv.org/abs/2312.06681)
    """

    Args = CAAArgs
    supports_batching = True

    def _configure(self):
        if self.steering_vector is not None:
            if isinstance(self.steering_vector, SteeringVector):
                artifact = self.steering_vector.clone()
                if self.normalize_vector:
                    artifact = artifact.normalized()
                source = _Precomputed(artifact)
            else:
                source = self.steering_vector
        else:
            source = ContrastiveFit(
                data=self.data,
                method=self.train_spec.method,
                accumulate=self.train_spec.accumulate,
                batch_size=self.train_spec.batch_size,
                prompt_format=self.train_spec.prompt_format,
                location=self.train_spec.location,
                normalize=self.normalize_vector,
            )
        transform = AdditiveTransform(source, strength=self.multiplier)
        if self.use_norm_preservation:
            transform = NormPreservingTransform(transform)

        self._template = (Intervention(
            # heuristic default: ~40% depth (the paper finds layer 13/32 optimal)
            layers=(self.layer_id,) if self.layer_id 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 _layer_id(self) -> int | None:
        """The resolved behavior layer (None before `steer()`)."""
        return self.interventions[0].layers[0] if self.interventions else None

    @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 {},
        )

    def export_state(self) -> dict:
        """The bound 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 vector plus the resolved layer.

        The frozen args carry the exported `steering_vector`, the resolved `layer_id`, and the
        recipe's application parameters; `normalize_vector` is cleared since the exported
        vector is already in its applied form.
        """
        return "state_control/caa", {
            "steering_vector": state["steering_vector"],
            "layer_id": self._layer_id,
            "multiplier": self.multiplier,
            "token_scope": self.token_scope,
            "last_k": self.last_k,
            "from_position": self.from_position,
            "normalize_vector": False,
            "use_norm_preservation": self.use_norm_preservation,
        }

RUNTIME_KWARGS_SCHEMA = [] class-attribute instance-attribute

_gate property

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

_layer_id property

The resolved behavior layer (None 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 = None 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/caa/control.py

def _configure(self):
    if self.steering_vector is not None:
        if isinstance(self.steering_vector, SteeringVector):
            artifact = self.steering_vector.clone()
            if self.normalize_vector:
                artifact = artifact.normalized()
            source = _Precomputed(artifact)
        else:
            source = self.steering_vector
    else:
        source = ContrastiveFit(
            data=self.data,
            method=self.train_spec.method,
            accumulate=self.train_spec.accumulate,
            batch_size=self.train_spec.batch_size,
            prompt_format=self.train_spec.prompt_format,
            location=self.train_spec.location,
            normalize=self.normalize_vector,
        )
    transform = AdditiveTransform(source, strength=self.multiplier)
    if self.use_norm_preservation:
        transform = NormPreservingTransform(transform)

    self._template = (Intervention(
        # heuristic default: ~40% depth (the paper finds layer 13/32 optimal)
        layers=(self.layer_id,) if self.layer_id 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()

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

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

def export_state(self) -> dict:
    """The bound 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 vector plus the resolved layer.

The frozen args carry the exported steering_vector, the resolved layer_id, and the recipe's application parameters; normalize_vector is cleared since the exported vector is already in its applied form.

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

def frozen_form(self, state: dict) -> tuple[str, dict]:
    """A same-class frozen form: the bound vector plus the resolved layer.

    The frozen args carry the exported `steering_vector`, the resolved `layer_id`, and the
    recipe's application parameters; `normalize_vector` is cleared since the exported
    vector is already in its applied form.
    """
    return "state_control/caa", {
        "steering_vector": state["steering_vector"],
        "layer_id": self._layer_id,
        "multiplier": self.multiplier,
        "token_scope": self.token_scope,
        "last_k": self.last_k,
        "from_position": self.from_position,
        "normalize_vector": False,
        "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()

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)