Contextual Representation Flow Across Depth in State-Space and Transformer Architectures

Nhat M. Hoang; Do Xuan Long; Cong-Duy Nguyen; Chunyan Miao; Min-Yen Kan; Luu Anh Tuan

cs.CLOct 8, 2025

Contextual Representation Flow Across Depth in State-Space and Transformer Architectures

Authors: Nhat M. Hoang, Do Xuan Long, Cong-Duy Nguyen, Chunyan Miao, Min-Yen Kan, Luu Anh Tuan

Organizations: Nanyang Technological University, Singapore · National University of Singapore

Abstract

State Space Models (SSMs) have emerged as efficient, linear-time alternatives to Transformer-Based Models (TBMs) for long-context processing. However, their different mechanisms raise an important question: how does contextual representation flow across depth differ between these model families and their hybrids? We examine this question across twenty-one pretrained models, combining token-wise spectral diversity, layer-wise variation on seven datasets, and answer-location probing on two retrieval tasks. This joint analysis reveals that TBMs retain most of their token diversity at the output, whereas SSMs reach richer intermediate diversity but retain less of it on average, with late-layer compression. Comparisons with reinitialized weights reveal that these profiles depend on pretrained weights, while hybrids combine strong answer-location accessibility with architecture-dependent geometry, showing that high accessibility need not require the steadiest representation flow. Finally, we conduct theoretical case studies connecting these contrasts to mechanism-level differences: query-time token mixing in TBMs, repeated recurrent updates in SSMs, and preserved attention updates in hybrids.

Figures & tables

Figure 1: Token geometry of all twenty-one models, averaging seven datasets at each native layer: (a) pretrained InterSim, (b) pretrained normalized rank, and (c) Gaussian-reinitialized rank.

	Rank $\uparrow$		Final rank $\uparrow$
Model	Peak	Ret. $\uparrow$	Pre.	Gau.	Xav.	Kai.
GPT-Neo-1.3B	0.329	96.5	0.321	0.214	0.086	0.100
GPT-Neo-2.7B	0.383	90.1	0.349	0.176	0.081	0.092
Pythia-1.4B	0.470	93.3	0.444	0.213	0.123	0.162
Pythia-2.8B	0.525	91.4	0.486	0.231	0.130	0.159
Llama-3.1-8B-Instruct	0.576	89.5	0.527	0.012	0.007	0.010

Table 1: Token diversity averaged over all seven datasets: peak normalized rank, rank retention (% of peak), final rank before and after reinitialization. Peaks and ratios are computed within each dataset before averaging; Mean rows then average models within families.

$L$	Family	Model	$V_{\mathrm{loc}}$	$V_{\mathrm{glob}}$	$\rho\downarrow$
24	TBM	GPT-Neo-1.3B	1.01	3.12	0.324
24	TBM	Pythia-1.4B	0.28	0.96	0.289
24	SSM	RWKV-7-1.5B	0.67	2.61	0.257
24	SSM	DeltaNet-1.3B	1.11	4.56	0.244
24	SSM	Mamba-3 SISO-1.5B	2.84	11.33	0.247
24	SSM	Mamba-3 MIMO-1.5B	0.53	2.17	0.245

Table 2: Models ordered by native depth

L

, then family. All measures include embedding and final normalization;

\rho

averages per-input ratios, then seven datasets equally. Family means span all listed depths.

Panel of Table 3 — Table 3: Location-probe accuracy (%, mean and $95\%$ interval, five seeds; $\uparrow$ ), peak-to-final drop ( $\Delta\downarrow$ ), and Spearman rank correlations ( $S_{\mathrm{rank}}$ ; green/red: positive/negative). Qwen2.5-32B is excluded from probing family means.

Figure 2: Peak-layer probing accuracy by item position: family means and model ranges exclude Qwen2.5-32B. MDQA targets occupy positions

1,3,\ldots,13

Appendix figures & tables14 assets

Supplementary material from the paper’s appendix.

Appendix

Model	KVPR	MDQA
GPT-Neo-1.3B	9.3	14.7
GPT-Neo-2.7B	9.5	14.7
Pythia-1.4B	9.5	14.7
Pythia-2.8B	9.6	14.9
Llama-3.1-8B-Instruct	9.5	14.9
Qwen2.5-32B Base	9.6	14.8

Appendix

Table 4: Permuted-label control accuracy (%). Each entry is the maximum over layers of the five-seed mean, not the control accuracy at the unpermuted probe peak. Chance is

9.1\%

for KVPR and

14.3\%

for thirteen-document MDQA.

	KVPR		MDQA
Model	$S_{\mathrm{rank}}$	$S_{\mathrm{cos}}$	$S_{\mathrm{rank}}$	$S_{\mathrm{cos}}$
GPT-Neo-1.3B	$+0.54$	$-0.14$	$-0.09$	$-0.21$
GPT-Neo-2.7B	$+0.56$	$+0.07$	$+0.16$	$+0.07$
Pythia-1.4B	$+0.64$	$-0.76$	$+0.62$	$-0.33$
Pythia-2.8B	$+0.85$	$-0.82$	$+0.81$	$-0.53$
Llama-3.1-8B-Instruct	$+0.31$	$-0.37$	$+0.41$	$+0.27$

Appendix

Table 5: Within-model, across-layer Spearman correlations of location-probe accuracy with rank (

S_{\mathrm{rank}}

) and InterSim (

S_{\mathrm{cos}}

). Geometry and probing use separate inputs.

Task	Model	Rank ret. (%)
KVPR	Mamba2-2.7B	10.0
	Mamba2Attn-2.7B	11.7
MDQA	Mamba2-2.7B	50.2
	Mamba2Attn-2.7B	42.8

Appendix

Table 6: Task-specific Mamba2 rank retention (%) for Table 3 ’s probe comparison. Geometry and probing use separate inputs.

Model	Rank retention (%)	InterSim final/max (%)
GPT-Neo-1.3B	96.5	77.6
GPT-Neo-2.7B	90.1	74.5
Pythia-1.4B	93.3	75.8
Pythia-2.8B	91.4	72.7
Llama-3.1-8B-Instruct	89.5	86.9
Qwen2.5-32B Base	63.2	87.3

Appendix

Table 7: Rank retention and final-to-maximum InterSim (%), computed within each dataset, then averaged equally over seven datasets. Mean rows average models within families.

	TBM			SSM			Hybrid
Dataset	Peak	Ret.	Gau.	Peak	Ret.	Gau.	Peak	Ret.	Gau.
KVPR	0.332	61.8	0.146	0.374	24.1	0.707	0.351	33.6	0.680
MDQA	0.568	94.2	0.273	0.562	67.0	0.737	0.551	66.1	0.705
AG News	0.436	93.0	0.235	0.444	62.0	0.630	0.454	58.9	0.655
SAMSum	0.477	88.2	0.232	0.487	55.8	0.703	0.506	55.0	0.690
CNN/DailyMail	0.546	92.8	0.267	0.561	62.6	0.732	0.558	62.5	0.707

Appendix

Table 8: Dataset-wise family means over six TBMs, nine SSMs and six hybrids: pretrained peak normalized rank, rank retention (% of peak), and Gaussian-reinitialized final rank.

Figure 3: Dataset-wise representation flow on seven input sets (rows): pretrained InterSim, pretrained normalized effective rank, and Gaussian-reinitialized rank (columns). All twenty-one models appear in every row, with shared column-wise axes and the colors and line styles of Figure 1 .

	Rank		$E({\bm{H}})$	InterSim	Final rank				Variation
Model	Peak	Ret.	Ret.	F/M	Pre.	Gau.	Xav.	Kai.	$V_{\mathrm{loc}}$	$V_{\mathrm{glob}}$
GPT-Neo-1.3B	.294	100.0	100.0	84.0	.294	.213	.087	.097	0.97	2.98
GPT-Neo-2.7B	.341	94.2	100.0	77.5	.321	.173	.084	.094	0.84	3.12
Pythia-1.4B	.424	100.0	100.0	80.2	.424	.210	.121	.166	0.30	1.05
Pythia-2.8B	.459	100.0	100.0	76.1	.459	.221	.123	.162	0.28	1.21
Llama-3.1-8B-Instruct	.568	95.3	81.9	81.9	.541	.012	.007	.010	0.07	0.36

Appendix

Table 9: AG News geometry: peak rank, rank and sharpness retention (%), InterSim final-to-maximum ratio (%), final rank before and after reinitialization, and native-depth local and global variation.

	Rank		$E({\bm{H}})$	InterSim	Final rank				Variation
Model	Peak	Ret.	Ret.	F/M	Pre.	Gau.	Xav.	Kai.	$V_{\mathrm{loc}}$	$V_{\mathrm{glob}}$
GPT-Neo-1.3B	.315	98.9	100.0	88.1	.311	.215	.086	.096	0.98	3.02
GPT-Neo-2.7B	.389	86.4	100.0	85.5	.336	.172	.084	.092	0.85	3.23
Pythia-1.4B	.464	94.7	88.0	89.6	.439	.198	.113	.158	0.28	1.01
Pythia-2.8B	.521	92.3	100.0	87.5	.481	.209	.116	.156	0.27	1.15
Llama-3.1-8B-Instruct	.605	93.8	74.1	89.8	.567	.012	.007	.010	0.07	0.35

Appendix

Table 10: SAMSum geometry: peak rank, rank and sharpness retention (%), InterSim final-to-maximum ratio (%), final rank before and after reinitialization, and native-depth local and global variation.

	Rank		$E({\bm{H}})$	InterSim	Final rank				Variation
Model	Peak	Ret.	Ret.	F/M	Pre.	Gau.	Xav.	Kai.	$V_{\mathrm{loc}}$	$V_{\mathrm{glob}}$
GPT-Neo-1.3B	.397	100.0	100.0	85.3	.397	.237	.099	.111	0.99	3.08
GPT-Neo-2.7B	.459	93.1	100.0	80.9	.427	.192	.092	.099	0.86	3.29
Pythia-1.4B	.535	100.0	100.0	81.1	.535	.252	.148	.189	0.26	0.91
Pythia-2.8B	.589	99.1	100.0	78.0	.584	.267	.150	.181	0.26	1.07
Llama-3.1-8B-Instruct	.652	95.7	77.9	85.8	.624	.012	.007	.009	0.07	0.36

Appendix

Table 11: CNN/DailyMail geometry: peak rank, rank and sharpness retention (%), InterSim final-to-maximum ratio (%), final rank before and after reinitialization, and native-depth local and global variation.

	Rank		$E({\bm{H}})$	InterSim	Final rank				Variation
Model	Peak	Ret.	Ret.	F/M	Pre.	Gau.	Xav.	Kai.	$V_{\mathrm{loc}}$	$V_{\mathrm{glob}}$
GPT-Neo-1.3B	.320	100.0	100.0	48.6	.320	.215	.088	.101	0.97	2.95
GPT-Neo-2.7B	.386	94.1	100.0	37.8	.363	.173	.085	.095	0.83	3.14
Pythia-1.4B	.480	95.1	100.0	46.0	.457	.205	.120	.165	0.28	1.00
Pythia-2.8B	.548	92.5	100.0	42.1	.506	.217	.123	.160	0.27	1.17
Llama-3.1-8B-Instruct	.595	93.0	69.8	80.5	.553	.012	.007	.009	0.07	0.35

Appendix

Table 12: Qasper geometry: peak rank, rank and sharpness retention (%), InterSim final-to-maximum ratio (%), final rank before and after reinitialization, and native-depth local and global variation.

	Rank		$E({\bm{H}})$	InterSim	Final rank				Variation
Model	Peak	Ret.	Ret.	F/M	Pre.	Gau.	Xav.	Kai.	$V_{\mathrm{loc}}$	$V_{\mathrm{glob}}$
GPT-Neo-1.3B	.341	100.0	86.3	91.4	.341	.218	.089	.104	0.96	2.99
GPT-Neo-2.7B	.392	96.5	100.0	87.6	.378	.173	.083	.092	0.84	3.22
Pythia-1.4B	.482	98.8	90.2	86.9	.477	.211	.123	.167	0.26	0.93
Pythia-2.8B	.548	96.4	100.0	84.1	.528	.223	.123	.155	0.26	1.09
Llama-3.1-8B-Instruct	.589	94.9	64.1	94.4	.559	.012	.007	.009	0.07	0.36

Appendix

Table 13: GovReport geometry: peak rank, rank and sharpness retention (%), InterSim final-to-maximum ratio (%), final rank before and after reinitialization, and native-depth local and global variation.

		Final sigma-linear rank		Change in Mamba–Pythia gap (%)
Task	Initializer	Pythia keep/reinit.	Mamba keep/reinit.	Sigma-linear	Energy
KVPR	Gaussian	0.112/0.136	0.477/0.629	+35.1	+227.7
KVPR	Xavier	0.049/0.073	0.321/0.571	+83.1	+369.4
KVPR	Kaiming	0.103/0.139	0.672/0.707	-0.1	+26.1
MDQA	Gaussian	0.274/0.299	0.619/0.682	+10.9	+89.2
MDQA	Xavier	0.151/0.184	0.523/0.628	+19.5	+132.0

Appendix

Table 14: Matched embedding effects on raw normalized final rank. Gap changes compare both-reinitialized with both-retained embeddings on the same reinitialized backbones, divided by the both-retained gap. Values are descriptive means over the three seeds and fixed inputs.

Model	Pooling	KVPR	MDQA
GPT-Neo-1.3B	last	$62.1{\hlsecondarytab{($ \downarrow $27.6)}}$	$29.3{\hlsecondarytab{($ \downarrow $35.7)}}$
GPT-Neo-1.3B	mean	$54.8{\hlsecondarytab{($ \downarrow $14.4)}}$	$26.5{\hlsecondarytab{($ \downarrow $10.9)}}$
GPT-Neo-2.7B	last	$61.1{\hlsecondarytab{($ \downarrow $21.0)}}$	$38.9{\hlsecondarytab{($ \downarrow $34.5)}}$
GPT-Neo-2.7B	mean	$57.8{\hlsecondarytab{($ \downarrow $7.6)}}$	$23.5{\hlsecondarytab{($ \downarrow $15.9)}}$
Pythia-1.4B	last	$82.8{\hlsecondarytab{($ \downarrow $16.5)}}$	$39.4{\hlsecondarytab{($ \downarrow $23.5)}}$
Pythia-1.4B	mean	$71.8{\hlsecondarytab{($ \downarrow $15.7)}}$	$26.5{\hlsecondarytab{($ \downarrow $7.2)}}$

Appendix

Table 15: Final-layer location-probe accuracy (%), averaged over five seeds, with last-token or mean pooling;

\downarrow

: peak-to-final gap in percentage points.

Figure 4: Numerical illustration of diversity loss through repeated Mamba2 updates. A synthetic two-token input starts with coefficients

\bm{\xi}=(1,-1)^{\top}

and passes through 1,500 copies of the complete block in Eq. 29 , at three update scales

\eta

. (a) Effective rank, defined in Eq. 1 , approaches one as the token-specific component disappears. (b) The raw update-change norm

\|{\bm{H}}_{\mathrm{raw}}^{(l+2)}-2{\bm{H}}_{\mathrm{raw}}^{(l+1)}+{\bm{H}}_{\mathrm{raw}}^{(l)}\|_{F}

is shown on a logarithmic scale. This raw update-change norm measures the unnormalized second difference; the empirical ratio

\rho

additionally divides coordinate-averaged local variation by global depth dispersion. All curves use the same block count, so changing

\eta

also changes the accumulated update amount. These are computed trajectories of the construction; Theorem 1 guarantees convergence locally.

State Space Models

State space models represent sequences through evolving latent states, enabling efficient temporal modeling, system identification, and forecasting.

See all 297 papers

Global Divergence, Local Convergence: Representation Geometry in SSMs and Transformers
Sep 8, 2026Amit Ben-Artzy, Roy SchwartzTransformer Representation Geometry in Language Models
Computer Science, The Hebrew University of Jerusalem
What Does Layer-Importance Reveal About Transformers and State-Space Models?
Sep 15, 2026Istabrak Abbes, Nizar Islah, Irina Rish +1Transformer Interpretability Selective SSMs
Chandar Research Lab · Mila - Quebec AI Institute · Université de Montréal +2
The Attention Within: Consensus Dynamics in Selective State Space Models
Sep 16, 2026João Pedro Silvestre, Álvaro Rodríguez Abella, Paulo TabuadaSelective SSMs
Electrical and Computer Engineering Department, University of California at Los Angeles, Los Angeles, CA 90095 USA · Department of Applied Mathematics, Comillas Pontifical University, Madrid, 28015 - Madrid, Spain

Sep 8, 2026J/K move · Enter open · S save

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

Amit Ben-Artzy, Roy Schwartz

Computer Science, The Hebrew University of Jerusalem

Transformer Representation Geometry in Language Models State Space Models

Open paper AskPDF arXiv

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.

Abstract

State Space Models

Explore similar work

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

What Does Layer-Importance Reveal About Transformers and State-Space Models?

The Attention Within: Consensus Dynamics in Selective State Space Models