A Comparative Analysis of Contextual Representation Flow in State-Space and Transformer Architectures
Extracted source text · arXiv 2510.06640
% This must be in the first 5 lines to tell arXiv to use pdfLaTeX, which is strongly recommended.
\pdfoutput=1
% In particular, the hyperref package requires pdfLaTeX in order to break URLs across lines.
\documentclass[11pt]{article}
% Change "review" to "final" to generate the final (sometimes called camera-ready) version.
% Change to "preprint" to generate a non-anonymous version with page numbers.
\usepackage[preprint]{acl}
% Standard package includes
\usepackage{times}
\usepackage{latexsym}
% For proper rendering and hyphenation of words containing Latin characters (including in bib files)
\usepackage[T1]{fontenc}
% For Vietnamese characters
% \usepackage[T5]{fontenc}
% See https://www.latex-project.org/help/documentation/encguide.pdf for other character sets
% This assumes your files are encoded as UTF8
\usepackage[utf8]{inputenc}
% This is not strictly necessary, and may be commented out,
% but it will improve the layout of the manuscript,
% and will typically save some space.
\usepackage{microtype}
% This is also not strictly necessary, and may be commented out.
% However, it will improve the aesthetics of text in
% the typewriter font.
\usepackage{inconsolata}
%Including images in your LaTeX document requires adding
%additional package(s)
\usepackage{graphicx}
% If the title and author information does not fit in the area allocated, uncomment the following
%
%\setlength\titlebox{<dim>}
%
% and set <dim> to something 5cm or larger.
\title{Instructions for *ACL Proceedings}
% Author information can be set in various styles:
% For several authors from the same institution:
% \author{Author 1 \and ... \and Author n \\
% Address line \\ ... \\ Address line}
% if the names do not fit well on one line use
% Author 1 \\ {\bf Author 2} \\ ... \\ {\bf Author n} \\
% For authors from different institutions:
% \author{Author 1 \\ Address line \\ ... \\ Address line
% \And ... \And
% Author n \\ Address line \\ ... \\ Address line}
% To start a separate ``row'' of authors use \AND, as in
% \author{Author 1 \\ Address line \\ ... \\ Address line
% \AND
% Author 2 \\ Address line \\ ... \\ Address line \And
% Author 3 \\ Address line \\ ... \\ Address line}
\author{First Author \\
Affiliation / Address line 1 \\
Affiliation / Address line 2 \\
Affiliation / Address line 3 \\
\texttt{email@domain} \\\And
Second Author \\
Affiliation / Address line 1 \\
Affiliation / Address line 2 \\
Affiliation / Address line 3 \\
\texttt{email@domain} \\}
%\author{
% \textbf{First Author\textsuperscript{1}},
% \textbf{Second Author\textsuperscript{1,2}},
% \textbf{Third T. Author\textsuperscript{1}},
% \textbf{Fourth Author\textsuperscript{1}},
%\\
% \textbf{Fifth Author\textsuperscript{1,2}},
% \textbf{Sixth Author\textsuperscript{1}},
% \textbf{Seventh Author\textsuperscript{1}},
% \textbf{Eighth Author \textsuperscript{1,2,3,4}},
%\\
% \textbf{Ninth Author\textsuperscript{1}},
% \textbf{Tenth Author\textsuperscript{1}},
% \textbf{Eleventh E. Author\textsuperscript{1,2,3,4,5}},
% \textbf{Twelfth Author\textsuperscript{1}},
%\\
% \textbf{Thirteenth Author\textsuperscript{3}},
% \textbf{Fourteenth F. Author\textsuperscript{2,4}},
% \textbf{Fifteenth Author\textsuperscript{1}},
% \textbf{Sixteenth Author\textsuperscript{1}},
%\\
% \textbf{Seventeenth S. Author\textsuperscript{4,5}},
% \textbf{Eighteenth Author\textsuperscript{3,4}},
% \textbf{Nineteenth N. Author\textsuperscript{2,5}},
% \textbf{Twentieth Author\textsuperscript{1}}
%\\
%\\
% \textsuperscript{1}Affiliation 1,
% \textsuperscript{2}Affiliation 2,
% \textsuperscript{3}Affiliation 3,
% \textsuperscript{4}Affiliation 4,
% \textsuperscript{5}Affiliation 5
%\\
% \small{
% \textbf{Correspondence:} \href{mailto:email@domain}{email@domain}
% }
%}
%%%%%%%%%%%%%%%% ADDITION %%%%%%%%%%%%%%%%
\usepackage{booktabs}
\usepackage{multirow}
\usepackage{tcolorbox}
\usepackage{amssymb}
\usepackage{amsmath}
\usepackage{amsthm}
\usepackage[tikz]{bclogo}
\usepackage{tabularray}
\usepackage{float}
\usepackage{array}
\usepackage{arydshln}
\usepackage{hhline}
\usepackage{cleveref}
\usepackage{caption}
\usepackage{soul}
\newcommand\BibTeX{B\textsc{ib}\TeX}
\newcommand{\supersmall}[1]{\fontsize{8}{5}\selectfont #1}
\newtheorem{corollary}{Corollary}[section]
\newtheorem{hypothesis}{Hypothesis}[section]
\newtheorem{remark}{Remark}[section]
\newtheorem{definition}{Definition}[section]
\crefformat{subsection}{\S#2#1#3}
\crefformat{section}{\S#2#1#3}
\crefformat{appendix}{\S#2#1#3}
\newtheorem{assumptioninner}{Assumption} % inner counter
\newcounter{assumption}[section]
\newenvironment{assumption}[1][]{
\refstepcounter{assumption} % increment counter
\begin{assumptioninner}[#1]
\renewcommand{\theassumptioninner}{\theassumption} % set inner number
}{
\end{assumptioninner}
}
\newtheorem{lemmainner}{Lemma} % inner counter
\newcounter{lemma}[section]
\newenvironment{lemma}[1][]{
\refstepcounter{lemma} % increment counter
\begin{lemmainner}[#1]
\renewcommand{\thelemmainner}{\thelemma} % set inner number
}{
\end{lemmainner}
}
\newtheorem{theoreminner}{Theorem} % inner counter
\newcounter{theorem}[section]
\newenvironment{theorem}[1][]{
\refstepcounter{theorem} % increment counter
\begin{theoreminner}[#1]
\renewcommand{\thetheoreminner}{\thetheorem} % set inner number
}{
\end{theoreminner}
}
\newtheorem{propositioninner}{Proposition} % inner counter
\newcounter{proposition}[section]
\newenvironment{proposition}[1][]{
\refstepcounter{proposition} % increment counter
\begin{propositioninner}[#1]
\renewcommand{\thepropositioninner}{\theproposition} % set inner number
}{
\end{propositioninner}
}
\definecolor{c3}{cmyk}{0.3081,0,0.7209,0.3255}
\newtcbox{\hlprimarytab}{on line, rounded corners, box align=base, colback=c3!10,colframe=white,size=fbox,arc=3pt, before upper=\strut, top=-2pt, bottom=-4pt, left=-2pt, right=-2pt, boxrule=0pt}
\newtcbox{\hlsecondarytab}{on line, box align=base, colback=red!10,colframe=white,size=fbox,arc=3pt, before upper=\strut, top=-2pt, bottom=-4pt, left=-2pt, right=-2pt, boxrule=0pt}
\newcommand{\ua}[1]{\tiny{\hlprimarytab{($\uparrow$#1)}}}
\newcommand{\da}[1]{\tiny{\hlsecondarytab{($\downarrow$#1)}}}
% Define a counter for findings
\newcounter{findingcounter}[subsection] % reset each subsection
\renewcommand{\thefindingcounter}{\thesubsection.\arabic{findingcounter}}
\newcommand{\finding}[1]{%
\refstepcounter{findingcounter}%
\begin{bclogo}[couleur=black!10, epBord=1, arrondi=0.1, logo=\bclampe, marge=2, ombre=true, blur, couleurBord=black!20, tailleOndu=3, sousTitre={\em \small \textbf{F \thefindingcounter.} #1}]{}
\end{bclogo}%
}
\newcommand{\xlong}[1]{\textcolor{red}{(Long: #1)}}
\newcommand{\nhat}[1]{\begin{color}{blue} (Nhat: #1) \end{color}}
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\title{A Comparative Analysis of Contextual Representation Flow in \\State-Space and Transformer Architectures}
\author{
Nhat M. Hoang$^{1}$, Do Xuan Long$^{2}$, Cong-Duy Nguyen$^{1}$, Min-Yen Kan$^{2}$, Luu Anh Tuan$^{1}$\\
$^{1}$Nanyang Technological University, Singapore,\\
$^{2}$National University of Singapore\\
\texttt{\{hoangmin003, nguyentr003\}@e.ntu.edu.sg}, \\
\texttt{xuanlong.do@u.nus.edu}, \texttt{kanmy@comp.nus.edu.sg}, \texttt{anhtuan.luu@ntu.edu.sg}\\
}
% Author information can be set in various styles:
% For several authors from the same institution:
% \author{Author 1 \and ... \and Author n \\
% Address line \\ ... \\ Address line}
% if the names do not fit well on one line use
% Author 1 \\ {\bf Author 2} \\ ... \\ {\bf Author n} \\
% For authors from different institutions:
% \author{Author 1 \\ Address line \\ ... \\ Address line
% \And ... \And
% Author n \\ Address line \\ ... \\ Address line}
% To start a seperate ``row'' of authors use \AND, as in
% \author{Author 1 \\ Address line \\ ... \\ Address line
% \AND
% Author 2 \\ Address line \\ ... \\ Address line \And
% Author 3 \\ Address line \\ ... \\ Address line}
% \author{First Author \\
% Affiliation / Address line 1 \\
% Affiliation / Address line 2 \\
% Affiliation / Address line 3 \\
% \texttt{email@domain} \\\And
% Second Author \\
% Affiliation / Address line 1 \\
% Affiliation / Address line 2 \\
% Affiliation / Address line 3 \\
% \texttt{email@domain} \\}
\begin{document}
\maketitle
\begin{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.\end{abstract}
\section{Introduction}
Long-context processing remains a critical challenge in natural language processing, with applications spanning document analysis, retrieval systems, and multi-turn dialogue \citep{Beltagy2020Longformer,liu-etal-2024-lost,goldman-etal-2024-really,liu2025comprehensive}. While Transformer-Based Models (TBMs) \citep{vaswani2017attention} have established strong performance baselines, their quadratic attention complexity poses significant scalability limitations for extended context \cite{mamba1}. State Space Models (SSMs) like Mamba \cite{mamba1} have emerged as promising linear-complexity alternatives, yet recent work has highlighted specific limitations in their long-context capabilities \cite{jelassi2024repeat, chen2024stuffed}.
Recent work has begun to probe the internal dynamics of these architectures. For example, \citet{skean2025layer} showed that intermediate layers often outperform final layers for task-relevant information in both TBMs and SSMs, challenging the conventional focus on final-layer outputs. Similarly, \citet{wang2025understanding} identified oversmoothing and recency bias in SSMs, where token representations converge as models favor local over distant context. {However, it remains unclear how TBMs and SSMs fundamentally differ in propagating and transforming contextual representations across layers, particularly when token- and layer-level perspectives are considered together. Understanding these differences is crucial for diagnosing their inductive biases and guiding the design of more effective long-context models. Prior studies have examined some of these aspects in isolation, but a unified characterization is lacking.}
% Recent investigations have begun to uncover key aspects of these architectures' internal dynamics. For example, \citet{skean2025layer} showed that intermediate layers consistently outperform final layers for task-relevant information extraction in both TBMs and SSMs, challenging conventional emphasis on final-layer outputs. Aditionally, \citet{wang2025understanding} identified oversmoothing and recency bias in SSMs, with token representations becoming similar as they prioritze local context over distant information.
% \textcolor{blue}{Despite these advances, it is still unclear how TBMs and SSMs fundamentally differ in their approach to propagating and transforming contextual representations across layers, especially when considering both token-level and layer-level representations in tandem. Most prior studies have explored these issues in isolation; a unified, side-by-side characterization remains lacking.}
{To bridge this gap, we present the first comprehensive pairwise comparison of representation flow in TBMs and SSMs, both empirically and theoretically. Our analysis spans local (token-level) and global (layer-level) perspectives, with aligned setups and evaluation tasks enabling direct, side-by-side comparison. This reveals the architectural fingerprints that drive success or failure on long-context tasks and offers actionable guidance for future model and training design.} {Taken together, these findings lay the groundwork for hybrid architectures and model-specific optimizations, paving the way for more robust and efficient long-range reasoning.} In summary, our main contributions are:
\begin{enumerate}
\item \textbf{Unified layer-wise representation propagation.} We characterize token- and layer-level dynamics, revealing opposing trends of diversity and homogenization in TBMs and SSMs.
\item \textbf{Architectural bias.} We show that oversmoothing in TBMs stems from architectural design, whereas in SSMs it arises primarily from training dynamics.
\item \textbf{Intermediate-layer effectiveness.} Intermediate layers outperform final layers across tasks, model scales, and context lengths.
\item \textbf{Theoretical analysis.} We provide formal stability results demonstrating that, under practical conditions, representation propagation is inherently more stable in SSMs than in TBMs.
\end{enumerate}
\section{Related Work}
\subsection{Long-Context Processing}
While TBMs \citep{vaswani2017attention} remain the dominant NLP architecture, their quadratic attention complexity limits scalability to long contexts \cite{mamba1}. SSMs offer a linear-complexity alternative, achieving efficiency gains in long-sequence tasks through compact state representations \citep{gu2022efficiently, mamba1}. However, recent studies reveal distinct architectural biases: SSMs often emphasize recency and local information, whereas TBMs maintain a broader contextual focus \citep{jelassi2024repeat, chen2024stuffed, wang2025understanding}. These contrasting inductive biases motivate systematic analysis of how representations propagate within each family.
\subsection{Representation Flow, Intermediate Layers, and Oversmoothing}
Recent works showed that intermediate layers often outperform final layers in both TBMs and SSMs, challenging the conventional reliance on final representations \citep{knowbutdonttell, skean2025layer}. Another study emphasized oversmoothing and recency bias in SSMs, where token representations gradually homogenize with depth \citep{wang2025understanding}. These findings suggest that models may fail to fully leverage their depth, raising a question about how representations propagate across layers.
Layer-wise analyses commonly use Centered Kernel Alignment (CKA) \citep{cka} to measure representational similarity in TBMs \citep{conneau-etal-2020-emerging, vulic-etal-2020-probing}, while cosine similarity and variance-based metrics track feature evolution across layers. Probing further reveals where task-relevant information resides \citep{vulic-etal-2020-probing,knowbutdonttell}. For example, probing on SSMs \citep{paulo2024doestransformerinterpretabilitytransfer} revealed that simple probes can recover correct knowledge even when fine-tuned outputs are incorrect, underscoring the richness of SSMs' internal representations.
\section{Empirical Analysis}
We conduct a mechanistic analysis to compare how different architectures maintain information flow during long-context processing. Our goal is to uncover how contextualized representations propagate and evolve across layers. We focus on three complementary perspectives: \textbf{(i) token-level analysis} (\Cref{sssec:token-analysis}) to track the evolution and correlation of {token} representations across depth; \textbf{(ii) layer-level analysis} (\Cref{sssec:layer-analysis}) to examine of layer representation; and \textbf{(iii) probing analysis} (\Cref{sssec:probing-analysis}) to identify layers encoding task-relevant features most critical for downstream performance. The first two analyses provide insights into how models transform information, while probing analysis focuses on downstream performance.
\subsection{Common Experimental Setups}
\paragraph{Models.}
We evaluated models pre-trained on the Pile dataset \cite{pile} for a fair comparison, covering both TBM and SSM families. Our comparison includes two TBMs, \texttt{GPT-Neo-2.7B} \cite{gpt-neo} and \texttt{Pythia-2.7B} \cite{pythia}; alongside three SSMs: \texttt{Mamba2-2.7B} \cite{mamba2}, \texttt{Mamba-2.8B} \cite{mamba1}, and a smaller \texttt{Mamba2-130M}. This selection enables both cross-architecture comparison and scaling effects within SSMs.
\paragraph{Tasks.}
We follow \citet{liu-etal-2024-lost} to adopt two benchmarks emphasizing long-range reasoning where models must process and retrieve information from extended contexts. These include: \textit{(i) Multi-Document Question Answering} (MDQA), where the input consists of multiple documents and a question, requiring the model to identify the relevant document and generate an answer; and \textit{(ii) Key-Value Pair Retrieval} (KVPR), where the input consists of multiple KV pairs represented as 128-bit randomly generated UUIDs. The number of documents or KV pairs is chosen so that the total context length $n \in \{300, 1K, 2K, 4K\}$. While absolute performance varies across models, consistent representational trends are observed. Thus, we report averaged results across tasks.
\paragraph{Probing Classifiers.}
For each input (instruction, documents/KV pairs, question), we extract its final token representation, $h^{(l)}_T \in \mathbb{R}^d$, from each layer $l$ where $d$ is the hidden size, following prior work \cite{knowbutdonttell}. In SSMs, the final token summarizes the entire context due to sequential state propagation. For each layer, we train a linear probe $f^l: \mathbb{R}^d \to \mathbb{R}^C$ to minimize the cross-entropy loss:
\begin{equation}
\mathcal{L}^l := -\frac{1}{M} \sum_{i=1}^M \sum_{c=1}^C y_{i,c} \log(\hat{y}_{i,c}^l)
\end{equation}
\noindent where $C$ is the number of classes (e.g., number of KV pairs in KVPR), $M$ is the number of train samples, $y_{i,c}$ is the true label and $\hat{y}_{i,c}^l$ is the softmax prediction of data sample $i$.
\paragraph{Implementation Details.}
We use pre-trained checkpoints of TBMs and SSMs from the Hugging Face \cite{wolf-etal-2020-transformers}. All probes are trained from frozen representations for 150 epochs using Adam optimizer \cite{adam} with a learning rate of 0.05 on the full $20K$-sample training set (no batching), with evaluation on a held-out validation split. Reported results are mean accuracies over five random seeds across both tasks, under consistent hyperparameters and evaluation protocols on NVIDIA L40S GPUs.
\subsection{Token-Level Analysis}
\label{sssec:token-analysis}
Our token-level analysis aims to understand: (i) how smoothly representations evolve across layers; (ii) whether tokens maintain their distinctiveness or become homogenized within layers; and (iii) whether these behaviors arise from architectural biases or training dynamics.
\paragraph{Setups.}
To track token representations dynamics in TBMs and SSMs, we employ three complementary cosine-similarity–based analyses capturing distinct properties of representation flow.
First, we calculate \textit{layerwise cosine similarity} to quantify the \textbf{temporal evolution} of each token across adjacent layers. For a token representation $h^{(l)}_t \in \mathbb{R}^d$ at layer $l$ and position $t$, we define:
\[
\mathrm{Sim}(h^{(l)}_t, h^{(l+1)}_t) = \frac{h^{(l)}_t \cdot h^{(l+1)}_t}{\|h^{(l)}_t\| \|h^{(l+1)}_t\|}
\]
\noindent which measures how much each token changes as it propagates through the network. Higher similarity indicates gradual, stable transitions, whereas lower values indicate substantial changes or instability.
Second, we calculate \textit{inter-token cosine similarity} within layers to assess the \textbf{spatial cohesiveness} of token features, a proxy for oversmoothing \cite{ali2024hidden}. Given $h^{(l)} \in \mathbb{R}^{n \times d}$, we compute:
\begin{equation}
\mathrm{InterSim}^{(l)} := \frac{2}{n(n-1)} \sum_{i=1}^{n} \sum_{j=i+1}^{n} \frac{h^{(l)}_i \cdot h^{(l)}_j}{\|h^{(l)}_i\| \|h^{(l)}_j\|}
\end{equation}
\noindent which reflects the average pairwise similarity among tokens, excluding self-similarity. Higher values indicate homogenization and loss of token distinctiveness, indicative of oversmoothing.
Finally, to \textbf{disentangle training artifacts from architectural priors}, we analyze both pretrained and randomly initialized models. For random initialization, we explore multiple schemes including Gaussian, Xavier \cite{xavier}, and He \cite{he}, under varying parameter settings. This allows us to isolate oversmoothing tendencies inherent to the architecture itself.
All analyses are performed on MDQA and KVPR with context length $n = 2K$ tokens, matching the effective window size of \texttt{GPT-Neo-2.7B} and \texttt{Pythia-2.8B}.
\paragraph{(3.2.1) Token evolution across layers.}
\begin{figure}[t!]
\centering
\includegraphics[width=\columnwidth]{imgs/cosine_maps.png}
\caption{
\textbf{Token-level cosine similarity evolution across layers.} We observe that \textbf{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, \textbf{SSMs (right column)} reveal greater variability and exploratory changes in early layers, with gradual convergence later, highlighting distinct representation dynamics per architecture.
}
\label{fig:tokenwise-cosine}
\end{figure}
Figure~\ref{fig:tokenwise-cosine} reveals clear differences in how token representations evolve through layers within the two architectures, though the overall trend remains consistent across tokens within each model (i.e., all rows follow a similar pattern). For \texttt{GPT-Neo-2.7B}, similarity starts at nearly $100\%$, dips slightly to $90\%$ by layer 3, quickly recovers to $100\%$, and holds steady until a sharp decline to $70\%$ at layer 32. Similarly, \texttt{Pythia-2.8B} begins at $\sim 90\%$, gradually increases to $\sim 100\%$ by layer 23, but also drops from layer 29 to 70\% by the final layer, mirroring \texttt{GPT-Neo-2.7B}'s late shift. In contrast, \texttt{Mamba2-2.7B} fluctuates between $20\%$ and $40\%$ until layer 51, reflecting diverse directional changes, before rising steadily to $80\%$ by the last layer 64. Meanwhile, \texttt{Mamba-2.8B} displays a unique elbow pattern: starting at 30\%, it drops to $0\%$ or even negative values at layers 30 and 41, then gradually rises to $\sim 60\%$ by the final layer. This suggests that while TBMs prioritize stability and preservation, SSMs promote continual token evolution and refinement. Such dynamism may be critical for maintaining expressivity in deep models, particularly for long-context tasks.
\finding{\small
TBMs exhibit stable token evolution until a sharp final-layer shift, contrasting with SSMs’ varied directions that are converged in later layers.\label{F1}
}
\paragraph{(3.2.2) Token uniqueness within layers.}
\begin{figure}[t!]
\centering
\includegraphics[width=\columnwidth]{imgs/inter_cosine_similarity.png}
\caption{
\textbf{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. \textbf{TBMs} rapidly increase token similarity early and maintain high homogenization until a late drop, reflecting oversmoothing of features. Conversely, \textbf{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.
}
\label{fig:inter-cosine}
\end{figure}
Figure~\ref{fig:inter-cosine} highlights differences in token distinctiveness inside layers. TBMs, \texttt{GPT-Neo-2.7B} and \texttt{Pythia-2.8B}, show high inter-token similarity (around $90\%$ and $70\%$ respectively) throughout most layers, indicative of oversmoothing where token representations become increasingly alike. Notably, both models experience a significant reduction to roughly $30\%$ near the final layer, signaling a late resurgence of token diversity. Conversely, SSMs such as \texttt{Mamba2-2.7B} and \texttt{Mamba-2.8B} maintain substantially lower inter-token similarity (around $30\%$ and $10\%$ respectively) across their layers, preserving token uniqueness longer. Particularly, \texttt{Mamba-2.8B} shows an increase in similarity beginning around layer 40, reaching $70\%$ by the final layer, while \texttt{Mamba2-2.7B} sustains lower similarity throughout. These trends underscore SSMs' ability to preserve token individuality, contrasting with TBMs' early oversmoothing.
\finding{\small
TBMs collapse token distinctions early and recover diversity late, while SSMs maintain the uniqueness longer.\label{F2}
}
\paragraph{(3.2.3) Architectural bias on the oversmoothing problem.}
\begin{figure}[t!]
\centering
\includegraphics[width=\columnwidth]{imgs/cosine_random_init.png}
\caption{
\textbf{Token-level cosine similarity across layers for both pre-trained and randomly initialized models.} Comparing pretrained and random initializations reveals an \textbf{architectural bias towards oversmoothing in TBM (left column)} where high similarity exists regardless of training, while \textbf{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.
}
\label{fig:random-init}
\end{figure}
Figure~\ref{fig:random-init} demonstrates the role of architecture in oversmoothing behavior by comparing pretrained and randomly initialized models. \texttt{GPT-Neo-2.7B} exhibits consistently high token similarity (approximately $75-100\%$) layer-by-layer regardless of whether weights are pretrained or randomly initialized, substantiating that oversmoothing arises fundamentally from the TBM architecture. On the other hand, \texttt{Mamba2-2.7B} displays markedly low similarity (around $0–25\%$) across layers when randomly initialized, developing oversmoothing patterns only after training. This contrast clearly shows that while oversmoothing in TBMs is largely an architectural artifact, in SSMs it is a learned phenomenon linked to training dynamics and optimization processes.
\finding{\small
Oversmoothing in TBMs is an intrinsic architectural, while in SSMs it is primarily a consequence of training or optimization process.\label{F3}
}
\subsection{Layer-Level Analysis} \label{sssec:layer-analysis}
This analysis aims to answer two key questions: (i) how TBMs and SSMs maintain and reshape coherent global feature manifolds across layers; and (ii) whether these manifolds evolve smoothly or undergo abrupt shifts as depth increases. Understanding these clarifies whether long-range dependencies are preserved or degraded as representations propagate through the network.
\paragraph{Setups.} To investigate how holistic layer-level features evolve across layers, we adopt the CKA metric \cite{cka}, which measures the similarity between layer representations in a geometry-aware manner. CKA provides insight into how the overall feature manifold transforms as the model deepens. To complement CKA, we define two additional metrics: \textit{Smoothness (Sm)}, measuring the local consistency of layerwise changes, and \textit{Stability (St)}, quantifying the global variability of representations through the network. Lower values of Sm and St indicate smoother, more stable evolution of feature spaces across layers. Together, these metrics offer a comprehensive view of representation dynamics at the layer-level. The formulas for Sm and St are:
\begin{equation}
\begin{aligned}
\text{Sm} := \frac{1}{n \cdot d} \sum_{t=1}^{n}\sum^{d} \bigg( \frac{1}{L-2}
\sum_{l=0}^{L-2} \\ \left|
h^{(l+1)} - \frac{h^{(i)} + h^{(i+2)}}{2}
\right| \bigg)
\end{aligned}
\label{eq:sm-formula}
\end{equation}
\begin{equation}
\text{St} := \frac{1}{n \cdot d} \sum_{t=1}^{n}\sum^{d} \sqrt{\frac{1}{L} \sum_{l=1}^{L} (h^{(l)} - \bar{h})^2}
\label{eq:st-formula}
\end{equation}
\noindent where $\bar{h}$ represents the mean representation across all $L$ layers. The analyses are conducted under the same task setup and context length as in the token-level experiments to ensure comparability.
%$h^{(l)} \in \mathbb{R}^{n \times d}$ denotes the representation of $n$ tokens with $d$ dimensions at layer $l$,
\begin{figure}[t!]
\centering
\includegraphics[width=\columnwidth]{imgs/cka_maps.png}
\caption{
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 \texttt{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.
}
\label{fig:cka}
\end{figure}
\begin{table*}[t!]
\centering
\small{\resizebox{0.99\textwidth}{!}{
\begin{tabular}{l ccc ccc ccc ccc}
\toprule
\multirow{2}{*}{\textbf{Model}} & \multicolumn{3}{c}{$n = 300$} & \multicolumn{3}{c}{$n = 1K$} & \multicolumn{3}{c}{$n = 2K$} & \multicolumn{3}{c}{$n = 4K$} \\
\cmidrule(lr){2-4} \cmidrule(lr){5-7} \cmidrule(lr){8-10} \cmidrule(lr){11-13}
& \textbf{Prob Acc. $\uparrow$} & \textbf{Sm. $\downarrow$} & \textbf{St. $\downarrow$}
& \textbf{Prob Acc. $\uparrow$} & \textbf{Sm. $\downarrow$} & \textbf{St. $\downarrow$}
& \textbf{Prob Acc. $\uparrow$} & \textbf{Sm. $\downarrow$} & \textbf{St. $\downarrow$}
& \textbf{Prob Acc. $\uparrow$} & \textbf{Sm. $\downarrow$} & \textbf{St. $\downarrow$} \\
\midrule
GPT-Neo-2.7B & $88.7 \da{10.9}$ & $0.938$ & $3.635$
& $70.7 \da{18.7}$ & $0.980$ & $3.833$
& $53.1 \da{25.3}$ & $0.979$ & $3.821$
& - & - & - \\
Pythia-2.8B & $95.9 \textbf{\da{3.8}}$ & $0.265$ & $1.059$
& $72.3 \da{14.0}$ & $0.277$ & $1.122$
& $60.9 \da{15.4}$ & $0.280$ & $1.147$
& - & - & - \\
Mamba2-130M & $72.5 \da{17.1}$ & $3.572$ & $5.395$
& $42.3 \da{13.8}$ & $3.704$ & $5.665$
& $28.6 \textbf{\da{8.0}}$ & $3.752$ & $5.735$
& $20.6 \textbf{\da{7.6}}$ & $3.846$ & $5.890$ \\
Mamba-2.8B & $93.0 \da{6.4}$ & $\textbf{0.173}$ & $\textbf{0.315}$
& $62.5 \textbf{\da{11.7}}$ & $\textbf{0.189}$ & $\textbf{0.354}$
& $41.7 \da{14.0}$ & $\textbf{0.193}$ & $\textbf{0.363}$
& $36.7 \da{12.2}$ & $\textbf{0.190}$ & $\textbf{0.354}$ \\
Mamba2-2.7B & $86.0 \da{13.5}$ & $1.945$ & $3.238$
& $57.2 \da{17.6}$ & $1.939$ & $3.288$
& $38.2 \da{20.6}$ & $1.932$ & $3.275$
& $29.7 \da{13.3}$ & $1.897$ & $3.193$ \\
\bottomrule
\end{tabular}
}}
\caption{
This table shows the probing accuracy (\%) using the \textbf{last layer's representation}. We run all the evaluation 5 times and report the average results. $\da{x}$ is the accuracy difference between the probe trained on the last layer and on the \textbf{peak layer}. The best results are \textbf{bolded}.
}
\label{tab:results}
\end{table*}
\paragraph{(3.3.1) Holistic feature manifold dynamics.} Figure~\ref{fig:cka} shows that \texttt{GPT-Neo-2.7B} maintains near-perfect CKA similarity ($\sim 100\%$) in early layers, with even the 1st layer retaining 90\% similarity with layer 20. Consistent with observations in Section~\ref{sssec:token-analysis}, a feature shift occurs at layer 29, where the similarity score drops to 80\% by the final layer. This suggests the manifold formed by token representations is stable and smoothly transforms before undergoing reconfiguration in late layers. In contrast, \texttt{Pythia-2.8B} starts at 90\% similarity but drops sharply to 30\% by the second layer. However, its propagation pattern in the remaining layers (3-32) closely resembles that of \texttt{GPT-Neo-2.7B}. On the other hand, both \texttt{Mamba2-2.7B} and \texttt{Mamba-2.8B} exhibit fluctuating similarity scores between 20\% and 70\% until layer 51, after which they steadily increase, reaching 90\% by the final layer 64. This difference signals distinct architectural strategies in maintaining and reshaping feature spaces at a global level, revealing their contrasting approaches to integrating and refining context beyond individual token trajectories.
\finding{\small
TBMs maintain coherent feature manifolds early while SSMs gradually refine global structure towards late-layer stabilization.\label{F4}
}
\paragraph{(3.3.2) Comparative analysis of Smoothness and Stability.} Table~\ref{tab:results} reveals that among SSMs, only \texttt{Mamba-2.8B} consistently achieves lower Sm ($\sim 0.19$) and St ($\sim 0.34$) values compared to the TBMs (\texttt{Pythia-2.8B}: $\sim 0.27$ Sm and $\sim 1.11$ St; \texttt{GPT-Neo-2.7B}: $\sim 0.97$ Sm and $\sim 3.76$ St), indicating a more gradual and stable evolution of its representations. This behavior may stem from \texttt{Mamba-2.8B}'s smaller state size, which may constrain its capacity to model diverse token features, resulting in smoother transitions across layers. In contrast, \texttt{Mamba2-2.7B}, with an $8\times$ larger state size, shows higher Sm and St values than \texttt{Pythia-2.8B}, suggesting that increased state dimensions amplify representation variability and reduce stability, potentially leading to more abrupt changes. Among TBMs, \texttt{Pythia-2.8B} exhibits Sm and St values roughly three times lower than \texttt{GPT-Neo-2.7B}, indicating that \texttt{Pythia-2.8B}'s architectural refinements foster smoother and more stable feature propagation. These findings underscore how design choices, even within the same model family, significantly influence representation dynamics.
\finding{\small
Smoothness and stability of representation evolution depend on model-specific factors like scale and parameterization rather than architecture type.\label{F5}
}
\subsection{Probing Analysis} \label{sssec:probing-analysis}
We investigate whether the final layer contains the most task-relevant representation, as commonly assumed in model design.
\paragraph{(3.4.1) Intermediate layers outperform the final layer.} Table~\ref{tab:results} and Figure~\ref{fig:layerwise-acc} show that \texttt{GPT-Neo-2.7B} and \texttt{Pythia-2.8B} achieve peak accuracy around layer 10 (out of 32) before dropping by up to 26\% in the final layer. \texttt{Mamba2-2.7B} peaks between layers 4 and 14 (out of 64) and declines by up to 13.9\%. In contrast, \texttt{Mamba-2.8B} reaches its peak later, around layer 28 or beyond (out of 64), with a more gradual drop of at most 10.5\% by the final layer.
\finding{\small
Task-relevant representations peak in intermediate layers for both architectures, with SSMs having a smaller gap to the last layer than TBMs.\label{F6}
}
% \paragraph{(2) Task-specific trends in performance gap.} \xlong{a sentence why doing this.} In the MDQA task, SSMs exhibit smaller performance drops between their peak and final layer (at most 13.9\%) compared to Transformers (up to 23.4\%). This aligns with our earlier observation that SSMs, such as \texttt{Mamba2-2.7B}, gradually align token- and layer-level features. Their progressive homogenization (Finding~\ref{F2}) may help preserve long-range contextual details until the final layers.
% Conversely, in the KVPR task, which relies on precise token-level information, Transformer models (e.g., \texttt{GPT-Neo-2.7B} and \texttt{Pythia-2.8B}) show smaller performance gaps (at most 26.4\%) than SSMs (up to 36.4\%). This suggests that Transformers' stable early token-level representations better retaixtxtxtxtxtxn the precise information needed for key-value retrieval, even as their final layers undergo sharper transformations (Finding~\ref{F1}).
% \finding{\small
% SSMs better preserve long-range context, while Transformers retain precise token associations.\label{F6}
% } \xlong{``preserve'' is not a good word, choose another one such as ``learn/represent''.}
\paragraph{(3.4.2) Effect of context length.} Table~\ref{tab:results} shows that while smoothness and stability metrics (Sm and St) remain stable across context lengths ranging from 300 to 4K tokens, probing accuracy steadily declines with increasing input length. This suggests that although internal representations evolve predictably, model capacity constraints limit the ability to retain task-relevant information in longer sequences. This trend holds consistently across models and tasks, illustrating a key trade-off between representational stability and capacity to capture extended context.
\finding{\small
Despite stable representation evolution across varied context lengths, probing accuracy decreases as context length increases, indicating capacity limitations rather than representation instability.\label{F7}
}
\paragraph{(3.4.3) Effect of Model Size.} Comparing the smaller (\texttt{Mamba2-130M}) and larger (\texttt{Mamba2-2.7B}) SSM variants reveals that the larger model significantly improves final-layer probing accuracy, particularly for longer contexts (by 2.3\% in the $4K$ MDQA task and 15.9\% in the $4K$ KVPR task). While its intermediate layers capture richer representations, they also exhibit a larger accuracy drop toward the final layer (increasing from 6.7\% to 13.3\% in MDQA and from 24.0\% in KVPR), indicating a more aggressive transformation that may discard fine-grained details. Nonetheless, the overall improvement in task performance suggests that deeper processing benefits outweigh this loss.
\finding{\small
While larger model size generally improves task performance, final-layer representations tend to become more abstract, which might reduce accessibility to certain fine-grained intermediate features.\label{F8}
}
%%%%%%%%%%%%%%%%%%%%%%%%
\section{Theoretical Analysis}
Our theoretical analysis aims to formally characterize the stability of representation propagation in TBMs and SSMs, providing mechanistic insight into their empirical differences. We focus on the expected value of the stability metric $\mathbb{E}[\text{St}^2]$ (Eq.~\ref{eq:st-formula}) over a single layer ($L=1$) to compare the stability of TBM and SSM, and we then generalize the result to depth-$L$ in Appendix~\ref{appdx:theory-generalization}. For SSMs, we use \texttt{Mamba}'s formulas as a representative case study for SSM to compare their behavior rigorously.
\subsection{Backgrounds and Setups} \label{subsec:theory-setups}
\paragraph{Backgrounds.} Given an input sequence $x = [x_1, x_2, \dots, x_n]$ of length $n$, where $x_i$ are tokens coming from a finite vocabulary $\mathcal{V}$, the Transformer \citep{vaswani2017attention} processes it as follows. Initially, each token $x_i$ is mapped to a $d$-dimensional vector $v_i = \text{Embed}(x_i) \in \mathbb{R}^d$ using an embedding layer. To encode positional information, a positional embedding $p_i \in \mathbb{R}^d$ is added to the token representation. The resulting embedded sequence is expressed as a matrix $h^{(0)} = [h^{(0)}_1, \dots, h^{(0)}_n]^T \in \mathbb{R}^{n\times d}$ where $h^{(0)}_i = v_i + p_i$. Subsequently, the sequence passes through $L_t$ transformer blocks, each of which applies the following transformation:
\begin{equation}\label{eq:transformer-eq}
\begin{aligned}
h^{(l)} := h^{(l-1)} + \text{Attn} (h^{(l-1)})
+ \text{FFN} \big(h^{(l-1)} + \\ \text{Attn} (h^{(l-1)})\big)
\end{aligned}
\end{equation}
\noindent where Attn, FFN stand for single-head self-attention, and feed-forward layers. Note that we use the single-head self-attention instead of multi-head attention, and we skip the layer normalization layers for simplicity following \citet{feng2023towards}. We also skip the layer normalization in the SSM formulations in Equation~\ref{eq:ssm-eq} below.
With the same input matrix $h^{(0)}$, Mamba \citep{mamba1} also passes it through $L_m$ layers where $L_m$ typically larger than $L_t$:
\begin{align}\label{eq:ssm-eq}
h^{(l)} &:= h^{(l-1)} + g^{(l)} \Big( \text{S6} \big( f^{(l)}, h^{(l-1)} \big) \circ z^{(l)} \Big)
\end{align}
\noindent here S6 is the selective SSM transformation, \( g^{(l)} \) is a linear transformation, and $z^{(l)} = \text{SiLU}(\text{Linear}(h^{(l-1)}))$.
\paragraph{Setups.} We consider the case where $L=1$ following \citet{feng2023towards,kajitsuka2024are}. The stability metric $\text{St}^2$ (Eq.~\ref{eq:st-formula}) is then defined as:
\begin{equation}
\text{St}^2 \propto \|h^{(1)} - h^{(0)}\|^2,
\end{equation}
where $h^{(l)}$ is the representation at layer $l$. To enable our analysis, we make the following assumptions:
\begin{assumption}[Random initial representations]\label{asm:A1}
The initial representation $h^{(0)} \in \mathbb{R}^{n \times d}$ is a Gaussian matrix with $\mathbb{E}[h^{(0)}] = 0$ and covariance $\text{Cov}(h^{(0)}) = \sigma^2 I$, where each $h^{(0)}_i \sim N(0, \sigma^2 I)$.
\end{assumption}
\begin{assumption}[Independence]\label{asm:A2}
The token representations $h^{(0)}_j$ and $h^{(0)}_t$ are independent for $j \neq t$ (i.e., $\mathbb{E}[h^{(0)}_j (h^{(0)}_t)^T] = 0$).
\end{assumption}
\begin{assumption}[Deterministic parameters]\label{asm:A3}
Model parameters (e.g., weight matrices) are deterministic, with randomness stemming solely from $h^{(0)}$.
\end{assumption}
\begin{assumption}[Non-linearity approximations]\label{asm:A4}
Non-linear functions (e.g., GELU in TBMs, SiLU in Mamba) are approximated via Taylor expansion as $\text{GELU}(x) \approx \frac{1}{2}x$ and $\text{SiLU}(x) \approx \frac{1}{2}x$ to simplify the computations.
\end{assumption}
\subsection{Theoretical Results}
Let \( F: \mathbb{R}^{n \times d} \to \mathbb{R}^{n \times d} \) denote a function that transforms an input matrix \( h^{(0)} \in \mathbb{R}^{n \times d} \) into an output matrix \( h^{(1)} = F(h^{(0)}) \), where \( n \) represents the number of rows (e.g., sequence length) and \( d \) represents the number of columns (e.g., feature dimension). For a TBM, let \( \text{St}^2_{\text{Trans}} \) denote the squared stability metric, with \( \Sigma_{F,\text{Trans}} \) as the covariance matrix and \( \mu_{F,\text{Trans}} \) as the mean vector of the output of the Transformer function \( F \). For a Mamba model, let \( \text{St}^2_{\text{Mamba}} \) denote the squared stability metric, with \( \Sigma_{F,\text{Mamba}} \) as the covariance matrix and \( \mu_{F,\text{Mamba}} \) as the mean vector of the output of the Mamba function \( F \). We aim to prove that $\mathbb{E}[\text{St}^2_{\text{Trans}}] > \mathbb{E}[\text{St}^2_{\text{Mamba}}]$.
\vspace{3mm}
\begin{proposition}\label{prop:prop-common-formula}
Consider $F$ (e.g., Mamba or Transformer) with input matrix $h^{(0)}$. Then, the expected squared stability satisfies
\begin{equation}
\mathbb{E}\big[\text{St}^2\big] \propto \mathbb{E}\Big[ \| F(h^{(0)}) \|_F^2 \Big]
= \operatorname{Tr}(\Sigma_F) + \| \mu_F \|_F^2,
\end{equation}
% \begin{align*}
% \mathbb{E}\big[\text{St}^2\big] &= \frac{1}{L} \sum_{l=1}^L \mathbb{E}\Big[ \| F(h^{(l)}) \|_F^2 \Big], \\
% \mathbb{E}\Big[ \| F(h^{(l)}) \|_F^2 \Big] &= \operatorname{Tr}(\Sigma_F) + \| \mu_F \|_F^2
% \end{align*}
where $\mu_F = \mathbb{E}[F(h^{(0)})]$ is the mean representation, $\Sigma_F = \operatorname{Cov}[F(h^{(0)})]$ is the covariance, and $\|\cdot\|_F$ denotes the Frobenius norm.
\end{proposition}
\vspace{1mm}
\begin{proposition}\label{prop:prop-att-odd}
The attention function is odd, i.e.,
\begin{equation}
\text{Attn}(-x) = -\text{Attn}(x).
\end{equation}
\end{proposition}
\vspace{1mm}
\begin{proposition}\label{propssm:prop10}
For zero-mean inputs $h^{(0)}$, the expected attention output vanishes:
\begin{equation}
\mathbb{E}[\text{Attn}(h^{(0)})] = 0.
\end{equation}
\end{proposition}
\vspace{1mm}
\begin{theorem}\label{theorem:thrm-1}
Let $\tilde{h}^{(0)} = h^{(0)} + \text{Attn}(h^{(0)})$ and consider the feed-forward layer
\[
\text{FFN}(\tilde{h}^{(0)}) = W_2 \, \text{GELU}(W_1 \tilde{h}^{(0)} + b_1) + b_2,
\]
where $W_1, W_2, b_1, b_2$ are independent. Then, the expected output of the feed-forward block is approximately
\begin{equation}
\mu_{F,\text{Trans}} \approx \frac{1}{2} W_2 b_1 + b_2.
\end{equation}
\end{theorem}
\vspace{1mm}
\begin{theorem}\label{theorem:thrm-2}
Let $F_{\text{Trans}} = \text{Attn}(h^{(0)}) + \text{FFN}(\tilde{h}^{(0)})$ and given that $\text{Attn} (h^{(0)})$ is odd and $h^{(0)}$ is symmetrically distributed around zero, the bias terms do not affect the covariance. We have:
\begin{equation}
\begin{split}
\text{Tr}(\Sigma_{F,\text{Trans}}) &= \sigma^2 \text{Tr}(T_1 T_1^\top) + n \sigma^2 \text{Tr}(T_2 T_2^\top) \\
&\quad + 2 \sigma^2 \text{Tr}(T_1 W_V T_2^\top)
\end{split}
\end{equation}
where $T_1 = I + \frac{1}{2}W_2 W_1$, $T_2 = \frac{1}{2}W_2 W_1$
\end{theorem}
\vspace{1mm}
\begin{theorem}\label{theorem:thrm-3}
Let $W_{h'_j} = W_{c_j} W_h$, where $W_{c_j}$ is a causal convolution of kernel size $j$ and $W_h$ is a linear projection. Then, the expected output of the Mamba block satisfies
\begin{equation}
\mu_{F,\text{Mamba}} = \frac{\sigma^2}{4n} \sum_{t=1}^n \mathrm{diag}\big(C_t \bar{B}_t W_{h'_0} W_z^\top\big),
\end{equation}
where $C_t$, $\bar{B}_t$, $W_{h'}$, and $W_z$ are Mamba parameters computed via the SiLU approximation and SSM recursion.
\end{theorem}
\vspace{1mm}
\begin{theorem}\label{theorem:thrm-4}
For $1 \le t, j \le n$, let $M_{t,j} = C_t \left( \prod_{k=j+1}^t \bar{A}_k \right) \bar{B}_j W_{h'}$. Then, the trace of the covariance of the Mamba block is
\begin{equation}
\begin{split}
\operatorname{Tr}(\Sigma_{F,\text{Mamba}}) &= \frac{\sigma^4}{4} \Big\{ \operatorname{Tr}\big[(W_z W_z^\top) \circ S_t\big] \\
&\quad + \operatorname{Tr}\big[(M_{t,t} W_z^\top) \circ (W_z M_{t,t}^\top)\big] \Big\}
\end{split}
\end{equation}
where $S_t = \sum_{j=1}^{t-1} M_{t,j} M_{t,j}^\top$
and $\bar{A}_k$ is a Mamba parameter.
\end{theorem}
\vspace{1mm}
\begin{theorem} \label{theorem:thrm-5}
In addition to Assumptions~\ref{asm:A1}--\ref{asm:A4}, we impose an extra condition to bound the Mamba parameters. Specifically, we assume that Mamba is uniformly contractive: there exists $\rho \in (0,1)$ such that $\|\bar{A}_t\|_2 \le \rho$ for all $t$, and that the operator and Frobenius norms are bounded as $\|C_t\|_2 \le c$, $\|\bar{B}_t\|_2 \le b$, $\|W_{h'}\|_2 \le h$, and $\|W_z\|_F \le z$.
Under these conditions, it holds that
\begin{equation}
\mathbb{E}[St^2_{\mathrm{Trans}}] > \mathbb{E}[St^2_{\mathrm{Mamba}}] \quad \forall n \ge 1.
\end{equation}
\end{theorem}
\noindent Proofs for the above propositions and theorems are provided in \Cref{appdx:theory-proofs}. In summary, our analysis formalizes the architectural tendencies observed empirically, explaining why TBMs exhibit early oversmoothing and abrupt final-layer shifts, while SSMs retain variability longer and converge later (Findings \ref{F1}–\ref{F4}). The bounds in \Cref{theorem:thrm-5} further show how model scale and parameterization influence these effects (Finding \ref{F5}), confirming that these dynamics stem from fundamental architectural properties rather than training artifacts.
\section{Conclusions}
In this work, we present a unified, token- and layer-wise comparison of representation flow in TBMs and SSMs, revealing opposing oversmoothing trajectories: TBMs homogenize early then recover, while SSMs preserve early token uniqueness but converge to homogenization deeper. This divergence explains why both architectures perform well on language tasks yet fail differently as context grows, reframing the question from ``which architecture is better'' to how each routes and re-encodes information across layers. Our analysis also provides practical diagnostics, showing that intermediate layers often contain the most usable knowledge, and that layerwise probes, combined with Sm/St and inter-token similarity metrics,
can detect failures before deployment. These insights motivate targeted interventions, including intermediate supervision, contrastive regularization, contractive constraints, and hybrid SSM-TBM designs, and inform stable scaling strategies. Finally, we introduce a reproducible toolkit of similarity measures, stability bounds, and layerwise analysis frameworks, offering a foundation for future work on multi-modal tasks, hybrid architectures, and robust long-context modeling.
\section*{Limitations}
While our study sheds light on representation flow in SSMs and TBMs, several limitations should be acknowledged. First, our analysis focuses on two tasks (MDQA and KVPR) that capture important aspects of long-context reasoning but may not fully represent the breadth of applications where these models are deployed. Extending the evaluation to a wider set of tasks would help assess the robustness of our observations. Second, the metrics we employ—cosine similarity, CKA, and probing accuracy—offer informative but partial perspectives on internal dynamics. Complementary techniques, such as circuit-level analysis, could provide a more detailed account of the mechanisms driving the observed patterns. {Finally, although we evaluated four representative models under controlled settings with shared training data to ensure fairness, this limited scope may restrict generalizability.} Future work could examine a broader range of architectures, training regimes, and data sources to test the consistency of our findings.
Taken together, these considerations point to the need for continued investigation, while our results provide an initial step toward systematically understanding the trade-offs in representation evolution between SSMs and TBMs.
\section*{Ethics Statement}
Our research does not involve human subjects, personal data collection, or direct societal applications that could cause harm. However, we acknowledge that our findings about representation flow and layer-wise performance could potentially be misused for adversarial purposes, though they primarily enable beneficial applications such as more efficient model design. Our analysis uses models trained on the Pile dataset, which may contain societal biases that could be reflected in the representation patterns we observe, and future work should consider how architectural differences interact with bias mitigation strategies. While this research contributes to fundamental understanding that may inform more efficient language models with societal benefits, it also raises broader considerations about computational resource concentration and the responsible development of increasingly capable AI systems.
\section*{Significance of Our Findings}
Our analysis uncovers a fundamental divergence in how SSMs and TBMs propagate information across depth, revealing an architecture-level trade-off that reframes long-context modeling. By showing that TBMs and SSMs follow opposite oversmoothing trajectories—TBMs exhibit early homogenization followed by late recovery, while SSMs preserve early uniqueness but suffer late homogenization—we explain why architectures that both perform well on language tasks nonetheless fail in systematically different ways as context grows. This insight transforms debates about "which family is better" into the more productive question of how each family routes and re-encodes information across depth.
Our findings provide concrete diagnostics for common failure modes. The empirical and theoretical link between layerwise representation flow and probing accuracy explains why final-layer readouts can be misleading and why intermediate layers often contain the most usable knowledge.
This enables a reliable diagnostic pipeline: layerwise probes combined with Sm/St and inter-token similarity metrics can detect whether a model will lose token-level detail or over-compress context before deployment.
These insights point to immediate, practical interventions. Because representational collapse occurs at predictable depths, we can target fixes precisely: intermediate supervision, contrastive regularization to maintain token distinctiveness, contractive constraints for SSM dynamics, and hybrid architectures that route early processing through SSMs and global reconfiguration through attention. Our analysis also informs scaling strategies, showing that larger state dimensions can destabilize representations, providing principled guidance for when and where to scale capacity.
Finally, our work establishes a new experimental toolkit. The Sm/St similarity measures, comparative stability bounds, and layerwise analysis framework provide the community with reproducible metrics and theory-backed baselines for rigorous architecture comparison, opening avenues for provable regularizers, architecture-aware training, and benchmarks tailored to layerwise information retention.
\section*{Future Directions}
Our findings open several promising directions. First, while our analysis focused on text-based sequence modeling, it would be valuable to extend the framework to multimodal domains such as video or speech, where long-context fidelity is equally critical. Second, integrating our similarity measures into the training loop may enable adaptive regularization schemes that detect and counteract oversmoothing in real time. Third, hybrid architectures that combine the local retention of SSMs with the global reconfiguration capacity of TBMs remain largely unexplored; systematic exploration of such designs could yield models with both efficiency and fidelity. Finally, future research could investigate scaling laws under these new diagnostics, asking how model size, depth, and state dimensionality interact with representation collapse across architectures. Together, these directions point toward a principled roadmap for developing the next generation of long-context models.
\section*{Acknowledgements}
We thank members in NAIL (NTU) and WING (NUS) labs for their valuable feedback. Do Xuan Long is supported by the A$^*$STAR Computing and Information Science Scholarship.
% Entries for the entire Anthology, followed by custom entries
% \bibliographystyle{acl_natbib}
\bibliography{anthology,custom}
\newpage
\onecolumn
\appendix
\section{Mathematical Derivations}
\subsection{Mathematical Setups}
The metric \( \text{St}^2 \) measuring the variance of layer representations with \( L=1 \) is defined as:
\begin{equation} \label{eqssm:appx-eq4}
\text{St}^2 \propto ||h^{(1)} - h^{(0)}||^2
\end{equation}
where \( h^{(l)} \) is the representation at layer \( l \). Our goal is to compute the expected value \( \mathbb{E}[\text{St}^2] \). For both Transformers and Mamba, the representation evolves as:
\begin{equation}
h^{(1)} = h^{(0)} + F(h^{(0)})
\end{equation}
From our assumptions in \Cref{subsec:theory-setups}, we have:
\begin{itemize}
\item \( h^{(0)} \), \( h^{(1)} \) are Gaussian matrices, $\mathbb{E}[h^{(0)}] = 0$ and $\text{Cov}(h^{(0)}) = \Sigma$. \( h_t \sim N(0, \sigma^2 I) \).
\item \( F\) is the update at layer \( 1 \) with deterministic parameters. We treat \( F(x) \) as random matrices with {\( \mathbb{E}[F] = \mu_F \) and \( \text{Cov}(F) = \Sigma_F \)}.
\item \( h^{(0)}_j \) and \( h^{(0)}_t \) are independent for \( j \neq t \) and \( \mathbb{E}[h^{(0)}_j] = 0 \) for all \( j \), and $\mathbb{E}[h_ih_i^T] = \sigma^2I_d$.
\item \( \mathbb{E}[h^{(0)}_j (h^{(0)}_t)^T] = 0 \) for \( j \neq t \).
\end{itemize}
\paragraph{Objective:} Under some mild assumptions, we simplify and compare St$^2$ from both architectures.
\subsection{Proofs} \label{appdx:theory-proofs}
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\begin{proof}[Proof of Proposition~\ref{prop:prop-common-formula}]
Substitute into \Cref{eqssm:appx-eq4}:
\[
h^{(1)} - h^{(0)} = (h^{(0)} + F(h^{(0)})) - h^{(0)} = F(h^{(0)})
\]
Thus:
\[
\text{St}^2 \propto ||h^{(1)} - h^{(0)}||^2 = ||F(h^{(0)})||^2
\]
Denote \(h^{(0)} \in \mathbb{R}^{n \times d}\), and \(F: \mathbb{R}^{n \times d} \to \mathbb{R}^{n \times d}\), the squared Frobenius norm is:
\[
||F(h^{(0)})||^2 = \sum_{i=1}^n \sum_{j=1}^d |F(h^{(0)})_{ij}|^2
\]
The expectation is:
\[
\mathbb{E}[||F(h^{(0)})||^2] = \mathbb{E}\left[ \sum_{i,j} |F(h^{(0)})_{ij}|^2 \right] = \sum_{i,j} \mathbb{E}[|F(h^{(0)})_{ij}|^2]
\]
Since \(F\) has deterministic parameters, the randomness in \(F(h^{(0)})\) comes from \(h^{(0)}\). The assumption states \(\mathbb{E}[F(h^{(0)})] = \mu_F\), where \(\mu_F \in \mathbb{R}^{n \times d}\), so:
\[
\mathbb{E}[F(h^{(0)})_{ij}] = (\mu_F)_{ij}
\]
The variance of \(F(h^{(0)})\) is given by a covariance matrix \(\Sigma\), but for a matrix-valued random variable, we interpret \(\mathbb{E}[|F(h^{(0)})_{ij}|^2] = \text{Var}(F(h^{(0)})_{ij}) + |\mathbb{E}[F(h^{(0)})_{ij}]|^2\). Thus:
\[
\mathbb{E}[|F(h^{(0)})_{ij}|^2] = \text{Var}(F(h^{(0)})_{ij}) + |(\mu_F)_{ij}|^2
\]
Summing over all elements:
\[
\mathbb{E}[||F(h^{(0)})||^2] = \sum_{i,j} \left( \text{Var}(F(h^{(0)})_{ij}) + |(\mu_F)_{ij}|^2 \right).
\]
Thus, we have:
\begin{equation}\label{eqssm:eqssm100}
\mathbb{E}[\text{St}^2] \propto \mathbb{E}\left[ ||F(h^{(0)})||^2 \right] = \text{Tr}(\Sigma_F) + ||(\mu_F)||^2
\end{equation}
\end{proof}
From \Cref{eqssm:eqssm100}, for Transformers:
\begin{equation}
\mathbb{E}[\text{St}^2_{\text{Trans}}] \propto \text{Tr}(\Sigma_{F,\text{Trans}}) + ||\mu_{F,\text{Trans}}||^2
\end{equation}
where:
\begin{itemize}
\item
$F_{\text{Trans}} = \text{Attn}(h^{(0)}) + \text{FFN}(\tilde{h}^{(0)}), \quad \tilde{h}^{(0)} = h^{(0)} + \text{Attn}(h^{(0)})$
\item $\mu_{F,\text{Trans}} = \mathbb{E}[F_{\text{Trans}}] = \mathbb{E}[\text{Attn}(h^{(0)})] + \mathbb{E}[\text{FFN}(\tilde{h}^{(0)})]$
\item $\Sigma_{F,\text{Trans}} = \text{Cov}(\text{Attn} + \text{FFN})$
\end{itemize}
\noindent For Mamba (a state-space model with selective mechanism):
\begin{equation}
\mathbb{E}[\text{St}^2_{\text{Mamba}}] \propto \text{Tr}(\Sigma_{F,\text{Mamba}}) + ||\mu_{F,\text{Mamba}}||^2
\end{equation}
where:
\begin{itemize}
\item
$F_{\text{Mamba}} = \text{S6}(h') \circ z(h^{(0)}), \quad z = \text{SiLU}(\text{Linear}(h^{(0)})), \quad h' = \text{SiLU}(\text{Conv1D}(\text{Linear}(h^{(0}))$
\item $\mu_{F,\text{Mamba}} = \mathbb{E}[\text{S6} \circ z] = \mathbb{E}[\text{S6}] \cdot \mathbb{E}[z] + \text{Cov}(S6, z)$
\item $\Sigma_{F,\text{Mamba}} = \text{Cov}(\text{S6} \circ z)$
\end{itemize}
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\begin{proof}[Proof of \Cref{prop:prop-att-odd}]
Consider the effect of negating the input: if we replace \( x \) with \( -x \):
\( Q_{\text{new}} = (-x) W_Q = -x W_Q = -Q, \quad
K_{\text{new}} = (-x) W_K = -x W_K = -K, \quad
V_{\text{new}} = (-x) W_V = -x W_V = -V \)
The new attention scores become:
\[
\frac{(-Q) (-K)^T}{\sqrt{d_k}} = \frac{(-Q) (-K^T)}{\sqrt{d_k}} = \frac{Q K^T}{\sqrt{d_k}}
\]
Thus, the softmax remains unchanged:
\[
\text{softmax}\left( \frac{(-Q) (-K)^T}{\sqrt{d_k}} \right) = \text{softmax}\left( \frac{Q K^T}{\sqrt{d_k}} \right)
\]
The output is:
\[
\text{Attn}(-x) = \text{softmax}\left( \frac{Q_{new} K_{new}^T}{\sqrt{d_k}} \right) (V_{new}) = -\text{softmax}\left( \frac{Q K^T}{\sqrt{d_k}} \right) V = -\text{Attn}(x)
\]
This shows that the attention mechanism $Attn(x)$ is odd.
\end{proof}
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\begin{proof}[Proof of \Cref{propssm:prop10}]
Because \( \text{Attn}(h^{(0)}) \) is odd and \( h^{(l)} \) is symmetrically distributed, we have $\mathbb{E}[\text{Attn}(h^{(0)})] = \mathbb{E}[\text{Attn}(h^{(1)})] = 0$.
\end{proof}
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\begin{proof}[Proof of \Cref{theorem:thrm-1}]
We have $\mu_{F,\text{Trans}} = \mathbb{E}[F_{\text{Trans}}] = \mathbb{E}[\text{Attn}(h^{(0)})] + \mathbb{E}[\text{FFN}(\tilde{h}^{(0)})]$ and \( \mu_{h^{(0)}} = 0\). From \Cref{propssm:prop10}, we have $\mathbb{E}[\text{Attn}(h^{(0)})] = 0$ and $\mathbb{E}[\tilde{h}^{(0)})] = \mathbb{E}[{h}^{(0)} + \text{Attn}({h}^{(0)})] = 0$.
We have $\text{FFN}(\tilde{h}^{(0)}) = W_2 (\text{GELU}(W_1 \tilde{h}^{(0)} + b_1)) + b_2$ and assume that $W_1$, $W_2$, $b_1$, $b_2$ are independent, we approximate $\text{GELU}(a) \approx a\cdot\sigma(1.702a)$ following \citet{hendrycks2016gaussian}, and further approximate $a\cdot\sigma(1.702a) \approx \frac{1}{2}a$ via Taylor expansion via removing higher-order terms, we have:
\begin{align} \label{eqssm:eqssm11}
\mu_{F,Trans} = \mathbb{E}[\text{FFN}(\tilde{h}^{(0)})] &= \mathbb{E}[W_2 \text{GELU}(W_1 \tilde{h}^{(0)} + b_1) + b_2] \\
&\approx \mathbb{E}[W_2 \frac{1}{2}(W_1 \tilde{h}^{(0)} + b_1) + b_2] = \frac{1}{2}W_2 b_1 + b_2
\end{align}
\end{proof}
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\begin{proof}[Proof of \Cref{theorem:thrm-2}]
We need to compute \( \text{Tr}(\Sigma_{F,\text{Trans}}) \) for the Transformer-based model's update function \( F_{\text{Trans}}(h^{(0)}) = \text{Attn}(h^{(0)}) + \text{FFN}(\tilde{h}^{(0)}) \). We have:
\begin{align}
F(h^{(0)}_t) &= \text{Attn}(h^{(0)}_t) + \text{FFN}(\tilde{h}^{(0)}_t) ß\approx \text{Attn}(h^{(0)}_t) + \frac{1}{2}W_2(W_1 \tilde{h}^{(0)}_t + b_1) + b_2 \\
&= \text{Attn}(h^{(0)}_t) + \frac{1}{2}W_2 W_1 (h^{(0)_t} + \text{Attn}(h^{(0)}_t)) + \frac{1}{2}W_2 b_1 + b_2 \\
&= \left(I + \frac{1}{2}W_2 W_1\right) \text{Attn}(h^{(0)}_t) + \frac{1}{2}W_2 W_1 h^{(0)}_t + \frac{1}{2}W_2 b_1 + b_2 \\
\therefore F_t &= T_1 \text{Attn}_t + T_2 h_t + \mu
\end{align}
where $T_1 = I + \frac{1}{2}W_2 W_1$, $T_2 = \frac{1}{2}W_2 W_1$, and $\mu = \frac{1}{2}W_2 b_1 + b_2$. We also denote $F_t, \text{Attn}_t, h_t$ refers to $F(h^{(0)}_t), \text{Attn}(h^{(0)}_t), h^{(0)}_t$ respectively for ease of notation. We aim to compute $\Sigma_{F_t} = \mathbb{E}[F_t F_t^T] - \mathbb{E}[F_t]\mathbb{E}[F_t]^T$. Thus:
\begin{align}
F_t F_t^T &= (T_1 \text{Attn}_t + T_2 h_t + \mu) (T_1 \text{Attn}_t + T_2 h_t + \mu)^T \\
\therefore\quad \mathbb{E}[F_t F_t^T] &= T_1 \mathbb{E}[\text{Attn}_t\text{Attn}_t^T] T_1^T + T_2 \mathbb{E}[h h^T] T_2^T + \mu \mu^T \\
&= T_1 \Sigma_{\text{Attn}_t} T_1^T + T_2 \Sigma_{h_t} T_2^T + \mu \mu^T
\end{align}
by using $\mathbb{E}[\text{Attn}_t] = \mathbb{E}[h_t] = 0$ and $\text{Attn}_t$ depends on $h_t$, we have:
\begin{align}
\Sigma_{F_t} &= \mathbb{E}[F_t F_t^T] - \mathbb{E}[F_t]\mathbb{E}[F_t]^T \\
&= T_1 \Sigma_{\text{Attn}_t} T_1^T + T_2 \Sigma_{h_t} T_2^T + T_1 \mathrm{Cov}(\text{Attn}_t, h_t) T_2^T + T_2 \mathrm{Cov}(\text{Attn}_t, h_t)^T T_1^T
\end{align}
Now we need to compute the covariance of the attention output, $\Sigma_{\text{Attn}_t}$. Given $h^{(0)} \in \mathbb{R}^{n \times d}$, the attention output for a single token $t$ is:
\begin{equation}
\text{Attn}_t = \sum_{j=1}^n a_{tj} h_j^{(0)} W_V, \quad a_{tj} = \text{softmax} \Bigg( \frac{(h_t^{(0)} W_Q)(h_j^{(0)} W_K)^T}{\sqrt{d_k}} \Bigg)
\end{equation}
As $\mathbb{E}[\text{Attn}_t] = 0$, the covariance matrix is:
\begin{align}
\Sigma_{\text{Attn}_t} &= \mathbb{E}[\text{Attn}_t \text{Attn}_t^T] = \mathbb{E} \Bigg[ \Bigg( \sum_{j=1}^n a_{tj} h_j^{(0)} W_V \Bigg) \Bigg( \sum_{k=1}^n a_{tk} h_k^{(0)} W_V \Bigg)^T \Bigg] \\
&= \mathbb{E} \Bigg[ \sum_{j, k=1}^n a_{tj} a_{tk} (h_j^{(0)} W_V)(h_k^{(0)} W_V)^T \Bigg]
\end{align}
Since $h_j^{(0)}$ and $h_k^{(0)}$ are independent for $j \neq k$, and $\mathbb{E}[h_j^{(0)}] = \mathbb{E}[h_k^{(0)}] = 0$, the cross-terms vanish unless $j = k$:
\begin{align}
\Sigma_{\text{Attn}_t} &= \sum_{j=1}^n \mathbb{E}[a_{tj}^2] \mathbb{E}[(h_j^{(0)} W_V)(h_k^{(0)} W_V)^T] = \sum_{j=1}^n \mathbb{E}[a_{tj}^2] (W_V \Sigma_{h^{(0)}} W_V^T) \\
&= \sum_{j=1}^n \mathbb{E}[a_{tj}^2] (\sigma^2 W_V W_V^T) = \sigma^2 \left( \sum_{j=1}^n \mathbb{E}[a_{tj}^2] \right) W_V W_V^T
\end{align}
The attention weights $a_{tj}$ depend on the softmax output. For large $d_k$, the dot product $(h^{(0)}_t W_Q)(h^{(0)}_j W_K)^T / \sqrt{d_k}$ is approximately Gaussian. Assuming standard initialization $(W_Q, W_K \sim \mathcal{N}(0, 1/\sqrt{d}))$, the variance of the dot product before scaling is:
\begin{align}
\text{Var}((h_t^{(0)}W_Q)(h_j^{(0)}W_K)^T) &\approx \sigma^2 \cdot \frac{1}{d} \cdot \sigma^2 \cdot d = \sigma^4 \\
\therefore\quad \text{Var}\Bigg( \frac{(h_t^{(0)}W_Q)(h_j^{(0)}W_K)^T)}{\sqrt{d_k}} \Bigg) &\approx \frac{\sigma^4}{\sqrt{d_k}}
\end{align}
The softmax normalizes these scores, and for symmetrically distributed inputs, $\mathbb{E}[a_{tj}] \approx 1/n$. The variance of the softmax output is small, and we approximate $\mathbb{E}[a_{tj}^2] \approx 1/n^2$. And by assuming that $W_V W_V^T \approx I_d$ (standard initialization ensures $\text{Tr}(W_V W_V^T) \approx d$). The covariance $\Sigma_{\text{Attn}_t}$ becomes:
\begin{equation}
\Sigma_{\text{Attn}_t} \approx \sigma^2 \left(\sum_{j=1}^n \frac{1}{n^2}\right) W_V W_V^T \approx \frac{\sigma^2}{n} I_d
\end{equation}
Now we need to compute $\mathrm{Cov}(\text{Attn}_t, h_t) = \mathbb{E}[\text{Attn}_t h_t] = \sum_{j=1}^t \mathbb{E}[a_{tj} v_j h_t^T]$. For $j \neq t$, $v_j$ is independent of $h_t$ and zero-mean so those terms are vanish. Thus, with $j = t$ and the same ``mean-field'' decoupling, we have:
\begin{align}
\mathrm{Cov}(\text{Attn}_t, h_t) &\approx \mathbb{E}[a_{tt}] \mathbb{E}[v_t h_t^T] \\
&= \frac{1}{n} W_V \mathbb{E}[h_t h_t^T] \\
&= \frac{\sigma^2}{n} W_V
\end{align}
Put everything together, we have the covariance trace for token $t$ as:
\begin{align}
\text{Tr}(\Sigma_t) &= \text{Tr}(T_1 \Sigma_{\text{Attn}_t} T_1^T) + \text{Tr}(T_2 \Sigma_{h_t} T_2^T) + 2 \text{Tr}(T_1 \mathrm{Cov}(\text{Attn}_t, h_t) T_2^T) \\
&\approx \frac{\sigma^2}{n} \text{Tr}(T_1 T_1^T) + \sigma^2 \text{Tr}(T_2 T_2^T) + \frac{2\sigma^2}{n} \text{Tr}(T_1 W_V T_2^T)
\end{align}
Sum over $t$ across $n$ tokens, we finally have:
\begin{equation}
\text{Tr}(\Sigma_{F,\text{Trans}}) = \sigma^2 \text{Tr}(T_1 T_1^T) + n \sigma^2 \text{Tr}(T_2 T_2^T) + 2\sigma^2 \text{Tr}(T_1 W_V T_2^T)
\end{equation}
\end{proof}
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\begin{proof}[Proof of \Cref{theorem:thrm-3}]
We have $\mu_{F,\text{Mamba}} = \mathbb{E}[\text{S6} \circ z]$. The S6 layer can be represented as a data-controlled linear operator \citep{poli2023hyena,ali2024hidden}. Specifically, for a sequence of inputs \( h' = [h'_1, h'_2, \ldots, h'_n]^T \), the outputs \( o = [o_1, o_2, \ldots, o_n]^T \) are computed through the following recursive equations. Assuming that \( s_0 = 0 \):
\begin{align}
s_1 &= \bar{B}_1 h'_1, \quad o_1 = C_1 \bar{B}_1 h'_1, \\
s_2 &= \bar{A}_2 \bar{B}_1 h'_1 + \bar{B}_2 h'_2, \quad o_2 = C_2 \bar{A}_2 \bar{B}_1 h'_1 + C_2 \bar{B}_2 h'_2,
\end{align}
These equations define \( s_t \) and \( h^{(1)}_t \) recursively using matrices \( \bar{A}_t \), \( \bar{B}_t \), and \( C_t \), applied to input vectors \( h'_t \). The general form, given in Equation (10), is:
\begin{align}\label{eqssm:eqssm200}
s_t &= \sum_{j=1}^t \left( \prod_{k=j+1}^t \bar{A}_k \right) \bar{B}_j h'_j, \quad o_t = C_t \sum_{j=1}^t \left( \prod_{k=j+1}^t \bar{A}_k \right) \bar{B}_j h'_j.
\end{align}
In matrix form, this can be expressed as:
\begin{equation}
o = \hat{\alpha} h',
\end{equation}
where \( \hat{\alpha} \) is the matrix:
\begin{equation}
\hat{\alpha} = \begin{bmatrix}
C_1 \bar{B}_1 & 0 & 0 & \cdots & 0 \\
C_2 \bar{A}_2 \bar{B}_1 & C_2 \bar{B}_2 & 0 & \cdots & 0 \\
\vdots & \vdots & \ddots & \ddots & \vdots \\
C_n \bar{A}_n \bar{A}_{n-1} \cdots \bar{A}_2 \bar{B}_1 & C_n \bar{A}_n \bar{A}_{n-1} \cdots \bar{A}_3 \bar{B}_2 & \cdots & C_n \bar{B}_n
\end{bmatrix}.
\end{equation}
The element at row \( i \) and column \( j \) of \( \hat{\alpha} \), as specified in Equation (12), is computed as:
\begin{equation}
\hat{\alpha}_{i,j} =
\begin{cases}
C_i \left( \prod_{k=j+1}^i \bar{A}_k \right) \bar{B}_j & \text{if } i \geq j, \\
0 & \text{if } i < j.
\end{cases}
\end{equation}
This matrix \( \hat{\alpha} \in \mathbb{R}^{n \times n} \) is a function of the input and the parameters \( \bar{A}, \bar{B}, \bar{C}, \bar{\Delta} \), encapsulating the data-controlled linear transformations applied by the S6 layer at layer \( l \).
Next, we have $z = \text{SiLU}(\text{Linear}(h^{(0)}))$. Assume that \(\text{Linear}(h^{(0)}) = W_z h^{(0)}\), where \(W_z \in \mathbb{R}^{d \times ed}\) is a weight matrix and $e$ is the expansion factor (typically 2). Denote \(u = \text{Linear}(h^{(0)}) = W_z h^{(0)}\), by approximating SiLU function as $\text{SiLU}(x) \approx \frac{x}{2}$ via Taylor expansion, we have $z_{i} \approx \frac{u_{i}}{2}$. In other words, we have $ \mathbb{E}[o_{i} z_{i}] \approx \mathbb{E}\left[ o_{i} \frac{u_{i}}{2} \right] = \frac{1}{2} \mathbb{E}[o_{i} u_{i}]$. Computing \(\mathbb{E}[o \circ z]\), we have:
\begin{equation}
\mathbb{E}[o \circ z] = \mathbb{E}\left[ \begin{bmatrix} o_{1} z_{1} \\ \vdots \\ o_{n} z_{n} \end{bmatrix} \right] = \frac{1}{2}\mathbb{E}\left[ \begin{bmatrix} o_{1} u_{1} \\ \vdots \\ o_{n} u_{n} \end{bmatrix} \right] = \frac{1}{2}\mathrm{diag}(\mathbb{E}[o u^T])
\end{equation}
Substituting \Cref{eqssm:eqssm200}, we have:
\begin{equation}
\mathbb{E}[o_t u_t^T] = \mathbb{E}\left[ C_t \sum_{j=1}^t \left( \prod_{k=j+1}^t \bar{A}_k \right) \bar{B}_j h'_j (W_z h^{(0)}_t)^T \right]
\end{equation}
Now, for $h' = \text{SiLU}(\text{Conv1D}(\text{Linear}(h^{(0)})))$, assume that the linear transformation parameter is $W_h \in \mathbb{R}^{d \times ed}$ and the causal convolution with kernel size \(K\) (typically 4) is $W_c$, for the time $t \in \{1,2,\dots,n\}$, we have:
\begin{align}
h'_t &= \text{SiLU}(\text{Conv1D}(\text{Linear}(h^{(0)}_t))) = \text{SiLU}(W_c (W_h h^{(0)}_t)) \\ &= \text{SiLU} \Bigg(\sum_{j=0}^{K-1} W_{c_j} (W_h h^{(0)}_{t-j}) \Bigg) \approx \frac{1}{2} W_{c_j} W_h \sum_{j=0}^{K-1} h^{(0)}_{t-j}
\end{align}
Define \( W_{h'_j} = W_{c_j} W_h \), where \( W_{h'_j} \) is the combined transformation matrix. Thus:
\begin{align}
\mathbb{E}[o_t u_t^T] &= \mathbb{E}\left[ C_t \sum_{i=1}^t \left( \prod_{k=i+1}^t \bar{A}_k \right) \bar{B}_i \left( \frac{1}{2} W_{h'_j} \sum_{j=0}^{K-1} h^{(0)}_{i-j} \right) (W_z h^{(0)}_t)^T \right] \\
&= \frac{1}{2} C_t \sum_{i=1}^t \sum_{j=0}^{K-1} \left( \prod_{k=i+1}^t \bar{A}_k \right) \bar{B}_i W_{h'_j} \mathbb{E}\left[h^{(0)}_{i-j} (h^{(0)}_t)^T\right] W_z^T
\end{align}
Under assumptions, in the sum over \( i = 1 \) to \( t \), the expectation \( \mathbb{E}[h^{(0)}_{i-j} (h^{(0)}_t)^T] \) is either 0 when \( i-j \neq t \), or \( \sigma^2 I_d \) when \( i-j = t \), which means \(i=t, j=0\). Thus:
\begin{align}
\mathbb{E}[o_t u_t^T] &= \frac{1}{2} C_t \sum_{i=1}^t \sum_{j=0}^{K-1} \left( \prod_{k=i+1}^t \bar{A}_k \right) \bar{B}_i W_{h'_j} \mathbb{E}\left[h^{(0)}_{i-j} (h^{(0)}_t)^T\right] W_z^T \\
&= \frac{1}{2} C_t \left( \prod_{k=i+1}^t \bar{A}_k \right) \bar{B}_t W_{h'_0} \mathbb{E}\left[h^{(0)}_t (h^{(0)}_t)^T\right] W_z^T\\
&= \frac{1}{2} C_t \left( \prod_{k=i+1}^t \bar{A}_k \right) \bar{B}_t W_{h'_0} \sigma^2 W_z^T = \frac{\sigma^2}{2} C_t \bar{B}_t W_{h'_0} W_z^T
\end{align}
and $\mu_{F,\text{Mamba}} = \mathbb{E}[\text{S6} \circ z] = \frac{1}{2}\mathrm{diag}(\mathbb{E}[o u^T]) = \frac{\sigma^2}{4n} \sum_{t=1}^n \mathrm{diag} (C_t \bar{B}_t W_{h'_0} W_z^T)$.
\end{proof}
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\begin{proof}[Proof of \Cref{theorem:thrm-4}]
We need to compute \( \text{Tr}(\Sigma_{F,\text{Mamba}}^2) \) for the Mamba model's update function \( F_{\text{Mamba}}(h^{(0)}) = o \circ z \), where \( o \) is the output of the S6 layer and \( z \) is the gating term, both dependent on the input \( h^{(0)} \).
Under the first-order nonlinearity approximation (S4), the Mamba update at time $t$ is:
\begin{align}
F(h^{(0)}_t) &= \frac{1}{2} (o_t \circ z_t) \\
&= \frac{1}{2} \mathrm{diag}(W_z h^{(0)}_t) C_t \sum_{j=1}^t \left( \prod_{k=j+1}^t \bar{A}_k \right) \bar{B}_j W_{h'} h^{(0)}_j \\
&= \frac{1}{2} \mathrm{diag}(W_z h^{(0)}_t) \sum_{j=1}^t M_{t,j} h^{(0)}_j, \quad \text{where} \quad M_{t,j} = C_t \left( \prod_{k=j+1}^t \bar{A}_k \right) \bar{B}_j W_{h'} \in \mathbb{R}^{ed \times d} \\
\mathbb{E}[F_t] &= \frac{1}{2} \left( \sum_{j=1}^t \mathbb{E}[ \mathrm{diag}(W_z h^{(0)}_t) M_{t,j} h^{(0)}_j \right)
\end{align}
where $F_t$ refers to $F(h^{(0)}_t)$ for ease of notation. For $j < t$, $h^{(0)}_j$ and $h^{(0)}_t$ are independent, so expectation is zero. For $j = t$, both terms involves $h^{(0)}_t$. Writing row $r$ of $M_{t,t}$ as $m^T_{t,t,r}$ and row $r$ of $W_z$ as $w^T_r$:
\begin{align}
\mathbb{E}[F_{t,r}] &= \frac{1}{2} \mathbb{E}[(m^T_{t,t,r} h^{(0)}_t)(w^T_r h^{(0)}_t)] = \frac{\sigma^2}{2} m^T_{t,t,r} w_r \\
\therefore\quad \mathbb{E}[F_t] &= \frac{\sigma^2}{2} \mathrm{diag}(M_{t,t} W^T_z)
\end{align}
For the second moment, use the identity $\mathrm{diag}(a) X \mathrm{diag}(b) = (ab^T) \circ X$, we can compute:
\begin{align}
F_t F^T_t &= \frac{1}{4} \mathrm{diag}(W_z h^{(0)}_t) \left( \sum_{j=1}^t M_{t,j} h^{(0)}_j \right) \left( \sum_{m=1}^t M_{t,m} h^{(0)}_m \right)^T \mathrm{diag}(W_z h^{(0)}_t) \\
&= \frac{1}{4} \left[ (W_z h^{(0)}_t) (W_z h^{(0)}_t)^T \right] \circ \left[ \left( \sum_{j=1}^t M_{t,j} h^{(0)}_j \right) \left( \sum_{m=1}^t M_{t,m} h^{(0)}_m \right)^T \right]
\end{align}
Now take expectation: (1) for $j \neq m$, cross terms vanish by independence and zero mean; (2) for $j = m \neq t$, contributions are $\sigma^4 (W_z W^T_z) \circ (\sum_{j=1}^{t-1} M_{t,j} M^T_{t,j})$; (3) for $j = m = t$, Isserlis's theorem gives an extra correction term. For $x\sim \mathcal N(0,\sigma^2 I)$ and vectors $a,b,c,d$,
\[
\mathbb E\left[(a^T x)(b^T x)(c^T x)(d^T x)\right] = \sigma^4 \left[(a \cdot b)(c \cdot d) + (a \cdot c)(b \cdot d) + (a \cdot d)(b \cdot c) \right]
\]
With $a=m_{t,t,r}$, $b=w_r$, $c=m_{t,t,s}$, $d=w_s$, the $(r,s)$-entry of $\mathbb E\big[(W_z h_t)(W_z h_t)^T\circ (M_{t,t} h_t h_t^T M_{t,t}^T)\big]$ equals
\[
\sigma^4 \left[(m_{t,t,r} \cdot w_r)(m_{t,t,s} \cdot w_s) + (m_{t,t,r} \cdot m_{t,t,s})(w_r \cdot w_s) + (m_{t,t,r} \cdot w_s)(w_r \cdot m_{t,t,s})\right]
\]
In matrix form, the three terms correspond to
\[
\sigma^4 \left[ v v^T + (M_{t,t}M_{t,t}^T) \circ (W_z W_z^T) + (M_{t,t} W_z^T)\circ (W_z M_{t,t}^T) \right]
\]
where $v = \mathrm{diag}(M_{t,t} W_z^T)$. Putting everything together, we have:
\begin{equation}
\mathbb{E}[F_t F^T_t] = \frac{\sigma^4}{4} \left[ vv^T + S_t \circ (W_z W^T_z) + (M_{t,t} W^T_z) \circ (W_z M^T_{t,t,})\right]
\end{equation}
where $S_t = \left( \sum_{j=1}^{t-1} M_{t,j} M^T_{t,j} \right)$ and $v = \mathrm{diag}(M_{t,t}W^T_z)$.
On the other hand, we have:
\begin{equation}
\mathbb{E}[F_t] \mathbb{E}[F_t]^T = \frac{\sigma^4}{4} \mathrm{diag}(M_{t,t} W^T_z) [\mathrm{diag}(M_{t,t} W^T_z)]^T = \frac{\sigma^4}{4} vv^T
\end{equation}
Subtracting this, we have the covariance as:
\begin{align}
\Sigma_{F_t} &= \text{Cov}(F_t) = \mathbb{E}[F_t F_t^T] - \mathbb{E}[F_t]\mathbb{E}[F_t]^T \\
&= \frac{\sigma^4}{4} \left[ (W_z W^T_z) \circ S_t + (M_{t,t}W^T_z) \circ (W_z M^T_{t,t}) \right]
\end{align}
Take the trace and sum over $t$:
\begin{equation}
\text{Tr}(\Sigma_{F,Mamba}) = \frac{\sigma^4}{4} \{ \text{Tr}[(W_z W^T_z) \circ S_t] + \text{Tr}[(M_{t,t}W^T_z) \circ (W_z M^T_{t,t})] \}
\end{equation}
with $S_t = \left( \sum_{j=1}^{t-1} M_{t,j} M^T_{t,j} \right)$.
\end{proof}
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\begin{proof}[Proof of \Cref{theorem:thrm-5}]
Given that $\mathbb{E}[\text{St}^2] = ||\mu_F||^2 + \mathrm{Tr}(\Sigma_F)$, we first rearrange all the equations we obtained.
\paragraph{(1) Transformer's mean.} By \Cref{theorem:thrm-1},
\[
\mu_{F,\mathrm{Trans}} \approx \tfrac{1}{2} W_2 b_1 + b_2.
\]
Using independence and centering,
\[
\mathbb{E}\!\left[\|\mu_{F,\mathrm{Trans}}\|_2^2\right] = \tfrac{1}{4}\,\mathbb{E}[\|W_2 b_1\|_2^2] + \mathbb{E}[\|b_2\|_2^2] = \alpha_T > 0
\]
\paragraph{(2) Transformer's covariance.} By \Cref{theorem:thrm-2},
\[
\mathrm{Tr}(\Sigma_{F,\mathrm{Trans}}) = \sigma^2 \mathrm{Tr}(T_1T_1^T) + n\sigma^2 \mathrm{Tr}(T_2T_2^T) + 2\sigma^2 \mathrm{Tr}(T_1 W_V T_2^T).
\]
The cross term is linear in $W_V$ and independent of $W_1,W_2$. Since $\mathbb{E}[W_V] = 0$, we have:
\[
\mathbb{E}[\mathrm{Tr}(\Sigma_{F,\mathrm{Trans}})] = \sigma^2 \mathbb{E}[\mathrm{Tr}(T_1 T_1^T)] + n\sigma^2 \mathbb{E}[\mathrm{Tr}(T_2 T_2^T)].
\]
Now computing:
\begin{align*}
\mathrm{Tr}(T_1 T_1^T) &= \mathrm{Tr}\left[ \left( I + \frac{1}{2}W_2W_1 \right) \left( I + \frac{1}{2}W_2W_1 \right)^T \right] \\
&= \mathrm{Tr}\left[ I + \frac{1}{2}W_2W_1 + \frac{1}{2}(W_2W_1)^T + \frac{1}{4}(W_2W_1)(W_2W_1)^T \right] \\
% &= \mathrm{Tr}(I) + \mathrm{Tr}(W_2W_1) + \frac{1}{4}\mathrm{Tr}[(W_2W_1)(W_2W_1)^T] \\
&= d + \mathrm{Tr}(W_2W_1) + \frac{1}{4} \| W_2W_1 \|_F^2 \\
\mathrm{Tr}(T_2 T_2^T) &= \frac{1}{4} \| W_2W_1 \|_F^2
\end{align*}
Putting together, we get:
\begin{align*}
\mathbb{E}[\mathrm{Tr}(\Sigma_{F,\mathrm{Trans}})] &= \sigma^2 \mathbb{E}[\mathrm{Tr}(T_1 T_1^T)] + n\sigma^2 \mathbb{E}[\mathrm{Tr}(T_2 T_2^T)] \\
&= \sigma^2\mathbb{E}\left[ d + \mathrm{Tr}(W_2W_1) + \frac{(n+1)}{4} \| W_2W_1 \|_F^2 \right] \\
&= \frac{\sigma^2n}{4} \mathbb{E}[\| W_2W_1 \|_F^2] + \sigma^2\left[ \mathbb{E}[\mathrm{Tr}(W_2W_1)] + \frac{1}{4} \mathbb{E}[\| W_2W_1 \|_F^2] + d \right] \\
&= \frac{\sigma^2n}{4} \beta_T + \sigma^2\gamma_T \\
% &= \sigma^2 \left[ \frac{n}{4} \mathbb{E}[\| W_2W_1 \|_F^2] + \beta_T \right]
\end{align*}
where $\beta_T > 0$ and $\gamma_T \ge d > 0$ are constants, given that $W_1, W_2$ are independent and centered.
\paragraph{(3) Mamba's mean.} By \Cref{theorem:thrm-3},
\[
\mu_{F,\mathrm{Mamba}} \approx \frac{\sigma^2}{4n}\sum_t \mathrm{diag}(C_t \bar B_t W_{h'}^0 W_z^T).
\]
Using Jensen and $\|\mathrm{diag}(X)\|_2 \le \|X\|_F$,
\[
\|\mu_{F,\mathrm{Mamba}}\|_2^2 \le \frac{\sigma^4}{16n^2}\sum_t \|C_t \bar B_t W_{h'}^0 W_z^T\|_F^2 \le \frac{\sigma^4}{16n^2} \alpha_M.
\]
where $\alpha_M > 0$ is a constant.
\paragraph{(4) Mamba's covariance.} By \Cref{theorem:thrm-4}, with $M_{t,j} = C_t(\prod_{k=j+1}^t \bar A_k)\bar B_j W_{h'}$ and $S_t = \sum_{j<t} M_{t,j} M_{t,j}^T$,
\[
\mathrm{Tr}(\Sigma_{F,\mathrm{Mamba}}) = \frac{\sigma^4}{4} \left\{\mathrm{Tr}[(W_z W_z^T) \circ S_t] + \mathrm{Tr}[(M_{t,t} W_z^T) \circ (W_z M_{t,t}^T)]\right\}.
\]
Assume that Mamba model is uniformly contractive: there exists $\rho \in (0,1)$ with $\|\bar{A_t}\|_2 \le \rho$ for all $t$, and bounded operator/Frobenius norms $\|C_t\|_2 \le c, \|\bar B_t\|_2 \le b, \|W_{h'}\|_2 \le h, \|W_z\|_F \le z$.
Using $\mathrm{Tr}(XY) \le \|X\|_F \|Y\|_F$ and contractivity bounds,
\[
\sum_{j<t}\|M_{t,j}\|_F^2 \le \frac{c^2 b^2 h^2}{1-\rho^2},
\qquad
\|M_{t,t}\|_F^2 \le c^2 b^2 h^2.
\]
Therefore,
\[
\mathrm{Tr}(\Sigma_{F,\mathrm{Mamba}}) \le \frac{\sigma^4}{4} z^2 c^2 b^2 h^2 \left (1 + \tfrac{1}{1-\rho^2}\right) = \frac{\sigma^4}{4} \beta_M.
\]
where $\beta_M > 0$ is a constant.
\paragraph{(5) Comparison and Threshold.} Collecting the bounds above, we have:
\begin{align*}
\mathbb{E}[St^2_{\mathrm{Trans}}] &= \mathbb{E}\|\mu_{F,\mathrm{Trans}}\|_2^2 + \mathbb{E}[\mathrm{Tr}(\Sigma_{F,\mathrm{Trans}}) = \frac{\sigma^2n}{4} \beta_T + \sigma^2\gamma_T + \alpha_T \\
\mathbb{E}[St^2_{\mathrm{Mamba}}] &= \mathbb{E}\|\mu_{F,\mathrm{Mamba}}\|_2^2 + \mathbb{E}[\mathrm{Tr}(\Sigma_{F,\mathrm{Mamba}}) \le \frac{\sigma^4}{16n^2} \alpha_M + \frac{\sigma^4}{4} \beta_M
\end{align*}
%%%%%%%%%%%%%%%%%%%%%%
\textbf{Cubic form.} Define the cubic polynomial
\[
Q(n) \;=\; a n^3 + b n^2 + d,
\]
with
\[
a := 4\sigma^2\beta_T,\qquad
b := 16\sigma^2\gamma_T + 16\alpha_T - 4\sigma^4\beta_M,\qquad
d := -\sigma^4\alpha_M.
\]
\textbf{Special case $n=1$:} We aim to prove that $Q(1) > 0$ in practice:
\begin{align*}
Q(1) &= 4\sigma^2\beta_T + 16\sigma^2\gamma_T + 16\alpha_T - 4\sigma^4\beta_M -\sigma^4\alpha_M \\
&= 16\alpha_T + 4\sigma^2(\beta_T + 4\gamma_T) - \sigma^4(4\beta_M + \alpha_M)
\end{align*}
Let $x = \sigma^2$, then
\begin{align*}
Q(1) = -&(4\beta_M + \alpha_M)x^2 + 4(\beta_T + 4\gamma_T)x + 16\alpha_T > 0 \\
\therefore\qquad\qquad\qquad &(4\beta_M + \alpha_M)x^2 - 4(\beta_T + 4\gamma_T)x - 16\alpha_T < 0
\end{align*}
Solve quadratic equality for $x$, we have:
\begin{align*}
x &= \frac{4(\beta_T + 4\gamma_T) \pm \sqrt{16(\beta_T + 4\gamma_T)^2 + 64(4\beta_M + \alpha_M)\alpha_T}}{2(4\beta_M + \alpha_M)} \\
&= \frac{4(\beta_T + 4\gamma_T) \pm 4\sqrt{(\beta_T + 4\gamma_T)^2 + 4(4\beta_M + \alpha_M)\alpha_T}}{2(4\beta_M + \alpha_M)} \\
&= \frac{(\beta_T + 4\gamma_T) + \sqrt{(\beta_T + 4\gamma_T)^2 + 4(4\beta_M + \alpha_M)\alpha_T}}{2\beta_M + \alpha_M/2} \\
\end{align*}
Note that we take the positive root with ``+'' because the parabola opens downward (coefficient of $x^2$ negative in original $Q(1)$) and the soluition with ``+'' gives the upper bound.
Hence, $Q(1) > 0$ if and only if:
\[
\sigma^2 \le x_{\max} := \frac{(\beta_T + 4\gamma_T) + \sqrt{(\beta_T + 4\gamma_T)^2 + 4(4\beta_M + \alpha_M)\alpha_T}}{2\beta_M + \alpha_M/2}
\]
Given that $\alpha_T > 0, \beta_T > 0, \alpha_M > 0, \beta_M > 0$, we have:
\[
\frac{ (\beta_T + 4\gamma_T) + \sqrt{ (\beta_T + 4\gamma_T)^2 + 4 (4 \beta_M + \alpha_M) \alpha_T }}{4 \beta_M + \alpha_M }
\ge \frac{ 4 \gamma_T + \sqrt{ (4 \gamma_T)^2 + 0} }{ 4 \beta_M + \alpha_M }
= \frac{ 8 \gamma_T }{ 4 \beta_M + \alpha_M }.
\]
Since $\gamma_T \ge d > 0$ and $\beta_M, \alpha_M > 0$ are small, this implies
\[
\sigma^2_{\max} \gg 1.
\]
Hence, for any typical choice of $\sigma^2 \in (0,1)$, we have
\[
\sigma^2 < 1 \ll \sigma^2_{\max} \quad \Longrightarrow \quad Q(1) > 0.
\]
\textbf{General case $n \ge 1$:} As $Q(1) = a + b + d> 0, d < 0$, then $a + b > 0$. We now compare $Q(n)$ vs $Q(1)$.
\begin{align*}
Q(n) - Q(1) &= an^3 + bn^2 - a - b \\
&= a(n-1)(n^2+n+1) + b(n-1)(n+1) \\
&= (n-1)[an^2 + (a+b)(n+1)] > 0 \qquad\forall n > 1 \\
\therefore\qquad\qquad Q(n) &> Q(1) > 0 \qquad\forall n > 1
\end{align*}
Therefore, we conclude that $Q(n) > 0 \quad\forall n \ge 1 \Leftrightarrow \mathbb{E}[St^2_{\mathrm{Trans}}] > \mathbb{E}[St^2_{\mathrm{Mamba}}] \quad\forall n \ge 1$.
\end{proof}
\subsection{Multi-layer Stability: From $L{=}1$ to Depth-$L$}\label{appdx:theory-generalization}
\paragraph{Setup and notation.}
Let $h^{(0)},h^{(1)},\dots,h^{(L)}\in\mathbb{R}^{n\times d}$ be the layer activations with
$h^{(l+1)} \!=\! F_l(h^{(l)})$ for blocks $F_l$ (Transformer or Mamba), and write layer increments
$\Delta^{(l)} \!\triangleq\! h^{(l+1)}{-}h^{(l)}$ and the layerwise mean $\bar h \!\triangleq\! \frac{1}{L+1}\sum_{l=0}^{L} h^{(l)}$.
Define the path energy $\mathcal{E}_{\mathrm{path}} \!\triangleq\! \sum_{l=0}^{L-1} \|\Delta^{(l)}\|_F^2$ and the depth-$L$ stability
\[
\mathrm{St}_L^2 \;\triangleq\; \frac{1}{nd}\cdot \frac{1}{L+1}\sum_{l=0}^{L} \|h^{(l)}-\bar h\|_F^2.
\]
\begin{assumption}[Layer-wise centering and sub-Gaussianity]\label{asm:LGauss}
Each layer input is centered and standardized by normalization (e.g., LayerNorm), so rows of $h^{(l)}$ are zero-mean, isotropic, sub-Gaussian with bounded second/fourth moments; learned affine shifts are tracked in means and do not affect covariances after centering.
\end{assumption}
\begin{assumption}[Weak dependence across tokens]\label{asm:weakdep}
Token rows form a weakly dependent process (e.g., $\alpha$-mixing or $\Psi$-weak dependence) with summable coefficients, so cross-token covariance terms are bounded by mixing coefficients times Lipschitz moduli of $F_l$.
\end{assumption}
\begin{assumption}[Block Lipschitz/contractivity]\label{asm:lipschitz}
Each block $F_l$ is (piecewise) Lipschitz with $\mathrm{Lip}(F_l)\!\le\!\Lambda_l$, and Mamba satisfies uniform contractivity for state updates with $\|\bar A_t\|\le \rho<1$ and bounded operators $\|C_t\|\le c$, $\|\bar B_t\|\le b$, $\|W_{h'}\|\le h$, $\|W_z\|_F \le z$, as in your one-layer bounds.
\end{assumption}
\begin{lemma}[Discrete Poincaré on a chain]\label{lem:dpoincare}
For the sequence $\{h^{(l)}\}_{l=0}^{L}$ it holds that
\[
\sum_{l=0}^{L}\|h^{(l)}-\bar h\|_F^2 \;\le\; \frac{1}{4\sin^2 \left( \frac{\pi}{2(L+1)} \right)}\sum_{l=0}^{L-1}\|h^{(l+1)}-h^{(l)}\|_F^2,
\]
and thus $\mathrm{St}_L^2 \;\le\; \frac{1}{4nd \sin^2 \left( \frac{\pi}{2(L+1)} \right)}\,\mathcal{E}_{\mathrm{path}}$.
\end{lemma}
\begin{lemma}[Propagation of one-layer surrogates]\label{lem:propagate}
Write $\Delta^{(l)} \!=\! F_l(h^{(l)}){-}h^{(l)}$. Then for any $0\!\le\! l\!\le\! L{-}1$,
\[
\|h^{(l+1)}-h^{(l)}\|_F \;\le\; \Big(\prod_{k=l+1}^{L}\Lambda_k\Big)\,\|F_l(h^{(l)})\|_F,
\]
and in SSM blocks the product admits $\prod_{k=l+1}^{L}\Lambda_k \!\le\! C\,\rho^{L-l}$ by uniform contractivity.
\end{lemma}
\noindent\textbf{Remark.} The Lipschitz constants $\{\Lambda_l\}$ control how single-layer surrogates propagate to later layers.
If $\prod_{k=l+1}^L\Lambda_k$ remains bounded by a modest constant, the per-layer bounds in Lemma~\ref{lem:propagate} give tight control on path increments. For SSMs, uniform contractivity $\|\bar A_t\|\le\rho<1$ yields the geometric bound
$\prod_{k=l+1}^L\Lambda_k \le C\,\rho^{\,L-l}$ used below.
\begin{proof}
By Lipschitz continuity of the blocks, for any $r>l$,
\[
\|h^{(r)}-h^{(r-1)}\|_F
= \|F_{r-1}(h^{(r-1)})-F_{r-1}(h^{(r-2)})\|_F
\le \Lambda_{r-1}\|h^{(r-1)}-h^{(r-2)}\|_F,
\]
and iterating this bound from $r=l+1$ up to $r=L$ gives
\[
\|h^{(l+1)}-h^{(l)}\|_F \le \Big(\prod_{k=l+1}^L \Lambda_k\Big)\,\|F_l(h^{(l)})\|_F.
\]
The SSM contractivity statement follows by replacing $\Lambda_k$ with $\rho$ for SSM blocks.
\end{proof}
\begin{assumption}[Uniform one-layer gap - quantitative form]\label{asm:gap}
There exists $\delta>0$ such that under matched capacity/initialization and Assumptions~\ref{asm:LGauss}--\ref{asm:lipschitz},
\[
\mathbb{E}\Big[\|F_l(h^{(l)})\|_F^2\Big]_{\mathrm{Trans}}
\ge (1+\delta)\,
\mathbb{E}\Big[\|F_l(h^{(l)})\|_F^2\Big]_{\mathrm{Mamba}}
\quad\text{for all }l\in\{0,\dots,L-1\}.
\]
\end{assumption}
\begin{theorem}[Depth-$L$ stability ordering]\label{thm:depthL}
Under Assumptions~\ref{asm:LGauss}–\ref{asm:gap}, the expected path energies obey
\[
\mathbb{E}\!\left[\mathcal{E}_{\mathrm{path}}\right]_{\mathrm{Trans}}
\;\ge\;
\mathbb{E}\!\left[\mathcal{E}_{\mathrm{path}}\right]_{\mathrm{Mamba}},
\]
and the depth-$L$ stabilities satisfy
\[
\mathbb{E}\!\left[\mathrm{St}_L^2\right]_{\mathrm{Trans}}
\;\ge\;
\mathbb{E}\!\left[\mathrm{St}_L^2\right]_{\mathrm{Mamba}}.
\]
\end{theorem}
\begin{proof}
From Lemma~\ref{lem:propagate} we have $\|\Delta^{(l)}\|_F^2\le \kappa_l^2\|F_l(h^{(l)})\|_F^2$ where $\kappa_l:=\prod_{k=l+1}^L\Lambda_k$.
Taking expectations and summing over $l$ yields
\[
\mathbb{E}[\mathcal{E}_{\mathrm{path}}]
\le \sum_{l=0}^{L-1}\kappa_l^2\,\mathbb{E}\big[\|F_l(h^{(l)})\|_F^2\big].
\]
By Assumption~\ref{asm:gap} (quantitative form) each term for the Transformer dominates the Mamba counterpart by factor $(1+\delta)$,
hence $\mathbb{E}[\mathcal{E}_{\mathrm{path}}]_{\mathrm{Trans}}\ge\mathbb{E}[\mathcal{E}_{\mathrm{path}}]_{\mathrm{Mamba}}$,
up to the (uniformly controlled) factors $\{\kappa_l\}$. Finally apply Lemma~\ref{lem:dpoincare} and normalize by $nd(L+1)$ to obtain the stated ordering for $\mathrm{St}_L^2$.
\end{proof}
% \begin{proof}[Proof sketch]
% By Lemma~\ref{lem:propagate} and Assumption~\ref{asm:gap}, summing over $l$ gives
% $\mathbb{E}[\mathcal{E}_{\mathrm{path}}]_{\mathrm{Trans}} \!\ge\! \sum_l \kappa_l\,\mathbb{E}[\|F_l(h^{(l)})\|_F^2]_{\mathrm{Trans}}
% \!\ge\! \sum_l \kappa_l\,\mathbb{E}[\|F_l(h^{(l)})\|_F^2]_{\mathrm{Mamba}}
% \!\ge\! \mathbb{E}[\mathcal{E}_{\mathrm{path}}]_{\mathrm{Mamba}}$ with layer constants $\kappa_l$ that are uniformly controlled, and contractive for SSM; then apply Lemma~\ref{lem:dpoincare} in expectation to obtain the stability inequality.
% \end{proof}
\begin{corollary}[Transformer vs. Mamba]\label{cor:TvM}
If each Transformer block’s one-layer bound dominates the Mamba counterpart as established in your L=1 analysis, then the depth-$L$ ordering $\mathbb{E}[\mathrm{St}_L^2]_{\mathrm{Trans}} \!\ge\! \mathbb{E}[\mathrm{St}_L^2]_{\mathrm{Mamba}}$ holds for all $L\!\ge\!1$, with constants depending on $\{\Lambda_l\}$ and SSM contractivity $\rho$.
\end{corollary}
\paragraph{Remarks.}
(i) The discrete Poincaré constant $(L{+}1)^2/\pi^2$ is tight for a chain and can be replaced with any equivalent spectral constant of the path graph; constants do not affect the ordering conclusion, only prefactors; (ii) Assumptions~\ref{asm:LGauss}–\ref{asm:weakdep} replace per-layer Gaussianity/independence with practical layer-wise centering and weak dependence; (iii) Learned affine shifts from normalization affect means but not covariances after centering, and can be tracked separately in $\mu_F$ terms as in your L=1 derivations.
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\section{Prompting Details}
\label{sec:prompt}
Following setup by \citet{liu-etal-2024-lost} and \citet{knowbutdonttell}, we construct key-value pairs retrieval and multi-document question answering prompting dataset.
\paragraph{Key-Value pairs retrieval (kv-pairs)}
We generate $n$ pairs of 128-bit randomly generated UUID.
\begin{tcolorbox}[colframe=gray, colback=white, colframe=black!75!white, arc=0mm]
{\small \textbf{Example Key-Value pair}} \\
{\small "7f666c61-573f-4212-a0a9-6f90d487cd4a" : "2a1d0ba0-cfe4-4df5-987a-6ee1be2c6ac0"}
\end{tcolorbox}
\noindent The $n$ kv-pairs are composed into one single JSON object. To test at ID $k$, we choose one pair as gold, insert it at ID $k$, and then construct as a prompt in the format:
\begin{tcolorbox}[colframe=gray, colback=white, colframe=black!75!white, arc=0mm]
{\small Extract the value corresponding to the specified key in the JSON object below.\\\\
JSON data:\\ \{
"key$^1$: "value$^1$",\\
“key$^2$": "value$^2$",\\
...\\
\textbf{"key$^k$": "value$^k$"},\\
...\\
"key$^n$": "value$^n$",\\
\} \\\\
Key: "key$^k$"\\Corresponding value:}
\end{tcolorbox}
\paragraph{Multi-document question answering (MDQA)}
In the $n$ document setting, we randomly select one question answer pair from the dataset by \citet{liu-etal-2024-lost}. Subsequently we retrieve the document containing this answer and mark it as gold.
\begin{tcolorbox}[colframe=gray, colback=white, colframe=black!75!white, arc=0mm]
{\small \textbf{Example retrieval}} \\
{\small \textbf{Question}: who got the first nobel prize in physics} \\
{\small \textbf{Answer}: Wilhelm Conrad Röntgen} \\
{\small \textbf{Document}: (Title: List of Nobel laureates in Physics) The first Nobel Prize in Physics was awarded in 1901 to Wilhelm Conrad Röntgen, of Germany, who received...}
\end{tcolorbox}
\noindent We then sample $n-1$ distractors, relevant documents that do not contain the answer. To test at ID $k$, we randomly shuffle the distractors and then insert the gold document at ID $k$. Example prompt with gold document at ID $k$ is like:
\begin{tcolorbox}[colframe=gray, colback=white, colframe=black!75!white, arc=0mm]
{\small Write a high-quality answer for the given question using only the provided search results (some of which might be irrelevant). \\
Document [1](Title: Asian Americans in science and technology) Prize in physics for discovery of the subatomic... \\
... \\
\textbf{Document [$k$](Title: List of Nobel laureates in Physics) The first Nobel Prize in Physics was awarded in 1901...} \\
... \\
Document [$n$] (Title: Scientist) and pursued through a unique method, was essentially in place. Ramón y Cajal won ... \\
Question: who got the first nobel prize in physics\\
Answer:
}
\end{tcolorbox}
\section{Additional Discussions}
\begin{figure}[t!]
\centering
\includegraphics[width=\columnwidth]{imgs/tokenwise_cosine_evolution_kv.png}
\caption{
Token-wise cosine similarity across layers for \texttt{GPT-Neo-2.7B} (left) and \texttt{Mamba2-2.7B} (right) on the KVPR task with $n = 2K$ tokens.
}
\label{fig:tokenwise-cosine-kv}
\end{figure}
\begin{figure}[t!]
\centering
\includegraphics[width=\columnwidth]{imgs/cka_maps_kv.png}
\caption{
The CKA between layers of \texttt{GPT-Neo-2.7B} (left) and \texttt{Mamba2-2.7B} (right) on KVPR task with $n = 2K$ tokens.
}
\label{fig:cka-kv}
\end{figure}
\begin{figure*}[t!]
\centering
\includegraphics[width=0.99\textwidth]{imgs/layerwise_probing.png}
\caption{
The layer-wise probing accuracy of TBMs (left) and SSMs (right) on MDQA task with $n = 2K$ tokens.
}
\label{fig:layerwise-acc}
\end{figure*}
\subsubsection{Probing Analysis}
% Our probing analysis explores where models encode the information that benefits the performance through
We examine whether the last-layer yields the best-performing representation.
\begin{table*}[t!]
\centering
\small{\resizebox{0.99\textwidth}{!}{
\begin{tabular}{l ccc ccc ccc ccc}
\toprule
\multirow{2}{*}{\textbf{Model}} & \multicolumn{3}{c}{$n = 300$} & \multicolumn{3}{c}{$n = 1K$} & \multicolumn{3}{c}{$n = 2K$} & \multicolumn{3}{c}{$n = 4K$} \\
\cmidrule(lr){2-4} \cmidrule(lr){5-7} \cmidrule(lr){8-10} \cmidrule(lr){11-13}
& \textbf{Prob Acc. $\uparrow$} & \textbf{Sm. $\downarrow$} & \textbf{St. $\downarrow$}
& \textbf{Prob Acc. $\uparrow$} & \textbf{Sm. $\downarrow$} & \textbf{St. $\downarrow$}
& \textbf{Prob Acc. $\uparrow$} & \textbf{Sm. $\downarrow$} & \textbf{St. $\downarrow$}
& \textbf{Prob Acc. $\uparrow$} & \textbf{Sm. $\downarrow$} & \textbf{St. $\downarrow$} \\
\midrule
\multicolumn{13}{c}{\textbf{\textit{Multi-Document Question Answering}}} \\
\midrule
GPT-Neo-2.7B & $97.5 \da{2.3}$ & $0.844$ & $3.256$
& $56.0 \da{23.4}$ & $0.858$ & $3.321$
& $43.5 \da{26.0}$ & $0.878$ & $3.386$
& - & - & - \\
Pythia-2.8B & $97.5 \da{2.1}$ & $0.245$ & $1.003$
& $57.4 \da{18.7}$ & $0.250$ & $1.033$
& $47.4 \da{18.6}$ & $0.254$ & $1.053$
& - & - & - \\
% mamba2-130m (89.1, 5.9) (29.4, 6.7) (20.4, 3.0) (16.4, 6.8)
Mamba2-130M & $89.1 \da{5.9}$ & $3.276$ & $4.864$
& $29.4 \da{6.7}$ & $3.273$ & $4.881$
& $20.4 \da{3.0}$ & $3.313$ & $4.932$
& $10.4 \da{6.7}$ & $3.410$ & $5.082$ \\
Mamba-2.8B & $97.5 \textbf{\da{1.7}}$ & $\textbf{0.167}$ & $\textbf{0.282}$
& $44.1 \da{10.5}$ & $\textbf{0.174}$ & $\textbf{0.293}$
& $27.7 \da{6.1}$ & $\textbf{0.175}$ & $\textbf{0.295}$
& $14.6 \da{9.3}$ & $\textbf{0.175}$ & $\textbf{0.293}$ \\
Mamba2-2.7B & $97.7 \da{1.8}$ & $1.964$ & $3.156$
& $41.5 \da{13.9}$ & $1.987$ & $3.155$
& $24.2 \da{10.3}$ & $2.013$ & $3.174$
& $12.7 \da{12.3}$ & $1.974$ & $3.065$ \\
\midrule
\multicolumn{13}{c}{\textbf{\textit{Key-Value Pairs Retrieval}}} \\
\midrule
GPT-Neo-2.7B & $79.9 \da{19.6}$ & $0.996$ & $3.970$
& $85.5 \da{14.1}$ & $1.057$ & $4.284$
& $62.7 \da{26.4}$ & $1.061$ & $4.287$
& - & - & - \\
Pythia-2.8B & $94.4 \textbf{\da{5.6}}$ & $0.278$ & $1.113$
& $95.2 \textbf{\da{4.8}}$ & $0.292$ & $1.177$
& $90.5 \textbf{\da{9.4}}$ & $0.300$ & $1.227$
& - & - & - \\
Mamba2-130M & $55.9 \da{29.1}$ & $3.692$ & $5.704$
& $55.2 \da{22.7}$ & $3.938$ & $6.216$
& $36.9 \da{14.3}$ & $3.950$ & $6.234$
& $30.7 \textbf{\da{6.7}}$ & $4.058$ & $6.431$ \\
Mamba-2.8B & $88.5 \da{11.4}$ & $\textbf{0.172}$ & $\textbf{0.337}$
& $78.9 \da{18.5}$ & $\textbf{0.204}$ & $\textbf{0.413}$
& $55.0 \da{26.5}$ & $\textbf{0.209}$ & $\textbf{0.425}$
& $58.7 \da{23.8}$ & $\textbf{0.203}$ & $\textbf{0.412}$ \\
Mamba2-2.7B & $74.3 \da{25.5}$ & $1.836$ & $3.249$
& $72.8 \da{26.2}$ & $1.853$ & $3.379$
& $52.2 \da{36.4}$ & $1.834$ & $3.355$
& $46.6 \da{24.0}$ & $1.811$ & $3.328$ \\
\bottomrule
\end{tabular}
}}
\caption{
This table shows the probing accuracy (\%) using the \textbf{last layer's representation}. We run all the evaluation 5 times and report the average results. $\da{x}$ is the accuracy difference between the probe trained on the last layer and on the \textbf{peak layer}. The best results are \textbf{bolded}.
}
\label{tab:results-full}
\end{table*}
\end{document}
\begin{thebibliography}{29}
\expandafter\ifx\csname natexlab\endcsname\relax\def\natexlab#1{#1}\fi
\bibitem[{Ali et~al.(2024)Ali, Zimerman, and Wolf}]{ali2024hidden}
Ameen Ali, Itamar Zimerman, and Lior Wolf. 2024.
\newblock The hidden attention of mamba models.
\newblock \emph{arXiv preprint arXiv:2403.01590}.
\bibitem[{Beltagy et~al.(2020)Beltagy, Peters, and Cohan}]{Beltagy2020Longformer}
Iz~Beltagy, Matthew~E. Peters, and Arman Cohan. 2020.
\newblock Longformer: The long-document transformer.
\newblock \emph{arXiv:2004.05150}.
\bibitem[{Biderman et~al.(2023)Biderman, Schoelkopf, Anthony, Bradley, O'Brien, Hallahan, Khan, Purohit, Prashanth, Raff, Skowron, Sutawika, and Van Der~Wal}]{pythia}
Stella Biderman, Hailey Schoelkopf, Quentin~Gregory Anthony, Herbie Bradley, Kyle O'Brien, Eric Hallahan, Mohammad~Aflah Khan, Shivanshu Purohit, Usvsn~Sai Prashanth, Edward Raff, Aviya Skowron, Lintang Sutawika, and Oskar Van Der~Wal. 2023.
\newblock \href {https://proceedings.mlr.press/v202/biderman23a.html} {Pythia: A suite for analyzing large language models across training and scaling}.
\newblock In \emph{Proceedings of the 40th International Conference on Machine Learning}, volume 202 of \emph{Proceedings of Machine Learning Research}, pages 2397--2430. PMLR.
\bibitem[{Black et~al.(2021)Black, Gao, Wang, Leahy, and Biderman}]{gpt-neo}
Sid Black, Leo Gao, Phil Wang, Connor Leahy, and Stella Biderman. 2021.
\newblock \href {https://doi.org/10.5281/zenodo.5297715} {{GPT-Neo: Large Scale Autoregressive Language Modeling with Mesh-Tensorflow}}.
\bibitem[{Chen et~al.(2024)Chen, Zhang, Hu, Han, Liu, and Sun}]{chen2024stuffed}
Yingfa Chen, Xinrong Zhang, Shengding Hu, Xu~Han, Zhiyuan Liu, and Maosong Sun. 2024.
\newblock \href {https://arxiv.org/abs/2410.07145} {Stuffed mamba: State collapse and state capacity of rnn-based long-context modeling}.
\newblock \emph{arXiv preprint arXiv:2410.07145}.
\bibitem[{Conneau et~al.(2020)Conneau, Wu, Li, Zettlemoyer, and Stoyanov}]{conneau-etal-2020-emerging}
Alexis Conneau, Shijie Wu, Haoran Li, Luke Zettlemoyer, and Veselin Stoyanov. 2020.
\newblock \href {https://doi.org/10.18653/v1/2020.acl-main.536} {Emerging cross-lingual structure in pretrained language models}.
\newblock In \emph{Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics}, pages 6022--6034, Online. Association for Computational Linguistics.
\bibitem[{Dao and Gu(2024)}]{mamba2}
Tri Dao and Albert Gu. 2024.
\newblock \href {https://openreview.net/forum?id=ztn8FCR1td} {Transformers are {SSM}s: Generalized models and efficient algorithms through structured state space duality}.
\newblock In \emph{Forty-first International Conference on Machine Learning}.
\bibitem[{Feng et~al.(2023)Feng, Zhang, Gu, Ye, He, and Wang}]{feng2023towards}
Guhao Feng, Bohang Zhang, Yuntian Gu, Haotian Ye, Di~He, and Liwei Wang. 2023.
\newblock \href {https://openreview.net/forum?id=qHrADgAdYu} {Towards revealing the mystery behind chain of thought: A theoretical perspective}.
\newblock In \emph{Thirty-seventh Conference on Neural Information Processing Systems}.
\bibitem[{Gao et~al.(2020)Gao, Biderman, Black, Golding, Hoppe, Foster, Phang, He, Thite, Nabeshima et~al.}]{pile}
Leo Gao, Stella Biderman, Sid Black, Laurence Golding, Travis Hoppe, Charles Foster, Jason Phang, Horace He, Anish Thite, Noa Nabeshima, et~al. 2020.
\newblock The pile: An 800gb dataset of diverse text for language modeling.
\newblock \emph{arXiv preprint arXiv:2101.00027}.
\bibitem[{Gao et~al.(2024)Gao, Lu, Yu, Byerly, and Khashabi}]{knowbutdonttell}
Muhan Gao, TaiMing Lu, Kuai Yu, Adam Byerly, and Daniel Khashabi. 2024.
\newblock \href {https://doi.org/10.18653/v1/2024.findings-emnlp.447} {Insights into {LLM} long-context failures: When transformers know but don`t tell}.
\newblock In \emph{Findings of the Association for Computational Linguistics: EMNLP 2024}, pages 7611--7625, Miami, Florida, USA. Association for Computational Linguistics.
\bibitem[{Glorot and Bengio(2010)}]{xavier}
Xavier Glorot and Yoshua Bengio. 2010.
\newblock \href {https://proceedings.mlr.press/v9/glorot10a.html} {Understanding the difficulty of training deep feedforward neural networks}.
\newblock In \emph{Proceedings of the Thirteenth International Conference on Artificial Intelligence and Statistics}, volume~9 of \emph{Proceedings of Machine Learning Research}, pages 249--256, Chia Laguna Resort, Sardinia, Italy. PMLR.
\bibitem[{Goldman et~al.(2024)Goldman, Jacovi, Slobodkin, Maimon, Dagan, and Tsarfaty}]{goldman-etal-2024-really}
Omer Goldman, Alon Jacovi, Aviv Slobodkin, Aviya Maimon, Ido Dagan, and Reut Tsarfaty. 2024.
\newblock \href {https://doi.org/10.18653/v1/2024.emnlp-main.924} {Is it really long context if all you need is retrieval? towards genuinely difficult long context {NLP}}.
\newblock In \emph{Proceedings of the 2024 Conference on Empirical Methods in Natural Language Processing}, pages 16576--16586, Miami, Florida, USA. Association for Computational Linguistics.
\bibitem[{Gu and Dao(2024)}]{mamba1}
Albert Gu and Tri Dao. 2024.
\newblock \href {https://openreview.net/forum?id=tEYskw1VY2} {Mamba: Linear-time sequence modeling with selective state spaces}.
\newblock In \emph{First Conference on Language Modeling}.
\bibitem[{Gu et~al.(2022)Gu, Goel, and Re}]{gu2022efficiently}
Albert Gu, Karan Goel, and Christopher Re. 2022.
\newblock \href {https://openreview.net/forum?id=uYLFoz1vlAC} {Efficiently modeling long sequences with structured state spaces}.
\newblock In \emph{International Conference on Learning Representations}.
\bibitem[{He et~al.(2015)He, Zhang, Ren, and Sun}]{he}
Kaiming He, X.~Zhang, Shaoqing Ren, and Jian Sun. 2015.
\newblock \href {https://api.semanticscholar.org/CorpusID:13740328} {Delving deep into rectifiers: Surpassing human-level performance on imagenet classification}.
\newblock In \emph{2015 IEEE International Conference on Computer Vision (ICCV)}, pages 1026--1034.
\bibitem[{Hendrycks and Gimpel(2016)}]{hendrycks2016gaussian}
Dan Hendrycks and Kevin Gimpel. 2016.
\newblock Gaussian error linear units (gelus).
\newblock \emph{arXiv preprint arXiv:1606.08415}.
\bibitem[{Jelassi et~al.(2024)Jelassi, Brandfonbrener, Kakade, and Malach}]{jelassi2024repeat}
Samy Jelassi, David Brandfonbrener, Sham~M. Kakade, and Eran Malach. 2024.
\newblock \href {https://proceedings.mlr.press/v235/jelassi24a.html} {Repeat after me: Transformers are better than state space models at copying}.
\newblock In \emph{Proceedings of the 41st International Conference on Machine Learning}, volume 235 of \emph{Proceedings of Machine Learning Research}, pages 21502--21521. PMLR.
\bibitem[{Kajitsuka and Sato(2024)}]{kajitsuka2024are}
Tokio Kajitsuka and Issei Sato. 2024.
\newblock \href {https://openreview.net/forum?id=nJnky5K944} {Are transformers with one layer self-attention using low-rank weight matrices universal approximators?}
\newblock In \emph{The Twelfth International Conference on Learning Representations}.
\bibitem[{Kingma and Ba(2015)}]{adam}
Diederik~P. Kingma and Jimmy Ba. 2015.
\newblock \href {http://arxiv.org/abs/1412.6980} {Adam: A method for stochastic optimization}.
\newblock In \emph{The Third International Conference on Learning Representations}.
\bibitem[{Kornblith et~al.(2019)Kornblith, Norouzi, Lee, and Hinton}]{cka}
Simon Kornblith, Mohammad Norouzi, Honglak Lee, and Geoffrey Hinton. 2019.
\newblock \href {https://proceedings.mlr.press/v97/kornblith19a.html} {Similarity of neural network representations revisited}.
\newblock In \emph{Proceedings of the 36th International Conference on Machine Learning}, volume~97 of \emph{Proceedings of Machine Learning Research}, pages 3519--3529. PMLR.
\bibitem[{Liu et~al.(2025)Liu, Zhu, Bai, He, Liao, Que, Wang, Zhang, Zhang, Zhang et~al.}]{liu2025comprehensive}
Jiaheng Liu, Dawei Zhu, Zhiqi Bai, Yancheng He, Huanxuan Liao, Haoran Que, Zekun Wang, Chenchen Zhang, Ge~Zhang, Jiebin Zhang, et~al. 2025.
\newblock \href {https://arxiv.org/abs/2503.17407} {A comprehensive survey on long context language modeling}.
\newblock \emph{arXiv preprint arXiv:2503.17407}.
\bibitem[{Liu et~al.(2024)Liu, Lin, Hewitt, Paranjape, Bevilacqua, Petroni, and Liang}]{liu-etal-2024-lost}
Nelson~F. Liu, Kevin Lin, John Hewitt, Ashwin Paranjape, Michele Bevilacqua, Fabio Petroni, and Percy Liang. 2024.
\newblock \href {https://doi.org/10.1162/tacl_a_00638} {Lost in the middle: How language models use long contexts}.
\newblock \emph{Transactions of the Association for Computational Linguistics}, 12:157--173.
\bibitem[{Paulo et~al.(2024)Paulo, Marshall, and Belrose}]{paulo2024doestransformerinterpretabilitytransfer}
Gonçalo Paulo, Thomas Marshall, and Nora Belrose. 2024.
\newblock \href {http://arxiv.org/abs/2404.05971} {Does transformer interpretability transfer to rnns?}
\bibitem[{Poli et~al.(2023)Poli, Massaroli, Nguyen, Fu, Dao, Baccus, Bengio, Ermon, and R{\'e}}]{poli2023hyena}
Michael Poli, Stefano Massaroli, Eric Nguyen, Daniel~Y Fu, Tri Dao, Stephen Baccus, Yoshua Bengio, Stefano Ermon, and Christopher R{\'e}. 2023.
\newblock Hyena hierarchy: Towards larger convolutional language models.
\newblock In \emph{International Conference on Machine Learning}, pages 28043--28078. PMLR.
\bibitem[{Skean et~al.(2025)Skean, Arefin, Zhao, Patel, Naghiyev, LeCun, and Shwartz-Ziv}]{skean2025layer}
Oscar Skean, Md~Rifat Arefin, Dan Zhao, Niket~Nikul Patel, Jalal Naghiyev, Yann LeCun, and Ravid Shwartz-Ziv. 2025.
\newblock \href {https://openreview.net/forum?id=WGXb7UdvTX} {Layer by layer: Uncovering hidden representations in language models}.
\newblock In \emph{Forty-second International Conference on Machine Learning}.
\bibitem[{Vaswani et~al.(2017)Vaswani, Shazeer, Parmar, Uszkoreit, Jones, Gomez, Kaiser, and Polosukhin}]{vaswani2017attention}
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan~N Gomez, \L~ukasz Kaiser, and Illia Polosukhin. 2017.
\newblock \href {https://proceedings.neurips.cc/paper_files/paper/2017/file/3f5ee243547dee91fbd053c1c4a845aa-Paper.pdf} {Attention is all you need}.
\newblock In \emph{Advances in Neural Information Processing Systems}, volume~30. Curran Associates, Inc.
\bibitem[{Vuli{\'c} et~al.(2020)Vuli{\'c}, Ponti, Litschko, Glava{\v{s}}, and Korhonen}]{vulic-etal-2020-probing}
Ivan Vuli{\'c}, Edoardo~Maria Ponti, Robert Litschko, Goran Glava{\v{s}}, and Anna Korhonen. 2020.
\newblock \href {https://doi.org/10.18653/v1/2020.emnlp-main.586} {Probing pretrained language models for lexical semantics}.
\newblock In \emph{Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP)}, pages 7222--7240, Online. Association for Computational Linguistics.
\bibitem[{Wang et~al.(2025)Wang, Cai, Wang, Zhu, Srivastava, Wang, and Li}]{wang2025understanding}
Peihao Wang, Ruisi Cai, Yuehao Wang, Jiajun Zhu, Pragya Srivastava, Zhangyang Wang, and Pan Li. 2025.
\newblock \href {https://openreview.net/forum?id=pymXpl4qvi} {Understanding and mitigating bottlenecks of state space models through the lens of recency and over-smoothing}.
\newblock In \emph{The Thirteenth International Conference on Learning Representations}.
\bibitem[{Wolf et~al.(2020)Wolf, Debut, Sanh, Chaumond, Delangue, Moi, Cistac, Rault, Louf, Funtowicz, Davison, Shleifer, von Platen, Ma, Jernite, Plu, Xu, Le~Scao, Gugger, Drame, Lhoest, and Rush}]{wolf-etal-2020-transformers}
Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, Remi Louf, Morgan Funtowicz, Joe Davison, Sam Shleifer, Patrick von Platen, Clara Ma, Yacine Jernite, Julien Plu, Canwen Xu, Teven Le~Scao, Sylvain Gugger, Mariama Drame, Quentin Lhoest, and Alexander Rush. 2020.
\newblock \href {https://doi.org/10.18653/v1/2020.emnlp-demos.6} {Transformers: State-of-the-art natural language processing}.
\newblock In \emph{Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing: System Demonstrations}, pages 38--45, Online. Association for Computational Linguistics.
\end{thebibliography}