What happens when you scale a hybrid SSM-Mamba-transformer: tmax-2B → 4B → 9B → 27B

R

juiceb0xc0de

Follow

Cross-post: I ran the same brain atlas pipeline on all four sizes of the allenai/tmax family. This is not a downstream benchmark; it is a look at how the internal geometry changes as you double the parameters and stack more hybrid layers.

models: allenai/tmax-2b, allenai/tmax-4b, allenai/tmax-9b, allenai/tmax-27b
atlas type: activation census + Sub-Zero brain atlas + OV-circuit SVD + logit lens + compliance/behavior contrast
corpus: 8,965 prompts for every run
attention cadence: one full multi-head attention block every four layers
datasets: juiceb0xc0de/tmax-2b-atlas, tmax-4b-atlas, tmax-9b-atlas, tmax-27b-atlas

What this is

Tmax is a hybrid SSM/Mamba/transformer family. Only one layer in four 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 scaling the recipe would change how the model splits work between attention and the hybrid path, or whether it would just make the same patterns louder.

Rather than relying on a single interpretability method, this analysis compares four views of the same models: activation census, OV-circuit SVD, logit lens, and Sub-Zero surgery with a capability fence. The interesting result is not any individual metric, but that they converge on the same scaling trends.

Short answer: within the tmax family, the signature stays, but scale pushes it in predictable directions. More width makes features broader. More depth spreads late-stage work over many layers instead of cramming it into the end. And the bigger models are more editable, because redundancy covers for the few load-bearing directions you should not touch.

One pattern that stands out immediately: doubling the parameters does not mean doubling attention. The 4B and 9B have the same 32 layers and the same 8 attention blocks; the 9B is wider. The 27B doubles depth to 64 but keeps only 16 attention blocks. So almost all the new capacity goes into the hybrid SSM/Mamba path, not into more attention heads. If you expected scaling to lean harder on attention, the data says the opposite.

The high-level picture is simple: width broadens the feature mix, depth spreads the late-stage routing over more layers, and both together give the network enough spare capacity that surgical edits stop being dangerous.

The shape of the family

The 4B and 9B share the same depth and attention cadence; the 9B is wider. The 27B doubles the depth while keeping the same 25%-attention cadence, so it ends up with 48 hybrid layers and only 16 attention blocks.

What the numbers suggest

Feature taxonomy shifts toward broad and partial

The bigger the model, the fewer dead directions and the more partially-shared ones. Universal all-shared directions also shrink as the model learns to specialize broad directions to narrower prompt types without going fully hyperspecific.

The specific_* tail is still tiny as a fraction, but the raw counts grow: 94 rows in the 2B, 312 in the 4B, 312 in the 9B, and 1,162 in the 27B. The 27B finally shows meaningful counts for things like tool_use (351), ml_ai (284), and creative_writing (272). Still a small fraction of the atlas, but the model is starting to build domain-flavored directions rather than just broad prompt detectors.

Even at 27B, hyperspecific directions are under 0.03% of the census. The real scaling trend is not an explosion of niche detectors; it is the broad and partial-shared classes eating away at dead and universally-shared directions.

Attention stays distributed across the whole family

OV-circuit spectral concentration stays low and effective rank stays high. The heads are doing weighted high-dimensional computation, not sparse token-to-token lookup.

The 4B and 9B move toward more distributed attention than the 2B. The 27B is roughly in line with the 9B, maybe slightly less concentrated because depth spreads the same attention budget thinner. QK and FC paths are more structured than OV, which is the same pattern you see in dense transformers: the routing machinery is more regular than the value computation.

One caveat: the 9B Sub-Zero run is a late-layer snapshot, concentrated in layers 16–30 with sparse early and very sparse tail coverage. Its cleaner-looking FC and QK numbers are partly a function of which layers got measured, not necessarily a statement about the full model.

v heads carry the cleanest signal, but depth spreads it out

Per-head F-stats show the same content-vs-routing pattern everywhere. v heads are strongest, k heads weakest.

The 4B and 9B are cleaner than the 2B. The 27B drops a bit on v and heads because 64 layers means value storage is distributed across more heads rather than concentrated in a few. q and k stay flat, which means the routing machinery is consistent across depth.

What this means is that value storage scales by adding more places to put content, while query/key routing stays a relatively fixed operation per head. The v heads are doing the semantic lifting; q and k are deciding where to look.

MLP/SSM gates dominate the logit lens, and depth spreads the peaks

The strongest output-vocabulary predictors are gate features in the middle-to-late layers. As depth grows, the peaks move deeper and spread over a wider range.

The 27B is wild. Its strongest vocabulary-routing directions are spread from layer 36 to layer 60, not crammed into the final few layers. That is a very different late-stage organization than a dense transformer.

The peak F-stat also rises with scale: 591 at 2B, 717 at 4B, 722 at 9B, and 863 at 27B. The gates do not just become more numerous; the cleanest ones get stronger as the model grows.

Why gates rather than attention heads? One plausible read is that the hybrid path has far more parameters per layer than the sparse attention blocks, and the sequential state updates let it accumulate vocabulary-routing context continuously. The attention blocks may still handle global routing corrections, but the bulk of late-stage token prediction gets built inside the gate-rich hybrid layers.

Code directions are mostly selective

Across the family, most code-analysis rows are labeled selective, meaning the directions that do tokenize language tend to have focused jobs rather than being entangled with everything.

The bigger models get slightly cleaner component roles. Scale helps the hybrid layers specialize. The selective share rises steadily from 80.9% at 2B to 87.0% at 27B, while entangled drops from 19.1% to 13.0%. That is a smooth trend, not a one-off from any single layer.

The worst surgical surprise moves deeper with scale

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 damage numbers are about compliance/behavior candidate axes, not every load-bearing direction in the model.

The 2B and 4B hit early linear-attention projections. The 9B hits a late gate. The 27B barely flinches. The 10.86 on the 4B layer 0 axis looks absurd next to the 0.6 on the 2B, but it is a real cross-entropy loss delta in nats per token. Sub-Zero's capability fence rejects anything above 0.15; 10.86 means the ablated model assigned roughly e^10.6 (40,000×) less probability to the correct next token on average. The fence correctly refused to freeze that axis. So the number is catastrophic, but it is also the system working as intended.

That progression is the clearest evidence that the model changes how it uses the hybrid path as it scales. Small tmax models lean on early linear-attention projections as load-bearing universal machinery, similar to how embeddings and first-attention layers matter in dense transformers. By 27B, that load has moved deeper, spread out, and become redundant enough that even the single worst compliance/behavior axis only does 0.031 damage.

Surgical headroom improves dramatically with scale

The 27B is so redundant that removing a top compliance/behavior axis barely harms any capability. The 4B looks worse than the 2B, but that is mostly the layer 0 outlier. The 9B has sparse early coverage, so its headroom number is more of a late-layer snapshot.

For anyone actually trying to edit these models: small tmax wants you to keep your hands off early linattn_out and linattn_z projections. Big tmax is safer to touch, but the gate-driven routing is spread across a 25-layer late region, so a single-layer edit is unlikely to catch the whole story.

Compliance/behavior directions stay only partly separated

Peak compliance-behaviour SV fraction stays low across the family. Style/behavior directions remain entangled with capability directions.

The 4B shows the highest peak at ~18.2%, while the 9B peak is only ~1.2%. That 9B number is almost certainly a coverage artifact: the Sub-Zero pass there is concentrated in layers 16–30, and the compliance/behavior singular value is measured per layer, so a sparse snapshot can miss the layer where the separation is strongest. Treat the 9B row as a lower bound, not a fair comparison.

Even with that caveat, none of these models have clean orthogonal style subspaces at the top singular value level.

What is novel vs a dense transformer

1. Depth grows while attention density stays fixed. A dense 27B transformer would usually be ~40 layers. Tmax doubles depth to 64 and pushes computation into the hybrid path. The attention cadence does not change.
2. MLP/SSM gates do the heavy output-vocabulary routing. The strongest logit-lens directions are gate features, not late attention heads. This holds at every scale.
3. Output routing spreads over a deep late region. In dense transformers, late attention and final MLP layers dominate. In tmax, gate features from layer 36 to layer 60 drive the strongest output-token predictors.
4. Redundancy grows with scale. The 27B is far more editable than the smaller models. More parameters do not just mean bigger matrices; they mean the model has room to keep style/behavior directions away from the truly load-bearing machinery.
5. Early hybrid projections are the real landmine. At every scale, the worst surgical surprise is an early linattn_* or gate direction that the rest of the network depends on. In a dense transformer that role is usually handled by the embedding and first few attention layers; here it is inside the hybrid path.

The stuff I deliberately skipped

Hybrid SSM layers that do not tokenize language were not probed because the activations would just be noise. The current atlas is built around language-aligned feature projections. Several SSM state updates do not project cleanly into token-aligned representations, so applying the same probe would mostly measure projection artifacts rather than meaningful internal state. I am working on a way to capture whatever those hybrid components are actually doing, but it is not in this atlas yet.

Caveats

• The Sub-Zero pass measures compliance/behavior candidate axes, not every important direction in the model.
• The 4B layer 0 linattn_out_proj damage score is a real cross-entropy loss delta in nats/token. The magnitude reflects genuine catastrophic entanglement, not a data error. The fence rejected the axis.
• The 9B Sub-Zero coverage is concentrated in layers 16–30, with sparse early and very sparse tail data. Treat its headroom number as a late-layer snapshot.
• The 27B Sub-Zero coverage is present across layers 0–31 and nearly absent from layers 33–52 except for isolated points. Claims about the very deepest layers lean more on census and logit-lens data than on full SV surgery.

Bottom line

Scaling the tmax recipe does not break interpretability. It makes the same patterns more pronounced: broader features, distributed attention, gate-driven output routing, and enough redundancy at 27B that compliance/behavior edits barely hurt capabilities. The hybrid architecture moves some load-bearing directions earlier than a dense transformer would, but it also gives the model more places to hide spare capacity.

Tmax is one recipe, and it is an open question whether other hybrid SSM families will follow the same trajectory. The bet this atlas lays out is simple: fixed attention density + scaled hybrid depth → broader features, deeper routing, and more redundancy. The next step is to run the same four-probe pipeline on another hybrid family and see whether the gates, the depth, and the editability scale the same way. This family is a clean testbed for that question.

If neural network atlases are something you're into make sure to check out my collection