transformer_lens.tools.analysis.backward_lens module¶

Backward Lens gradient-factor capture and vocabulary projection.

The Backward Lens represents a linear weight gradient as a sum of token-position outer products and projects residual-width factors into the model vocabulary. The public API currently supports raw GPT-2 TransformerBridge models.

class transformer_lens.tools.analysis.backward_lens.BackwardLens(model: Any)¶

Bases: object

Analyze GPT-2 MLP weight gradients in the output vocabulary basis.

The analyzer accepts a fresh, raw GPT-2 TransformerBridge. Results retain no model or tokenizer reference and contain detached CPU-owned tensors. Raw backward signals are loss gradients; gradient descent subtracts them.

__init__(model: Any)¶: Validate and retain the raw GPT-2 Bridge used for analyses.

analyze(prompt: str, target_token: str, layers: Sequence[int], *, normalized: bool = False, top_k: int = 10, return_full_logits: bool = False) → BackwardLensResult¶

Analyze one final-position, one-token target loss.

Parameters:

prompt – Non-empty unbatched prompt text.
target_token – Text encoding to exactly one token without BOS.
layers – Unique GPT-2 layer indices in desired result order.
normalized – Also project unit-normalized nonzero factors using the Normalized Logit Lens. Raw projections are always returned.
top_k – Number of largest and smallest values and token ids retained per matrix and position. Defaults to 10.
return_full_logits – Also retain full vocabulary tensors on CPU. Defaults to False to keep result size bounded.

Returns:

Detached gradient factors, bounded vocabulary rankings, norms, reconstruction errors, target metadata, and optional full logits.

class transformer_lens.tools.analysis.backward_lens.BackwardLensLayerResult(layer: int, input_projection: BackwardLensMatrixResult, output_projection: BackwardLensMatrixResult)¶

Bases: object

Vocabulary-facing input/output MLP matrix results for one indexed layer.

input_projection: BackwardLensMatrixResult¶

layer: int¶

output_projection: BackwardLensMatrixResult¶

class transformer_lens.tools.analysis.backward_lens.BackwardLensMatrixResult(factors: LinearGradientFactors, projected_factor: Literal['forward_inputs', 'output_gradients'], factor_norms: Float[Tensor, 'position'], zero_norm_mask: Bool[Tensor, 'position'], vocabulary_size: int, target_token_id: int, top_ranking: VocabularyRanking, bottom_ranking: VocabularyRanking, target_largest_ranks: Int[Tensor, 'position'], target_smallest_ranks: Int[Tensor, 'position'], normalized_top_ranking: VocabularyRanking | None = None, normalized_bottom_ranking: VocabularyRanking | None = None, normalized_target_largest_ranks: Int[Tensor, 'position'] | None = None, normalized_target_smallest_ranks: Int[Tensor, 'position'] | None = None, vocabulary_logits: Float[Tensor, 'position d_vocab'] | None = None, normalized_vocabulary_logits: Float[Tensor, 'position d_vocab'] | None = None)¶

Bases: object

Factors and vocabulary readouts for one GPT-2 MLP weight matrix.

factors contains the full linear factorization. projected_factor says whether its residual-width forward_inputs or raw-gradient output_gradients were decoded. factor_norms and zero_norm_mask have shape [position] with float32 and bool dtypes. Largest and smallest signed rankings are always retained. Full vocabulary_logits are present only when explicitly requested. Normalized rankings and optional full logits are present when the Normalized Logit Lens is requested. Every retained tensor is a detached CPU-owned value; gradient descent subtracts raw gradients.

bottom(*, k: int, normalized: bool = False) → VocabularyRanking¶: Return up to the retained smallest signed logits and token ids.

bottom_ranking: VocabularyRanking¶

bottom_tokens(tokenizer: Any, *, k: int, normalized: bool = False) → list[list[str]]¶: Decode the smallest-k vocabulary ids for every position.

factor_norms: Float[Tensor, 'position']¶

factors: LinearGradientFactors¶

gradient_descent_target_ranks(target_token_id: int, *, normalized: bool = False) → Int[Tensor, 'position']¶: Return ascending raw-gradient target ranks (rank zero is smallest).

logits(*, normalized: bool = False) → Float[Tensor, 'position d_vocab']¶: Return opted-in raw or Normalized Logit Lens full logits.

normalized_bottom_ranking: VocabularyRanking | None = None¶

normalized_target_largest_ranks: Int[Tensor, 'position'] | None = None¶

normalized_target_smallest_ranks: Int[Tensor, 'position'] | None = None¶

normalized_top_ranking: VocabularyRanking | None = None¶

normalized_vocabulary_logits: Float[Tensor, 'position d_vocab'] | None = None¶

projected_factor: Literal['forward_inputs', 'output_gradients']¶

target_largest_ranks: Int[Tensor, 'position']¶

target_ranks(target_token_id: int, *, largest: bool, normalized: bool = False) → Int[Tensor, 'position']¶

Return zero-based target ranks per position in the requested ordering.

largest=True gives rank zero to the largest logit. largest=False gives rank zero to the smallest, which is the useful raw-gradient convention for the second MLP matrix because gradient descent subtracts it. Ties receive the same competition rank. The analyzed target’s ranks are always retained; other token ids require opted-in full logits.

target_smallest_ranks: Int[Tensor, 'position']¶

target_token_id: int¶

top(*, k: int, normalized: bool = False) → VocabularyRanking¶: Return up to the retained largest signed logits and token ids.

top_ranking: VocabularyRanking¶

top_tokens(tokenizer: Any, *, k: int, normalized: bool = False) → list[list[str]]¶: Decode the largest-k vocabulary ids for every position.

vocabulary_logits: Float[Tensor, 'position d_vocab'] | None = None¶

vocabulary_size: int¶

zero_norm_mask: Bool[Tensor, 'position']¶

class transformer_lens.tools.analysis.backward_lens.BackwardLensResult(prompt: str, prompt_token_ids: Int[Tensor, 'position'], target_token: str, target_token_id: int, loss: float, layers: tuple[BackwardLensLayerResult, ...], max_absolute_reconstruction_error: float, max_relative_reconstruction_error: float, includes_normalized_logits: bool, includes_full_logits: bool)¶

Bases: object

Detached result of one BackwardLens.analyze() call.

prompt and target_token echo the analyzed inputs; target_token_id is the single vocabulary id the target text encodes to. loss is the raw scalar cross-entropy of the final-position next-token prediction against the target; it preserves the d(loss)/d(...) sign convention and is not negated. prompt_token_ids is an owned CPU int64 tensor with shape [position]; every residual-width factor in layers is aligned to these same positions. Position zero is a prepended BOS only when the model and tokenizer configuration requests one. layers preserves requested order. Maximum errors summarize both matrices over every requested layer. includes_normalized_logits records whether the Normalized Logit Lens was computed. includes_full_logits records whether full vocabulary tensors were retained in addition to bounded rankings. No model or tokenizer reference is retained.

includes_full_logits: bool¶

includes_normalized_logits: bool¶

layer(layer: int) → BackwardLensLayerResult¶: Return one requested layer result or raise KeyError.

layers: tuple[BackwardLensLayerResult, ...]¶

loss: float¶

max_absolute_reconstruction_error: float¶

max_relative_reconstruction_error: float¶

prompt: str¶

prompt_token_ids: Int[Tensor, 'position']¶

target_token: str¶

target_token_id: int¶

class transformer_lens.tools.analysis.backward_lens.LinearGradientFactors(forward_inputs: Float[Tensor, 'position in_features'], output_gradients: Float[Tensor, 'position out_features'], weight_gradient: Float[Tensor, 'weight_dim_0 weight_dim_1'], reconstructed_gradient: Float[Tensor, 'weight_dim_0 weight_dim_1'], absolute_reconstruction_error: float, relative_reconstruction_error: float, weight_layout: Literal['in_out', 'out_in'])¶

Bases: object

Detached factors and reconstruction for one linear weight gradient.

forward_inputs and output_gradients have shapes [position, in] and [position, out]. output_gradients and weight_gradient preserve the raw d(loss)/d(tensor) sign; they are not negated into update directions. Gradient tensors use the requested storage layout. All tensors are cloned to CPU in float32 so the result owns no autograd graph.

absolute_reconstruction_error: float¶

forward_inputs: Float[Tensor, 'position in_features']¶

output_gradients: Float[Tensor, 'position out_features']¶

reconstructed_gradient: Float[Tensor, 'weight_dim_0 weight_dim_1']¶

relative_reconstruction_error: float¶

weight_gradient: Float[Tensor, 'weight_dim_0 weight_dim_1']¶

weight_layout: Literal['in_out', 'out_in']¶

class transformer_lens.tools.analysis.backward_lens.VocabularyRanking(values: Float[Tensor, '*leading k'], indices: Int[Tensor, '*leading k'])¶

Bases: object

Owned CPU copies of signed vocabulary rankings with shape [..., k].

values preserves the floating dtype and sign of logits; indices has dtype torch.int64. Both tensors are detached.

indices: Int[Tensor, '*leading k']¶

values: Float[Tensor, '*leading k']¶