Steps — Sae

Sparse-Autoencoder steps for Murano.

This module lets you encode prompts through a pre-trained SAE loaded from HuggingFace via sae-lens and then read the resulting features two ways.

Building blocks

:class:SAEModel wraps the underlying sae-lens SAE with lazy loading. Its one user-facing method is feature_direction(feature_id) (the SAE’s decoder row for that feature, used for steering and labeling); encode and n_features are plumbing that :class:SAEEncode drives internally.
:func:top_sae_features_per_prompt is a plain helper that picks each prompt’s top firing feature(s), either at the last token (reduce="last") or by mean over the sentence’s content tokens (reduce="mean"), filtering out the BOS-anchored sink and broad high-frequency features.
:func:sae_steer is a plain helper that builds an :class:~murano.steps.intervene.Intervene step which adds a single feature’s decoder direction to the residual stream at the exact site the SAE was trained on, so steering works for any SAE hook point without the caller picking a layer or module. It reuses the existing steering machinery (:func:~murano.steps.intervene.steer_direction) rather than adding a new one.

Pipeline steps

:class:SAEEncode runs prompts through the SAE and writes :class:SAEActivationStore ([N, seq, n_features] activations).
:class:SAETopActivations is the input-side interpretation: given feature ids, it returns the prompt contexts where each feature fires hardest as :class:SAEFeatureExamples.
:class:SAEFeatureLabel is the output-side interpretation: given feature ids, it projects each feature’s decoder direction through the model’s final norm + unembedding via model.project_on_vocab (the same ln_final + lm_head projection :class:~murano.steps.logit_lens.LogitLens applies to per-layer residuals) and returns the top tokens each feature promotes in next-token space, as :class:SAEFeatureLabels.

Typical flow::

results = Pipeline([
    LoadPrompts(prompts),
    SAEEncode(model, release=..., sae_id=...),     # -> sae_record
]).run()
feats = top_sae_features_per_prompt(results["sae_record"], n=1)
feat_ids = sorted({f for prompt in feats for f in prompt})

results = Pipeline([
    SAEFeatureLabel(model, feat_ids=feat_ids, k_tokens=5),
]).run(results)                                     # -> feature_labels

`SAEActivationStore`

Per-token SAE encodings of one layer’s residual stream.

Attributes:

activations: Tensor [N, seq, n_features] of SAE encoder outputs.
tokens: Tensor [N, seq] of input token ids.
attention_mask: Tensor [N, seq] marking real (1) vs padding (0) positions.
texts: Input texts paired by index with the N dimension.
hook: Component address (:class:Node) the SAE was applied at. Feature
indices are a separate space and stay plain ints.
release: HuggingFace SAE release identifier.
sae_id: SAE id within the release.
n_features: SAE width (equals activations.shape[-1]).

`init`

def __init__(self, activations: Tensor, tokens: Tensor, attention_mask: Tensor, texts: list[str], hook: Node, release: str, sae_id: str, n_features: int) -> None:

`top_sae_features_per_prompt`

def top_sae_features_per_prompt(record: SAEActivationStore, n: int = 1, reduce: str = 'last', sink_position: int = 1, max_density: float = 0.4) -> list[list[int]]:

Top SAE features per prompt, excluding BOS sinks.

Two ways to summarize a prompt into one feature vector before ranking, selected by reduce:

"last" (default): the activations at the prompt’s last real token. For a fact-completion prompt this is the concept the model has committed to at the end (often a local “what comes next” feature, e.g. a generic “city” feature for “…located in Paris”).
"mean": the mean activation across the prompt’s content tokens (real tokens, skipping the BOS-sink region). This asks what the sentence is about as a whole. Best with a few diverse prompts: a concept feature then fires on only its own prompt and stands out, whereas broad high-frequency features fire everywhere and are dropped (see max_density).

Residual-stream SAEs reliably learn a sink feature that fires huge at <bos> and the first content token. "last" removes those by checking each feature’s global-peak position (<= sink_position); "mean" removes them by excluding the first sink_position + 1 positions from the average. Both default to treating positions 0-1 as the sink region, which relies on right padding (enforced by :class:SAEEncode).

Args:

record: SAEActivationStore from a prior :class:SAEEncode run.
n: Features to keep per prompt. Must be in [1, n_features].
reduce: "last" or "mean" (see above).
sink_position: Index of the last position treated as the BOS-sink
region. Defaults to 1 (positions 0-1).
max_density: Only used when reduce="mean". Features active on more
than this fraction of the batch’s real tokens are dropped as broad,
non-specific features (these are usually high-frequency, poorly
interpretable features that otherwise dominate the average). Set to
1.0 to keep all features.

Returns:

[[feat_ids_for_prompt_0], ...]. Each inner list holds up to n
feature ids; it is shorter (possibly empty) for a prompt whose top
features were all removed as sinks or as broad features.

Raises:

ValueError: If n is outside [1, n_features], reduce is not
"last"/"mean", or any prompt has no real tokens.

`top_sae_features_for_tokens`

def top_sae_features_for_tokens(record: SAEActivationStore, model: ModelBackend, target_tokens: set[str], n: int = 5, skip_positions: int = 2) -> list[int]:

Top SAE features by mean activation on a set of target tokens.

Unlike :func:top_sae_features_per_prompt (which reads each prompt’s last token), this scans every position whose decoded token is in target_tokens and ranks features by their mean activation there. Use it to discover the feature for a concept that appears mid-sentence, e.g. to find a “Golden Gate Bridge” feature from prompts that mention it.

The first skip_positions positions of every prompt are excluded to avoid the BOS sink (see :func:top_sae_features_per_prompt). Because it decodes token ids, this helper takes model where :func:top_sae_features_per_prompt does not.

Args:

record: SAEActivationStore from a prior :class:SAEEncode run.
model: Model whose tokenizer decodes token ids.
target_tokens: Decoded, stripped token strings to match, e.g.
{"Golden", "Gate", "Bridge"}.
n: Number of features to return. Must be >= 1.
skip_positions: Leading positions to exclude per prompt to avoid the
BOS sink. Defaults to 2 (<bos> plus the first content token).

Returns:

The n feature ids with the highest mean activation across the
matched positions, highest first.

Raises:

ValueError: If n < 1 or no position matches target_tokens.

`SAEFeatureLabels`

Top-promoted vocabulary tokens per SAE feature.

Produced by projecting each feature’s decoder direction through the model’s final norm and unembedding (the standard logit-lens shortcut). Gives a fast, corpus-free semantic handle on what concept each feature carries: the tokens the feature pushes the model toward.

This is an approximation: it skips the intermediate layers between the SAE’s layer and the unembedding. For ground-truth input-side labels (the contexts that make the feature fire), pair with :class:SAETopActivations.

Attributes:

feat_ids: SAE feature indices included.
tokens: {feat_id: list[str]} top-K promoted tokens per feature.
logits: {feat_id: list[float]} top-K logit scores per feature.
k_tokens: Cap on tokens kept per feature.
layer: Layer the SAE was trained on.
release: HuggingFace SAE release identifier.
sae_id: SAE id within the release.

`init`

def __init__(self, feat_ids: list[int], tokens: dict[int, list[str]], logits: dict[int, list[float]], k_tokens: int, layer: int, release: str, sae_id: str) -> None:

`SAEFeatureExamples`

Top-K activating contexts per SAE feature, sorted descending by activation.

Attributes:

feat_ids: Features included.
contexts: {feat_id: list[str]} up to K context strings per feature
(fewer if the batch has fewer non-padding positions).
tokens: {feat_id: list[str]} up to K triggering tokens per feature.
act_vals: {feat_id: list[float]} up to K activation values per feature.
hook: Component address (:class:Node) the SAE was applied at.
release: HuggingFace SAE release identifier.
sae_id: SAE id within the release.
k: Top-K cap per feature.

`init`

def __init__(self, feat_ids: list[int], contexts: dict[int, list[str]], tokens: dict[int, list[str]], act_vals: dict[int, list[float]], hook: Node, release: str, sae_id: str, k: int) -> None:

`SAEModel`

Loaded SAE encoder, applied to a model’s residual stream.

Weights are pulled lazily from HuggingFace via sae-lens on first use. Sharing one instance across pipelines avoids repeated loads.

Requires the [sae] extra: pip install -e ".[sae]".

Attributes:

release: sae-lens release identifier.
sae_id: SAE id within the release.
device: Torch device the encoder runs on.

`init`

def __init__(self, release: str, sae_id: str, device: str = 'cpu'):

`encode`

def encode(self, residual: Tensor) -> Tensor:

Encode residual [N, seq, d_model] to SAE codes [N, seq, n_features].

`feature_direction`

def feature_direction(self, feature_id: int) -> Tensor:

Decoder direction ([d_model]) for a single SAE feature.

For steering prefer :func:sae_steer, which adds this direction at the exact site the SAE was trained on.

Args:

feature_id: SAE feature index in [0, n_features).

Returns:

Detached copy of the decoder row, on the SAE’s device.

`sae_steer`

def sae_steer(model: ModelBackend, sae_model: SAEModel, feature_id: int, alpha: float, gen_kwargs: dict | None = None) -> Intervene:

Build an :class:~murano.steps.intervene.Intervene step that steers a generation with a single SAE feature.

Adds alpha * feature_direction(feature_id) to the residual stream at the exact site the SAE was trained on, so the same call works whether the SAE is resid_post, mlp_out, attn_out or resid_pre without the caller having to choose a layer or module (which is easy to get wrong: a hardcoded modules="residual" only matches resid_post SAEs). The direction is normalized to unit norm internally (see :func:~murano.steps.intervene.steer_direction), so alpha is the absolute magnitude added; positive enhances the feature, negative suppresses it.

alpha is therefore scale-dependent, and it has no default worth quoting: the edit is applied at every decoded token, so what matters is its size relative to the residual stream at the steering site, which ranges from tens to thousands across models. Measure that norm and start near a tenth of it. Pushing much past it swamps the residual stream and the model repeats one token; see notebooks/applications/sae_steering.ipynb, where no strength both invokes the concept and preserves the text. Adding a fixed magnitude is only one way to steer a feature; clamping the feature’s own activation (Scaling Monosemanticity) leaves the rest of the residual stream intact and is not implemented here.

The returned step generates a clean (unsteered) and a steered response for each prompt and writes an :class:~murano.steps.intervene.InterveneResult::

results = Pipeline([
    LoadPrompts(prompts),
    sae_steer(model, encode.sae_model, feature_id, alpha=alpha),
]).run()
comparison = results["intervene"]    # clean vs steered

Args:

model: Model to generate with (must match the SAE’s base model).
sae_model: Loaded SAE providing the feature direction and hook point.
feature_id: SAE feature index in [0, n_features).
alpha: Steering strength, as an absolute magnitude added to the residual
stream; positive enhances, negative suppresses. Calibrate against the
residual norm at the steering site.
gen_kwargs: Generation kwargs forwarded to the model. Defaults to
```{“max_new_tokens”`: 256, “do_sample”: False}“.

Returns:

A configured: class:~murano.steps.intervene.Intervene step.

Raises:

ValueError: If the resolved steering layer is out of bounds for
model, or for a resid_pre SAE at layer 0.
NotImplementedError: For SAE hook kinds with no clean steering site
(e.g. resid_mid).

`SAEEncode`

Encode residual-stream activations through an SAE loaded from HuggingFace.

Constructs an SAEModel internally from release + sae_id and auto-detects the target layer and hook point from the SAE’s own config, so the same step works against any sae-lens release without the caller having to know its training-time conventions. The loaded SAE is reachable via self.sae_model for reuse or inspection.

Supported hook points: resid_pre, resid_post, resid_mid, mlp_out, attn_out. Per-head hook_z SAEs are not handled because reshape semantics vary per release.

Reads from results:

results['prompts']: PromptBatch

Writes to results:

results['sae_record']: SAEActivationStore

Args:

model: MuranoModel to record from.
release: HuggingFace SAE release identifier.
sae_id: SAE id within the release.
max_length: Truncate prompts to this many tokens before tracing.
None disables truncation (default; suitable for short prompts).

Raises:

ValueError: If the SAE’s hook layer is out of bounds for model,
or the SAE layer cannot be determined from the config or sae_id.
NotImplementedError: If the SAE was trained on an unsupported hook
point (e.g. hook_z).

`init`

def __init__(self, model: ModelBackend, release: str, sae_id: str, max_length: int | None = None):

`SAETopActivations`

Input-side interpretation: rank the top-K activating contexts per SAE feature.

For each feature, scans every (text, token) position in the SAEActivationStore and keeps the K positions with the largest activation. Padded tokens are excluded. BOS-token positions are excluded by default, since residual-stream SAEs tend to develop strong BOS-anchored features that dominate the top-K and crowd out content-bearing features.

Reads from results:

results['sae_record']: SAEActivationStore

Writes to results:

results['feature_examples']: SAEFeatureExamples

Args:

model: MuranoModel, used to decode triggering tokens.
k: Number of top contexts per feature; must be >= 1.
feat_ids: Specific features to rank. None ranks every feature.
skip_bos: If True, mask out positions whose token id equals the
tokenizer’s bos_token_id before ranking. Has no effect when
the tokenizer has no BOS token.

Raises:

ValueError: If k < 1, the SAE activations are not
[N, seq, n_features]-shaped, or any requested feat_id
is out of range.

`init`

def __init__(self, model: ModelBackend, k: int = 10, feat_ids: list[int] | None = None, skip_bos: bool = True):

`SAEFeatureLabel`

Output-side interpretation: label SAE features by the tokens they promote.

For each requested feature, projects W_dec[feature_id] through the model’s standardized final norm and unembedding via model.project_on_vocab (the same logit-lens projection :class:~murano.steps.logit_lens.LogitLens applies to per-layer residuals). The top-K tokens by resulting logit are returned as a label per feature.

This is the fast, corpus-free counterpart to :class:SAETopActivations. Pair the two when possible: SAETopActivations tells you what makes a feature fire (input side), this tells you what the feature pushes the model toward (output side).

Reads from results:

results['sae_record']: SAEActivationStore. The SAE’s release
and sae_id are taken from here so the same SAE used for
encoding is used for labeling.

Writes to results:

results['feature_labels']: SAEFeatureLabels

Args:

model: Model providing project_on_vocab (final norm + unembedding).
feat_ids: SAE feature indices to label.
k_tokens: How many top-promoted tokens to keep per feature.
Must be >= 1.
sae_model: Optional pre-built :class:SAEModel to source decoder
directions from. When None (default), one is constructed
from sae_record.release and sae_record.sae_id. Pass an
existing instance to avoid reloading when chaining after
“: class:SAEEncode.

Raises:

ValueError: If k_tokens < 1 or any requested feat_id is out
of range for the SAE width recorded in sae_record.

`init`

def __init__(self, model: ModelBackend, feat_ids: list[int], k_tokens: int = 5, sae_model: SAEModel | None = None):