feat(jacobian_lens): add occupancy and fraction-of-variance

[janmenjayap]

Description

Adds occupancy and fraction-of-variance to the Jacobian lens: two J-space profiling statistics that sit on top of the sparse decomposition (#1596), following Gurnee et al. (2026), "Verbalizable Representations Form a Global Workspace in Language Models" (Transformer Circuits Thread). They quantify two of the paper's qualitative claims: that "only a small number of J-lens vectors are strongly active at a time" (occupancy) and that the J-space is "never more than about 10%" of activation variance, "a median of 6 to 7%" for concept vectors (fraction of variance).

This is the Tier-A follow-up the maintainer explicitly deferred from #1596; it builds directly on that PR's greedy selection and cached J-lens dictionary. Part of #1539 (Tier-A follow-up to #1596). No new dependencies.

New public surface (TransformerBridge only, matching the rest of the Jacobian lens)

Model-free core in transformer_lens/tools/analysis/jacobian_lens_decomposition.py:

• estimate_occupancy(x, dictionary, *, max_atoms=25, num_control_dictionaries=32, seed=0) returns a JSpaceOccupancy (occupancy, marginal_captured_variance, control_captured_variance, support).

On JacobianLens (additive; existing methods unchanged):

• occupancy(model, activation_or_prompt, layer, *, position=None, max_atoms=25, num_control_dictionaries=32, seed=0) resolves a raw [d_model] activation or a prompt position exactly as decompose, then calls estimate_occupancy on the cached full-vocabulary dictionary.
• fraction_of_variance(model, prompts, layers=None, *, k=25, skip_first=16, positions=None, show_progress=False) profiles a prompt corpus and returns a JSpaceVarianceProfile (layers, median, pooled, per_position).

JSpaceOccupancy, JSpaceVarianceProfile, and estimate_occupancy are exported from transformer_lens.tools.analysis.

Occupancy: max-separation versus a random-dictionary control

estimate_occupancy runs the same greedy selection as get_sparse_decomposition up to max_atoms, recording the per-step captured variance ||Pi_S x||² / ||x||² (orthogonal projection onto the selected span). It runs the same greedy on num_control_dictionaries random unit-norm dictionaries of equal size and averages their curves. The occupancy is the step of maximum separation between the real and control cumulative captured variance: the point past which further vectors add no more than random directions would. It is deterministic (seeded) and needs no threshold.

An earlier per-step quantile-cutoff variant was calibration-sensitive (planted structure wanted q ≤ 0.75, a clean noise floor wanted q ≥ 0.9, with no single robust value). The max-separation rule recovers planted sparsity exactly on orthonormal dictionaries (k=1 gives 1, k=4 gives 4, k=7 gives 7) with no tuning, so quantile was dropped.

Fraction of variance: the span-projection operationalization

For each (layer, position) at or past skip_first (mirroring the fit's early-position skip), fraction_of_variance records ||j_space_component||² / ||activation||² and reports, per layer, the median of those fractions and the pooled ratio Σ||j_space||² / Σ||activation||². The numerator is the selected_support span projection (the paper's appendix operationalization), not the nonnegative reconstruction; the two coincide only when every selected atom stays active. This is named explicitly in the docstring and docs so the number is not misread. A layer that samples no positions (every prompt shorter than skip_first, or only zero-norm activations) yields NaN median and pooled values and an empty per_position. This is a documented, tested contract.

Tests

• Model-free unit (test_jacobian_lens_decomposition.py): planted-sparsity recovery (occupancy == k), ordering by planted density, determinism, and input validation.
• Wrapper unit (test_jacobian_lens.py): occupancy raw and prompt paths, fraction_of_variance median and pooled over a corpus, unfitted-layer and empty-corpus rejection, and the no-sample NaN contract.
• gpt2-small integration (test_jacobian_lens.py): occupancy is a small positive integer with [0, 1]-bounded cumulative captured-variance curves and is seed-deterministic; fraction_of_variance median and pooled land in [0, 1], including the positions= override path over a multi-prompt corpus.

Honesty note

The paper's quantitative occupancy and variance figures are closed-model results (Sonnet, Haiku, Opus). Tests and docs assert shape (occupancy is a small positive integer; the variance fraction is a small ratio in [0, 1]), not the paper's exact numbers, which are not expected to transfer to open weights.

Docs

A new "Occupancy and fraction of variance" subsection in the "Sparse decomposition" section of jacobian_lens_fitting.md, with runnable snippets and the open-weight, shape-only caveat.

Type of change

• New feature (non-breaking change which adds functionality)
• This change requires a documentation update

Screenshots

Not applicable. This adds an analysis API plus a docs page, with no visual output.

Checklist:

• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my feature works
• New and existing unit tests pass locally with my changes
• I have not rewritten tests relating to key interfaces which would affect backward compatibility

Verification run locally

The change is confined to transformer_lens/tools/analysis/jacobian_lens*, so the local run was scoped to the Jacobian-lens surface; CI re-runs the full suite across Python 3.10, 3.11, and 3.12.

• check-format: clean under the CI-pinned black 23.12.1 plus isort (profile=black) on all changed files.
• mypy: Success, no issues in the two analysis modules.
• unit plus integration plus docstring (not slow): one pytest invocation over tests/unit/tools/test_jacobian_lens.py, test_jacobian_lens_decomposition.py, tests/integration/test_jacobian_lens.py, and (docstring tier, --doctest-modules --doctest-plus) both analysis modules: 246 passed, 4 deselected (the deselected are the slow gemma-2-2b-it cases).
• build-docs (Sphinx): builds clean (exit 0) with notebooks excluded; the new page and API entries render. The only warnings touching the new dataclasses are the pre-existing duplicate object description re-export pattern (identical to JSpaceDecomposition; -W is disabled in conf.py for exactly this), so no new warning class is introduced. A full local build-docs additionally needs pandoc for the demo notebooks, absent on this box and present in CI.
• Real-model integration: GPT-2 occupancy returns a small positive integer with [0, 1]-bounded captured-variance curves, and fraction_of_variance returns per-layer median and pooled ratios in [0, 1].