Skip to content
Murano

Steps — Attention

Attention-pattern analysis and intervention.

Captures the per-head softmax attention weights [batch, n_heads, query, key] that :class:~murano.steps.record.Record cannot reach (it captures head outputs, not the weights), and offers task-agnostic readouts over them: entropy, attention-sink mass, mean attention distance, and the attention along a fixed query-to-key offset. These are the generic building blocks; task-specific interpretation (labelling induction, name-mover, or duplicate-token heads) is left to the caller, who picks the offsets and applies the labels.

The module also exposes the head’s OV map (:func:ov_circuit) and an intervention step (:class:AblateAttention) that overwrites attention weights through nnterp’s settable accessor. All of it requires a model loaded with enable_attention_probs=True (eager attention); the steps fail with a clear message otherwise.

Attributes:

  • RecordAttention: Capture attention patterns for chosen layers.
  • AttentionResult: Captured patterns plus the generic reduction methods.
  • AblateAttention: Overwrite attention weights and read the resulting logits.
  • ov_circuit: The effective per-head value-output map W_O_h @ W_V_h.

AttentionResult

Section titled “AttentionResult”

Captured per-head attention weights and the generic readouts over them.

The reduction methods each return a [n_captured_layers, n_heads] tensor, rows in :attr:layers order, averaged over the batch and valid query positions.

Attributes:

  • patterns: {layer: tensor} of attention weights “[batch, n_heads,
  • query, key]“ on CPU.
  • attention_mask: Tensor [batch, seq] marking real (1) vs padding (0)
  • positions, used to exclude padding from the reductions.
  • str_tokens: Decoded input tokens, shape [batch][seq], for labelling a
  • pattern heatmap.
  • layers: The captured layer indices, in row order of every reduction.
  • addresses: One :class:Node per captured layer (self_attn nodes).
  • metadata: Free-form provenance (e.g. the prompts key read).

entropy

Section titled “entropy”
def entropy(self) -> Tensor:

Return the per-head mean attention entropy [n_layers, n_heads].

sink

Section titled “sink”
def sink(self, index: int = 0) -> Tensor:

Return the per-head mean attention mass on a sink key [n_layers, n_heads].

index counts from each example’s first real token (default 0, the usual attention sink), so padding never shifts the sink column.

distance

Section titled “distance”
def distance(self) -> Tensor:

Return the per-head mean attention distance [n_layers, n_heads].

at_offset

Section titled “at_offset”
def at_offset(self, offset: int) -> Tensor:

Return the per-head mean attention at key - query == offset [n_layers, n_heads].

attention_to

Section titled “attention_to”
def attention_to(self, query = None, key = None) -> Tensor:

Return per-head mean attention from query to key [n_layers, n_heads].

Averages pattern[query, key] over the batch for each head, so it reads “how much does this head, at the query position, attend to the key position.” Both positions accept the :func:~murano.steps.metrics._answer_positions form (an int, a per-example sequence or tensor, negatives from the end); None uses each example’s last real token (the natural query for next-token prediction). Explicit positive positions are absolute columns, so under Murano’s default left padding use -1 for the last real token and prefer negative indices.

head_pattern

Section titled “head_pattern”
def head_pattern(self, layer: int, head: int, index: int = 0) -> Tensor:

Return one head’s attention matrix [query, key] for input index.

__init__

Section titled “__init__”
def __init__(self, patterns: dict[int, Tensor], attention_mask: Tensor, str_tokens: list[list[str]], layers: list[int], addresses: list[Node], metadata: dict[str, Any] = dict()) -> None:

RecordAttention

Section titled “RecordAttention”

Capture per-head attention weights across layers.

Runs one trace over the batched prompts and stores each requested layer’s [batch, n_heads, query, key] softmax weights in an :class:AttentionResult, ready for the generic reductions or a pattern heatmap. Requires a model loaded with enable_attention_probs=True.

Reads from results:

  • results[prompts_key]: PromptBatch (default prompts).

Writes to results:

  • results[output_key]: AttentionResult.

Args:

  • model: Model backend loaded with attention weights enabled.
  • layers: Layer indices to capture, or "all" for every layer.
  • prompts_key: Results key to read the prompt batch from.
  • output_key: Results key to write the AttentionResult under.

Raises:

  • ValueError: If layers is a string other than "all".

Note:

  • Every captured layer’s full [batch, n_heads, seq, seq] pattern is
  • held in memory, so capturing all layers on a large model or long inputs
  • is memory-bound; restrict layers or the prompt count when it is.

__init__

Section titled “__init__”
def __init__(self, model: ModelBackend, layers: list[int] | str = 'all', prompts_key: str = keys.PROMPTS, output_key: str = keys.ATTENTION_PATTERN):

ov_circuit

Section titled “ov_circuit”
def ov_circuit(model: ModelBackend, layer: int, head: int) -> Tensor:

Return the effective per-head OV map W_O_h @ W_V_h, shape [d_model, d_model].

The OV (value-output) circuit is the linear transform a head applies to the residual stream it reads before writing it back, ignoring the attention weighting. Composing it with the unembedding (a caller’s step) gives the head’s copy/writing behavior. The projection bias, when present, is omitted.

Works for separate q/k/v projections (grouped-query aware, mapping the query head to its shared key/value head) and for GPT-2’s fused c_attn. GPT-NeoX style interleaved query_key_value is not supported and raises.

Args:

  • model: Model backend to read weights from.
  • layer: Decoder layer index.
  • head: Query-head index in [0, n_heads).

Returns:

  • The OV matrix on CPU as float32, mapping a residual vector x to the
  • head’s output contribution OV @ x.

Raises:

  • ValueError: If head is out of range.
  • NotImplementedError: If the value weights cannot be resolved for the
  • architecture.

AblateAttention

Section titled “AblateAttention”

Overwrite per-head attention weights and read the resulting logits.

Runs one intervened forward pass: the addressed heads’ attention weights are zeroed, replaced by the batch-mean pattern, or resampled from a second run, and the model’s output logits are stored. Comparing them against a clean :class:~murano.steps.logits.Logits run with a metric step scores how much those heads’ attention pattern carries the behavior. Requires a model loaded with enable_attention_probs=True.

Zeroing a head’s weights makes it read nothing (it writes a zero-weighted sum of values, i.e. nothing); mean/resample keep the rows normalized. The edit is restricted to the query rows in positions (every row by default); the full key axis of each edited row is replaced.

Reads from results:

  • results[prompts_key]: PromptBatch (default prompts), the batch to run.
  • results[source_key]: PromptBatch, only with source_key set: the
  • resample source whose patterns are injected.

Writes to results:

  • results[logits_key]: Tensor [B, S, vocab] of the intervened logits.
  • results[mask_key]: Tensor [B, S] marking real (1) vs padding (0).

Args:

  • model: Model backend loaded with attention weights enabled.
  • targets: Heads to ablate: a :class:NodeSet, a single address, or an
  • iterable. Each must be an attention node “Node(layer, “self_attn”,
  • head=h); a head-less Node(layer, “self_attn”)“ selects every head
  • at that layer.
  • method: "zero", "mean", or "resample".
  • source: For method="resample", corrupted prompts paired one-to-one
  • with the inputs, token-length-matched so positions align.
  • source_key: For method="resample", a Results key naming the source
  • “: class:~murano.artifacts.PromptBatch; mutually exclusive with
  • source.
  • positions: Query position(s) to overwrite, in the
  • “: func:~murano.steps.metrics._answer_positions form (an int or
  • per-example sequence, negatives allowed); None overwrites every
  • query row. These are absolute columns (negatives count from the end),
  • so under Murano’s default left padding use -1 for the last real
  • token.
  • prompts_key: Results key to read the batch to run from.
  • logits_key: Results key to write the intervened logits under.
  • mask_key: Results key to write the attention mask under.

Raises:

  • ValueError: If method is unknown, targets is empty or malformed,
  • a method-specific argument is misused, or resample has no source.

__init__

Section titled “__init__”
def __init__(self, model: ModelBackend, targets: NodeSet | AddressLike | Iterable[AddressLike], method: Literal['zero', 'mean', 'resample'] = 'zero', source: Sequence[str] | None = None, source_key: str | None = None, positions: int | Sequence[int] | Tensor | None = None, prompts_key: str = keys.PROMPTS, logits_key: str = keys.ATTN_ABLATED_LOGITS, mask_key: str = keys.ATTN_ABLATED_MASK):