The Effect of Layer Ordering in Hybrid Architectures

John Cooper; Mingchen Ma; Jiaqi Mo; Shengqi Qiu; Nathaniel Bastian; Ilias Diakonikolas; Frederic Sala

cs.LGOct 13, 2026

The Effect of Layer Ordering in Hybrid Architectures

Authors: John Cooper, Mingchen Ma, Jiaqi Mo, Shengqi Qiu, Nathaniel Bastian, Ilias Diakonikolas, Frederic Sala

Organizations: Department of Computer Sciences, University of Wisconsin-Madison · Department of Computer Sciences, John Hopkins University

Abstract

Hybrid attention is rapidly becoming the dominant architectural choice for frontier models. Despite their successes, such architectures complicate model design by introducing additional choices, including the types and locations of layers. This combinatorially large design space has been studied empirically, but given the vast expense of training and evaluating models with different configurations, a principled understanding of the effects of layer selection and ordering is needed. We provide the first such analysis, showing that the dominant hybrid architecture is highly task-specific. Concretely, for any layer configuration, we prove that there exists a task where we can construct a model with this configuration that provably outperforms any other model of any different configuration under fixed resource constraints. Empirically, we observe that for standard training, learned hybrid architectures exhibit separations between layer orderings that match our theoretical findings. For instance, under the same memory budget, the optimal architecture predicted by our theory for a segmented retrieval task consistently achieves perfect accuracy, whereas the next strongest plateaus below 0.5. These results extend to deeper models, where ``containing'' the dominant architecture for a task indicates whether the model attains high performance. In addition, we investigate the mechanisms learned by hybrids, observing that Transformer layers contribute the largest share of performance on associative recall across architectures. This work establishes insights into task-dependent hybrid architecture design relevant to the training and scaling of frontier-scale hybrids.

Figures & tables

Figure 1: Structure of tasks that exhibit performance differences among different hybrid architectures. Large gaps need SSM layers to communicate each segment’s query, while the dense computation per segment is done by the Transformers.

Panel of Figure 2 — Figure 2: Results for the TS tasks where an inference memory proxy is held constant. The annotations represent the relationship between window storage and state size for Transformer-Mamba hybrids. The only architecture with high performance on the task is the TS architecture. Due to fixed memory, models with larger token dimension have smaller window and state size, making the tasks difficult. Conversely, small token dimension limits expressivity.

Appendix figures & tables12 assets

Supplementary material from the paper’s appendix.

Appendix

Symbol	Meaning
${\mathcal{V}}$	The vocabulary space.
${\mathcal{U}}$	The intermediate space used between architecture layers.
${\mathcal{Y}}$	The output space.
S	An SSM layer.
T	A TF layer.
$\Sigma$	Defined to be $\{\textmd{{{S}}},\textmd{{{T}}}\}$ , the possible letters in an architecture.

Appendix

Table 2: Notation

Figure 6: Extended results for segmented retrieval.

Panel of Figure 7 — Figure 7: Hybrid performance on exact copy for 1024 dimensions of memory.

Panel of Figure 8 — Figure 8: Hybrid performance on exact copy for 256 dimensions of memory.

Model	Performance	Model	Performance
TTTST	0.911	TTTTT	0.533
TTSTT	0.934	TTTTS	0.840
TTSTS	0.884	TTTSS	0.776
TTSST	0.890	TTSSS	0.787
TSTTT	0.922	STTTT	0.527
TSTTS	0.923	STTTS	0.842

Appendix

Table 3: Quantitative results for a deeper hybrids on the TST task. The left architectures contain TST as a sub-architecture, while the other three do not.

Figure 9: Results for training various architectures with different token dimensions on the TST task. All models used the same window size and state size. As expected, only the TST model attains a high accuracy.

Figure 10: Performance on segmented retrieval, including Gated DeltaNet ( G ) hybrids. Only the 15:1 ratio of window to SSM state for TS and ST hybrids is shown.

Figure 11: Performance of architectures trained on selective copy. We observe that the ST architecture is largely dominant, as expected.

Panel of Figure 12 — Figure 12: Token accuracy of $\textmd{{{T}}}^{*}\textmd{{{S}}}^{*}$ and $\textmd{{{S}}}^{*}\textmd{{{T}}}^{*}$ architectures on $(2,1)$ associative recall. Small differences in architectures affect performance dramatically.

Panel of Figure 13 — Figure 13: Learning rate sweep for select architectures trained on associative recall.

Figure 14: Structure of Exact Copy. At each position $t>l$ , the TF layer retrieves $\mathbf{x}_{t-l}$ . The first SSM layer counts consecutive matches $\mathbf{x}_{t}=\mathbf{x}_{t-l}$ in $H_{t}$ , while the second sets $y_{t}=1$ once $H_{t}=(c-1)l$ and retains this value thereafter. Here, $abc$ repeats three times consecutively, giving $H_{10}=6$ and $y_{L}=1$ despite subsequent mismatches.

Figure 15: Structure of Segmented Retrieval. At each internal boundary $t_{j}$ , the TF layer retrieves $v_{j}=\mathbf{x}_{t_{j}-\lambda}$ . The SSM layer retains the most recent retrieved token belonging to $\mathcal{N}$ and outputs it at position $L$ . Here, the retrieved tokens are $A,u,B,v$ , so the output is $B$ .

Explore similar work

Expressivity-Efficiency Tradeoffs for Hybrid Sequence Models
Mar 9, 2026John Cooper, Ilias Diakonikolas, Mingchen Ma +1Sequence Modeling Transformer Expressivity
Department of Computer Sciences, University of Wisconsin-Madison
Where Should LoRA Go? Component-Type Placement in Hybrid Language Models
Apr 24, 2026Hector Borobia, Elies Seguí-Mas, Guillermina Tormo-CarbóLow-Rank Adaptation Recurrent Transformers
VRAIN – Valencian Research Institute for Artificial Intelligence, Universitat Politècnica de València, Valencia, Spain · Department of Economics and Social Sciences, Universitat Politècnica de València, Valencia, Spain · Department of Business Organisation, Universitat Politècnica de València, Valencia, Spain
A Systematic Analysis of Hybrid Linear Attention
Jul 8, 2025Dustin Wang, Rui-Jie Zhu, Steven Abreu +9Linear Attention Hybrid Attention
1UC Santa Cruz · 2ByteDance Seed · University of Groningen +3

Mar 9, 2026J/K move · Enter open · S save

Expressivity-Efficiency Tradeoffs for Hybrid Sequence Models

John Cooper, Ilias Diakonikolas, Mingchen Ma, Frederic Sala

Department of Computer Sciences, University of Wisconsin-Madison

Sequence Modeling Transformer Expressivity

Open paper AskPDF arXiv

Hybrid sequence models--combining Transformer and state-space model layers--seek to gain the expressive versatility of attention as well as the computational efficiency of state-space model layers. Despite burgeoning interest in hybrid models, we lack a basic understanding of the settings where--and underlying mechanisms through which--they offer benefits over their constituent models. In this paper, we study this question, focusing on a broad family of core synthetic tasks. For this family of tasks, we prove the existence of fundamental limitations for non-hybrid models. Specifically, any Transformer or state-space model that solves the underlying task requires either a large number of parameters or a large working memory. On the other hand, for two prototypical tasks within this family--namely selective copying and associative recall--we construct hybrid models of small size and working memory that provably solve these tasks, thus achieving the best of both worlds. Our experimental evaluation empirically validates our theoretical findings. Importantly, going beyond the settings in our theoretical analysis, we empirically show that learned--rather than constructed--hybrids outperform non-hybrid models with up to 6x as many parameters. We additionally demonstrate that hybrid models exhibit stronger length generalization and out-of-distribution robustness than non-hybrids.

Abstract

Explore similar work

Expressivity-Efficiency Tradeoffs for Hybrid Sequence Models

Where Should LoRA Go? Component-Type Placement in Hybrid Language Models

A Systematic Analysis of Hybrid Linear Attention