What we learned
How the per-layer embedding (PLE) works
Half of E2B’s text weights are one giant lookup table. Each token gets a 35 × 256 block of numbers, and every layer gates its own 256-number slice into the residual stream. Hover the boxes; pick a layer in the profile below.
Click a column in the profile below to change it. The boxes on the left read their numbers for this layer.
- Rows → mean keeps the branch but erases token identity in that layer’s slice.
- Branch off removes the whole PLE path in that layer (its knob set to 0).
- KL = how far next-token predictions move from the original model (average over 12 corpora).
Per-layer profile
Hover for values; click a column to select that layer. Compare the second and fourth rows: layer 14’s PLE write is the biggest (11.6× the stream), yet switching it off barely moves predictions; layer 8’s modest write is one of the two whose loss moves them most (with layer 33’s). Size isn’t importance. (Layer 8’s write is 97% one fixed vector: replace it with its average instead of zero and predictions barely move. Puzzle 4 solves why: L14’s write is diluted to 0.26× by L14’s 47.2-RMS pre-PLE state, 59% collinear with it, and runs after shared K/V are already saved.)
Which languages lean on which layer’s PLE
One layer at a time. Rows alone are cheap to remove (KL ≤ 0.35; layer 13 matters most; low-resource languages lean on layers 0–3). The whole branch matters a lot in layers 8, 13, 23 and 33, but in 8, 23 and 33 only as a constant offset: replace the branch’s write with its average instead of switching it off and every skill survives. Layer 13’s token-specific part does matter (KL 0.40 even when averaged).
Removing token identity everywhere
What the table stores
Compression: table size vs damage
int4 is free: 4.70 → 1.32 GB with KL 0.013 and no change in skills. int4 plus the 128k most common tokens gives 0.66 GB for −0.37 pt. Low rank fails: the table is genuinely high-rank, as ID-like codes would be.
The instruction-strength dial
Blend every text-model tensor of the base model (PT) and the instruct model (IT): W = (1 − α)·PT + α·IT. Chat behaviour switches on abruptly, around α ≈ 0.4 (Puzzle 7 explains why). At α = 0.75 the model matches IT on every chat probe while predicting raw text better. Past IT it degrades, and at 1.5 it breaks. Safety, multi-turn and tool use were not measured, so treat blends as a research result, not a deployable model.
Grafts: swap parts between the two models
Same lineage, but the parts were trained together. Swap a group of tensors and see what survives.
How far each block moved
Is it a low-rank change?
Puzzles: results that didn’t fit
Each puzzle below is a measured result that doesn’t fit the simple story “each block does one job”, with a guess or two and an experiment that would settle it. Follow-up experiments (including a pre-registered 3-pillar validation suite across E2B & E4B plus a constructive Freeze-PLE SFT training study) have solved seven of the eight puzzles (Puzzles 1, 2, 3, 4, 5, 7 and 8) and turned Puzzle 5 into a 48.8%-lower-VRAM fine-tuning recipe. All of them are good starting points for the next round of edits.
Cross-scale test: google/gemma-4-E4B-it (42 layers, 5:1 local/global)
Do the architectural laws found on E2B (35 layers, 4:1 local/global, 1,536 hidden) hold on a deeper, wider model with a different attention ratio? We ran the full 126-block single-block zero-ablation atlas, KV-sharing boundary surgical edits, and PLE barcode/compression suite on Gemma 4 E4B IT (42 layers, 2,560 hidden, global attention every 6th layer at ◆ 5, 11, 17, 23, 29, 35, 41, with layers 24–41 reusing K/V from local layer 22 and global layer 23). Every law held—and int4 PLE compression got even tighter (KL 0.007).
Side-by-side: E2B (35L, 4:1) vs E4B (42L, 5:1)
Pre-registered predictions P6–P9 tested out-of-the-box on E4B before any forward pass.
E4B’s 5.64 GB PLE table: barcode R² & int4 / int2
Held-out linear R² from E4B’s 2,560-dim main embedding stays near zero across all 42 slices (median −0.15, max 0.19), while int4 (group 64) shrinks 5.64 → 1.59 GB (saving 4.05 GB) at KL 0.0071 and 0.9995 retention—even int2 (0.88 GB) keeps 0.993.
Constructive proof: Freeze-int4-PLE supervised fine-tuning from raw PT
Puzzles 2 & 5 showed that embed_tokens_per_layer (2.35B parameters, 4.70 GB in bf16—51% of the entire text model) acts as a static token-identity dictionary while instruction-following lives in the 2.28B backbone. To prove this constructively, we started from raw pretrained google/gemma-4-E2B (PT, where chat = 0.00) and ran controlled 400-step SFT across identical seeds, microbatches (3,200 UltraChat + Alpaca + GSM8K turns), and mixed-precision AdamW—replacing PLE before training with a frozen, packed uint8 group-64 Int4FrozenPLE module (1.32 GB, fp16 scales).
1-GPU SFT instead of 2 GPUs: peak VRAM & speed
Standard full SFT allocates gradients + fp32/bf16 AdamW states across all 4.63B parameters (64.81 GB), OOMing a 40 GB A100 and forcing 2-GPU sharding. Freezing & int4-packing PLE upfront drops peak VRAM to 33.21 GB on a single GPU at 2.25× throughput per GPU. Hover any segment.
Loss parity & 13-skill profile
At Epoch 1 (step 200), SFT-Frozen-int4-PLE matches SFT-Unfrozen-bf16 to 0.0001 nats on 256 held-out validation turns (1.1073 vs 1.1074) and reaches 1.00 on strict <turn|> closure.
One 1.32 GB dictionary, many fine-tuned personas
Because a frozen PLE table never drifts (exact KL = 0.0000; vs 36.2% drift in official E2B-it), multiple full-SFT backbones can point to the exact same 1.32 GB GPU buffer. We trained a second zero-Markdown persona (SFT-Voice-Concise) on the same buffer and co-loaded PT + SFT-General + SFT-Voice live on one GPU in 15.00 GB instead of 27.78 GB.
Side-by-side outputs sharing the same 1.32 GB PT PLE buffer
Select a prompt to compare raw PT, SFT-Unfrozen-bf16, and two distinct backbones (SFT-General vs zero-Markdown SFT-Voice-Concise) reading the exact same frozen Int4FrozenPLE table.
The edits
Most experiments are direct weight edits on google/gemma-4-E2B-it (and google/gemma-4-E2B for the base model), bf16, Hugging Face transformers, one GPU on GCP. The Puzzle 1 and 2 follow-ups also edit activations: a block’s output is replaced during the forward pass. No model weights were trained; the only fitting is a small least-squares map for the R² and linear-prediction tests.
- Switch off a block (attention, MLP or PLE branch): set its post-norm weight, the “volume knob”, to 0. Gemma’s RMSNorm multiplies by this weight directly, so the block’s write into the residual stream disappears exactly. This is “Zero” on tab 1.
- Replace a block’s write with its average (“Average” on tab 1): at every position except the first (BOS) token, the write becomes its average over 49k tokens of held-out text from the same 12 corpora, separate from the text used for scoring. This keeps the part of the write that’s the same for every token and removes only the token-specific part.
- Skip a layer: switch off all three of its blocks. Its
layer_scalarstill applies, so for layers with small scalars the skip also shrinks the old state relative to later writes. Puzzle 1 also tried an identity skip (output = input), the input rescaled to the normal output size, the input × c, and one fixed vector at every position after the first. - Path edits (Puzzle 1): remove attention’s write from only the layer’s output, or only from what the rest of the layer reads (its MLP and PLE branch). Puzzle 2 also hid one block’s output, minus its average, from the input of layer 14’s key or value projection only, or kept it only there.
- Flatten keys: set
k_norm’s weight to 0, so every key is 0 and attention spreads evenly over everything visible. Layers 15–34 reuse layer 13’s and 14’s keys, so this edit reaches them too. Switching off layer 13’s or 14’s attention, by contrast, leaves the shared keys and values intact. - Reset a scalar: set
layer_scalarto 1. - PLE table edits: replace the rows looked up for each token (average row, zero, linear prediction, low rank, quantised, pruned vocabulary). Mathematically the same as editing the table.
- Grafts and blends: swap whole tensors between the base and instruct checkpoints, or set W = (1 − α)·WPT + α·WIT for every text-model tensor. Puzzle 6 also swapped one layer’s norms or linear weights at a time.
- Partial changes (Puzzle 8): PT plus the top-k singular directions of each linear layer’s change ΔW, a random slice of it, or a fraction of the whole of it; or IT minus its top directions. Norms and tables stay IT.
- Turn-token rows (Puzzle 7): replace the embedding rows of the chat template’s two turn tokens (
<|turn>and<turn|>) with the other model’s, alone or inside a blend. The embedding is tied to the output layer, so this also changes how closing a turn is scored. - Bypass, decoding & E4B validation (9 pre-registered predictions): masked diagonal
t → tself-attention and averaged upper per-layer inputs in layers 15–34 underL14.const_only; ran 240 KV-cached greedy rollouts separating prompt prefill from 320-token decoding; traced sub-layer normalised-input shifts across all 35 layers (solving Puzzle 4, Spearman ρ = 0.912); and replicated the 126-block atlas, K/V-donor sub-layer copy test (p22mlp.notK23.mean), and int4/int2 PLE compression on 42-layergoogle/gemma-4-E4B-it. - Constructive Freeze-
int4-PLE SFT (5 pre-registered predictions): starting from rawgoogle/gemma-4-E2B(PT), trained 400-step SFT runs under identical seeds, microbatches (3,200 UltraChat + Alpaca + GSM8K turns), and mixed-precision AdamW comparing unfrozen bf16 PLE (64.81 GB, 2 GPUs) vs frozen bf16 PLE (36.58 GB, 1 GPU) vs packeduint8group-64Int4FrozenPLEwith fp16 scales (33.21 GB, 1 GPU), plus a 200-step zero-MarkdownVoice-Concisepersona co-loaded alongsidePTandSFT-Generalon one shared 1.32 GB PLE buffer (15.00 GB total VRAM).
The measurements
- Retained: the share of items the unedited model got right that are still right after the edit. For induction and raw text: the top-1 accuracy ratio. 1 = unchanged, 0 = destroyed. “Mean retention” averages the 13 skills, capping each at 1.
- KL: KL(original ‖ edited) on next-token predictions, averaged over 12 corpora (11 Wikipedia languages + Python): 16k tokens per corpus for whole-table edits and compression, 4k for the per-layer sweeps and Puzzle 1, 6k for the pruning curves. The text is held out from our token-frequency counts, but likely not from pretraining.
- Δ top-1: change in next-token accuracy on that text, in percentage points.
- Volume knob: RMS of a block’s post-norm weight vector. A pure weight statistic.
- Write size: RMS of a block’s write ÷ RMS of the stream entering the layer, on English text. Old state carried over: RMS of the stream entering a layer ÷ RMS of that stream plus the layer’s three writes.
- Constant part: the share of a write’s energy in its average vector (‖average‖² ÷ average ‖write‖²), on the same held-out text as the averages. Near 100% means the write is effectively a fixed bias.
- Copying score (Puzzle 2): on 32 prompts that repeat a list of 48 random words, the share of a head’s attention that lands on the word right after the earlier copy of the current word. The recall copy score is attention to the digit being copied. Measured with eager attention; the map shows each layer’s best head.
- Changed in tuning: ‖WIT − WPT‖ / ‖WPT‖. It counts rescaling as change, even where a following norm cancels it.
- Lead (Puzzle 7): a token’s logit minus the best other token’s at one position, with the earlier tokens fed in. Above 0, greedy decoding would pick it.
The 13 skills
Scored from log-probabilities on fixed prompts; nothing is sampled. “Base” is the unedited model’s accuracy.
Caveats
- Small probes: 8–116 items each (raw text: 8–20 passages), mostly near 100% at baseline.
- Switching a block off shows what depends on it, not where a skill is computed. And effects interact: with zeroing, removing more could even hurt less (Puzzle 1).
- Zeroing a write also deletes its constant part, which overstates many blocks: tab 1’s Average toggle shows how much. Averaging isn’t neutral either: the averages come from Wikipedia and code, while most probes are short prompts (few-shot, chat or synthetic).
- Quantisation is simulated (round-to-nearest, per-group scales), not real low-bit kernels.
- Copying scores measure attention to one exact position; a head could copy in ways they miss.
- The mixed-precision recipes picked their “sensitive” layers on the same text they were scored on.
- The instruct model is very overconfident on raw text, so the atlas uses accuracy, not NLL, for text.
- Generation checks use 12 prompts per model, so one prompt is 8 points of stop rate. A model that stops by emitting <eos> at once counts as stopping.
- Blended models were not checked for safety, multi-turn or tool use.
- Puzzle 7’s prompts are short, templated questions; varied real chats may switch on more gradually.
Background
Inspired by weight-space programming: reading and editing a network’s weights directly, as in arXiv 2611.05341 (an LLM auditing and repairing small clinical networks) and “Sparks of Weight-Space Programming” (a small LLM’s looping bug fixed by editing weights).