tmax-4b brain atlas activation census, OV-circuits, and capability fence (32 layers)
allenai/tmax-4b Brain Atlas — Mid-Size Hybrid, Full Atlas
Cross-post: I ran a full GWIQ-style brain atlas on the middle member of the tmax family. It's not a downstream benchmark, just a look at what the tensors are actually doing.
model: allenai/tmax-4b
atlas type: activation census + Sub-Zero brain atlas + OV-circuit SVD + logit lens + compliance/behavior contrast
corpus: 8,965 prompts
layers: 32
attention layers: 3, 7, 11, 15, 19, 23, 27, 31
hybrid layers: literally everything else
sacred (fully probed) layers: all 32
datasets: juiceb0xc0de/tmax-4b-atlas
What this is
TMax-4B is a hybrid SSM/Mamba/transformer architecture. Only one in four layers is a full multi-head attention block. The rest are linear-attention / SSM-style layers with components like linattn_qkv, linattn_z, linattn_out, gate, up, and mlp. I wanted to see whether a 4B hybrid could still build clean late-layer subspaces and surgical directions, or whether the extra width just makes interpretability noisy.
Short answer: it is surprisingly clean. The census is complete, the Sub-Zero pass finished, and the model is mostly editable once you know where not to cut.
What was run
- Activation census over 8,965 prompts.
- Per-layer feature taxonomy for every component that tokenizes language.
- Per-head analysis on the 8 full attention layers.
- OV-circuit SVD (
W_V @ W_O) on every head. - Logit-lens pass to see which internal directions predict output tokens.
- Coactivation and code-analysis passes.
- Compliance/behavior contrast pass.
- Sub-Zero surgery pass with capability fence across
code,math,reasoning,factual, andmultilingual.
The shape of the thing
| Property | Value |
|---|---|
| Layers | 32 |
| Attention layers | 8 |
| KV heads | 4 |
| Head dim | 256 |
| d_model | 4096 implied |
| Hybrid components | gate, up, mlp, linattn_out, linattn_qkv, linattn_z |
What the tables contain
| Table | Rows | What it gives you |
|---|---|---|
layers |
32 | layer metadata and completion flags |
features |
1,376,256 | feature taxonomy + activation stats per (layer, component, feature_idx) |
per_head |
448 | per-head selectivity on the 8 attention layers |
ov_circuits |
128 | SVD over W_V @ W_O plus QK/FC spectral metrics |
logit_lens |
8,704 | promoted/suppressed output tokens per feature |
coactivation |
17,356 | feature-pair correlations |
code_analysis |
6,240 | entangled vs selective role labels |
compliance_behaviour_features |
1,376,256 | authentic-vs-corporate contrast per feature |
compliance_behaviour_per_head |
448 | per-head compliance/behavior contrast |
subzero_layer |
32 | classifier accuracy and SV summary per layer |
subzero_svs |
309 | bouncer singular values per projection |
subzero_capability |
795 | DAS-axis damage and capability-fence decisions |
What the numbers suggest
Attention is doing distributed computation, not copy-paste
OV-circuit spectral concentration is 0.043 with effective rank around 89.5. That is a broad signature, not a memorized token-to-token circuit. The QK path is a bit more structured at 0.195, but the attention itself looks like weighted high-dimensional computation, not sparse lookup.
MLP/SSM gates dominate the logit lens
The strongest logit-lens directions are gate features in the middle-to-late layers. Layer 23 gate feature 1326 hits F-stat 717, and layer 31 gate feature 4626 is around 657. In a dense transformer you might expect late attention to dominate the logit lens. Here the hybrid MLP gates are doing a lot of output vocabulary routing.
Features are broad, with a growing specific tail
The taxonomy is dominated by partial_shared and broadly_shared, with very few specific_* directions. Most dimensions respond to many prompts rather than one weird niche trigger.
| Class | Count | Share |
|---|---|---|
partial_shared |
448,419 | 32.6% |
non_activated |
372,827 | 27.1% |
broadly_shared |
345,429 | 25.1% |
all_shared |
209,269 | 15.2% |
specific_* |
312 | <0.03% |
The specific_* tail is still small as a fraction, but the model is beginning to show specialized directions for creative_writing (81), ml_ai (55), introspection (47), and tool_use (43).
v heads carry the cleanest signal
Per-head F-stats follow the same content-vs-routing pattern:
| Component | Avg F-stat best |
|---|---|
v |
263.5 |
heads |
217.8 |
q |
174.7 |
k |
153.8 |
v heads are the cleanest, k heads are the weakest. This is consistent with value vectors storing task-relevant content while query/key machinery mostly routes.
Code directions are mostly selective
Of the 6,240 code-analysis rows, 86.0% are labeled selective and 14.0% are entangled. The directions that do tokenize language tend to have focused jobs rather than being tangled up in everything.
The worst surgical surprise is at the very start
Classifier accuracy is solid across all 32 layers: 0.969–1.000, average 0.981. But the worst damage from removing a Sub-Zero direction is layer 0 linattn_out_proj axis 0, which fails the fence across all five capability domains with up to 10.86 damage. Layer 1 linattn_out_proj axis 0 is the next worst at 1.95 damage.
"Damage" here is a cross-entropy loss delta in nats per token. Sub-Zero's capability fence rejects anything above 0.15. A score of 10.86 means the ablated model assigned roughly e^10.6 (40,000×) less probability to the correct next token on average. That is not a data error; it is a real catastrophic entanglement between the compliance/behavior direction and core capability/language directions. The fence correctly refused to freeze that axis, which is exactly what it is supposed to do.
The structural conclusion does not change: an early linear-attention output projection is structurally unsafe to edit. Remove it and code, factual, math, reasoning, and multilingual all take a hit simultaneously.
Surgical headroom is decent
134 of 159 tested Sub-Zero axes pass the capability fence (84.3%). Average damage is 0.105. So aside from those early linattn directions, the network is reasonably editable.
What Sub-Zero is actually measuring here
The Sub-Zero pass is not a generic "find all important directions" sweep. It specifically looks for directions that separate corporate style from authentic style, then uses DAS rotation and the capability fence to check whether removing those directions damages code, math, reasoning, factual, or multilingual ability. So the 159 tested axes are compliance/behavior candidate axes, not a census of every load-bearing direction in the model. The layer 0 linattn axis happens to be a compliance/behavior direction that is also load-bearing, which is why the fence rejects it.
Compliance/behavior directions are only partly separated
Peak compliance-behaviour SV fraction is around 18.2%. Style/behavior directions are still partially entangled with capability directions at this scale.
The stuff I deliberately skipped
The atlas does not probe hybrid SSM layers that do not tokenize language. Measuring activations on those components just produced noisy numbers with no real structure, so I left them out. I am working on a way to capture whatever those hybrid components are actually doing, but that is a future post.
Caveats
- The layer 0
linattn_out_projdamage number is a real cross-entropy loss delta in nats/token, not a data error. The fence rejected the axis. - Specific-feature counts are still <0.03%, so do not over-read specialization.
- No Qwen3.5-4B base comparison here; that needs a separate atlas.
Bottom line
TMax-4B is a complete brain atlas. Stable classifiers, editable late layers, distributed attention, MLP/SSM gates that route to output tokens, and an early linear-attention projection that you should not touch. The hybrid architecture does not break interpretability. It keeps some of the load-bearing directions earlier than you would expect in a dense transformer.