PaperlayerView exact revision on arXiv

Source-linked AI summary

The Geometry of Inference in Transformer Residual Streams

Timur Mudarisov, Mikhail Burtsev, Radu State

arXiv:2609.37824v1cs.CLcs.LG

TL;DR

The paper asks how selectively intermediate residual states distinguish their own final representations from alternatives. It compares trajectories with empirical endpoint banks and models how endpoint geometry shapes competition. Across six models, own-endpoint preference appears early despite many closer alternatives, and competitor sets can change as they shrink.

  • Problem

    Prior work identifies predictive information in intermediate states and output-related structure in final representations, but how selectively an intermediate state distinguishes its own final representation from alternatives across depth remains open.

    Source: §1 · p. 1

  • Method

    The paper compares intermediate residual states with their own final endpoints and a bank of endpoints from other contexts, and models how endpoint geometry determines competition.

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

  • Results

    Across six models, the own endpoint is preferred to the average alternative from the earliest measured layers although many individual alternatives remain closer, and cosine competitor sets show entries as well as exits as their mean count falls.

    Source: §1 · p. 1, §4 · p. 9

  • Takeaways & Limitations

    Empirical results and theoretical models support the view that residual updates refine early endpoint preference into greater selectivity while individual comparisons can be revised.

    Source: §4 · p. 9

    Competition depends on the endpoint bank and metric, endpoint geometry captures only part of output variation, and observed turnover does not establish an explicit internal search procedure.

    Source: §4 · p. 9

Abstract

from arXiv · show

Transformer language models build predictions through successive residual updates, but how their representations become specific to an eventual outcome remains unclear. We study this process by comparing intermediate residual states with their own final states and an empirical bank of final states from other contexts. Across six pretrained language models, the own endpoint becomes preferable to the average alternative early, while many individual endpoints remain closer. These competing sets generally shrink with depth, but their membership changes and their surviving endpoints need not become more similar to one another. Directional alignment and endpoint rank can therefore improve while Euclidean distance to the final state changes little. We develop a simple high-dimensional model that separates the roles of norm, alignment, and endpoint geometry, showing how gradual directional changes can produce sharp reductions in competition. We also prove that a straight path toward the own endpoint cannot introduce new competitors under either Euclidean or cosine distance; observed entries thus establish departures from straight-line convergence. Finally, endpoints associated with lower-ranked output tokens tend to lie farther away in cosine distance across all studied models, connecting residual geometry to output organization. Together, these findings characterize increasing geometric specificity during transformer inference and explain why distance, competitor count, and concentration of the surviving endpoints provide distinct views of that process.

1 Introduction

The paper studies how intermediate residual states distinguish their own final representations from alternatives, using an endpoint bank to track that geometry. Across six pretrained models, early preference for the own endpoint coexists with many closer alternatives, while theoretical analysis explains how competition can later decline and change membership.

Geometric inference hypothesis

  • The endpoint bank treats another context’s final residual state as an alternative, with closer-than-own endpoints forming the competitor set.The paper’s geometric inference hypothesis is that residual updates refine partial distinctions among possible final states through changes in relative geometry.Source: §1 · p. 1

Contributions

  • Across six pretrained models, early preference for the own endpoint coexists with many closer alternatives, and cosine competitor sets can gain and lose members as their mean size declines.The paper also relates cosine endpoint distance to output-token rank and analyzes how gradual alignment can sharply reduce competition.Source: §1 · pp. 1–2
  • A straight path toward the own endpoint yields nested competitor sets under Euclidean and cosine distance, so observed entries indicate departures from that baseline.A structured endpoint distribution fitted only to final states predicts mean competition along held-out trajectories more accurately than a uniform-sphere baseline.Source: §1 · p. 2

2 Geometry of Residual Inference in Language Models

The section probes inference by comparing intermediate residual states with endpoints from other contexts, measuring both average preference for the own endpoint and the number of closer alternatives. Across models, preference appears early, while competition generally declines with depth through changing membership; cosine and Euclidean measures can tell different parts of that story.

2.1 An Endpoint-Based Probe of Inference

  • The probe compares each intermediate residual state with its own final state and a bank of other contexts’ final states, treating closer alternatives as competitors.The study uses six pretrained models and 1,024 non-overlapping 256-token contexts per model, with 1,023 alternatives per query.Source: §2.1 · pp. 2–3
  • Endpoints whose top predictions are less probable under the query context tend to lie farther from its own endpoint in cosine distance across all six models.Endpoints sharing a top prediction are more compact on average; endpoint distance is associated with output-distribution divergence, but explains only part of output variation.Source: §2.1 · p. 3, Figure 2 · p. 4

2.2 Early Preference with Many Competitors

  • Across all six models and both metrics, the document-averaged margin favors the own endpoint from the earliest measured post-block state, although many individual alternatives remain closer.Under Euclidean distance, five models retain hundreds of competitors across substantial portions of depth.Source: §2.2 · p. 3, Figure 3 · p. 5

2.3 Competitor Sets Shrink with Changing Membership

  • Cosine competition often decreases earlier than Euclidean competition, and directional alignment can improve while Euclidean endpoint distance remains nearly constant.This pattern is especially noted in intermediate layers of Mistral-7B and Llama-3-8B.Source: §2.3 · pp. 3–4
  • Excluding same-document endpoints leaves cosine preference and competitor-fraction curves essentially unchanged, while restricting alternatives to the same final input token preserves preference with smaller margins and more competitors.Removing the input-token embedding component also preserves first-layer cosine preference in all six models.Source: §2.3 · p. 4
  • Cosine competitor sets show both entries and exits over depths where their mean count decreases, whereas Euclidean sets are closer to nested for much of the trajectory.Entries include first-time entries and possible returns, which the statistic does not distinguish.Source: §2.3 · p. 4
  • Fewer competitors do not establish a tighter group: an exact fixed-bank construction shows that counts can fall while mean pairwise Euclidean and cosine distances increase.The authors therefore describe shrinkage as a change in competitor number, not as an empirical cohesion trend.Source: §2.3 · pp. 4–5

3 Simple Models of Residual Trajectories

The section models how residual movement and endpoint geometry jointly shape competition, then tests whether endpoint distributions predict competition along held-out trajectories. Its results distinguish competitor count, distance, and competitor identity as separate geometric measurements.

How an update changes endpoint preference

  • Residual updates can strengthen preference for some endpoints while weakening it for others, allowing competitors to enter as others leave.A negative margin identifies a competitor; crossing zero changes set membership.Source: §3 · p. 6

A straight-path baseline for turnover

  • Under either Euclidean or cosine distance, a monotone straight path toward the own endpoint produces nested competitor sets, so observed entries reject that baseline on affected trajectories.Cosine comparisons require nonzero states and endpoints.Source: §3 · p. 6, Figure 4 · p. 6

Why directional progress can coexist with flat distance

  • Increasing alignment can coexist with constant normalized Euclidean distance when norm changes offset directional progress.With b = ρ = 0 and r/R = 2a, normalized distance remains one as a increases within (0, 1/2).Source: §3 · pp. 6–7

From gradual alignment to sharp count reductions

  • For small positive alignment a, the expected competitor fraction can remain near one half, while sharp count reductions depend on alternative mass near the comparison boundary.The model relates the rate of count reduction to both alignment speed and the distribution of alternatives near the boundary.Source: §3 · p. 7

Predicting competition from endpoint geometry

  • On held-out trajectories with intermediate states and own endpoints fixed, projected-normal achieves lowest synthetic error in five of six models; vMF mixture is slightly better for Gemma-7B and ranks second elsewhere.Both outperform the cap and single vMF across all six architectures, while the uniform sphere has the largest error in five.Source: Figure 5 · p. 8, §3 · pp. 8–9

Fitting endpoint distributions

  • Endpoint distributions are fitted using final endpoints, then sampled banks are evaluated along held-out trajectories with intermediate states and own endpoints fixed.The model families use the same empirical endpoint-norm distribution, so prediction differences reflect fitted directional structure.Source: §3 · p. 8

4 Discussion and Conclusion

The discussion presents inference as increasing selectivity among contextual endpoints, while individual comparisons can be revised. The derivations explain how straight-path behavior, norm, and alignment constrain the interpretation of geometric measurements.

4 Discussion and Conclusion

  • Across six models, early average preference for the own endpoint coexists with many closer alternatives, and cosine competitors can enter as others exit.A straight path would produce nested sets under either metric, so observed entries depart from monotone straight-line convergence on affected trajectories.Source: §4 · p. 9

A Geometric Derivations

  • Endpoint-population geometry determines how residual movement translates into competition, linking final representations quantitatively to increasing selectivity across depth.The fitted distributions predict competition along observed trajectories but do not explain the origin or timing of residual updates.Source: §4 · p. 9

A.1 Endpoint margins and straight-line convergence

  • A resolved competitor entry excludes monotone straight-line convergence, but does not identify the update mechanism or distinguish curvature from backtracking.Ties remain outside the strict competitor set, and entries close to a tie depend on stored-state numerical precision.Source: §A.1 · p. 12

A.2 Norm and alignment during a distance plateau

  • Increasing alignment can be offset by norm changes, leaving normalized Euclidean distance constant while cosine distance decreases.The stated construction uses 0 < c < 1/2, so the intermediate-state norm remains below the endpoint norm.Source: §A.2 · pp. 12–13

A.3 High-dimensional competition with a shared component

The shared-axis model separates average preference from the number of remaining competitors and distinguishes trajectory-specific alignment from population spread. It also shows that norm and direction can jointly affect Euclidean competition.

Average preference and near-isolation

  • With M independent alternatives, expected competitor count is M F_D(a), where F_D(a) is the probability an alternative direction scores above alignment a.A small probability of any remaining competitor is ensured when M F_D(a) is much less than one, without requiring independence among alternatives for that bound.Source: §A.3 · p. 13
  • Any positive alignment advantage makes the own endpoint closer on average, yet a small advantage can coexist with a substantial competitor fraction.The mean-preference claim applies to cosine and squared Euclidean distance, not automatically to mean Euclidean distance.Source: §A.3 · p. 13

Population spread and own-endpoint alignment

  • Competitor counts can decline as own-endpoint alignment rises even while marginal directional population spread remains unchanged.The construction varies alignment while preserving the marginal directional population and common state norm across contexts.Source: §A.3 · p. 14

Unequal endpoint norms

  • With unequal endpoint norms, Euclidean competition depends on both norm and direction, whereas cosine competition is invariant to positive radial rescaling.The fitted endpoint models retain radial variation through a common empirical norm distribution.Source: §A.3 · p. 14

B A Shrinking Competitor Set Can Become Less Cohesive

A mathematical example separates shrinking competitor-set size from cohesion: fewer competitors need not be closer to one another. The experimental setup uses document splits to test endpoint-distribution predictions on held-out trajectories.

B A Shrinking Competitor Set Can Become Less Cohesive

  • The construction shows that nested competitor sets can lose an endpoint even as mean pairwise distance among the survivors increases.Equal endpoint norms make the example’s rankings identical under Euclidean and cosine distance.Source: B A Shrinking Competitor Set Can… · p. 14
  • The example’s competitor count falls from three to two while the surviving set becomes less cohesive under both Euclidean and cosine distance.It is a mathematical counterexample, not evidence of a cohesion trend in observed transformer trajectories.Source: B A Shrinking Competitor Set Can… · p. 14

C Experimental Setup and Endpoint-Distribution Fitting

  • The distribution-fitting experiment uses document splits to test whether final-state geometry predicts competition along held-out trajectories.The section specifies data, observables, uncertainty conventions, and endpoint-model protocol for the later analyses.Source: C Experimental Setup and… · p. 15

C.1 Language models, data, and residual trajectories

The section defines how intermediate residual states are compared with endpoint populations and how endpoint geometry is fitted and evaluated. Euclidean preference can coexist with large, architecture-dependent competitor sets.

C.1 Language models, data, and residual trajectories

  • For each model, residual states are measured at the final input position after blocks 1–L, before final normalization, across 1,024 non-overlapping 256-token FineWeb sample-10BT contexts; each context’s endpoint is its block-L state.There is no pre-block embedding state in the measured trajectory.Source: §C.1 · p. 15

C.2 Preference and competitor fractions

  • A competitor is an eligible endpoint closer to the current state than its own endpoint, while preference compares own-endpoint distance with the average alternative distance.Strict competitor counts vanish at the final state by construction; these statistics are not calibrated probabilities over output tokens.Source: §C.2 · p. 16

C.3 From endpoint counts to competitive mass

  • Competitive mass is the probability that an endpoint drawn from a population falls inside the region where it competes with the own endpoint.For the empirical endpoint distribution, this mass equals the competitor fraction; independence is needed for a binomial count law, but not for the expected count.Source: §C.3 · p. 16

C.4 Common radial law and directional models

  • The five fitted endpoint families share the fitting split’s empirical norm distribution, so their prediction differences arise from directional models.The shared norm-direction independence assumption may fail in the real population, and fitted axes have no established semantic interpretation.Source: §C.4 · pp. 16–17

C.5 Document-level fitting and model selection

  • Endpoint models are selected using document-level training and validation splits, with intermediate states, output labels, and competitor curves excluded from fitting and selection.The protocol tests endpoint structure’s ability to predict competitive mass without fitting observed competitor curves.Source: §C.5 · p. 17, Figure 6 · p. 18

C.6 Matched endpoint-distribution diagnostics

  • Model selection compares real and synthetic endpoint distributions using pairwise cosine similarity, Euclidean distance, nearest-neighbor spacing, and similarity to a fitted global axis.Lower average KS discrepancies indicate closer agreement; the scores are distributional discrepancies, not hypothesis-test p-values.Source: §C.6 · p. 18

C.7 Predicting competitor fractions from fitted endpoints

  • Fitted endpoint distributions are tested on held-out trajectories for predicting mean competitor fractions, with uncertainty intervals conditional on fitted distributions and sampled banks.The real-bank reference is not an error lower bound, and the intervals omit uncertainty from refitting endpoint models.Source: §C.7 · p. 19

D Euclidean Trajectories and Competitor Counts

  • In five models, Euclidean competitor counts remain in the hundreds across substantial depth before falling sharply, while Gemma-7B reaches a smaller set earlier.Mean distance gaps are small over much of depth in Gemma-2B, Qwen2.5-7B, and Llama-3-8B; timing depends on architecture and metric.Source: D Euclidean Trajectories and… · p. 19

E Early Own-Endpoint Preference

Average own-endpoint preference is detectable at the first measured layer, but this population-level pattern need not hold for every context. Controls show that it persists after excluding same-document alternatives and among contexts sharing a final input token.

E Early Own-Endpoint Preference

  • All six models meet the working p < 10−3 threshold for positive average preference at the first measured post-block layer under both metrics.The tests concern mean preference across documents, not positive preference for every context or a unique transition in computation.Source: E Early Own-Endpoint Preference · p. 20, Figure 8 · p. 21

F.1 Excluding alternatives from the query’s document

  • Excluding same-document alternatives changes cosine curves little and leaves first-layer average preference positive in every model.Euclidean effect sizes also change little, although the threshold crossing shifts by a few layers when early margins are small.Source: §F.1 · p. 21

F.2 Matching the final input token

  • Average preference remains positive from the first measured layer in all six models when alternatives share the final input token, though same-token margins are smaller.At the first layer, cosine fractions are 0.364 versus 0.058 for Qwen2.5-1.5B and 0.253 versus 0.059 for Gemma-7B, for same-token versus different-token banks.Source: §F.2 · p. 21

F.3 Removing the input-token embedding direction

The projection control removes the final input token’s embedding direction from each trajectory and its own endpoint. First-layer cosine preference remains positive in all six models, while Euclidean preference is more sensitive to the projection.

F.3 Removing the input-token embedding direction

  • The control linearly projects the query trajectory and own endpoint to remove the normalized embedding direction of the final input token.The projection preserves the comparison with the transformed own endpoint.Source: §F.3 · p. 21
  • After projection, first-layer cosine preference remains positive in all six models, with mean margins from 0.0159 in Gemma-2B to 0.1124 in Gemma-7B.The earliest positive Euclidean mean shifts to layer 5 in Qwen2.5-1.5B, layer 2 in Mistral-7B, and layer 13 in Qwen2.5-7B.Source: §F.3 · p. 21
  • The surviving cosine preference supports a broader-context contribution, but the projection removes only one direction, leaving token information represented elsewhere and residual-connection information possible.The Euclidean first-layer competitor fraction is near one half in five models.Source: §F.3 · p. 21

G Competitor Membership and Euclidean Turnover

Competitor sets can change through both entries and exits even as their mean size falls, while endpoint geometry’s links to output organization are strongest for cosine distance. The Euclidean rank comparison is less consistent, and output-related geometric associations do not make geometric competition equivalent to token probability.

G Competitor Membership and Euclidean Turnover

  • Cosine turnover persists as mean competitor counts decline, whereas Euclidean entries remain near zero for much of depth and late changes are dominated by exits.Cosine entries include first appearances and possible returns; the entry and exit fractions use different denominators, so their difference is not the count change.Source: G Competitor Membership and Euclidean… · p. 23, Figure 14 · p. 25
  • Same-input-token alternatives preserve early Euclidean preference, while removing embedding direction delays positive Euclidean preference in both Qwen models and Mistral-7B.After embedding-direction removal, cosine preference remains positive from the first measured layer in all six models.Source: Figure 11 · p. 23, Figure 10 · p. 23, Figure 13 · p. 24

H.1 Distance and output-distribution divergence

  • Endpoint distance is positively associated with directed output-distribution divergence, though substantial variation in divergence remains unexplained.For cosine distance, Pearson correlations range from 0.25 in Mistral-7B to 0.63 in Qwen2.5-1.5B.Source: H Endpoint Geometry and Output… · p. 24, §H.1 · p. 25

H.2 Compactness of endpoints sharing a top prediction

  • Endpoints sharing a top predicted token are more compact on average under both Euclidean and cosine distance, but do not form disjoint classes.The within-to-between distance ratio is below one in every model.Source: §H.2 · p. 25

H.3 Output-token rank and Euclidean endpoint distance

  • The Euclidean rank trend is less consistent than the cosine result, with positive-trend tests nonsignificant for Gemma-2B and Qwen2.5-1.5B.Rank-specific curves use only queries with a represented alternative endpoint for the relevant token; they do not assign endpoint distributions to every vocabulary token.Source: §H.3 · p. 25

I Population Geometry and Endpoint-Model Evaluation

The paper evaluates how endpoint population geometry relates to competitor fractions along held-out trajectories. Directional models improve on a uniform-sphere baseline for cosine prediction, while family rankings vary for Euclidean prediction and cosine patterns remain stable under bank subsampling.

I.1 Layerwise population geometry

  • Population geometry is described using normalized mean pairwise Euclidean and mean pairwise cosine distances, but these means do not determine modes or competitor-set cohesion.The Euclidean reference of 2 assumes equal-norm orthogonal states; the bootstrap bands hold the original normalizer fixed.Source: §I.1 · pp. 26–27, Figure 18 · p. 28

I.2 Complexity selection and held-out endpoint fit

  • Several validation searches reach their tested complexity boundary, leaving open whether higher complexity would improve those fits.All reported configurations were fixed before test evaluation.Source: §I.2 · p. 28, Figure 19 · p. 28
  • The best reported held-out endpoint diagnostic scores improve on the sphere by factors of 4.6–7.5, but endpoint fit does not determine trajectory-prediction error.For Qwen2.5-1.5B, projected normal and the mixture are nearly tied at 0.151 and 0.153.Source: §I.2 · p. 28, Table 2 · p. 29

I.3 Cosine competitor-fraction prediction

  • Cosine competitor-fraction prediction favors directional structure: projected normal has the lowest synthetic point estimate in five models, while the vMF mixture is slightly better for Gemma-7B.Both families outperform the cap, single vMF, and sphere in all six models; the results do not establish one family as uniformly superior.Source: §I.3 · p. 29, Table 3 · p. 29

J Sensitivity to Endpoint-Bank Size

  • Cosine competitor-fraction curves overlap closely across bank sizes, and the main depth patterns persist across samples.The largest visible differences are below 0.01 and within the sampling range of the smallest bank; the check covers cosine distance only.Source: J Sensitivity to Endpoint-Bank Size · p. 29, Figure 20 · p. 30

K.1 Euclidean competitor-fraction prediction

  • Euclidean competitor-fraction errors show no consistently preferred fitted family, with the smallest errors varying by architecture and several ties at reported precision.Euclidean membership depends jointly on endpoint norms and directions, and all fitted families assume independence between them.Source: §K.1 · p. 30, Table 4 · p. 30

K.2 Cosine uncertainty and paired comparisons

Paired cosine-error comparisons favor projected normal over the vMF mixture for Gemma-2B and both Qwen models, but not for the other three models. Reproduction is also limited by unavailable analysis materials, and the authors report language-model use in preparing the manuscript.

K.2 Cosine uncertainty and paired comparisons

  • Paired comparisons favor projected normal over the vMF mixture for Gemma-2B and both Qwen models, but do not distinguish the families for Gemma-7B, Mistral-7B, or Llama-3-8B.The paired analysis uses the same query-document weights, but leaves endpoint-model fitting and new bank-draw uncertainty outside the comparison.Source: §K.2 · pp. 30–31, Table 5 · p. 31
  • The reported uncertainty intervals concern the real-bank reference, not differences between fitted families.Numerical intervals for paired error differences are unavailable in the accompanying source materials.Source: Table 5 · p. 31, §K.2 · p. 31

L Reproducibility and Source Materials

  • Independent reproduction is constrained because extraction and analysis code, residual arrays, fitted models, and random seeds are not included.The manuscript and figure PDFs are present, but the missing materials are needed to reproduce numerical results and resolve implementation details.Source: L Reproducibility and Source Materials · pp. 31–32

AI use statement

  • The authors report using language-model tools for language editing, restructuring, mathematical derivation, and source debugging.They state responsibility for the mathematical statements, experimental design, data analysis, citations, and final text.Source: AI use statement · p. 32
Loading 2609.37824v1…