Info

Last Execution: 2026-08-28

Package	Version
nnsight	0.7.1.dev41+gd901da3ed
Python	3.12.13
torch	2.13.0+cu126
transformers	5.15.0.dev0

State-Space Models: Reading and Patching the Recurrent State¶

Introduction¶

Every other tutorial on this site assumes a transformer. This one is for the reader whose model is not one.

A transformer moves information across a sequence by re-reading the past. At position t, attention looks back over every earlier position, and the KV cache is a per-position record that grows with the sequence. Almost all of our tooling inherits that shape: an attention pattern is a matrix over pairs of positions, and "which layer, which token" is a well-posed question because the value at (layer, position) is a tensor you can index and overwrite.

A state-space model (SSM) moves information across a sequence by updating a fixed-size state. In Mamba, each layer holds a recurrent state h_t of shape (d_inner, d_state) and steps it forward,

h_t = A_t * h_{t-1} + B_t * x_t, y_t = C_t . h_t

where A_t, B_t, C_t are produced from the current input alone. The output at position t is read out of h_t and nothing else. Three things follow, and all three change how you do interpretability:

1. There is no attention pattern. Mixing across positions is a linear recurrence, not a weight over pairs of positions. There is nothing to plot as a matrix, and no head to ablate.
2. There is no KV cache. The per-layer analogue of "everything the model knows about the prefix" is h_t — one tensor per layer, the same size whether the prefix is 5 tokens or 5000.
3. The state is not a module output. backbone.layers[i].output gives you the residual stream, which does exist and behaves exactly as you would expect. But h_t is a local variable inside the mixer's forward, and the only place it surfaces is the cache object the model returns.

Point 3 is the practical obstacle, and most of this tutorial is about getting past it. Point 2 is the scientific payoff: a fixed-size state is a sufficient statistic for the model's entire future, and we will check that claim to the last bit.

We use state-spaces/mamba-130m-hf throughout — small enough for a free Colab GPU, real enough that its recurrent state actually carries facts. A second model, AntonV/mamba2-130m-hf, appears in the section on in-trace access.

What we will do:

• Route A — read the recurrent state out of the returned cache object, edit it, and feed it back in. This rests on a fact nothing in the documentation states, so we make it explicit and test it.
• The anchor, in both directions — a no-op write must leave the logits bit-identical, and a real write must move them. On a recurrent model the likely failure is a write that is silently discarded, and only the no-op half catches that.
• Route B — reach the state in-trace with .source, at the point in the forward where it is written. This is where a genuinely nasty trap lives.
• The experiment — activation patching on the recurrent state, across layers and across split points, with a self-patch anchor and two destructive controls.
• Breadth — which recurrent and hybrid architectures trace, which do not, and how to tell whose fault a failure is.

Setup¶

In [1]:

try:
    import google.colab
    is_colab = True
except ImportError:
    is_colab = False

if is_colab:
    !pip install -U nnsight

try: import google.colab is_colab = True except ImportError: is_colab = False if is_colab: !pip install -U nnsight

In [2]:

import statistics
import warnings

warnings.filterwarnings("ignore", message=".*IProgress not found.*")

from IPython.display import clear_output
import torch
import pandas as pd
import nnsight
from nnsight import TransformersModel

import plotly.express as px
import plotly.io as pio
pio.renderers.default = "colab" if is_colab else "plotly_mimetype+notebook_connected+colab+notebook"

DEVICE = "cuda:0" if torch.cuda.is_available() else "cpu"

import statistics import warnings warnings.filterwarnings("ignore", message=".*IProgress not found.*") from IPython.display import clear_output import torch import pandas as pd import nnsight from nnsight import TransformersModel import plotly.express as px import plotly.io as pio pio.renderers.default = "colab" if is_colab else "plotly_mimetype+notebook_connected+colab+notebook" DEVICE = "cuda:0" if torch.cuda.is_available() else "cpu"

state-spaces/mamba-130m-hf is a 130M-parameter Mamba-1 with 24 layers. We load it exactly as we would load GPT-2 — there is no SSM-specific loading path, and there does not need to be.

In [3]:

model = TransformersModel(
    "state-spaces/mamba-130m-hf", task="text-generation", device_map="auto", dispatch=True
)
clear_output()
print(model)

model = TransformersModel( "state-spaces/mamba-130m-hf", task="text-generation", device_map="auto", dispatch=True ) clear_output() print(model)

MambaForCausalLM(
  (backbone): MambaModel(
    (embeddings): Embedding(50280, 768)
    (layers): ModuleList(
      (0-23): 24 x MambaBlock(
        (norm): MambaRMSNorm(768, eps=1e-05)
        (mixer): MambaMixer(
          (conv1d): Conv1d(1536, 1536, kernel_size=(4,), stride=(1,), padding=(3,), groups=1536)
          (act): SiLUActivation()
          (in_proj): Linear(in_features=768, out_features=3072, bias=False)
          (x_proj): Linear(in_features=1536, out_features=80, bias=False)
          (dt_proj): Linear(in_features=48, out_features=1536, bias=True)
          (out_proj): Linear(in_features=1536, out_features=768, bias=False)
        )
      )
    )
    (norm_f): MambaRMSNorm(768, eps=1e-05)
  )
  (lm_head): Linear(in_features=768, out_features=50280, bias=False)
  (generator): Generator(
    (streamer): Streamer()
  )
)

The envoy tree is discoverable in the usual way. backbone.layers[i] is a MambaBlock — the analogue of a transformer's decoder layer — and inside it mixer is the SSM, with conv1d (a short causal convolution over the last few positions), in_proj, x_proj and dt_proj (which produce B, C and the step size from the input), and out_proj.

What is not in that tree is any module whose output is h_t. Let's confirm what the modules do give us.

In [4]:

tok = model.tokenizer
tok.clean_up_tokenization_spaces = False   # BPE tokenizer; the cleanup step would strip our spaces
N_LAYERS = len(model.backbone.layers)

prompt = "The Eiffel Tower is in the city of"

# Inside one invoke, modules must be requested in forward-pass order. The mixer runs
# *inside* the block, so asking for the block's output first raises OutOfOrderError.
with torch.no_grad(), model.trace(prompt):
    mixer_out = model.backbone.layers[0].mixer.output.save()
    block_out = model.backbone.layers[0].output.save()
    logits = model.output.logits.save()

print("mixer output ", tuple(mixer_out.shape))
print("block output ", tuple(block_out.shape), type(block_out).__name__)
print("top-1        ", repr(tok.decode(logits[0, -1].argmax())))

tok = model.tokenizer tok.clean_up_tokenization_spaces = False # BPE tokenizer; the cleanup step would strip our spaces N_LAYERS = len(model.backbone.layers) prompt = "The Eiffel Tower is in the city of" # Inside one invoke, modules must be requested in forward-pass order. The mixer runs # *inside* the block, so asking for the block's output first raises OutOfOrderError. with torch.no_grad(), model.trace(prompt): mixer_out = model.backbone.layers[0].mixer.output.save() block_out = model.backbone.layers[0].output.save() logits = model.output.logits.save() print("mixer output ", tuple(mixer_out.shape)) print("block output ", tuple(block_out.shape), type(block_out).__name__) print("top-1 ", repr(tok.decode(logits[0, -1].argmax())))

mixer output  (1, 10, 768)
block output  (1, 10, 768) Tensor
top-1         ' Paris'

Two things to note. First, a MambaBlock's .output is a bare tensor (batch, seq, hidden), not a tuple — the same convention as a transformers-5 decoder layer. Do not write .output[0] here; that indexes the batch axis.

Second, both of those tensors are (1, 10, 768) — the residual stream, one vector per position. That is the thing an SSM shares with a transformer, and everything you already know about the logit lens, steering vectors and probing applies to it unchanged. The recurrent state is a different object with a different shape, and it is nowhere in this tree.

Route A: the state lives on the model's output¶

MambaModel.forward takes a use_cache flag. With it set, the model returns a cache object alongside the logits, and the recurrent state of every layer is an attribute of that object. So we save the whole model output rather than just the logits.

In [5]:

with torch.no_grad(), model.trace(prompt, use_cache=True):
    out = model.output.save()

print("output keys:", list(out.keys()))

cache = out["cache_params"]
print("cache type :", type(cache).__name__, "with", len(cache.layers), "layers")

layer0 = cache.layers[0]
print("layer 0    :", type(layer0).__name__)
print("  recurrent_states[0]", tuple(layer0.recurrent_states[0].shape), "= (batch, d_inner, d_state)")
print("  conv_states[0]     ", tuple(layer0.conv_states[0].shape), "= (batch, d_inner, conv_kernel)")

with torch.no_grad(), model.trace(prompt, use_cache=True): out = model.output.save() print("output keys:", list(out.keys())) cache = out["cache_params"] print("cache type :", type(cache).__name__, "with", len(cache.layers), "layers") layer0 = cache.layers[0] print("layer 0 :", type(layer0).__name__) print(" recurrent_states[0]", tuple(layer0.recurrent_states[0].shape), "= (batch, d_inner, d_state)") print(" conv_states[0] ", tuple(layer0.conv_states[0].shape), "= (batch, d_inner, conv_kernel)")

output keys: ['logits', 'cache_params']
cache type : DynamicCache with 24 layers
layer 0    : LinearAttentionLayer
  recurrent_states[0] (1, 1536, 16) = (batch, d_inner, d_state)
  conv_states[0]      (1, 1536, 4) = (batch, d_inner, conv_kernel)

There it is: cache.layers[i].recurrent_states[0] is h_T for layer i after the whole prompt. Alongside it, conv_states[0] is the mixer's other piece of memory — the last few positions of the pre-SSM signal, which the depthwise conv1d needs in order to keep producing the same outputs when we resume. Together they are the model's entire carried state.

Note the shape: (1, 1536, 16) does not depend on the prompt length. That is the whole point of a recurrent model, and later it will let us patch between prompts of different lengths with no alignment step at all.

The undocumented part. Reading the state is half the job; we also need to put an edited state back. model.trace(...) forwards any extra keyword argument straight to the model's forward, so cache_params= works:

with model.trace(next_ids, cache_params=cache, use_cache=True):
    ...

Nothing in the nnsight documentation says this — the trace page describes how inputs are assembled and does not mention that other keywords ride along. It is worth knowing generally: use_cache=, past_key_values=, trust_remote_code=, and anything else your model's forward accepts, all pass through untouched. For recurrent models it is the entire intervention API, because the state is not something an envoy can reach.

Let's define the two operations we will use for the rest of the notebook.

In [6]:

def encode(text):
    return tok(text, return_tensors="pt")["input_ids"].to(DEVICE)

def run(ids, cache=None):
    """One forward pass. Returns (cache, logits); `cache` is passed to the forward as `cache_params`."""
    with torch.no_grad(), model.trace(ids, cache_params=cache, use_cache=True):
        out = model.output.save()
    return out["cache_params"], out["logits"]

def rec(cache, layer):
    """The recurrent state of one layer: (batch, d_inner, d_state)."""
    return cache.layers[layer].recurrent_states[0]

def conv(cache, layer):
    """The convolutional state of one layer: (batch, d_inner, conv_kernel)."""
    return cache.layers[layer].conv_states[0]

def encode(text): return tok(text, return_tensors="pt")["input_ids"].to(DEVICE) def run(ids, cache=None): """One forward pass. Returns (cache, logits); `cache` is passed to the forward as `cache_params`.""" with torch.no_grad(), model.trace(ids, cache_params=cache, use_cache=True): out = model.output.save() return out["cache_params"], out["logits"] def rec(cache, layer): """The recurrent state of one layer: (batch, d_inner, d_state).""" return cache.layers[layer].recurrent_states[0] def conv(cache, layer): """The convolutional state of one layer: (batch, d_inner, conv_kernel).""" return cache.layers[layer].conv_states[0]

Before we intervene, check that the mechanism reproduces the model. Running the full prompt in one pass and running it as a prefill plus single-token steps carrying the cache should give the same prediction. If they don't, every number that follows is measuring our plumbing rather than the model.

In [7]:

ids = encode(prompt)

_, logits_full = run(ids)                        # one 10-token forward

cache, _ = run(ids[:, :-2])                      # prefill 8 tokens
for t in range(ids.shape[1] - 2, ids.shape[1]):  # then two single-token steps
    cache, logits_stepped = run(ids[:, t : t + 1], cache)

print("same top-1 :", logits_full[0, -1].argmax().item() == logits_stepped[0, -1].argmax().item())
print("max |delta|:", (logits_full[0, -1] - logits_stepped[0, -1]).abs().max().item())

ids = encode(prompt) _, logits_full = run(ids) # one 10-token forward cache, _ = run(ids[:, :-2]) # prefill 8 tokens for t in range(ids.shape[1] - 2, ids.shape[1]): # then two single-token steps cache, logits_stepped = run(ids[:, t : t + 1], cache) print("same top-1 :", logits_full[0, -1].argmax().item() == logits_stepped[0, -1].argmax().item()) print("max |delta|:", (logits_full[0, -1] - logits_stepped[0, -1]).abs().max().item())

same top-1 : True
max |delta|: 0.0002040863037109375

Same prediction, and the logits agree to about 2e-4. They are not bit-identical because HuggingFace's Mamba implementation runs two different code paths — a scan over the whole prefill, and a single-step recurrence during decoding — and in fp32 those differ in the last few bits. That is a property of the reference implementation, not of nnsight or of our cache surgery, and it sits three orders of magnitude below the effects we are about to measure. It is worth knowing it is there so you don't chase it later.

Feed continuation tokens one at a time. With a populated cache, HuggingFace's Mamba mixer takes its decode branch, which assumes seq_len == 1. Handing it a two-token chunk on top of an existing cache silently produces the wrong convolution state. This is upstream, but it shapes every experiment on this page: after a patch, step one token at a time.

The anchor, in both directions¶

Before reading anything into an intervention, we check it two ways.

The converse anchor: write the state back unchanged and confirm the output is bit-identical. On a recurrent model this is not a formality. The state is a local variable that gets copied into a cache object; a write that lands in the wrong place, or after the value has already been consumed, is silently discarded. A patch that produces no change is indistinguishable from a patch that never happened — unless you have first shown that a no-op write is exactly a no-op.

The forward anchor: write a different state in and confirm the output moves. Together the two say "the write lands, and only the write lands."

Our task is factual recall with a counterfactual pair, the same setup as the activation patching tutorial:

• clean: "The Eiffel Tower is in the city of" → " Paris"
• corrupt: "The Colosseum is in the city of" → " Rome"

and the metric is the logit difference logit(" Paris") - logit(" Rome") at the last position. Note the two prompts have different lengths (10 and 9 tokens). For a transformer that would force an alignment decision before we could patch anything. For an SSM it costs nothing: the state has the same shape either way.

In [8]:

clean_prompt = "The Eiffel Tower is in the city of"
corrupt_prompt = "The Colosseum is in the city of"

clean_ids, corrupt_ids = encode(clean_prompt), encode(corrupt_prompt)

# Leading spaces: our prompts do not end in one, so the space+word token is what we want.
PARIS = tok(" Paris")["input_ids"][0]
ROME = tok(" Rome")["input_ids"][0]

def logit_diff(logits, good=PARIS, bad=ROME):
    return (logits[0, -1, good] - logits[0, -1, bad]).item()

print("clean  ", tuple(clean_ids.shape), tok.convert_ids_to_tokens(clean_ids[0]))
print("corrupt", tuple(corrupt_ids.shape), tok.convert_ids_to_tokens(corrupt_ids[0]))

clean_prompt = "The Eiffel Tower is in the city of" corrupt_prompt = "The Colosseum is in the city of" clean_ids, corrupt_ids = encode(clean_prompt), encode(corrupt_prompt) # Leading spaces: our prompts do not end in one, so the space+word token is what we want. PARIS = tok(" Paris")["input_ids"][0] ROME = tok(" Rome")["input_ids"][0] def logit_diff(logits, good=PARIS, bad=ROME): return (logits[0, -1, good] - logits[0, -1, bad]).item() print("clean ", tuple(clean_ids.shape), tok.convert_ids_to_tokens(clean_ids[0])) print("corrupt", tuple(corrupt_ids.shape), tok.convert_ids_to_tokens(corrupt_ids[0]))

clean   (1, 10) ['The', 'ĠE', 'iff', 'el', 'ĠTower', 'Ġis', 'Ġin', 'Ġthe', 'Ġcity', 'Ġof']
corrupt (1, 9) ['The', 'ĠCol', 'osse', 'um', 'Ġis', 'Ġin', 'Ġthe', 'Ġcity', 'Ġof']

The patching driver: prefill all but the last k tokens, apply an edit to the cache, then feed the remaining k tokens one at a time. With patch=None it is just a run.

In [9]:

def continuation(ids, k, patch=None):
    """Prefill ids[:, :-k], optionally edit the cache, then step the last k tokens."""
    split = ids.shape[1] - k
    cache, logits = run(ids[:, :split])
    if patch is not None:
        patch(cache)
    for t in range(split, ids.shape[1]):
        cache, logits = run(ids[:, t : t + 1], cache)
    return logits

K = 2

clean_logits = continuation(clean_ids, K)
corrupt_logits = continuation(corrupt_ids, K)

print(f"clean   LD {logit_diff(clean_logits):+.3f}  top-1 {tok.decode(clean_logits[0, -1].argmax())!r}")
print(f"corrupt LD {logit_diff(corrupt_logits):+.3f}  top-1 {tok.decode(corrupt_logits[0, -1].argmax())!r}")

def continuation(ids, k, patch=None): """Prefill ids[:, :-k], optionally edit the cache, then step the last k tokens.""" split = ids.shape[1] - k cache, logits = run(ids[:, :split]) if patch is not None: patch(cache) for t in range(split, ids.shape[1]): cache, logits = run(ids[:, t : t + 1], cache) return logits K = 2 clean_logits = continuation(clean_ids, K) corrupt_logits = continuation(corrupt_ids, K) print(f"clean LD {logit_diff(clean_logits):+.3f} top-1 {tok.decode(clean_logits[0, -1].argmax())!r}") print(f"corrupt LD {logit_diff(corrupt_logits):+.3f} top-1 {tok.decode(corrupt_logits[0, -1].argmax())!r}")

clean   LD +5.508  top-1 ' Paris'
corrupt LD -6.239  top-1 ' Rome'

Anchor 1 — the no-op write. We take a snapshot of the corrupt run's own state at the split point and write it straight back into every layer, then continue. If our cache surgery is real, this changes nothing at all — not approximately, exactly.

In [10]:

snap_cache, _ = run(corrupt_ids[:, : corrupt_ids.shape[1] - K])
corrupt_state = [(rec(snap_cache, L).clone(), conv(snap_cache, L).clone()) for L in range(N_LAYERS)]

def write_state(cache, states, layers=None, what="both"):
    for L in (range(N_LAYERS) if layers is None else layers):
        h, c = states[L]
        if what in ("both", "recurrent"):
            rec(cache, L).copy_(h)
        if what in ("both", "conv"):
            conv(cache, L).copy_(c)

noop_logits = continuation(corrupt_ids, K, lambda c: write_state(c, corrupt_state))

print("bit-identical:", torch.equal(noop_logits, corrupt_logits))
print("max |delta|  :", (noop_logits - corrupt_logits).abs().max().item())

snap_cache, _ = run(corrupt_ids[:, : corrupt_ids.shape[1] - K]) corrupt_state = [(rec(snap_cache, L).clone(), conv(snap_cache, L).clone()) for L in range(N_LAYERS)] def write_state(cache, states, layers=None, what="both"): for L in (range(N_LAYERS) if layers is None else layers): h, c = states[L] if what in ("both", "recurrent"): rec(cache, L).copy_(h) if what in ("both", "conv"): conv(cache, L).copy_(c) noop_logits = continuation(corrupt_ids, K, lambda c: write_state(c, corrupt_state)) print("bit-identical:", torch.equal(noop_logits, corrupt_logits)) print("max |delta| :", (noop_logits - corrupt_logits).abs().max().item())

bit-identical: True
max |delta|  : 0.0

torch.equal is True and the maximum absolute difference is exactly 0.0. The write reaches the tensors the forward pass reads, and it perturbs nothing else. Now the other direction.

Anchor 2 — the real write. We take the clean run's state at its own split point and write it into the corrupt run. Because the state is fixed-size, "its own split point" needs no alignment: both runs have two tokens left to go, so both states summarise everything up to "... is in the city", one having read the Eiffel Tower and one the Colosseum.

We do this three ways — recurrent state only, convolutional state only, and both — because that decomposition is itself informative.

In [11]:

clean_cache, _ = run(clean_ids[:, : clean_ids.shape[1] - K])
clean_state = [(rec(clean_cache, L).clone(), conv(clean_cache, L).clone()) for L in range(N_LAYERS)]

for what in ("recurrent", "conv", "both"):
    patched = continuation(corrupt_ids, K, lambda c, w=what: write_state(c, clean_state, what=w))
    print(f"patch {what:>9} -> LD {logit_diff(patched):+7.3f}   top-1 {tok.decode(patched[0, -1].argmax())!r}")

print(f"      {'(clean)':>9} -> LD {logit_diff(clean_logits):+7.3f}   top-1 {tok.decode(clean_logits[0, -1].argmax())!r}")
print(f"      {'(corrupt)':>9} -> LD {logit_diff(corrupt_logits):+7.3f}   top-1 {tok.decode(corrupt_logits[0, -1].argmax())!r}")

clean_cache, _ = run(clean_ids[:, : clean_ids.shape[1] - K]) clean_state = [(rec(clean_cache, L).clone(), conv(clean_cache, L).clone()) for L in range(N_LAYERS)] for what in ("recurrent", "conv", "both"): patched = continuation(corrupt_ids, K, lambda c, w=what: write_state(c, clean_state, what=w)) print(f"patch {what:>9} -> LD {logit_diff(patched):+7.3f} top-1 {tok.decode(patched[0, -1].argmax())!r}") print(f" {'(clean)':>9} -> LD {logit_diff(clean_logits):+7.3f} top-1 {tok.decode(clean_logits[0, -1].argmax())!r}") print(f" {'(corrupt)':>9} -> LD {logit_diff(corrupt_logits):+7.3f} top-1 {tok.decode(corrupt_logits[0, -1].argmax())!r}")

patch recurrent -> LD  +5.614   top-1 ' Paris'
patch      conv -> LD  -6.125   top-1 ' Rome'

patch      both -> LD  +5.508   top-1 ' Paris'
        (clean) -> LD  +5.508   top-1 ' Paris'
      (corrupt) -> LD  -6.239   top-1 ' Rome'

The recurrent state alone carries the fact: writing it in flips the prediction from " Rome" all the way to " Paris" and takes the logit difference slightly past the clean value. The convolutional state alone does essentially nothing — it is a four-position sliding window over "... is in the city of", which is identical in both prompts, so there is nothing in it to swap. The anchor passes in both directions.

The state is a sufficient statistic¶

The last row of that table deserves its own claim. If recurrent_states and conv_states are everything the model carries forward, then overwriting both should not merely move the prediction towards the clean run — it should reproduce the clean run exactly, from a different prompt, of a different length.

In [12]:

both_logits = continuation(corrupt_ids, K, lambda c: write_state(c, clean_state, what="both"))

print("allclose to the clean continuation:", torch.allclose(both_logits, clean_logits))
print("bit-identical                     :", torch.equal(both_logits, clean_logits))
print("max |delta|                       :", (both_logits - clean_logits).abs().max().item())

both_logits = continuation(corrupt_ids, K, lambda c: write_state(c, clean_state, what="both")) print("allclose to the clean continuation:", torch.allclose(both_logits, clean_logits)) print("bit-identical :", torch.equal(both_logits, clean_logits)) print("max |delta| :", (both_logits - clean_logits).abs().max().item())

allclose to the clean continuation: True
bit-identical                     : True
max |delta|                       : 0.0

Bit-identical. Having read "The Colosseum" and then had its state replaced, the model's continuation is indistinguishable — to the last bit of every one of the 50,280 logits — from a model that read "The Eiffel Tower".

This is worth pausing on, because it is a property no KV-cache model has. In a transformer, "everything the model knows about the prefix" is a record that grows with the prefix, and swapping it between two prompts of different lengths is not even well defined. Here, two tensors of fixed shape per layer are the complete interface between the past and the future. Every intervention on an SSM's memory is an edit to those tensors, and this check tells us the edit surface is complete.

Route B: reaching the state in-trace with .source¶

Route A works between traces. It is complete — any layer, any split point — but it costs one forward pass per token you step through, and it can only see the state at moments where the model writes a cache. The in-trace route uses .source, which exposes the individual operations inside a module's forward -- every call and every assignment -- as hookable locations. That second half matters here, because the recurrence is written with operators, not function calls.

In [13]:

print(model.backbone.layers[0].mixer.source)

print(model.backbone.layers[0].mixer.source)

                                 * @force_accelerate_hooks("conv1d")
                                 0 def forward(
                                 1     self,
                                 2     hidden_states,
                                 3     cache_params: Cache | None = None,
                                 4     attention_mask: torch.LongTensor | None = None,
                                 5     **kwargs,
                                 6 ):
 all_0                       ->  7     is_fast_path_available = all(
 is_fast_path_available_0    ->  +     ...
                                 8         (selective_state_update, selective_scan_fn, causal_conv1d_fn, causal_conv1d_update, mamba_inner_fn)
                                 9     )
 is_tracing_0                -> 10     if is_fast_path_available and "cuda" in self.x_proj.weight.device.type and not is_tracing(hidden_states):
 self_cuda_kernels_forward_0 -> 11         return self.cuda_kernels_forward(hidden_states, cache_params, attention_mask)
 self_slow_forward_0         -> 12     return self.slow_forward(hidden_states, cache_params, attention_mask)
                                13 

MambaMixer.forward is wrapped in transformers' @force_accelerate_hooks, a decorator that does not use functools.wraps. nnsight peels it anyway -- it reads the wrapper's own source to find the function it calls -- so the listing is the real forward. What it shows is a dispatcher: fused CUDA kernels if they are installed, slow_forward otherwise. We have no kernels, so slow_forward runs, and a second .source reaches into it. Recursive .source resolves the callee from the live value flowing through the call, so it only works inside a trace.

In [14]:

with torch.no_grad(), model.trace(prompt, use_cache=True) as tracer:
    sf = model.backbone.layers[0].mixer.source.self_slow_forward_0.source
    n_tokens = len(sf.range_0.output)              # the recurrence loop's own range(seq_len)
    per_token = nnsight.save([])
    for _ in tracer.iter[:n_tokens]:               # the loop-body assignment fires once per token
        per_token.append(sf.ssm_state_3.output)    # ssm_state = discrete_A[:, :, i] * ssm_state + deltaB_u[:, :, i]
    out = model.output.save()

print(f"{len(per_token)} per-token states of shape {tuple(per_token[0].shape)} = (batch, d_inner, d_state)")
print("the last one is the state Route A reads from the cache:",
      torch.equal(per_token[-1], out["cache_params"].layers[0].recurrent_states[0]))

with torch.no_grad(), model.trace(prompt, use_cache=True) as tracer: sf = model.backbone.layers[0].mixer.source.self_slow_forward_0.source n_tokens = len(sf.range_0.output) # the recurrence loop's own range(seq_len) per_token = nnsight.save([]) for _ in tracer.iter[:n_tokens]: # the loop-body assignment fires once per token per_token.append(sf.ssm_state_3.output) # ssm_state = discrete_A[:, :, i] * ssm_state + deltaB_u[:, :, i] out = model.output.save() print(f"{len(per_token)} per-token states of shape {tuple(per_token[0].shape)} = (batch, d_inner, d_state)") print("the last one is the state Route A reads from the cache:", torch.equal(per_token[-1], out["cache_params"].layers[0].recurrent_states[0]))

10 per-token states of shape (1, 1536, 16) = (batch, d_inner, d_state)
the last one is the state Route A reads from the cache: True

The recurrence is one line -- ssm_state = discrete_A[:, :, i, :] * ssm_state + deltaB_u[:, :, i, :] -- a multiply and an add, not a call. It is addressable because every assignment in a forward is an operation named after its target (ssm_state_3: the fourth binding of that name in slow_forward), and because it sits inside for i in range(seq_len) it fires once per token, so tracer.iter[k] selects the token. The loop's bound is the range(...) call one line above it, so the count comes from the trace rather than from the tokenizer. Route B sees every intermediate state; Route A only ever sees the last.

Mamba-2's mixer is written the same way, a dispatcher over a pure-PyTorch torch_forward, and its cache API is the one the decode-step experiments below need, so we switch to it for the rest of the section.

In [15]:

model2 = TransformersModel(
    "AntonV/mamba2-130m-hf", task="text-generation", device_map="auto", dispatch=True
)
clear_output()

mixer2 = model2.backbone.layers[0].mixer
print(mixer2.source)

model2 = TransformersModel( "AntonV/mamba2-130m-hf", task="text-generation", device_map="auto", dispatch=True ) clear_output() mixer2 = model2.backbone.layers[0].mixer print(mixer2.source)

                                 * def forward(
                                 0     self,
                                 1     hidden_states,
                                 2     cache_params: Cache | None = None,
                                 3     attention_mask: torch.Tensor | None = None,
                                 4     **kwargs,
                                 5 ):
 is_torchdynamo_compiling_0  ->  6     if is_fast_path_available and "cuda" in self.in_proj.weight.device.type and not is_torchdynamo_compiling():
 self_cuda_kernels_forward_0 ->  7         return self.cuda_kernels_forward(hidden_states, cache_params, attention_mask)
 self_torch_forward_0        ->  8     return self.torch_forward(hidden_states, cache_params, attention_mask)
                                 9 

forward is a dispatcher: it picks a fused CUDA kernel if one is installed, and the pure-PyTorch torch_forward otherwise. We have no fused kernels here, so torch_forward runs, and we reach into it with a second .source. Recursive .source resolves the callee from the live value flowing through the call, so it only works inside a trace.

In [16]:

def run2(ids, cache=None):
    with torch.no_grad(), model2.trace(ids, cache_params=cache, use_cache=True):
        out = model2.output.save()
    return out["cache_params"], out["logits"]

prompt2_ids = model2.tokenizer(prompt, return_tensors="pt")["input_ids"].to(DEVICE)

names = {}
with torch.no_grad(), model2.trace(prompt2_ids, use_cache=True):
    names["ops"] = nnsight.save(mixer2.source.self_torch_forward_0.source.names)

print(len(names["ops"]), "addressable operations in torch_forward")
print([n for n in names["ops"] if "recurrent" in n or "ssm_state" in n])

def run2(ids, cache=None): with torch.no_grad(), model2.trace(ids, cache_params=cache, use_cache=True): out = model2.output.save() return out["cache_params"], out["logits"] prompt2_ids = model2.tokenizer(prompt, return_tensors="pt")["input_ids"].to(DEVICE) names = {} with torch.no_grad(), model2.trace(prompt2_ids, use_cache=True): names["ops"] = nnsight.save(mixer2.source.self_torch_forward_0.source.names) print(len(names["ops"]), "addressable operations in torch_forward") print([n for n in names["ops"] if "recurrent" in n or "ssm_state" in n])

173 addressable operations in torch_forward
['ssm_states_0', 'cache_params_update_recurrent_state_0', 'ssm_states_1', 'ssm_states_to_0', 'ssm_states_2', 'ssm_states_view_0', 'ssm_states_reshaped_0', 'cache_params_layers_recurrent_states_to_0', 'ssm_state_0', 'cache_params_update_recurrent_state_1']

Several handles mention the state. Three matter here. ssm_states_0 is the assignment that computes h_t -- an operator expression, recurrent_states[0] * dA + dBx. cache_params_update_recurrent_state_0 is the call one line later that stores it in the single-token decode branch. ssm_states_to_0 is a dtype cast on the way into the readout. We will use all three, and the differences between them are the point of this section.

We set up a decode step: prefill all but the last token, then trace a single-token step with the cache in place.

In [17]:

prefill = prompt2_ids[:, :-1]
last_token = prompt2_ids[:, -1:]

def fresh_cache():
    cache, _ = run2(prefill)
    return cache

# Warm-up. The *first* `.source` access on a module raises OutOfOrderError: instrumenting an
# operation rewrites the forward, and that can only happen before the forward runs. One
# throwaway trace that reads nothing else gets the instrumentation installed.
with torch.no_grad(), model2.trace(last_token, cache_params=fresh_cache(), use_cache=True):
    _ = mixer2.source.self_torch_forward_0.source.cache_params_update_recurrent_state_0.input.save()

cache_b = fresh_cache()
with torch.no_grad(), model2.trace(last_token, cache_params=cache_b, use_cache=True):
    sf = mixer2.source.self_torch_forward_0.source
    h_in = sf.cache_params_update_recurrent_state_0.input.clone().save()
    base_logits = model2.output.logits.clone().save()

stored = cache_b.layers[0].recurrent_states[0].clone()
print("h read in-trace  ", tuple(h_in.shape), "= (batch, heads, head_dim, d_state)")
print("equals the cache ", torch.equal(h_in, stored))

prefill = prompt2_ids[:, :-1] last_token = prompt2_ids[:, -1:] def fresh_cache(): cache, _ = run2(prefill) return cache # Warm-up. The *first* `.source` access on a module raises OutOfOrderError: instrumenting an # operation rewrites the forward, and that can only happen before the forward runs. One # throwaway trace that reads nothing else gets the instrumentation installed. with torch.no_grad(), model2.trace(last_token, cache_params=fresh_cache(), use_cache=True): _ = mixer2.source.self_torch_forward_0.source.cache_params_update_recurrent_state_0.input.save() cache_b = fresh_cache() with torch.no_grad(), model2.trace(last_token, cache_params=cache_b, use_cache=True): sf = mixer2.source.self_torch_forward_0.source h_in = sf.cache_params_update_recurrent_state_0.input.clone().save() base_logits = model2.output.logits.clone().save() stored = cache_b.layers[0].recurrent_states[0].clone() print("h read in-trace ", tuple(h_in.shape), "= (batch, heads, head_dim, d_state)") print("equals the cache ", torch.equal(h_in, stored))

h read in-trace   (1, 24, 64, 128) = (batch, heads, head_dim, d_state)
equals the cache  True

The value we read at the call site is bit-identical to the one that ends up in the cache — Route A and Route B are looking at the same tensor. (The shape differs from Mamba-1's because Mamba-2 splits d_inner into heads.)

Now the same two anchors, in-trace.

In [18]:

# Anchor 1: write the value back unchanged.
c_noop = fresh_cache()
with torch.no_grad(), model2.trace(last_token, cache_params=c_noop, use_cache=True):
    sf = mixer2.source.self_torch_forward_0.source
    same = sf.cache_params_update_recurrent_state_0.input
    sf.cache_params_update_recurrent_state_0.input = same
    noop2 = model2.output.logits.clone().save()

print("no-op write  -> logits identical", torch.equal(noop2, base_logits),
      "| state identical", torch.equal(c_noop.layers[0].recurrent_states[0], stored))

# Anchor 2: zero the state on its way into the store.
c_zero = fresh_cache()
with torch.no_grad(), model2.trace(last_token, cache_params=c_zero, use_cache=True):
    sf = mixer2.source.self_torch_forward_0.source
    sf.cache_params_update_recurrent_state_0.input = torch.zeros_like(h_in)
    zero_logits = model2.output.logits.clone().save()

print(f"zeroed store -> max |delta logits| {(zero_logits - base_logits).abs().max().item():.4f}"
      f" | stored state max |.| {c_zero.layers[0].recurrent_states[0].abs().max().item():.4f}")

# Anchor 1: write the value back unchanged. c_noop = fresh_cache() with torch.no_grad(), model2.trace(last_token, cache_params=c_noop, use_cache=True): sf = mixer2.source.self_torch_forward_0.source same = sf.cache_params_update_recurrent_state_0.input sf.cache_params_update_recurrent_state_0.input = same noop2 = model2.output.logits.clone().save() print("no-op write -> logits identical", torch.equal(noop2, base_logits), "| state identical", torch.equal(c_noop.layers[0].recurrent_states[0], stored)) # Anchor 2: zero the state on its way into the store. c_zero = fresh_cache() with torch.no_grad(), model2.trace(last_token, cache_params=c_zero, use_cache=True): sf = mixer2.source.self_torch_forward_0.source sf.cache_params_update_recurrent_state_0.input = torch.zeros_like(h_in) zero_logits = model2.output.logits.clone().save() print(f"zeroed store -> max |delta logits| {(zero_logits - base_logits).abs().max().item():.4f}" f" | stored state max |.| {c_zero.layers[0].recurrent_states[0].abs().max().item():.4f}")

no-op write  -> logits identical True | state identical True
zeroed store -> max |delta logits| 1.2358 | stored state max |.| 0.0000

The trap¶

Now the same experiment aimed one operation later — at ssm_states_to_0, the cast that feeds the readout. It looks like the same intervention: we are writing zeros over the recurrent state, at a call site whose value is the recurrent state.

In [19]:

c_trap = fresh_cache()
with torch.no_grad(), model2.trace(last_token, cache_params=c_trap, use_cache=True):
    sf = mixer2.source.self_torch_forward_0.source
    sf.ssm_states_to_0.output = torch.zeros_like(h_in)
    trap_logits = model2.output.logits.clone().save()

print(f"max |delta logits|                     {(trap_logits - base_logits).abs().max().item():.4f}")
print(f"stored state bit-identical to baseline {torch.equal(c_trap.layers[0].recurrent_states[0], stored)}")

c_trap = fresh_cache() with torch.no_grad(), model2.trace(last_token, cache_params=c_trap, use_cache=True): sf = mixer2.source.self_torch_forward_0.source sf.ssm_states_to_0.output = torch.zeros_like(h_in) trap_logits = model2.output.logits.clone().save() print(f"max |delta logits| {(trap_logits - base_logits).abs().max().item():.4f}") print(f"stored state bit-identical to baseline {torch.equal(c_trap.layers[0].recurrent_states[0], stored)}")

max |delta logits|                     1.2358
stored state bit-identical to baseline True

The logits move by exactly as much as the real intervention did — and the stored state is bit-identical to the unpatched run. We changed the model's output for this token without changing its memory at all.

Everything about the experience says the intervention worked: no error, no warning, and a large output change. If we had stopped at "the write moved the logits", we would have written up a state intervention that never touched the state.

The mechanism is visible in the source printed above. Three consecutive lines of torch_forward do this:

ssm_states = cache_params.layers[self.layer_idx].recurrent_states[0] * dA + dBx   # ssm_states_0
ssm_states = cache_params.update_recurrent_state(ssm_states, layer_idx=...)       # cache_params_update_recurrent_state_0
ssm_states = ssm_states.to(device=C.device, dtype=C.dtype)                         # ssm_states_to_0

The cast on the third line makes a copy. Writing to it changes what the readout sees for this token, not the tensor the store put in the cache one line earlier -- and nothing in the experience says so. Two handles show the same tensor at different moments, and only one of them is the memory. The intervention that means what it says is the first line: ssm_states_0 is the expression that computes the state, and a write there is seen by the store and the readout alike.

The check that separates the two cases is not in the logits at all — it is in the cache. Route A hands us every layer's state after the pass, so we can ask each write what it actually changed.

In [20]:

# Anchor 3: write the state where it is computed. `ssm_states_0` is the operator
# expression itself, one line above the store, so the cache and the readout both see it.
c_expr = fresh_cache()
with torch.no_grad(), model2.trace(last_token, cache_params=c_expr, use_cache=True):
    sf = mixer2.source.self_torch_forward_0.source
    sf.ssm_states_0.output = torch.zeros_like(h_in)
    expr_logits = model2.output.logits.clone().save()

print(f"zeroed the expression -> max |delta logits| {(expr_logits - base_logits).abs().max().item():.4f}",
      "| stored state is zero:", bool((c_expr.layers[0].recurrent_states[0] == 0).all()))

# Anchor 3: write the state where it is computed. `ssm_states_0` is the operator # expression itself, one line above the store, so the cache and the readout both see it. c_expr = fresh_cache() with torch.no_grad(), model2.trace(last_token, cache_params=c_expr, use_cache=True): sf = mixer2.source.self_torch_forward_0.source sf.ssm_states_0.output = torch.zeros_like(h_in) expr_logits = model2.output.logits.clone().save() print(f"zeroed the expression -> max |delta logits| {(expr_logits - base_logits).abs().max().item():.4f}", "| stored state is zero:", bool((c_expr.layers[0].recurrent_states[0] == 0).all()))

zeroed the expression -> max |delta logits| 1.2358 | stored state is zero: True

In [21]:

c_ref = fresh_cache()
_, _ = run2(last_token, c_ref)          # the same decode step, unpatched, for reference

for name, c in (("no-op write", c_noop), ("zeroed the store", c_zero), ("zeroed the readout cast", c_trap),
                ("zeroed the expression", c_expr)):
    d = [(c.layers[L].recurrent_states[0] - c_ref.layers[L].recurrent_states[0]).abs().max().item()
         for L in range(len(c.layers))]
    print(f"{name:>24}: layer 0 state moved by {d[0]:.4f}, layers 1-23 by up to {max(d[1:]):.4f}")

c_ref = fresh_cache() _, _ = run2(last_token, c_ref) # the same decode step, unpatched, for reference for name, c in (("no-op write", c_noop), ("zeroed the store", c_zero), ("zeroed the readout cast", c_trap), ("zeroed the expression", c_expr)): d = [(c.layers[L].recurrent_states[0] - c_ref.layers[L].recurrent_states[0]).abs().max().item() for L in range(len(c.layers))] print(f"{name:>24}: layer 0 state moved by {d[0]:.4f}, layers 1-23 by up to {max(d[1:]):.4f}")

             no-op write: layer 0 state moved by 0.0000, layers 1-23 by up to 0.0000
        zeroed the store: layer 0 state moved by 8.3147, layers 1-23 by up to 4.1748
 zeroed the readout cast: layer 0 state moved by 0.0000, layers 1-23 by up to 4.1748
   zeroed the expression: layer 0 state moved by 8.3147, layers 1-23 by up to 4.1748

Both writes changed the model, and by the same amount downstream — layers 1-23 move by at most 4.17 either way. What separates them is layer 0 itself: zeroing the store moves layer 0's recurrent state by 8.31, the full magnitude of what it replaced, while zeroing the readout cast leaves it at exactly 0.0000 and instead perturbs layer 0's output, which then propagates into every later layer's state.

So the trap is not a no-op. It is a real intervention on a different quantity, wearing the same clothes. If the question was "what does this model do when layer 0 forgets everything?", it answers a different question and looks identical while doing so.

On a stateful model, verify a write against the state, not against the logits. Route A is what makes that cheap: the cache is right there on the output.

The experiment: does information localise in a recurrent model?¶

We now run the SSM version of causal tracing. The design is the denoising direction: run the corrupt prompt, replace one layer's recurrent state with the clean run's at a chosen split point, feed the remaining tokens one at a time, and measure how much of the clean logit difference comes back.

restoration = (LD_patched - LD_corrupt) / (LD_clean - LD_corrupt)

so 0 means the patch did nothing and 1 means it fully restored the clean behaviour.

Two axes, and only one of them is the familiar one. Depth is the same question as in a transformer: which layer holds the entity. The other axis is not token position, because there is no per-position value to patch — the state at layer L exists at exactly one moment per forward pass. What we can vary is where we cut: how many tokens the model had consumed when we swapped its memory. We index that by how far the split point sits past the end of the subject, so that the two prompts are compared at the same place even though their subjects tokenise to different lengths.

The prompt set, and how it was chosen. Five factual-recall pairs. The screening rule was fixed before any patching: keep a pair only if the model's top-1 continuation is correct for both prompts and the base logit gap exceeds 5. Five of thirteen candidates passed; the rest were dropped for answering something else entirely (for "The Sahara desert is in the continent of" / "The Andes mountains are in the continent of" the model prefers " Africa" for both, so its "clean" run is not clean). Dropping items after seeing a patching result would be a different and much worse procedure.

In [22]:

# (clean prompt, corrupt prompt, clean subject, corrupt subject, clean answer, corrupt answer)
PAIRS = [
    ("The Eiffel Tower is in the city of", "The Colosseum is in the city of",
     "The Eiffel Tower", "The Colosseum", " Paris", " Rome"),
    ("The Louvre museum is in the city of", "The Acropolis is in the city of",
     "The Louvre museum", "The Acropolis", " Paris", " Athens"),
    ("Berlin is the capital city of", "Madrid is the capital city of",
     "Berlin", "Madrid", " Germany", " Spain"),
    ("Tokyo is a city in the country of", "Cairo is a city in the country of",
     "Tokyo", "Cairo", " Japan", " Egypt"),
    ("Seoul is a city in the country of", "Lisbon is a city in the country of",
     "Seoul", "Lisbon", " Korea", " Portugal"),
]

def subject_len(prompt, subject):
    n = len(tok(subject)["input_ids"])
    assert tok(prompt)["input_ids"][:n] == tok(subject)["input_ids"], "subject is not a token prefix"
    return n

for cp, xp, cs, xs, a, b in PAIRS:
    print(f"{len(tok(cp)['input_ids']):2d}/{len(tok(xp)['input_ids']):2d} tokens, "
          f"subject ends at {subject_len(cp, cs)}/{subject_len(xp, xs)}  "
          f"{cs!r} vs {xs!r} -> {a!r}/{b!r}")

# (clean prompt, corrupt prompt, clean subject, corrupt subject, clean answer, corrupt answer) PAIRS = [ ("The Eiffel Tower is in the city of", "The Colosseum is in the city of", "The Eiffel Tower", "The Colosseum", " Paris", " Rome"), ("The Louvre museum is in the city of", "The Acropolis is in the city of", "The Louvre museum", "The Acropolis", " Paris", " Athens"), ("Berlin is the capital city of", "Madrid is the capital city of", "Berlin", "Madrid", " Germany", " Spain"), ("Tokyo is a city in the country of", "Cairo is a city in the country of", "Tokyo", "Cairo", " Japan", " Egypt"), ("Seoul is a city in the country of", "Lisbon is a city in the country of", "Seoul", "Lisbon", " Korea", " Portugal"), ] def subject_len(prompt, subject): n = len(tok(subject)["input_ids"]) assert tok(prompt)["input_ids"][:n] == tok(subject)["input_ids"], "subject is not a token prefix" return n for cp, xp, cs, xs, a, b in PAIRS: print(f"{len(tok(cp)['input_ids']):2d}/{len(tok(xp)['input_ids']):2d} tokens, " f"subject ends at {subject_len(cp, cs)}/{subject_len(xp, xs)} " f"{cs!r} vs {xs!r} -> {a!r}/{b!r}")

10/ 9 tokens, subject ends at 5/4  'The Eiffel Tower' vs 'The Colosseum' -> ' Paris'/' Rome'
 9/ 8 tokens, subject ends at 4/3  'The Louvre museum' vs 'The Acropolis' -> ' Paris'/' Athens'
 7/ 7 tokens, subject ends at 2/2  'Berlin' vs 'Madrid' -> ' Germany'/' Spain'
 9/ 9 tokens, subject ends at 2/2  'Tokyo' vs 'Cairo' -> ' Japan'/' Egypt'
 9/10 tokens, subject ends at 2/3  'Seoul' vs 'Lisbon' -> ' Korea'/' Portugal'

Three of the five pairs have prompts of different lengths, and every pair has subjects that tokenise to different lengths. None of that needs handling: the state we copy has the same shape whatever produced it.

Step 1: a layer sweep, with three controls. At a fixed split point — three tokens past the subject — we patch each of the 24 layers in turn and record restoration. Alongside the real patch we run:

• self-patch: write the corrupt run's own state back. This must be exactly zero at every layer; if it isn't, the driver is leaking something.
• norm-matched random: a Gaussian state rescaled to the norm of the clean state. This is what "any large perturbation at this layer" buys you.
• zero: erase the layer's state entirely. Together with the random control this brackets how much of the effect is mere damage.
• conv-state only: swap the convolutional memory instead of the recurrent memory, as a check that the effect is in the state we think it is in.

In [23]:

CONDITIONS = ("recurrent", "self", "random", "zero", "conv")

def restoration_sweep(offset=3, conditions=CONDITIONS, layers=None, seed=0):
    """Patch one layer's state at `offset` tokens past the subject; return restoration per pair, per layer."""
    layers = list(range(N_LAYERS)) if layers is None else list(layers)
    g = torch.Generator(device=DEVICE).manual_seed(seed)
    out = {c: [[0.0] * len(layers) for _ in PAIRS] for c in conditions}

    for pi, (cp, xp, cs, xs, a, b) in enumerate(PAIRS):
        ci, xi = encode(cp), encode(xp)
        good, bad = tok(a)["input_ids"][0], tok(b)["input_ids"][0]
        split_c = subject_len(cp, cs) + offset
        split_x = subject_len(xp, xs) + offset
        k_x = xi.shape[1] - split_x
        ld = lambda lg: logit_diff(lg, good, bad)

        ld_clean = ld(continuation(ci, ci.shape[1] - split_c))
        ld_corrupt = ld(continuation(xi, k_x))

        cc, _ = run(ci[:, :split_c])
        xc, _ = run(xi[:, :split_x])
        clean_snap = [(rec(cc, L).clone(), conv(cc, L).clone()) for L in range(N_LAYERS)]
        corrupt_snap = [rec(xc, L).clone() for L in range(N_LAYERS)]

        for li, L in enumerate(layers):
            h = clean_snap[L][0]
            r = torch.randn(h.shape, generator=g, device=h.device, dtype=h.dtype)
            r = r * (h.norm() / r.norm())
            edits = {
                "recurrent": lambda c, v=h, L=L: rec(c, L).copy_(v),
                "self": lambda c, v=corrupt_snap[L], L=L: rec(c, L).copy_(v),
                "random": lambda c, v=r, L=L: rec(c, L).copy_(v),
                "zero": lambda c, L=L: rec(c, L).zero_(),
                "conv": lambda c, v=clean_snap[L][1], L=L: conv(c, L).copy_(v),
            }
            for cond in conditions:
                lg = continuation(xi, k_x, edits[cond])
                out[cond][pi][li] = (ld(lg) - ld_corrupt) / (ld_clean - ld_corrupt)
    return out

layer_sweep = restoration_sweep(offset=3)
means = {c: [statistics.mean(layer_sweep[c][p][L] for p in range(len(PAIRS))) for L in range(N_LAYERS)]
         for c in layer_sweep}

for c in CONDITIONS:
    print(f"{c:>9}", " ".join(f"{v:+.2f}" for v in means[c]))

CONDITIONS = ("recurrent", "self", "random", "zero", "conv") def restoration_sweep(offset=3, conditions=CONDITIONS, layers=None, seed=0): """Patch one layer's state at `offset` tokens past the subject; return restoration per pair, per layer.""" layers = list(range(N_LAYERS)) if layers is None else list(layers) g = torch.Generator(device=DEVICE).manual_seed(seed) out = {c: [[0.0] * len(layers) for _ in PAIRS] for c in conditions} for pi, (cp, xp, cs, xs, a, b) in enumerate(PAIRS): ci, xi = encode(cp), encode(xp) good, bad = tok(a)["input_ids"][0], tok(b)["input_ids"][0] split_c = subject_len(cp, cs) + offset split_x = subject_len(xp, xs) + offset k_x = xi.shape[1] - split_x ld = lambda lg: logit_diff(lg, good, bad) ld_clean = ld(continuation(ci, ci.shape[1] - split_c)) ld_corrupt = ld(continuation(xi, k_x)) cc, _ = run(ci[:, :split_c]) xc, _ = run(xi[:, :split_x]) clean_snap = [(rec(cc, L).clone(), conv(cc, L).clone()) for L in range(N_LAYERS)] corrupt_snap = [rec(xc, L).clone() for L in range(N_LAYERS)] for li, L in enumerate(layers): h = clean_snap[L][0] r = torch.randn(h.shape, generator=g, device=h.device, dtype=h.dtype) r = r * (h.norm() / r.norm()) edits = { "recurrent": lambda c, v=h, L=L: rec(c, L).copy_(v), "self": lambda c, v=corrupt_snap[L], L=L: rec(c, L).copy_(v), "random": lambda c, v=r, L=L: rec(c, L).copy_(v), "zero": lambda c, L=L: rec(c, L).zero_(), "conv": lambda c, v=clean_snap[L][1], L=L: conv(c, L).copy_(v), } for cond in conditions: lg = continuation(xi, k_x, edits[cond]) out[cond][pi][li] = (ld(lg) - ld_corrupt) / (ld_clean - ld_corrupt) return out layer_sweep = restoration_sweep(offset=3) means = {c: [statistics.mean(layer_sweep[c][p][L] for p in range(len(PAIRS))) for L in range(N_LAYERS)] for c in layer_sweep} for c in CONDITIONS: print(f"{c:>9}", " ".join(f"{v:+.2f}" for v in means[c]))

recurrent +0.00 +0.00 +0.01 -0.00 -0.00 +0.00 -0.00 -0.00 -0.00 +0.00 -0.00 -0.01 -0.00 -0.00 -0.01 -0.00 -0.00 +0.00 +0.04 +0.86 +0.11 +0.00 +0.00 +0.00
     self +0.00 +0.00 +0.00 +0.00 +0.00 +0.00 +0.00 +0.00 +0.00 +0.00 +0.00 +0.00 +0.00 +0.00 +0.00 +0.00 +0.00 +0.00 +0.00 +0.00 +0.00 +0.00 +0.00 +0.00
   random +0.05 +0.01 +0.01 +0.00 +0.00 +0.01 +0.01 -0.01 -0.00 -0.01 +0.02 -0.00 +0.01 -0.01 +0.02 -0.01 +0.03 +0.01 +0.05 +0.29 +0.04 +0.05 -0.01 +0.01
     zero +0.07 +0.01 +0.01 +0.01 +0.00 +0.00 +0.01 -0.00 -0.01 -0.01 +0.02 +0.00 +0.01 -0.01 +0.02 -0.01 +0.02 +0.01 +0.04 +0.30 +0.01 +0.04 -0.00 -0.00
     conv +0.00 +0.00 +0.00 -0.00 +0.00 -0.00 -0.00 -0.00 -0.00 +0.00 -0.00 +0.00 -0.00 +0.00 +0.00 +0.00 +0.00 -0.00 +0.00 +0.00 +0.00 +0.00 -0.00 +0.00

In [24]:

LABELS = {
    "recurrent": "recurrent state (clean into corrupt)",
    "self": "self-patch (anchor)",
    "random": "norm-matched random (control)",
    "zero": "zeroed state (control)",
    "conv": "conv state only (control)",
}

df = pd.DataFrame(
    [{"layer": L, "restoration": means[c][L], "condition": LABELS[c]}
     for c in CONDITIONS for L in range(N_LAYERS)]
)
fig = px.line(df, x="layer", y="restoration", color="condition", markers=True,
              title="Restoration from patching one layer's state<br>"
                    "<sup>mean over 5 pairs, split point 3 tokens past the subject</sup>")
fig.update_layout(width=880, height=460, legend=dict(orientation="h", y=-0.22))
fig.show()

LABELS = { "recurrent": "recurrent state (clean into corrupt)", "self": "self-patch (anchor)", "random": "norm-matched random (control)", "zero": "zeroed state (control)", "conv": "conv state only (control)", } df = pd.DataFrame( [{"layer": L, "restoration": means[c][L], "condition": LABELS[c]} for c in CONDITIONS for L in range(N_LAYERS)] ) fig = px.line(df, x="layer", y="restoration", color="condition", markers=True, title="Restoration from patching one layer's state<br>" "^{mean over 5 pairs, split point 3 tokens past the subject}") fig.update_layout(width=880, height=460, legend=dict(orientation="h", y=-0.22)) fig.show()