Writeup - Co-Activation Coverage Diagnostic for SAE Hierarchies

Motivation

Several recent SAE variants (Matryoshka, Temporal) aim to recover hierarchical structure from language-model activations, but there is little quantitative testing of whether the recovered features form coherent parent-child relationships. This project builds a coverage-based diagnostic for that question, validates it on synthetic ground truth and on negative controls, and applies it to a real Gemma-2-2b Matryoshka SAE (SOAR project I-6).

The diagnostic

A coherent hierarchy has three measurable signatures: coverage (P(parent|child) ~ 1, the child implies its parent), asymmetry (coverage parent|child >> child|parent, distinguishing hierarchy from symmetric co-occurrence), and sibling exclusivity (children of one parent partition its activations, so high sibling overlap signals feature absorption). Because a frequently-firing parent covers everything by chance, the diagnostic always reports lift (P(parent|child)/P(parent)) and compares against a shuffled-pair null and a permutation p-value.

Part 1 - Synthetic validation

A planted 2-level tree (5 parents, 4 children each; 25 features, 20 edges), 4,000 samples, base rate 0.069, with tunable noise (p_drop, unrelated firing, and a p_absorb knob that injects absorption).

Recovers planted hierarchy. True edges vs random pairs: coverage 0.814 vs 0.042, lift 4.33, asymmetry 0.646, separation AUC 0.980, significant by permutation (observed 0.646 vs null 0.150, p = 0.003). Each child's best candidate parent is its true parent in 100% of cases, and unsupervised discovery recovers the full 20-edge tree at recall 1.00.

Does not hallucinate hierarchy. On independent features the diagnostic reports chance (asymmetry 0.00, AUC 0.54). On symmetrically co-firing pairs it reports high coverage (1.00) but ~zero asymmetry - the concrete demonstration that raw coverage alone would be fooled and the asymmetry term does real work.

Detects absorption / robust. Sibling overlap rises monotonically with injected absorption (0.006 -> 0.283); coverage falls gracefully under parent-drop noise while AUC holds ~0.98; a 3-level hierarchy shows transitive ancestor->descendant coverage (AUC 0.918); AUC is stable across thresholds.

Part 2 - Real SAE: Gemma-2-2b layer-12 Matryoshka

Applied to the layer-12 Matryoshka SAE (chanind/...-w-32k-l0-40), coarsest 2,048 latents, ~100 C4 documents. Activations use continuous tokenization (BOS only on the first chunk, to avoid flooding the statistics with BOS-token features), with stop_at_layer=13 and bfloat16 to fit a single GPU.

The raw signal is confounded. The aggregate cross-group test is weak (outer->inner coverage 0.091 vs null 0.056, AUC 0.605), and that weak signal is itself partly an artifact: coarse Matryoshka latents fire more often, and P(parent|child) - P(child|parent) is positive whenever the parent's base rate exceeds the child's, regardless of any real relationship. So coverage and asymmetry cannot be read directly.

Lift correction + best-parent. Correcting by base rate (lift) and taking each fine latent's single best coarse parent (1,920 children with a candidate parent):

The shuffle control permutes each latent independently - destroying co-activation while preserving base rates and sparsity. The real median best-parent lift (28.6) is about 9x the control (3.07), so coarse latents cover fine ones far above what base rate and sparsity alone produce. A genuine, above-chance co-activation hierarchy exists in the layer-12 Matryoshka latents.

Key takeaways

1. Asymmetry and sibling-exclusivity are what make this a hierarchy test, not a co-occurrence test; the symmetric-pairs control makes that concrete.
2. On the real SAE, the honest result lives in the median best-parent lift vs the shuffle control. The mean (83.7) is inflated by a heavy right tail (top-k sparsity caps lift at 1/P(parent)), and frac(>2) is uninformative because the max-over-candidates selection pushes even the shuffle to 0.77. Reporting the single big number without the control would have overstated the effect.

Layout

src/
  coverage.py      # coverage, lift, asymmetry, evaluate_hierarchy, sibling_overlap,
                   # best_parent, parent_recovery, permutation_pvalue, cross_group_coverage, propose_edges
  synthetic.py     # planted hierarchy (+ absorption knob), flat & symmetric controls, deep 3-level tree
  sae_loading.py   # SAELens loader for the Gemma-2-2b Matryoshka SAE + group boundaries  -- real path
  activations.py   # TransformerLens hooks -> SAE feature activations                      -- real path
tests/
  test_diagnostic.py   # 9 property tests (recovery, controls, absorption, p-value, robustness)
notebooks/
  01_mvp.ipynb     # full synthetic validation + controls (+ gated real-SAE path)
figures/

Quickstart

Tests: python tests/test_diagnostic.py (or pytest tests/)

Colab: open notebooks/01_mvp.ipynb, run the clone + pip cell, run all. Synthetic by default (USE_REAL = False, no downloads). Set USE_REAL = True for the SAE path (network + GPU + HF access).

Local:

pip install -r requirements.txt
jupyter notebook notebooks/01_mvp.ipynb

Limitations

• Establishes statistical hierarchy, not semantic: it does not yet show the high-lift pairs are meaningful parent-child concepts.
• One layer, one SAE, ~100 documents; lift estimates would tighten with a larger corpus.
• Binarizes at a fixed threshold (shown robust on synthetic, but magnitude-aware coverage is untested).

What I'd do next

• Semantic check: inspect the highest-lift parent->child pairs on Neuronpedia for coherence.
• Matryoshka vs standard SAE at the same layer - does the nested objective buy more hierarchy than a plain SAE? This is the falsifiable form of the project's title question.
• Add Temporal Feature Analysis as a third method; extend to PCFG-trained small transformers to isolate which distributional properties make hierarchy recovery fail.