Internals

steerability.algorithms.core.internals

Tools for reading and characterizing model internals (the model's activations and statistics) without steering them.

The root modules provide the shared signal-reading layer used by detection and by steering estimation:

• capture: hooked/forward capture of hidden states at module boundaries.
• encoding: text-to-tensor tokenization entry points.
• pooling: token, span, and position aggregation.
• render: ContrastivePairs to model-ready text.
• data: contrastive and labeled-example data containers.
• stats: ambient activation statistics and residual-norm characterization.
• fingerprint: digests identifying the model a fitted artifact was estimated on.

Layering rules:

1. The root is signal acquisition and characterization; subpackages are capabilities built from the readings. Root modules never import from subpackages, and each __init__.py re-exports its own level only.
2. Modules here may import the core foundations (base_args, types, core/utils) and steerability.utils; they never import the *_control category packages at module level.
3. New signals (attention patterns, per-head capture, MLP activations) extend root modules; new capabilities become sibling subpackages.

HiddenStateLocation = Literal['layer_output', 'layer_input'] module-attribute

ActivationStats dataclass

Per-layer mean and variance of a model's ambient activation distribution.

The statistics support per-coordinate centering and scaling of pooled activations (standardize), the whitening used by diagonal-LDA probe fitting. Covariance is diagonal, since the dominant effects in residual-stream geometry (a large shared component in every activation and rogue dimensions dominating inner products) are axis-aligned, so per-coordinate centering and scaling removes most of the effect, and a full covariance is hard to estimate reliably and expensive to invert at a few thousand samples.

Under pooling="tokens", every retained real token position is one sample of the ambient distribution, and exclude_first_n (default 1) drops the leading real positions per sequence, since the first real position (BOS or an opening delimiter) carries known massive-activation outliers that distort the mean disproportionately. "mean" and "last" pool per prompt first, yielding one sample per prompt, and ignore exclude_first_n.

count is the number of pooled samples accumulated, i.e., tokens under pooling="tokens" and prompts otherwise. estimate warns when count falls below its min_samples argument; a few hundred varied texts clears the default under token pooling, while "mean" and "last" need a few thousand prompts.

Attributes:

Name	Type	Description
model_type	str	HuggingFace model_type of the model the statistics were estimated on.
model_fingerprint	str	Identity digest of that model, recorded at estimate time.
location	str	Residual-stream boundary the statistics were measured at.
count	int	Number of pooled samples accumulated.
mean	dict[int, Tensor]	Per-layer mean, [H] float32.
var	dict[int, Tensor]	Per-layer variance, [H] float32, floored at var_floor.
var_floor	float	Lower bound applied to every variance coordinate.

Source code in steerability/algorithms/core/internals/stats.py

@dataclass
class ActivationStats:
    """Per-layer mean and variance of a model's ambient activation distribution.

    The statistics support per-coordinate centering and scaling of pooled activations
    (`standardize`), the whitening used by diagonal-LDA probe fitting. Covariance is diagonal,
    since the dominant effects in residual-stream geometry (a large shared component in every
    activation and rogue dimensions dominating inner products) are axis-aligned, so per-coordinate
    centering and scaling removes most of the effect, and a full covariance is hard to estimate
    reliably and expensive to invert at a few thousand samples.

    Under `pooling="tokens"`, every retained real token position is one sample of the ambient
    distribution, and `exclude_first_n` (default 1) drops the leading real positions per sequence,
    since the first real position (BOS or an opening delimiter) carries known massive-activation
    outliers that distort the mean disproportionately. `"mean"` and `"last"` pool per prompt first,
    yielding one sample per prompt, and ignore `exclude_first_n`.

    `count` is the number of pooled samples accumulated, i.e., tokens under `pooling="tokens"` and
    prompts otherwise. `estimate` warns when `count` falls below its `min_samples` argument; a few
    hundred varied texts clears the default under token pooling, while `"mean"` and `"last"` need a
    few thousand prompts.

    Attributes:
        model_type: HuggingFace `model_type` of the model the statistics were estimated on.
        model_fingerprint: Identity digest of that model, recorded at estimate time.
        location: Residual-stream boundary the statistics were measured at.
        count: Number of pooled samples accumulated.
        mean: Per-layer mean, `[H]` float32.
        var: Per-layer variance, `[H]` float32, floored at `var_floor`.
        var_floor: Lower bound applied to every variance coordinate.
    """

    model_type: str
    model_fingerprint: str
    location: str
    count: int
    mean: dict[int, torch.Tensor]
    var: dict[int, torch.Tensor]
    var_floor: float = 1e-6

    @classmethod
    def estimate(
        cls,
        model: PreTrainedModel,
        tokenizer: PreTrainedTokenizerBase,
        texts: Sequence[str],
        *,
        layer_ids: Sequence[int] | None = None,
        location: str = "layer_input",
        pooling: Literal["mean", "last", "tokens"] = "tokens",
        exclude_first_n: int = 1,
        batch_size: int = 8,
        min_samples: int = 5000,
        session=None,
    ) -> "ActivationStats":
        """Estimate per-layer activation statistics over `texts`.

        Texts are tokenized in chunks of `batch_size`, captured at `location`, and pooled per
        `pooling`. Sequences with no real tokens contribute no samples, and duplicate entries in
        `layer_ids` are collapsed. Accumulation is streaming (per-chunk Welford merge) in float32
        on CPU, so the full feature matrix is never materialized. Capture runs under `torch.no_grad()` with
        `use_cache=False`, wrapped in `auxiliary_pass(aligned=True)`, so co-resident state
        controls' condition scorers, gates, and position counters ignore the passes while
        `"all"`-scoped behavior transforms still apply (the statistics measure the stream as
        deployed).

        Args:
            model: Model to capture activations from.
            tokenizer: Tokenizer for encoding `texts`.
            texts: Texts whose activations define the ambient distribution.
            layer_ids: Layers to estimate (0-based). None selects all decoder layers.
            location: Residual-stream boundary to measure at.
            pooling: `"tokens"` treats every retained real token position as one sample; `"mean"`
                and `"last"` pool per prompt first.
            exclude_first_n: Leading real token positions dropped per sequence under
                `pooling="tokens"`; ignored for `"mean"` and `"last"`.
            batch_size: Batch size for the capture forward passes.
            min_samples: Pooled-sample count below which a low-sample warning is emitted.

        Returns:
            An `ActivationStats` with float32 CPU tensors and the model's fingerprint recorded.

        Raises:
            ValueError: If `texts` is empty, `pooling` is unsupported, `exclude_first_n` is
                negative, `batch_size < 1`, a requested layer id is out of range, or no pooled
                samples remain after exclusion.

        Warns:
            UserWarning: When the accumulated pooled-sample `count` is below `min_samples`.
        """
        if len(texts) == 0:
            raise ValueError("texts must contain at least one text.")
        if pooling not in ("mean", "last", "tokens"):
            raise ValueError(f"pooling must be 'mean', 'last', or 'tokens', got {pooling!r}.")
        if exclude_first_n < 0:
            raise ValueError(f"exclude_first_n must be >= 0, got {exclude_first_n}.")
        if batch_size < 1:
            raise ValueError("batch_size must be >= 1.")

        device = next(model.parameters()).device
        model_type = getattr(model.config, "model_type", "unknown")
        fingerprint = model_fingerprint(model)

        target_layers: list[int] | None = None
        count = 0
        running_mean: dict[int, torch.Tensor] = {}
        running_m2: dict[int, torch.Tensor] = {}

        for start in range(0, len(texts), batch_size):
            chunk = list(texts[start:start + batch_size])
            enc = tokenize_texts(tokenizer, chunk, device)
            with auxiliary_pass(aligned=True):
                hidden, chunk_mask = capture_hidden(
                    enc, model=model, session=session, batch_size=len(chunk), location=location
                )

            if target_layers is None:
                num_layers = len(hidden)
                if layer_ids is None:
                    target_layers = sorted(hidden.keys())
                else:
                    target_layers = list(dict.fromkeys(int(lid) for lid in layer_ids))
                    for lid in target_layers:
                        if not 0 <= lid < num_layers:
                            raise ValueError(f"layer id {lid} out of range [0, {num_layers}).")

            mask = (
                chunk_mask.to("cpu", torch.bool)
                if chunk_mask is not None
                else torch.ones(hidden[target_layers[0]].shape[:2], dtype=torch.bool)
            )

            if pooling == "tokens":
                # rank real positions per row and drop the first exclude_first_n of them
                rank = mask.long().cumsum(dim=1)
                keep = mask & (rank > exclude_first_n)  # [B, T]
                row_keep = None
            else:
                keep = None
                row_keep = mask.any(dim=1)  # rows with no real tokens contribute no sample

            chunk_count = 0
            for lid in target_layers:
                states = hidden[lid].to(torch.float32)  # [B, T, H]
                if pooling == "tokens":
                    features = states[keep]  # [K, H]
                else:
                    row_states = states[row_keep]
                    row_mask = mask[row_keep]
                    if pooling == "mean":
                        features = masked_mean(row_states, row_mask.to(row_states.dtype))
                    else:  # last
                        positions = get_last_token_positions(row_mask.long(), row_states.size(1), row_states.size(0))
                        features = select_at_positions(row_states, positions)

                chunk_count = features.size(0)
                if chunk_count == 0:
                    break

                batch_mean = features.mean(dim=0)
                batch_m2 = (features - batch_mean).pow(2).sum(dim=0)
                if lid not in running_mean:
                    running_mean[lid] = batch_mean
                    running_m2[lid] = batch_m2
                else:
                    delta = batch_mean - running_mean[lid]
                    total = count + chunk_count
                    running_mean[lid] = running_mean[lid] + delta * (chunk_count / total)
                    running_m2[lid] = running_m2[lid] + batch_m2 + delta.pow(2) * (count * chunk_count / total)

            count += chunk_count

        if count == 0:
            raise ValueError(
                "No pooled samples were accumulated; every sequence was shorter than "
                f"exclude_first_n={exclude_first_n} real tokens."
            )
        if count < min_samples:
            suggestion = (
                "supply more texts"
                if pooling == "tokens"
                else "supply more texts, or use pooling='tokens' to draw one sample per token"
            )
            warnings.warn(
                f"ActivationStats accumulated {count} pooled samples, below min_samples={min_samples}. "
                f"Estimates of per-coordinate variance may be unstable; {suggestion}.",
                UserWarning,
                stacklevel=2,
            )

        var_floor = cls.var_floor
        variances = {
            lid: (running_m2[lid] / (count - 1) if count > 1 else torch.zeros_like(running_m2[lid]))
            .clamp_min(var_floor)
            for lid in running_mean
        }
        return cls(
            model_type=model_type,
            model_fingerprint=fingerprint,
            location=location,
            count=count,
            mean=running_mean,
            var=variances,
            var_floor=var_floor,
        )

    def standardize(self, features: torch.Tensor, layer_id: int) -> torch.Tensor:
        """Center and scale features per coordinate with this artifact's layer statistics.

        Args:
            features: Tensor of shape `[..., H]` in raw activation coordinates.
            layer_id: Layer whose statistics to apply.

        Returns:
            `(features - mean) / sqrt(var)` as a float32 tensor of the same shape.

        Raises:
            KeyError: If `layer_id` has no recorded statistics.
        """
        if layer_id not in self.mean:
            raise KeyError(
                f"ActivationStats has no statistics for layer {layer_id}; "
                f"available layers: {sorted(self.mean)}."
            )
        features = features.to(torch.float32)
        return (features - self.mean[layer_id]) / self.var[layer_id].sqrt()

    def save(self, dir_path: str | Path) -> None:
        """Save the artifact to a directory (safetensors for tensors, JSON sidecar for metadata).

        Args:
            dir_path: Directory to write to; created if missing.
        """
        dir_path = Path(dir_path)
        dir_path.mkdir(parents=True, exist_ok=True)

        tensors = {}
        for lid in sorted(self.mean):
            tensors[f"mean.{lid}"] = self.mean[lid].to(torch.float32).contiguous()
            tensors[f"var.{lid}"] = self.var[lid].to(torch.float32).contiguous()
        save_file(tensors, str(dir_path / _TENSOR_FILE))

        metadata = {
            "format_version": STATS_FORMAT_VERSION,
            "model_type": self.model_type,
            "model_fingerprint": self.model_fingerprint,
            "location": self.location,
            "count": self.count,
            "var_floor": self.var_floor,
            "layer_ids": sorted(self.mean),
        }
        with open(dir_path / _META_FILE, "w") as f:
            json.dump(metadata, f, indent=2, sort_keys=True)
        logger.debug("Saved ActivationStats to %s", dir_path)

    @classmethod
    def load(cls, dir_path: str | Path) -> "ActivationStats":
        """Load an artifact saved by `save`.

        Args:
            dir_path: Directory containing the safetensors file and JSON sidecar.

        Returns:
            The loaded `ActivationStats`.

        Raises:
            ValueError: If the artifact-format version is unsupported.
        """
        dir_path = Path(dir_path)
        with open(dir_path / _META_FILE) as f:
            metadata = json.load(f)

        version = metadata.get("format_version")
        if version != STATS_FORMAT_VERSION:
            raise ValueError(
                f"Unsupported ActivationStats format version {version!r}; expected {STATS_FORMAT_VERSION}."
            )

        tensors = load_file(str(dir_path / _TENSOR_FILE))
        mean: dict[int, torch.Tensor] = {}
        var: dict[int, torch.Tensor] = {}
        for lid in metadata["layer_ids"]:
            mean[int(lid)] = tensors[f"mean.{lid}"]
            var[int(lid)] = tensors[f"var.{lid}"]

        logger.debug("Loaded ActivationStats from %s with layers %s", dir_path, sorted(mean))
        return cls(
            model_type=metadata["model_type"],
            model_fingerprint=metadata["model_fingerprint"],
            location=metadata["location"],
            count=int(metadata["count"]),
            mean=mean,
            var=var,
            var_floor=float(metadata["var_floor"]),
        )

    def fingerprint(self) -> str:
        """Digest of this artifact (not of the model it was estimated on).

        Returns:
            A 16-character lowercase hex digest over the metadata and tensor contents.
        """
        digest = hashlib.sha256()
        header = json.dumps(
            {
                "model_type": self.model_type,
                "model_fingerprint": self.model_fingerprint,
                "location": self.location,
                "count": self.count,
                "var_floor": self.var_floor,
            },
            sort_keys=True,
        )
        digest.update(header.encode("utf-8"))
        for lid in sorted(self.mean):
            digest.update(str(lid).encode("utf-8"))
            digest.update(self.mean[lid].to(torch.float32).contiguous().numpy().tobytes())
            digest.update(self.var[lid].to(torch.float32).contiguous().numpy().tobytes())
        return digest.hexdigest()[:16]

count instance-attribute

location instance-attribute

mean instance-attribute

model_fingerprint instance-attribute

model_type instance-attribute

var instance-attribute

var_floor = 1e-06 class-attribute instance-attribute

estimate(model, tokenizer, texts, *, layer_ids=None, location='layer_input', pooling='tokens', exclude_first_n=1, batch_size=8, min_samples=5000, session=None) classmethod

Estimate per-layer activation statistics over texts.

Texts are tokenized in chunks of batch_size, captured at location, and pooled per pooling. Sequences with no real tokens contribute no samples, and duplicate entries in layer_ids are collapsed. Accumulation is streaming (per-chunk Welford merge) in float32 on CPU, so the full feature matrix is never materialized. Capture runs under torch.no_grad() with use_cache=False, wrapped in auxiliary_pass(aligned=True), so co-resident state controls' condition scorers, gates, and position counters ignore the passes while "all"-scoped behavior transforms still apply (the statistics measure the stream as deployed).

Parameters:

Name	Type	Description	Default
model	PreTrainedModel	Model to capture activations from.	required
tokenizer	PreTrainedTokenizerBase	Tokenizer for encoding texts.	required
texts	Sequence[str]	Texts whose activations define the ambient distribution.	required
layer_ids	Sequence[int] | None	Layers to estimate (0-based). None selects all decoder layers.	None
location	str	Residual-stream boundary to measure at.	'layer_input'
pooling	Literal['mean', 'last', 'tokens']	"tokens" treats every retained real token position as one sample; "mean" and "last" pool per prompt first.	'tokens'
exclude_first_n	int	Leading real token positions dropped per sequence under pooling="tokens"; ignored for "mean" and "last".	1
batch_size	int	Batch size for the capture forward passes.	8
min_samples	int	Pooled-sample count below which a low-sample warning is emitted.	5000

Returns:

Type	Description
ActivationStats	An ActivationStats with float32 CPU tensors and the model's fingerprint recorded.

Raises:

Type	Description
ValueError	If texts is empty, pooling is unsupported, exclude_first_n is negative, batch_size < 1, a requested layer id is out of range, or no pooled samples remain after exclusion.

Warns:

Type	Description
UserWarning	When the accumulated pooled-sample count is below min_samples.

Source code in steerability/algorithms/core/internals/stats.py

@classmethod
def estimate(
    cls,
    model: PreTrainedModel,
    tokenizer: PreTrainedTokenizerBase,
    texts: Sequence[str],
    *,
    layer_ids: Sequence[int] | None = None,
    location: str = "layer_input",
    pooling: Literal["mean", "last", "tokens"] = "tokens",
    exclude_first_n: int = 1,
    batch_size: int = 8,
    min_samples: int = 5000,
    session=None,
) -> "ActivationStats":
    """Estimate per-layer activation statistics over `texts`.

    Texts are tokenized in chunks of `batch_size`, captured at `location`, and pooled per
    `pooling`. Sequences with no real tokens contribute no samples, and duplicate entries in
    `layer_ids` are collapsed. Accumulation is streaming (per-chunk Welford merge) in float32
    on CPU, so the full feature matrix is never materialized. Capture runs under `torch.no_grad()` with
    `use_cache=False`, wrapped in `auxiliary_pass(aligned=True)`, so co-resident state
    controls' condition scorers, gates, and position counters ignore the passes while
    `"all"`-scoped behavior transforms still apply (the statistics measure the stream as
    deployed).

    Args:
        model: Model to capture activations from.
        tokenizer: Tokenizer for encoding `texts`.
        texts: Texts whose activations define the ambient distribution.
        layer_ids: Layers to estimate (0-based). None selects all decoder layers.
        location: Residual-stream boundary to measure at.
        pooling: `"tokens"` treats every retained real token position as one sample; `"mean"`
            and `"last"` pool per prompt first.
        exclude_first_n: Leading real token positions dropped per sequence under
            `pooling="tokens"`; ignored for `"mean"` and `"last"`.
        batch_size: Batch size for the capture forward passes.
        min_samples: Pooled-sample count below which a low-sample warning is emitted.

    Returns:
        An `ActivationStats` with float32 CPU tensors and the model's fingerprint recorded.

    Raises:
        ValueError: If `texts` is empty, `pooling` is unsupported, `exclude_first_n` is
            negative, `batch_size < 1`, a requested layer id is out of range, or no pooled
            samples remain after exclusion.

    Warns:
        UserWarning: When the accumulated pooled-sample `count` is below `min_samples`.
    """
    if len(texts) == 0:
        raise ValueError("texts must contain at least one text.")
    if pooling not in ("mean", "last", "tokens"):
        raise ValueError(f"pooling must be 'mean', 'last', or 'tokens', got {pooling!r}.")
    if exclude_first_n < 0:
        raise ValueError(f"exclude_first_n must be >= 0, got {exclude_first_n}.")
    if batch_size < 1:
        raise ValueError("batch_size must be >= 1.")

    device = next(model.parameters()).device
    model_type = getattr(model.config, "model_type", "unknown")
    fingerprint = model_fingerprint(model)

    target_layers: list[int] | None = None
    count = 0
    running_mean: dict[int, torch.Tensor] = {}
    running_m2: dict[int, torch.Tensor] = {}

    for start in range(0, len(texts), batch_size):
        chunk = list(texts[start:start + batch_size])
        enc = tokenize_texts(tokenizer, chunk, device)
        with auxiliary_pass(aligned=True):
            hidden, chunk_mask = capture_hidden(
                enc, model=model, session=session, batch_size=len(chunk), location=location
            )

        if target_layers is None:
            num_layers = len(hidden)
            if layer_ids is None:
                target_layers = sorted(hidden.keys())
            else:
                target_layers = list(dict.fromkeys(int(lid) for lid in layer_ids))
                for lid in target_layers:
                    if not 0 <= lid < num_layers:
                        raise ValueError(f"layer id {lid} out of range [0, {num_layers}).")

        mask = (
            chunk_mask.to("cpu", torch.bool)
            if chunk_mask is not None
            else torch.ones(hidden[target_layers[0]].shape[:2], dtype=torch.bool)
        )

        if pooling == "tokens":
            # rank real positions per row and drop the first exclude_first_n of them
            rank = mask.long().cumsum(dim=1)
            keep = mask & (rank > exclude_first_n)  # [B, T]
            row_keep = None
        else:
            keep = None
            row_keep = mask.any(dim=1)  # rows with no real tokens contribute no sample

        chunk_count = 0
        for lid in target_layers:
            states = hidden[lid].to(torch.float32)  # [B, T, H]
            if pooling == "tokens":
                features = states[keep]  # [K, H]
            else:
                row_states = states[row_keep]
                row_mask = mask[row_keep]
                if pooling == "mean":
                    features = masked_mean(row_states, row_mask.to(row_states.dtype))
                else:  # last
                    positions = get_last_token_positions(row_mask.long(), row_states.size(1), row_states.size(0))
                    features = select_at_positions(row_states, positions)

            chunk_count = features.size(0)
            if chunk_count == 0:
                break

            batch_mean = features.mean(dim=0)
            batch_m2 = (features - batch_mean).pow(2).sum(dim=0)
            if lid not in running_mean:
                running_mean[lid] = batch_mean
                running_m2[lid] = batch_m2
            else:
                delta = batch_mean - running_mean[lid]
                total = count + chunk_count
                running_mean[lid] = running_mean[lid] + delta * (chunk_count / total)
                running_m2[lid] = running_m2[lid] + batch_m2 + delta.pow(2) * (count * chunk_count / total)

        count += chunk_count

    if count == 0:
        raise ValueError(
            "No pooled samples were accumulated; every sequence was shorter than "
            f"exclude_first_n={exclude_first_n} real tokens."
        )
    if count < min_samples:
        suggestion = (
            "supply more texts"
            if pooling == "tokens"
            else "supply more texts, or use pooling='tokens' to draw one sample per token"
        )
        warnings.warn(
            f"ActivationStats accumulated {count} pooled samples, below min_samples={min_samples}. "
            f"Estimates of per-coordinate variance may be unstable; {suggestion}.",
            UserWarning,
            stacklevel=2,
        )

    var_floor = cls.var_floor
    variances = {
        lid: (running_m2[lid] / (count - 1) if count > 1 else torch.zeros_like(running_m2[lid]))
        .clamp_min(var_floor)
        for lid in running_mean
    }
    return cls(
        model_type=model_type,
        model_fingerprint=fingerprint,
        location=location,
        count=count,
        mean=running_mean,
        var=variances,
        var_floor=var_floor,
    )

fingerprint()

Digest of this artifact (not of the model it was estimated on).

Returns:

Type	Description
str	A 16-character lowercase hex digest over the metadata and tensor contents.

Source code in steerability/algorithms/core/internals/stats.py

def fingerprint(self) -> str:
    """Digest of this artifact (not of the model it was estimated on).

    Returns:
        A 16-character lowercase hex digest over the metadata and tensor contents.
    """
    digest = hashlib.sha256()
    header = json.dumps(
        {
            "model_type": self.model_type,
            "model_fingerprint": self.model_fingerprint,
            "location": self.location,
            "count": self.count,
            "var_floor": self.var_floor,
        },
        sort_keys=True,
    )
    digest.update(header.encode("utf-8"))
    for lid in sorted(self.mean):
        digest.update(str(lid).encode("utf-8"))
        digest.update(self.mean[lid].to(torch.float32).contiguous().numpy().tobytes())
        digest.update(self.var[lid].to(torch.float32).contiguous().numpy().tobytes())
    return digest.hexdigest()[:16]

load(dir_path) classmethod

Load an artifact saved by save.

Parameters:

Name	Type	Description	Default
dir_path	str | Path	Directory containing the safetensors file and JSON sidecar.	required

Returns:

Type	Description
ActivationStats	The loaded ActivationStats.

Raises:

Type	Description
ValueError	If the artifact-format version is unsupported.

Source code in steerability/algorithms/core/internals/stats.py

@classmethod
def load(cls, dir_path: str | Path) -> "ActivationStats":
    """Load an artifact saved by `save`.

    Args:
        dir_path: Directory containing the safetensors file and JSON sidecar.

    Returns:
        The loaded `ActivationStats`.

    Raises:
        ValueError: If the artifact-format version is unsupported.
    """
    dir_path = Path(dir_path)
    with open(dir_path / _META_FILE) as f:
        metadata = json.load(f)

    version = metadata.get("format_version")
    if version != STATS_FORMAT_VERSION:
        raise ValueError(
            f"Unsupported ActivationStats format version {version!r}; expected {STATS_FORMAT_VERSION}."
        )

    tensors = load_file(str(dir_path / _TENSOR_FILE))
    mean: dict[int, torch.Tensor] = {}
    var: dict[int, torch.Tensor] = {}
    for lid in metadata["layer_ids"]:
        mean[int(lid)] = tensors[f"mean.{lid}"]
        var[int(lid)] = tensors[f"var.{lid}"]

    logger.debug("Loaded ActivationStats from %s with layers %s", dir_path, sorted(mean))
    return cls(
        model_type=metadata["model_type"],
        model_fingerprint=metadata["model_fingerprint"],
        location=metadata["location"],
        count=int(metadata["count"]),
        mean=mean,
        var=var,
        var_floor=float(metadata["var_floor"]),
    )

save(dir_path)

Save the artifact to a directory (safetensors for tensors, JSON sidecar for metadata).

Parameters:

Name	Type	Description	Default
dir_path	str | Path	Directory to write to; created if missing.	required

Source code in steerability/algorithms/core/internals/stats.py

def save(self, dir_path: str | Path) -> None:
    """Save the artifact to a directory (safetensors for tensors, JSON sidecar for metadata).

    Args:
        dir_path: Directory to write to; created if missing.
    """
    dir_path = Path(dir_path)
    dir_path.mkdir(parents=True, exist_ok=True)

    tensors = {}
    for lid in sorted(self.mean):
        tensors[f"mean.{lid}"] = self.mean[lid].to(torch.float32).contiguous()
        tensors[f"var.{lid}"] = self.var[lid].to(torch.float32).contiguous()
    save_file(tensors, str(dir_path / _TENSOR_FILE))

    metadata = {
        "format_version": STATS_FORMAT_VERSION,
        "model_type": self.model_type,
        "model_fingerprint": self.model_fingerprint,
        "location": self.location,
        "count": self.count,
        "var_floor": self.var_floor,
        "layer_ids": sorted(self.mean),
    }
    with open(dir_path / _META_FILE, "w") as f:
        json.dump(metadata, f, indent=2, sort_keys=True)
    logger.debug("Saved ActivationStats to %s", dir_path)

standardize(features, layer_id)

Center and scale features per coordinate with this artifact's layer statistics.

Parameters:

Name	Type	Description	Default
features	Tensor	Tensor of shape [..., H] in raw activation coordinates.	required
layer_id	int	Layer whose statistics to apply.	required

Returns:

Type	Description
Tensor	(features - mean) / sqrt(var) as a float32 tensor of the same shape.

Raises:

Type	Description
KeyError	If layer_id has no recorded statistics.

Source code in steerability/algorithms/core/internals/stats.py

def standardize(self, features: torch.Tensor, layer_id: int) -> torch.Tensor:
    """Center and scale features per coordinate with this artifact's layer statistics.

    Args:
        features: Tensor of shape `[..., H]` in raw activation coordinates.
        layer_id: Layer whose statistics to apply.

    Returns:
        `(features - mean) / sqrt(var)` as a float32 tensor of the same shape.

    Raises:
        KeyError: If `layer_id` has no recorded statistics.
    """
    if layer_id not in self.mean:
        raise KeyError(
            f"ActivationStats has no statistics for layer {layer_id}; "
            f"available layers: {sorted(self.mean)}."
        )
    features = features.to(torch.float32)
    return (features - self.mean[layer_id]) / self.var[layer_id].sqrt()

ContrastivePairs dataclass

Paired positive/negative text data for contrastive estimation.

The token sequence the model sees for each example is produced by render_for_model according to the consumer's prompt format, and the rendered string is tokenized with add_special_tokens=False for chat-templated text.

Attributes:

Name	Type	Description
positives	Sequence[str]	Texts exhibiting the target behavior. Treated as completions under prompt_format="chat_completion", and as standalone prompts under prompt_format="chat_prompt".
negatives	Sequence[str]	Texts not exhibiting the target behavior (see positives).
prompts	Sequence[str] | None	Optional shared prompts. Used as the user turn under prompt_format="chat_completion" and as the prefix under prompt_format="raw". Required when accumulate == "suffix-only". Ignored under prompt_format="chat_prompt".

Source code in steerability/algorithms/core/internals/data.py

@dataclass(frozen=True)
class ContrastivePairs:
    """Paired positive/negative text data for contrastive estimation.

    The token sequence the model sees for each example is produced by
    `render_for_model` according to the consumer's prompt format, and the
    rendered string is tokenized with `add_special_tokens=False` for
    chat-templated text.

    Attributes:
        positives: Texts exhibiting the target behavior. Treated as completions
            under `prompt_format="chat_completion"`, and as standalone prompts
            under `prompt_format="chat_prompt"`.
        negatives: Texts not exhibiting the target behavior (see `positives`).
        prompts: Optional shared prompts. Used as the user turn under
            `prompt_format="chat_completion"` and as the prefix under
            `prompt_format="raw"`. Required when `accumulate == "suffix-only"`.
            Ignored under `prompt_format="chat_prompt"`.
    """

    positives: Sequence[str]
    negatives: Sequence[str]
    prompts: Sequence[str] | None = None

    def __post_init__(self):
        if len(self.positives) == 0 or len(self.negatives) == 0:
            raise ValueError("positives and negatives must each have at least one entry.")
        if len(self.positives) != len(self.negatives):
            raise ValueError(
                f"positives ({len(self.positives)}) and negatives ({len(self.negatives)}) "
                f"must have equal length."
            )
        if self.prompts is not None and len(self.prompts) != len(self.positives):
            raise ValueError("prompts must have the same length as positives/negatives.")

negatives instance-attribute

positives instance-attribute

prompts = None class-attribute instance-attribute

HeadGeometry dataclass

Attention head geometry of one decoder layer.

Attributes:

Name	Type	Description
num_heads	int	Number of attention heads.
head_dim	int	Per-head dimension.

Source code in steerability/algorithms/core/internals/model_layout.py

@dataclass(frozen=True)
class HeadGeometry:
    """Attention head geometry of one decoder layer.

    Attributes:
        num_heads: Number of attention heads.
        head_dim: Per-head dimension.
    """

    num_heads: int
    head_dim: int

head_dim instance-attribute

num_heads instance-attribute

LabeledExamples dataclass

Independent positive/negative text data with binary labels.

The positive and negative lists need not be the same length. Applies to methods where positive and negative examples are independent and unpaired, and the estimator concatenates them. Optional group keys mark statements that must not straddle a train/validation split (e.g. the source question of each answer); estimators that do not split ignore them.

Attributes:

Name	Type	Description
positives	Sequence[str]	Texts exhibiting the target behavior (label=1).
negatives	Sequence[str]	Texts not exhibiting the target behavior (label=0).
positive_groups	Sequence[str | int] | None	Group key per positive, same length as positives, or None.
negative_groups	Sequence[str | int] | None	Group key per negative, same length as negatives, or None.

Source code in steerability/algorithms/core/internals/data.py

@dataclass(frozen=True)
class LabeledExamples:
    """Independent positive/negative text data with binary labels.

    The positive and negative lists need not be the same length. Applies to methods where
    positive and negative examples are independent and unpaired, and the estimator concatenates
    them. Optional group keys mark statements that must not straddle a train/validation split
    (e.g. the source question of each answer); estimators that do not split ignore them.

    Attributes:
        positives: Texts exhibiting the target behavior (label=1).
        negatives: Texts not exhibiting the target behavior (label=0).
        positive_groups: Group key per positive, same length as `positives`, or None.
        negative_groups: Group key per negative, same length as `negatives`, or None.
    """

    positives: Sequence[str]
    negatives: Sequence[str]
    positive_groups: Sequence[str | int] | None = None
    negative_groups: Sequence[str | int] | None = None

    def __post_init__(self):
        if len(self.positives) == 0 or len(self.negatives) == 0:
            raise ValueError("positives and negatives must each have at least one entry.")

        given = [self.positive_groups is not None, self.negative_groups is not None]
        if any(given) and not all(given):
            raise ValueError(
                "positive_groups and negative_groups must both be given or both omitted."
            )

        for name, groups, texts in (
            ("positive_groups", self.positive_groups, self.positives),
            ("negative_groups", self.negative_groups, self.negatives),
        ):
            if groups is None:
                continue
            if len(groups) != len(texts):
                raise ValueError(
                    f"{name} has length {len(groups)} but its text list has length {len(texts)}."
                )
            for key in groups:
                if not isinstance(key, (str, int)) or isinstance(key, bool):
                    raise ValueError(
                        f"{name} keys must be str or int; got {type(key).__name__}."
                    )

    @property
    def groups(self) -> bool:
        """Whether the data carries group keys (positive/negative groups both present)."""
        return self.positive_groups is not None

groups property

Whether the data carries group keys (positive/negative groups both present).

negative_groups = None class-attribute instance-attribute

negatives instance-attribute

positive_groups = None class-attribute instance-attribute

positives instance-attribute

ModelLayout dataclass

Architecture-specific module paths for one model family.

Attributes:

Name	Type	Description
family	str	Family label ("gemma_style", "llama_style", or "gpt2_style").
layer_prefix	str	Dotted prefix of the decoder layer list (e.g. "model.layers", "model.language_model.layers", "transformer.h", prefixed with "base_model.model." per PEFT wrapper).
num_layers	int	Number of decoder layers.
attn_suffix	str	Suffix of the per-layer attention module (".self_attn" or ".attn").
oproj_suffix	str	Suffix of the attention output projection (".self_attn.o_proj" or ".attn.c_proj").
norm_attrs	tuple[str, ...]	The per-layer normalization sub-module attribute names whose input is the residual stream (("input_layernorm", "pre_feedforward_layernorm") on Gemma, ("input_layernorm", "post_attention_layernorm") on Llama, ("ln_1", "ln_2") on GPT-2).
attention_layer_ids	tuple[int, ...] | None	The decoder layers carrying the attention module at attn_suffix, ascending. None means every layer (homogeneous stacks, and layouts built by hand); a hybrid stack lists only its attention layers, and an empty tuple declares that no layer carries attention.

Source code in steerability/algorithms/core/internals/model_layout.py

@dataclass(frozen=True)
class ModelLayout:
    """Architecture-specific module paths for one model family.

    Attributes:
        family: Family label (`"gemma_style"`, `"llama_style"`, or `"gpt2_style"`).
        layer_prefix: Dotted prefix of the decoder layer list (e.g. `"model.layers"`,
            `"model.language_model.layers"`, `"transformer.h"`, prefixed with
            `"base_model.model."` per PEFT wrapper).
        num_layers: Number of decoder layers.
        attn_suffix: Suffix of the per-layer attention module (`".self_attn"` or `".attn"`).
        oproj_suffix: Suffix of the attention output projection (`".self_attn.o_proj"` or
            `".attn.c_proj"`).
        norm_attrs: The per-layer normalization sub-module attribute names whose input is the
            residual stream (`("input_layernorm", "pre_feedforward_layernorm")` on Gemma,
            `("input_layernorm", "post_attention_layernorm")` on Llama, `("ln_1", "ln_2")` on
            GPT-2).
        attention_layer_ids: The decoder layers carrying the attention module at `attn_suffix`,
            ascending. None means every layer (homogeneous stacks, and layouts built by hand);
            a hybrid stack lists only its attention layers, and an empty tuple declares that
            no layer carries attention.
    """

    family: str
    layer_prefix: str
    num_layers: int
    attn_suffix: str
    oproj_suffix: str
    norm_attrs: tuple[str, ...]
    attention_layer_ids: tuple[int, ...] | None = None

    @property
    def layer_names(self) -> list[str]:
        """Dotted paths of every decoder layer, `[f"{layer_prefix}.{i}"]`."""
        return [f"{self.layer_prefix}.{i}" for i in range(self.num_layers)]

    @property
    def oproj_names(self) -> list[str]:
        """Dotted paths of every layer's attention output projection, index-aligned with
        `layer_names`. On a hybrid stack, entries outside `attention_layers` name modules the
        model does not have."""
        return [name + self.oproj_suffix for name in self.layer_names]

    @property
    def attn_names(self) -> list[str]:
        """Dotted paths of every layer's attention module, index-aligned with `layer_names`. On a
        hybrid stack, entries outside `attention_layers` name modules the model does not have."""
        return [name + self.attn_suffix for name in self.layer_names]

    @property
    def attention_layers(self) -> tuple[int, ...]:
        """The decoder layers carrying an attention module, ascending; every layer when
        `attention_layer_ids` is None."""
        if self.attention_layer_ids is None:
            return tuple(range(self.num_layers))
        return self.attention_layer_ids

    @property
    def is_hybrid(self) -> bool:
        """True when some decoder layer carries no attention module."""
        return len(self.attention_layers) < self.num_layers

    def has_attention(self, layer_id: int) -> bool:
        """True when decoder layer `layer_id` carries the attention module at `attn_suffix`."""
        return layer_id in self.attention_layers

attention_layer_ids = None class-attribute instance-attribute

attention_layers property

The decoder layers carrying an attention module, ascending; every layer when attention_layer_ids is None.

attn_names property

Dotted paths of every layer's attention module, index-aligned with layer_names. On a hybrid stack, entries outside attention_layers name modules the model does not have.

attn_suffix instance-attribute

family instance-attribute

is_hybrid property

True when some decoder layer carries no attention module.

layer_names property

Dotted paths of every decoder layer, [f"{layer_prefix}.{i}"].

layer_prefix instance-attribute

norm_attrs instance-attribute

num_layers instance-attribute

oproj_names property

Dotted paths of every layer's attention output projection, index-aligned with layer_names. On a hybrid stack, entries outside attention_layers name modules the model does not have.

oproj_suffix instance-attribute

has_attention(layer_id)

True when decoder layer layer_id carries the attention module at attn_suffix.

Source code in steerability/algorithms/core/internals/model_layout.py

def has_attention(self, layer_id: int) -> bool:
    """True when decoder layer `layer_id` carries the attention module at `attn_suffix`."""
    return layer_id in self.attention_layers

RenderedContrastive dataclass

Rendered text for a ContrastivePairs, plus tokenization policy.

Attributes:

Name	Type	Description
pos_texts	list[str]	Rendered positive examples.
neg_texts	list[str]	Rendered negative examples.
prompt_texts	list[str] | None	User-turn-only renders for suffix-only span computation, or None when there is no separate prompt/completion split.
add_special_tokens	bool	Flag callers must pass to the tokenizer for all of the above (False for chat modes, True for raw).
effective_mode	PromptFormat	The mode actually applied after raw fallbacks.

Source code in steerability/algorithms/core/internals/render.py

@dataclass(frozen=True)
class RenderedContrastive:
    """Rendered text for a `ContrastivePairs`, plus tokenization policy.

    Attributes:
        pos_texts: Rendered positive examples.
        neg_texts: Rendered negative examples.
        prompt_texts: User-turn-only renders for suffix-only span computation, or
            None when there is no separate prompt/completion split.
        add_special_tokens: Flag callers must pass to the tokenizer for all of
            the above (False for chat modes, True for raw).
        effective_mode: The mode actually applied after raw fallbacks.
    """

    pos_texts: list[str]
    neg_texts: list[str]
    prompt_texts: list[str] | None
    add_special_tokens: bool
    effective_mode: PromptFormat

add_special_tokens instance-attribute

effective_mode instance-attribute

neg_texts instance-attribute

pos_texts instance-attribute

prompt_texts instance-attribute

StatsSpec dataclass

Deferred recipe for ActivationStats; holds every argument except the model.

Attributes:

Name	Type	Description
texts	Sequence[str]	Texts whose activations define the ambient distribution.
layer_ids	Sequence[int] | None	Layers to estimate (0-based). None selects all decoder layers.
location	str	Residual-stream boundary the statistics are measured at.
pooling	Literal['mean', 'last', 'tokens']	Sample construction. "tokens" treats every retained real token position as one sample; "mean" and "last" pool per prompt first.
exclude_first_n	int	Number of leading real token positions dropped per sequence under pooling="tokens"; ignored for "mean" and "last".
batch_size	int	Batch size for the capture forward passes.

Source code in steerability/algorithms/core/internals/stats.py

@dataclass
class StatsSpec:
    """Deferred recipe for `ActivationStats`; holds every argument except the model.

    Attributes:
        texts: Texts whose activations define the ambient distribution.
        layer_ids: Layers to estimate (0-based). None selects all decoder layers.
        location: Residual-stream boundary the statistics are measured at.
        pooling: Sample construction. `"tokens"` treats every retained real token position as one
            sample; `"mean"` and `"last"` pool per prompt first.
        exclude_first_n: Number of leading real token positions dropped per sequence under
            `pooling="tokens"`; ignored for `"mean"` and `"last"`.
        batch_size: Batch size for the capture forward passes.
    """

    texts: Sequence[str]
    layer_ids: Sequence[int] | None = None  # None selects all decoder layers
    location: str = "layer_input"
    pooling: Literal["mean", "last", "tokens"] = "tokens"
    exclude_first_n: int = 1
    batch_size: int = 8

    def estimate(
        self, model: PreTrainedModel | None, tokenizer: PreTrainedTokenizerBase, session=None,
    ) -> "ActivationStats":
        """Estimate `ActivationStats` on `model` with this recipe's settings."""
        return ActivationStats.estimate(
            model,
            tokenizer,
            self.texts,
            session=session,
            layer_ids=self.layer_ids,
            location=self.location,
            pooling=self.pooling,
            exclude_first_n=self.exclude_first_n,
            batch_size=self.batch_size,
        )

batch_size = 8 class-attribute instance-attribute

exclude_first_n = 1 class-attribute instance-attribute

layer_ids = None class-attribute instance-attribute

location = 'layer_input' class-attribute instance-attribute

pooling = 'tokens' class-attribute instance-attribute

texts instance-attribute

estimate(model, tokenizer, session=None)

Estimate ActivationStats on model with this recipe's settings.

Source code in steerability/algorithms/core/internals/stats.py

def estimate(
    self, model: PreTrainedModel | None, tokenizer: PreTrainedTokenizerBase, session=None,
) -> "ActivationStats":
    """Estimate `ActivationStats` on `model` with this recipe's settings."""
    return ActivationStats.estimate(
        model,
        tokenizer,
        self.texts,
        session=session,
        layer_ids=self.layer_ids,
        location=self.location,
        pooling=self.pooling,
        exclude_first_n=self.exclude_first_n,
        batch_size=self.batch_size,
    )

aggregate_condition_hidden(hidden, mode, attention_mask=None)

Aggregate [B, T, H] hidden states to [B, H] using non-pad tokens only.

Parameters:

Name	Type	Description	Default
hidden	Tensor	Shape [B, T, H].	required
mode	Literal['mean', 'last']	"mean" pools over all real tokens; "last" selects the last real token per row.	required
attention_mask	Tensor | None	Shape [B, T] (1 for real tokens, 0 for pads), or None. When None, "mean" averages all positions and "last" uses the final position.	None

Returns:

Type	Description
Tensor	Aggregated tensor of shape [B, H].

Raises:

Type	Description
ValueError	If a row has no real tokens, or the mode is unsupported.

Source code in steerability/algorithms/core/internals/pooling.py

@torch.no_grad()
def aggregate_condition_hidden(
    hidden: torch.Tensor,
    mode: Literal["mean", "last"],
    attention_mask: torch.Tensor | None = None,
) -> torch.Tensor:
    """Aggregate `[B, T, H]` hidden states to `[B, H]` using non-pad tokens only.

    Args:
        hidden: Shape `[B, T, H]`.
        mode: "mean" pools over all real tokens; "last" selects the last real token per row.
        attention_mask: Shape `[B, T]` (1 for real tokens, 0 for pads), or None. When None, "mean"
            averages all positions and "last" uses the final position.

    Returns:
        Aggregated tensor of shape `[B, H]`.

    Raises:
        ValueError: If a row has no real tokens, or the mode is unsupported.
    """
    if mode == "mean":
        return masked_mean(hidden, attention_mask)

    if mode != "last":
        raise ValueError(f"Unsupported condition comparison mode: {mode!r}.")

    batch_size = hidden.size(0)
    if attention_mask is None:
        return hidden[:, -1, :]

    mask = attention_mask.to(hidden.device).bool()
    if not mask.any(dim=1).all():
        raise ValueError("aggregate_condition_hidden received a row with no real tokens.")

    sequence_length = mask.size(1)
    positions = torch.arange(sequence_length, device=hidden.device).unsqueeze(0)
    last_positions = positions.masked_fill(~mask, -1).max(dim=1).values  # [B]
    return hidden[torch.arange(batch_size, device=hidden.device), last_positions]

as_contrastive_pairs(x)

Normalize input to ContrastivePairs.

Accepts:

- An existing `ContrastivePairs` instance (returned as-is).
- A dict with keys `"positives"`, `"negatives"`, and optionally `"prompts"`.

Parameters:

Name	Type	Description	Default
x	ContrastivePairs | dict	Input to normalize.	required

Returns:

Type	Description
ContrastivePairs	A ContrastivePairs instance.

Raises:

Type	Description
TypeError	If input is neither ContrastivePairs nor a suitable dict.

Source code in steerability/algorithms/core/internals/data.py

def as_contrastive_pairs(x: ContrastivePairs | dict) -> ContrastivePairs:
    """Normalize input to `ContrastivePairs`.

    Accepts:

        - An existing `ContrastivePairs` instance (returned as-is).
        - A dict with keys `"positives"`, `"negatives"`, and optionally `"prompts"`.

    Args:
        x: Input to normalize.

    Returns:
        A `ContrastivePairs` instance.

    Raises:
        TypeError: If input is neither `ContrastivePairs` nor a suitable dict.
    """
    if isinstance(x, ContrastivePairs):
        return x
    if isinstance(x, dict):
        return ContrastivePairs(**x)
    raise TypeError("Expected ContrastivePairs or dict with positives/negatives[/prompts].")

as_labeled_examples(x)

Normalize input to LabeledExamples.

Accepts:

- An existing `LabeledExamples` instance (returned as-is).
- A `ContrastivePairs` instance (converted; pairing is dropped).
- A dict with keys `"positives"` and `"negatives"`, and optionally `"positive_groups"`
  and `"negative_groups"`.

Parameters:

Name	Type	Description	Default
x	LabeledExamples | ContrastivePairs | dict	Input to normalize.	required

Returns:

Type	Description
LabeledExamples	A LabeledExamples instance.

Raises:

Type	Description
TypeError	If input is not LabeledExamples, ContrastivePairs, or a suitable dict.

Source code in steerability/algorithms/core/internals/data.py

def as_labeled_examples(x: LabeledExamples | ContrastivePairs | dict) -> LabeledExamples:
    """Normalize input to `LabeledExamples`.

    Accepts:

        - An existing `LabeledExamples` instance (returned as-is).
        - A `ContrastivePairs` instance (converted; pairing is dropped).
        - A dict with keys `"positives"` and `"negatives"`, and optionally `"positive_groups"`
          and `"negative_groups"`.

    Args:
        x: Input to normalize.

    Returns:
        A `LabeledExamples` instance.

    Raises:
        TypeError: If input is not `LabeledExamples`, `ContrastivePairs`, or a suitable dict.
    """
    if isinstance(x, LabeledExamples):
        return x
    if isinstance(x, ContrastivePairs):
        return LabeledExamples(positives=x.positives, negatives=x.negatives)
    if isinstance(x, dict):
        return LabeledExamples(**x)
    raise TypeError("Expected LabeledExamples, ContrastivePairs, or dict with positives/negatives.")

get_last_token_positions(attention_mask, seq_len, num_samples)

Find the last non-pad token position for each sample.

Parameters:

Name	Type	Description	Default
attention_mask	Tensor | None	Shape [N, T] or None.	required
seq_len	int	Sequence length T.	required
num_samples	int	Number of samples N.	required

Returns:

Type	Description
LongTensor	Tensor of shape [N] with last token positions.

Source code in steerability/algorithms/core/internals/pooling.py

def get_last_token_positions(
    attention_mask: torch.Tensor | None,
    seq_len: int,
    num_samples: int,
) -> torch.LongTensor:
    """Find the last non-pad token position for each sample.

    Args:
        attention_mask: Shape `[N, T]` or None.
        seq_len: Sequence length `T`.
        num_samples: Number of samples `N`.

    Returns:
        Tensor of shape `[N]` with last token positions.
    """
    if attention_mask is None:
        # no padding, last token is at seq_len - 1
        return torch.full((num_samples,), seq_len - 1, dtype=torch.long)

    # for each sample, find the last position where attention_mask == 1
    # this handles both left-padded and right-padded sequences
    positions = torch.arange(seq_len, device=attention_mask.device).unsqueeze(0).expand(num_samples, -1)
    # mask out padded positions with -1
    masked_positions = torch.where(attention_mask == 1, positions, torch.tensor(-1, device=attention_mask.device))
    return masked_positions.max(dim=1).values

head_geometry(model, layout, layer_id)

Per-layer attention head geometry, read from the module tree.

head_dim comes from the attention module's head_dim attribute, else the text config; num_heads is the output projection's input width divided by head_dim. GPT-2's Conv1D stores its weight as [in, out], so the width is read from weight.shape[0] when in_features is absent.

Parameters:

Name	Type	Description	Default
model	PreTrainedModel	The live model (or PEFT wrapper) carrying the attention modules.	required
layout	ModelLayout	The resolved ModelLayout naming the per-layer attention and output projections.	required
layer_id	int	The decoder layer index.	required

Returns:

Type	Description
HeadGeometry	The layer's HeadGeometry.

Raises:

Type	Description
ValueError	If layer_id carries no attention module (a non-attention layer of a hybrid stack), or the projection width is not a multiple of head_dim.

Source code in steerability/algorithms/core/internals/model_layout.py

def head_geometry(model: PreTrainedModel, layout: ModelLayout, layer_id: int) -> HeadGeometry:
    """Per-layer attention head geometry, read from the module tree.

    `head_dim` comes from the attention module's `head_dim` attribute, else the text config;
    `num_heads` is the output projection's input width divided by `head_dim`. GPT-2's `Conv1D`
    stores its weight as `[in, out]`, so the width is read from `weight.shape[0]` when
    `in_features` is absent.

    Args:
        model: The live model (or PEFT wrapper) carrying the attention modules.
        layout: The resolved `ModelLayout` naming the per-layer attention and output projections.
        layer_id: The decoder layer index.

    Returns:
        The layer's `HeadGeometry`.

    Raises:
        ValueError: If `layer_id` carries no attention module (a non-attention layer of a hybrid
            stack), or the projection width is not a multiple of `head_dim`.
    """
    if not layout.has_attention(layer_id):
        raise ValueError(
            f"Decoder layer {layer_id} carries no attention module ({layout.attn_suffix!r}); "
            f"attention layers of this model are {list(layout.attention_layers)}."
        )
    attn = model.get_submodule(layout.attn_names[layer_id])
    head_dim = getattr(attn, "head_dim", None)
    if head_dim is None:
        text_cfg = text_config(model)
        head_dim = getattr(text_cfg, "head_dim", None)
        if head_dim is None:
            head_dim = text_cfg.hidden_size // text_cfg.num_attention_heads

    oproj = model.get_submodule(layout.oproj_names[layer_id])
    width = getattr(oproj, "in_features", None)
    if width is None:
        width = oproj.weight.shape[0]  # gpt-2 Conv1D stores weight as [in, out]

    if width % head_dim != 0:
        raise ValueError(
            f"Output-projection input width {width} at layer {layer_id} is not a multiple of "
            f"head_dim {head_dim}; cannot infer the head count."
        )
    return HeadGeometry(num_heads=width // head_dim, head_dim=head_dim)

layerwise_tokenwise_hidden(model, enc, batch_size=8, on_batch=None, *, location='layer_output')

Extract per-layer hidden states for all tokens.

outputs.hidden_states is a tuple of num_layers + 1 tensors: index 0 is the embedding output (the input to layer 0) and index i is the output of layer i - 1, except the final entry, which transformers returns with the model's final norm already applied.

• location="layer_output": key l maps to the raw output boundary of layer l (hidden_states[l + 1] for l < num_layers - 1; the final layer is re-captured pre-norm by a forward hook on the last decoder layer, matching hook runtimes and engine capture).
• location="layer_input": key l maps to the input of layer l, i.e. the output of layer l - 1 (hidden_states[l]), the boundary a layer pre-hook observes.

Parameters:

Name	Type	Description	Default
model	PreTrainedModel	The model to extract from.	required
enc	dict[str, Tensor]	Tokenized input with input_ids and attention_mask.	required
batch_size	int	Batch size for forward passes.	8
on_batch	Callable[[], None] | None	Optional callable invoked after each batch finishes. Used by callers to surface progress to the UI.	None
location	HiddenStateLocation	Which residual-stream boundary each layer key maps to.	'layer_output'

Returns:

Type	Description
dict[int, Tensor]	Dict mapping layer_id (0 .. num_layers - 1) to tensor of shape [N, T, H]. Rows
dict[int, Tensor]	outside attention_mask are zeroed when a mask is provided.

Raises:

Type	Description
ValueError	If location is unsupported or the number of mapped states does not equal the model's layer count.

Source code in steerability/algorithms/core/internals/capture.py

@torch.no_grad()
def layerwise_tokenwise_hidden(
    model: PreTrainedModel,
    enc: dict[str, torch.Tensor],
    batch_size: int = 8,
    on_batch: Callable[[], None] | None = None,
    *,
    location: HiddenStateLocation = "layer_output",
) -> dict[int, torch.Tensor]:
    """Extract per-layer hidden states for all tokens.

    `outputs.hidden_states` is a tuple of `num_layers + 1` tensors: index 0 is the embedding output
    (the input to layer 0) and index `i` is the output of layer `i - 1`, except the final entry,
    which transformers returns with the model's final norm already applied.

    - `location="layer_output"`: key `l` maps to the raw output boundary of layer `l`
        (`hidden_states[l + 1]` for `l < num_layers - 1`; the final layer is re-captured pre-norm
        by a forward hook on the last decoder layer, matching hook runtimes and engine capture).
    - `location="layer_input"`: key `l` maps to the input of layer `l`, i.e. the output of layer
        `l - 1` (`hidden_states[l]`), the boundary a layer pre-hook observes.

    Args:
        model: The model to extract from.
        enc: Tokenized input with `input_ids` and `attention_mask`.
        batch_size: Batch size for forward passes.
        on_batch: Optional callable invoked after each batch finishes. Used by callers to surface
            progress to the UI.
        location: Which residual-stream boundary each layer key maps to.

    Returns:
        Dict mapping `layer_id` (`0 .. num_layers - 1`) to tensor of shape `[N, T, H]`. Rows
        outside `attention_mask` are zeroed when a mask is provided.

    Raises:
        ValueError: If `location` is unsupported or the number of mapped states does not equal the
            model's layer count.
    """
    if location not in ("layer_output", "layer_input"):
        raise ValueError(f"Unsupported hidden-state location: {location!r}.")

    input_ids = enc["input_ids"]
    attention_mask = enc.get("attention_mask")
    N = input_ids.size(0)
    final_layer_module = None
    if location == "layer_output":
        # the last `hidden_states` entry is post-final-norm; recover the final layer's raw output
        # boundary with a forward hook on the last decoder layer
        final_layer_module = model.get_submodule(resolve_model_layout(model).layer_names[-1])

    # collect states per layer
    all_hidden: dict[int, list[torch.Tensor]] = {}
    num_layers: int | None = None

    for start in range(0, N, batch_size):
        end = min(start + batch_size, N)
        batch_ids = input_ids[start:end]
        batch_mask = attention_mask[start:end] if attention_mask is not None else None

        final_boundary: list[torch.Tensor] = []
        handle = None
        if final_layer_module is not None:
            def _grab_final(module, args, output):
                raw = output[0] if isinstance(output, tuple) else output
                final_boundary.append(raw)

            handle = final_layer_module.register_forward_hook(_grab_final)
        try:
            outputs = model(
                input_ids=batch_ids,
                attention_mask=batch_mask,
                output_hidden_states=True,
                return_dict=True,
                use_cache=False,
                # transformers v5 threads `cache_position` into decoder-layer kwargs only when the
                # caller passes it, and aligned auxiliary passes are positioned by that kwarg (see
                # `TransformHookRuntime._position_offset`), so it is passed explicitly here.
                cache_position=torch.arange(batch_ids.size(1), device=batch_ids.device),
            )
        finally:
            if handle is not None:
                handle.remove()

        num_layers = len(outputs.hidden_states) - 1
        layer_states = list(outputs.hidden_states[1:]) if location == "layer_output" else list(outputs.hidden_states[:-1])
        if final_layer_module is not None:
            if len(final_boundary) != 1:
                raise RuntimeError(
                    f"Expected exactly one final-layer forward per batch, observed {len(final_boundary)}."
                )
            layer_states[-1] = final_boundary[0]
        batch_row_mask = batch_mask.unsqueeze(-1) if batch_mask is not None else None
        for layer_idx, hs in enumerate(layer_states):
            if batch_row_mask is not None:
                # zero rows outside the attention mask; their values are computed but unused
                hs = hs * batch_row_mask.to(device=hs.device)
            all_hidden.setdefault(layer_idx, []).append(hs.cpu())

        if on_batch is not None:
            on_batch()

    result = {layer_idx: torch.cat(tensors, dim=0) for layer_idx, tensors in all_hidden.items()}

    if num_layers is not None and len(result) != num_layers:
        raise ValueError(
            f"Expected {num_layers} mapped hidden states for location={location!r}, got {len(result)}."
        )

    return result

masked_mean(hidden, attention_mask)

Mean-pool [N, T, H] hidden states over real (non-pad) positions.

Falls back to an unmasked mean over all positions when no attention mask is available.

Parameters:

Name	Type	Description	Default
hidden	Tensor	Shape [N, T, H].	required
attention_mask	Tensor | None	Shape [N, T] (1 for real tokens, 0 for pads), or None.	required

Returns:

Type	Description
Tensor	Pooled tensor of shape [N, H].

Source code in steerability/algorithms/core/internals/pooling.py

def masked_mean(hidden: torch.Tensor, attention_mask: torch.Tensor | None) -> torch.Tensor:
    """Mean-pool `[N, T, H]` hidden states over real (non-pad) positions.

    Falls back to an unmasked mean over all positions when no attention mask is available.

    Args:
        hidden: Shape `[N, T, H]`.
        attention_mask: Shape `[N, T]` (1 for real tokens, 0 for pads), or None.

    Returns:
        Pooled tensor of shape `[N, H]`.
    """
    if attention_mask is None:
        return hidden.mean(dim=1)
    m = attention_mask.to(hidden.dtype).unsqueeze(-1)  # [N, T, 1]
    return (hidden * m).sum(dim=1) / m.sum(dim=1).clamp_min(1e-8)

measure_residual_norms(model, tokenizer, layer_ids, prompts, *, location='layer_output', stat='median', prompt_format='chat_prompt', batch_size=8)

Typical per-token residual-stream norm at each requested layer boundary.

The statistic is aggregated over the real (non-pad) tokens of prompts at the residual-stream boundary named by location. Residual norms grow with depth, so the returned per-layer norms let a caller rescale a direction to a fixed fraction of the norm at each layer (see SteeringVector.scaled_to_norms).

Prompts are rendered via render_for_model(tokenizer, prompt=p, mode=prompt_format). The default prompt_format="chat_prompt" renders with add_generation_prompt=True, matching what inference produces for a prompt. The rendered string is tokenized with add_special_tokens=False whenever a chat template was applied, matching how the estimators consume contrastive text.

Parameters:

Name	Type	Description	Default
model	PreTrainedModel	Model to extract hidden states from.	required
tokenizer	PreTrainedTokenizerBase	Tokenizer for rendering and encoding the prompts.	required
layer_ids	Sequence[int] | None	Layers to measure (0-based). None measures every layer (no extra cost, since all layers are extracted regardless).	required
prompts	Sequence[str]	Calibration prompts. Use fit-distribution, inference-style prompts rather than a held-out evaluation set.	required
location	HiddenStateLocation	Residual-stream boundary each layer key maps to. "layer_output" (default) measures the output of each layer; "layer_input" measures the input of each layer, the boundary a layer pre-hook observes.	'layer_output'
stat	Literal['median', 'mean']	Aggregation over the pooled real-token norms, "median" (default) or "mean".	'median'
prompt_format	PromptFormat	How each prompt is rendered into model-ready text (via render_for_model).	'chat_prompt'
batch_size	int	Batch size for the extraction forward passes.	8

Returns:

Type	Description
dict[int, float]	Mapping from layer id to the aggregated per-token residual norm, as a Python float.

Raises:

Type	Description
ValueError	If prompts is empty, stat is unsupported, or any requested layer id is out of range.

Source code in steerability/algorithms/core/internals/stats.py

@torch.no_grad()
def measure_residual_norms(
    model: PreTrainedModel,
    tokenizer: PreTrainedTokenizerBase,
    layer_ids: Sequence[int] | None,
    prompts: Sequence[str],
    *,
    location: HiddenStateLocation = "layer_output",
    stat: Literal["median", "mean"] = "median",
    prompt_format: PromptFormat = "chat_prompt",
    batch_size: int = 8,
) -> dict[int, float]:
    """Typical per-token residual-stream norm at each requested layer boundary.

    The statistic is aggregated over the real (non-pad) tokens of `prompts` at the residual-stream
    boundary named by `location`. Residual norms grow with depth, so the returned per-layer norms
    let a caller rescale a direction to a fixed fraction of the norm at each layer (see
    `SteeringVector.scaled_to_norms`).

    Prompts are rendered via `render_for_model(tokenizer, prompt=p, mode=prompt_format)`. The default
    `prompt_format="chat_prompt"` renders with `add_generation_prompt=True`, matching what inference
    produces for a prompt. The rendered string is tokenized with `add_special_tokens=False` whenever
    a chat template was applied, matching how the estimators consume contrastive text.

    Args:
        model: Model to extract hidden states from.
        tokenizer: Tokenizer for rendering and encoding the prompts.
        layer_ids: Layers to measure (0-based). `None` measures every layer (no extra cost, since
            all layers are extracted regardless).
        prompts: Calibration prompts. Use fit-distribution, inference-style prompts rather than a
            held-out evaluation set.
        location: Residual-stream boundary each layer key maps to. `"layer_output"` (default)
            measures the output of each layer; `"layer_input"` measures the input of each layer, the
            boundary a layer pre-hook observes.
        stat: Aggregation over the pooled real-token norms, `"median"` (default) or `"mean"`.
        prompt_format: How each prompt is rendered into model-ready text (via `render_for_model`).
        batch_size: Batch size for the extraction forward passes.

    Returns:
        Mapping from layer id to the aggregated per-token residual norm, as a Python `float`.

    Raises:
        ValueError: If `prompts` is empty, `stat` is unsupported, or any requested layer id is out of
            range.
    """
    if len(prompts) == 0:
        raise ValueError("prompts must contain at least one prompt.")
    if stat not in ("median", "mean"):
        raise ValueError(f"stat must be 'median' or 'mean', got {stat!r}.")

    device = next(model.parameters()).device

    # render then tokenize with add_special_tokens=False when a template was applied
    rendered = [render_for_model(tokenizer, prompt=p, mode=prompt_format) for p in prompts]
    template_applied = has_chat_template(tokenizer) and prompt_format != "raw"
    enc = tokenize_texts(tokenizer, rendered, device, add_special_tokens=not template_applied)

    hidden = layerwise_tokenwise_hidden(model, enc, batch_size=batch_size, location=location)
    num_layers = len(hidden)

    if layer_ids is None:
        target_layers = sorted(hidden.keys())
    else:
        target_layers = [int(lid) for lid in layer_ids]
        for lid in target_layers:
            if not 0 <= lid < num_layers:
                raise ValueError(f"layer id {lid} out of range [0, {num_layers}).")

    # real-token mask
    attention_mask = enc.get("attention_mask")
    if attention_mask is not None:
        keep = attention_mask.to("cpu", torch.bool)  # [N, T]
    else:
        keep = None

    norms: dict[int, float] = {}
    for lid in target_layers:
        per_token = hidden[lid].to(torch.float32).norm(dim=-1)  # [N, T]
        values = per_token[keep] if keep is not None else per_token.reshape(-1)
        agg = values.median() if stat == "median" else values.mean()
        norms[lid] = float(agg)

    return norms

model_fingerprint(model)

Compute a deterministic digest identifying a model, from its configuration, dtype, and a sample of its weights.

The digest is a sha256 over (a) model.config.to_json_string(), (b) str(model.dtype), and (c) for up to 8 evenly spaced named parameters, the parameter name plus its first 64 elements flattened and cast to float32 bytes, truncated to 16 hex characters.

Properties:

• The same checkpoint at the same dtype yields equal digests in any process or device placement, since parameter values are deterministic under from_pretrained.
• Any weight edit that touches the sampled slices (SFT, DPO, merging, LoRA merge) changes the digest; so does attaching unmerged adapters, since parameter names change.
• dtype is part of the identity, since a calibrated operating point is a numeric artifact and the same checkpoint loaded at a different dtype is a different runtime model for detection purposes.

Parameters:

Name	Type	Description	Default
model	PreTrainedModel	The model to fingerprint.	required

Returns:

Type	Description
str	A 16-character lowercase hex digest.

Source code in steerability/algorithms/core/internals/fingerprint.py

def model_fingerprint(model: PreTrainedModel) -> str:
    """Compute a deterministic digest identifying a model, from its configuration, dtype, and a sample of its weights.

    The digest is a sha256 over (a) `model.config.to_json_string()`, (b) `str(model.dtype)`, and
    (c) for up to 8 evenly spaced named parameters, the parameter name plus its first 64 elements
    flattened and cast to float32 bytes, truncated to 16 hex characters.

    Properties:

    - The same checkpoint at the same dtype yields equal digests in any process or device
        placement, since parameter values are deterministic under `from_pretrained`.
    - Any weight edit that touches the sampled slices (SFT, DPO, merging, LoRA merge) changes the
        digest; so does attaching unmerged adapters, since parameter names change.
    - dtype is part of the identity, since a calibrated operating point is a numeric artifact and
        the same checkpoint loaded at a different dtype is a different runtime model for detection
        purposes.

    Args:
        model: The model to fingerprint.

    Returns:
        A 16-character lowercase hex digest.
    """
    digest = hashlib.sha256()
    digest.update(model.config.to_json_string().encode("utf-8"))
    digest.update(str(model.dtype).encode("utf-8"))

    named_parameters = list(model.named_parameters())
    n = len(named_parameters)
    if n > 0:
        k = min(8, n)
        indices = [0] if k == 1 else [round(i * (n - 1) / (k - 1)) for i in range(k)]
        for idx in dict.fromkeys(indices):
            name, param = named_parameters[idx]
            digest.update(name.encode("utf-8"))
            sample = param.detach().reshape(-1)[:64].to(device="cpu", dtype=torch.float32).contiguous()
            digest.update(sample.numpy().tobytes())

    return digest.hexdigest()[:16]

pool_over_spans(hidden, spans)

Mean-pool hidden states over specified spans.

A degenerate span (start >= end) falls back to the sample's last token position.

Parameters:

Name	Type	Description	Default
hidden	Tensor	Shape [N, T, H].	required
spans	list[tuple[int, int]]	List of (start, end) tuples.	required

Returns:

Type	Description
Tensor	Pooled tensor of shape [N, H].

Source code in steerability/algorithms/core/internals/pooling.py

def pool_over_spans(
    hidden: torch.Tensor,
    spans: list[tuple[int, int]],
) -> torch.Tensor:
    """Mean-pool hidden states over specified spans.

    A degenerate span (`start >= end`) falls back to the sample's last token position.

    Args:
        hidden: Shape `[N, T, H]`.
        spans: List of `(start, end)` tuples.

    Returns:
        Pooled tensor of shape `[N, H]`.
    """
    N, T, H = hidden.shape
    pooled = []
    for i, (start, end) in enumerate(spans):
        if start >= end:
            # fallback: use last token
            pooled.append(hidden[i, -1, :])
        else:
            pooled.append(hidden[i, start:end, :].mean(dim=0))
    return torch.stack(pooled, dim=0)

register_layout_detector(detector, *, prepend=True)

Register a callable mapping a model to a ModelLayout (or None to decline).

Registered detectors run before the built-in root and convention probing, in registration order (prepend=True puts the new detector first). Registration is process-global.

Parameters:

Name	Type	Description	Default
detector	Callable[[PreTrainedModel], ModelLayout | None]	A callable that returns a ModelLayout for models it recognizes, or None to fall through to the next detector and then the built-in resolution.	required
prepend	bool	When True (the default), the detector runs before previously registered ones.	True

Source code in steerability/algorithms/core/internals/model_layout.py

def register_layout_detector(
    detector: Callable[[PreTrainedModel], ModelLayout | None], *, prepend: bool = True
) -> None:
    """Register a callable mapping a model to a `ModelLayout` (or None to decline).

    Registered detectors run before the built-in root and convention probing, in registration
    order (`prepend=True` puts the new detector first). Registration is process-global.

    Args:
        detector: A callable that returns a `ModelLayout` for models it recognizes, or None to
            fall through to the next detector and then the built-in resolution.
        prepend: When True (the default), the detector runs before previously registered ones.
    """
    if prepend:
        _DETECTORS.insert(0, detector)
    else:
        _DETECTORS.append(detector)

render_contrastive(tokenizer, data, mode)

Render both sides of a contrastive dataset under mode.

Accepts ContrastivePairs (paired, with optional shared prompts) or LabeledExamples (unpaired classes, rendered independently). Resolves the effective mode (with raw fallbacks and warnings), renders positives/negatives, renders the prompt-only strings used for suffix-only span computation, and reports the add_special_tokens flag the tokenizer must use for all of the above.

Parameters:

Name	Type	Description	Default
tokenizer	PreTrainedTokenizerBase	Tokenizer whose chat template defines the rendering.	required
data	ContrastivePairs | LabeledExamples	ContrastivePairs (with positives, negatives, and optional prompts) or LabeledExamples.	required
mode	PromptFormat	Requested rendering policy.	required

Returns:

Type	Description
RenderedContrastive	A RenderedContrastive with rendered texts and tokenization policy.

Raises:

Type	Description
ValueError	If mode is "chat_completion" and data is LabeledExamples (that format's shared prompts align per example, which requires paired data).

Source code in steerability/algorithms/core/internals/render.py

def render_contrastive(
    tokenizer: PreTrainedTokenizerBase,
    data: ContrastivePairs | LabeledExamples,
    mode: PromptFormat,
) -> RenderedContrastive:
    """Render both sides of a contrastive dataset under `mode`.

    Accepts `ContrastivePairs` (paired, with optional shared prompts) or
    `LabeledExamples` (unpaired classes, rendered independently). Resolves the
    effective mode (with raw fallbacks and warnings), renders
    positives/negatives, renders the prompt-only strings used for suffix-only
    span computation, and reports the `add_special_tokens` flag the tokenizer
    must use for all of the above.

    Args:
        tokenizer: Tokenizer whose chat template defines the rendering.
        data: `ContrastivePairs` (with `positives`, `negatives`, and optional
            `prompts`) or `LabeledExamples`.
        mode: Requested rendering policy.

    Returns:
        A `RenderedContrastive` with rendered texts and tokenization policy.

    Raises:
        ValueError: If `mode` is `"chat_completion"` and `data` is
            `LabeledExamples` (that format's shared prompts align per example,
            which requires paired data).
    """
    if mode == "chat_completion" and isinstance(data, LabeledExamples):
        raise ValueError(
            "prompt_format='chat_completion' requires paired data with shared prompts "
            "(ContrastivePairs); LabeledExamples classes are unpaired."
        )
    prompts = getattr(data, "prompts", None)

    effective: PromptFormat = mode

    if mode in ("chat_completion", "chat_prompt") and not has_chat_template(tokenizer):
        logger.warning("render_contrastive: no chat_template; falling back to raw.")
        effective = "raw"

    if mode == "chat_completion" and prompts is None:
        logger.warning(
            "prompt_format='chat_completion' requires `prompts` (positives/negatives "
            "are treated as completions); none provided. Falling back to raw."
        )
        effective = "raw"

    add_special = effective == "raw"

    if effective == "raw":
        if prompts is not None:
            pos = [p + c for p, c in zip(prompts, data.positives)]
            neg = [p + c for p, c in zip(prompts, data.negatives)]
            prompt_texts = list(prompts)
        else:
            pos = list(data.positives)
            neg = list(data.negatives)
            prompt_texts = None
    elif effective == "chat_completion":
        pos = [
            render_for_model(tokenizer, prompt=p, completion=c, mode="chat_completion")
            for p, c in zip(prompts, data.positives)
        ]
        neg = [
            render_for_model(tokenizer, prompt=p, completion=c, mode="chat_completion")
            for p, c in zip(prompts, data.negatives)
        ]
        prompt_texts = [render_for_model(tokenizer, prompt=p, mode="chat_prompt") for p in prompts]
    else:  # chat_prompt mode, where each positive/negative is a standalone prompt
        pos = [render_for_model(tokenizer, prompt=t, mode="chat_prompt") for t in data.positives]
        neg = [render_for_model(tokenizer, prompt=t, mode="chat_prompt") for t in data.negatives]
        prompt_texts = None  # no separate prompt/completion split

    return RenderedContrastive(pos, neg, prompt_texts, add_special, effective)

resolve_model_layout(model)

Resolve the ModelLayout for a HuggingFace causal LM.

Consults registered detectors first, then locates the decoder stack at one of the built-in roots and selects a naming convention whose norm markers exist on the first decoder layer and whose attention module exists on at least one layer. A hybrid stack that interleaves attention layers with another token mixer resolves to its attention layers' family, with those layers recorded in attention_layer_ids. PEFT wrappers are peeled before resolution and their prefix is carried on the resolved layer_prefix.

Parameters:

Name	Type	Description	Default
model	PreTrainedModel	A HuggingFace causal LM, a multimodal wrapper whose text decoder sits at a known root, or a PEFT-wrapped model.	required

Returns:

Type	Description
ModelLayout	The matching ModelLayout.

Raises:

Type	Description
ValueError	If no detector and no built-in root/convention matches the model.

Source code in steerability/algorithms/core/internals/model_layout.py

def resolve_model_layout(model: PreTrainedModel) -> ModelLayout:
    """Resolve the `ModelLayout` for a HuggingFace causal LM.

    Consults registered detectors first, then locates the decoder stack at one of the built-in
    roots and selects a naming convention whose norm markers exist on the first decoder layer and
    whose attention module exists on at least one layer. A hybrid stack that interleaves attention
    layers with another token mixer resolves to its attention layers' family, with those layers
    recorded in `attention_layer_ids`. PEFT wrappers are peeled before resolution and their prefix
    is carried on the resolved `layer_prefix`.

    Args:
        model: A HuggingFace causal LM, a multimodal wrapper whose text decoder sits at a known
            root, or a PEFT-wrapped model.

    Returns:
        The matching `ModelLayout`.

    Raises:
        ValueError: If no detector and no built-in root/convention matches the model.
    """
    for detector in _DETECTORS:
        layout = detector(model)
        if layout is not None:
            return layout
    layout = _resolve_builtin(model)
    if layout is not None:
        return layout
    raise ValueError(
        f"Cannot determine model layout for {type(model).__name__}. Probed decoder-stack roots "
        f"{_DECODER_ROOTS} for a stack matching one of the families "
        f"{tuple(convention.family for convention in _CONVENTIONS)} (the family's norm "
        "sub-modules on the first layer and its attention module on at least one layer); "
        "register a detector with register_layout_detector() for other architectures."
    )

select_at_positions(hidden, positions)

Select hidden states at specified positions for each sample.

Parameters:

Name	Type	Description	Default
hidden	Tensor	Shape [N, T, H].	required
positions	LongTensor	Shape [N] with position indices.	required

Returns:

Type	Description
Tensor	Tensor of shape [N, H].

Source code in steerability/algorithms/core/internals/pooling.py

def select_at_positions(
    hidden: torch.Tensor,
    positions: torch.LongTensor,
) -> torch.Tensor:
    """Select hidden states at specified positions for each sample.

    Args:
        hidden: Shape `[N, T, H]`.
        positions: Shape `[N]` with position indices.

    Returns:
        Tensor of shape `[N, H]`.
    """
    N, _, H = hidden.shape
    # gather at the specified positions
    idx = positions.view(N, 1, 1).expand(N, 1, H)
    return hidden.gather(dim=1, index=idx).squeeze(1)

select_spans(enc, prompt_enc, accumulate)

Determine token spans to pool over for each sample.

Span bounds are derived from the attention mask, so pads on either padding side are excluded. Under accumulate="suffix-only", the per-sample prompt length (the number of real tokens in prompt_enc) is skipped at the start of each span.

Parameters:

Name	Type	Description	Default
enc	dict[str, Tensor]	Tokenized full sequences (prompts + completions).	required
prompt_enc	dict[str, Tensor] | None	Tokenized prompts only (if accumulate == "suffix-only").	required
accumulate	str	"all", "suffix-only", or "last_token".	required

Returns:

Type	Description
list[tuple[int, int]]	List of (start, end) tuples, one per sample.

Raises:

Type	Description
ValueError	If accumulate is unsupported.

Source code in steerability/algorithms/core/internals/pooling.py

def select_spans(
    enc: dict[str, torch.Tensor],
    prompt_enc: dict[str, torch.Tensor] | None,
    accumulate: str,
) -> list[tuple[int, int]]:
    """Determine token spans to pool over for each sample.

    Span bounds are derived from the attention mask, so pads on either padding side are excluded.
    Under `accumulate="suffix-only"`, the per-sample prompt length (the number of real tokens in
    `prompt_enc`) is skipped at the start of each span.

    Args:
        enc: Tokenized full sequences (prompts + completions).
        prompt_enc: Tokenized prompts only (if accumulate == "suffix-only").
        accumulate: "all", "suffix-only", or "last_token".

    Returns:
        List of (start, end) tuples, one per sample.

    Raises:
        ValueError: If `accumulate` is unsupported.
    """
    if accumulate not in ("all", "suffix-only", "last_token"):
        raise ValueError(
            f"select_spans does not support accumulate='{accumulate}'. "
            "Expected one of: 'all', 'suffix-only', 'last_token'."
        )

    input_ids = enc["input_ids"]
    attention_mask = enc.get("attention_mask")
    N, T = input_ids.shape

    spans = []
    for i in range(N):
        # number of prompt tokens to skip for suffix-only pooling (a count, not an absolute index)
        if accumulate == "suffix-only" and prompt_enc is not None:
            prompt_len = (
                int(prompt_enc["attention_mask"][i].sum().item())
                if "attention_mask" in prompt_enc
                else prompt_enc["input_ids"].size(1)
            )
        else:
            prompt_len = 0

        # derive both bounds from the mask so the span excludes pads on either padding side
        if attention_mask is not None:
            non_pad = (attention_mask[i] == 1).nonzero(as_tuple=True)[0]
            if len(non_pad) > 0:
                first = int(non_pad[0].item())
                last = int(non_pad[-1].item())
            else:
                first, last = 0, T - 1
            if accumulate == "last_token":
                # one-token span at the final non-pad position (derived from the mask, so it works
                # regardless of padding side)
                start, end = last, last + 1
            else:
                start = first + prompt_len
                end = last + 1
        else:
            if accumulate == "last_token":
                start, end = T - 1, T
            else:
                start = prompt_len
                end = T

        spans.append((start, end))

    return spans

text_config(model_or_config)

The config carrying the text decoder's structural facts.

Composite multimodal configs (text plus vision or audio sub-configs) return their text sub-config; plain configs return themselves. Numeric structural facts (hidden_size, num_attention_heads, head_dim, num_hidden_layers, intermediate_size) are read from the returned config; identity facts (model_type, the fingerprint JSON, _attn_implementation) stay on the composite model.config.

Parameters:

Name	Type	Description	Default
model_or_config	PreTrainedModel | PretrainedConfig	A model (its config is used) or a config.	required

Returns:

Type	Description
PretrainedConfig	The text config.

Source code in steerability/algorithms/core/internals/model_layout.py

def text_config(model_or_config: PreTrainedModel | PretrainedConfig) -> PretrainedConfig:
    """The config carrying the text decoder's structural facts.

    Composite multimodal configs (text plus vision or audio sub-configs) return their text
    sub-config; plain configs return themselves. Numeric structural facts (`hidden_size`,
    `num_attention_heads`, `head_dim`, `num_hidden_layers`, `intermediate_size`) are read from the
    returned config; identity facts (`model_type`, the fingerprint JSON, `_attn_implementation`)
    stay on the composite `model.config`.

    Args:
        model_or_config: A model (its `config` is used) or a config.

    Returns:
        The text config.
    """
    config = getattr(model_or_config, "config", model_or_config)
    return config.get_text_config()

tokenize_pairs(tokenizer, pos_texts, neg_texts, device, *, add_special_tokens=True)

Tokenize positive/negative pairs together to ensure consistent padding.

Interleaves pairs before tokenization so each (pos, neg) pair shares the same padding length, keeping shared prefixes token-aligned.

Parameters:

Name	Type	Description	Default
tokenizer	PreTrainedTokenizerBase	Tokenizer to use.	required
pos_texts	Sequence[str]	List of positive text strings.	required
neg_texts	Sequence[str]	List of negative text strings (same length as pos_texts).	required
device	device | str	Target device.	required
add_special_tokens	bool	Whether to add special tokens (e.g. BOS). Pass False for chat-templated text that already contains them.	True

Returns:

Type	Description
tuple[dict[str, Tensor], dict[str, Tensor]]	Tuple of (enc_pos, enc_neg) dictionaries with input_ids and attention_mask.

Source code in steerability/algorithms/core/internals/encoding.py

def tokenize_pairs(
    tokenizer: PreTrainedTokenizerBase,
    pos_texts: Sequence[str],
    neg_texts: Sequence[str],
    device: torch.device | str,
    *,
    add_special_tokens: bool = True,
) -> tuple[dict[str, torch.Tensor], dict[str, torch.Tensor]]:
    """Tokenize positive/negative pairs together to ensure consistent padding.

    Interleaves pairs before tokenization so each (pos, neg) pair shares the same
    padding length, keeping shared prefixes token-aligned.

    Args:
        tokenizer: Tokenizer to use.
        pos_texts: List of positive text strings.
        neg_texts: List of negative text strings (same length as `pos_texts`).
        device: Target device.
        add_special_tokens: Whether to add special tokens (e.g. BOS). Pass False
            for chat-templated text that already contains them.

    Returns:
        Tuple of `(enc_pos, enc_neg)` dictionaries with `input_ids` and `attention_mask`.
    """
    # interleave: [pos0, neg0, pos1, neg1, ...]
    interleaved = []
    for pos, neg in zip(pos_texts, neg_texts):
        interleaved.append(pos)
        interleaved.append(neg)

    enc = tokenizer(
        interleaved,
        return_tensors="pt",
        padding=True,
        truncation=True,
        add_special_tokens=add_special_tokens,
    )
    enc = {k: v.to(device) for k, v in enc.items()}

    # split the interleaved batch back apart: even indices are positive, odd indices are negative
    enc_pos = {k: v[0::2] for k, v in enc.items()}
    enc_neg = {k: v[1::2] for k, v in enc.items()}

    return enc_pos, enc_neg

tokenize_texts(tokenizer, texts, device=None, *, add_special_tokens=True, max_length=None)

Tokenize a flat list of texts independently.

Each text is tokenized on its own, without interleaving. Use this when positive and negative examples are independent and do not need co-padding for token alignment.

Parameters:

Name	Type	Description	Default
tokenizer	PreTrainedTokenizerBase	Tokenizer to use.	required
texts	Sequence[str]	List of text strings.	required
device	device | str | None	Target device for the returned tensors. When None, tensors are returned on the device the tokenizer produces them on (CPU).	None
add_special_tokens	bool	Whether to add special tokens (e.g. BOS). Pass False for chat-templated text that already contains them.	True
max_length	int | None	Truncation bound. When None, truncation falls back to the tokenizer's model maximum length.	None

Returns:

Type	Description
dict[str, Tensor]	Dictionary with input_ids and attention_mask tensors.

Source code in steerability/algorithms/core/internals/encoding.py

def tokenize_texts(
    tokenizer: PreTrainedTokenizerBase,
    texts: Sequence[str],
    device: torch.device | str | None = None,
    *,
    add_special_tokens: bool = True,
    max_length: int | None = None,
) -> dict[str, torch.Tensor]:
    """Tokenize a flat list of texts independently.

    Each text is tokenized on its own, without interleaving. Use this when positive and negative
    examples are independent and do not need co-padding for token alignment.

    Args:
        tokenizer: Tokenizer to use.
        texts: List of text strings.
        device: Target device for the returned tensors. When None, tensors are returned on the
            device the tokenizer produces them on (CPU).
        add_special_tokens: Whether to add special tokens (e.g. BOS). Pass False
            for chat-templated text that already contains them.
        max_length: Truncation bound. When None, truncation falls back to the tokenizer's model
            maximum length.

    Returns:
        Dictionary with `input_ids` and `attention_mask` tensors.
    """
    enc = tokenizer(
        list(texts),
        return_tensors="pt",
        padding=True,
        truncation=True,
        max_length=max_length,
        add_special_tokens=add_special_tokens,
    )
    if device is None:
        return dict(enc)
    return {k: v.to(device) for k, v in enc.items()}

capture

Hooked/forward capture of hidden states at module boundaries.

HiddenStateLocation = Literal['layer_output', 'layer_input'] module-attribute

capture_hidden(enc, *, model=None, session=None, batch_size=8, on_batch=None, location='layer_output')

Extract per-layer hidden states for enc using a live model or a capture session.

With a live model (given directly, or reachable through an in-process session), the extraction runs layerwise_tokenwise_hidden with the caller's batch size and progress callback, preserving the in-process layout, i.e., the returned mask is enc's own. With a remote capture-capable session, each row of enc becomes a token-id prompt (padding positions dropped) served by session.capture over every decoder layer; the returned tensors are right-padded to the batch's longest prompt and the returned mask describes that layout, so pooling must use the returned mask rather than enc's.

Parameters:

Name	Type	Description	Default
enc	dict[str, Tensor]	Tokenized input with input_ids and optionally attention_mask.	required
model	PreTrainedModel | None	A live model, used when no session provides one.	None
session	SteeringSession | None	A SteeringSession; in-process sessions expose their live model, remote sessions serve capture.	None
batch_size	int	Forward batch size on the in-process path.	8
on_batch	Callable[[], None] | None	Per-batch progress callback on the in-process path.	None
location	HiddenStateLocation	Which residual-stream boundary each layer key maps to.	'layer_output'

Returns:

Type	Description
tuple[dict[int, Tensor], Tensor | None]	A (hidden, attention_mask) pair with hidden[l] of shape [N, T, H] on CPU.

Raises:

Type	Description
ValueError	If neither a live model nor a session is available.

Source code in steerability/algorithms/core/internals/capture.py

def capture_hidden(
    enc: dict[str, torch.Tensor],
    *,
    model: PreTrainedModel | None = None,
    session: SteeringSession | None = None,
    batch_size: int = 8,
    on_batch: Callable[[], None] | None = None,
    location: HiddenStateLocation = "layer_output",
) -> tuple[dict[int, torch.Tensor], torch.Tensor | None]:
    """Extract per-layer hidden states for `enc` using a live model or a capture session.

    With a live model (given directly, or reachable through an in-process session), the
    extraction runs `layerwise_tokenwise_hidden` with the caller's batch size and progress
    callback, preserving the in-process layout, i.e., the returned mask is `enc`'s own. With a
    remote capture-capable session, each row of `enc` becomes a token-id prompt (padding
    positions dropped) served by `session.capture` over every decoder layer; the returned
    tensors are right-padded to the batch's longest prompt and the returned mask describes
    that layout, so pooling must use the returned mask rather than `enc`'s.

    Args:
        enc: Tokenized input with `input_ids` and optionally `attention_mask`.
        model: A live model, used when no session provides one.
        session: A `SteeringSession`; in-process sessions expose their live model, remote
            sessions serve `capture`.
        batch_size: Forward batch size on the in-process path.
        on_batch: Per-batch progress callback on the in-process path.
        location: Which residual-stream boundary each layer key maps to.

    Returns:
        A `(hidden, attention_mask)` pair with `hidden[l]` of shape `[N, T, H]` on CPU.

    Raises:
        ValueError: If neither a live model nor a session is available.
    """
    live_model = model
    if live_model is None and session is not None:
        try:
            live_model = session.model
        except (AttributeError, RuntimeError):
            live_model = None
    if live_model is not None:
        hidden = layerwise_tokenwise_hidden(
            live_model, enc, batch_size=batch_size, on_batch=on_batch, location=location
        )
        attention_mask = enc.get("attention_mask")
        return hidden, attention_mask.cpu() if attention_mask is not None else None
    if session is None or not callable(getattr(session, "capture", None)):
        raise ValueError("Hidden-state extraction requires a live model or a capture-capable session.")

    from steerability.algorithms.core.execution.payloads import PreparedPrompt

    input_ids = enc["input_ids"]
    attention_mask = enc.get("attention_mask")
    prompts = [
        PreparedPrompt.from_token_ids(
            input_ids[index:index + 1],
            attention_mask[index:index + 1] if attention_mask is not None else None,
        )
        for index in range(input_ids.size(0))
    ]
    result = session.capture(
        prompts, layers=list(range(session.layout.num_layers)), mode="all_tokens", location=location
    )
    if on_batch is not None:
        on_batch()
    return dict(result.hidden), result.attention_mask

layerwise_tokenwise_hidden(model, enc, batch_size=8, on_batch=None, *, location='layer_output')

Extract per-layer hidden states for all tokens.

outputs.hidden_states is a tuple of num_layers + 1 tensors: index 0 is the embedding output (the input to layer 0) and index i is the output of layer i - 1, except the final entry, which transformers returns with the model's final norm already applied.

• location="layer_output": key l maps to the raw output boundary of layer l (hidden_states[l + 1] for l < num_layers - 1; the final layer is re-captured pre-norm by a forward hook on the last decoder layer, matching hook runtimes and engine capture).
• location="layer_input": key l maps to the input of layer l, i.e. the output of layer l - 1 (hidden_states[l]), the boundary a layer pre-hook observes.

Parameters:

Name	Type	Description	Default
model	PreTrainedModel	The model to extract from.	required
enc	dict[str, Tensor]	Tokenized input with input_ids and attention_mask.	required
batch_size	int	Batch size for forward passes.	8
on_batch	Callable[[], None] | None	Optional callable invoked after each batch finishes. Used by callers to surface progress to the UI.	None
location	HiddenStateLocation	Which residual-stream boundary each layer key maps to.	'layer_output'

Returns:

Type	Description
dict[int, Tensor]	Dict mapping layer_id (0 .. num_layers - 1) to tensor of shape [N, T, H]. Rows
dict[int, Tensor]	outside attention_mask are zeroed when a mask is provided.

Raises:

Type	Description
ValueError	If location is unsupported or the number of mapped states does not equal the model's layer count.

Source code in steerability/algorithms/core/internals/capture.py

@torch.no_grad()
def layerwise_tokenwise_hidden(
    model: PreTrainedModel,
    enc: dict[str, torch.Tensor],
    batch_size: int = 8,
    on_batch: Callable[[], None] | None = None,
    *,
    location: HiddenStateLocation = "layer_output",
) -> dict[int, torch.Tensor]:
    """Extract per-layer hidden states for all tokens.

    `outputs.hidden_states` is a tuple of `num_layers + 1` tensors: index 0 is the embedding output
    (the input to layer 0) and index `i` is the output of layer `i - 1`, except the final entry,
    which transformers returns with the model's final norm already applied.

    - `location="layer_output"`: key `l` maps to the raw output boundary of layer `l`
        (`hidden_states[l + 1]` for `l < num_layers - 1`; the final layer is re-captured pre-norm
        by a forward hook on the last decoder layer, matching hook runtimes and engine capture).
    - `location="layer_input"`: key `l` maps to the input of layer `l`, i.e. the output of layer
        `l - 1` (`hidden_states[l]`), the boundary a layer pre-hook observes.

    Args:
        model: The model to extract from.
        enc: Tokenized input with `input_ids` and `attention_mask`.
        batch_size: Batch size for forward passes.
        on_batch: Optional callable invoked after each batch finishes. Used by callers to surface
            progress to the UI.
        location: Which residual-stream boundary each layer key maps to.

    Returns:
        Dict mapping `layer_id` (`0 .. num_layers - 1`) to tensor of shape `[N, T, H]`. Rows
        outside `attention_mask` are zeroed when a mask is provided.

    Raises:
        ValueError: If `location` is unsupported or the number of mapped states does not equal the
            model's layer count.
    """
    if location not in ("layer_output", "layer_input"):
        raise ValueError(f"Unsupported hidden-state location: {location!r}.")

    input_ids = enc["input_ids"]
    attention_mask = enc.get("attention_mask")
    N = input_ids.size(0)
    final_layer_module = None
    if location == "layer_output":
        # the last `hidden_states` entry is post-final-norm; recover the final layer's raw output
        # boundary with a forward hook on the last decoder layer
        final_layer_module = model.get_submodule(resolve_model_layout(model).layer_names[-1])

    # collect states per layer
    all_hidden: dict[int, list[torch.Tensor]] = {}
    num_layers: int | None = None

    for start in range(0, N, batch_size):
        end = min(start + batch_size, N)
        batch_ids = input_ids[start:end]
        batch_mask = attention_mask[start:end] if attention_mask is not None else None

        final_boundary: list[torch.Tensor] = []
        handle = None
        if final_layer_module is not None:
            def _grab_final(module, args, output):
                raw = output[0] if isinstance(output, tuple) else output
                final_boundary.append(raw)

            handle = final_layer_module.register_forward_hook(_grab_final)
        try:
            outputs = model(
                input_ids=batch_ids,
                attention_mask=batch_mask,
                output_hidden_states=True,
                return_dict=True,
                use_cache=False,
                # transformers v5 threads `cache_position` into decoder-layer kwargs only when the
                # caller passes it, and aligned auxiliary passes are positioned by that kwarg (see
                # `TransformHookRuntime._position_offset`), so it is passed explicitly here.
                cache_position=torch.arange(batch_ids.size(1), device=batch_ids.device),
            )
        finally:
            if handle is not None:
                handle.remove()

        num_layers = len(outputs.hidden_states) - 1
        layer_states = list(outputs.hidden_states[1:]) if location == "layer_output" else list(outputs.hidden_states[:-1])
        if final_layer_module is not None:
            if len(final_boundary) != 1:
                raise RuntimeError(
                    f"Expected exactly one final-layer forward per batch, observed {len(final_boundary)}."
                )
            layer_states[-1] = final_boundary[0]
        batch_row_mask = batch_mask.unsqueeze(-1) if batch_mask is not None else None
        for layer_idx, hs in enumerate(layer_states):
            if batch_row_mask is not None:
                # zero rows outside the attention mask; their values are computed but unused
                hs = hs * batch_row_mask.to(device=hs.device)
            all_hidden.setdefault(layer_idx, []).append(hs.cpu())

        if on_batch is not None:
            on_batch()

    result = {layer_idx: torch.cat(tensors, dim=0) for layer_idx, tensors in all_hidden.items()}

    if num_layers is not None and len(result) != num_layers:
        raise ValueError(
            f"Expected {num_layers} mapped hidden states for location={location!r}, got {len(result)}."
        )

    return result

data

Contrastive data containers shared by detection and steering estimation.

ContrastivePairs dataclass

Paired positive/negative text data for contrastive estimation.

The token sequence the model sees for each example is produced by render_for_model according to the consumer's prompt format, and the rendered string is tokenized with add_special_tokens=False for chat-templated text.

Attributes:

Name	Type	Description
positives	Sequence[str]	Texts exhibiting the target behavior. Treated as completions under prompt_format="chat_completion", and as standalone prompts under prompt_format="chat_prompt".
negatives	Sequence[str]	Texts not exhibiting the target behavior (see positives).
prompts	Sequence[str] | None	Optional shared prompts. Used as the user turn under prompt_format="chat_completion" and as the prefix under prompt_format="raw". Required when accumulate == "suffix-only". Ignored under prompt_format="chat_prompt".

Source code in steerability/algorithms/core/internals/data.py

@dataclass(frozen=True)
class ContrastivePairs:
    """Paired positive/negative text data for contrastive estimation.

    The token sequence the model sees for each example is produced by
    `render_for_model` according to the consumer's prompt format, and the
    rendered string is tokenized with `add_special_tokens=False` for
    chat-templated text.

    Attributes:
        positives: Texts exhibiting the target behavior. Treated as completions
            under `prompt_format="chat_completion"`, and as standalone prompts
            under `prompt_format="chat_prompt"`.
        negatives: Texts not exhibiting the target behavior (see `positives`).
        prompts: Optional shared prompts. Used as the user turn under
            `prompt_format="chat_completion"` and as the prefix under
            `prompt_format="raw"`. Required when `accumulate == "suffix-only"`.
            Ignored under `prompt_format="chat_prompt"`.
    """

    positives: Sequence[str]
    negatives: Sequence[str]
    prompts: Sequence[str] | None = None

    def __post_init__(self):
        if len(self.positives) == 0 or len(self.negatives) == 0:
            raise ValueError("positives and negatives must each have at least one entry.")
        if len(self.positives) != len(self.negatives):
            raise ValueError(
                f"positives ({len(self.positives)}) and negatives ({len(self.negatives)}) "
                f"must have equal length."
            )
        if self.prompts is not None and len(self.prompts) != len(self.positives):
            raise ValueError("prompts must have the same length as positives/negatives.")

negatives instance-attribute

positives instance-attribute

prompts = None class-attribute instance-attribute

LabeledExamples dataclass

Independent positive/negative text data with binary labels.

The positive and negative lists need not be the same length. Applies to methods where positive and negative examples are independent and unpaired, and the estimator concatenates them. Optional group keys mark statements that must not straddle a train/validation split (e.g. the source question of each answer); estimators that do not split ignore them.

Attributes:

Name	Type	Description
positives	Sequence[str]	Texts exhibiting the target behavior (label=1).
negatives	Sequence[str]	Texts not exhibiting the target behavior (label=0).
positive_groups	Sequence[str | int] | None	Group key per positive, same length as positives, or None.
negative_groups	Sequence[str | int] | None	Group key per negative, same length as negatives, or None.

Source code in steerability/algorithms/core/internals/data.py

@dataclass(frozen=True)
class LabeledExamples:
    """Independent positive/negative text data with binary labels.

    The positive and negative lists need not be the same length. Applies to methods where
    positive and negative examples are independent and unpaired, and the estimator concatenates
    them. Optional group keys mark statements that must not straddle a train/validation split
    (e.g. the source question of each answer); estimators that do not split ignore them.

    Attributes:
        positives: Texts exhibiting the target behavior (label=1).
        negatives: Texts not exhibiting the target behavior (label=0).
        positive_groups: Group key per positive, same length as `positives`, or None.
        negative_groups: Group key per negative, same length as `negatives`, or None.
    """

    positives: Sequence[str]
    negatives: Sequence[str]
    positive_groups: Sequence[str | int] | None = None
    negative_groups: Sequence[str | int] | None = None

    def __post_init__(self):
        if len(self.positives) == 0 or len(self.negatives) == 0:
            raise ValueError("positives and negatives must each have at least one entry.")

        given = [self.positive_groups is not None, self.negative_groups is not None]
        if any(given) and not all(given):
            raise ValueError(
                "positive_groups and negative_groups must both be given or both omitted."
            )

        for name, groups, texts in (
            ("positive_groups", self.positive_groups, self.positives),
            ("negative_groups", self.negative_groups, self.negatives),
        ):
            if groups is None:
                continue
            if len(groups) != len(texts):
                raise ValueError(
                    f"{name} has length {len(groups)} but its text list has length {len(texts)}."
                )
            for key in groups:
                if not isinstance(key, (str, int)) or isinstance(key, bool):
                    raise ValueError(
                        f"{name} keys must be str or int; got {type(key).__name__}."
                    )

    @property
    def groups(self) -> bool:
        """Whether the data carries group keys (positive/negative groups both present)."""
        return self.positive_groups is not None

groups property

Whether the data carries group keys (positive/negative groups both present).

negative_groups = None class-attribute instance-attribute

negatives instance-attribute

positive_groups = None class-attribute instance-attribute

positives instance-attribute

as_contrastive_pairs(x)

Normalize input to ContrastivePairs.

Accepts:

- An existing `ContrastivePairs` instance (returned as-is).
- A dict with keys `"positives"`, `"negatives"`, and optionally `"prompts"`.

Parameters:

Name	Type	Description	Default
x	ContrastivePairs | dict	Input to normalize.	required

Returns:

Type	Description
ContrastivePairs	A ContrastivePairs instance.

Raises:

Type	Description
TypeError	If input is neither ContrastivePairs nor a suitable dict.

Source code in steerability/algorithms/core/internals/data.py

def as_contrastive_pairs(x: ContrastivePairs | dict) -> ContrastivePairs:
    """Normalize input to `ContrastivePairs`.

    Accepts:

        - An existing `ContrastivePairs` instance (returned as-is).
        - A dict with keys `"positives"`, `"negatives"`, and optionally `"prompts"`.

    Args:
        x: Input to normalize.

    Returns:
        A `ContrastivePairs` instance.

    Raises:
        TypeError: If input is neither `ContrastivePairs` nor a suitable dict.
    """
    if isinstance(x, ContrastivePairs):
        return x
    if isinstance(x, dict):
        return ContrastivePairs(**x)
    raise TypeError("Expected ContrastivePairs or dict with positives/negatives[/prompts].")

as_labeled_examples(x)

Normalize input to LabeledExamples.

Accepts:

- An existing `LabeledExamples` instance (returned as-is).
- A `ContrastivePairs` instance (converted; pairing is dropped).
- A dict with keys `"positives"` and `"negatives"`, and optionally `"positive_groups"`
  and `"negative_groups"`.

Parameters:

Name	Type	Description	Default
x	LabeledExamples | ContrastivePairs | dict	Input to normalize.	required

Returns:

Type	Description
LabeledExamples	A LabeledExamples instance.

Raises:

Type	Description
TypeError	If input is not LabeledExamples, ContrastivePairs, or a suitable dict.

Source code in steerability/algorithms/core/internals/data.py

def as_labeled_examples(x: LabeledExamples | ContrastivePairs | dict) -> LabeledExamples:
    """Normalize input to `LabeledExamples`.

    Accepts:

        - An existing `LabeledExamples` instance (returned as-is).
        - A `ContrastivePairs` instance (converted; pairing is dropped).
        - A dict with keys `"positives"` and `"negatives"`, and optionally `"positive_groups"`
          and `"negative_groups"`.

    Args:
        x: Input to normalize.

    Returns:
        A `LabeledExamples` instance.

    Raises:
        TypeError: If input is not `LabeledExamples`, `ContrastivePairs`, or a suitable dict.
    """
    if isinstance(x, LabeledExamples):
        return x
    if isinstance(x, ContrastivePairs):
        return LabeledExamples(positives=x.positives, negatives=x.negatives)
    if isinstance(x, dict):
        return LabeledExamples(**x)
    raise TypeError("Expected LabeledExamples, ContrastivePairs, or dict with positives/negatives.")

encoding

Text-to-tensor tokenization entry points.

tokenize_pairs(tokenizer, pos_texts, neg_texts, device, *, add_special_tokens=True)

Tokenize positive/negative pairs together to ensure consistent padding.

Interleaves pairs before tokenization so each (pos, neg) pair shares the same padding length, keeping shared prefixes token-aligned.

Parameters:

Name	Type	Description	Default
tokenizer	PreTrainedTokenizerBase	Tokenizer to use.	required
pos_texts	Sequence[str]	List of positive text strings.	required
neg_texts	Sequence[str]	List of negative text strings (same length as pos_texts).	required
device	device | str	Target device.	required
add_special_tokens	bool	Whether to add special tokens (e.g. BOS). Pass False for chat-templated text that already contains them.	True

Returns:

Type	Description
tuple[dict[str, Tensor], dict[str, Tensor]]	Tuple of (enc_pos, enc_neg) dictionaries with input_ids and attention_mask.

Source code in steerability/algorithms/core/internals/encoding.py

def tokenize_pairs(
    tokenizer: PreTrainedTokenizerBase,
    pos_texts: Sequence[str],
    neg_texts: Sequence[str],
    device: torch.device | str,
    *,
    add_special_tokens: bool = True,
) -> tuple[dict[str, torch.Tensor], dict[str, torch.Tensor]]:
    """Tokenize positive/negative pairs together to ensure consistent padding.

    Interleaves pairs before tokenization so each (pos, neg) pair shares the same
    padding length, keeping shared prefixes token-aligned.

    Args:
        tokenizer: Tokenizer to use.
        pos_texts: List of positive text strings.
        neg_texts: List of negative text strings (same length as `pos_texts`).
        device: Target device.
        add_special_tokens: Whether to add special tokens (e.g. BOS). Pass False
            for chat-templated text that already contains them.

    Returns:
        Tuple of `(enc_pos, enc_neg)` dictionaries with `input_ids` and `attention_mask`.
    """
    # interleave: [pos0, neg0, pos1, neg1, ...]
    interleaved = []
    for pos, neg in zip(pos_texts, neg_texts):
        interleaved.append(pos)
        interleaved.append(neg)

    enc = tokenizer(
        interleaved,
        return_tensors="pt",
        padding=True,
        truncation=True,
        add_special_tokens=add_special_tokens,
    )
    enc = {k: v.to(device) for k, v in enc.items()}

    # split the interleaved batch back apart: even indices are positive, odd indices are negative
    enc_pos = {k: v[0::2] for k, v in enc.items()}
    enc_neg = {k: v[1::2] for k, v in enc.items()}

    return enc_pos, enc_neg

tokenize_texts(tokenizer, texts, device=None, *, add_special_tokens=True, max_length=None)

Tokenize a flat list of texts independently.

Each text is tokenized on its own, without interleaving. Use this when positive and negative examples are independent and do not need co-padding for token alignment.

Parameters:

Name	Type	Description	Default
tokenizer	PreTrainedTokenizerBase	Tokenizer to use.	required
texts	Sequence[str]	List of text strings.	required
device	device | str | None	Target device for the returned tensors. When None, tensors are returned on the device the tokenizer produces them on (CPU).	None
add_special_tokens	bool	Whether to add special tokens (e.g. BOS). Pass False for chat-templated text that already contains them.	True
max_length	int | None	Truncation bound. When None, truncation falls back to the tokenizer's model maximum length.	None

Returns:

Type	Description
dict[str, Tensor]	Dictionary with input_ids and attention_mask tensors.

Source code in steerability/algorithms/core/internals/encoding.py

def tokenize_texts(
    tokenizer: PreTrainedTokenizerBase,
    texts: Sequence[str],
    device: torch.device | str | None = None,
    *,
    add_special_tokens: bool = True,
    max_length: int | None = None,
) -> dict[str, torch.Tensor]:
    """Tokenize a flat list of texts independently.

    Each text is tokenized on its own, without interleaving. Use this when positive and negative
    examples are independent and do not need co-padding for token alignment.

    Args:
        tokenizer: Tokenizer to use.
        texts: List of text strings.
        device: Target device for the returned tensors. When None, tensors are returned on the
            device the tokenizer produces them on (CPU).
        add_special_tokens: Whether to add special tokens (e.g. BOS). Pass False
            for chat-templated text that already contains them.
        max_length: Truncation bound. When None, truncation falls back to the tokenizer's model
            maximum length.

    Returns:
        Dictionary with `input_ids` and `attention_mask` tensors.
    """
    enc = tokenizer(
        list(texts),
        return_tensors="pt",
        padding=True,
        truncation=True,
        max_length=max_length,
        add_special_tokens=add_special_tokens,
    )
    if device is None:
        return dict(enc)
    return {k: v.to(device) for k, v in enc.items()}

fingerprint

Digests identifying the model a fitted artifact was estimated on.

artifact_provenance_meta(model, tokenizer=None)

Provenance fingerprints for a fitted steering artifact.

Always records the toolkit model fingerprint; when vllm_hook_plugins is installed, adds the plugin's config and chat-template fingerprint recipes, which cross-check against a serving engine's discovery payload.

Parameters:

Name	Type	Description	Default
model	PreTrainedModel	The model the artifact was fitted on.	required
tokenizer	PreTrainedTokenizerBase | None	The tokenizer used during fitting, for the chat-template fingerprint.	None

Returns:

Type	Description
dict	The provenance mapping. Keys: "model_fingerprint", and when the plugin is installed,
dict	"config_fingerprint" and (with a tokenizer) "chat_template_fingerprint".

Source code in steerability/algorithms/core/internals/fingerprint.py

def artifact_provenance_meta(model: PreTrainedModel, tokenizer: PreTrainedTokenizerBase | None = None) -> dict:
    """Provenance fingerprints for a fitted steering artifact.

    Always records the toolkit model fingerprint; when `vllm_hook_plugins` is installed, adds
    the plugin's config and chat-template fingerprint recipes, which cross-check against a
    serving engine's discovery payload.

    Args:
        model: The model the artifact was fitted on.
        tokenizer: The tokenizer used during fitting, for the chat-template fingerprint.

    Returns:
        The provenance mapping. Keys: `"model_fingerprint"`, and when the plugin is installed,
        `"config_fingerprint"` and (with a tokenizer) `"chat_template_fingerprint"`.
    """
    meta = {"model_fingerprint": model_fingerprint(model)}
    try:
        from vllm_hook_plugins.core.fingerprints import chat_template_fingerprint, config_fingerprint
    except ImportError:
        return meta
    try:
        meta["config_fingerprint"] = config_fingerprint(model.config.to_dict())
    except (TypeError, ValueError):
        pass
    if tokenizer is not None:
        meta["chat_template_fingerprint"] = chat_template_fingerprint(
            getattr(tokenizer, "chat_template", None)
        )
    return meta

is_absent_chat_template_fingerprint(fingerprint)

Whether fingerprint denotes an absent chat template under the plugin recipe.

A serving engine that exposes no chat template reports the fingerprint of an empty template, so a mismatch against such a value reflects exposure rather than divergence and comparisons should skip it. Uses the plugin recipe when vllm_hook_plugins is installed and falls back to the recipe's stable digest of an empty template otherwise.

Parameters:

Name	Type	Description	Default
fingerprint	str | None	The reported fingerprint, or None when none was reported.	required

Returns:

Type	Description
bool	True when the fingerprint denotes an absent chat template.

Source code in steerability/algorithms/core/internals/fingerprint.py

def is_absent_chat_template_fingerprint(fingerprint: str | None) -> bool:
    """Whether `fingerprint` denotes an absent chat template under the plugin recipe.

    A serving engine that exposes no chat template reports the fingerprint of an empty
    template, so a mismatch against such a value reflects exposure rather than divergence
    and comparisons should skip it. Uses the plugin recipe when `vllm_hook_plugins` is
    installed and falls back to the recipe's stable digest of an empty template otherwise.

    Args:
        fingerprint: The reported fingerprint, or None when none was reported.

    Returns:
        True when the fingerprint denotes an absent chat template.
    """
    if not fingerprint:
        return True
    try:
        from vllm_hook_plugins.core.fingerprints import chat_template_fingerprint
    except ImportError:
        return fingerprint in (f"sha256:{_ABSENT_TEMPLATE_DIGEST}", _ABSENT_TEMPLATE_DIGEST)
    return fingerprint in (chat_template_fingerprint(None), chat_template_fingerprint(""))

model_fingerprint(model)

Compute a deterministic digest identifying a model, from its configuration, dtype, and a sample of its weights.

The digest is a sha256 over (a) model.config.to_json_string(), (b) str(model.dtype), and (c) for up to 8 evenly spaced named parameters, the parameter name plus its first 64 elements flattened and cast to float32 bytes, truncated to 16 hex characters.

Properties:

• The same checkpoint at the same dtype yields equal digests in any process or device placement, since parameter values are deterministic under from_pretrained.
• Any weight edit that touches the sampled slices (SFT, DPO, merging, LoRA merge) changes the digest; so does attaching unmerged adapters, since parameter names change.
• dtype is part of the identity, since a calibrated operating point is a numeric artifact and the same checkpoint loaded at a different dtype is a different runtime model for detection purposes.

Parameters:

Name	Type	Description	Default
model	PreTrainedModel	The model to fingerprint.	required

Returns:

Type	Description
str	A 16-character lowercase hex digest.

Source code in steerability/algorithms/core/internals/fingerprint.py

def model_fingerprint(model: PreTrainedModel) -> str:
    """Compute a deterministic digest identifying a model, from its configuration, dtype, and a sample of its weights.

    The digest is a sha256 over (a) `model.config.to_json_string()`, (b) `str(model.dtype)`, and
    (c) for up to 8 evenly spaced named parameters, the parameter name plus its first 64 elements
    flattened and cast to float32 bytes, truncated to 16 hex characters.

    Properties:

    - The same checkpoint at the same dtype yields equal digests in any process or device
        placement, since parameter values are deterministic under `from_pretrained`.
    - Any weight edit that touches the sampled slices (SFT, DPO, merging, LoRA merge) changes the
        digest; so does attaching unmerged adapters, since parameter names change.
    - dtype is part of the identity, since a calibrated operating point is a numeric artifact and
        the same checkpoint loaded at a different dtype is a different runtime model for detection
        purposes.

    Args:
        model: The model to fingerprint.

    Returns:
        A 16-character lowercase hex digest.
    """
    digest = hashlib.sha256()
    digest.update(model.config.to_json_string().encode("utf-8"))
    digest.update(str(model.dtype).encode("utf-8"))

    named_parameters = list(model.named_parameters())
    n = len(named_parameters)
    if n > 0:
        k = min(8, n)
        indices = [0] if k == 1 else [round(i * (n - 1) / (k - 1)) for i in range(k)]
        for idx in dict.fromkeys(indices):
            name, param = named_parameters[idx]
            digest.update(name.encode("utf-8"))
            sample = param.detach().reshape(-1)[:64].to(device="cpu", dtype=torch.float32).contiguous()
            digest.update(sample.numpy().tobytes())

    return digest.hexdigest()[:16]

session_artifact_identity(session)

(model_type, meta) recorded for a session-fitted artifact, from the session layout.

The meta carries the layout's model_fingerprint and model_ref when present, so identity checks against the venue that will read the artifact stay possible without a live model. Returns ("unknown", {}) when no session layout is available.

Parameters:

Name	Type	Description	Default
session	SteeringSession | None	The SteeringSession the artifact was fitted through, or None.	required

Returns:

Type	Description
tuple[str, dict]	The model type and the provenance mapping.

Source code in steerability/algorithms/core/internals/fingerprint.py

def session_artifact_identity(session: SteeringSession | None) -> tuple[str, dict]:
    """`(model_type, meta)` recorded for a session-fitted artifact, from the session layout.

    The meta carries the layout's `model_fingerprint` and `model_ref` when present, so
    identity checks against the venue that will read the artifact stay possible without a
    live model. Returns `("unknown", {})` when no session layout is available.

    Args:
        session: The `SteeringSession` the artifact was fitted through, or None.

    Returns:
        The model type and the provenance mapping.
    """
    layout = None
    if session is not None:
        try:
            layout = session.layout
        except Exception:
            layout = None
    if layout is None:
        return "unknown", {}
    meta: dict = {}
    if layout.model_fingerprint:
        meta["model_fingerprint"] = layout.model_fingerprint
    if layout.model_ref:
        meta["model_ref"] = layout.model_ref
    return layout.model_type or "unknown", meta

model_layout

Architecture-specific module paths for reading and steering a model's decoder stack.

ModelLayout names the decoder layer prefix, attention/output-projection suffixes, and normalization sub-module attributes for one model family. resolve_model_layout maps a PreTrainedModel to its layout by locating the decoder stack at one of an ordered set of dotted roots and probing the first decoder layer's attributes to select a naming convention.

Decoder-stack roots, in probe order:

• model.layers: text-only decoder language models (Llama, Mistral, Qwen, Gemma text).
• model.language_model.layers: composite multimodal wrappers (the transformers >= 5 layout, e.g. Gemma 3/4, Llava, Qwen-VL), where the toolkit steers the text decoder under text-only prompting.
• transformer.h: GPT-2-style models.

Per-layer naming conventions, in probe order (gemma_style precedes llama_style because a Gemma decoder layer also carries the Llama markers):

• gemma_style: self_attn, residual-stream norms input_layernorm and pre_feedforward_layernorm.
• llama_style: self_attn, residual-stream norms input_layernorm and post_attention_layernorm.
• gpt2_style: attn, residual-stream norms ln_1 and ln_2.

A convention matches when its norm markers exist on the first decoder layer and its attention module exists on at least one decoder layer. Hybrid stacks that interleave attention layers with another token mixer (Qwen3.5 and Qwen3-Next, where three Gated DeltaNet linear_attn layers precede each self_attn layer) therefore resolve to the convention of their attention layers, and the layout records which layers carry attention in attention_layer_ids. Residual-stream sites (decoder-layer boundaries and the norm inputs) exist on every layer of such a stack; attention sites (attn_names, oproj_names, head_geometry) exist only on the attention layers, and consumers check has_attention before resolving them.

PEFT wrappers are peeled before resolution: PeftModel.base_model is the tuner and the tuner's model attribute is the wrapped model, both registered submodules, so each wrapper contributes the prefix "base_model.model." and the resulting paths resolve through get_submodule without relying on attribute forwarding.

register_layout_detector adds a callable consulted before the built-in resolution, so a user on an unlisted family is not blocked on a toolkit release.

logger = logging.getLogger(__name__) module-attribute

HeadGeometry dataclass

Attention head geometry of one decoder layer.

Attributes:

Name	Type	Description
num_heads	int	Number of attention heads.
head_dim	int	Per-head dimension.

Source code in steerability/algorithms/core/internals/model_layout.py

@dataclass(frozen=True)
class HeadGeometry:
    """Attention head geometry of one decoder layer.

    Attributes:
        num_heads: Number of attention heads.
        head_dim: Per-head dimension.
    """

    num_heads: int
    head_dim: int

head_dim instance-attribute

num_heads instance-attribute

ModelLayout dataclass

Architecture-specific module paths for one model family.

Attributes:

Name	Type	Description
family	str	Family label ("gemma_style", "llama_style", or "gpt2_style").
layer_prefix	str	Dotted prefix of the decoder layer list (e.g. "model.layers", "model.language_model.layers", "transformer.h", prefixed with "base_model.model." per PEFT wrapper).
num_layers	int	Number of decoder layers.
attn_suffix	str	Suffix of the per-layer attention module (".self_attn" or ".attn").
oproj_suffix	str	Suffix of the attention output projection (".self_attn.o_proj" or ".attn.c_proj").
norm_attrs	tuple[str, ...]	The per-layer normalization sub-module attribute names whose input is the residual stream (("input_layernorm", "pre_feedforward_layernorm") on Gemma, ("input_layernorm", "post_attention_layernorm") on Llama, ("ln_1", "ln_2") on GPT-2).
attention_layer_ids	tuple[int, ...] | None	The decoder layers carrying the attention module at attn_suffix, ascending. None means every layer (homogeneous stacks, and layouts built by hand); a hybrid stack lists only its attention layers, and an empty tuple declares that no layer carries attention.

Source code in steerability/algorithms/core/internals/model_layout.py

@dataclass(frozen=True)
class ModelLayout:
    """Architecture-specific module paths for one model family.

    Attributes:
        family: Family label (`"gemma_style"`, `"llama_style"`, or `"gpt2_style"`).
        layer_prefix: Dotted prefix of the decoder layer list (e.g. `"model.layers"`,
            `"model.language_model.layers"`, `"transformer.h"`, prefixed with
            `"base_model.model."` per PEFT wrapper).
        num_layers: Number of decoder layers.
        attn_suffix: Suffix of the per-layer attention module (`".self_attn"` or `".attn"`).
        oproj_suffix: Suffix of the attention output projection (`".self_attn.o_proj"` or
            `".attn.c_proj"`).
        norm_attrs: The per-layer normalization sub-module attribute names whose input is the
            residual stream (`("input_layernorm", "pre_feedforward_layernorm")` on Gemma,
            `("input_layernorm", "post_attention_layernorm")` on Llama, `("ln_1", "ln_2")` on
            GPT-2).
        attention_layer_ids: The decoder layers carrying the attention module at `attn_suffix`,
            ascending. None means every layer (homogeneous stacks, and layouts built by hand);
            a hybrid stack lists only its attention layers, and an empty tuple declares that
            no layer carries attention.
    """

    family: str
    layer_prefix: str
    num_layers: int
    attn_suffix: str
    oproj_suffix: str
    norm_attrs: tuple[str, ...]
    attention_layer_ids: tuple[int, ...] | None = None

    @property
    def layer_names(self) -> list[str]:
        """Dotted paths of every decoder layer, `[f"{layer_prefix}.{i}"]`."""
        return [f"{self.layer_prefix}.{i}" for i in range(self.num_layers)]

    @property
    def oproj_names(self) -> list[str]:
        """Dotted paths of every layer's attention output projection, index-aligned with
        `layer_names`. On a hybrid stack, entries outside `attention_layers` name modules the
        model does not have."""
        return [name + self.oproj_suffix for name in self.layer_names]

    @property
    def attn_names(self) -> list[str]:
        """Dotted paths of every layer's attention module, index-aligned with `layer_names`. On a
        hybrid stack, entries outside `attention_layers` name modules the model does not have."""
        return [name + self.attn_suffix for name in self.layer_names]

    @property
    def attention_layers(self) -> tuple[int, ...]:
        """The decoder layers carrying an attention module, ascending; every layer when
        `attention_layer_ids` is None."""
        if self.attention_layer_ids is None:
            return tuple(range(self.num_layers))
        return self.attention_layer_ids

    @property
    def is_hybrid(self) -> bool:
        """True when some decoder layer carries no attention module."""
        return len(self.attention_layers) < self.num_layers

    def has_attention(self, layer_id: int) -> bool:
        """True when decoder layer `layer_id` carries the attention module at `attn_suffix`."""
        return layer_id in self.attention_layers

attention_layer_ids = None class-attribute instance-attribute

attention_layers property

The decoder layers carrying an attention module, ascending; every layer when attention_layer_ids is None.

attn_names property

Dotted paths of every layer's attention module, index-aligned with layer_names. On a hybrid stack, entries outside attention_layers name modules the model does not have.

attn_suffix instance-attribute

family instance-attribute

is_hybrid property

True when some decoder layer carries no attention module.

layer_names property

Dotted paths of every decoder layer, [f"{layer_prefix}.{i}"].

layer_prefix instance-attribute

norm_attrs instance-attribute

num_layers instance-attribute

oproj_names property

Dotted paths of every layer's attention output projection, index-aligned with layer_names. On a hybrid stack, entries outside attention_layers name modules the model does not have.

oproj_suffix instance-attribute

has_attention(layer_id)

True when decoder layer layer_id carries the attention module at attn_suffix.

Source code in steerability/algorithms/core/internals/model_layout.py

def has_attention(self, layer_id: int) -> bool:
    """True when decoder layer `layer_id` carries the attention module at `attn_suffix`."""
    return layer_id in self.attention_layers

head_geometry(model, layout, layer_id)

Per-layer attention head geometry, read from the module tree.

head_dim comes from the attention module's head_dim attribute, else the text config; num_heads is the output projection's input width divided by head_dim. GPT-2's Conv1D stores its weight as [in, out], so the width is read from weight.shape[0] when in_features is absent.

Parameters:

Name	Type	Description	Default
model	PreTrainedModel	The live model (or PEFT wrapper) carrying the attention modules.	required
layout	ModelLayout	The resolved ModelLayout naming the per-layer attention and output projections.	required
layer_id	int	The decoder layer index.	required

Returns:

Type	Description
HeadGeometry	The layer's HeadGeometry.

Raises:

Type	Description
ValueError	If layer_id carries no attention module (a non-attention layer of a hybrid stack), or the projection width is not a multiple of head_dim.

Source code in steerability/algorithms/core/internals/model_layout.py

def head_geometry(model: PreTrainedModel, layout: ModelLayout, layer_id: int) -> HeadGeometry:
    """Per-layer attention head geometry, read from the module tree.

    `head_dim` comes from the attention module's `head_dim` attribute, else the text config;
    `num_heads` is the output projection's input width divided by `head_dim`. GPT-2's `Conv1D`
    stores its weight as `[in, out]`, so the width is read from `weight.shape[0]` when
    `in_features` is absent.

    Args:
        model: The live model (or PEFT wrapper) carrying the attention modules.
        layout: The resolved `ModelLayout` naming the per-layer attention and output projections.
        layer_id: The decoder layer index.

    Returns:
        The layer's `HeadGeometry`.

    Raises:
        ValueError: If `layer_id` carries no attention module (a non-attention layer of a hybrid
            stack), or the projection width is not a multiple of `head_dim`.
    """
    if not layout.has_attention(layer_id):
        raise ValueError(
            f"Decoder layer {layer_id} carries no attention module ({layout.attn_suffix!r}); "
            f"attention layers of this model are {list(layout.attention_layers)}."
        )
    attn = model.get_submodule(layout.attn_names[layer_id])
    head_dim = getattr(attn, "head_dim", None)
    if head_dim is None:
        text_cfg = text_config(model)
        head_dim = getattr(text_cfg, "head_dim", None)
        if head_dim is None:
            head_dim = text_cfg.hidden_size // text_cfg.num_attention_heads

    oproj = model.get_submodule(layout.oproj_names[layer_id])
    width = getattr(oproj, "in_features", None)
    if width is None:
        width = oproj.weight.shape[0]  # gpt-2 Conv1D stores weight as [in, out]

    if width % head_dim != 0:
        raise ValueError(
            f"Output-projection input width {width} at layer {layer_id} is not a multiple of "
            f"head_dim {head_dim}; cannot infer the head count."
        )
    return HeadGeometry(num_heads=width // head_dim, head_dim=head_dim)

register_layout_detector(detector, *, prepend=True)

Register a callable mapping a model to a ModelLayout (or None to decline).

Registered detectors run before the built-in root and convention probing, in registration order (prepend=True puts the new detector first). Registration is process-global.

Parameters:

Name	Type	Description	Default
detector	Callable[[PreTrainedModel], ModelLayout | None]	A callable that returns a ModelLayout for models it recognizes, or None to fall through to the next detector and then the built-in resolution.	required
prepend	bool	When True (the default), the detector runs before previously registered ones.	True

Source code in steerability/algorithms/core/internals/model_layout.py

def register_layout_detector(
    detector: Callable[[PreTrainedModel], ModelLayout | None], *, prepend: bool = True
) -> None:
    """Register a callable mapping a model to a `ModelLayout` (or None to decline).

    Registered detectors run before the built-in root and convention probing, in registration
    order (`prepend=True` puts the new detector first). Registration is process-global.

    Args:
        detector: A callable that returns a `ModelLayout` for models it recognizes, or None to
            fall through to the next detector and then the built-in resolution.
        prepend: When True (the default), the detector runs before previously registered ones.
    """
    if prepend:
        _DETECTORS.insert(0, detector)
    else:
        _DETECTORS.append(detector)

resolve_model_layout(model)

Resolve the ModelLayout for a HuggingFace causal LM.

Consults registered detectors first, then locates the decoder stack at one of the built-in roots and selects a naming convention whose norm markers exist on the first decoder layer and whose attention module exists on at least one layer. A hybrid stack that interleaves attention layers with another token mixer resolves to its attention layers' family, with those layers recorded in attention_layer_ids. PEFT wrappers are peeled before resolution and their prefix is carried on the resolved layer_prefix.

Parameters:

Name	Type	Description	Default
model	PreTrainedModel	A HuggingFace causal LM, a multimodal wrapper whose text decoder sits at a known root, or a PEFT-wrapped model.	required

Returns:

Type	Description
ModelLayout	The matching ModelLayout.

Raises:

Type	Description
ValueError	If no detector and no built-in root/convention matches the model.

Source code in steerability/algorithms/core/internals/model_layout.py

def resolve_model_layout(model: PreTrainedModel) -> ModelLayout:
    """Resolve the `ModelLayout` for a HuggingFace causal LM.

    Consults registered detectors first, then locates the decoder stack at one of the built-in
    roots and selects a naming convention whose norm markers exist on the first decoder layer and
    whose attention module exists on at least one layer. A hybrid stack that interleaves attention
    layers with another token mixer resolves to its attention layers' family, with those layers
    recorded in `attention_layer_ids`. PEFT wrappers are peeled before resolution and their prefix
    is carried on the resolved `layer_prefix`.

    Args:
        model: A HuggingFace causal LM, a multimodal wrapper whose text decoder sits at a known
            root, or a PEFT-wrapped model.

    Returns:
        The matching `ModelLayout`.

    Raises:
        ValueError: If no detector and no built-in root/convention matches the model.
    """
    for detector in _DETECTORS:
        layout = detector(model)
        if layout is not None:
            return layout
    layout = _resolve_builtin(model)
    if layout is not None:
        return layout
    raise ValueError(
        f"Cannot determine model layout for {type(model).__name__}. Probed decoder-stack roots "
        f"{_DECODER_ROOTS} for a stack matching one of the families "
        f"{tuple(convention.family for convention in _CONVENTIONS)} (the family's norm "
        "sub-modules on the first layer and its attention module on at least one layer); "
        "register a detector with register_layout_detector() for other architectures."
    )

text_config(model_or_config)

The config carrying the text decoder's structural facts.

Composite multimodal configs (text plus vision or audio sub-configs) return their text sub-config; plain configs return themselves. Numeric structural facts (hidden_size, num_attention_heads, head_dim, num_hidden_layers, intermediate_size) are read from the returned config; identity facts (model_type, the fingerprint JSON, _attn_implementation) stay on the composite model.config.

Parameters:

Name	Type	Description	Default
model_or_config	PreTrainedModel | PretrainedConfig	A model (its config is used) or a config.	required

Returns:

Type	Description
PretrainedConfig	The text config.

Source code in steerability/algorithms/core/internals/model_layout.py

def text_config(model_or_config: PreTrainedModel | PretrainedConfig) -> PretrainedConfig:
    """The config carrying the text decoder's structural facts.

    Composite multimodal configs (text plus vision or audio sub-configs) return their text
    sub-config; plain configs return themselves. Numeric structural facts (`hidden_size`,
    `num_attention_heads`, `head_dim`, `num_hidden_layers`, `intermediate_size`) are read from the
    returned config; identity facts (`model_type`, the fingerprint JSON, `_attn_implementation`)
    stay on the composite `model.config`.

    Args:
        model_or_config: A model (its `config` is used) or a config.

    Returns:
        The text config.
    """
    config = getattr(model_or_config, "config", model_or_config)
    return config.get_text_config()

pooling

Token, span, and position aggregation primitives for captured hidden states.

aggregate_condition_hidden(hidden, mode, attention_mask=None)

Aggregate [B, T, H] hidden states to [B, H] using non-pad tokens only.

Parameters:

Name	Type	Description	Default
hidden	Tensor	Shape [B, T, H].	required
mode	Literal['mean', 'last']	"mean" pools over all real tokens; "last" selects the last real token per row.	required
attention_mask	Tensor | None	Shape [B, T] (1 for real tokens, 0 for pads), or None. When None, "mean" averages all positions and "last" uses the final position.	None

Returns:

Type	Description
Tensor	Aggregated tensor of shape [B, H].

Raises:

Type	Description
ValueError	If a row has no real tokens, or the mode is unsupported.

Source code in steerability/algorithms/core/internals/pooling.py

@torch.no_grad()
def aggregate_condition_hidden(
    hidden: torch.Tensor,
    mode: Literal["mean", "last"],
    attention_mask: torch.Tensor | None = None,
) -> torch.Tensor:
    """Aggregate `[B, T, H]` hidden states to `[B, H]` using non-pad tokens only.

    Args:
        hidden: Shape `[B, T, H]`.
        mode: "mean" pools over all real tokens; "last" selects the last real token per row.
        attention_mask: Shape `[B, T]` (1 for real tokens, 0 for pads), or None. When None, "mean"
            averages all positions and "last" uses the final position.

    Returns:
        Aggregated tensor of shape `[B, H]`.

    Raises:
        ValueError: If a row has no real tokens, or the mode is unsupported.
    """
    if mode == "mean":
        return masked_mean(hidden, attention_mask)

    if mode != "last":
        raise ValueError(f"Unsupported condition comparison mode: {mode!r}.")

    batch_size = hidden.size(0)
    if attention_mask is None:
        return hidden[:, -1, :]

    mask = attention_mask.to(hidden.device).bool()
    if not mask.any(dim=1).all():
        raise ValueError("aggregate_condition_hidden received a row with no real tokens.")

    sequence_length = mask.size(1)
    positions = torch.arange(sequence_length, device=hidden.device).unsqueeze(0)
    last_positions = positions.masked_fill(~mask, -1).max(dim=1).values  # [B]
    return hidden[torch.arange(batch_size, device=hidden.device), last_positions]

get_last_token_positions(attention_mask, seq_len, num_samples)

Find the last non-pad token position for each sample.

Parameters:

Name	Type	Description	Default
attention_mask	Tensor | None	Shape [N, T] or None.	required
seq_len	int	Sequence length T.	required
num_samples	int	Number of samples N.	required

Returns:

Type	Description
LongTensor	Tensor of shape [N] with last token positions.

Source code in steerability/algorithms/core/internals/pooling.py

def get_last_token_positions(
    attention_mask: torch.Tensor | None,
    seq_len: int,
    num_samples: int,
) -> torch.LongTensor:
    """Find the last non-pad token position for each sample.

    Args:
        attention_mask: Shape `[N, T]` or None.
        seq_len: Sequence length `T`.
        num_samples: Number of samples `N`.

    Returns:
        Tensor of shape `[N]` with last token positions.
    """
    if attention_mask is None:
        # no padding, last token is at seq_len - 1
        return torch.full((num_samples,), seq_len - 1, dtype=torch.long)

    # for each sample, find the last position where attention_mask == 1
    # this handles both left-padded and right-padded sequences
    positions = torch.arange(seq_len, device=attention_mask.device).unsqueeze(0).expand(num_samples, -1)
    # mask out padded positions with -1
    masked_positions = torch.where(attention_mask == 1, positions, torch.tensor(-1, device=attention_mask.device))
    return masked_positions.max(dim=1).values

masked_mean(hidden, attention_mask)

Mean-pool [N, T, H] hidden states over real (non-pad) positions.

Falls back to an unmasked mean over all positions when no attention mask is available.

Parameters:

Name	Type	Description	Default
hidden	Tensor	Shape [N, T, H].	required
attention_mask	Tensor | None	Shape [N, T] (1 for real tokens, 0 for pads), or None.	required

Returns:

Type	Description
Tensor	Pooled tensor of shape [N, H].

Source code in steerability/algorithms/core/internals/pooling.py

def masked_mean(hidden: torch.Tensor, attention_mask: torch.Tensor | None) -> torch.Tensor:
    """Mean-pool `[N, T, H]` hidden states over real (non-pad) positions.

    Falls back to an unmasked mean over all positions when no attention mask is available.

    Args:
        hidden: Shape `[N, T, H]`.
        attention_mask: Shape `[N, T]` (1 for real tokens, 0 for pads), or None.

    Returns:
        Pooled tensor of shape `[N, H]`.
    """
    if attention_mask is None:
        return hidden.mean(dim=1)
    m = attention_mask.to(hidden.dtype).unsqueeze(-1)  # [N, T, 1]
    return (hidden * m).sum(dim=1) / m.sum(dim=1).clamp_min(1e-8)

pool_over_spans(hidden, spans)

Mean-pool hidden states over specified spans.

A degenerate span (start >= end) falls back to the sample's last token position.

Parameters:

Name	Type	Description	Default
hidden	Tensor	Shape [N, T, H].	required
spans	list[tuple[int, int]]	List of (start, end) tuples.	required

Returns:

Type	Description
Tensor	Pooled tensor of shape [N, H].

Source code in steerability/algorithms/core/internals/pooling.py

def pool_over_spans(
    hidden: torch.Tensor,
    spans: list[tuple[int, int]],
) -> torch.Tensor:
    """Mean-pool hidden states over specified spans.

    A degenerate span (`start >= end`) falls back to the sample's last token position.

    Args:
        hidden: Shape `[N, T, H]`.
        spans: List of `(start, end)` tuples.

    Returns:
        Pooled tensor of shape `[N, H]`.
    """
    N, T, H = hidden.shape
    pooled = []
    for i, (start, end) in enumerate(spans):
        if start >= end:
            # fallback: use last token
            pooled.append(hidden[i, -1, :])
        else:
            pooled.append(hidden[i, start:end, :].mean(dim=0))
    return torch.stack(pooled, dim=0)

select_at_positions(hidden, positions)

Select hidden states at specified positions for each sample.

Parameters:

Name	Type	Description	Default
hidden	Tensor	Shape [N, T, H].	required
positions	LongTensor	Shape [N] with position indices.	required

Returns:

Type	Description
Tensor	Tensor of shape [N, H].

Source code in steerability/algorithms/core/internals/pooling.py

def select_at_positions(
    hidden: torch.Tensor,
    positions: torch.LongTensor,
) -> torch.Tensor:
    """Select hidden states at specified positions for each sample.

    Args:
        hidden: Shape `[N, T, H]`.
        positions: Shape `[N]` with position indices.

    Returns:
        Tensor of shape `[N, H]`.
    """
    N, _, H = hidden.shape
    # gather at the specified positions
    idx = positions.view(N, 1, 1).expand(N, 1, H)
    return hidden.gather(dim=1, index=idx).squeeze(1)

select_spans(enc, prompt_enc, accumulate)

Determine token spans to pool over for each sample.

Span bounds are derived from the attention mask, so pads on either padding side are excluded. Under accumulate="suffix-only", the per-sample prompt length (the number of real tokens in prompt_enc) is skipped at the start of each span.

Parameters:

Name	Type	Description	Default
enc	dict[str, Tensor]	Tokenized full sequences (prompts + completions).	required
prompt_enc	dict[str, Tensor] | None	Tokenized prompts only (if accumulate == "suffix-only").	required
accumulate	str	"all", "suffix-only", or "last_token".	required

Returns:

Type	Description
list[tuple[int, int]]	List of (start, end) tuples, one per sample.

Raises:

Type	Description
ValueError	If accumulate is unsupported.

Source code in steerability/algorithms/core/internals/pooling.py

def select_spans(
    enc: dict[str, torch.Tensor],
    prompt_enc: dict[str, torch.Tensor] | None,
    accumulate: str,
) -> list[tuple[int, int]]:
    """Determine token spans to pool over for each sample.

    Span bounds are derived from the attention mask, so pads on either padding side are excluded.
    Under `accumulate="suffix-only"`, the per-sample prompt length (the number of real tokens in
    `prompt_enc`) is skipped at the start of each span.

    Args:
        enc: Tokenized full sequences (prompts + completions).
        prompt_enc: Tokenized prompts only (if accumulate == "suffix-only").
        accumulate: "all", "suffix-only", or "last_token".

    Returns:
        List of (start, end) tuples, one per sample.

    Raises:
        ValueError: If `accumulate` is unsupported.
    """
    if accumulate not in ("all", "suffix-only", "last_token"):
        raise ValueError(
            f"select_spans does not support accumulate='{accumulate}'. "
            "Expected one of: 'all', 'suffix-only', 'last_token'."
        )

    input_ids = enc["input_ids"]
    attention_mask = enc.get("attention_mask")
    N, T = input_ids.shape

    spans = []
    for i in range(N):
        # number of prompt tokens to skip for suffix-only pooling (a count, not an absolute index)
        if accumulate == "suffix-only" and prompt_enc is not None:
            prompt_len = (
                int(prompt_enc["attention_mask"][i].sum().item())
                if "attention_mask" in prompt_enc
                else prompt_enc["input_ids"].size(1)
            )
        else:
            prompt_len = 0

        # derive both bounds from the mask so the span excludes pads on either padding side
        if attention_mask is not None:
            non_pad = (attention_mask[i] == 1).nonzero(as_tuple=True)[0]
            if len(non_pad) > 0:
                first = int(non_pad[0].item())
                last = int(non_pad[-1].item())
            else:
                first, last = 0, T - 1
            if accumulate == "last_token":
                # one-token span at the final non-pad position (derived from the mask, so it works
                # regardless of padding side)
                start, end = last, last + 1
            else:
                start = first + prompt_len
                end = last + 1
        else:
            if accumulate == "last_token":
                start, end = T - 1, T
            else:
                start = prompt_len
                end = T

        spans.append((start, end))

    return spans

render

Render ContrastivePairs into model-ready text.

The general, tokenizer-level rendering primitives live in steerability.utils.rendering; this module adds the data-specific renderer that turns a ContrastivePairs into model-ready text under a PromptFormat.

logger = logging.getLogger(__name__) module-attribute

RenderedContrastive dataclass

Rendered text for a ContrastivePairs, plus tokenization policy.

Attributes:

Name	Type	Description
pos_texts	list[str]	Rendered positive examples.
neg_texts	list[str]	Rendered negative examples.
prompt_texts	list[str] | None	User-turn-only renders for suffix-only span computation, or None when there is no separate prompt/completion split.
add_special_tokens	bool	Flag callers must pass to the tokenizer for all of the above (False for chat modes, True for raw).
effective_mode	PromptFormat	The mode actually applied after raw fallbacks.

Source code in steerability/algorithms/core/internals/render.py

@dataclass(frozen=True)
class RenderedContrastive:
    """Rendered text for a `ContrastivePairs`, plus tokenization policy.

    Attributes:
        pos_texts: Rendered positive examples.
        neg_texts: Rendered negative examples.
        prompt_texts: User-turn-only renders for suffix-only span computation, or
            None when there is no separate prompt/completion split.
        add_special_tokens: Flag callers must pass to the tokenizer for all of
            the above (False for chat modes, True for raw).
        effective_mode: The mode actually applied after raw fallbacks.
    """

    pos_texts: list[str]
    neg_texts: list[str]
    prompt_texts: list[str] | None
    add_special_tokens: bool
    effective_mode: PromptFormat

add_special_tokens instance-attribute

effective_mode instance-attribute

neg_texts instance-attribute

pos_texts instance-attribute

prompt_texts instance-attribute

render_contrastive(tokenizer, data, mode)

Render both sides of a contrastive dataset under mode.

Accepts ContrastivePairs (paired, with optional shared prompts) or LabeledExamples (unpaired classes, rendered independently). Resolves the effective mode (with raw fallbacks and warnings), renders positives/negatives, renders the prompt-only strings used for suffix-only span computation, and reports the add_special_tokens flag the tokenizer must use for all of the above.

Parameters:

Name	Type	Description	Default
tokenizer	PreTrainedTokenizerBase	Tokenizer whose chat template defines the rendering.	required
data	ContrastivePairs | LabeledExamples	ContrastivePairs (with positives, negatives, and optional prompts) or LabeledExamples.	required
mode	PromptFormat	Requested rendering policy.	required

Returns:

Type	Description
RenderedContrastive	A RenderedContrastive with rendered texts and tokenization policy.

Raises:

Type	Description
ValueError	If mode is "chat_completion" and data is LabeledExamples (that format's shared prompts align per example, which requires paired data).

Source code in steerability/algorithms/core/internals/render.py

def render_contrastive(
    tokenizer: PreTrainedTokenizerBase,
    data: ContrastivePairs | LabeledExamples,
    mode: PromptFormat,
) -> RenderedContrastive:
    """Render both sides of a contrastive dataset under `mode`.

    Accepts `ContrastivePairs` (paired, with optional shared prompts) or
    `LabeledExamples` (unpaired classes, rendered independently). Resolves the
    effective mode (with raw fallbacks and warnings), renders
    positives/negatives, renders the prompt-only strings used for suffix-only
    span computation, and reports the `add_special_tokens` flag the tokenizer
    must use for all of the above.

    Args:
        tokenizer: Tokenizer whose chat template defines the rendering.
        data: `ContrastivePairs` (with `positives`, `negatives`, and optional
            `prompts`) or `LabeledExamples`.
        mode: Requested rendering policy.

    Returns:
        A `RenderedContrastive` with rendered texts and tokenization policy.

    Raises:
        ValueError: If `mode` is `"chat_completion"` and `data` is
            `LabeledExamples` (that format's shared prompts align per example,
            which requires paired data).
    """
    if mode == "chat_completion" and isinstance(data, LabeledExamples):
        raise ValueError(
            "prompt_format='chat_completion' requires paired data with shared prompts "
            "(ContrastivePairs); LabeledExamples classes are unpaired."
        )
    prompts = getattr(data, "prompts", None)

    effective: PromptFormat = mode

    if mode in ("chat_completion", "chat_prompt") and not has_chat_template(tokenizer):
        logger.warning("render_contrastive: no chat_template; falling back to raw.")
        effective = "raw"

    if mode == "chat_completion" and prompts is None:
        logger.warning(
            "prompt_format='chat_completion' requires `prompts` (positives/negatives "
            "are treated as completions); none provided. Falling back to raw."
        )
        effective = "raw"

    add_special = effective == "raw"

    if effective == "raw":
        if prompts is not None:
            pos = [p + c for p, c in zip(prompts, data.positives)]
            neg = [p + c for p, c in zip(prompts, data.negatives)]
            prompt_texts = list(prompts)
        else:
            pos = list(data.positives)
            neg = list(data.negatives)
            prompt_texts = None
    elif effective == "chat_completion":
        pos = [
            render_for_model(tokenizer, prompt=p, completion=c, mode="chat_completion")
            for p, c in zip(prompts, data.positives)
        ]
        neg = [
            render_for_model(tokenizer, prompt=p, completion=c, mode="chat_completion")
            for p, c in zip(prompts, data.negatives)
        ]
        prompt_texts = [render_for_model(tokenizer, prompt=p, mode="chat_prompt") for p in prompts]
    else:  # chat_prompt mode, where each positive/negative is a standalone prompt
        pos = [render_for_model(tokenizer, prompt=t, mode="chat_prompt") for t in data.positives]
        neg = [render_for_model(tokenizer, prompt=t, mode="chat_prompt") for t in data.negatives]
        prompt_texts = None  # no separate prompt/completion split

    return RenderedContrastive(pos, neg, prompt_texts, add_special, effective)

stats

Estimate per-layer activation moments and characterize residual-stream norms.

STATS_FORMAT_VERSION = 1 module-attribute

logger = logging.getLogger(__name__) module-attribute

ActivationStats dataclass

Per-layer mean and variance of a model's ambient activation distribution.

The statistics support per-coordinate centering and scaling of pooled activations (standardize), the whitening used by diagonal-LDA probe fitting. Covariance is diagonal, since the dominant effects in residual-stream geometry (a large shared component in every activation and rogue dimensions dominating inner products) are axis-aligned, so per-coordinate centering and scaling removes most of the effect, and a full covariance is hard to estimate reliably and expensive to invert at a few thousand samples.

Under pooling="tokens", every retained real token position is one sample of the ambient distribution, and exclude_first_n (default 1) drops the leading real positions per sequence, since the first real position (BOS or an opening delimiter) carries known massive-activation outliers that distort the mean disproportionately. "mean" and "last" pool per prompt first, yielding one sample per prompt, and ignore exclude_first_n.

count is the number of pooled samples accumulated, i.e., tokens under pooling="tokens" and prompts otherwise. estimate warns when count falls below its min_samples argument; a few hundred varied texts clears the default under token pooling, while "mean" and "last" need a few thousand prompts.

Attributes:

Name	Type	Description
model_type	str	HuggingFace model_type of the model the statistics were estimated on.
model_fingerprint	str	Identity digest of that model, recorded at estimate time.
location	str	Residual-stream boundary the statistics were measured at.
count	int	Number of pooled samples accumulated.
mean	dict[int, Tensor]	Per-layer mean, [H] float32.
var	dict[int, Tensor]	Per-layer variance, [H] float32, floored at var_floor.
var_floor	float	Lower bound applied to every variance coordinate.

Source code in steerability/algorithms/core/internals/stats.py

@dataclass
class ActivationStats:
    """Per-layer mean and variance of a model's ambient activation distribution.

    The statistics support per-coordinate centering and scaling of pooled activations
    (`standardize`), the whitening used by diagonal-LDA probe fitting. Covariance is diagonal,
    since the dominant effects in residual-stream geometry (a large shared component in every
    activation and rogue dimensions dominating inner products) are axis-aligned, so per-coordinate
    centering and scaling removes most of the effect, and a full covariance is hard to estimate
    reliably and expensive to invert at a few thousand samples.

    Under `pooling="tokens"`, every retained real token position is one sample of the ambient
    distribution, and `exclude_first_n` (default 1) drops the leading real positions per sequence,
    since the first real position (BOS or an opening delimiter) carries known massive-activation
    outliers that distort the mean disproportionately. `"mean"` and `"last"` pool per prompt first,
    yielding one sample per prompt, and ignore `exclude_first_n`.

    `count` is the number of pooled samples accumulated, i.e., tokens under `pooling="tokens"` and
    prompts otherwise. `estimate` warns when `count` falls below its `min_samples` argument; a few
    hundred varied texts clears the default under token pooling, while `"mean"` and `"last"` need a
    few thousand prompts.

    Attributes:
        model_type: HuggingFace `model_type` of the model the statistics were estimated on.
        model_fingerprint: Identity digest of that model, recorded at estimate time.
        location: Residual-stream boundary the statistics were measured at.
        count: Number of pooled samples accumulated.
        mean: Per-layer mean, `[H]` float32.
        var: Per-layer variance, `[H]` float32, floored at `var_floor`.
        var_floor: Lower bound applied to every variance coordinate.
    """

    model_type: str
    model_fingerprint: str
    location: str
    count: int
    mean: dict[int, torch.Tensor]
    var: dict[int, torch.Tensor]
    var_floor: float = 1e-6

    @classmethod
    def estimate(
        cls,
        model: PreTrainedModel,
        tokenizer: PreTrainedTokenizerBase,
        texts: Sequence[str],
        *,
        layer_ids: Sequence[int] | None = None,
        location: str = "layer_input",
        pooling: Literal["mean", "last", "tokens"] = "tokens",
        exclude_first_n: int = 1,
        batch_size: int = 8,
        min_samples: int = 5000,
        session=None,
    ) -> "ActivationStats":
        """Estimate per-layer activation statistics over `texts`.

        Texts are tokenized in chunks of `batch_size`, captured at `location`, and pooled per
        `pooling`. Sequences with no real tokens contribute no samples, and duplicate entries in
        `layer_ids` are collapsed. Accumulation is streaming (per-chunk Welford merge) in float32
        on CPU, so the full feature matrix is never materialized. Capture runs under `torch.no_grad()` with
        `use_cache=False`, wrapped in `auxiliary_pass(aligned=True)`, so co-resident state
        controls' condition scorers, gates, and position counters ignore the passes while
        `"all"`-scoped behavior transforms still apply (the statistics measure the stream as
        deployed).

        Args:
            model: Model to capture activations from.
            tokenizer: Tokenizer for encoding `texts`.
            texts: Texts whose activations define the ambient distribution.
            layer_ids: Layers to estimate (0-based). None selects all decoder layers.
            location: Residual-stream boundary to measure at.
            pooling: `"tokens"` treats every retained real token position as one sample; `"mean"`
                and `"last"` pool per prompt first.
            exclude_first_n: Leading real token positions dropped per sequence under
                `pooling="tokens"`; ignored for `"mean"` and `"last"`.
            batch_size: Batch size for the capture forward passes.
            min_samples: Pooled-sample count below which a low-sample warning is emitted.

        Returns:
            An `ActivationStats` with float32 CPU tensors and the model's fingerprint recorded.

        Raises:
            ValueError: If `texts` is empty, `pooling` is unsupported, `exclude_first_n` is
                negative, `batch_size < 1`, a requested layer id is out of range, or no pooled
                samples remain after exclusion.

        Warns:
            UserWarning: When the accumulated pooled-sample `count` is below `min_samples`.
        """
        if len(texts) == 0:
            raise ValueError("texts must contain at least one text.")
        if pooling not in ("mean", "last", "tokens"):
            raise ValueError(f"pooling must be 'mean', 'last', or 'tokens', got {pooling!r}.")
        if exclude_first_n < 0:
            raise ValueError(f"exclude_first_n must be >= 0, got {exclude_first_n}.")
        if batch_size < 1:
            raise ValueError("batch_size must be >= 1.")

        device = next(model.parameters()).device
        model_type = getattr(model.config, "model_type", "unknown")
        fingerprint = model_fingerprint(model)

        target_layers: list[int] | None = None
        count = 0
        running_mean: dict[int, torch.Tensor] = {}
        running_m2: dict[int, torch.Tensor] = {}

        for start in range(0, len(texts), batch_size):
            chunk = list(texts[start:start + batch_size])
            enc = tokenize_texts(tokenizer, chunk, device)
            with auxiliary_pass(aligned=True):
                hidden, chunk_mask = capture_hidden(
                    enc, model=model, session=session, batch_size=len(chunk), location=location
                )

            if target_layers is None:
                num_layers = len(hidden)
                if layer_ids is None:
                    target_layers = sorted(hidden.keys())
                else:
                    target_layers = list(dict.fromkeys(int(lid) for lid in layer_ids))
                    for lid in target_layers:
                        if not 0 <= lid < num_layers:
                            raise ValueError(f"layer id {lid} out of range [0, {num_layers}).")

            mask = (
                chunk_mask.to("cpu", torch.bool)
                if chunk_mask is not None
                else torch.ones(hidden[target_layers[0]].shape[:2], dtype=torch.bool)
            )

            if pooling == "tokens":
                # rank real positions per row and drop the first exclude_first_n of them
                rank = mask.long().cumsum(dim=1)
                keep = mask & (rank > exclude_first_n)  # [B, T]
                row_keep = None
            else:
                keep = None
                row_keep = mask.any(dim=1)  # rows with no real tokens contribute no sample

            chunk_count = 0
            for lid in target_layers:
                states = hidden[lid].to(torch.float32)  # [B, T, H]
                if pooling == "tokens":
                    features = states[keep]  # [K, H]
                else:
                    row_states = states[row_keep]
                    row_mask = mask[row_keep]
                    if pooling == "mean":
                        features = masked_mean(row_states, row_mask.to(row_states.dtype))
                    else:  # last
                        positions = get_last_token_positions(row_mask.long(), row_states.size(1), row_states.size(0))
                        features = select_at_positions(row_states, positions)

                chunk_count = features.size(0)
                if chunk_count == 0:
                    break

                batch_mean = features.mean(dim=0)
                batch_m2 = (features - batch_mean).pow(2).sum(dim=0)
                if lid not in running_mean:
                    running_mean[lid] = batch_mean
                    running_m2[lid] = batch_m2
                else:
                    delta = batch_mean - running_mean[lid]
                    total = count + chunk_count
                    running_mean[lid] = running_mean[lid] + delta * (chunk_count / total)
                    running_m2[lid] = running_m2[lid] + batch_m2 + delta.pow(2) * (count * chunk_count / total)

            count += chunk_count

        if count == 0:
            raise ValueError(
                "No pooled samples were accumulated; every sequence was shorter than "
                f"exclude_first_n={exclude_first_n} real tokens."
            )
        if count < min_samples:
            suggestion = (
                "supply more texts"
                if pooling == "tokens"
                else "supply more texts, or use pooling='tokens' to draw one sample per token"
            )
            warnings.warn(
                f"ActivationStats accumulated {count} pooled samples, below min_samples={min_samples}. "
                f"Estimates of per-coordinate variance may be unstable; {suggestion}.",
                UserWarning,
                stacklevel=2,
            )

        var_floor = cls.var_floor
        variances = {
            lid: (running_m2[lid] / (count - 1) if count > 1 else torch.zeros_like(running_m2[lid]))
            .clamp_min(var_floor)
            for lid in running_mean
        }
        return cls(
            model_type=model_type,
            model_fingerprint=fingerprint,
            location=location,
            count=count,
            mean=running_mean,
            var=variances,
            var_floor=var_floor,
        )

    def standardize(self, features: torch.Tensor, layer_id: int) -> torch.Tensor:
        """Center and scale features per coordinate with this artifact's layer statistics.

        Args:
            features: Tensor of shape `[..., H]` in raw activation coordinates.
            layer_id: Layer whose statistics to apply.

        Returns:
            `(features - mean) / sqrt(var)` as a float32 tensor of the same shape.

        Raises:
            KeyError: If `layer_id` has no recorded statistics.
        """
        if layer_id not in self.mean:
            raise KeyError(
                f"ActivationStats has no statistics for layer {layer_id}; "
                f"available layers: {sorted(self.mean)}."
            )
        features = features.to(torch.float32)
        return (features - self.mean[layer_id]) / self.var[layer_id].sqrt()

    def save(self, dir_path: str | Path) -> None:
        """Save the artifact to a directory (safetensors for tensors, JSON sidecar for metadata).

        Args:
            dir_path: Directory to write to; created if missing.
        """
        dir_path = Path(dir_path)
        dir_path.mkdir(parents=True, exist_ok=True)

        tensors = {}
        for lid in sorted(self.mean):
            tensors[f"mean.{lid}"] = self.mean[lid].to(torch.float32).contiguous()
            tensors[f"var.{lid}"] = self.var[lid].to(torch.float32).contiguous()
        save_file(tensors, str(dir_path / _TENSOR_FILE))

        metadata = {
            "format_version": STATS_FORMAT_VERSION,
            "model_type": self.model_type,
            "model_fingerprint": self.model_fingerprint,
            "location": self.location,
            "count": self.count,
            "var_floor": self.var_floor,
            "layer_ids": sorted(self.mean),
        }
        with open(dir_path / _META_FILE, "w") as f:
            json.dump(metadata, f, indent=2, sort_keys=True)
        logger.debug("Saved ActivationStats to %s", dir_path)

    @classmethod
    def load(cls, dir_path: str | Path) -> "ActivationStats":
        """Load an artifact saved by `save`.

        Args:
            dir_path: Directory containing the safetensors file and JSON sidecar.

        Returns:
            The loaded `ActivationStats`.

        Raises:
            ValueError: If the artifact-format version is unsupported.
        """
        dir_path = Path(dir_path)
        with open(dir_path / _META_FILE) as f:
            metadata = json.load(f)

        version = metadata.get("format_version")
        if version != STATS_FORMAT_VERSION:
            raise ValueError(
                f"Unsupported ActivationStats format version {version!r}; expected {STATS_FORMAT_VERSION}."
            )

        tensors = load_file(str(dir_path / _TENSOR_FILE))
        mean: dict[int, torch.Tensor] = {}
        var: dict[int, torch.Tensor] = {}
        for lid in metadata["layer_ids"]:
            mean[int(lid)] = tensors[f"mean.{lid}"]
            var[int(lid)] = tensors[f"var.{lid}"]

        logger.debug("Loaded ActivationStats from %s with layers %s", dir_path, sorted(mean))
        return cls(
            model_type=metadata["model_type"],
            model_fingerprint=metadata["model_fingerprint"],
            location=metadata["location"],
            count=int(metadata["count"]),
            mean=mean,
            var=var,
            var_floor=float(metadata["var_floor"]),
        )

    def fingerprint(self) -> str:
        """Digest of this artifact (not of the model it was estimated on).

        Returns:
            A 16-character lowercase hex digest over the metadata and tensor contents.
        """
        digest = hashlib.sha256()
        header = json.dumps(
            {
                "model_type": self.model_type,
                "model_fingerprint": self.model_fingerprint,
                "location": self.location,
                "count": self.count,
                "var_floor": self.var_floor,
            },
            sort_keys=True,
        )
        digest.update(header.encode("utf-8"))
        for lid in sorted(self.mean):
            digest.update(str(lid).encode("utf-8"))
            digest.update(self.mean[lid].to(torch.float32).contiguous().numpy().tobytes())
            digest.update(self.var[lid].to(torch.float32).contiguous().numpy().tobytes())
        return digest.hexdigest()[:16]

count instance-attribute

location instance-attribute

mean instance-attribute

model_fingerprint instance-attribute

model_type instance-attribute

var instance-attribute

var_floor = 1e-06 class-attribute instance-attribute

estimate(model, tokenizer, texts, *, layer_ids=None, location='layer_input', pooling='tokens', exclude_first_n=1, batch_size=8, min_samples=5000, session=None) classmethod

Estimate per-layer activation statistics over texts.

Texts are tokenized in chunks of batch_size, captured at location, and pooled per pooling. Sequences with no real tokens contribute no samples, and duplicate entries in layer_ids are collapsed. Accumulation is streaming (per-chunk Welford merge) in float32 on CPU, so the full feature matrix is never materialized. Capture runs under torch.no_grad() with use_cache=False, wrapped in auxiliary_pass(aligned=True), so co-resident state controls' condition scorers, gates, and position counters ignore the passes while "all"-scoped behavior transforms still apply (the statistics measure the stream as deployed).

Parameters:

Name	Type	Description	Default
model	PreTrainedModel	Model to capture activations from.	required
tokenizer	PreTrainedTokenizerBase	Tokenizer for encoding texts.	required
texts	Sequence[str]	Texts whose activations define the ambient distribution.	required
layer_ids	Sequence[int] | None	Layers to estimate (0-based). None selects all decoder layers.	None
location	str	Residual-stream boundary to measure at.	'layer_input'
pooling	Literal['mean', 'last', 'tokens']	"tokens" treats every retained real token position as one sample; "mean" and "last" pool per prompt first.	'tokens'
exclude_first_n	int	Leading real token positions dropped per sequence under pooling="tokens"; ignored for "mean" and "last".	1
batch_size	int	Batch size for the capture forward passes.	8
min_samples	int	Pooled-sample count below which a low-sample warning is emitted.	5000

Returns:

Type	Description
ActivationStats	An ActivationStats with float32 CPU tensors and the model's fingerprint recorded.

Raises:

Type	Description
ValueError	If texts is empty, pooling is unsupported, exclude_first_n is negative, batch_size < 1, a requested layer id is out of range, or no pooled samples remain after exclusion.

Warns:

Type	Description
UserWarning	When the accumulated pooled-sample count is below min_samples.

Source code in steerability/algorithms/core/internals/stats.py

@classmethod
def estimate(
    cls,
    model: PreTrainedModel,
    tokenizer: PreTrainedTokenizerBase,
    texts: Sequence[str],
    *,
    layer_ids: Sequence[int] | None = None,
    location: str = "layer_input",
    pooling: Literal["mean", "last", "tokens"] = "tokens",
    exclude_first_n: int = 1,
    batch_size: int = 8,
    min_samples: int = 5000,
    session=None,
) -> "ActivationStats":
    """Estimate per-layer activation statistics over `texts`.

    Texts are tokenized in chunks of `batch_size`, captured at `location`, and pooled per
    `pooling`. Sequences with no real tokens contribute no samples, and duplicate entries in
    `layer_ids` are collapsed. Accumulation is streaming (per-chunk Welford merge) in float32
    on CPU, so the full feature matrix is never materialized. Capture runs under `torch.no_grad()` with
    `use_cache=False`, wrapped in `auxiliary_pass(aligned=True)`, so co-resident state
    controls' condition scorers, gates, and position counters ignore the passes while
    `"all"`-scoped behavior transforms still apply (the statistics measure the stream as
    deployed).

    Args:
        model: Model to capture activations from.
        tokenizer: Tokenizer for encoding `texts`.
        texts: Texts whose activations define the ambient distribution.
        layer_ids: Layers to estimate (0-based). None selects all decoder layers.
        location: Residual-stream boundary to measure at.
        pooling: `"tokens"` treats every retained real token position as one sample; `"mean"`
            and `"last"` pool per prompt first.
        exclude_first_n: Leading real token positions dropped per sequence under
            `pooling="tokens"`; ignored for `"mean"` and `"last"`.
        batch_size: Batch size for the capture forward passes.
        min_samples: Pooled-sample count below which a low-sample warning is emitted.

    Returns:
        An `ActivationStats` with float32 CPU tensors and the model's fingerprint recorded.

    Raises:
        ValueError: If `texts` is empty, `pooling` is unsupported, `exclude_first_n` is
            negative, `batch_size < 1`, a requested layer id is out of range, or no pooled
            samples remain after exclusion.

    Warns:
        UserWarning: When the accumulated pooled-sample `count` is below `min_samples`.
    """
    if len(texts) == 0:
        raise ValueError("texts must contain at least one text.")
    if pooling not in ("mean", "last", "tokens"):
        raise ValueError(f"pooling must be 'mean', 'last', or 'tokens', got {pooling!r}.")
    if exclude_first_n < 0:
        raise ValueError(f"exclude_first_n must be >= 0, got {exclude_first_n}.")
    if batch_size < 1:
        raise ValueError("batch_size must be >= 1.")

    device = next(model.parameters()).device
    model_type = getattr(model.config, "model_type", "unknown")
    fingerprint = model_fingerprint(model)

    target_layers: list[int] | None = None
    count = 0
    running_mean: dict[int, torch.Tensor] = {}
    running_m2: dict[int, torch.Tensor] = {}

    for start in range(0, len(texts), batch_size):
        chunk = list(texts[start:start + batch_size])
        enc = tokenize_texts(tokenizer, chunk, device)
        with auxiliary_pass(aligned=True):
            hidden, chunk_mask = capture_hidden(
                enc, model=model, session=session, batch_size=len(chunk), location=location
            )

        if target_layers is None:
            num_layers = len(hidden)
            if layer_ids is None:
                target_layers = sorted(hidden.keys())
            else:
                target_layers = list(dict.fromkeys(int(lid) for lid in layer_ids))
                for lid in target_layers:
                    if not 0 <= lid < num_layers:
                        raise ValueError(f"layer id {lid} out of range [0, {num_layers}).")

        mask = (
            chunk_mask.to("cpu", torch.bool)
            if chunk_mask is not None
            else torch.ones(hidden[target_layers[0]].shape[:2], dtype=torch.bool)
        )

        if pooling == "tokens":
            # rank real positions per row and drop the first exclude_first_n of them
            rank = mask.long().cumsum(dim=1)
            keep = mask & (rank > exclude_first_n)  # [B, T]
            row_keep = None
        else:
            keep = None
            row_keep = mask.any(dim=1)  # rows with no real tokens contribute no sample

        chunk_count = 0
        for lid in target_layers:
            states = hidden[lid].to(torch.float32)  # [B, T, H]
            if pooling == "tokens":
                features = states[keep]  # [K, H]
            else:
                row_states = states[row_keep]
                row_mask = mask[row_keep]
                if pooling == "mean":
                    features = masked_mean(row_states, row_mask.to(row_states.dtype))
                else:  # last
                    positions = get_last_token_positions(row_mask.long(), row_states.size(1), row_states.size(0))
                    features = select_at_positions(row_states, positions)

            chunk_count = features.size(0)
            if chunk_count == 0:
                break

            batch_mean = features.mean(dim=0)
            batch_m2 = (features - batch_mean).pow(2).sum(dim=0)
            if lid not in running_mean:
                running_mean[lid] = batch_mean
                running_m2[lid] = batch_m2
            else:
                delta = batch_mean - running_mean[lid]
                total = count + chunk_count
                running_mean[lid] = running_mean[lid] + delta * (chunk_count / total)
                running_m2[lid] = running_m2[lid] + batch_m2 + delta.pow(2) * (count * chunk_count / total)

        count += chunk_count

    if count == 0:
        raise ValueError(
            "No pooled samples were accumulated; every sequence was shorter than "
            f"exclude_first_n={exclude_first_n} real tokens."
        )
    if count < min_samples:
        suggestion = (
            "supply more texts"
            if pooling == "tokens"
            else "supply more texts, or use pooling='tokens' to draw one sample per token"
        )
        warnings.warn(
            f"ActivationStats accumulated {count} pooled samples, below min_samples={min_samples}. "
            f"Estimates of per-coordinate variance may be unstable; {suggestion}.",
            UserWarning,
            stacklevel=2,
        )

    var_floor = cls.var_floor
    variances = {
        lid: (running_m2[lid] / (count - 1) if count > 1 else torch.zeros_like(running_m2[lid]))
        .clamp_min(var_floor)
        for lid in running_mean
    }
    return cls(
        model_type=model_type,
        model_fingerprint=fingerprint,
        location=location,
        count=count,
        mean=running_mean,
        var=variances,
        var_floor=var_floor,
    )

fingerprint()

Digest of this artifact (not of the model it was estimated on).

Returns:

Type	Description
str	A 16-character lowercase hex digest over the metadata and tensor contents.

Source code in steerability/algorithms/core/internals/stats.py

def fingerprint(self) -> str:
    """Digest of this artifact (not of the model it was estimated on).

    Returns:
        A 16-character lowercase hex digest over the metadata and tensor contents.
    """
    digest = hashlib.sha256()
    header = json.dumps(
        {
            "model_type": self.model_type,
            "model_fingerprint": self.model_fingerprint,
            "location": self.location,
            "count": self.count,
            "var_floor": self.var_floor,
        },
        sort_keys=True,
    )
    digest.update(header.encode("utf-8"))
    for lid in sorted(self.mean):
        digest.update(str(lid).encode("utf-8"))
        digest.update(self.mean[lid].to(torch.float32).contiguous().numpy().tobytes())
        digest.update(self.var[lid].to(torch.float32).contiguous().numpy().tobytes())
    return digest.hexdigest()[:16]

load(dir_path) classmethod

Load an artifact saved by save.

Parameters:

Name	Type	Description	Default
dir_path	str | Path	Directory containing the safetensors file and JSON sidecar.	required

Returns:

Type	Description
ActivationStats	The loaded ActivationStats.

Raises:

Type	Description
ValueError	If the artifact-format version is unsupported.

Source code in steerability/algorithms/core/internals/stats.py

@classmethod
def load(cls, dir_path: str | Path) -> "ActivationStats":
    """Load an artifact saved by `save`.

    Args:
        dir_path: Directory containing the safetensors file and JSON sidecar.

    Returns:
        The loaded `ActivationStats`.

    Raises:
        ValueError: If the artifact-format version is unsupported.
    """
    dir_path = Path(dir_path)
    with open(dir_path / _META_FILE) as f:
        metadata = json.load(f)

    version = metadata.get("format_version")
    if version != STATS_FORMAT_VERSION:
        raise ValueError(
            f"Unsupported ActivationStats format version {version!r}; expected {STATS_FORMAT_VERSION}."
        )

    tensors = load_file(str(dir_path / _TENSOR_FILE))
    mean: dict[int, torch.Tensor] = {}
    var: dict[int, torch.Tensor] = {}
    for lid in metadata["layer_ids"]:
        mean[int(lid)] = tensors[f"mean.{lid}"]
        var[int(lid)] = tensors[f"var.{lid}"]

    logger.debug("Loaded ActivationStats from %s with layers %s", dir_path, sorted(mean))
    return cls(
        model_type=metadata["model_type"],
        model_fingerprint=metadata["model_fingerprint"],
        location=metadata["location"],
        count=int(metadata["count"]),
        mean=mean,
        var=var,
        var_floor=float(metadata["var_floor"]),
    )

save(dir_path)

Save the artifact to a directory (safetensors for tensors, JSON sidecar for metadata).

Parameters:

Name	Type	Description	Default
dir_path	str | Path	Directory to write to; created if missing.	required

Source code in steerability/algorithms/core/internals/stats.py

def save(self, dir_path: str | Path) -> None:
    """Save the artifact to a directory (safetensors for tensors, JSON sidecar for metadata).

    Args:
        dir_path: Directory to write to; created if missing.
    """
    dir_path = Path(dir_path)
    dir_path.mkdir(parents=True, exist_ok=True)

    tensors = {}
    for lid in sorted(self.mean):
        tensors[f"mean.{lid}"] = self.mean[lid].to(torch.float32).contiguous()
        tensors[f"var.{lid}"] = self.var[lid].to(torch.float32).contiguous()
    save_file(tensors, str(dir_path / _TENSOR_FILE))

    metadata = {
        "format_version": STATS_FORMAT_VERSION,
        "model_type": self.model_type,
        "model_fingerprint": self.model_fingerprint,
        "location": self.location,
        "count": self.count,
        "var_floor": self.var_floor,
        "layer_ids": sorted(self.mean),
    }
    with open(dir_path / _META_FILE, "w") as f:
        json.dump(metadata, f, indent=2, sort_keys=True)
    logger.debug("Saved ActivationStats to %s", dir_path)

standardize(features, layer_id)

Center and scale features per coordinate with this artifact's layer statistics.

Parameters:

Name	Type	Description	Default
features	Tensor	Tensor of shape [..., H] in raw activation coordinates.	required
layer_id	int	Layer whose statistics to apply.	required

Returns:

Type	Description
Tensor	(features - mean) / sqrt(var) as a float32 tensor of the same shape.

Raises:

Type	Description
KeyError	If layer_id has no recorded statistics.

Source code in steerability/algorithms/core/internals/stats.py

def standardize(self, features: torch.Tensor, layer_id: int) -> torch.Tensor:
    """Center and scale features per coordinate with this artifact's layer statistics.

    Args:
        features: Tensor of shape `[..., H]` in raw activation coordinates.
        layer_id: Layer whose statistics to apply.

    Returns:
        `(features - mean) / sqrt(var)` as a float32 tensor of the same shape.

    Raises:
        KeyError: If `layer_id` has no recorded statistics.
    """
    if layer_id not in self.mean:
        raise KeyError(
            f"ActivationStats has no statistics for layer {layer_id}; "
            f"available layers: {sorted(self.mean)}."
        )
    features = features.to(torch.float32)
    return (features - self.mean[layer_id]) / self.var[layer_id].sqrt()

StatsSpec dataclass

Deferred recipe for ActivationStats; holds every argument except the model.

Attributes:

Name	Type	Description
texts	Sequence[str]	Texts whose activations define the ambient distribution.
layer_ids	Sequence[int] | None	Layers to estimate (0-based). None selects all decoder layers.
location	str	Residual-stream boundary the statistics are measured at.
pooling	Literal['mean', 'last', 'tokens']	Sample construction. "tokens" treats every retained real token position as one sample; "mean" and "last" pool per prompt first.
exclude_first_n	int	Number of leading real token positions dropped per sequence under pooling="tokens"; ignored for "mean" and "last".
batch_size	int	Batch size for the capture forward passes.

Source code in steerability/algorithms/core/internals/stats.py

@dataclass
class StatsSpec:
    """Deferred recipe for `ActivationStats`; holds every argument except the model.

    Attributes:
        texts: Texts whose activations define the ambient distribution.
        layer_ids: Layers to estimate (0-based). None selects all decoder layers.
        location: Residual-stream boundary the statistics are measured at.
        pooling: Sample construction. `"tokens"` treats every retained real token position as one
            sample; `"mean"` and `"last"` pool per prompt first.
        exclude_first_n: Number of leading real token positions dropped per sequence under
            `pooling="tokens"`; ignored for `"mean"` and `"last"`.
        batch_size: Batch size for the capture forward passes.
    """

    texts: Sequence[str]
    layer_ids: Sequence[int] | None = None  # None selects all decoder layers
    location: str = "layer_input"
    pooling: Literal["mean", "last", "tokens"] = "tokens"
    exclude_first_n: int = 1
    batch_size: int = 8

    def estimate(
        self, model: PreTrainedModel | None, tokenizer: PreTrainedTokenizerBase, session=None,
    ) -> "ActivationStats":
        """Estimate `ActivationStats` on `model` with this recipe's settings."""
        return ActivationStats.estimate(
            model,
            tokenizer,
            self.texts,
            session=session,
            layer_ids=self.layer_ids,
            location=self.location,
            pooling=self.pooling,
            exclude_first_n=self.exclude_first_n,
            batch_size=self.batch_size,
        )

batch_size = 8 class-attribute instance-attribute

exclude_first_n = 1 class-attribute instance-attribute

layer_ids = None class-attribute instance-attribute

location = 'layer_input' class-attribute instance-attribute

pooling = 'tokens' class-attribute instance-attribute

texts instance-attribute

estimate(model, tokenizer, session=None)

Estimate ActivationStats on model with this recipe's settings.

Source code in steerability/algorithms/core/internals/stats.py

def estimate(
    self, model: PreTrainedModel | None, tokenizer: PreTrainedTokenizerBase, session=None,
) -> "ActivationStats":
    """Estimate `ActivationStats` on `model` with this recipe's settings."""
    return ActivationStats.estimate(
        model,
        tokenizer,
        self.texts,
        session=session,
        layer_ids=self.layer_ids,
        location=self.location,
        pooling=self.pooling,
        exclude_first_n=self.exclude_first_n,
        batch_size=self.batch_size,
    )

measure_residual_norms(model, tokenizer, layer_ids, prompts, *, location='layer_output', stat='median', prompt_format='chat_prompt', batch_size=8)

Typical per-token residual-stream norm at each requested layer boundary.

The statistic is aggregated over the real (non-pad) tokens of prompts at the residual-stream boundary named by location. Residual norms grow with depth, so the returned per-layer norms let a caller rescale a direction to a fixed fraction of the norm at each layer (see SteeringVector.scaled_to_norms).

Prompts are rendered via render_for_model(tokenizer, prompt=p, mode=prompt_format). The default prompt_format="chat_prompt" renders with add_generation_prompt=True, matching what inference produces for a prompt. The rendered string is tokenized with add_special_tokens=False whenever a chat template was applied, matching how the estimators consume contrastive text.

Parameters:

Name	Type	Description	Default
model	PreTrainedModel	Model to extract hidden states from.	required
tokenizer	PreTrainedTokenizerBase	Tokenizer for rendering and encoding the prompts.	required
layer_ids	Sequence[int] | None	Layers to measure (0-based). None measures every layer (no extra cost, since all layers are extracted regardless).	required
prompts	Sequence[str]	Calibration prompts. Use fit-distribution, inference-style prompts rather than a held-out evaluation set.	required
location	HiddenStateLocation	Residual-stream boundary each layer key maps to. "layer_output" (default) measures the output of each layer; "layer_input" measures the input of each layer, the boundary a layer pre-hook observes.	'layer_output'
stat	Literal['median', 'mean']	Aggregation over the pooled real-token norms, "median" (default) or "mean".	'median'
prompt_format	PromptFormat	How each prompt is rendered into model-ready text (via render_for_model).	'chat_prompt'
batch_size	int	Batch size for the extraction forward passes.	8

Returns:

Type	Description
dict[int, float]	Mapping from layer id to the aggregated per-token residual norm, as a Python float.

Raises:

Type	Description
ValueError	If prompts is empty, stat is unsupported, or any requested layer id is out of range.

Source code in steerability/algorithms/core/internals/stats.py

@torch.no_grad()
def measure_residual_norms(
    model: PreTrainedModel,
    tokenizer: PreTrainedTokenizerBase,
    layer_ids: Sequence[int] | None,
    prompts: Sequence[str],
    *,
    location: HiddenStateLocation = "layer_output",
    stat: Literal["median", "mean"] = "median",
    prompt_format: PromptFormat = "chat_prompt",
    batch_size: int = 8,
) -> dict[int, float]:
    """Typical per-token residual-stream norm at each requested layer boundary.

    The statistic is aggregated over the real (non-pad) tokens of `prompts` at the residual-stream
    boundary named by `location`. Residual norms grow with depth, so the returned per-layer norms
    let a caller rescale a direction to a fixed fraction of the norm at each layer (see
    `SteeringVector.scaled_to_norms`).

    Prompts are rendered via `render_for_model(tokenizer, prompt=p, mode=prompt_format)`. The default
    `prompt_format="chat_prompt"` renders with `add_generation_prompt=True`, matching what inference
    produces for a prompt. The rendered string is tokenized with `add_special_tokens=False` whenever
    a chat template was applied, matching how the estimators consume contrastive text.

    Args:
        model: Model to extract hidden states from.
        tokenizer: Tokenizer for rendering and encoding the prompts.
        layer_ids: Layers to measure (0-based). `None` measures every layer (no extra cost, since
            all layers are extracted regardless).
        prompts: Calibration prompts. Use fit-distribution, inference-style prompts rather than a
            held-out evaluation set.
        location: Residual-stream boundary each layer key maps to. `"layer_output"` (default)
            measures the output of each layer; `"layer_input"` measures the input of each layer, the
            boundary a layer pre-hook observes.
        stat: Aggregation over the pooled real-token norms, `"median"` (default) or `"mean"`.
        prompt_format: How each prompt is rendered into model-ready text (via `render_for_model`).
        batch_size: Batch size for the extraction forward passes.

    Returns:
        Mapping from layer id to the aggregated per-token residual norm, as a Python `float`.

    Raises:
        ValueError: If `prompts` is empty, `stat` is unsupported, or any requested layer id is out of
            range.
    """
    if len(prompts) == 0:
        raise ValueError("prompts must contain at least one prompt.")
    if stat not in ("median", "mean"):
        raise ValueError(f"stat must be 'median' or 'mean', got {stat!r}.")

    device = next(model.parameters()).device

    # render then tokenize with add_special_tokens=False when a template was applied
    rendered = [render_for_model(tokenizer, prompt=p, mode=prompt_format) for p in prompts]
    template_applied = has_chat_template(tokenizer) and prompt_format != "raw"
    enc = tokenize_texts(tokenizer, rendered, device, add_special_tokens=not template_applied)

    hidden = layerwise_tokenwise_hidden(model, enc, batch_size=batch_size, location=location)
    num_layers = len(hidden)

    if layer_ids is None:
        target_layers = sorted(hidden.keys())
    else:
        target_layers = [int(lid) for lid in layer_ids]
        for lid in target_layers:
            if not 0 <= lid < num_layers:
                raise ValueError(f"layer id {lid} out of range [0, {num_layers}).")

    # real-token mask
    attention_mask = enc.get("attention_mask")
    if attention_mask is not None:
        keep = attention_mask.to("cpu", torch.bool)  # [N, T]
    else:
        keep = None

    norms: dict[int, float] = {}
    for lid in target_layers:
        per_token = hidden[lid].to(torch.float32).norm(dim=-1)  # [N, T]
        values = per_token[keep] if keep is not None else per_token.reshape(-1)
        agg = values.median() if stat == "median" else values.mean()
        norms[lid] = float(agg)

    return norms