transformer_lens.tools.analysis.svd_circuits module

Singular-vector decomposition of a single attention head’s QK and OV maps.

An attention head is characterized by two low-rank linear maps: the query-key map W_Q W_K^T that scores source positions, and the output-value map W_V W_O that writes the attended value back into the residual stream. This tool takes the singular value decomposition of each map for one head and exposes the singular values (how much each direction matters) together with the left and right singular vectors (the output and input directions they act on).

The decomposition is weight-space only: it reads W_Q/W_K/W_V/W_O and needs no forward pass, no activation cache, and no compatibility mode. Each map is kept factored through FactoredMatrix, so the d_model x d_model product is never materialized and the returned rank is bounded by d_head.

Right singular vectors are read from V (its columns are the right singular vectors). The historical .Vh alias is deprecated and returns the same tensor, so it is never used here.

Adjacent singular values closer than a relative gap eps leave their singular directions defined only up to a rotation, so the result carries a per-direction degeneracy report. Directions are grouped into contiguous blocks that end only at a gap of at least eps; callers attribute such a block as a subspace instead of trusting a single, rotation-dependent direction inside it.

A singular value near zero relative to the top of the spectrum is null rather than near-equal: its singular vector is an arbitrary null-space direction, not a rotation of a comparable neighbour. Null directions are flagged under their own tolerance, null_rtol, keyed to the spectrum’s top value the way torch.linalg.matrix_rank() keys its default tolerance.

Example:

from transformer_lens.model_bridge import TransformerBridge
from transformer_lens.tools.analysis.svd_circuits import decompose_head

model = TransformerBridge.boot_transformers("gpt2", device="cpu")
decomposition = decompose_head(model, layer=0, head=0)
ov = decomposition.OV
for row in ov.rank_report:
    print(f"direction {row.idx}: sigma={row.sigma:.3f} ratio={row.sigma_ratio:.3f}")
exception transformer_lens.tools.analysis.svd_circuits.DegenerateDirectionError

Bases: ValueError

Raised when per-direction attribution is requested for a degenerate direction.

Subclasses ValueError so callers that already except ValueError keep working, mirroring how the other analysis tools raise ValueError for their input guards.

class transformer_lens.tools.analysis.svd_circuits.HeadDecomposition(layer: int, head: int, QK: HeadSVD | None = None, OV: HeadSVD | None = None)

Bases: object

Container returned by decompose_head().

QK and OV hold the HeadSVD for each requested map, or None when that map was not requested.

OV: HeadSVD | None = None
QK: HeadSVD | None = None
head: int
layer: int
class transformer_lens.tools.analysis.svd_circuits.HeadSVD(which: Literal['QK', 'OV'], layer: int, head: int, U: Float[Tensor, 'd_model rank'], S: Float[Tensor, 'rank'], V: Float[Tensor, 'd_model rank'], rank_report: List[RankReportRow], eps: float, null_rtol: float)

Bases: object

Factored SVD of one head map ("QK" or "OV") with a degeneracy report.

which

Which map this decomposes, "QK" or "OV".

Type:

Literal[‘QK’, ‘OV’]

layer

Layer of the decomposed head.

Type:

int

head

Head index within the layer.

Type:

int

U

Left singular vectors, [d_model, rank] (column i is output direction i).

Type:

jaxtyping.Float[Tensor, ‘d_model rank’]

S

Singular values, [rank], sorted descending.

Type:

jaxtyping.Float[Tensor, ‘rank’]

V

Right singular vectors, [d_model, rank] (column i is input direction i). The reconstruction is U @ S.diag() @ V.transpose(-2, -1).

Type:

jaxtyping.Float[Tensor, ‘d_model rank’]

rank_report

Per-direction RankReportRow list, aligned with the columns of U/V.

Type:

List[transformer_lens.tools.analysis.svd_circuits.RankReportRow]

eps

Relative gap below which adjacent directions share a block; every block boundary sits at a gap of at least eps.

Type:

float

null_rtol

Relative-to-top-singular-value tolerance below which a direction is numerically null.

Type:

float

S: Float[Tensor, 'rank']
U: Float[Tensor, 'd_model rank']
V: Float[Tensor, 'd_model rank']
block_of(i: int) List[int]

Return every direction index sharing direction i’s degeneracy block.

The result is a singleton [i] for an isolated direction and the full run for a degenerate one.

degenerate_blocks() List[List[int]]

Return every degenerate block’s indices: the subspaces to attribute whole or skip.

eps: float
head: int
is_degenerate(i: int) bool

Whether direction i is refused: rotation-ambiguous inside a block, or null.

layer: int
null_rtol: float
rank_report: List[RankReportRow]
require_isolated(i: int) None

Raise DegenerateDirectionError unless direction i is attributable alone.

The message names the cause, since a rotation-ambiguous block is still a subspace worth attributing while a null block carries no signal.

which: Literal['QK', 'OV']
class transformer_lens.tools.analysis.svd_circuits.RankReportRow(idx: int, sigma: float, sigma_ratio: float, is_degenerate: bool, is_null: bool, block_id: int)

Bases: object

One singular direction’s summary, aligned with column idx of U/V.

idx

Position of the direction, matching the column index in U and V.

Type:

int

sigma

The singular value for this direction.

Type:

float

sigma_ratio

sigma normalized by the largest singular value, in [0, 1].

Type:

float

is_degenerate

True when this direction is not attributable on its own: it shares a block with a neighbour (defined only up to a rotation within that block) or it is numerically null.

Type:

bool

is_null

True when sigma_ratio falls below null_rtol, so the singular vector is an arbitrary direction from the map’s null space.

Type:

bool

block_id

Index of the contiguous block this direction belongs to.

Type:

int

block_id: int
idx: int
is_degenerate: bool
is_null: bool
sigma: float
sigma_ratio: float
transformer_lens.tools.analysis.svd_circuits.decompose_head(model, layer: int, head: int, *, which: Sequence[str] = ('QK', 'OV'), eps: float = 0.01, null_rtol: float | None = None) HeadDecomposition

Decompose a head’s QK (W_Q W_K^T) and/or OV (W_V W_O) maps via SVD.

Weight-space only: this reads the head’s per-block weights via the bridge’s model.blocks[layer].attn accessors and needs no forward pass and no compatibility mode. The returned factors are detached from the model.

Parameters:
  • model – A TransformerBridge.

  • layer – Layer of the head to decompose.

  • head – Head index within the layer.

  • which – Which maps to decompose, a non-empty sequence drawn from ("QK", "OV"). A bare string is rejected rather than iterated.

  • eps – Relative gap at which a block of adjacent directions ends; directions closer than this share a block.

  • null_rtol – Relative-to-top-singular-value tolerance below which a direction counts as numerically null. Defaults to None, which resolves to d_model * torch.finfo(S.dtype).eps per map, matching torch.linalg.matrix_rank()’s default tolerance.

Returns:

A HeadDecomposition whose QK/OV fields hold a HeadSVD for each requested map.

Raises:

ValueError – If layer or head is out of range, or which is empty, a bare string, or contains an unknown entry.