Natural Language Autoencoders: explain and reconstruct an activation

A Natural Language Autoencoder (NLA) is a pair of fine-tuned language models trained to send a residual-stream activation through natural language and back. The activation verbalizer (AV) reads one activation, injected as a single input embedding, and writes an explanation. The activation reconstructor (AR) reads that explanation and predicts the original activation. The two were trained together, so the reconstruction quality measures how much of the vector survives the trip through text.

This notebook runs that round trip with TransformerBridge hooks and generation, without a model server. The mechanics section below executes the same helpers on a tiny Qwen2 checkpoint and checks them against plain Hugging Face. It demonstrates the inference protocol of the released models; it is not a replication of the paper’s training or evaluation.

Reconstruction fidelity shows that a vector is preserved, not that the explanation is faithful to what a human would read into it. A reconstruction can be accurate while the text still misleads.

The method and the released models are described in Natural Language Autoencoders Produce Unsupervised Explanations of LLM Activations. The authors’ reference implementation serves the AV through SGLang embedding injection; this notebook does not import or depend on it.

The round trip

Step	Model	Input	Output
Extract	target model	a prompt and one token position	the raw output of block \(L\) at that position, a vector \(h \in \mathbb{R}^{d}\)
Verbalize	AV	a fixed chat prompt with one placeholder token	an explanation inside <explanation>...</explanation>
Reconstruct	AR	the explanation text, in a fixed raw template	a vector \(\hat h \in \mathbb{R}^{d}\)
Score	none	\(h\) and \(\hat h\)	cosine similarity and direction-normalized MSE

The AV is a full copy of the target architecture. Before its forward pass, the embedding row of the placeholder token is replaced by \(h\), rescaled to a fixed norm recorded with the artifact. The AR is the target architecture truncated to its first \(L + 1\) blocks, followed by a separately released bias-free linear value head applied to the final token’s block-\(L\) output. Its final norm and LM head are never used.

The Patchscopes demo also patches an activation into a generating model, but uses an untrained inspection prompt and has no reconstruction step.

Requirements

The mechanics section runs on CPU in seconds and downloads only a tiny Qwen2 checkpoint. Every helper is defined in this notebook on top of public TransformerBridge methods: run_with_cache, hooks, and generate.

The released section at the end runs the authors’ Qwen2.5-7B-Instruct models and is opt-in: set the environment variable TL_RUN_RELEASED_NLA=1 before starting the kernel. It downloads three checkpoints totalling 38.6 GiB and loads one at a time in bfloat16. On the machine that produced the saved outputs (an Apple M3 Max with 48 GiB of memory, CPU only), the section took about 3.2 minutes, excluding downloads, with a peak resident memory of 14.1 GiB. Its cells are skipped by the notebook tests, so the outputs shown are from that run.

# NBVAL_IGNORE_OUTPUT
import gc
import math
import os
import re
import sys

os.environ["HF_HUB_DISABLE_PROGRESS_BARS"] = "1"
os.environ["TOKENIZERS_PARALLELISM"] = "false"

import pandas as pd
import torch
import yaml
from huggingface_hub import hf_hub_download
from huggingface_hub.utils import logging as hf_hub_logging
from IPython.display import display
from safetensors.torch import load_file
from transformers import AutoModelForCausalLM, AutoTokenizer
from transformers.utils import logging as hf_logging

from transformer_lens.model_bridge import TransformerBridge

# Loading messages and progress bars would otherwise land in the checked cell outputs.
hf_hub_logging.set_verbosity_error()
hf_logging.set_verbosity_error()
hf_logging.disable_progress_bar()

DEVICE = "cpu"  # Stable saved outputs and portable notebook validation.
RUN_RELEASED = os.environ.get("TL_RUN_RELEASED_NLA") == "1"
EXPLANATION_OPEN, EXPLANATION_CLOSE = "<explanation>", "</explanation>"
pd.set_option("display.max_colwidth", None)
_ = torch.manual_seed(0)

Reading the artifact sidecar

Each released AV and AR repository ships an nla_meta.yaml sidecar that records the contract the models were trained with: the extraction layer, the width, the placeholder character and its neighboring token ids, the injection norm, the prompt templates, and the AR’s expected prompt suffix. Every later step reads these values instead of hardcoding them, so check_sidecar validates them against the live tokenizer first. A tokenizer that splits the placeholder differently, or a template that drifted, would otherwise produce plausible-looking text from the wrong input.

The reconstruction metric is direction-normalized: both vectors are rescaled to norm \(\sqrt{d}\) before comparison. The sidecar records this as mse_scale, which is checked to equal \(\sqrt{d}\) to within its serialized precision.

def check_sidecar(meta, role, tokenizer, d_model, layer):
    # Validate a sidecar against its role and the live tokenizer, and return its key values.
    assert meta["kind"] == "nla_model" and meta["schema_version"] == 2, "unknown sidecar schema"
    assert meta["role"] == role, f"expected a {role} sidecar, got {meta['role']}"
    assert meta["d_model"] == d_model and meta["extraction_layer_index"] == layer
    extraction, tokens, templates = meta["extraction"], meta["tokens"], meta["prompt_templates"]
    assert math.isclose(extraction["mse_scale"], math.sqrt(d_model), rel_tol=1e-12)

    char = tokens["injection_char"]
    char_ids = tokenizer.encode(char, add_special_tokens=False)
    assert char_ids == [tokens["injection_token_id"]], "placeholder is not one recorded token"
    assert char_ids[0] != tokenizer.unk_token_id, "placeholder encodes to the unknown token"
    assert templates["av"].count("{injection_char}") == 1
    assert templates["ar"].count("{explanation}") == 1

    if role == "av":
        norm = extraction["injection_scale"]
        assert isinstance(norm, (int, float)) and math.isfinite(norm) and norm > 0
        assert tokens["critic_suffix_ids"] is None
    else:
        assert extraction["injection_scale"] is None
        probe = tokenizer.encode(templates["ar"].format(explanation="x"), add_special_tokens=True)
        suffix = tokens["critic_suffix_ids"]
        assert probe[-len(suffix) :] == suffix, "reconstructor prompt suffix drifted"
    return meta

Extracting an activation

blocks.{L}.hook_out is the raw residual stream after zero-based block \(L\), the same tensor as Hugging Face’s hidden_states[L + 1]. The Bridge must be booted without compatibility mode so that this is the basis the NLA was trained on. Only that one hook is cached, and the selected row is copied to the CPU in float32 so the target model can be released before the AV is loaded.

def extract(bridge, tokens, layer, position=-1):
    # Return the detached block-`layer` output at one token position, and the resolved position.
    assert not bridge.compatibility_mode, "NLAs read the raw residual basis; boot a raw Bridge"
    assert tokens.ndim == 2 and tokens.shape[0] == 1 and tokens.shape[1] > 0, "one unpadded row"
    seq_len = tokens.shape[1]
    assert -seq_len <= position < seq_len, f"position {position} outside {seq_len} tokens"
    hook_name = f"blocks.{layer}.hook_out"
    with torch.no_grad():
        _, cache = bridge.run_with_cache(tokens, names_filter=hook_name)
    vector = cache[hook_name][0, position].detach().to("cpu", torch.float32).clone()
    del cache
    return vector, position % seq_len

Injecting it into the verbalizer

The AV prompt contains the placeholder character exactly once. make_injection_hook rescales the activation to the sidecar’s injection norm in float64, casts it to the embedding’s dtype, and overwrites only the placeholder row at hook_embed.

Generation needs care. With the KV cache, the first forward pass (the prefill) covers the whole prompt and later passes embed only the newest token. Without the cache, every pass recomputes the full prefix from position 0. The hook therefore injects on every pass that reaches the placeholder position and leaves the others unchanged, without being told which mode generate uses. This matters because generate turns the cache off by itself for architectures that have none, so use_past_kv_cache=True does not guarantee cached passes. The hook records each call’s sequence length and whether it injected, so the mechanics checks below can confirm both patterns. A new hook is built for every generation, so no state carries over between calls.

def make_injection_hook(vector, position, norm):
    # Build a hook_embed hook that writes the rescaled vector at the placeholder position.
    vector = vector.detach().to("cpu", torch.float64)
    vector_norm = vector.norm()
    assert torch.isfinite(vector).all() and vector_norm > 0, "need a finite nonzero vector"
    scaled = vector / vector_norm * norm
    calls = []

    def hook(embeddings, hook):
        # Every pass that starts at position 0 and reaches the placeholder must carry the vector:
        # the prefill, and each full-prefix pass when generation runs uncached. A cached decode
        # step embeds only the newest token, so it never reaches the placeholder.
        inject = embeddings.shape[1] > position
        calls.append((embeddings.shape[1], inject))
        if inject:
            embeddings = embeddings.clone()
            embeddings[:, position] = scaled.to(embeddings.device, embeddings.dtype)
        return embeddings

    return hook, calls

Verbalizing

render_av_prompt applies the AV’s chat template once, with the generation prompt appended, and locates the single placeholder token. It checks the neighboring token ids against the sidecar, which catches a template or tokenizer change that would shift the injection site. verbalize then samples a continuation inside bridge.hooks(...), which removes the injection hook on exit, including when generation raises. Sampling settings are explicit and bounded; pass do_sample=False for greedy decoding.

def render_av_prompt(tokenizer, meta):
    # Tokenize the AV chat prompt once and return its ids and the placeholder position.
    tokens = meta["tokens"]
    content = meta["prompt_templates"]["av"].format(injection_char=tokens["injection_char"])
    rendered = tokenizer.apply_chat_template(
        [{"role": "user", "content": content}], tokenize=True, add_generation_prompt=True
    )
    ids = list(rendered["input_ids"] if hasattr(rendered, "keys") else rendered)
    matches = [i for i, token in enumerate(ids) if token == tokens["injection_token_id"]]
    assert len(matches) == 1, "the AV prompt must contain the placeholder exactly once"
    position = matches[0]
    assert 0 < position < len(ids) - 1, "the placeholder needs a token on each side"
    assert ids[position - 1] == tokens["injection_left_neighbor_id"], "left neighbor drifted"
    assert ids[position + 1] == tokens["injection_right_neighbor_id"], "right neighbor drifted"
    return torch.tensor([ids]), position


def verbalize(
    av,
    tokenizer,
    meta,
    vector,
    *,
    max_new_tokens=256,
    do_sample=True,
    temperature=1.0,
    use_past_kv_cache=True,
    seed=0,
):
    # Generate an explanation of `vector` and return the raw continuation and the parse result.
    assert not av.compatibility_mode, "the AV must run on raw weights"
    assert 0 < max_new_tokens and (not do_sample or temperature > 0)
    prompt, position = render_av_prompt(tokenizer, meta)
    assert prompt.shape[1] + max_new_tokens <= av.cfg.n_ctx, "prompt plus budget exceeds n_ctx"
    norm = meta["extraction"]["injection_scale"]
    hook, calls = make_injection_hook(vector, position, norm)
    torch.manual_seed(seed)
    with torch.no_grad(), av.hooks(fwd_hooks=[("hook_embed", hook)]):
        output = av.generate(
            prompt.to(av.cfg.device),
            max_new_tokens=max_new_tokens,
            do_sample=do_sample,
            temperature=temperature,
            use_past_kv_cache=use_past_kv_cache,
            return_type="tokens",
            verbose=False,
        )
    assert calls and calls[0][1], "the prefill never reached the placeholder"
    continuation = output[0, prompt.shape[1] :].tolist()
    raw_text = tokenizer.decode(continuation, skip_special_tokens=True)
    text, status = parse_explanation(raw_text)
    return {
        "raw_text": raw_text,
        "token_ids": continuation,
        "text": text,
        "status": status,
        "hook_calls": calls,
    }

Parsing the explanation

The AV is trained to wrap its explanation in <explanation>...</explanation>. parse_explanation accepts the stripped inner text of exactly one complete region and reports every other outcome as a named failure: missing, multiple, nested, incomplete, or empty. A failure returns None rather than falling back to the raw continuation, because reconstructing arbitrary model output would report a score for text the AR was never trained on. Failures are results and are reported, not dropped.

This is stricter than the authors’ reference client, which takes the first <explanation> region it finds and, when there is none, returns the raw continuation with a warning. Parse-failure counts from the two are therefore not directly comparable.

def parse_explanation(text):
    # Return (inner text, "ok") for exactly one complete region, otherwise (None, reason).
    tags = sorted(
        [(m.start(), +1) for m in re.finditer(re.escape(EXPLANATION_OPEN), text)]
        + [(m.start(), -1) for m in re.finditer(re.escape(EXPLANATION_CLOSE), text)]
    )
    if not tags:
        return None, "missing"
    depth, regions, nested = 0, 0, False
    for _, step in tags:
        depth += step
        if depth < 0:
            return None, "incomplete"
        nested |= depth > 1
        regions += step < 0 and depth == 0
    if depth != 0:
        return None, "incomplete"
    if nested:
        return None, "nested"
    if regions > 1:
        return None, "multiple"
    start = text.index(EXPLANATION_OPEN) + len(EXPLANATION_OPEN)
    inner = text[start : text.index(EXPLANATION_CLOSE)].strip()
    return (inner, "ok") if inner else (None, "empty")

Reconstructing

The AR reads the explanation in a fixed raw template, not a chat prompt, tokenized once with the tokenizer’s own special-token policy. reconstruct checks that the rendered prompt ends with the sidecar’s recorded suffix, reads the final token’s block-\(L\) output with extract, and applies the value head.

The AR checkpoint has only \(L + 1\) blocks. Its final norm and LM head are absent from the released weights, so Hugging Face initializes them itself (the norm weight to ones, the LM head randomly); reading blocks.{L}.hook_out never touches them. Hugging Face’s last hidden_states entry gives the wrong vector, because it has the final norm applied. stop_at_layer=L + 1 stops after the last block and returns the same raw residual as blocks.{L}.hook_out; the helpers read the hook so that the target and the AR share extract. Input that still contains explanation tags is rejected instead of silently stripped, since the AR was trained on the inner text only.

def reconstruct(ar, tokenizer, meta, value_head, text):
    # Predict the activation described by `text` with the AR and its value head.
    assert text.strip(), "need a nonempty explanation"
    assert EXPLANATION_OPEN not in text and EXPLANATION_CLOSE not in text, "pass the inner text"
    layer = meta["extraction_layer_index"]
    assert ar.cfg.n_layers == layer + 1, "the AR must be truncated to layer + 1 blocks"
    rendered = meta["prompt_templates"]["ar"].format(explanation=text)
    ids = tokenizer.encode(rendered, add_special_tokens=True)
    suffix = meta["tokens"]["critic_suffix_ids"]
    assert ids[-len(suffix) :] == suffix, "the explanation changed the prompt suffix"
    assert len(ids) <= ar.cfg.n_ctx, "explanation exceeds the AR context"
    hidden, _ = extract(ar, torch.tensor([ids]), layer, position=-1)
    weight = value_head.weight
    with torch.no_grad():
        return value_head(hidden.to(weight.device, weight.dtype)).float().cpu()

Scoring

Both vectors are rescaled to norm \(\sqrt{d}\), matching the sidecar’s mse_scale, and compared in float64:

\[\cos(h, \hat h) = \frac{h \cdot \hat h}{\lVert h\rVert\,\lVert \hat h\rVert}, \qquad \mathrm{nMSE}(h, \hat h) = \frac{1}{d}\left\lVert \sqrt{d}\,\frac{h}{\lVert h\rVert} - \sqrt{d}\,\frac{\hat h}{\lVert \hat h\rVert}\right\rVert^2 = 2\,(1 - \cos).\]

score asserts that identity. These are directional diagnostics for one vector: they are not the raw reconstruction error, and they are not a dataset-level fraction of variance explained.

def score(original, reconstructed):
    # Cosine similarity and direction-normalized MSE of two activation vectors.
    a, b = original.double().flatten(), reconstructed.double().flatten()
    assert a.shape == b.shape, "vectors must have the same width"
    assert torch.isfinite(a).all() and torch.isfinite(b).all() and a.norm() > 0 and b.norm() > 0
    scale = math.sqrt(a.numel())
    a_hat, b_hat = a / a.norm() * scale, b / b.norm() * scale
    # Rounding can push the ratio slightly outside [-1, 1] for (anti)parallel vectors.
    cosine = min(1.0, max(-1.0, float(torch.dot(a, b) / (a.norm() * b.norm()))))
    normalized_mse = float(((a_hat - b_hat) ** 2).mean())
    assert math.isclose(normalized_mse, 2 * (1 - cosine), rel_tol=1e-9, abs_tol=1e-12)
    return {"cosine": cosine, "normalized_mse": normalized_mse}

Checking the mechanics on a tiny model

The released models are three 7B checkpoints, too large for notebook CI. This section runs every helper above on the two-block trl-internal-testing/tiny-Qwen2ForCausalLM-2.5 checkpoint, which shares Qwen2.5’s architecture and tokenizer, and compares extraction, embedding injection, and the reconstructor read-out against plain Hugging Face, and cached against uncached generation. The tiny model was never trained as an NLA, so its generated text is meaningless and its scores carry no information; only the assertions matter.

The sidecars below are synthetic but follow the released schema. The verbalizer and target share the tiny checkpoint, and the reconstructor is the same checkpoint truncated to \(L + 1 = 1\) block with a seeded random value head. Each Hugging Face reference is loaded separately with eager attention, because boot_transformers(hf_model=...) routes the wrapped module’s own forward pass through the Bridge’s hooks. The tiny checkpoint is pinned to a commit, like the released artifacts, because the checks print values derived from it, such as prompt lengths and the placeholder position.

TINY = "trl-internal-testing/tiny-Qwen2ForCausalLM-2.5"
TINY_REVISION = "ce8d0bf270b28c8fab026ced69fe8aa14b0f0eda"
TINY_LAYER = 0
tiny_tokenizer = AutoTokenizer.from_pretrained(TINY, revision=TINY_REVISION)


def load_hf(**kwargs):
    # Plain Hugging Face reference of the pinned tiny checkpoint.
    model = AutoModelForCausalLM.from_pretrained(
        TINY, revision=TINY_REVISION, attn_implementation="eager", dtype=torch.float32, **kwargs
    )
    return model.to(DEVICE).eval()


def boot_tiny(**kwargs):
    bridge = TransformerBridge.boot_transformers(
        TINY, revision=TINY_REVISION, device=DEVICE, dtype=torch.float32, **kwargs
    )
    return bridge.eval()


tiny_target = boot_tiny()
tiny_av = boot_tiny()
tiny_ar = boot_tiny(hf_model=load_hf(num_hidden_layers=TINY_LAYER + 1))
hf_reference = load_hf()
TINY_D_MODEL = tiny_target.cfg.d_model

print(f"Tiny target and AV: {tiny_target.cfg.n_layers} blocks, d_model={TINY_D_MODEL}.")
print(f"Tiny AR: {tiny_ar.cfg.n_layers} block (truncated after block {TINY_LAYER}).")

Tiny target and AV: 2 blocks, d_model=8.
Tiny AR: 1 block (truncated after block 0).

def synthetic_sidecar(role, tokenizer, d_model, layer, av_template, ar_template, char):
    # Build a sidecar with the released schema, deriving token ids from the live tokenizer.
    (char_id,) = tokenizer.encode(char, add_special_tokens=False)
    rendered = tokenizer.apply_chat_template(
        [{"role": "user", "content": av_template.format(injection_char=char)}],
        tokenize=True,
        add_generation_prompt=True,
    )
    ids = list(rendered["input_ids"] if hasattr(rendered, "keys") else rendered)
    position = ids.index(char_id)
    # The recorded suffix is the token tail every explanation shares: punctuation at the end of
    # an explanation can merge with the start of the closing tag.
    probes = [
        tokenizer.encode(ar_template.format(explanation=end), add_special_tokens=True)
        for end in ("x", "x.", "x!", "x)")
    ]
    length = 0
    while length < min(map(len, probes)) and len({p[-length - 1] for p in probes}) == 1:
        length += 1
    suffix = probes[0][-length:]
    return {
        "kind": "nla_model",
        "schema_version": 2,
        "role": role,
        "d_model": d_model,
        "extraction_layer_index": layer,
        "extraction": {
            "injection_scale": 150.0 if role == "av" else None,
            "mse_scale": math.sqrt(d_model),
        },
        "tokens": {
            "injection_char": char,
            "injection_token_id": char_id,
            "injection_left_neighbor_id": ids[position - 1],
            "injection_right_neighbor_id": ids[position + 1],
            "critic_suffix_ids": suffix if role == "ar" else None,
        },
        "prompt_templates": {"av": av_template, "ar": ar_template},
    }


SYNTHETIC_TEMPLATES = {
    "av_template": "Explain this concept:\n{injection_char}\nAnswer inside explanation tags.",
    "ar_template": "Summary of the following text: <text>{explanation}</text> <summary>",
    "char": "\u2605",
}
tiny_meta = {
    role: check_sidecar(
        synthetic_sidecar(role, tiny_tokenizer, TINY_D_MODEL, TINY_LAYER, **SYNTHETIC_TEMPLATES),
        role,
        tiny_tokenizer,
        TINY_D_MODEL,
        TINY_LAYER,
    )
    for role in ("av", "ar")
}
torch.manual_seed(0)
tiny_value_head = torch.nn.Linear(TINY_D_MODEL, TINY_D_MODEL, bias=False).eval()

tiny_prompt = tiny_tokenizer("The auditor asked about reasoning.", return_tensors="pt").input_ids
av_prompt, av_position = render_av_prompt(tiny_tokenizer, tiny_meta["av"])
print("Both synthetic sidecars pass check_sidecar.")
print(f"AV prompt: {av_prompt.shape[1]} tokens, placeholder at position {av_position}.")

Both synthetic sidecars pass check_sidecar.
AV prompt: 41 tokens, placeholder at position 29.

(a) Extraction and injection match Hugging Face

extract must return Hugging Face’s hidden_states[L + 1] at the chosen position, and the AV’s prefill logits with the injection hook must equal Hugging Face logits computed from inputs_embeds with the same row replaced. Both sides use the same rescaling formula and cast, so any difference would come from the hook placement itself.

tiny_vector, resolved = extract(tiny_target, tiny_prompt, TINY_LAYER, position=-1)
with torch.no_grad():
    hf_hidden = hf_reference(tiny_prompt, output_hidden_states=True).hidden_states
print(f"Resolved position -1 to {resolved} of {tiny_prompt.shape[1]} tokens.")
extract_matches = torch.equal(tiny_vector, hf_hidden[TINY_LAYER + 1][0, -1])
print(f"extract equals HF hidden_states[L + 1]: {extract_matches}")

norm = tiny_meta["av"]["extraction"]["injection_scale"]
hook, calls = make_injection_hook(tiny_vector, av_position, norm)
with torch.no_grad(), tiny_av.hooks(fwd_hooks=[("hook_embed", hook)]):
    bridge_logits = tiny_av(av_prompt)
with torch.no_grad():
    embeddings = hf_reference.get_input_embeddings()(av_prompt)
    injected = embeddings.clone()
    injected[:, av_position] = (tiny_vector.double() / tiny_vector.double().norm() * norm).float()
    hf_logits = hf_reference(inputs_embeds=injected).logits
changed_rows = (injected != embeddings).any(-1)[0].nonzero().flatten().tolist()
print(f"Only the placeholder row is replaced: {changed_rows == [av_position]}")
logits_match = torch.equal(bridge_logits, hf_logits)
print(f"Bridge prefill logits equal HF inputs_embeds logits: {logits_match}")

Resolved position -1 to 5 of 6 tokens.
extract equals HF hidden_states[L + 1]: True
Only the placeholder row is replaced: True
Bridge prefill logits equal HF inputs_embeds logits: True

(b) Cached and uncached generation agree

Greedy generation with and without the KV cache must produce the same tokens. The hook’s call log shows the two injection patterns: one injecting prefill followed by one-token decode steps with the cache, and an injecting full-prefix pass at every step without it. The hook is not told the caching mode, so both patterns follow only from which passes reach the placeholder; the uncached pattern is also what happens when generate turns the cache off by itself.

runs = {
    cached: verbalize(
        tiny_av,
        tiny_tokenizer,
        tiny_meta["av"],
        tiny_vector,
        max_new_tokens=6,
        do_sample=False,
        use_past_kv_cache=cached,
    )
    for cached in (True, False)
}
prompt_len = av_prompt.shape[1]
cached_calls, uncached_calls = runs[True]["hook_calls"], runs[False]["hook_calls"]
tokens_agree = runs[True]["token_ids"] == runs[False]["token_ids"]
print(f"Greedy tokens agree with and without the KV cache: {tokens_agree}")
print(
    "Cached: one injecting prefill, then non-injecting one-token steps: "
    f"{cached_calls[0] == (prompt_len, True) and all(c == (1, False) for c in cached_calls[1:])}"
)
print(
    "Uncached: every pass covers the growing prefix and injects: "
    f"{all(c == (prompt_len + i, True) for i, c in enumerate(uncached_calls))}"
)
print(f"Tiny-model parse status (meaningless text is expected): {runs[True]['status']}")

Greedy tokens agree with and without the KV cache: True
Cached: one injecting prefill, then non-injecting one-token steps: True
Uncached: every pass covers the growing prefix and injects: True
Tiny-model parse status (meaningless text is expected): missing

(c) The truncated reconstructor reads the raw block output

The AR’s block-\(L\) output at the final token must equal the full model’s hidden_states[L + 1], because the truncated checkpoint shares those blocks. The truncated model’s own last hidden_states entry differs: Hugging Face applies the final norm to it. stop_at_layer=L + 1 stops after the last block and returns the same raw residual as the hook. reconstruct must equal the value head applied to the raw read-out.

canonical = "A question about how an auditor reasons."
rendered = tiny_meta["ar"]["prompt_templates"]["ar"].format(explanation=canonical)
ar_tokens = torch.tensor([tiny_tokenizer.encode(rendered, add_special_tokens=True)])
ar_hidden, _ = extract(tiny_ar, ar_tokens, TINY_LAYER)
truncated_hf = load_hf(num_hidden_layers=TINY_LAYER + 1)
with torch.no_grad():
    full_hidden = hf_reference(ar_tokens, output_hidden_states=True).hidden_states
    truncated_last = truncated_hf(ar_tokens, output_hidden_states=True).hidden_states[-1]
    stopped = tiny_ar(ar_tokens, stop_at_layer=TINY_LAYER + 1)
readout_matches = torch.equal(ar_hidden, full_hidden[TINY_LAYER + 1][0, -1])
normed_differs = not torch.equal(ar_hidden, truncated_last[0, -1])
print(f"AR read-out equals full-model hidden_states[L + 1]: {readout_matches}")
print(f"Truncated HF last hidden_states is normed, so it differs: {normed_differs}")
stopped_matches = torch.equal(stopped[0, -1], ar_hidden)
print(f"stop_at_layer=L + 1 returns the same block-L read-out: {stopped_matches}")

reconstruction = reconstruct(tiny_ar, tiny_tokenizer, tiny_meta["ar"], tiny_value_head, canonical)
with torch.no_grad():
    expected = tiny_value_head(ar_hidden)
print(f"reconstruct equals value_head(read-out): {torch.equal(reconstruction, expected)}")
del truncated_hf

AR read-out equals full-model hidden_states[L + 1]: True
Truncated HF last hidden_states is normed, so it differs: True
stop_at_layer=L + 1 returns the same block-L read-out: True
reconstruct equals value_head(read-out): True

(d) Explanation parsing

Every parse outcome on a fixed set of continuations, including the tagged-input rejection in reconstruct.

parse_cases = [
    ("<explanation> A legal audit. </explanation>", "ok"),
    ("No tags at all.", "missing"),
    ("<explanation>a</explanation> and <explanation>b</explanation>", "multiple"),
    ("<explanation>a <explanation>b</explanation></explanation>", "nested"),
    ("<explanation> cut off by the token budget", "incomplete"),
    ("</explanation> closed before opening <explanation>", "incomplete"),
    ("<explanation>   </explanation>", "empty"),
]
parse_rows = []
for text, expected_status in parse_cases:
    parsed, status = parse_explanation(text)
    assert status == expected_status, (text, status)
    parse_rows.append((text, status, parsed))
display(pd.DataFrame(parse_rows, columns=["continuation", "status", "parsed text"]))

try:
    reconstruct(tiny_ar, tiny_tokenizer, tiny_meta["ar"], tiny_value_head, parse_cases[0][0])
except AssertionError as error:
    print(f"Tagged input rejected by reconstruct: {error}")

	continuation	status	parsed text
0	<explanation> A legal audit. </explanation>	ok	A legal audit.
1	No tags at all.	missing	None
2	<explanation>a</explanation> and <explanation>b</explanation>	multiple	None
3	<explanation>a <explanation>b</explanation></explanation>	nested	None
4	<explanation> cut off by the token budget	incomplete	None
5	</explanation> closed before opening <explanation>	incomplete	None
6	<explanation> </explanation>	empty	None

Tagged input rejected by reconstruct: pass the inner text

(e) Hooks are removed, including after a failure

After every helper, no hook may remain on any model. To check the failure path, a caller-owned hook that raises is attached to block 0 before verbalize runs. The injection hook must be gone after the exception, while the caller’s own hook stays until the caller removes it.

def hook_count(bridge):
    # Count torch-level forward hooks on unique hook points (aliases share one HookPoint).
    # The private `_forward_hooks` also sees hooks registered outside the HookPoint API.
    points = {id(point): point for point in bridge.hook_dict.values()}
    return sum(len(point._forward_hooks) for point in points.values())


def failing_hook(activation, hook):
    raise RuntimeError("simulated failure during generation")


hooks_left = [hook_count(m) for m in (tiny_target, tiny_av, tiny_ar)]
print(f"Hooks left after the checks above: {hooks_left}")
caller_point = tiny_av.hook_dict["blocks.0.hook_out"]
caller_point.add_hook(failing_hook)
try:
    verbalize(tiny_av, tiny_tokenizer, tiny_meta["av"], tiny_vector, max_new_tokens=2)
except RuntimeError as error:
    print(f"verbalize raised: {error}")
print(f"Only the caller's hook remains after the failure: {hook_count(tiny_av) == 1}")
caller_point.remove_hooks("fwd")
print(f"No hooks remain once the caller removes it: {hook_count(tiny_av) == 0}")

Hooks left after the checks above: [0, 0, 0]
verbalize raised: simulated failure during generation
Only the caller's hook remains after the failure: True
No hooks remain once the caller removes it: True

Tiny round trip

The full pipeline end to end. Because the tiny model cannot produce a tagged explanation, the fixed canonical explanation stands in for the verbalizer’s output. The score only shows that the pipeline composes; it measures nothing about the tiny model.

tiny_reconstruction = reconstruct(
    tiny_ar, tiny_tokenizer, tiny_meta["ar"], tiny_value_head, canonical
)
tiny_score = score(tiny_vector, tiny_reconstruction)
finite = all(math.isfinite(v) for v in tiny_score.values())
print(f"Score is finite and satisfies nMSE = 2 (1 - cos): {finite}")

del tiny_target, tiny_av, tiny_ar, hf_reference
_ = gc.collect()

Score is finite and satisfies nMSE = 2 (1 - cos): True

The released Qwen2.5-7B-Instruct round trip

This section runs the released block-20 pair at pinned revisions:

Role	Repository	Revision
target	`Qwen/Qwen2.5-7B-Instruct <https://huggingface.co/Qwen/Qwen2.5-7B-Instruct>`__	a09a354
AV	`kitft/nla-qwen2.5-7b-L20-av <https://huggingface.co/kitft/nla-qwen2.5-7b-L20-av>`__	b884691
AR	`kitft/nla-qwen2.5-7b-L20-ar <https://huggingface.co/kitft/nla-qwen2.5-7b-L20-ar>`__	e2c9e57

Everything runs in bfloat16, the precision of the released weights and of the authors’ reference. Each model is loaded, used, and released before the next one, so only one 7B model is resident at a time. Every stage compares the Bridge against plain Hugging Face at the same dtype, device, tokens, and eager attention: the Hugging Face reference is computed first, then the same module is wrapped with boot_transformers(hf_model=...), so no second copy is needed. These parity numbers measure what the Bridge changes, not the gap between bfloat16 and float32.

The cells below are skipped by the notebook tests. Run them by setting TL_RUN_RELEASED_NLA=1 before starting the kernel.

# NBVAL_SKIP
assert RUN_RELEASED, "Set TL_RUN_RELEASED_NLA=1 to download and run the released 7B models."
import time
from contextlib import contextmanager

RELEASED = {
    "target": ("Qwen/Qwen2.5-7B-Instruct", "a09a35458c702b33eeacc393d103063234e8bc28"),
    "av": ("kitft/nla-qwen2.5-7b-L20-av", "b88469162777ae6553bc14208eb0cb579336f8f4"),
    "ar": ("kitft/nla-qwen2.5-7b-L20-ar", "e2c9e57eac213d37a31612087f645ab6332c1bb6"),
}
LAYER, D_MODEL = 20, 3584
RELEASED_DTYPE = torch.bfloat16
RELEASED_DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
# Declared before running; position -1 selects each prompt's final token. The reference
# client treats the first ~10 positions as noisy, so every read position is well past them.
PROMPTS = [
    "Last summer my family finally took the trip we had been planning for years. We flew into "
    "Nice, spent a week on the coast, and then took the long train north to the capital city of",
    "def fibonacci(n):\n    # Return the n-th Fibonacci number, with fibonacci(0) == 0.\n"
    '    if n < 0:\n        raise ValueError("n must be non-negative")\n'
    "    if n <= 1:\n        return",
    "She had waited three months for an answer from the university. When the envelope finally "
    "arrived, she tore it open with shaking hands, read the first line, and burst into",
]
# Defaults of the reference client's `NLAClient._sglang_generate`.
GENERATION = {"max_new_tokens": 200, "do_sample": True, "temperature": 1.0}
SEED = 0
stage_log = []


def peak_memory_gib():
    # Process peak resident memory, plus the CUDA allocator peak when a GPU is used.
    peak = {}
    try:
        import resource

        maxrss = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
        peak["peak RSS (GiB)"] = maxrss / 2**30 if sys.platform == "darwin" else maxrss / 2**20
    except ImportError:
        pass
    if torch.cuda.is_available():
        peak["peak CUDA (GiB)"] = torch.cuda.max_memory_allocated() / 2**30
    return peak


@contextmanager
def stage(name):
    start = time.perf_counter()
    yield
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
    stage_log.append({"stage": name, "seconds": time.perf_counter() - start, **peak_memory_gib()})


def load_released(role):
    # The HF reference runs on RELEASED_DEVICE so that it and the wrapping Bridge share a device.
    repo, revision = RELEASED[role]
    model, loading_info = AutoModelForCausalLM.from_pretrained(
        repo,
        revision=revision,
        dtype=RELEASED_DTYPE,
        attn_implementation="eager",
        output_loading_info=True,
    )
    return model.to(RELEASED_DEVICE).eval(), loading_info


def wrap_released(role, hf_model):
    repo, revision = RELEASED[role]
    bridge = TransformerBridge.boot_transformers(
        repo,
        hf_model=hf_model,
        tokenizer=released_tokenizers[role],
        revision=revision,
        device=RELEASED_DEVICE,
        dtype=RELEASED_DTYPE,
    )
    return bridge.eval()


def max_abs_error(a, b):
    return float((a.detach().float().cpu() - b.detach().float().cpu()).abs().max())


print(f"Device: {RELEASED_DEVICE}, dtype: {RELEASED_DTYPE}, torch {torch.__version__}.")

Device: cpu, dtype: torch.bfloat16, torch 2.11.0.

Sidecars and value head

Both sidecars go through check_sidecar against their own repository’s tokenizer. The checks below add the released values the round trip depends on: the injection norm of 150, identical prompt templates in both sidecars, and a bias-free \(3584 \times 3584\) value head. The placeholder is U+320E PARENTHESIZED HANGUL KIYEOK A, which encodes to one token in Qwen’s vocabulary and almost never appears in ordinary text.

# NBVAL_SKIP
with stage("sidecars and value head"):
    released_tokenizers = {
        role: AutoTokenizer.from_pretrained(repo, revision=revision)
        for role, (repo, revision) in RELEASED.items()
    }
    released_meta = {}
    for role in ("av", "ar"):
        repo, revision = RELEASED[role]
        with open(hf_hub_download(repo, "nla_meta.yaml", revision=revision), encoding="utf-8") as f:
            meta = yaml.safe_load(f)
        released_meta[role] = check_sidecar(meta, role, released_tokenizers[role], D_MODEL, LAYER)
    assert released_meta["av"]["extraction"]["injection_scale"] == 150.0
    assert released_meta["av"]["prompt_templates"] == released_meta["ar"]["prompt_templates"]
    assert released_meta["ar"].get("critic", {}).get("extraction_layer_index") == LAYER

    ar_repo, ar_revision = RELEASED["ar"]
    head_state = load_file(hf_hub_download(ar_repo, "value_head.safetensors", revision=ar_revision))
    assert set(head_state) == {"weight"}, f"unexpected value-head keys {sorted(head_state)}"
    assert tuple(head_state["weight"].shape) == (D_MODEL, D_MODEL)
    head_dtype = head_state["weight"].dtype
    value_head = torch.nn.Linear(D_MODEL, D_MODEL, bias=False, dtype=RELEASED_DTYPE)
    value_head.load_state_dict(head_state)
    value_head = value_head.to(RELEASED_DEVICE).eval()
    del head_state

av_tokens, ar_tokens_meta = released_meta["av"]["tokens"], released_meta["ar"]["tokens"]
ar_suffix = released_tokenizers["ar"].decode(ar_tokens_meta["critic_suffix_ids"])
checked = [
    ("extraction layer", released_meta["av"]["extraction_layer_index"]),
    ("d_model", released_meta["av"]["d_model"]),
    (
        "placeholder",
        f"U+{ord(av_tokens['injection_char']):04X} (token {av_tokens['injection_token_id']})",
    ),
    ("injection norm (AV)", released_meta["av"]["extraction"]["injection_scale"]),
    (
        "mse_scale / sqrt(d_model)",
        released_meta["ar"]["extraction"]["mse_scale"] / math.sqrt(D_MODEL),
    ),
    ("AR template", released_meta["ar"]["prompt_templates"]["ar"]),
    ("AR suffix tokens", f"{ar_tokens_meta['critic_suffix_ids']} = {ar_suffix!r}"),
    ("value head", f"{D_MODEL} x {D_MODEL}, no bias, stored as {head_dtype}"),
]
display(pd.DataFrame(checked, columns=["checked value", "released"]).set_index("checked value"))

	released
checked value
extraction layer	20
d_model	3584
placeholder	U+320E (token 149705)
injection norm (AV)	150.0
mse_scale / sqrt(d_model)	1.0
AR template	Summary of the following text: <text>{explanation}</text> <summary>
AR suffix tokens	[1318, 29, 366, 1708, 29] = 'text> <summary>'
value head	3584 x 3584, no bias, stored as torch.bfloat16

Target: extract block-20 activations

The raw prompts are tokenized with the target’s tokenizer and no chat template; Qwen adds no BOS token. The Hugging Face reference reads hidden_states[21], which is not the last entry for the 28-block target, so it carries no final norm. The extracted vectors are kept on the CPU, and the target is released before the AV is loaded.

# NBVAL_SKIP
with stage("target: load and extract"):
    hf_target, _ = load_released("target")
    target_tokens = [
        released_tokenizers["target"](prompt, return_tensors="pt").input_ids for prompt in PROMPTS
    ]
    with torch.no_grad():
        hf_vectors = [
            hf_target(t.to(hf_target.device), output_hidden_states=True).hidden_states[LAYER + 1][
                0, -1
            ]
            for t in target_tokens
        ]
    target = wrap_released("target", hf_target)
    extractions = [extract(target, t.to(RELEASED_DEVICE), LAYER) for t in target_tokens]
    vectors = [vector for vector, _ in extractions]
    target_errors = [max_abs_error(v, h) for v, h in zip(vectors, hf_vectors)]
    final_tokens = [
        released_tokenizers["target"].decode(t[0, position])
        for t, (_, position) in zip(target_tokens, extractions)
    ]
    del target, hf_target, hf_vectors

parity = [("target block-20 output vs HF hidden_states[21]", max(target_errors))]
display(
    pd.DataFrame(
        {
            "prompt": PROMPTS,
            "tokens": [t.shape[1] for t in target_tokens],
            "final token": final_tokens,
            "activation norm": [round(float(v.norm()), 1) for v in vectors],
        }
    )
)
print(f"Largest Bridge vs HF difference over the three activations: {max(target_errors)}")

	prompt	tokens	final token	activation norm
0	Last summer my family finally took the trip we had been planning for years. We flew into Nice, spent a week on the coast, and then took the long train north to the capital city of	39	of	112.1
1	def fibonacci(n):\n # Return the n-th Fibonacci number, with fibonacci(0) == 0.\n if n < 0:\n raise ValueError("n must be non-negative")\n if n <= 1:\n return	48	return	116.8
2	She had waited three months for an answer from the university. When the envelope finally arrived, she tore it open with shaking hands, read the first line, and burst into	34	into	121.5

Largest Bridge vs HF difference over the three activations: 0.0

AV: injection parity and verbalization

The parity check injects the first activation and compares the Bridge’s prefill logits, at every prompt position, with Hugging Face logits from inputs_embeds with the same rescaled row. The authors’ client rescales in float32 while make_injection_hook rescales in float64; the check also counts how many of the 3584 injected bfloat16 values the two choices change.

Each activation is then verbalized once with its own fixed seed and the sampling defaults of the reference client’s server request (NLAClient._sglang_generate: temperature 1.0, 200 new tokens); the reference’s command-line entry point uses temperature 0.7 instead. With the KV cache, the injection applies on the prefill only, as in the reference’s single server request.

# NBVAL_SKIP
with stage("AV: load, parity, verbalize"):
    hf_av, _ = load_released("av")
    av_tokenizer, av_meta = released_tokenizers["av"], released_meta["av"]
    av_prompt, av_position = render_av_prompt(av_tokenizer, av_meta)
    injection_norm = av_meta["extraction"]["injection_scale"]

    first = vectors[0].to(torch.float64)
    row64 = (first / first.norm() * injection_norm).to(RELEASED_DTYPE)
    first32 = vectors[0].float()
    row32 = (first32 / (first32.norm() / injection_norm)).to(RELEASED_DTYPE)
    rescale_differences = int((row64 != row32).sum())

    with torch.no_grad():
        embeddings = hf_av.get_input_embeddings()(av_prompt.to(hf_av.device))
        injected = embeddings.clone()
        injected[:, av_position] = row64.to(injected.device)
        hf_av_logits = hf_av(inputs_embeds=injected).logits
    av = wrap_released("av", hf_av)
    hook, _ = make_injection_hook(vectors[0], av_position, injection_norm)
    with torch.no_grad(), av.hooks(fwd_hooks=[("hook_embed", hook)]):
        bridge_av_logits = av(av_prompt.to(RELEASED_DEVICE))
    av_error = max_abs_error(bridge_av_logits, hf_av_logits)
    argmax_agrees = bool((bridge_av_logits.argmax(-1) == hf_av_logits.argmax(-1)).all())
    del hf_av_logits, bridge_av_logits, embeddings, injected

    explanations = [
        verbalize(av, av_tokenizer, av_meta, vector, seed=SEED + i, **GENERATION)
        for i, vector in enumerate(vectors)
    ]
    del av, hf_av

parity.append(("AV prefill logits vs HF inputs_embeds", av_error))
print(f"AV prompt: {av_prompt.shape[1]} tokens, placeholder at position {av_position}.")
print(f"Largest Bridge vs HF prefill-logit difference: {av_error}")
print(f"Prefill argmax agrees at every position: {argmax_agrees}")
print(f"Injected values that differ between float64 and float32 rescaling: {rescale_differences}")
print(f"Parse statuses: {[e['status'] for e in explanations]}")

AV prompt: 125 tokens, placeholder at position 111.
Largest Bridge vs HF prefill-logit difference: 0.0
Prefill argmax agrees at every position: True
Injected values that differ between float64 and float32 rescaling: 0
Parse statuses: ['ok', 'ok', 'ok']

AR: read-out parity and reconstruction

The released AR checkpoint omits the final norm and LM head, so Hugging Face reports exactly those two as missing and initializes them itself. The reference read-out follows the authors’ NLACritic: the final norm is replaced by an identity and last_hidden_state is read at the final token. It is computed on a fixed explanation, the norm is restored, and the module is wrapped; the Bridge never calls either missing module.

# NBVAL_SKIP
with stage("AR: load, parity, reconstruct"):
    hf_ar, loading_info = load_released("ar")
    missing = set(loading_info["missing_keys"])
    assert missing == {"model.norm.weight", "lm_head.weight"}, missing
    assert hf_ar.config.num_hidden_layers == LAYER + 1
    ar_tokenizer, ar_meta = released_tokenizers["ar"], released_meta["ar"]

    fixed_explanation = "A sentence about the location of a famous landmark in a European capital."
    rendered = ar_meta["prompt_templates"]["ar"].format(explanation=fixed_explanation)
    fixed_ids = torch.tensor([ar_tokenizer.encode(rendered, add_special_tokens=True)])
    final_norm = hf_ar.model.norm
    hf_ar.model.norm = torch.nn.Identity()
    with torch.no_grad():
        reference_hidden = hf_ar.model(
            fixed_ids.to(hf_ar.device), use_cache=False
        ).last_hidden_state[0, -1]
    hf_ar.model.norm = final_norm
    ar = wrap_released("ar", hf_ar)
    bridge_hidden, _ = extract(ar, fixed_ids.to(RELEASED_DEVICE), LAYER)
    ar_error = max_abs_error(bridge_hidden, reference_hidden)

    reconstructions = [
        reconstruct(ar, ar_tokenizer, ar_meta, value_head, e["text"]) if e["text"] else None
        for e in explanations
    ]
    del ar, hf_ar, reference_hidden

parity.append(("AR block-20 read-out vs NLACritic-style HF", ar_error))
print(f"Missing AR weights reported by Hugging Face: {sorted(missing)}")
print(f"Largest Bridge vs HF read-out difference: {ar_error}")

Missing AR weights reported by Hugging Face: ['lm_head.weight', 'model.norm.weight']
Largest Bridge vs HF read-out difference: 0.0

Results

One row per declared prompt, in order, with nothing filtered. A parse failure leaves the scores empty; its raw continuation is printed below the table. Explanations are single samples, so another seed gives different text and a different score.

# NBVAL_SKIP
rows = []
for prompt, token, explanation, original, reconstruction in zip(
    PROMPTS, final_tokens, explanations, vectors, reconstructions
):
    scores = score(original, reconstruction) if reconstruction is not None else {}
    rows.append(
        {
            "prompt": prompt,
            "final token": token,
            "status": explanation["status"],
            "explanation": explanation["text"],
            "cosine": round(scores["cosine"], 3) if scores else None,
            "normalized MSE": round(scores["normalized_mse"], 3) if scores else None,
        }
    )
display(pd.DataFrame(rows))
for prompt, explanation in zip(PROMPTS, explanations):
    if explanation["status"] != "ok":
        print(f"Raw continuation for {prompt!r}:\n{explanation['raw_text']}\n")

	prompt	final token	status	explanation	cosine	normalized MSE
0	Last summer my family finally took the trip we had been planning for years. We flew into Nice, spent a week on the coast, and then took the long train north to the capital city of	of	ok	Travel blog format with casual American English tone and adventure blog structure, introducing a food dish (Paris Baguette) with cultural context cues.\n\nThe phrase "On a recent trip I visited the beautiful city of Brussels and the capital city of" implies a list of attractions or historical facts about Belgium's capital Brussels, with "Royal Belgium" signaling the official government's capital.\n\nFinal token "of" ends an incomplete noun phrase ("the capital city of"), part of a parallel compound phrase ("the capital city and the capital city of"), requiring a proper noun like "Brussels" or "Belgium" or "the Kingdom," completing the phrase about the government's headquarters.	0.902	0.195
1	def fibonacci(n):\n # Return the n-th Fibonacci number, with fibonacci(0) == 0.\n if n < 0:\n raise ValueError("n must be non-negative")\n if n <= 1:\n return	return	ok	Python code format with structured documentation conventions showing `math.fib` function definition, establishing pattern of defining behavior for Fib function.\n\nThe sentence "For n < n n >: def fib(n): if n <= 0: return 0" mirrors a pattern asserting the function function's's initial values, implying the return value for n <= n numbers is the number itself.\n\nFinal token "return" ends an incomplete return statement ("if n <= n return"), directly continuing the pattern " "return 0" or " the integer n" — expects "n" or "n" or "0" to complete the integer result expression, completing the definition specification around small inputs.	0.785	0.430
2	She had waited three months for an answer from the university. When the envelope finally arrived, she tore it open with shaking hands, read the first line, and burst into	into	ok	Formal Chinese language context with mysterious narrative structure suggests a short story or dialogue fragment about a woman's emotions and situation.\n\nThe phrase "When she heard the shocking news, she froze, tears streamed down her face and she burst into" mirrors a reported clause pattern, strongly implying the common expression "cried of joy/excitement, completing the emotional reaction of the unexpected reveal about her pregnancy.\n\nFinal token "into" ends an incomplete verb phrase ("she burst into"), part of a listed emotional symptoms ("She burst into tears and burst into"), expecting "tears" or "joyful sobs" or "a cry," completing the negative action describing her emotional outburst.	0.909	0.183

Parity and resources

The parity table lists the largest absolute difference between the Bridge and plain Hugging Face at each stage, at the same dtype and device. The resource table lists each stage’s wall time and the process’s peak memory once that stage finished; the peak never decreases, so a flat value across stages shows the previous model was released. Downloads are excluded: the checkpoints were already in the local Hugging Face cache.

# NBVAL_SKIP
display(pd.DataFrame(parity, columns=["comparison", "max abs difference"]))
resources = pd.DataFrame(stage_log).round(2)
display(resources)
total_minutes = resources["seconds"].sum() / 60
print(f"Total: {total_minutes:.1f} minutes on {RELEASED_DEVICE} in {RELEASED_DTYPE}.")

	comparison	max abs difference
0	target block-20 output vs HF hidden_states[21]	0.0
1	AV prefill logits vs HF inputs_embeds	0.0
2	AR block-20 read-out vs NLACritic-style HF	0.0

	stage	seconds	peak RSS (GiB)
0	sidecars and value head	3.85	1.16
1	target: load and extract	22.56	13.42
2	AV: load, parity, verbalize	132.64	14.05
3	AR: load, parity, reconstruct	33.03	14.05

Total: 3.2 minutes on cpu in torch.bfloat16.

Reading these results

All three explanations parsed. Every Bridge-versus-Hugging Face comparison above is exact in bfloat16 on this CPU: the block-20 extraction, the AV’s prefill logits with the injected row, and the AR’s block-20 read-out. Generation itself was compared only between the cached and uncached paths, in the mechanics section. The float32 and float64 rescalings produce identical injected values here.

Each explanation identifies the prompt’s format (a travel narrative, Python code, a narrative sentence) and the role of the final token in an incomplete phrase. Each also contains content that is not in the prompt: the first names Brussels and Belgium as the continuation, although nothing in the prompt mentions Belgium; the second refers to a nonexistent math.fib; the third mentions a pregnancy and a Chinese-language context. The cosines are 0.902, 0.785, and 0.909, so high cosines and incorrect details occur in the same explanations. A score does not certify the individual claims in the text.

These are three single samples, one per prompt. They illustrate the protocol and its failure modes; they are not an estimate of how well the released NLA explains Qwen’s activations in general.

Interpretation checklist

• A high cosine says the AR could recover the activation’s direction from the text. It does not say the explanation is what a human would conclude from it, or that the model uses the activation for what the text describes.

• A low cosine can come from a poor explanation, a hard activation, or an AR that ignores parts of the text. One vector cannot separate these.

• Sampled explanations vary between runs. Compare several samples, with fixed seeds, before reading meaning into one.

• Parse failures are outcomes of the AV, not noise to filter. Report them alongside successful round trips.

• The scores are directional. They ignore the activation’s norm, and they are not a dataset-level fraction of variance explained.

• An NLA pair is trained for one target model, one layer, and one residual basis. Activations from another model, another layer, or a compatibility-mode Bridge are outside that contract.

• Hooks already installed on a model still apply during these helpers; remove them to analyze the unmodified computation.

Scope

This notebook covers the inference round trip for one activation at a time. It does not cover steering with edited explanations, corpus-level benchmarks, other released NLA pairs, training, or quantized execution.

References

• Natural Language Autoencoders Produce Unsupervised Explanations of LLM Activations.

• Authors’ reference implementation. This notebook does not import, vendor, or depend on that repository’s code.

• Patchscopes generation demo, a related activation-patching technique without a trained reconstructor.