Interpretability with lumen

tensorplay.lumen is TensorPlay’s native explainability suite. It covers the classic attribution algorithms (integrated gradients, deep lift, LRP, SHAP, TCAV), layer and neuron attribution, visualization and statistics, and a set of research-grade LLM tools (sparse autoencoders, cross-layer transcoders, attribution graphs, feature circuits, logit lens, causal tracing). Everything runs on tensorplay primitives — no third-party runtime is required.

import tensorplay.lumen as lumen

What you can reproduce

The suite implements methods from the interpretability literature, so you can rerun their core experiments on your own models:

Paper / technique	lumen API
Integrated Gradients — Sundararajan et al., ICML 2017	`IntegratedGradients`
DeepLIFT — Shrikumar et al., 2017	`DeepLift`, `DeepLiftShap`
Layer-wise Relevance Propagation — Bach et al., PLOS ONE 2015	`LRP`
SHAP — Lundberg & Lee, NeurIPS 2017	`KernelShap`
TCAV — Kim et al., ICML 2018	`tensorplay.lumen.concept.TCAV`
Causal tracing / knowledge localization — Vig et al., 2020; Meng et al., ICLR 2022	`lumen.llm.ActivationPatching.causal_trace`
Logit lens — nostalgebraist, 2020	`lumen.llm.LogitLens`
Sparse autoencoder features — Anthropic, 2023–2024	`lumen.llm.SparseAutoencoder`
Cross-layer transcoders — Anthropic, 2025	`lumen.llm.CrossLayerTranscoder`
Attribution graphs — Anthropic, 2023	`lumen.llm.AttributionGraph`
Feature circuits — Anthropic, 2024	`lumen.llm.FeatureCircuit`

Feature attribution on a small model

Start with a tiny network and ask “which input features drive the output?”

import tensorplay as tp
from tensorplay import nn
from tensorplay.lumen import IntegratedGradients, DeepLift, LRP

class MLP(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(4, 8)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(8, 3)

    def forward(self, x):
        return self.fc2(self.relu(self.fc1(x)))

model = MLP()
x = tp.tensor([[1.0, 2.0, -1.0, 0.5]])

ig = IntegratedGradients(model).attribute(x, target=1, n_steps=50)
dl = DeepLift(model).attribute(x, target=1, baselines=0.0)
lrp = LRP(model).attribute(x, target=1)

print("IG   :", ig)
print("DeepLift:", dl)
print("LRP     :", lrp)

IntegratedGradients satisfies completeness: the attributions sum to the model’s output difference between the input and a baseline.
DeepLift and LRP are backpropagation-based and give fine-grained per-feature relevance. All three return tensors with the same shape as the input.

Layer and neuron attribution

When you want to know what a layer or a single neuron is responding to, use the layer and neuron series:

from tensorplay.lumen import LayerActivation, LayerIntegratedGradients
from tensorplay.lumen import NeuronConductance, NeuronIntegratedGradients

layer = model.fc2
layer_act = LayerActivation(model, layer).attribute(x)          # (1, 3)
layer_ig = LayerIntegratedGradients(model, layer).attribute(x, target=1, baselines=0.0)

neuron_ig = NeuronIntegratedGradients(model, layer).attribute(
    x, neuron_selector=(0,), baselines=0.0
)  # gradient of neuron (0) w.r.t. the input

Concept-based attribution (TCAV)

TCAV answers “does the model use concept X (e.g. stripes) when making a decision?” It trains a linear classifier between concept activation vectors and measures the sensitivity of the prediction to that concept direction:

from tensorplay.lumen.concept import TCAV, Concept
from tensorplay.lumen.concept._utils.data_iterator import dataset_to_dataloader
from tensorplay.utils.data import TensorDataset

concept_data = dataset_to_dataloader(TensorDataset(tp.randn(32, 4)))
random_data = dataset_to_dataloader(TensorDataset(tp.randn(32, 4)))
tcav = TCAV(model, layers=["fc1"])
tcav.interpret(
    inputs=x,
    experimental_sets=[
        [Concept(0, "concept", concept_data)],
        [Concept(1, "random", random_data)],
    ],
)

TCAV trains its default classifier with sklearn; install it to run this example. Without sklearn, TCAV raises a friendly ValueError.

LLM interpretability

The LLM tools operate on any transformer built with tensorplay primitives. The examples below use a tiny native GPT so they run offline, end to end.

A tiny native transformer

from types import SimpleNamespace

def causal_mask(seq):
    return tp.zeros(seq, seq, dtype=tp.float32).masked_fill(
        tp.triu(tp.ones(seq, seq, dtype=tp.bool), diagonal=1), float("-inf")
    )

class Block(nn.Module):
    def __init__(self, d=16, heads=4, ff=32):
        super().__init__()
        self.ln1 = nn.LayerNorm(d)
        self.attn = nn.MultiheadAttention(d, heads, batch_first=True)
        self.ln2 = nn.LayerNorm(d)
        self.mlp = nn.Sequential(nn.Linear(d, ff), nn.GELU(), nn.Linear(ff, d))

    def forward(self, x, mask):
        h = self.ln1(x)
        o, w = self.attn(h, h, h, attn_mask=mask, need_weights=True,
                        average_attn_weights=False)
        x = x + o
        h = self.ln2(x)
        x = x + self.mlp(h)
        return x, w

class TinyGPT(nn.Module):
    def __init__(self, vocab=11, d=16, heads=4, layers=2, ff=32):
        super().__init__()
        self.transformer = nn.Module()
        self.transformer.wte = nn.Embedding(vocab, d)
        self.transformer.wpe = nn.Embedding(16, d)
        self.transformer.h = nn.ModuleList([Block(d, heads, ff) for _ in range(layers)])
        self.transformer.ln_f = nn.LayerNorm(d)
        self.lm_head = nn.Linear(d, vocab, bias=False)

    def forward(self, input_ids, output_attentions=False):
        seq = input_ids.shape[1]
        x = self.transformer.wte(input_ids) + self.transformer.wpe(
            tp.arange(0, seq, dtype=tp.int64).unsqueeze(0)
        )
        mask = causal_mask(seq)
        attentions = []
        for blk in self.transformer.h:
            x, w = blk(x, mask)
            attentions.append(w)
        x = self.transformer.ln_f(x)
        return SimpleNamespace(
            logits=self.lm_head(x),
            attentions=tuple(attentions) if output_attentions else None,
        )

VOCAB = ["<pad>", "<unk>", "The", "capital", "of", "France", "England",
         "is", "Paris", "London", "<eos>"]

class TinyTokenizer:
    def __init__(self):
        self.id2tok = {i: t for i, t in enumerate(VOCAB)}
        self.tok2id = {t: i for i, t in enumerate(VOCAB)}

    def encode(self, text, return_tensors="pt", add_special_tokens=False):
        ids = [self.tok2id.get(w, 1) for w in text.split()]
        return tp.tensor([ids], dtype=tp.int64) if return_tensors == "pt" else ids

    def convert_ids_to_tokens(self, ids):
        return [self.id2tok[int(i)] for i in ids]

    def decode(self, ids):
        return " ".join(self.convert_ids_to_tokens(ids))

    def __call__(self, text, return_offsets_mapping=False, **kwargs):
        words = text.split()
        ids = [self.tok2id.get(w, 1) for w in words]
        res = {"input_ids": ids}
        if return_offsets_mapping:
            offsets, pos = [], 0
            for w in words:
                start = text.index(w, pos)
                offsets.append((start, start + len(w)))
                pos = start + len(w)
            res["offset_mapping"] = offsets
        return res

model = TinyGPT()
tokenizer = TinyTokenizer()
prompt = "The capital of France is"
ids = tokenizer.encode(prompt)

Attention heatmaps

Extract per-head attention weights and render them with the visualization helper. Highlight important tokens (e.g. a subject) to make the attended-to pattern legible:

from tensorplay.lumen.llm import (
    extract_attention,
    visualize_attention_multi_layer,
    token_attention,
    visualize_token_attention,
)

result = extract_attention(model, tokenizer, prompt, layers=[0, 1])

# paper-style grid: layers x heads, what "France" attends to outlined
fig, axes = visualize_attention_multi_layer(result, query_token="France")

# single-token view: France's attention at every (layer, head)
ta = token_attention(result, query_token="France")
fig, ax = visualize_token_attention(ta)

Causal tracing: where is a fact stored?

The causal-tracing protocol from knowledge-editing research corrupts a prompt, then restores the clean hidden state at every (layer, position) and measures how much of the clean prediction is recovered:

from tensorplay.lumen.llm import ActivationPatching

clean = tokenizer.encode("The capital of France is")[0]
corrupted = tokenizer.encode("The capital of England is")[0]
trace = ActivationPatching(model).causal_trace(clean, corrupted, target_position=-1)

for layer, row in zip(trace["layers"], trace["scores"]):
    print("L%d:" % layer, ["%.2f" % v for v in row])

High recovery at the subject position (“France”) and the final layer localizes where the model stores the fact. This is the diagnostic used before applying knowledge-editing weight updates.

Logit lens

The logit lens decodes every layer’s hidden state through the LM head, showing which token the model “would have predicted” at each depth:

from tensorplay.lumen.llm import LogitLens

result = LogitLens(model, tokenizer).lens(ids, position=-1, topk=3)
for layer, preds in result["predictions"].items():
    print("L%d:" % layer, ["%s=%.2f" % (t, p) for t, p in preds])

Sparse autoencoder features

Sparse autoencoders learn a dictionary of sparse, interpretable features from hidden states — the dictionary-learning step used in feature research. Train one on activations collected from a layer:

from tensorplay.utils.data import DataLoader, TensorDataset
from tensorplay.lumen.llm import (
    SparseAutoencoder,
    collect_hidden_states,
    train_sae,
    feature_frequency,
)

corpus = ["The capital of France is", "The capital of England is",
          "The capital of France is Paris", "The capital of England is London"]
seqs = [tokenizer.encode(s)[0] for s in corpus]
hidden = collect_hidden_states(model, layer=0, input_ids_list=seqs, batch_size=2)

sae = SparseAutoencoder(input_dim=16, hidden_dim=32)
loader = DataLoader(TensorDataset(hidden), batch_size=8)
train_sae(sae, loader, n_epochs=20, lr=1e-2, l1_coef=1e-3)

features = sae.encode(hidden)
print("active fraction per feature:", feature_frequency(features))

Each learned feature is a direction in activation space; top_activating_indices lists the corpus examples that most activate each feature, which is how you would label a feature (e.g. “the token right before a country name”).

Cross-layer transcoders

A cross-layer transcoder predicts the update a transformer component adds to the residual stream from the incoming residual stream, decomposing the component’s computation into sparse features:

from tensorplay.lumen.llm import (
    CrossLayerTranscoder,
    collect_component_updates,
    train_clt,
)

residuals, deltas = collect_component_updates(model, layer=0, input_ids_list=seqs, batch_size=2)
clt = CrossLayerTranscoder(input_dim=16, hidden_dim=32)
loader = DataLoader(TensorDataset(residuals, deltas), batch_size=8)
train_clt(clt, loader, n_epochs=20, lr=1e-2, l1_coef=1e-3)

Attribution graphs

Attribution graphs show which MLP features at each layer positively contribute to an output and how they connect across layers. The nodes are features weighted by gradient * activation; the edges are positive linearized feature-to-feature paths:

from tensorplay.lumen.llm import AttributionGraph

graph = AttributionGraph(model).build(ids, target_position=-1, topk=5)
print("nodes:", [(n.layer, n.kind, n.index) for n in graph.nodes])
print("edges:", [(e.source, e.target, round(e.weight, 4)) for e in graph.edges[:5]])

Feature circuits

Feature circuits measure the direct effect of every candidate feature node (MLP features and attention heads) by zero-ablating it, then estimate pairwise interaction edges via joint ablation:

from tensorplay.lumen.llm import (
    FeatureCircuit,
    top_mlp_features,
    top_attention_heads,
)

clean_logits = model(ids).logits
target_id = int(clean_logits[0, -1].argmax())
target = lambda out: float(out.logits[0, -1, target_id])

nodes = top_mlp_features(model, ids, target_position=-1, topk=4)
nodes += top_attention_heads(model, ids, target_position=-1, topk=2)
circuit = FeatureCircuit(model, target).build(ids, nodes=nodes, topk=6)
print("circuit nodes:", [n.name for n in circuit.nodes])

Visualization and statistics

Attribution results can be rendered with the visualization helpers (images, text, timeseries, attention heatmaps) and aggregated with Summarizer / ClassSummarizer:

from tensorplay.lumen import Summarizer, Mean, StdDev, Count

stats = Summarizer([Mean(), StdDev(), Count()])
for _ in range(10):
    stats.update(tp.randn(4))
print(stats.summary)

Optional dependencies

sklearn is used only by Lime, KernelShap, and the TCAV default classifier. Without it, those classes raise a friendly ValueError when you run them.
matplotlib is used for visualization; the helpers accept either numpy arrays or tensorplay tensors directly.

Where to go next

tensorplay.lumen reference — the full public API (80+ symbols) is re-exported from tensorplay/lumen/__init__.py.
Models — build the nn.Module models the attributions run on.
Autograd — how the gradient computations behind the attributions work.