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:
ValueErrorRaised when per-direction attribution is requested for a degenerate direction.
Subclasses
ValueErrorso callers that alreadyexcept ValueErrorkeep working, mirroring how the other analysis tools raiseValueErrorfor their input guards.
- class transformer_lens.tools.analysis.svd_circuits.HeadDecomposition(layer: int, head: int, QK: HeadSVD | None = None, OV: HeadSVD | None = None)¶
Bases:
objectContainer returned by
decompose_head().QKandOVhold theHeadSVDfor each requested map, orNonewhen that map was not requested.- 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:
objectFactored 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 isU @ S.diag() @ V.transpose(-2, -1).- Type:
jaxtyping.Float[Tensor, ‘d_model rank’]
- rank_report¶
Per-direction
RankReportRowlist, aligned with the columns ofU/V.
- 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
iis refused: rotation-ambiguous inside a block, or null.
- layer: int¶
- null_rtol: float¶
- rank_report: List[RankReportRow]¶
- require_isolated(i: int) None¶
Raise
DegenerateDirectionErrorunless directioniis 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:
objectOne singular direction’s summary, aligned with column
idxofU/V.- idx¶
Position of the direction, matching the column index in
UandV.- Type:
int
- sigma¶
The singular value for this direction.
- Type:
float
- sigma_ratio¶
sigmanormalized 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_ratiofalls belownull_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].attnaccessors 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 tod_model * torch.finfo(S.dtype).epsper map, matchingtorch.linalg.matrix_rank()’s default tolerance.
- Returns:
A
HeadDecompositionwhoseQK/OVfields hold aHeadSVDfor each requested map.- Raises:
ValueError – If
layerorheadis out of range, orwhichis empty, a bare string, or contains an unknown entry.