Interpretation tools: SAE weights (english only) brain atlas SQLite database

#24
by juiceb0xc0de - opened

SAE weights + a full brain atlas for MiniCPM5-1B

Cross-post: I trained a sparse autoencoder on every layer of MiniCPM5-1B and then ran a full activation-census brain atlas on it. Neither is a downstream benchmark. Both are a look at what the tensors are actually doing. Full disclaimer: The SAE weights I trained were completed while trouble shooting a scheduler I'm currently developing for future release.


model: openbmb/MiniCPM5-1B
SAE weights: juiceb0xc0de/minicpm5-1b-SAE
atlas: juiceb0xc0de/MiniCPM5-1B-atlas
corpus: 8,965 prompts, English only
layers: all 24


The SAEs

One JumpReLU SAE per layer. d_in 1536, 49,152 features (32x expansion), streamed FineWeb-Edu, target sparsity L0=50. Same trainer, same settings, no hyperparameter changes across the whole run.

Layer EV L0 Dead Residual RMS
0 0.650 48.91 0.00% 0.133
1 0.621 48.81 0.01% 0.172
2 0.648 48.61 0.04% 0.228
3 0.728 49.16 0.43% 0.327
+0.262 EV 7.0x norm
4 0.990 50.18 0.06% 2.301
5 0.988 47.71 0.11% 2.372
6 0.987 49.77 0.22% 2.544
7 0.983 48.82 0.07% 2.356
8 0.983 50.41 0.09% 2.382
9 0.972 49.58 0.09% 2.408
10 0.970 50.22 0.58% 2.650
11 0.963 49.55 0.09% 2.493
12 0.954 47.69 0.09% 2.581
13 0.952 49.27 0.85% 2.765
14 0.948 48.32 0.29% 2.848
15 0.933 48.13 0.30% 3.061
16 0.924 47.14 0.03% 3.413
17 0.925 46.10 0.00% 3.770
18 0.918 46.70 0.14% 4.234
19 0.910 47.79 0.30% 4.776
20 0.903 47.42 0.00% 5.275
21 0.910 47.92 0.87% 6.133
22 0.905 48.93 0.51% 6.577
23 0.872 50.22 0.30% 7.087

All 24 layers, no gaps. Sparsity holds everywhere: every layer between 46.1 and 50.4 against a target of 50, and dead features under 1% on all twenty-four.

EV tracks the residual norm across the whole model

Put the two right-hand columns next to each other and the shape of this model is obvious. Residual RMS climbs monotonically from 0.133 at layer 0 to 7.087 at layer 23, a 53x increase end to end, and past the layer-4 cliff explained variance walks down in step with it: 0.990 at layer 4, 0.954 at 12, 0.933 at 15, 0.903 at 20, 0.872 at 23.

One number goes up, the other goes down, for nineteen straight layers. Global-variance EV is easier to satisfy when a few directions dominate the stream, so as the stream gets louder and more concentrated with depth, fifty features cover progressively less of it. There are two small ticks upward, layer 17 and layer 21, neither big enough that I could try to explain why. The brain atlas wasn't able to shed any light that I would call interpretable proof either.

The last layer is the extreme of both: highest norm in the model, lowest EV. I wouldn't say it shows the SAE giving up, it is layer 23 carrying more energy in fewer directions than anything before it.

Worth knowing if you use these: 0.87 to 0.91 in the back half is a usable basis, it is just not the 0.99 you get between layers 4 and 8. If you want the cleanest features in this model, work in the front half of the post-cliff range.

The floor is four stories up

Layers 0 through 3 do not decompose. 0.62 to 0.73 explained variance, and they need 4,200 to 4,700 steps to get there. Layer 4 lands 0.990 in 1,751 steps.

Layer 4 is not special because 0.990 is a big number. Layer 5 is 0.988 and layer 6 is 0.987. It is special because of the size of the step to get there. Every other layer-to-layer transition in this model moves EV by less than 0.04. The L3 to L4 step moves it by 0.262.

That is a cliff, not a ramp, and it lines up with the model rather than the trainer. Look at the residual RMS column. It sits between 0.133 and 0.327 for the first four layers, then jumps to 2.301 at layer 4. Seven times, across the same boundary, exactly where the EV cliff is. In per-token terms the vector norm goes from roughly 5 to roughly 90 in four layers.

Below the cliff the stream is small and flat, lots of comparable-magnitude directions and nothing for fifty features to grab onto. Above it a handful of directions carry the energy and fifty is plenty.

Here is the part I did not expect. This model has a 130,560-token bilingual vocabulary, and untied embeddings account for about 401M of its 1,081M parameters. Roughly 37% of MiniCPM5-1B is vocabulary. Four layers of dense, hard-to-decompose stream in front of a vocabulary that large is not a defect. You do not get bilingual coverage in a 1B model easily, and the first four layers look like they put in the work that makes it possible.

Two bugs in my own trainer

The weights come out of a trainer that is still under development.

The deep layers broke it. Layers 13 through 15 initially finished at L0 around 1,700 against a target of 50, which is not a K=50 SAE at all. Tracing it through the step histories: once sparsity locks and the controller pins the sparsity pressure, the feature tail starves. Dead features went 0.8% to 8.8% in a single 250-step window while EV moved 0.949 to 0.948. Then the emergency revival dumped thousands of features back in at once, and with the pressure term frozen nothing the trainer tried to do in response was going to pull the sparsity back down.

Two fixes. One releases the frozen pressure term when sparsity escapes its band. The other was embarrassing: a cached dead-feature set was being cleared on 249 of every 250 steps, so the gentle revival path was running 0.4% of the time and a metric that should have shown it was structurally pinned at zero.

The fix that actually worked is a hard 1% dead-feature ceiling with rollback: buffer recent states, and when dead breaches the ceiling, walk back to the last state under it. On layer 13 that cost one thousandth of EV for 8.6x fewer dead features.

The interesting part is what happened next. Layers 16 through 20 never needed it. Dead came in at 0.03%, 0.00%, 0.14%, 0.30% and 0.00%, and every one of them stopped through the normal quality gate with no rollback at all. So the starvation is not a property of depth. It is a band, layers 13 to 15, and past it the model goes back to behaving.

I may not have found any of this on a shallower model, or on one whose deep layers were less demanding. So the weights in this repo are honest, the deep ones are rolled-back states where the log says so, and MiniCPM5-1B is the thing that surfaced two real bugs in my trainer. That is a genuinely useful model to have run into.

The atlas

Activation census over all 24 layers and eight component groups (attn, gate, heads, k, mlp, q, up, v), with a combination of feature taxonomy, per-head analysis, OV-circuit SVD, logit lens, coactivation, and a Sub-Zero surgery probe with a capability fence. I combine these methods for their ability to compliment and enrich each other.

Table Rows
features 479,232
compliance_behaviour_features 479,232
logit_lens 7,680
coactivation 11,546
code_analysis 5,760
per_head 864
ov_circuits 384
subzero_capability 635
subzero_svs 503

Most of the model is safely editable

635 Sub-Zero tests, 127 per capability domain. 111 of 127 axes pass in every single domain, 87.4%, with mean damage between 0.026 and 0.046 nats per token.

Domain Pass rate Mean damage Max damage
code 87.4% 0.038 0.694
factual 87.4% 0.029 0.755
math 87.4% 0.028 0.323
multilingual 87.4% 0.046 0.932
reasoning 87.4% 0.026 0.233

Multilingual is the most sensitive domain by both mean and max, which is what you would expect from a model carrying two languages in one vocabulary.

Every failure appears early, in clusters, in the MLP

Projection Failures
gate_proj 40
up_proj 20
down_proj 20

No attention projection failed a single test. And the failures concentrate at the front: layers 0 through 4 hold 55 of 80, about 69%. The five worst individual results are all layer 3 or 4, topping out at layer 4 gate_proj axis 1 doing 0.932 damage to multilingual. A fraction of this damage can cause punctuation repetition and a total loss of communication capability in my experience.

The practical read for anyone editing this model: the mid-stack is roomy, the attention machinery is safe, and the first five layers deserve a light touch.

Two unrelated instruments point at the same four layers

The SAE says layers 0 through 3 cannot be reconstructed by fifty sparse features. The capability fence says layers 0 through 4 are where removing a direction breaks things. Those are completely different pieces of math. One is an unsupervised dictionary asking "can fifty directions span this?" The other is a rotational probe asking "what breaks if I take this away?"

Both land on the same band. A third angle agrees: mean feature separation F-stat climbs monotonically with depth, from 15.5 at layer 1 to 42.6 at layer 23, so early-layer directions are the least linearly separable in the model.

My interpretation is that the front of this model is dense and load-bearing. In a 1B model there may be no room for redundancy early, and this is what a well-packed small model looks like from the inside.

The one specific feature in the whole model is a bilingual type concept

Feature taxonomy across all 479,232 directions:

Class Count Share
partial_shared 165,617 34.6%
broadly_shared 138,553 28.9%
all_shared 100,544 21.0%
non_activated 74,517 15.6%
specific_core_technical 1 <0.01%

Exactly one direction out of 479,232 reads as domain-specific. Layer 18, gate, feature 3488. I went and looked at what it promotes:

Ġtype   0.513
类型     0.511
Ġtypes  0.504
type    0.419
types   0.408
ĠTypes  0.406
类型的   0.400
ĠType   0.388

Every single token is the word "type," in English or Chinese. Nothing else. Its logits are around 0.51, where typical logit-lens projections in this model sit at 0.12 to 0.23, so it is roughly four times stronger than the field. And its style-contrast F-stat is only 42.4 against a model top of 1,744. To me it is not entangled with tone or register at all.

The single most specialized direction in a model built for coding agents and tool use is a clean, bilingual, four-times-stronger-than-normal representation of the programming concept type. I did not go looking for that. It is my favourite thing in this atlas.

The bilingual vocabulary was developed, not tacked on

I ran this census on an English-only corpus. So I expected Chinese tokens to be rare in the logit lens, and concentrated wherever the model does its lexical work.

Neither is true:

Layers Features promoting at least one CJK token Total
0-5 1,584 1,920
6-11 1,593 1,920
12-17 1,590 1,920
18-23 1,543 1,920

Around 82% of probed features promote a CJK token, and it is flat at every depth. On English input. The two languages are not living in separate neighbourhoods and Chinese is not parked in an early lexical stage. It is threaded through the entire stack.

For a model whose whole pitch is bilingual capability in a 1B footprint, this was loud. The integration is real rather than appended. It also means my English-only caveat has more impact on a complete representation than I first thought.

Layer 12 is a hub

The four strongest single heads in the model are all in layer 12:

Layer Component Head Best F-stat Top code dim
12 heads 8 216.1 45
12 heads 12 202.0 45
12 v 1 192.7 14
12 heads 9 189.8 45

Three of them lock onto the same code dimension. Layer 12 also carries the highest coactivation correlation in the model at 0.812. Several heads at one depth converging on the same category, with the most tightly correlated feature pairs in the network, looks like a routing waypoint.

One head is a dedicated copy engine

OV-circuit averages across all 384 head rows: spectral concentration 0.051, effective rank 54.4, QK concentration 0.119. Low concentration with rank in the fifties means these heads do weighted high-dimensional mixing rather than sparse token lookup.

Except one. Layer 17, head 9 has the top induction score in the model at 1.21, and it gets there with effective rank 27.7 and spectral concentration 0.115. Half the rank and twice the concentration of the average head. Everything else is a generalist; that one is a specialist.

The style subspace pivots at the exit

The corporate-versus-authentic angle holds between 145 and 179 degrees for layers 0 through 22, then drops to 98.3 at layer 23. The style subspace rotates sharply in the final layer. Sub-Zero classifier accuracy is a flat 0.922 through most of the stack and 0.984 at layer 23. I chose a stylistic feature over coding related due to the models size and bilingual nature.

Strongest style contrasts, for anyone who wants to go poking:

Layer Component Feature F-stat Delta
20 up 4063 1,744.9 1.428
17 mlp 220 1,687.8 -3.705
12 heads 1705 1,675.5 0.425
23 up 4471 1,656.8 1.793
14 gate 2179 1,641.8 0.538

High F-stat is not permission to edit. That is what the fence in the first section is for.

The MLP gate path is the cleanest I observed

Role labels, selective meaning a focused job:

Component Share selective
up 86.9%
gate 86.5%
mlp 78.9%
heads 76.5%
attn 71.5%
q 68.2%
v 65.4%
k 58.3%

The SwiGLU path is where interpretable single-purpose directions live. Key projections are the most entangled, which fits routing rather than content.

Logit lens agrees: gate leads on mean F-stat at 136.8, and the top projections sit around layers 13 to 19 rather than at the exit. Reading the promoted tokens by depth is its own method. Early layers push subword fragments and identifiers (refe, circu, IMATE, .removeEventListener), late layers push word-shaped tokens (scene, False, parse, isNaN). The lexical-to-semantic progression is right there in the table.

The stuff I deliberately skipped

I did not merge the SAE features into the atlas, so the sae_features table is empty. Cross-referencing "what a feature is" against "the basis it decomposes into," on the same model, is the obvious next move and it is not done here.

I also did not probe the <think> / no-think mode split. This model serves both from one checkpoint, and a contrast census between the two modes is a genuinely interesting experiment I have not run.

Caveats

English-only corpus on a bilingual model. 8,965 English prompts, and the CJK finding above says the Chinese half is threaded through every layer rather than parked somewhere I can ignore. So this atlas describes what English text pulls out of a bilingual model, and there is a whole second half unmapped. That is the biggest limitation here and it is not a footnote.

Layers 0 to 3 are weak SAE reconstructions and I shipped them anyway. 0.62 to 0.73 EV is what fifty features buys on that part of this model. That is a statement about the measurement, not about the model.

EV is a blunt ruler. Global-variance explained variance rewards streams with a few dominant directions, which is very likely part of what the residual RMS column is showing.

One corpus, one seed. FineWeb-Edu, seed 0. Feature labels are what this prompt set surfaces.

The SAE trainer is still under development. It found two of its own bugs on this model, which is the honest reason these weights exist in the shape they do. I am a bartender who taught himself this, so take the artifacts with curiosity and some caution.

Sub-Zero measures style axes, not every load-bearing direction. The 635 rows are domain-by-domain damage for corporate-versus-authentic candidate axes, not a census of everything that matters.

Find an interpretation on your own

import sqlite3, pandas as pd
conn = sqlite3.connect("atlas.sqlite")

# the one specific direction in the whole model
pd.read_sql_query("""
  SELECT layer_id, component, feature_idx, fstat, promoted_tokens
  FROM logit_lens WHERE layer_id=18 AND component='gate' AND feature_idx=3488
""", conn)

# where to be careful if you plan to edit
pd.read_sql_query("""
  SELECT layer_id, projection, das_axis, domain, damage
  FROM subzero_capability WHERE fence_passed=0 ORDER BY damage DESC LIMIT 20
""", conn)

SAEs need no trainer install, just torch against the state dict:

import torch
sd = torch.load("layer_04_s0/sae.pt", map_location="cpu")
pre = (acts - sd["b_dec"]) @ sd["W_enc.weight"].T + sd["W_enc.bias"]
feats = pre * (pre > sd["log_threshold"].exp())
print((feats > 0).sum(-1).float().mean())   # should sit right around 50

Bottom line

MiniCPM5-1B is a tidy model with an unusual front end. Four dense load-bearing layers, a 7x norm jump at layer 4, and then nineteen layers that decompose into fifty sparse features at 0.87 to 0.99 explained variance, sliding down in lockstep with a residual norm that climbs 53x from end to end. 87.4% of style axes are safe to remove, no attention projection failed a single capability test, and the one hyperspecific direction in 479,232 is a clean bilingual representation of type. Two languages threaded through every layer rather than partitioned.

Thank the OpenBMB team by downloading the model and using it yourself then go scan something.

Thanks for sharing this — really appreciate the work and the transparency behind it.

This looks like a valuable resource for the community, and I’m grateful you made it available.

Sign up or log in to comment