Source-linked AI summary

Global Divergence, Local Convergence: Representation Geometry in SSMs and Transformers

Amit Ben-Artzy, Roy Schwartz

arXiv:2609.08692v1cs.CL

TL;DR

The paper asks how transformer and SSM architectural differences affect internal representation geometry and function. It compares residual-stream representations across transformers, SSMs, and hybrids using geometric, compressibility, probing, and local-alignment analyses. The architectures diverge globally but converge in effective capacity, concept subspaces, and local semantic structure.

  • Problem

    The paper asks whether fundamentally different transformer and SSM architectures induce fundamentally different internal representations.

  • Method

    The study performs a multi-scale residual-stream analysis of transformers, SSMs, and hybrid architectures using geometric metrics, autoencoders, constrained probes, PCA ablations, CKA, and m-KNN.

  • Results

    Transformers and SSMs differ globally in latent-space geometry but show tightly matched effective capacity, similar concept-subspace dimensionality, and highly aligned local structures.

  • Takeaways & Limitations

    The architectures use latent space differently while converging functionally at the level of local semantic manifolds.

  • Takeaways & Limitations

    The probing and alignment analyses do not establish a direct causal link between concept encoding and downstream generation.

Abstract

from arXiv · show

Recent state-space models (SSMs) such as Mamba achieve language modeling performance comparable to transformers despite relying on fundamentally different architectures. This raises an important question: how do these structural differences influence the geometry and functional nature of their internal representations? We study this question through a multi-scale analysis of representations in transformers, SSMs, and hybrid architecture. First, we find that SSMs distribute their representational information evenly across all dimensions, whereas transformer representations are heavily dominated by a single principal direction. By evaluating hybrid architectures, we observe that the representation space becomes increasingly skewed toward a single dominant direction after each attention layer. Next, we explore how the different geometric spread of representations impacts representational capacity through compressibility. Surprisingly, we find that despite their contrasting geometric structures, both architectures exhibit tightly matched effective capacities. We further investigate whether this skewed geometry affects how concepts are encoded. Using rank-constrained probes, we demonstrate that both architectures encode concepts in subspaces of surprisingly similar dimensionality. Furthermore, we demonstrate that the transformers' dominant principal direction does not inherently encode more conceptual information. Finally, we zoom in and examine the alignment between manifolds, either by analyzing representations of specific topics or by looking at the nearest neighborhoods of tokens, and find that they are highly aligned. Ultimately, our analysis suggests that while transformers and SSMs induce different usage of latent space, they display a striking functional convergence at the level of local semantic manifolds.

1 Introduction

The paper asks whether fundamentally different transformer and SSM architectures produce different internal representations, using a multi-scale residual-stream analysis. It finds divergent global geometry but convergent local structure and concept-level function.

  • Introduction: The study compares transformer, SSM, and hybrid residual-stream representations across multiple geometric and functional scales.The models include Pythia, Falcon, Mamba, FalconMamba, and Jamba2.
  • Introduction: SSMs have isotropic, high-dimensional representations, whereas transformers have anisotropic representations dominated by a principal direction.Hybrid layers alternate between these geometric regimes.
  • Introduction: Both architectures show similar effective capacity and concept-decoding dimensionality despite their different global geometries.Transformer dominant directions do not capture more conceptual information than SSM representations.
  • Introduction: Local structural properties converge within the first 4% of training, while global geometry stabilizes later.This dynamic is also observed in hybrid architectures that oscillate between isotropic and anisotropic states.
  • Introduction: Local concept manifolds and token neighborhoods are highly aligned across architectures despite global geometric divergence.The paper uses structural alignment analyses to establish this local convergence.

2 Background: Transformer and SSM Architectures

The paper frames each architectural layer as a direct update to a shared residual stream, enabling structural comparison between transformers and SSMs. Transformers use attention, whereas Mamba replaces attention with a selective state-space mechanism and recurrent sequence state.

  • Shared residual-stream view: Each architectural layer is treated as a direct update to the residual stream.This shared representation framework supports comparison across model families.
  • Transformer architecture: A transformer block interleaves attention and feed-forward blocks while updating hidden state x_l through the residual stream.The supplied passage introduces the transformer update formulation.
  • SSM architecture: Mamba replaces attention with a selective state-space model that updates the residual stream.Its internal block includes projection, depthwise convolution, selective SSM mixing, and output projection.
  • SSM architecture: Mamba maintains a recurrent hidden state along the sequence dimension rather than across layers.The analysis does not examine this internal sequence-level state.
  • Compared models: The experiments compare two transformer families, two SSM families, and the hybrid Jamba2 model using shared tokenization.Jamba2 interleaves Mamba and transformer blocks.

3 Global Geometry of Representations

SSMs use a more isotropic and higher-dimensional global representation space than transformers, whose variance is concentrated in dominant directions. In Jamba2, attention drives transformer-like anisotropy while later Mamba layers reverse it.

  • Metrics: RankMe measures effective dimensionality, IsoScore measures isotropy, and PC1 explained variance indicates dominance of the first principal direction.IsoScore equals 1 when representations use all available dimensions equally.
  • Experimental setup: ≈500K FineWeb-Edu representations are collected from all layers for the geometric analysis.The evaluated models are Pythia-1.4B, Falcon-7B, Mamba-1.4B, FalconMamba-7B, and Jamba2-3B.
  • Global geometry: SSM models sustain higher RankMe and more uniform variance than same-size transformer counterparts across nearly all layers.The exception is the final layers.
  • Global geometry: Transformers show lower IsoScores and higher PC1 explained variance, indicating anisotropic, cone-like representations dominated by a few principal directions.These metrics jointly distinguish transformer and SSM latent-space organization.
  • Hybrid architecture: In Jamba2, attention layers reduce RankMe and increase IsoScore and PC1 explained variance, while subsequent Mamba layers reverse the trend.The hybrid representation therefore oscillates between transformer-like and SSM-like geometry.

4 Quantifying the Effective Dimensionality of Representations

Despite sharply different global geometry, transformers and SSMs have tightly matched effective representational capacity and require similarly dimensional subspaces for relational concept decoding. Transformer dominant directions also do not provide a clear semantic advantage.

  • Motivation: The analysis asks whether different global geometric spread changes representational capacity or concentrates concept information along transformer dominant axes.This motivates compressibility tests and rank- and PCA-constrained probing.
  • Compressibility: Two-layer autoencoders with rank-constrained bottlenecks measure capacity through reconstruction loss and KL divergence from the original output distribution.Each layer is evaluated across bottleneck ranks.
  • Compressibility: Both architectures show highly similar KL degradation across bottleneck ranks, except for the more sensitive final SSM layer.This indicates tightly matched effective representational capacity despite geometric differences.
  • Concept subspaces: Identical rank constraints produce closely matched decoding accuracies across architectures and layers for LAMA-TREx relational facts.The facts occupy subspaces of similar required rank.
  • Spectral ablation: Removing or retaining top PCA directions causes generally similar probing degradation across architectures.Transformer dominant directions therefore do not serve as the primary axes for factual knowledge.

5 Alignment of Semantic Concepts and Local Neighborhoods

Although transformers and SSMs differ globally, their semantic sub-manifolds and token neighborhoods are structurally aligned. ZCA-CKA controls for global anisotropy, while m-KNN compares local neighborhoods through relative distances.

  • Concept manifolds: ZCA-CKA compares whitened inner-product structures to measure alignment between representation spaces.Global whitening uses the inverse square root of covariance estimated from matched FineWeb-Edu tokens.
  • Concept manifolds: Global whitening lets CKA reveal higher similarity on individual semantic relations than across the general corpus.Relation-wise whitening is avoided because it would remove relation-specific covariance structure.
  • Local neighborhoods: At the token level, m-KNN measures overlap between cosine-similarity neighborhoods of identical tokens across models.The analysis uses 10,000 FineWeb-Edu token samples and k = 10, and is invariant to global linear transformations and scaling.
  • Local neighborhoods: Fully trained transformers and SSMs maintain high m-KNN overlap at corresponding depths despite different architectures.Their local topologies evolve through the residual stream in parallel.

6 Temporal Dynamics of Alignment

Local representational alignment emerges early during training, whereas global geometry stabilizes substantially later. Pythia matches Mamba’s local structures after roughly 4% of training, while global metrics settle near 25%.

  • Early local convergence: 5,000 training steps (∼4% of training) are sufficient for Pythia’s representations to match Mamba’s local CKA and m-KNN structures.Cross-architecture and cross-scale similarity is already high at this checkpoint.
  • Scope: The temporal analysis is anchored on Pythia because checkpoints for the other models are unavailable.
  • Early local convergence: Local similarity saturates early in both general and concept-specific contexts, with absolute similarity higher for concepts.
  • Global geometry: Approximately 25% of training is required before IsoScore and PC1 Explained Variance stabilize.RankMe instead fluctuates nonmonotonically, first increasing and then decreasing.

7 Related Work

Related work frames representational convergence as a recurring phenomenon across architectures, scales, and granularities, while also emphasizing anisotropy and the evolution of latent geometry.

  • Representational convergence: Prior work reports that independently trained models can share semantic subspaces while distributing information across different basis vectors.
  • Cross-architecture comparisons: Recent studies find similar interpretable features, numerical-token geometries, and tokenwise transformations in Mamba and transformer models.
  • Geometry and anisotropy: Transformer self-attention is associated with anisotropic representations occupying a narrow latent-space cone.Prior work also links isotropy measures with language-model efficacy and identifies low-dimensional attribute subspaces.
  • Dynamic evolution: Representational geometry changes during training and across layers, and prior work has examined token similarity across contexts and architectures.

8 Conclusion

The paper finds a sharp global geometric divide between SSMs and transformers but strong local functional and structural convergence. Local alignment is high for semantic relations and neighborhoods, and it appears before global geometry stabilizes.

  • Global geometry: SSM representations are highly isotropic, while transformer representations are heavily anisotropic; hybrid models oscillate between these states across layers.
  • Functional capacity: Both architectures have remarkably similar effective dimensionality, and the transformer’s dominant direction does not inherently encode more conceptual information.
  • Local convergence: Structural alignment is high for individual semantic relations but low across the general corpus.
  • Conclusion: Transformers and SSMs converge on highly aligned local representational structures despite organizing latent space differently.

Limitations

The supplementary analyses establish robustness across seeds, model families, constraints, and sampling, while the paper’s main scope remains bounded by correlational evidence for concept encoding.

  • Experimental scope: The experiments use bounded samples and training settings, including approximately 500K tokens for geometric metrics and 10,000 samples for autoencoder training.Additional analyses use fixed sample sizes for probing, ZCA-CKA, and m-KNN.
  • Training dynamics: RankMe, IsoScore, and PC1 explained variance follow distinct training trajectories: RankMe is nonmonotonic, while the other metrics converge monotonically.Stabilization occurs primarily within the first 25% of training steps.
  • Cross-model robustness: Across additional model families, SSMs retain higher RankMe than transformers at matched scale, while IsoScore and PC1 explained variance remain largely scale-invariant.These results extend the reported geometric contrast beyond the primary model families.
  • Variance analysis: The geometric metrics were re-evaluated across five random seeds and subsets of approximately 500K tokens to assess stability.Table 1 reports mean and maximum layer-wise coefficient of variation across seeds.

E Additional CKA Results

Additional CKA analyses test whether the reported local/global alignment gap is robust to spectral mismatch and sampling choices. The results preserve stronger alignment for LAMA TREx than FineWeb-Edu without whitening and show stable estimates.

  • Whitening analysis: ZCA-whitening is used to standardize transformer and SSM representations because their differing variance profiles can distort linear CKA comparisons.The analysis varies whitening strength α to test whether alignment differences reflect representation structure rather than spectral mismatch.
  • Evaluation design: ZCA-CKA is evaluated across LAMA TREx and FineWeb-Edu, with LAMA TREx sampled by one context sequence per unique factual triplet.The sampling controls for redundant semantic structures in the factual dataset.
  • Whitening analysis: Without whitening (α = 0), LAMA TREx representations have higher layer-mean CKA than FineWeb-Edu representations, 0.78 versus 0.52.This supports the conclusion that the alignment gap is not an artifact of ZCA whitening.
  • Sampling stability: For FineWeb-Edu, 2.5K-sample subsets yield mean CV% 0.15% and maximum CV% 1.34% across layer pairs.These values are reported for Mamba-1.4B versus Pythia-1.4B across varying seeds.
  • Sampling stability: For LAMA TREx, the corresponding mean CV% is 0.56% and maximum CV% is 3.31%, indicating stable estimates across seeds.The reported variance remains low despite the different dataset structure.

F Additional Results for m-KNN

Additional m-KNN analyses show that local neighborhood topology generalizes to an unseen architecture pair and remains stable under repeated sampling.

  • Generalization: Falcon-7B and FalconMamba-7B exhibit highly aligned local neighborhood topology at corresponding relative depths.This reproduces the earlier trend outside the original Mamba/Pythia comparison.
  • Sampling stability: m-KNN robustness is evaluated through five independently drawn batches of 2,000 samples for Mamba and Pythia models.The analysis covers both 1.4B and 2.8B model sizes.
  • Sampling stability: The mean SEM across layer-pair comparisons is 0.0047, with a range of 0.0003–0.0113.These estimates are computed across five batches.
  • Sampling stability: The mean standard deviation across layer-pair comparisons is 0.0104, ranging from 0.0008 to 0.0253.The reported variation supports stable m-KNN estimates under resampling.
Loading 2609.08692v1…