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

This article explains the key points of Your Transformer Can Hold Two Thoughts at Once: Evidence of Linear Superposition in LLMs.

• 2026-09-24 (arXiv)
• Tikhonov, Pavel, Korznikov, Anton, Mikhalchuk, Matvey, Dragunov, Nikita, Rahmatullaev, Temurbek, Druzhinina, Polina, Razzhigaev, Anton, Oseledets, Ivan, Tutubalina, Elena.
• Paper

Read this article in Korean

Summary

• The paper asks whether a large language model (LLM), which predicts the next text token, can preserve information from two unrelated texts when their embeddings, the vectors representing input tokens, are averaged. Compared with unmodified Pythia-2.8B, targeted fine-tuning reduces the reported Kullback–Leibler (KL) divergence, a measure of distribution mismatch, from 1.86 to 0.27 between the mixed-input output and the average of the two independent outputs.
• The Superposition Linearity Hypothesis proposes that averaging input representations approximately preserves the average of their independent next-token distributions. In the illustrated unmodified models, independent top predictions remain within the mixed distribution’s top-10 in approximately 30–40% of cases, although aggregate recovery disproportionately reflects predictable tokens rather than difficult content words.
• Hidden-state additivity, agreement between mixed internal representations and sums of independent representations, generally deteriorates during Pythia pre-training. Self-distillation, training a model against a frozen copy of itself, improves mixed-output fidelity and Qwen2.5-3B content-position median rank from 284 to 5, while reducing ordinary single-stream prediction quality.
• Joint Contrastive decoding uses a small auxiliary model to separate the streams by adding guidance for one and subtracting guidance for the other. On LAMBADA, a benchmark requiring broad narrative context to predict a word, Llama-3.2-3B mixed accuracy improves from 0.182 to 0.430, remaining below its small guide model’s single-stream accuracy of 0.540; separate generation from mixed inputs remains a proof of concept.

1 Introduction

The paper extends earlier observations of approximately affine transitions, linear transformations with an added offset, between Transformer layers to the complete input–output mapping. A Transformer processes token representations through attention and nonlinear transformations; the question here is whether two independent text streams survive element-wise embedding averaging without a dedicated mechanism for encoding and recovering stream identities.

• The authors distinguish an architectural tendency from a learned capability by tracking intermediate pre-training checkpoints.
• They report restoration using less than 0.025% of the original pre-training dataset size, then investigate whether the preserved signal supports separate continuations.

2 Intrinsic Linearity in Large Language Models

The initial experiments examine standard pre-trained Transformers without architectural changes. The analysis separates preservation of individual token rankings, similarity of complete output distributions, and changes in internal representations during training.

2.1 Problem Formulation

For two equal-length token sequences, the mixed embedding at position t is e_t(z) = 1/2(E(x_t) + E(y_t)), where E maps tokens to vectors. A frozen decoder-only Transformer, which predicts subsequent tokens from preceding context, processes this sequence with causal self-attention, restricting each position’s access to its current and preceding positions. It produces P_mix(z) = softmax(ℓ(z)), where ℓ denotes logits, the prediction scores that softmax converts into probabilities.

• The target is the arithmetic average P_target(x|A,B) = 1/2(P(x|A) + P(x|B)).
• The evaluated families include Pythia, Qwen, Llama, OLMo, and Gemma. Random text pairs come from TinyStories, which uses simplified grammar, and FineWeb, which contains natural web text, and are truncated to fixed lengths.

2.2 Rank Analysis

The operational rank metric follows each independent stream’s argmax token, its highest-scoring prediction, into the mixed output distribution. Although the prose sometimes calls these tokens “ground-truth,” Figure 1 and the formal definition use model predictions; their preservation measures agreement with independent inference rather than correctness against observed text.

• Across the illustrated unmodified models, approximately 30–40% of tracked tokens remain in the top-10, 50–60% in the top-50, and 60–65% in the top-100.
• These ranks show signal retention for vocabularies of at least 50,000 tokens, but do not establish that two correct continuations can be selected independently.

Independent top predictions often remain highly ranked after embedding mixing, and fine-tuning increases their recovery. The curves track agreement with independent predictions rather than correctness against observed next tokens.

2.3 Distributional Shape Preservation

Distributional preservation is measured using KL divergence, Jensen–Shannon (JS) divergence, a symmetric distribution mismatch measure, and Wasserstein distance, a transport-based distance using token-embedding similarity. KL and JS use temperature smoothing with τ = 1.5, which softens probabilities for numerical stability; Wasserstein distance uses raw probabilities over the top-256 tokens and cosine distance, which compares embedding directions.

• The Superposition Approximation Ratio divides the expected mixed-to-target mismatch by the expected mismatch between unrelated independent contexts. A ratio below 1 establishes improvement over that reference, rather than exact linearity.
• At context lengths L = 32 and L = 512 on FineWeb, every reported ratio is below 1. For Pythia-2.8B at L = 32, the KL, JS, and Wasserstein ratios are 0.42, 0.54, and 0.71.
• The main text says mismatch is initially higher, whereas Figure 4 shows lower initial Total Variation Distance, a measure of probability-mass difference, followed by stabilization.

All reported normalized distances are below the unrelated-context reference at both tested lengths. The remaining absolute divergences indicate approximate preservation; the displayed rows do not include the fine-tuned result cited later in the paper.

2.4 Linearity Dynamics During Training

For each text pair and checkpoint, the authors run both independent streams and the mixed stream, center hidden states by their layer means, and compare normalized mixed states with normalized sums of independent states. Lower layer-averaged distance indicates stronger additivity; Figure 2 shows an overall increase during Pythia pre-training, with local reversals and late declines in several curves.

• This trajectory supports an architectural contribution to additivity, but the metric does not establish useful language processing at initialization.
• Equal-weight mixing of three streams retains distributional similarity relative to unrelated contexts while increasing interference. The full tables show KL-ratio increases of +0.04 to +0.10, including the +0.10 case outside the main text’s stated +0.04 to +0.09 range.

Additivity error generally grows from early to later Pythia checkpoints, supporting deterioration during pre-training. Several curves reverse locally or decline late, qualifying the text’s claim of monotonic growth.

3 An attention-patching analysis

The attention-patching experiments use Qwen2.5-3B on FineWeb-Edu, the text corpus used for these controls, at sequence length T = 128 to separate attention structure from text-specific attention weights. Donor patching replaces every layer’s and head’s normalized attention weights with those from an unrelated text; permutation patching randomly rearranges weights within each causal prefix while preserving their row sums and values.

• For predictable positions, approximately 65% of positions, embedding mixing and donor patching give median ranks of 6 and 3. For content positions, approximately 35%, their median ranks deteriorate to 284 and 111; the appendix associates these positions mainly with content words and rare token fragments.
• Permutation patching produces median ranks of 3,079 on predictable positions and 19,246 on content positions. Preserving the distribution of attention weights alone therefore does not preserve the original predictions.
• After fine-tuning, mixed-input content-position median rank becomes 5 and exact top-1 agreement rises from 4.2% to 22.8%; predictable-position median rank changes from 6 to 8.

Splitting positions by token type exposes weak base-model retention of content-dependent predictions. Fine-tuning improves content median rank from 284 to 5 and top-1 agreement from 4.2% to 22.8%, while predictable-position metrics slightly worsen.

Aggregate agreement can conceal failure on context-dependent tasks. Donor patching preserves an overall median rank of 8 and a KL approximation ratio of 0.27, yet on 200 LAMBADA prompts truncated from the left to 128 tokens it reduces accuracy from the reported vanilla 73% to 0.5%.

• Permutation increases overall median rank to 8,148, lowers top-10 recovery from 53% to 10.1%, and raises the KL ratio to 0.68.
• The control supports a joint role for natural attention structure and the output layer’s token-frequency preferences; aggregate rank agreement does not establish task accuracy.

Permuting attention weights causes much greater damage than substituting natural donor attention. The difference supports a role for attention structure beyond the unchanged output layer’s token-frequency preferences.

Embedding mixing performs worse than donor patching on overall agreement, with median ranks of 19 versus 8, but better on difficult LAMBADA targets. Its raw accuracy is 2.25% versus 0.5%, with target median ranks of 339 versus 2,350.

• This comparison supports retention of some context-specific signal beyond the information surviving an unrelated attention-weight substitution.
• The absolute accuracies remain low, so this result establishes partial signal preservation rather than reliable task performance.

Embedding mixing performs worse than donor patching on aggregate agreement but better on LAMBADA accuracy and target rank. This reversal shows why aggregate prediction agreement can misrepresent retained context-specific information.

4 Improving Linearity with Finetuning

The fine-tuning objective uses a frozen pre-trained teacher to predict each stream independently, averages those probability distributions, and trains a student on averaged input embeddings. The loss is L = D_KL(P_target || M_student(z)); the paper describes applying this procedure to Pythia, Qwen, and Llama on a FineWeb subset for approximately 200k steps.

• For Pythia-2.8B, the reported mean KL divergence falls from 1.86 to 0.27 and its approximation ratio from 0.42 to 0.06.
• The text also reports top-5 recovery increasing from approximately 30% to more than 60%. Its citation to a fine-tuned row in Table 1 is inaccurate: the displayed table contains only the base-model comparisons.

Improving mixed-output fidelity incurs a single-stream quality cost. The text reports LAMBADA accuracy declining from 0.544 to 0.357 for Pythia-2.8B and from 0.602 to 0.460 for Qwen2.5-3B.

• Appendix G.4 reports increases in perplexity, the exponentiated average negative log-likelihood of observed tokens, after tuned distillation: Qwen2.5 rises from 12.44 to 65.52, Llama-3.2 from 11.20 to 71.44, and Pythia from 14.09 to 43.95.
• Joint Contrastive guidance has smaller penalties, with perplexities of 19.31, 14.75, and 29.72, respectively, compared with small-model baselines of 18.96, 13.53, and 29.53.

5 Decoding Two Streams Out of a Mixed Forward Pass

Separate decoding is harder than approximating an averaged output distribution: direct sampling can alternate between the two contexts. The paper describes a geometric-mean obstruction under the condition that logits are averaged: P′_target(t) ∝ exp(1/2(ℓ_A(t) + ℓ_B(t))) ∝ √(P_A(t)P_B(t)), which suppresses tokens strongly favored by only one stream.

• This conditional logit-averaging argument differs from the arithmetic probability mixture used as the distillation target; the equation alone does not establish the cause of the remaining decoding error.
• Joint Contrastive decoding forms each stream’s scores by adding its small-model scores and subtracting the other stream’s scores from the shared large-model scores. The coefficients α and β start at 1 and are trained jointly with the backbone, the main model, using symmetric per-stream cross-entropy, a loss that rewards the observed next tokens.

With Joint Contrastive decoding, Qwen2.5-3B mixed LAMBADA accuracy increases from 0.168 to 0.345 and Llama-3.2-3B from 0.182 to 0.430. Their FineWeb Jaccard token overlaps, the fraction of shared tokens relative to the union, decrease from 0.126 to 0.061 and from 0.094 to 0.067.

• Both remain below their small guide models’ independent accuracies of 0.437 and 0.540. Low token overlap measures separation, but does not independently establish coherent or correct continuations.
• The Pythia comparison changes backbone size: the pre-trained row uses Pythia-2.8B at 0.065, while Joint Contrastive uses Pythia-1.4B at 0.110. It is therefore not a controlled before-and-after comparison of the same backbone.

Joint Contrastive decoding improves mixed accuracy for the matched Qwen and Llama backbones and lowers inter-stream token overlap. Both accuracies remain below the corresponding independent small-model baselines, while the Pythia rows use different backbone sizes.

6 Conclusion

The conclusion presents recoverable superposition as a possible route to two continuations from one shared backbone pass, with a theoretical 2× throughput benefit and a halved shared key–value cache, the stored attention keys and values, per stream. The appendix measurements qualify this claim: the two-head mode approximately matches ordinary batch-size-2 throughput, while the guided mode adds auxiliary-model computation and memory.

7 Related Work

Related work includes DataMUX, MIMONets, and RevMUX, methods that combine inputs into a shared representation and recover separate outputs; Superposed Decoding mixes draft-token embeddings, while superposition prompting processes multiple document paths. This paper’s distinguishing empirical claim is that simple embedding averaging already retains measurable signal in standard pre-trained Transformers, and targeted optimization strengthens that behavior.

8 Limitations

The stated limitations concern short, predominantly monolingual text contexts and the absence of multimodal experiments, which would combine different input types such as text and images. Analytical experiments mainly use L ≤ 128, distributional comparisons extend to L = 512, and guided training also uses maximum length 512; substantially longer contexts, multilingual inputs, and mixing text with image representations remain untested.

• Generation accuracy remains below independent inference, and the strongest guided results require a separate small model.
• Appendix measurements show that throughput depends on the decoding variant: two-head separation is close to batched inference, while guidance is slower. The reported memory figures do not demonstrate a general halving of total deployment memory.

Intermediate-to-high confidence predictions generally retain better mixed ranks than low-confidence predictions. Sharp deterioration in the highest-confidence Pythia bins limits the caption’s broad claim about robust high-confidence recovery.

Mismatch rises from the earliest positions and then remains comparatively stable over the measured window. This supports positional stability but contradicts the main text’s statement that the earliest positions have higher mismatch.

Fine-tuning changes consecutive-layer linearity differently across models, with particularly large gains for Llama-3.2-3B. The base-model profiles do not uniformly support the accompanying description of near-linear final thirds.

Alternative separation objectives do not consistently improve mixed accuracy: Llama Two Heads falls below its pre-trained mixing baseline, whereas Pythia-2.8B Two Heads improves it. No separately labeled gradient-optimized inference row appears despite the appendix’s reference to one.

Fine-tuning recovers much of the consistency loss and some grammar quality under the automated grader, but creativity barely improves. The table reports three dimensions despite the evaluation protocol requesting additional outputs.

Tuned distillation substantially worsens single-stream likelihood across all three families. Joint Contrastive scores remain near the smaller independent models, showing a quality trade-off relative to the large independent backbones.

Separate runs at approximately twice sequential throughput and close to ordinary batch-size-2 inference. Guided decoding is slower than batched inference in every reported pair, so the fastest mode and the strongest guided accuracy results require separate assessment.

Three-stream mixing remains closer to its target than the unrelated-context reference under every reported metric and length. This establishes retained distributional structure without demonstrating independent three-stream generation.

Every reported approximation ratio worsens when moving from two streams to three. The largest KL increase is +0.10 for Pythia-2.8B at L = 512, exceeding the main text’s stated upper bound of +0.09.

Appendix

• A Compute Resources: Pythia-2.8B fine-tuning for 150,000 steps takes approximately 114 hours on 2×A100 80GB graphics processing units (GPUs). Qwen2.5-3B and Llama-3.2-3B distillation take 132 and 128 hours on the same setup; Pythia-1.4B Joint Contrastive training with a 160M guide takes 68 hours on one A100 80GB GPU.
• B Frequency-Baseline Control: Predicting an unrelated stream’s “ground-truth next token,” as described in the appendix, from single-stream Pythia-2.8B yields top-3, top-10, and top-100 recovery of 1.12%, 2.63%, and 10.41%, with 0.2% same-position token overlap. These controls are substantially below the reported mixed-input recovery, but the target terminology differs from Figure 1’s explicit use of independent model predictions.
• C Hidden-state additivity error: The metric centers each representation using its layer-specific evaluation-set mean, normalizes the mixed representation and the sum of the independent representations, then measures their Euclidean distance. Averaging over pairs, positions, and layers gives the summary error used to track pre-training dynamics.
• D Confidence-vs-Rank Curve: Figure 3 relates mixed-output median rank to the independent prediction’s probability. Its caption summarizes predictions with P > 0.5 as having median rank approximately 3, but the plotted Pythia curves worsen sharply near probability 1, so the relationship is not uniform across models and confidence bins.
• E Contextual Stability of Superposition: Figure 4 shows lower initial Total Variation Distance for t < 20, followed by a relatively stable, model-dependent level through approximately 128 positions. Stability across these positions supports persistence over the tested window, while the sizable mismatch does not establish a close match to the target distribution.
• F Layer-wise Linearization Dynamics: Consecutive-layer representations are centered and Frobenius-normalized, meaning that each representation matrix is divided by the square root of its summed squared entries, then compared using the best-fitting linear transformation. Figure 5 shows model-dependent profiles and improvements after fine-tuning, especially for Llama-3.2-3B; several displayed base-model curves do not support the text’s general description of near-linear final thirds.
• G Additional Decoding Variants and Throughput: This appendix compares mixed-distribution fitting, stream recovery, story quality, ordinary language modeling, and execution speed. Improvements in mixed-distribution fidelity do not consistently produce better separate decoding or preserve single-stream quality.
• G.1 Gradient-Optimized Inference (Logit Arithmetic): Both models remain frozen while scalar guidance coefficients are calibrated by gradient descent on next-token cross-entropy. The text states that this post-hoc method underperforms Joint Contrastive decoding, but the cited Table 6 does not display a separately identified gradient-optimized row.
• G.2 Two-Head Separation via Symmetry Breaking: Learnable input projections and two output heads, separate mappings from hidden states to prediction scores, distinguish stream identities, with projections initialized as W_k = I + N(0, σ²), σ = 10⁻³. Dynamic loss balancing adjusts the streams’ training weights to address collapse onto one stream; Table 6 reports mixed LAMBADA accuracy of 0.164 for Pythia-2.8B and 0.105 for Llama-3.2-3B, showing inconsistent benefits across backbones.
• G.3 TinyStories LLM-as-Judge: OpenAI/GPT-5.2 serves as an automated grader of story completions. For Pythia-2.8B, fine-tuned superposition improves grammar from 2.59 to 3.69 and consistency from 2.70 to 3.53, while creativity changes from 2.08 to 2.13 and remains below the independent baseline of 5.72; Table 7 omits the described plot-coherence and estimated-age results.
• G.4 Single-stream Language Modeling Quality on FineWeb: Table 8 reports negative log-likelihood, the average loss assigned to observed tokens, and perplexity after each objective. Joint Contrastive results remain near the small-model baselines, while tuned distillation substantially worsens ordinary language modeling in all three reported families.
• G.5 Throughput and Memory: Benchmarks use prompt length 128, generation length 128, and n = 100 runs, reporting mean tokens/s with 95% confidence intervals. For Pythia-2.8B, two-head Separate achieves 109.6 ± 0.8 tokens/s, Guided 79.4 ± 0.5, ordinary batch-size-2 inference 113.1 ± 0.6, and sequential inference 56.0 ± 0.3; Separate peak memory is 7.15 GB versus 6.34 GB for one vanilla pass, while Llama-3B+1B Guided peaks at 12.93 GB.
• H Attention-Patching: Stratified Examples: Uniform samples from rank buckets associate low ranks [2,10] mainly with punctuation and function words, middle ranks [50,500] mainly with content words, and high ranks [2000,20000] mainly with content words or rare token fragments. The examples support the qualitative distinction between predictable and content positions, but do not provide a fully specified general classification procedure.
• I Distributional Metrics for N = 3 Superposition: Equal-weight mixing of three unrelated contexts yields approximation ratios below 1 at L = 32 and L = 512 for all reported models and distances. Compared with two-stream mixing, KL ratios increase by +0.04 to +0.10, JS ratios by +0.03 to +0.09, and Wasserstein ratios by +0.01 to +0.08; these experiments evaluate distribution preservation rather than successful three-stream decoding.
• J Hyperparameters and Implementation Details: Distillation uses the AdamW optimizer with learning rates 10⁻³ for heads and 10⁻⁴ for the backbone, the ReduceLROnPlateau learning-rate scheduler with patience 10 and factor 0.5, temperature 2.0, and maximum length 128. Joint Contrastive training uses learning rate 5×10⁻⁵ and effective batch size 16 for Qwen and Llama, or 10⁻⁴ and batch size 8 for Pythia, with maximum length 512, gradient clipping at 1.0, bf16 reduced-precision computation, and scaled dot-product attention (SDPA); checkpoints are taken at 30,000 steps for Qwen and Llama and 700,000 for Pythia. Two-head training uses balancing strength γ = 5.0, exponential-moving-average momentum 0.99, 100 warmup steps, and coefficient clipping to [0.1,10.0].
• K LLM as a Judge Prompt: Listing 1 provides the story-completion grading template, separating the opening and continuation with ***. It requests grammar, creativity, consistency, and plot-coherence scores out of 10, plus an estimated age group; the listing documents the evaluation procedure.
• L Impact Statement: The authors propose possible infrastructure benefits from higher throughput and identify inadvertent mixing of unrelated or sensitive streams as a deployment concern. The paper does not experimentally evaluate downstream stream-isolation failures.

Brief Thoughts

The strongest evidence combines token-type stratification with targeted fine-tuning: content-position median rank improves from 284 to 5, whereas predictable positions do not improve. This identifies an effect that aggregate rank preservation alone cannot establish, since donor-patching controls show that high aggregate agreement can coexist with severe task failure.

The deployment case remains limited by lower generation accuracy, single-stream degradation, and auxiliary-model costs. Interpreting the mechanism and reproducing the gains also requires distinguishing the arithmetic-mixture training target from the conditional geometric-mean argument and accounting for inconsistent descriptions of early-position mismatch and absent cited result rows.