Your Transformer Can Hold Two Thoughts at Once: Evidence of Linear Superposition in LLMs

AI summary

Overview

Research area: Natural Language Processing — mechanistic interpretability and efficient inference for decoder-only Transformer language models.

Technical level: Advanced. The paper assumes familiarity with residual streams, self-attention, softmax, KL/Jensen–Shannon divergences, Wasserstein distance, and layer-wise linearity probing.

Scope: The paper tests whether standard pretrained LLMs process a token-wise average of two input embeddings as an approximate superposition of the two independent next-token distributions, tracks how that property changes during pretraining, restores it with lightweight fine-tuning, and builds a proof-of-concept decoder that separates the two streams.

What This Paper Is About

Decoder-only Transformers are built from nonlinear pieces (self-attention, MLP activations), so the default assumption is that two unrelated text streams mixed into one forward pass should destructively interfere and produce nonsense. The authors ask whether the opposite is true: when embeddings from two documents are averaged element-wise, does the model output approximate the arithmetic mean of the two models' separate next-token distributions? They call this the Superposition Linearity Hypothesis, and their goal is to establish whether it holds, where it comes from, whether it can be improved, and whether the two mixed streams can be decoded apart.

Key Contributions

1. Demonstration of intrinsic superposition in off-the-shelf models. Standard pretrained LLMs (Pythia, Qwen, Llama, OLMo, Gemma families) retain substantial probability mass on the tokens each independent stream would have predicted, despite no architectural modification.
2. Evidence that superposition is architectural, not learned. Hidden-state additivity error is smallest at the earliest pretraining checkpoints and grows monotonically as training proceeds, implying pretraining degrades rather than creates the property.
3. Recovery via lightweight fine-tuning. A self-distillation objective using less than 0.025% of the original pretraining dataset size substantially restores linearity, with the largest gains on hard "content" tokens rather than predictable ones.
4. A proof-of-concept disentangling decoder. A Joint Contrastive decoding variant using a small auxiliary guide model recovers two coherent continuations from a single mixed forward pass, though it does not close the gap to single-stream inference.

Main Findings

• Rank survival in unmodified models. Under embedding averaging, the true next token lands in the top-10 in approximately 30–40% of cases, in the top-50 in 50–60% of cases, and in the top-100 in 60–65% of cases, with vocabulary sizes |V| ≥ 50,000.
• Distributional approximation beats chance. Across Pythia-160M, Pythia-410M, Pythia-2.8B, and Llama-3.1-8B on FineWeb at context lengths L = 32 and L = 512, the Superposition Approximation Ratio R_D is below 1 for KL, Jensen–Shannon, and Wasserstein distances. For Pythia-2.8B at L = 32, KL is 1.86 with ratio 0.42; after fine-tuning the mean KL drops to 0.27 and R_KL to 0.06.
• A frequency-only baseline cannot explain it. Against a control that preserves token-frequency statistics but breaks the contextual link, cross-stream ground-truth tokens appear in the top-3 only 1.12% of the time, top-10 in 2.63%, and top-100 in 10.41%, with same-position token overlap of 0.2%.
• Superposition degrades during pretraining. The layer-averaged hidden-state superposition error is lowest at early Pythia checkpoints and rises monotonically, and layer-wise analysis shows a U-shaped profile where deep layers (l ≳ 2L/3) remain near-linear.
• It persists at three streams. Mixing three streams and measuring all three ground-truth ranks raises R_KL by only +0.04 to +0.09 across models.
• Attention shape matters more than attention content. On Qwen2.5-3B, permutation patching (which destroys attention structure) collapses median rank from 8 to 8,148 and drops top-10 agreement from 53% to 10.1%, raising E[KL] from 3.54 to 9.13 and R_KL from 0.27 to 0.68.
• A predictable/content split explains the aggregate numbers. Approximately 65% of positions are predictable and about 35% are content. Under base-model embedding mixing, predictable positions have median rank 6, while content positions sit at median rank 284. After fine-tuning, overall median rank improves to 6 and content positions drop to 5 with 22.8% exact top-1 agreement. Predictable positions barely change (median rank 8).
• Embedding mixing carries more than donor patching on hard tokens. On 200 LAMBADA prompts left-truncated to 128 tokens, embedding mixing reaches 2.25% argmax accuracy at median target rank 339, versus 0.5% at median rank 2,350 for donor patching, against a vanilla 73% accuracy — a 4.5x accuracy gap and 7x rank gap. LAMBADA baselines are 0.0% for LSTM and 0.1% for N-Gram.
• Restoration is not free. Fine-tuning reduces single-stream quality: Pythia-2.8B LAMBADA falls from 0.544 to 0.357, and Qwen2.5-3B from 0.602 to 0.460.
• Decoding is strictly harder than distribution matching. Because mixed probabilities scale with the geometric mean of the independent distributions, tokens strongly favored by one stream but unlikely in the other are penalized. Joint Contrastive decoding lifts LAMBADA mean accuracy on the mixed pass from the 0.06–0.18 pretrained range to 0.11–0.43 while lowering inter-stream Jaccard overlap (e.g., Qwen2.5-3B from 0.126 to 0.061; Llama-3.2-3B from 0.094 to 0.067). Llama-3.2-3B reaches 0.430 against a 0.540 single-stream small-model baseline.
• Contextual stability. Total Variation Distance between the mixed output and the target mixture is slightly higher for the first ~20 tokens and then stabilizes across the context window.
• Confidence predicts survival. Tokens predicted with probability above 0.5 almost always survive with median rank around 3 in the mixed distribution.

Methodology in Plain English

The core setup is simple. Take two unrelated texts A and B, tokenize them, look up each token's embedding vector, and average the two vectors position by position. Feed this averaged sequence through a frozen pretrained model as if it were normal text. Then ask: how highly does the model rank the token it would have predicted for A alone, and for B alone?

To measure more than just ranks, the authors compare the model's output distribution on the averaged input against the arithmetic mean of the two separate output distributions, using KL divergence, Jensen–Shannon divergence (with temperature smoothing τ = 1.5), and Wasserstein distance computed over the top-256 tokens with cosine distance between token embeddings as the ground metric. They normalize all of these into a Superposition Approximation Ratio so different model families and context lengths are comparable.

To separate architecture from training, they run the same test across the Pythia pretraining checkpoint trajectory and measure hidden-state additivity: after mean-centering and L2-normalizing, how far is the mixed hidden state from the normalized sum of the two single-stream hidden states?

To probe the role of attention, they run three variants of a single-stream forward pass on Qwen2.5-3B over FineWeb-Edu at T = 128: a vanilla pass, a donor-patched pass where attention weights from an unrelated text C replace A's weights at every layer and head (while Q/K/V, RoPE, and value paths stay A's), and a permutation-patched pass where A's own attention rows are randomly shuffled within the causal prefix.

To improve linearity, they do self-distillation: the teacher is a frozen copy of the model producing the average of its two independent predictions, the student sees the averaged embedding, and the loss is the KL divergence from teacher target to student output. This used a subset of FineWeb for approximately 200k steps.

Finally, to decode, they add a small guide model whose logits are combined with the mixed logits — adding the guide's logits for one stream and subtracting the guide's logits for the other — with learnable scalars α and β initialized to 1.

Why This Matters

Research impact. The paper reframes superposition from an engineered capability into an intrinsic property of standard pretrained Transformers. Prior multiplexing work (DataMUX, MIMONets, RevMUX) adds dedicated layers, binding keys, or isometry regularization; this work argues that simple embedding averaging in off-the-shelf LLMs already preserves signal, and that lightweight fine-tuning restores a property pretraining has degraded rather than creating a new one. It also connects residual-stream linearity work to end-to-end input–output behavior.

Real-world applications:

• Parallel inference: Generating two continuations from a single forward pass, which the authors describe as theoretically offering a 2x increase in inference throughput.
• Memory reduction: Because multiple streams compress into one vector representation, the KV-cache footprint per active stream is effectively halved.
• Retrieval-augmented generation: Processing multiple document paths within a single forward pass rather than sequential passes.
• Batch serving: Serving unrelated concurrent requests through a shared residual stream instead of separate batches.

Industry relevance. Serving cost for LLMs is dominated by KV-cache memory and forward-pass count. A method that halves per-stream cache footprint and packs two streams into one pass directly targets both bottlenecks, though the paper is explicit that the current decoder is a proof of concept and that the fine-tuning objective trades away single-stream accuracy.

Future Directions

• Closing the decoding gap. The residual gap between Joint Contrastive accuracy (0.43 on Llama-3.2-3B) and the single-stream small-model baseline (0.540) is attributed to the geometric-mean obstruction rather than to insufficient fine-tuning; overcoming it fully remains, in the authors' words, an open problem.
• Longer contexts and multilingual settings. Analytical experiments used L ≤ 128 with extensions to L = 512, on predominantly monolingual corpora; the authors flag more complex embedding geometries at longer contexts and multilingual coverage as untested.
• Multimodal superposition. Whether mixing embeddings from intrinsically different modalities, such as text tokens with image patches, preserves these properties is completely unexplored.
• Scaling the stream count. The N = 3 extension shows quantitative degradation without qualitative change; behavior at higher N is not reported.
• Recovering the lost single-stream accuracy. The fine-tuning trade-off (Pythia-2.8B LAMBADA 0.544 to 0.357; Qwen2.5-3B 0.602 to 0.460) raises the question of whether superposition and single-stream quality can be jointly preserved.

Target Audience

Interpretability researchers studying residual-stream geometry and feature superposition; efficiency and inference-systems engineers interested in KV-cache reduction and multi-stream batching; and graduate-level readers already comfortable with Transformer internals, distributional distance metrics, and causal attention patching. Readers looking for a benchmark-topping model or a production-ready decoding recipe will find the paper explicitly positions its decoder as a proof of concept.

Authors’ abstract

While Large Language Models (LLMs) rely on highly non-linear components, in this work we demonstrate that they exhibit fundamental linearity: when inputs from distinct text streams are linearly combined, the model outputs a superposition of the individual next-token distributions. We term this the \textit{Superposition Linearity Hypothesis}. We provide evidence that superposition is an intrinsic property of the Transformer architecture rather than an emergent consequence of training; in fact, we observe that it tends to diminish as pretraining progresses. However, we demonstrate that linearity can be substantially restored through lightweight fine-tuning, significantly reducing the divergence between the predicted next-token distribution and the average of the individual next-token distributions. Finally, we introduce a guided decoding procedure that disentangles superposed outputs, enabling the simultaneous generation of two coherent continuations from a single forward pass.