Lab · QKV swaps and rewiring

Rewiring trained transformers

Twelve open-weight models, changed at inference with no retraining. First, one pair of query, key and value vectors is exchanged inside attention, in all layers, the first half or the last half. Then, with attention left alone, each layer runs its MLP before its attention, or the layers run in a different order. We read what the models generate, their next-token distributions and, for the swaps, the residual stream layer by layer.

What the residual stream shows

Every model, every swap

Measure

Each model reads the same 8 wikitext articles (first 512 tokens each), teacher-forced, so every condition sees identical input and every number compares the same positions. The first position is left out. Darker cells are larger values, whatever the measure.

One model, layer by layer

Model

Measure

Layer 0 is the token embedding; layer n is the residual stream after decoder layer n. The shaded band marks the layers whose attention is swapped. Hover a chart to read values.

What the swapped model predicts most often

Next-token predictions, token by token

Swapped

The question

What does a trained transformer do if its attention uses the wrong vectors? We exchanged pairs of query, key and value vectors inside attention, with no retraining, and looked at the text, the next-token distribution and the residual stream. Then two follow-ups with attention left intact: what if each layer runs its MLP before its attention, and what if the layers run in a different order? Everything below ran on 2026-11-10 on the AICR cluster.

What we ran, in order

Swapped Q and K in every layer of Qwen2.5-0.5B-Instruct on a CPU node and generated text for 3 prompts under 4 conditions: baseline, all 24 layers, the first 12, and the last 12.
Picked every model in the cluster's model cache whose weights were still complete, checked each shard byte for byte, and fetched the one missing file (Qwen2.5-72B's config.json).
Wrote the general script, adding Q⇄V and K⇄V and a no-swap check against each model's own eager attention. Ran a CPU smoke test, then queued 12 models × 10 conditions × 4 prompts on RTX PRO 6000 GPUs.
Two models failed to load offline (cache refs missing). Restored the refs and reran them.
Published this page with the 480 generations.
Found that DeepSeek-Prover's tokenizer drops spaces under transformers 5, so its first runs had read garbled input. Checked all 12 tokenizers, fixed the loader, and reran Prover.
Built the residual and distribution probe. Extracted 8 wikitext-103 validation articles, ran a CPU smoke test on Qwen2.5-0.5B, then ran all 12 models.
Gemma-4's baseline scored far below the others on wikitext, so a separate job compared its stock attention, eager attention and our attention code. They agree, so the low score is the model's own.
Ran the probe on the two 70B-class models, two GPUs each.
MLP before attention: wrote the script, ran a CPU smoke test on Qwen2.5-0.5B, then queued all 12 models (4 conditions × 4 prompts, plus a wikitext score).
Layer order: the same steps, with 4 orderings.
Both arrays stalled behind low queue priority, so their pending tasks were allowed onto any RTX or B200 partition. Qwen3.5-9B's MLP-first task ran on a B200; everything else ran on RTX PRO 6000s.

How the swap works

The model's attention function is replaced with our own (registered through transformers' AttentionInterface). It receives the query, key and value tensors after projection, QK-norm and RoPE, exchanges one pair in the chosen layers, then computes causal softmax attention in fp32.

qk : score(i, j) = k_i · q_j          (instead of q_i · k_j)
qv : queries = v (no RoPE),  values = q (rotated)
kv : keys    = v (no RoPE),  values = k (rotated)

Q⇄K at this point is exactly a swap of W_Q and W_K, because RoPE gives (R_ia)·(R_jb) = a·R_j−ib, which depends only on j − i. Q⇄V and K⇄V are not weight swaps: the swapped-in values carry the position rotation and the swapped-in queries or keys have none.
Grouped-query attention: key and value heads are repeated up to the query head count first (repeat_interleave, the same order as transformers' repeat_kv), so any pair can swap.
No KV cache: the swapped score needs the queries of earlier tokens, which a KV cache does not keep. Every step recomputes the whole sequence.
Kept from each model: its attention scaling, Gemma-2's attention-logit soft-capping and the sliding windows (never reached at these lengths). The causal mask is built inside the function, with batch size 1.
Hybrid Qwen3.5 and Qwen3.8 only call softmax attention in every fourth layer. Their linear-attention layers are untouched, so "all layers" means all 8 or 16 attention layers.
Check: with no swap, our function matches eager attention bit for bit in fp32 (logit diff 0.0), and the eager reference makes zero calls into our function. In bf16 on GPU the gap is rounding; per-model numbers are in the model table.

Generation runs

Conditions: baseline, then Q⇄K, Q⇄V and K⇄V, each in all layers, layers 0 to ⌊L/2⌋−1, or ⌊L/2⌋ to L−1. That makes 10.
Prompts: "Write a short poem about the ocean.", "Explain what a black hole is in two sentences.", "What is the capital of France? Answer in one word." (chat template, enable_thinking=False where accepted) and "Once upon a time, in a small village," (raw text).
Greedy decoding with our own loop, up to 80 new tokens, stopping on any EOS id from the generation config or tokenizer.
bf16 weights, device_map="auto". One RTX PRO 6000 (96 GB) per model, two for the 70B-class models.

Residual and distribution probe

Text: the first 8 articles longer than 6,000 characters in wikitext-103-raw validation, detokenized (@-@ joins, heading markers and spacing before punctuation removed). Each is cut to its first 512 tokens with BOS. The same tokens go to every condition.
Hooks record the residual stream before layer 0 (the embeddings) and after every decoder layer, before the final norm.
Per layer: cosine and linear CKA to the baseline at the same positions, mean norm, effective rank (exp of the entropy of the normalized eigenvalues of the centered Gram matrix), mean token-to-token cosine, and a logit lens (final norm, unembedding and Gemma's final soft-cap) scored on the real next token. The lens on the last layer reproduces the model's logits (max diff 0.0, or 0.12 for the soft-capped Gemmas in bf16).
Next-token distribution: loss and top-1 on the real next token, entropy, KL both ways from the baseline, Jensen–Shannon, top-10 overlap with the baseline, the probability given to the baseline's top token and to repeating the current token, context sensitivity (mean JS between each position's distribution and the passage average) and the most frequent argmax tokens.
Linear recovery: ridge regression (dual form) from the swapped last-layer residual to the baseline's. Fitted on articles 1–5, λ picked on article 6 from {10⁻³, 10⁻², 10⁻¹, 1} × the mean squared norm, refitted on 1–6 and tested on 7–8, then pushed through final norm and unembedding. The floor fits the same map from the token embeddings, which carry no context.

MLP before attention

original  : h → h′ = h + A(h)  → h′ + M(h′)
MLP first : h → h′ = h + M(h)  → h′ + A(h′)

A and M are the whole attention and MLP branches with their own norms: pre-norm in Llama, Qwen and DeepSeek, pre- and post-norms in Gemma, post-norm in OLMo.
Each chosen layer's own forward is called twice. Pass 1 zeros the attention output with a forward hook, which gives h + M(h). Pass 2 starts from that and zeros the MLP output, which adds A. Masks, RoPE and the residual adds run as the model wrote them.
Gemma-4 multiplies each layer's output by a learned scalar. Pass 1's scalar is divided back out, so it is applied once.
Hybrid Qwen3.5 and Qwen3.8: in 3 of every 4 layers the token mixer is linear attention. It moves after the MLP the same way.
Conditions: baseline, all layers, layers 0 to ⌊L/2⌋−1, and ⌊L/2⌋ to L−1.
Check: one forward with every layer reordered must zero exactly 2L branch outputs, one attention and one MLP per layer.
Stock SDPA attention and no KV cache. With a cache, pass 1 would store keys and values for the wrong input.

Layer order

baseline         0, 1, …, L−1
reversed         L−1, …, 1, 0
ends swapped     L−1, 1, 2, …, L−2, 0
middle reversed  0 … q−1, then L−q−1 down to q, then L−q … L−1     (q = ⌊L/4⌋)

The decoder's layer list is rebuilt in the new order. Each layer keeps its weights and norms.
The model's list of layer types moves with the layers. The decoder uses it to pick each layer's mask (full or sliding window, or linear attention in the hybrids) and, in Gemma-4, its RoPE and head size. Gemma-4 shares no keys or values between layers, so nothing else is tied to position.
Check: putting the identity order back reproduces the untouched model exactly (max logit difference 0.0).
Stock SDPA attention, no KV cache, same prompts and decoding as the other runs.

Models

Problems found and how we fixed them

Purged and partial model cache. About 30 models in the shared model cache had been purged by its 30-day rule. Qwen2.5-72B's snapshot had every weight shard but no config.json (fetched, 663 bytes). Every shard of the 12 models was checked by comparing its safetensors header against the file size. Leftover .incomplete blobs in Qwen3.8-27B and Llama-3.3-70B are unused partial downloads.
Offline load failure. Qwen3.5-9B and R1-Distill had weights but no refs/main, so HF_HUB_OFFLINE=1 could not resolve them. We wrote refs/main pointing at each model's only snapshot and reran them.
DeepSeek-Prover tokenizer. transformers 5.12.1 loads it as LlamaTokenizer and drops spaces when encoding: "Once upon a time, in a small village," became Once|upon|at|ime|,|in|asm|all|village|,. The other 11 tokenizers match tokenizer.json exactly. Both scripts now compare the two and load PreTrainedTokenizerFast when they disagree. Prover's generations were rerun, and the probe used the fix from the start.
Replacement characters in model output. Some swapped models emit tokens that are partial UTF-8 bytes and decode to U+FFFD. They are real output, so they are escaped in the page data rather than removed.
Dataset download. hf download of wikitext stalled, so the single validation parquet (657,209 bytes) was fetched with curl.
Gemma-4's weak baseline. On wikitext its unswapped top-1 is 0.29 against 0.49–0.73 for the others. A separate job scored the same passages with stock SDPA, eager attention and our attention function with no swap: top-1 0.340, 0.342 and 0.337 at 512 tokens (0.17–0.19 at 128). Our code is faithful; the low score is the model. Its chat replies are normal, but as a raw-text continuer it repeats itself ("in a small village, in a small village, …"), and teacher-forced wikitext measures exactly that mode. Read its numbers against its own baseline.

Environment

Limits of what this shows

Q⇄V and K⇄V are defined at the attention op, after RoPE. Swapping the projection weights instead would give different results for those two.
Generations are greedy, single runs on 4 prompts. Read them as examples, not measurements.
The probe uses one domain (Wikipedia) and 4,080 scored positions per condition. Recovery tests a linear map only; information a nonlinear readout could use would not show up.
Weights are bf16 and attention runs in fp32, so the swapped models are compared against a baseline computed the same way, not against stock SDPA.
MLP-first and layer-order scores come from 2 articles (1,020 positions) and 4 greedy prompts per condition, with no residual probe.
Qwen3.5-9B's two baselines differ slightly (top-1 0.648 and 0.651, and two of its four greedy replies part ways after 27–30 tokens), most likely because its two tasks ran on different GPU types. Every other model's two baselines match token for token. Each tab shows the baseline from its own run.