Your Transformer Can Hold Two Thoughts at Once: Evidence of Linear Superposition in LLMs
Pavel Tikhonov, Anton Korznikov, Matvey Mikhalchuk, Nikita Dragunov, Temurbek Rahmatullaev, Polina Druzhinina, Anton Razzhigaev, Ivan Oseledets, Elena Tutubalina
cs.CL, cs.AI
2026-09-24
Averaging two text embeddings, a stock LLM still approximates both next-token distributions. Light distillation drops Pythia-2.8B mixed KL from 1.86 to 0.27.
Decoder-only Transformers are full of nonlinear parts: softmax attention, MLP activations. The usual picture is one forward pass, one semantic stream. Two unrelated texts mean two forwards, or an architecture change so the streams do not smash into each other.
Layer-to-layer maps in the residual stream are often close to affine. This paper asks the end-to-end version of that fact. Average two token embedding sequences position-wise. Does the model emit a superposition of the two independent next-token distributions? They call that the Superposition Linearity Hypothesis.
The mix is elementary. For equal-length sequences x and y, the input at position t is (E(xt)+E(yt))/2. Causal attention runs as usual. No mask change, no mux layer.
They score the hypothesis three ways.
Attention gets two controls. Donor patching swaps in post-softmax weights from an unrelated text C at every layer and head, leaving A's Q/K/V intact. Permutation patching shuffles A's own causal-prefix weights, keeping row sums and the weight multiset, destroying position structure. Tokens are split into predictable positions (about 65%, function words and punctuation) and content positions (about 35%).
To amplify linearity they self-distill: a frozen teacher runs the two streams separately, the student sees the mixed embeddings, and the loss is KL(0.5(Pteacher(A)+Pteacher(B)) || Pstudent(z)). FineWeb subset, about 200k steps, claimed at under 0.025% of the original pretraining set.
Decoding is a separate problem. If logits are approximately averaged, probabilities follow a geometric mean: P(t) ∝ √(PA(t) PB(t)). A token that is high in A and near-zero in B gets crushed. Sampling the mixed distribution hops between the two stories. Joint Contrastive decoding feeds the mix to the large model and the two originals to a small guide, then forms ℓlarge(z)+α ℓsmall(A)−β ℓsmall(B).
Off-the-shelf pretrained models already beat a frequency baseline by a wide margin. On Pythia-2.8B, Llama-3.2-3B, and Qwen2.5-3B, a single-stream predicted token lands in the mixed top-10 about 30-40% of the time, top-50 50-60%, top-100 60-65%. Vocabulary size is at least 50k. A control that looks up stream B's true token in stream A's logits gets top-10 only 2.63% of the time.
| Model | L=32 KL | RKL | KL after fine-tune |
| Pythia-2.8B | 1.86 | 0.42 | 0.27 |
| Llama-3.1-8B | 2.16 | 0.37 | not reported |
| Pythia-160M | 0.91 | 0.35 | not reported |
L=512 is slightly cleaner, not worse. Mixing N=3 streams raises RKL by 0.04 to 0.09; the qualitative picture holds.
Checkpoints point to an architectural prior. Across Pythia sizes, hidden-state superposition error is smallest at initialization and rises monotonically through pretraining. Layer-wise linearity is U-shaped: high in the first layers, down to about 0.65 in the middle, back above 0.95 in the last third. That terminal linearity is a plausible reason the superposed signal still reaches the logits.
The attention controls split the aggregate recovery numbers into pieces. On Qwen2.5-3B, embedding mixing keeps predictable positions at median rank 6 (top-1 24.8%) and content positions at median rank 284 (top-1 4.2%). Donor patching is gentler on predictable tokens (median 3) and still poor on content (111). Permutation collapses both: 3079 and 19246. On LAMBADA, which asks for a content word, mixing gets 2.25% argmax accuracy and median target rank 339; donor patching gets 0.5% and 2350. The mixed forward keeps more case-specific signal on hard positions than a wholesale swap of natural attention shape.
Fine-tuning repairs content positions. After the Qwen2.5-3B run, predictable median rank barely moves (8), while content median rank falls from 284 to 5 and top-1 hits 22.8%. On Pythia-2.8B, mixed KL drops from 1.86 to 0.27 and RKL from 0.42 to 0.06; the chance the true token sits in the top-5 rises from about 30% to above 60%.
Decoding still lags. Joint Contrastive lifts mean LAMBADA on a superposed forward from 0.06-0.18 (pretrained mix) to 0.11-0.43; Llama-3.2-3B reaches 0.430 against a 0.540 single-stream small-model baseline. Jaccard overlap falls as well. Single-stream quality pays: Pythia-2.8B LAMBADA 0.544 to 0.357, Qwen2.5-3B 0.602 to 0.460. Plain distillation is harsher, sending Qwen PPL from 12.44 to 65.52; Joint Contrastive stays at 19.31, next to the small model's 18.96. Throughput for a two-head Separate decoder matches batch-2 independent inference and is about 2x sequential; Pythia-2.8B Separate is 109.6 tok/s versus 56.0 sequential.
This is a different claim from DataMUX or RevMUX, which bolt mux/demux structure onto the network. A stock pretrained LLM, given averaged embeddings, already carries two streams in one residual stream. Pretraining weakens that property. A light distillation brings it back.
The inference reading is concrete: one forward, one KV cache, a theoretical 2x in throughput and half the cache per active stream. The Separate decoder in the appendix does touch batch-2 speed. The geometric-mean obstruction remains. Joint Contrastive still needs a small guide, and 0.43 on Llama does not match 0.54 from the small model alone. This is an existence proof that the residual stream can hold two thoughts, not a drop-in parallel decoder.
The authors bound the study themselves: most analytical runs use L≤128, distributional tests go to 512, the corpora are mostly monolingual. No multimodal models, and no mix of text tokens with image patches.
The larger gap is between fitting a mixed distribution and actually emitting two coherent continuations. Fine-tuning drives KL down; naive sampling still cross-talks. Content-position survival on LAMBADA was already poor in the base model; restoring it costs single-stream skill. The distillation set is under 0.025% of pretraining tokens, yet one Pythia-2.8B run still took 114 hours on 2xA100. N=3 is measured only as distributional distance, with no decoder.