transformer_lens.tools.analysis.projection_kernel module¶

Projection Kernel utilities for comparing linear subspaces.

The Projection Kernel (PK) between subspaces with orthonormal bases U and V is ||U.T @ V||_F^2. It is invariant to basis choices within either subspace and equals the sum of squared principal-angle cosines.

class transformer_lens.tools.analysis.projection_kernel.AttentionHeadRef(layer: int, head: int, role: Literal['Q', 'K', 'V', 'O'], kind: Literal['query', 'kv'])¶

Bases: object

Structured identity for one attention-head weight subspace.

head: int¶

kind: Literal['query', 'kv']¶

property label: str¶: Return the conventional TransformerLens layer/head label.

layer: int¶

role: Literal['Q', 'K', 'V', 'O']¶

class transformer_lens.tools.analysis.projection_kernel.HeadAffinityPair(source: AttentionHeadRef, target: AttentionHeadRef, score: float, normalized: float)¶

Bases: object

One ranked source-target head pair.

normalized: float¶

score: float¶

source: AttentionHeadRef¶

target: AttentionHeadRef¶

class transformer_lens.tools.analysis.projection_kernel.HeadAffinityResult(scores: Float[Tensor, 'source_layer source_head target_layer target_head'], normalized: Float[Tensor, 'source_layer source_head target_layer target_head'], valid_mask: Bool[Tensor, 'source_layer source_head target_layer target_head'], source_role: Literal['Q', 'K', 'V', 'O'], target_role: Literal['Q', 'K', 'V', 'O'], source_layer_indices: Tuple[int, ...], target_layer_indices: Tuple[int, ...], source_head_kind: Literal['query', 'kv'], target_head_kind: Literal['query', 'kv'], source_ranks: Int[Tensor, 'source_layer source_head'], target_ranks: Int[Tensor, 'target_layer target_head'], source_rank: int, target_rank: int, rank: int | None, rtol: float)¶

Bases: object

Projection Kernel affinities between two attention-head roles.

Score tensors have shape [source_layer, source_head, target_layer, target_head]. Layer index tuples map tensor positions to original model block numbers.

source_ranks and target_ranks are measured numerical ranks for each head before optional truncation. Scalar source_rank and target_rank are the retained basis widths used for their respective roles.

normalized: Float[Tensor, 'source_layer source_head target_layer target_head']¶

rank: int | None¶

rtol: float¶

scores: Float[Tensor, 'source_layer source_head target_layer target_head']¶

source_head_kind: Literal['query', 'kv']¶

source_layer_indices: Tuple[int, ...]¶

source_rank: int¶

source_ranks: Int[Tensor, 'source_layer source_head']¶

source_role: Literal['Q', 'K', 'V', 'O']¶

target_head_kind: Literal['query', 'kv']¶

target_layer_indices: Tuple[int, ...]¶

target_rank: int¶

target_ranks: Int[Tensor, 'target_layer target_head']¶

target_role: Literal['Q', 'K', 'V', 'O']¶

top_pairs(k: int = 20, *, normalized: bool = False) → List[HeadAffinityPair]¶: Return the highest-scoring valid pairs with deterministic tie order.

valid_mask: Bool[Tensor, 'source_layer source_head target_layer target_head']¶

class transformer_lens.tools.analysis.projection_kernel.ProjectionKernelResult(score: Float[Tensor, ''], normalized: Float[Tensor, ''], cosines: Float[Tensor, 'principal_angle'], angles: Float[Tensor, 'principal_angle'], rank_a: int, rank_b: int, ambient_dim: int)¶

Bases: object

Projection Kernel score and its principal-angle decomposition.

ambient_dim: int¶

angles: Float[Tensor, 'principal_angle']¶

cosines: Float[Tensor, 'principal_angle']¶

normalized: Float[Tensor, '']¶

rank_a: int¶

rank_b: int¶

score: Float[Tensor, '']¶

class transformer_lens.tools.analysis.projection_kernel.RandomSubspaceReference(ambient_dim: int, rank: int, mean: float, variance: float)¶

Bases: object

Analytic PK moments for independent random equal-rank subspaces.

ambient_dim: int¶

mean: float¶

rank: int¶

variance: float¶

class transformer_lens.tools.analysis.projection_kernel.SubspaceBasis(basis: Float[Tensor, 'ambient rank'], singular_values: Float[Tensor, 'spectrum'], rank: int, measured_rank: int, rtol: float, threshold: float, input_shape: Tuple[int, int])¶

Bases: object

An explicitly ranked orthonormal basis extracted from a matrix.

basis¶

Orthonormal column-space basis, [ambient_dim, rank].

Type:: jaxtyping.Float[Tensor, ‘ambient rank’]

singular_values¶

All reduced-SVD singular values, in descending order.

Type:: jaxtyping.Float[Tensor, ‘spectrum’]

rank¶

Number of retained basis directions.

Type:: int

measured_rank¶

Numerical rank before optional caller truncation.

Type:: int

rtol¶

Effective relative rank tolerance.

Type:: float

threshold¶

Absolute singular-value threshold used for rank measurement.

Type:: float

input_shape¶

Shape of the matrix from which the basis was extracted.

Type:: Tuple[int, int]

property ambient_dim: int¶: Dimension of the space containing the subspace.

basis: Float[Tensor, 'ambient rank']¶

input_shape: Tuple[int, int]¶

measured_rank: int¶

rank: int¶

rtol: float¶

singular_values: Float[Tensor, 'spectrum']¶

threshold: float¶

transformer_lens.tools.analysis.projection_kernel.attention_head_subspace_affinity(model: Any, *, source_role: str = 'O', target_role: str, layer_order: str = 'forward', rank: int | None = None, rtol: float | None = None) → HeadAffinityResult¶

Compute OQ, OK, or OV Projection Kernel affinities for a TransformerBridge.

K/V axes preserve native key-value heads on grouped-query attention models; they are never expanded to query-head count. Hybrid models include only attention blocks and report their original block indices.

Parameters:

model – A TransformerBridge exposing readable per-head attention weights.
source_role – Source role; v1 supports only "O".
target_role – One of "Q", "K", or "V".
layer_order – "forward" keeps strict earlier-to-later pairs; "all" keeps every pair.
rank – Optional common truncation rank. By default every head must be full column rank.
rtol – Optional relative numerical-rank tolerance.

Returns:

Affinity tensors, validity mask, original layer indices, and rank metadata.

transformer_lens.tools.analysis.projection_kernel.orthonormal_subspace(matrix: Float[Tensor, 'ambient width'], *, rank: int | None = None, rtol: float | None = None) → SubspaceBasis¶

Extract an explicitly ranked orthonormal column-space basis.

Low-precision inputs are promoted to float32 before the reduced SVD. With no explicit rtol, numerical rank uses the larger of the compute-SVD error scale and one input-storage epsilon, relative to the largest singular value. An explicit rank truncates the measured subspace but may not exceed its measured rank.

Parameters:

matrix – Finite floating-point matrix with shape [ambient_dim, width].
rank – Optional number of leading singular directions to retain.
rtol – Optional non-negative relative singular-value threshold.

Returns:

Basis, complete singular spectrum, and rank metadata.

Raises:

ValueError – If the matrix or rank policy is invalid.

transformer_lens.tools.analysis.projection_kernel.projection_kernel(subspace_a: SubspaceBasis, subspace_b: SubspaceBasis, *, check_orthonormal: bool = True) → ProjectionKernelResult¶

Measure overlap between two explicitly extracted subspaces.

Raw PK lies in [0, min(rank_a, rank_b)]. The normalized value is PK / sqrt(rank_a * rank_b), the cosine between the two projection matrices. Principal angles are returned in radians.

transformer_lens.tools.analysis.projection_kernel.random_projection_kernel_moments(ambient_dim: int, rank: int) → RandomSubspaceReference¶

Return PK moments for independent Haar-distributed rank-rank planes.

These idealized descriptive moments are not calibrated p-values for trained model weights, whose head subspaces are dependent and anisotropic.