1908-scope/reports/dynamic-selection.md

Dynamic Per-Token Feature Selection: Can Each Token Pick Its Own SAE Features?

Hypothesis

Every SAE arm so far selects features the same way: it pools activations across the whole training set, picks one fixed set of top features per layer, and freezes that set for every token, forever. That fixed set is chosen by total activation magnitude, which is topic-agnostic — and we saw its cost directly in the feature-budget experiment, where Financial never recovered no matter how many features we added, because its discriminative features are reliably present but not loud enough on average to survive the global pool.

This experiment asks the obvious counter-question: what if, instead of one frozen set, every token picks its own top features — the ones actually firing hardest for that token — while keeping each feature's identity intact? The hope is that letting each token surface its own evidence recovers the topics the global pool starves (Financial especially) and closes most of the remaining gap to raw hidden states, while every input dimension is still a named, interpretable SAE feature.

The Current Approach (what "fixed selection" does)

To make the contrast clear, here is exactly what the existing SAE caches (the 500- and 1500-wide ones) do, step by step, for one layer:

1. Take every token of every training record — one big pool of ~374,546 tokens. Sample boundaries are ignored; it is one flat bag of tokens.
2. Run each token through that layer's SAE → a 16,384-long activation vector.
3. Sum each feature's activation across the entire pool → one total per feature.
4. Take the top-K (50 or 150) features by that total. Freeze those indices.

Those K indices are then the same for every token, train and test. Each token keeps only its activation values at those fixed indices. So the representation is a fixed coordinate system: the same 500 (or 1500) columns for everyone, and the numbers in them vary per token.

The weakness is Step 3: selection by pooled magnitude keeps the features that are loud across the corpus, which biases toward generic high-frequency features and can drop features that fire reliably but only within one narrow topic. That is the mechanism behind the stuck Financial score.

What We Did (dynamic selection, explained)

Dynamic selection changes only which features each token keeps — nothing about the SAE or the base model:

• For every token, in every layer, keep that token's own top-150 features by JumpReLU magnitude. Different tokens keep different features.
• Crucially, keep each feature's identity by position: a value is placed in the column that belongs to its feature, not in an anonymous "slot 3." Everything else in that token's vector is zero.

Why identity-by-position matters (and why the input is wide)

The single most important design point — the one that trips everyone up — is what the model is fed. It is fed the activation values, and each value's column encodes which feature it is. The model never sees an index as a number.

A tiny worked example makes it concrete. Suppose a layer had only 6 features and we kept the top-2 per token, for two tokens:

• Token A's top-2: feature 2 (value 0.9), feature 5 (value 0.4)
• Token B's top-2: feature 5 (value 0.8), feature 1 (value 0.3)

If we packed these into a compact 2-wide vector — A = [0.9, 0.4], B = [0.8, 0.3] — then "slot 0" is feature 2 for A but feature 5 for B. The model's weight for slot 0 would have to mean two unrelated features at once, and it learns nothing. This is the wrong way, and it is why a narrow "just feed the 150 values" input does not work for a dynamic scheme.

The right way is to put each value at its feature's own column:

	feat1	feat2	feat5
Token A	0	0.9	0.4
Token B	0.3	0	0.8

Now the feat5 column is always feature 5. The model can learn a stable weight per feature. This is exactly bag-of-words: a sentence has ~20 words but the input is vocabulary-wide (say 50,000), because "word 3" differs across sentences — each word needs its own dimension. Here each feature needs its own dimension.

So the input width is the size of the feature vocabulary — how many distinct features any token might pick — not how many are active per token. That is why dynamic is wide and fixed is narrow: fixed shares one small set of features across all tokens (500/1500 columns), while dynamic lets any of thousands of features appear, each needing a permanent home.

Where the width number comes from

For each layer we take the union of every feature that appears in some training token's top-150. That union is not a pooling of token data — every token still keeps its own separate 150; the union only counts how many distinct columns the table needs. Measured per layer, the union ranges from ~6,500 (layer 16) to ~14,500 (mid layers), except the last layer (35). We cap each layer at its 6,000 most-frequently-picked features (an engineering bound; it keeps 99.3% of all per-token picks). The total width is:

$$\mathrm{6,000}\times 9\text{ layers}+35=\mathrm{54,035.}6\{,\}000\; \backslash times\; 9\; \backslash text\{\; layers\}\; +\; 35\; =\; 54\{,\}035.$$

Each token fills only ~1,136 of those 54,035 columns (its own picks); the rest are zero. The data is ~98% sparse, which is why we can store and train it cheaply (below).

Configuration

• Base model & features: Gemma 3 1B instruction-tuned, Gemma Scope SAEs (layers 16–25, JumpReLU), identical to the fixed arms. Per token, per layer: top-150 features by magnitude, placed by identity.
• Cache: stored sparse (col indices + values + per-token pointers), 3.66 GB, width 54,035, kept fraction 99.3%. Labels/roles/splits are byte-identical to the fixed 500/1500 caches, so results are directly comparable.
• Training: the sparse cache is loaded resident on the GPU and each batch is densified on-GPU, so there is zero per-batch CPU↔GPU transfer and no disk reads after load. This brought training from 180 s/epoch (dense mmap over a network disk) down to **6–10 s/epoch**. Same hyperparameters as the fixed arms (AdamW, 15 epochs, best-validation-macro-F1 checkpoint).
• Architectures: the same three — GRU (256×3), Transformer (256, 8 heads, 2 layers, ff 512), ConvNeXt (256 ch, depth 6, kernel 7). The input layer is now vocabulary-sized (54,035), so the models are much larger than the fixed versions.
• Scoring: response tokens only, identical to the fixed-arm scorer, so numbers line up with the no-SAE ablation report.

Two measured facts worth recording, because they shape the interpretation:

• These SAEs are not very sparse in early layers. Mean active features per token (L0) is ~2,000 at layer 16, tapering to ~26 by layer 25. So the fixed top-150 was discarding a great deal in early layers, and dynamic selection has real room to differ token-to-token.
• The 10 layers are independent dictionaries. Cross-layer decoder cosine similarity is ~0.02–0.10 index-aligned (essentially orthogonal): feature j in layer 16 is unrelated to feature j in layer 20. This is why the layers must be concatenated (each its own column block), not merged into one shared 16,384-wide dictionary.

What We Found

Dynamic selection beats fixed-150 on every architecture and every metric, with no regressions, and lands close to raw hidden. Overall, response tokens:

Model	Metric	SAE-50	SAE-150	Dynamic	Raw	dyn − 150
GRU	Accuracy	0.9350	0.9475	0.9757	0.9821	+0.0282
	Macro F1 (=Dice)	0.9339	0.9439	0.9746	0.9812	+0.0307
	Macro IoU	0.8774	0.8949	0.9507	0.9634	+0.0558
Transformer	Accuracy	0.9662	0.9703	0.9761	0.9806	+0.0059
	Macro F1 (=Dice)	0.9650	0.9687	0.9749	0.9803	+0.0063
	Macro IoU	0.9327	0.9396	0.9514	0.9615	+0.0118
ConvNeXt	Accuracy	0.9435	0.9695	0.9806	0.9856	+0.0110
	Macro F1 (=Dice)	0.9394	0.9691	0.9799	0.9854	+0.0107
	Macro IoU	0.8869	0.9405	0.9606	0.9713	+0.0201

Params (M)	SAE-50	SAE-150	Dynamic	Raw
GRU	1.37	2.14	42.49	9.84
Transformer	1.18	1.44	14.89	4.01
ConvNeXt	3.30	3.56	17.01	6.12

Measured as fraction of the SAE-50 → raw IoU gap closed, dynamic reaches ~85% (GRU), ~65% (Transformer), ~87% (ConvNeXt) — versus 20–64% for fixed-150. For the GRU especially, dynamic (0.951) is within 0.013 IoU of raw (0.963).

Per-topic IoU, fixed-150 → dynamic, with raw in brackets:

Topic	GRU	Transformer	ConvNeXt
Enterprise documents	0.810 → 0.908 [0.921]	0.890 → 0.903 [0.942]	0.915 → 0.927 [0.956]
General news & content	0.894 → 0.944 [0.953]	0.941 → 0.953 [0.965]	0.962 → 0.971 [0.972]
Customer service	0.885 → 0.966 [0.984]	0.942 → 0.967 [0.969]	0.964 → 0.969 [0.981]
Legal	0.927 → 0.958 [0.962]	0.954 → 0.958 [0.958]	0.956 → 0.963 [0.973]
Financial	0.873 → 0.947 [0.973]	0.927 → 0.951 [0.955]	0.891 → 0.958 [0.955]
HR & people operations	0.914 → 0.953 [0.966]	0.944 → 0.945 [0.961]	0.923 → 0.953 [0.977]
Healthcare	0.962 → 0.979 [0.984]	0.978 → 0.982 [0.980]	0.973 → 0.983 [0.984]

The headline per-topic result: Financial is rescued. It was the one topic that more features could not fix — ConvNeXt Financial IoU was stuck at 0.891 whether the fixed budget was 50 or 150, against raw's 0.955. Dynamic selection takes it to 0.958, matching raw. This is direct confirmation of the hypothesis: Financial's discriminative features are reliably present per-token but not loud enough to survive the global pool, so letting each token pick locally surfaces them. The diffuse/procedural domains the fixed arms starved (Customer service, Enterprise, General news) also improve most.

The one real cost is model size, not accuracy. Because the input is 54,035-wide, the input layer balloons — the dynamic models are larger than even the raw-hidden models (GRU 42.5M vs raw 9.84M). So dynamic is not a compactness win. Its value is different: near-raw-hidden accuracy while every active dimension is still an interpretable SAE feature, which raw hidden is not. If we want the accuracy and small models, the input projection is the thing to shrink (an embedding-table / learned down-projection over the active features — same information, ~30× less compute).

Next Steps

Two concrete directions for now (more will follow once we have more theory):

• Ablate the two budgets: how many layers, and how many features per token. We fixed both at "all ten layers" and "top-150 per layer." Both are free knobs. Sweep the per-token K (e.g. 50 / 150 / 300) and the number of layers used (e.g. last 4 vs last 10), to find where accuracy saturates and how small the vocabulary/width can get before it hurts. This also directly addresses the size cost above.
• Build more theory for why this works. We have an empirical result and an intuition (per-token selection surfaces topic-specific-but-quiet features the global pool drops). It is worth formalizing: characterize which features dynamic keeps that fixed drops, quantify per-topic feature signatures and their overlap across tokens, and connect the observed early-layer high-L0 / late-layer low-L0 structure to why dynamic helps some layers more than others.

References

• Gemma 3 1B, instruction-tuned base model (Google); Gemma Scope residual-stream SAEs, JumpReLU activations (Google DeepMind).
• Companion reports: the no-SAE / feature-budget ablation (SAE-50 / SAE-150 / raw hidden) and the decode ablation.
• Artifacts: scripts/build_dynamic_sparse.py (cache), scripts/train_dynamic_gpu.py (GPU-resident training), scripts/score_dynamic.py (scoring); cache at cache/dyn150_sparse/, results at results/dyn150_2k/.