A Comparative Analysis of Contextual Representation Flow in State-Space and Transformer Architectures

TL;DR

This work provides a unified, token- and layer-level comparison of how contextual representations flow in long-context Transformer-Based Models (TBMs) and State Space Models (SSMs). By combining empirical analyses (token- and layer-level dynamics, probing, and context-length effects) with theoretical stability results, it shows that TBMs rapidly homogenize token representations early and undergo a final-layer reconfiguration, while SSMs preserve token uniqueness early and homogenize deeper due to training dynamics. The study demonstrates that intermediate layers often contain the most task-relevant information and that oversmoothing is architecture-driven in TBMs but largely training-driven in SSMs, with implications for designing hybrids and regularizers. These findings offer practical diagnostics and interventions—such as intermediate supervision and contractive constraints—and establish a framework for robust long-context modeling and architecture-aware scaling. The results advance understanding of inductive biases in long-context models and provide a toolkit for systematic evaluation across architectures.

Abstract

State Space Models (SSMs) have recently emerged as efficient alternatives to Transformer-Based Models (TBMs) for long-sequence processing, offering linear scaling and lower memory use. Yet, how contextual information flows across layers and tokens in these architectures remains understudied. We present the first unified, token- and layer-level analysis of representation propagation in SSMs and TBMs. Using centered kernel alignment, stability metrics, and probing, we characterize how representations evolve within and across layers. We find a key divergence: TBMs rapidly homogenize token representations, with diversity reemerging only in later layers, while SSMs preserve token uniqueness early but converge to homogenization deeper. Theoretical analysis and parameter randomization further reveal that oversmoothing in TBMs stems from architectural design, whereas in SSMs it arises mainly from training dynamics. These insights clarify the inductive biases of both architectures and inform future model and training designs for long-context reasoning.

Figures (7)

• Figure 1: Token-level cosine similarity evolution across layers. We observe that TBMs (left column) maintain consistently high layerwise similarity until a sharp shift near the final layers, indicating stable token evolution followed by more abstract refinement. In contrast, SSMs (right column) reveal greater variability and exploratory changes in early layers, with gradual convergence later, highlighting distinct representation dynamics per architecture.
• Figure 2: Inter-token cosine similarity across layers for different pre-trained models. This figure shows average similarity among tokens in each layer, measuring token distinction and homogenization. TBMs rapidly increase token similarity early and maintain high homogenization until a late drop, reflecting oversmoothing of features. Conversely, SSMs sustain greater token diversity through most layers, with only a late-stage increase in similarity, emphasizing their tendency to preserve unique token features deeper in the network.
• Figure 3: Token-level cosine similarity across layers for both pre-trained and randomly initialized models. Comparing pretrained and random initializations reveals an architectural bias towards oversmoothing in TBM (left column) where high similarity exists regardless of training, while SSM (right column) exhibits near-zero similarity at random initialization. This distinction confirms that oversmoothing is intrinsic to TBM architecture but training-dependent in SSM.
• Figure 4: layer-level CKA similarity for every layer pairs, averaged over the MDQA and KVPR tasks with $n = 2K$ tokens. We find that TBMs (left column) exhibit stable alignment across the initial layers (except for the first two layers of Pythia-2.8B), followed by a representation shift towards the final layers, with lower similarity between the early and last layers. In contrast, SSMs (right column) show significant fluctuation in the lower layers, followed by a more consistent alignment in deeper layers, indicating a gradual stabilization of feature representations.
• Figure 5: Token-wise cosine similarity across layers for GPT-Neo-2.7B (left) and Mamba2-2.7B (right) on the KVPR task with $n = 2K$ tokens.

Theorems & Definitions (22)

• Proposition 1
• Proposition 2
• Proposition 3
• Theorem 1
• Theorem 2
• Theorem 3
• Theorem 4
• Theorem 5
• proof : Proof of Proposition \ref{['prop:prop-common-formula']}
• proof : Proof of \ref{['prop:prop-att-odd']}