PaperlayerView exact revision on arXiv

Source-linked AI summary

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

arXiv:2609.29845v1cs.CLcs.AI

TL;DR

The paper asks whether Transformers’ nonlinear components preclude linear end-to-end behavior when text streams are combined. It tests mixed inputs, finds that superposition diminishes during pretraining but can be restored by lightweight fine-tuning, and proposes decoding separate continuations from one mixed pass.

  • Problem

    Despite evidence of affine layer-to-layer structure, it is unclear whether a Transformer’s end-to-end output approximates the corresponding combination of outputs for separate streams.

    Source: §1 · p. 1

  • Method

    The authors average two streams’ token embeddings, train with self-distillation toward the frozen model’s average independent predictions, and use decoding to separate continuations.

    Source: §1 · p. 2, §2.1 · p. 3, §4 · p. 7

  • Results

    Superposition is intrinsic but diminishes during pretraining, and fine-tuning with less than 0.025% of the original pretraining-data size substantially restores it.

    Source: §1 · p. 2

  • Takeaways & Limitations

    Disentangling superposed outputs enables two coherent continuations from one forward pass, theoretically offering 2× inference throughput and halving KV-cache footprint per active stream.

    Source: §6 · p. 9

    Evaluations use relatively short contexts (L ≤128 for analytical experiments, extended to L = 512) and predominantly monolingual corpora, leaving longer-context and multilingual generalization open.

    Source: §8 · p. 9

Abstract

from arXiv · show

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.

1 Introduction

Although Transformers use highly nonlinear components, the paper asks whether end-to-end outputs preserve linear structure when distinct text streams are mixed. It reports intrinsic superposition that weakens during pretraining, can be restored with lightweight fine-tuning, and supports disentangling continuations from a mixed pass.

1 Introduction

  • Transformers’ nonlinear components motivate asking whether end-to-end responses to combined inputs approximate combinations of independent outputs.Earlier work found affine approximations to layer-to-layer residual-stream transitions, prompting the broader input–output question.Source: §1 · p. 1
  • Ground-truth next tokens from both streams frequently remain in the mixed distribution’s top 10, supporting the Superposition Linearity Hypothesis.The hypothesis describes mixed token streams as a superposition of their separate pathways.Source: §1 · p. 1
  • Superposition fidelity is strongest at initialization and diminishes as pretraining optimizes the language-modeling objective, suggesting it is architectural rather than learned.The paper also reports a strong correlation between hidden-state feature additivity and rank preservation.Source: §1 · pp. 1–2
  • Fine-tuning with less than 0.025% of the original pretraining-data size substantially restores linearity, and a decoding procedure recovers distinct continuations from one mixed pass.The contributions also include a mechanism to disentangle the mixed output distribution.Source: §1 · p. 2

2 Intrinsic Linearity in Large Language Models

The section tests whether standard Transformers preserve information and distributional structure when two token streams are averaged. It finds that valid tokens often remain highly ranked, mixed outputs resemble the target distribution, and linearity decreases during pretraining but persists with three streams.

2.1 Problem Formulation

  • They average paired token embeddings, pass the mixture through a frozen causal Transformer, and compare its output with the average of independent predictions across five model families.The evaluation uses paired sequences and tests whether both streams’ true next tokens retain mass in the mixed distribution.Source: §2.1 · pp. 2–3

2.2 Rank Analysis

  • 30–40% of true next tokens rank in the top 10, 50–60% in the top 50, and 60–65% or more in the top 100, despite vocabularies of at least 50,000 tokens.The mixed state narrows the search space around valid continuations rather than collapsing into gibberish.Source: §2.2 · p. 3

2.3 Distributional Shape Preservation

  • Across FineWeb, the mixed output has R_D < 1 for KL, JS, and Wasserstein distances, while total variation distance stabilizes after roughly 20 tokens.The reported pattern indicates preserved target-mixture structure and contextual stability across the window.Source: §2.3 · p. 4

2.4 Linearity Dynamics During Training

  • Hidden-state superposition error is lowest at the earliest Pythia checkpoints and grows monotonically during training; deep layers remain near-linear at ℓ≳2L/3.The complementary layer-wise analysis reports a U-shaped depth profile.Source: §2.4 · p. 4, Figure 2 · p. 5
  • With three equally weighted streams, superposition persists with moderate degradation: R_KL increases by +0.04 to +0.09 across models, without qualitative change.The authors interpret this as evidence that the same interference mechanisms operate across stream counts.Source: §2.4 · pp. 4–5

3 An attention-patching analysis

The attention-patching experiments separate the effects of attention structure and content, then compare those perturbations with embedding mixing. Results distinguish predictable from content positions: attention shape and frequency priors explain some rank survival, but embedding mixing retains more signal on hard LAMBADA predictions, and fine-tuning restores content-position survival.

3 An attention-patching analysis

  • The experiment separates attention shape from content: donor patching substitutes unrelated text's attention weights, while permutation shuffles A's causal-prefix weights.Donor patching leaves A's Q/K/V projections and value paths unchanged; permutation preserves causality, row sums, and per-row weight multisets.Source: §3 · p. 5
  • 284 is the content-position median rank under embedding mixing, versus 111 for donor patching; permutation collapses to 3,079 predictable and 19,246 content.On predictable positions, embedding mixing and donor patching retain median ranks of 3–6; exact agreement remains near 25–33%.Source: Table 2 · p. 6, §3 · p. 6
  • Permutation raises RKL from 0.27 to 0.68 and lowers top-10 agreement from 53% to 10%, showing the frequency prior alone is insufficient.Donor patching preserves natural attention shape, whereas permutation destroys it while retaining the LM-head frequency prior.Source: Table 3 · p. 6, §3 · p. 6
  • On LAMBADA, embedding mixing achieved 2.25% raw argmax accuracy and target median rank 339, versus donor patching’s 0.5% raw argmax accuracy and target median rank 2,350, a 4.5× accuracy and 7× rank gap.Source: §3 · p. 7
  • Median self-agreement rank is 19 for embedding mixing versus 8 for donor patching, despite embedding mixing's stronger LAMBADA performance.The aggregate self-agreement ranking reverses on LAMBADA, where the target is a content word at the final position of a long-narrative passage.Source: §3 · pp. 6–7
  • After fine-tuning, content-position median rank drops from 284 to 5, with exact top-1 agreement reaching 22.8%, while predictable-position median rank remains 8.The aggregate median rank improves from 19 to 6.Source: §3 · p. 7

4 Improving Linearity with Finetuning

Lightweight self-distillation aligns mixed-input predictions with averaged independent predictions and substantially restores superposition fidelity, including on hard content tokens, while reducing single-stream quality.

4 Improving Linearity with Finetuning

  • Self-distillation initializes a student from pretrained weights and trains its mixed-input output toward the mean of a frozen same-model teacher’s independent predictions.Source: §4 · p. 7
  • Experiments fine-tuned Pythia, Qwen, and Llama on a FineWeb subset for approximately 200k steps.Source: §4 · p. 7
  • On Pythia-2.8B, mean KL divergence drops from 1.86 to 0.27 and Superposition Approximation Ratio RKL drops from 0.42 to 0.06.The probability of a true next token appearing in the top-5 rises from approximately 30% to over 60%.Source: §4 · p. 7
  • On hard content tokens, median rank falls from 284 to 5 and exact top-1 agreement reaches 22.8%.The passage reports that this recovery is not limited to high-frequency predictable tokens.Source: §4 · p. 7
  • The restoration is not free: fine-tuning reduces single-stream next-token quality, including Pythia-2.8B LAMBADA from 0.544 to 0.357.Source: §4 · p. 7

5 Decoding Two Streams Out of a Mixed Forward Pass

Direct sampling from a fitted mixed distribution can alternate between streams because geometric-mean probabilities penalize stream-specific tokens. Joint Contrastive decoding offers a proof-of-concept separation method that improves LAMBADA accuracy but does not close the gap to a single-stream baseline.

5 Decoding Two Streams Out of a Mixed Forward Pass

  • Direct sampling can alternate between streams because approximately averaged logits yield geometric-mean probabilities that penalize tokens unlikely in the other stream.Source: §5 · p. 8
  • The decoder remains a proof of concept: it does not close the single-stream baseline gap, and fully overcoming the geometric-mean obstruction remains open.Source: §5 · pp. 8–9
  • Joint Contrastive decoding uses a small auxiliary model to guide per-stream logits, with α and β trained jointly with the backbone using symmetric per-stream cross-entropy.Source: §5 · p. 8
  • Joint Contrastive raises LAMBADA mean accuracy from 0.06–0.18 to 0.11–0.43; Llama-3.2-3B reaches 0.43 versus the single-stream small-model baseline of 0.54.The reported residual gap is consistent with the geometric-mean obstruction.Source: §5 · pp. 8–9

6 Conclusion

Transformers can process linearly combined text streams as superposed independent pathways, and lightweight fine-tuning can restore distinct semantic signals within the mixed hidden state.

6 Conclusion

  • The paper concludes that Transformers process linearly combined text streams as a superposition of independent pathways, a property that degrades during standard pre-training but can be restored through lightweight fine-tuning.Source: §6 · p. 9
  • Disentangling enables two coherent continuations from one forward pass, theoretically offering 2× inference throughput and halving the KV-cache footprint per active stream.Source: §6 · p. 9

7 Related Work

Prior work studies linear residual-stream geometry and engineered multiplexing or parallel generation; this paper distinguishes its focus on intrinsic input–output superposition in standard pretrained Transformers.

7 Related Work

  • Consecutive-layer mappings in decoder-only Transformers are often approximated by affine transforms, as reported by Razzhigaev et al. [2024].Source: §7 · p. 9
  • DataMUX [Murahari et al., 2022], MIMONets [Menet et al., 2023], and RevMUX [Xu et al., 2024] impose multiplexing through dedicated structures, binding and unbinding, or reversible adapters.Source: §7 · p. 9
  • Superposed Decoding [Shen et al., 2024] mixes draft embeddings for parallel generation, while superposition prompting [Merth et al., 2024] processes multiple retrieval paths in one pass.Source: §7 · p. 9
  • Unlike engineered multiplexing approaches, this work studies intrinsic input–output superposition under embedding averaging in standard pretrained models and its restoration through fine-tuning.Source: §7 · p. 9

8 Limitations

The paper identifies limits in the evaluated contexts, languages, and modalities, alongside compute requirements and a frequency-baseline control. It also reports a layer-averaged hidden-state additivity error, where lower values indicate stronger superposition linearity.

8 Limitations

  • Evaluations use contexts up to L = 128 for analytical experiments, with extensions to L = 512, predominantly monolingual text, and text-only models.Generalizing to substantially longer contexts or multilingual settings involves more complex embedding geometries; multimodal extensions remain untested.Source: §8 · p. 9

A Compute Resources

  • A Pythia-2.8B fine-tuning run took approximately 114 hours on 2×A100 80GB GPUs; Qwen2.5-3B took 132 hours and Llama-3.2-3B took 128 hours on the same hardware.Joint Contrastive guided decoding fine-tuning for Pythia-1.4B with a 160M guide model took 68 hours on one A100 80GB GPU.Source: A Compute Resources · p. 12

B Frequency-Baseline Control

  • Under the unrelated-stream control, cross-stream tokens ranked top-10 in 2.63% of cases, versus 30–40% under superposition, exceeding the frequency-only baseline by an order of magnitude.The control’s top-3 and top-100 rates were 1.12% and 10.41%, respectively; same-position token overlap was 0.2%.Source: B Frequency-Baseline Control · p. 12

C Hidden-state additivity error

  • The layer-averaged hidden-state additivity error summarizes the distance between mixed and combined single-stream hidden states, with lower error indicating stronger superposition linearity.The calculation uses mean-centered, normalized hidden states from stream A, stream B, and the mixed input.Source: C Hidden-state additivity error · p. 12

F Layer-wise Linearization Dynamics

Layer-wise linearity follows a U-shaped depth profile, and fine-tuning modestly improves local scores while reducing global divergence. The section also describes alternative separation methods and reports their evaluation scope and trade-offs.

F Layer-wise Linearization Dynamics

  • Linearity scores are high in layers 0–5, fall to roughly 0.65 in layers 6–20, then recover above 0.95 in the final third, where terminal linearity preserves the superposed signal before output logits.After fine-tuning, local score improvements are modest but consistent with reduced global divergence because per-layer nonlinear errors compound across depth.Source: F Layer-wise Linearization Dynamics · pp. 12–13, Figure 5 · p. 14
  • Tokens with original probability P > 0.5 almost always survive superposition at median rank ≈3, while divergence is slightly lower for the first 20 tokens before stabilizing across the context window.Source: Figure 3 · p. 13, Figure 4 · p. 13

G.1 Gradient-Optimized Inference (Logit Arithmetic)

  • Gradient-optimized logit arithmetic combines the large model’s mixed logits with small-model stream logits, but its post-hoc approach achieves lower accuracy than Joint Contrastive decoding.The large and small model weights remain frozen while scalar coefficients are optimized on a small calibration dataset.Source: §G.1 · pp. 13–14

G.2 Two-Head Separation via Symmetry Breaking

  • Two-Head separation learns stream-specific projections and output heads, using dynamic loss balancing because naïve joint cross-entropy often collapses onto a single stream.The projections break the permutation symmetry of averaging; the two-head results are reported in Table 6.Source: §G.2 · p. 14

G.3 TinyStories LLM-as-Judge

  • The TinyStories evaluation compares independent, raw-superposed, and guided-decoding Pythia-2.8B completions using mean judge ratings across streams on a 1–10 scale.The full separation and recovery metrics are reported in Table 6.Source: §G.3 · p. 14, Table 7 · p. 15, Table 6 · p. 15

G.4 Single-stream Language Modeling Quality on FineWeb

  • Joint Contrastive guidance preserves single-stream fluency close to the small-model baseline, whereas Tuned distillation incurs a substantially larger penalty.Source: Table 8 · p. 15

G.5 Throughput and Memory

The section compares fused-stream generation with independent decoding and describes qualitative token-rank examples. Separate decoding delivers throughput close to batch-two inference while generating two continuations from a fused embedding.

G.5 Throughput and Memory

  • Separate matches Big(b=2) within 3% across all pairs and delivers roughly 2× the throughput of sequential vanilla decoding while producing two coherent continuations from a fused embedding.Memory is modestly above one vanilla pass for Separate (Pythia-2.8B 7.15 vs. 6.34 GB; Llama-3B 9.26 vs. 6.72 GB); Guided is slower and can peak at 12.93 GB for Llama-3B+1B.Source: §G.5 · pp. 15–16, Table 9 · p. 16

H Attention-Patching: Stratified Examples

  • Across rank buckets, low-ranked examples are mostly punctuation and function words, mid-ranked examples are mostly content words, and high-ranked examples are almost entirely content or rare BPE fragments.The observed function/content split is described as a systematic property of the rank distribution, not an ad-hoc heuristic.Source: H Attention-Patching: Stratified… · p. 16

I Distributional Metrics for N = 3 Superposition

For three mixed streams, the output still approximates the mixture of independent distributions, though approximation degrades relative to two streams, especially at longer context length. The accompanying material gives implementation settings, judge-prompt criteria, and potential impacts and risks.

I Distributional Metrics for N = 3 Superposition

J Hyperparameters and Implementation Details

  • Distillation uses FineWeb or FineWeb-Edu at maximum sequence length 128, while Joint Contrastive guided-decoding runs use FineWeb at length 512.Distillation learning rates are 10^-3 for heads and 10^-4 for the backbone; guided-decoding rates are 5 × 10^-5 for Qwen and Llama and 10^-4 for Pythia.Source: J Hyperparameters and Implementation… · p. 17

K LLM as a Judge Prompt

  • The TinyStories grader assesses grammar, creativity, consistency, and plot coherence on 1–10 scales and estimates the writer’s age group.Source: K LLM as a Judge Prompt · p. 17

L Impact Statement

  • The authors suggest that throughput gains and lower infrastructure costs may benefit high-volume or real-time generation, while warning that mixing unrelated or sensitive streams could introduce risks.Source: L Impact Statement · p. 18

Related papers

Loading 2609.29845v1…