Facets¶

Facets are named views of a captured activation record. They provide one access surface for structural outputs, such as an LSTM's output, h_n, and c_n, and semantic recipe outputs, such as attention q, k, v, MLP activations, normalization parameters, and embedding weights.

Access¶

Use item access as the canonical form:

log = tl.trace(model, x)
q = log.modules["blocks.0.attn"].facets["q"]
h_n = log.modules["lstm"].facets["h_n"]

Facets live on a Trace. tl.record takes no facet keyword and a Recording has no modules[...].facets; cook it first with recording.to_trace() and read the facets there (a facet whose inputs the save= predicate did not keep fails explicitly).

Attribute access is convenience sugar for valid Python identifiers that do not collide with FacetView methods:

q = log.modules["blocks.0.attn"].facets.q

Item access always works for names that collide with methods, such as keys, and for dotted structural names such as out1.0. Typed container paths are also accepted when string names would collide.

Structural output names come from the captured ContainerSpec:

single tensor output -> out
named output         -> values, indices, output, h_n, c_n
positional output    -> out0, out1
nested positional    -> out1.0

module.outs remains a thin alias for raw structural outputs. module.facets is the unified named access layer. If a module was called more than once, module.facets raises; choose an explicit pass with module.calls[n].facets. A call's view runs the module's recipes for that call alone, and child modules resolve to their call inside it: for an attention layer shared across depth (ALBERT), module.calls[n].facets.q is the query projection of pass n. trace.modules_with_facet(name) and trace.attention_blocks() still yield such a reused Module when its recipe declares the facet and at least one of its calls has it available now; loop over module.calls.values() to read each pass.

keys(), iteration, in, has(name), and item/attribute access are available-now surfaces: recipes run once and only facets readable in the current capture appear. .get(name, default) returns default for any facet that is not available now. Use the provisional declared() accessor for declaration-level discovery; it reports every structural and declared recipe facet with status available_now, structurally_absent, needs_capture, or declared_not_produced, including any enabled aliases.

Structurally absent facets are omitted and raise plain KeyError on item access. Declared facets that need additional payload capture are omitted and raise MissingFacetError with a recapture hint naming what to save.

Specs¶

Recipes now return FacetSpec objects rather than computed tensors. A spec records:

home kind: op, module_output, module_input, parameter, computed
home label/address and pass/call index
output path
transform primitive chain
capability class and flags
value version
conflict or alias group
recipe id and version

The supported structural primitive vocabulary is:

spec[key]
spec.heads(n_heads, d_head)
spec.split(n_sections, dim=-1)
spec.reshape(...)
spec.transpose(dim0, dim1)
spec.select(dim, index)

Read is lazy: facet.value applies the transform chain to the home value. Facet objects are tensor-like, so most tensor reads can use the facet directly:

torch.equal(module.facets["q"], manual_q)
module.facets["q"].shape

Gradients¶

facet.grad applies the same transform chain to the home op's saved output gradient. It is available only for grad-capable op-anchored structural facets and only when the home op has a saved gradient.

Default traces do not save gradients. When a gradient is unavailable, facet.grad returns a MissingGrad sentinel with an exact recapture instruction. The sentinel does not raise merely because it was returned; tensor use raises:

log = tl.trace(model, x)
missing = log.modules["attn"].facets["q"].grad
print(missing.reason)

log = tl.trace(
    model,
    x,
    capture=tl.options.CaptureOptions(backward_ready=True, save_grads=True),
)
log.backward(log[log.output_layers[0]].out.sum())
q_grad = log.modules["attn"].facets["q"].grad

Parameter facets, such as LayerNorm gamma and Embedding weight, read the parameter value. Their parameter gradients are not exposed through facet.grad in this release; this path is intentionally read-only.

Registration¶

The process registry remains additive:

@tl.facets.register(class_name="MyAttention", target_scope="module")
def my_attention(module):
    ...

Registration affects future traces. Each trace captures an immutable registry snapshot with a version and provenance id, so lazy facet access is stable even if the global registry changes later.

Per-trace and context additions are also capture-time only:

log = tl.trace(model, x, facets=[my_attention])

with tl.facets.adding(my_attention):
    log = tl.trace(model, x)

The guide Adding conveniences for your model builds a custom facet next to a custom input detector (tl.preprocessing.register) and output decoder (tl.postprocessing.register); Semantic input and output covers the built-in routes those registries already hold.

tl.facets.reset() restores the process registry to built-ins. tl.facets.functions() and tl.facets.describe(class_name) expose registry metadata.

Conflict resolution is deterministic: class_qualname is more specific than class_name, which is more specific than a predicate; user recipes beat built-ins only within the same specificity tier. True same-tier facet collisions warn.

Reading module config in a recipe (head counts, dims)¶

A recipe receives the TorchLens Module record, not the live nn.Module. The record does not proxy arbitrary live attributes as direct attributes -- they are snapshotted into module.custom_attributes. Use tl.facets.config_value(module, *names), which checks record attributes and the custom_attributes snapshot, instead of module.num_heads directly. The recipe-authoring helpers are public under tl.facets: config_value, child_output_spec(module, child_name, recipe_id) (a child module's output as a spec) and parameter_spec(module, name, recipe_id) (a parameter as a spec); no private import is needed:

@tl.facets.register(class_name="MyAttention", target_scope="module")
def my_attention(module):
    n_heads = tl.facets.config_value(module, "num_heads", "n_heads", "num_attention_heads", "n_head")
    qkv = tl.facets.child_output_spec(module, "qkv", "my_attention")
    return {"q": qkv.split(3, dim=-1)[0].heads(n_heads)}  # d_head inferred from shape

FacetSpec.heads(num_heads, d_head=None) infers the head dimension from the projection's last axis when the facet is read (last_dim // num_heads); a last axis the head count does not divide refuses at read time. Pass d_head explicitly when the projection carries extra columns.

Head counts are inferred automatically for any module that stores a conventionally named count -- HF configs, torch.nn.MultiheadAttention (.num_heads), or a custom module using one of the common names. Note that num_heads cannot be recovered from tensor shapes alone (a q-projection's output dim is num_heads * d_head, an unfactorable product), so a truly bespoke module with a non-standard field name must pass the count explicitly.

Capability Classes¶

Facet transform chains fail closed. The chain capability is the intersection of its primitive flags.

bijective_view      reshape, transpose, permute-style exact views; read, grad, write
selection           slice, select, split; read, grad, write
aliasing_selection  grouped-query or broadcast-style aliases; read and grad only
computed            callable or multi-home values; read-only

Only bijective_view and selection chains are writable. aliasing_selection facets, such as grouped-query K/V heads selected through a query-head index, require an explicit alias policy and are refused by default. computed facets are never writable.

Intervention¶

Facet intervention reuses the existing TorchLens hook machinery. A facet target resolves to its home op, and TorchLens installs a normal whole-output hook on that home. The wrapper reads the facet slice from the home output, calls the user hook on that slice, then scatters the edited slice back into the full home output returned by the hook.

Use tl.facet(name) for a named facet and tl.head(index, name) or tl.facet(name).head(index) for attention-head slices:

# Facet recipes often need internal child activations, so this example
# intentionally saves all payloads. For ordinary selective capture, prefer
# save=tl.func(...) or save=tl.in_module(...).
log = tl.trace(
    model,
    x,
    capture=tl.options.CaptureOptions(save_args=True), save="all",
)
edited = log.fork("ablated")
edited.do(tl.head(3, "q"), tl.zero_ablate(), model=model, inputs=x)

Static patching uses the same scatter-back path:

patch = torch.zeros_like(log.modules["blocks.0.attn"].facets.head(3).q)
edited = log.fork("patched")
edited.do(tl.facet("q").head(3), patch, intervention=tl.options.InterventionOptions(engine="set_only"))
edited = edited.run(model, x)

tl.head(index) without a facet name targets the default attention projection facets q, k, and v wherever they are present:

edited.do(tl.head(3), tl.zero_ablate(), model=model, inputs=x)

Scatter-back preserves the home tensor dtype, device, shape, and memory format. bijective_view primitives write through exact view operations such as reshape and transpose. selection primitives scatter into the selected region, such as a split third of a fused QKV projection or one selected head.

When multiple facet edits share one home op, TorchLens applies them in attach order and checks boolean write masks before rerun. Non-overlapping writes, such as GPT-2 c_attn edits to q head 0 and k head 1, compose on the shared home. Overlapping same-home writes fail at attach time with a conflict error.

Facet intervention fails closed in these cases:

computed facet        refused; read-only
aliasing_selection    refused unless a future explicit alias policy is supplied
non-op home           refused
non-raw value version refused as not intervention-safe
overlapping writes    refused before rerun

Fused attention reconstruction facets are read-only. Editing internal fused facets requires recapturing with an eager attention implementation so the facet is backed by a real op.

Patching¶

Patching helpers live at tl.facets.patching. They run a clean trace and a corrupted trace, then reuse the facet intervention path (fork, do with model=/inputs= (a rerun), and tl.facet/tl.head selectors) to patch clean facet activations into the corrupted run.

The metric callable receives the patched Trace and must return a scalar tensor:

def metric(log):
    return log[log.output_layers[0]].out[..., answer_id].mean()

scores = tl.facets.patching.activation_patch_attention_heads(
    model, clean_tokens, corrupted_tokens, metric
)

Available activation-patching helpers:

activation_patch_residual_stream -> [layer, pos] by default, or [layer]
activation_patch_attention_output -> [layer]
activation_patch_attention_heads  -> [layer, head]
activation_patch_mlp_output       -> [layer]

Residual patching defaults to resid_pre and patches each position along axis 1. Attention-head patching defaults to the result facet, whose shape is [batch, pos, head, d_model]; tl.head(i, "result") selects one head's contribution. MLP patching defaults to the output facet on modules that expose the built-in MLP facet family.

Attribution patching is a faster linear approximation to activation patching:

effect ~= grad(metric wrt corrupted component) * (clean component - corrupted component)

summed over the component dimensions. The helper runs one clean forward/backward and one corrupted forward/backward, reads facet gradients through facet.grad, and returns [layer, head]:

approx = tl.facets.patching.attribution_patch_attention_heads(
    model, clean_tokens, corrupted_tokens, metric
)

Attribution patching requires gradient capture. By default the helper traces with save_grads=True. If gradients are not saved, it raises a clear error that includes the MissingGrad recapture instruction.

Reconstruction¶

Fused torch.nn.functional.scaled_dot_product_attention hides scores, pattern, and z. TorchLens reconstructs these as read-only computed facets from the exact inputs and options the fused kernel received. A default capture already records them: the SDPA call's replay argument template holds scale, is_causal, dropout_p and enable_gqa, and the Q, K, V and mask tensors are the saved outputs of the ops that produced them. So a plain trace serves the facets:

log = tl.trace(model, x)
pattern = log.modules["blocks.0.attn"].facets.pattern
z = log.modules["blocks.0.attn"].facets.z

When the default capture did not keep an input (a selective save= that does not save the op producing Q, K, V or the mask, or a structure-only capture), the facet fails closed: facets.declared() marks it needs_capture and access raises MissingFacetError naming the missing input. Saving the SDPA call's arguments verbatim lifts that requirement:

log = tl.trace(model, x, save=tl.in_module("blocks.0.attn"),
               capture=CaptureOptions(reconstruction_ready=True))

CaptureOptions(reconstruction_ready=True) enables save_args=True and save_rng=True. Either way the reconstruction uses the tensors that actually fed the fused kernel, so RoPE models use the post-RoPE Q and K, not the q_proj/k_proj outputs.

Validation is per facet and non-vacuous:

z                 compared directly to the fused SDPA op output
pattern, scores   recompute z = pattern @ V, then compare to SDPA output
result            sum per-head W_O contributions and compare to output projection

GQA/MQA uses PyTorch SDPA's enable_gqa repeat-interleave convention. Causal masks and boolean/additive tensor masks are supported. Nonzero SDPA dropout is fail-closed for now; the facet is marked needs_capture in facets.declared() and item or attribute access raises MissingFacetError naming dropout/RNG as the missing prerequisite rather than serving an unvalidated value.

Residuals¶

Transformer-like block recipes expose:

resid_pre   block input op
resid_mid   first identifiable post-attention residual add op
resid_post  block output op

These are op-structural facets when the anchors exist, so .grad works under the normal gradient contract. If a block structure does not expose a real post-attention add, resid_mid is absent instead of guessed.

Attention modules also expose read-only per-head result when the recipe can find an output projection. result[..., head, :] is that head's contribution to the projected residual stream; summing heads and adding projection bias matches the captured output-projection tensor.

Unembedding Head and Logit Lens¶

Every name in this section is experimental and may change in any release.

The language_model_head recipe matches a module with a conventional unembedding child (lm_head, embed_out, or output_projection) and anchors the facts a logit-lens projection needs:

logits            head output op (the model's real output logits)
unembed_weight    head weight parameter, shape (vocab, d_model)
unembed_bias      head bias parameter, or structurally absent
final_norm_kind   "layer_norm" | "rms_norm" for the norm feeding the head
final_norm_eps    that norm's epsilon
final_norm_gamma  that norm's weight parameter
final_norm_beta   that norm's bias parameter (absent for RMSNorm)
final_norm_input  op-anchored norm input, when the anchor is unambiguous

The final norm is derived structurally from the traced dataflow (the innermost classified norm module containing the head's input op), never by a name search over the module tree.

torchlens._semantic.logit_lens projects each block's residual-stream facet through the model's own final norm + unembedding:

from torchlens._semantic import logit_lens

result = logit_lens(log)                      # resid_post through the model's head
print(result.summary(tokenizer=tokenizer))    # per-layer top-1 table
result.entries[3].logits                      # one layer's projected logits
result.stacked()                              # (n_layers, *logits_shape)
result.top_tokens(5, tokenizer=tokenizer)     # per-layer top-k (token, prob)

The reconstructed lens is validated before use: applying it to the LAST block's captured resid_post must reproduce the model's captured logits, so a head the reconstruction cannot represent (nonstandard norm scaling such as a (1 + weight) RMSNorm, an extra projection, dropout before the head) refuses with LogitLensError instead of silently mislabeling layers. Pass verify=False to trust the reconstruction explicitly, lens= to supply your own callable (or per-address mapping, e.g. a tuned lens), and modules= / facet= to control what is projected. Architecture coverage extends through facet recipes producing the same facet names -- the appliance never special cases architectures.

Validation proves only that the reconstruction reproduces the LAST layer; it says nothing about how faithful the projection is at depth. The final norm and unembedding were trained on the last layer's residual stream, and earlier layers do not use that basis, so early-layer top tokens are weak evidence and often look like noise even when the layer carries the information (Belrose et al., 2023, introduced the tuned lens for this reason). Read the lens as a trend across depth, compare it with a probe or a tuned lens passed as lens=, and do not read one early layer's top token as what that layer "thinks". A Jacobian lens (the corpus-average Jacobian of the final residual with respect to a block's, passed as lens=) is fitted from public primitives in the Jacobian lens recipe.

torchlens._semantic.logit_lens_predictions is the streaming reduction of the same sweep: it projects ONE layer at a time, keeps only per-position top-k values, the full-vocabulary logsumexp, and requested-token logits / probabilities / one-based ranks, then discards that layer's full projection -- so long sequences never materialize every [batch, positions, vocab] tensor at once:

from torchlens._semantic import logit_lens_predictions

preds = logit_lens_predictions(log, k=5, positions=[-1], tokens=[the_id])
preds.rows[3].top_ids                         # one layer's top-k token ids
preds.rows[-1].provenance                     # "native output" for captured logits

mask_input=True bans the model input's own token ids from every row (each batch row's input tokens get -inf logits before the top-k, the logsumexp denominator, probabilities and ranks are taken), the one masking that cannot be redone on the retained rows afterwards; the trace's first input must be a saved integer token-id tensor (logit_lens_mask_input_unavailable otherwise). For entropy, KL to the final output, cross-entropy and rank along the trajectory in one call, see tl.mechinterp.measure_trajectory in docs/reference/stats.md.

It shares facet= / modules= / lens= / verify= semantics (and the fail-closed lens validation) with logit_lens. Probabilities always use the full-vocabulary softmax denominator. include_native=True (the default) appends a row served directly from the model's CAPTURED output logits, labelled "native output"; every reconstructed row is labelled "projected through final norm/head". On a logits_to_keep-sliced capture the native row covers only the positions the capture kept, and requested positions outside a row's covered range are dropped from that row rather than invented.

Vision heads (experimental)¶

A vision model has no language-model head. When a trace has none, the default lens falls back to a vision head recipe, and it is validated the same way (the last block's captured resid_post must reproduce the captured logits):

• classification_head matches a module with a linear classifier child, such as ViTForImageClassification: the final norm plus the classifier applied to the one token the head reads ([CLS]). result.lens_source is "classifier_head" and each layer's logits are [batch, num_labels], so result.top_tokens() answers "what does this layer already classify the image as?".
• contrastive_image_head matches CLIPModel. Capture the model on images AND the label prompts you want to score (input_ids of "a photo of a cat", "a photo of a dog", ...): each vision-tower block's [CLS] state goes through post_layernorm, visual_projection and L2 normalization, and is scored against the captured text features with CLIP's own logit_scale. result.lens_source is "contrastive_head", each layer's logits are [n_images, n_prompts], and the last layer equals CLIP's logits_per_image. Only vision-tower blocks are projected; passing a text-tower block in modules= refuses logit_lens_layer_outside_head.

classification_head facets          contrastive_image_head facets
logits                              logits_per_image
classifier_weight / classifier_bias image_projection_weight
pooled_position (token read)        text_features (text_projection output)
tower_address (blocks feeding it)   logit_scale
final_norm_* (as above)             pooled_position, tower_address, final_norm_*

Only a single-token read is reconstructed. A head that mean-pools, uses an attention pooler, or runs a pooler MLP leaves the final norm unanchored, and the lens refuses logit_lens_head_unavailable instead of guessing.

User lens identity (experimental)¶

A user lens may record which model it was built for: set model_identity on the callable (or on the mapping, or on each per-layer callable) to the model-identity record an extraction artifact carries (manifest["signature"]["model_identity"]). With the default validate=True:

• a measured identity that differs from the traced model (compared against the trace's live source model) refuses with logit_lens_identity_mismatch; verify=False runs the lens anyway;
• a matching identity runs silently;
• a lens with no identity, or one that cannot be compared (an asserted record, or the source model is gone), gets a last-layer agreement check: when the lens applied to the last block's resid_post does not reproduce the captured logits, the coded warning logit_lens_identity_unrecorded reports the largest difference and the top-1 agreement.

Fallback¶

Every module has structural facets even with no semantic recipe. This gives a nnsight-style floor for arbitrary models:

log.modules["custom.submodule"].facets["out"]

Semantic names are added when a recipe classifies the module; otherwise the module path plus structural output names remain the stable addressing surface.

TransformerLens Aliases¶

TransformerLens hook-name aliases are off by default. Enable them once per process:

tl.facets.enable_transformer_lens_aliases()

When enabled, aliases such as hook_pattern, hook_z, hook_result, hook_resid_pre, hook_resid_mid, and hook_resid_post resolve to native TorchLens facets when those native facets exist. torchlens.facets.transformer_lens_aliases_enabled() reports the current process-wide state.

Migration Cheat Sheet¶

TransformerLens hook_q           -> facets.q
TransformerLens hook_k           -> facets.k
TransformerLens hook_v           -> facets.v
TransformerLens hook_pattern     -> facets.pattern
TransformerLens hook_z           -> facets.z
TransformerLens hook_result      -> facets.result
TransformerLens hook_resid_pre   -> facets.resid_pre
TransformerLens hook_resid_mid   -> facets.resid_mid
TransformerLens hook_resid_post  -> facets.resid_post
nnsight module.path.output       -> log.modules["module.path"].facets["out"]

Coverage and Recipe Maintenance¶

torchlens.facets.coverage(log) (experimental) reports per-module recipe coverage for a completed trace: which recipes matched, which declared facets are readable, typed absence reasons, an inventory of structural-only module classes (recipe candidates), and disclosed rows for modules whose facet view refuses (e.g. multi-call reuse). It is the machine-readable input to the facet-maintenance pipeline in tools/facet_maintenance/ (DISCOVER_FACETS.md is the durable sweep prompt; run_facet_audit.py is the repeatable runner). The pipeline produces PROPOSALS FOR REVIEW only -- facet recipes are never auto-merged, because a wrong facet label is a confidently mislabelled part of someone's model.

Recipe Plugins¶

On FIRST USE of the facet-recipe subsystem (not at import torchlens -- the loader is module-level in the facet-recipes package initializer, and the facet subsystem is lazy), TorchLens loads installed setuptools entry points in the torchlens.recipes group. A broken entry point therefore warns when facets are first touched, not at import. Entry points are never scanned from local directories.

Built-In Inventory¶

TorchLens records built-in facet capabilities as data in BUILTIN_FACET_CAPABILITY_INVENTORY in the private facet-recipes package (no public spelling yet). Current classes are:

op_structural
parameter
module_input
computed_read_only

Only op_structural built-in facets may claim facet-gradient capability.

Glossary¶

FacetSpec: ABI record describing a facet home, transform chain, capability flags, value version, alias/conflict group, and recipe provenance.

tl.facet(name): selector for facet-level intervention on a named facet.

tl.head(index, name=None): selector for facet-level intervention on one attention head. With name=None, targets default q, k, and v facets.

scatter_update: ABI method that writes an edited facet slice back into the full home-op output and returns the edited home tensor.

MissingGrad: typed sentinel returned by facet.grad when the home op has no saved gradient or the facet is not grad-capable.

bijective_view: transform class for exact view operations such as reshape and transpose.

selection: transform class for slice/select/split operations.

aliasing_selection: transform class for repeated or shared storage selections, such as grouped-query K/V head access.

computed: transform class for callable, nonlinear, or multi-home values.

Write capability: a facet may be edited only when every primitive in its chain is scatter-back capable and the home is an op raw output.

Alias policy: explicit rule required before an aliasing_selection facet can be written; TorchLens provides no public policy and refuses these writes by default.

Conflict policy: same-home facet writes compose in attach order when their write masks do not overlap; overlapping writes raise before rerun.

facets=: tl.trace argument for additive per-trace facet recipes captured in that trace's registry snapshot.

adding(): tl.facets.adding(...) context manager for additive context-local recipes used only by traces captured inside the context.

reset(): tl.facets.reset() restores the process registry to built-in recipes.

Item-access naming: canonical record.facets[...] lookup for structural and recipe facets, including dotted names and method-name collisions.