The Effect of Layer Ordering in Hybrid Architectures

1 Introduction

• Hybrid architectures replace many self-attention layers with subquadratic alternatives, reducing memory needs and inference time for long-context tasks.The paper identifies state-space layers such as linear attention, Mamba, and Gated DeltaNet as popular alternatives.Source: §1 · p. 1
• Large gaps between segments require SSM layers to communicate each segment’s query, while Transformers perform dense computation within each segment.The paper uses this task structure to illustrate performance differences among hybrid architectures.Source: Figure 1 · p. 2
• Layer selection and ordering are difficult to evaluate because layer types impose different sequence-processing constraints and exhaustive empirical sweeps are intractable.Different orderings may enable early sequence reordering and reduce the model size and depth needed for a task.Source: §1 · p. 2
• For any ordering of generalized SSM and causal-windowed layers, the authors construct a task and model where any different ordering must use more resources at some layer.The analysis focuses on window width, state size, and embedding precision as constraints on layer types.Source: §1 · p. 2
• Learned hybrids exhibit the predicted separations on natural tasks, and depth-4 and depth-5 models perform well when they contain the task’s dominant TS or TST sub-architecture.TS is associated with segmented retrieval; TST supplies a task requiring layer alternation.Source: §1 · p. 2
• For the deeper task studied, T*S* architectures outperform all S*T* architectures, while linear probes show gradual decodability increases in TTSSSSSS’s SSM layers but not TTTSSSSS’s.The authors interpret this as evidence that slightly different architectures can fail to use SSM layers to improve performance.Source: §1 · p. 3

2 Theoretical Foundations for Hybrid Layer Ordering

• The layer abstractions distinguish causal windowed Transformer functions from finite-state SSMs, which update state from each input token and output from that state.The paper characterizes Transformer layers by window size W and SSM layers by state-space size at most 2^B.Source: §2 · p. 3

2.1 Expressivity and Efficiency Tradeoff under Budget Constraints

• For each target profile, its task is perfectly solvable with B,P=O(log m), W=m+1; wrong-order profiles require larger budgets, while incomplete profiles require a constant-factor increase.For wrong-order profiles, achieving error O(1/D) requires W≥cτm^(1/ϵ) or max{B,P}≥cτm; incomplete profiles require one resource to exceed a constant multiple of the target budget.Source: §2.1 · p. 4
• The TST tensor-lookup task makes TST dominant by passing an extracted lookup value through S to a later T, but a similar simple task cannot separate STS from TS.For the STS variant, a T-layer can copy the matrix after q2, changing the processing order so an S-layer with O(log m) bits can solve the task.Source: §2.1 · p. 5
• The size-2 cover narrows the competing profiles to maximal profiles with compressed length at most two, whose hard tasks are combined into one task for the target profile.The paper uses a meta-model to combine the target-profile submodels that solve the component tasks.Source: §2.1 · p. 6

2.2 Natural Tasks for 2-layer Hybrid Models

• For natural tasks with a sequence cut, ST requires a large state, while TS can process relevant information with T and carry it through the sequence with S.The TS construction uses T for the λ-back recall in segmented retrieval, then S carries the relevant tokens across the sequence.Source: §2.2 · p. 7

3.1 Two Layers - Segmented Retrieval and Exact Copy

• Two-layer hybrid sweeps test segmented retrieval and exact copy under a fixed proxy for inference memory, with results averaged across five seeds.The experiments examine whether the theoretical architecture separations also appear in learned hybrids.Source: §3.1 · p. 8, §3 · p. 8

3.2 Containing Subarchitectures Predicts Performance

• All tested 4-layer segmented-retrieval models containing TS achieve perfect performance, while none without TS does so across the three seeds.The table separates profiles with and without TS.Source: Table 1 · p. 9, §3.2 · p. 10
• On the 5-layer TST task, models containing TST perform notably better than other architectures, except for STSTS.The figure compares all possible 5-layer Transformer–Mamba hybrids.Source: Figure 3 · p. 9, §3.2 · p. 10

4 Observations on Associative Recall

• On (2, 1)-AR, T∗S∗ outperforms S∗T∗, and the two architecture families have different peak Transformer-to-SSM ratios.The task represents a compressed multidocument setting; the authors suggest this may explain the ordering difference.Source: §4 · p. 10
• Linear probes show a significant decodability gain at the first Transformer layer; stacked SSM layers add slower gains, hypothesized to reflect retrieval beyond the Transformer context window.The authors also note similar Transformer-dominant recall behavior in models trained on natural language [17].Source: §4 · p. 11

5 Conclusion

• For each hybrid design, the authors construct tasks on which that architecture is provably favored over other choices.The general theory uses deliberately constructed tasks and expressive abstractions of SSM and attention layers; optimal designs may not be realizable by standard practical models.Source: §5 · p. 11
• The authors identify extending the theory to practical layer classes, richer task mixtures, and natural-language settings as future work.They aim to identify task-level properties that predict effective layer orderings before expensive training.Source: §5 · p. 11

AI use statement

• Generative AI tools assisted with discussion, some proof construction and code, and grammar polishing; authors say they reviewed and verified the assisted work.The authors report that GPT 6 Astra established and proved the main analysis of Case 4 of Theorem 4 based on an author-provided hardness construction.Source: AI use statement · p. 12

A Notation

• The notation section defines Transformer layers through attention and an MLP, and SSM layers through recursive state updates and emissions.It also introduces attention heads and attention layers.Source: A Notation · pp. 17–18

B.1 Hybrid Models

• Prior empirical studies examine hybrid ratios and some placement choices, while this work studies broader ordering differences and deeper hybrids theoretically.The cited work on placement explores only a narrow slice of the design space.Source: §B.1 · p. 18

• Associative recall probes whether models maintain an associative memory of facts indexed by prior tokens.Source: §B.2 · p. 19
• The original task separated pure SSMs from pure Transformers on storing large memories, but now many similar architectures perform nearly perfectly.Source: §B.2 · p. 19
• The paper uses phonebook, where keys and values each contain multiple tokens, as a probe of hybrid-architecture differences.Source: §B.2 · p. 19

C Experiment Details

• Experiments use single-head GPTNeoX Transformer layers with RoPE; increasing head count by splitting key and query matrices into thin rectangles severely hindered performance at this scale.Source: C Experiment Details · p. 19
• Mamba is the SSM used in all experiments, and its state size can be varied independently of token dimension.Source: C Experiment Details · p. 19

C.1 Two- and Three- Layer Hybrid Details

• Segmented retrieval, exact copy, and TST use task-specific sequence, vocabulary, and lookup or copying settings, with 512 dimensions of memory for the first two tasks.Segmented retrieval uses length-150 sequences, 10 segments, and lookup distance 8; exact copy uses length-100 sequences and copies 5-token spans; TST uses length 150 and vocabulary size 7.Source: §C.1 · pp. 19–20

C.2 Learning Rate Selection

• Learning rates are selected with different sweep procedures: short loss-based sweeps for segmented retrieval and TST, full runs for exact copy, and two-epoch training for associative recall.Exact copy uses five seeds per learning rate; its sweep is more computationally expensive than the other task sweeps.Source: §C.2 · p. 20

C.3 Memory Proxy Calculation

• The inference-memory proxy sums Transformer context storage, Mamba state storage, and Gated DeltaNet square-state storage across the model's layers.With layer counts N_TF, N_SSM, and N_GDN, the stated total is M = N_TF dW + N_SSM dN + N_GDN d^2.Source: §C.3 · p. 20

D.1 Segmented Retrieval

• With additional, non-power-of-two token dimensions, only the TS architecture performs segmented retrieval, matching the main-body pattern.Source: §D.1 · p. 21
• TT reaches its best average segmented-retrieval performance at smaller token dimensions and larger windows than TS, whose state is used to increase computation per token.Source: §D.1 · p. 21

D.2 Different Quantities of Available Memory

• With double the memory, TS remains the highest-performing architecture; TT improves significantly on exact copy at small token dimensions, and segmented-retrieval performance rises for all architectures versus 512 dimensions.The authors attribute TT's exact-copy gain likely to the larger window capturing most copied subsequences.Source: §D.2 · p. 21, Figure 6 · p. 22, Figure 7 · p. 23

D.3 Testing DeltaNet Hybrids

• For Gated DeltaNet hybrids, TS mirrors TG closely, while TG and GT curves end earlier because GDN's larger state reaches memory constraints sooner.The authors describe this as confirming inference memory as a core bottleneck across multiple SSM architectures.Source: §D.3 · p. 23

D.4 Three Layers - Tensor Lookup

• On tensor lookup, TST exceeds 90% accuracy while competing architectures do not; among five-layer models, all but one with a TST motif outperform the remaining models at roughly equal memory.Source: Figure 9 · p. 23, §D.4 · p. 24

D.5 An ST Task

• The ST task’s hybrid performance curves have two local maxima, with the favored algorithm depending on token dimension and window and state-expansion settings.Mamba continues improving as dimension increases and state expansion shrinks; the authors suggest pure SSM may be favored in that regime, while TS may resemble pure TF at moderate dimensions.Source: §D.5 · pp. 24, 26

D.6 Associative Recall Variants

• Both associative-recall variants favor low Transformer-to-SSM ratios, but (1, 2)-AR performs better with more Transformer layers than (2, 1)-AR.The authors suggest the (1, 2)-AR variant’s smaller key space may explain its higher average performance.Source: §D.6 · p. 26

D.7 Learning Rate Sensitivity

• SSMs show resilience to learning-rate selection, although some SSM-heavy architectures have a narrow high-performance band and their optimal-performance ordering mostly persists between 10^-4 and 10^-3.The cited passage describes this range as the optimal learning-rate band.Source: §D.7 · pp. 26–27

D.5 An ST Task

• ST architectures achieve perfect performance on selective copy, while the other architectures do not.Figure 11 is described as showing ST as largely dominant.Source: §D.5 · p. 24, Figure 11 · p. 27

E Proving Theorem 1

• Theorem 1’s proof handles wrong-order hybrids and correctly ordered hybrids that have insufficient layers as separate cases.The section first introduces prior lemmas for combining neighboring layers of the same type.Source: E Proving Theorem 1 · p. 27
• A composition of state-space layers can be represented by one equivalent SSM layer.The lemma bounds the state space of the equivalent layer.Source: E Proving Theorem 1 · p. 27
• A composition of causal windowed functions can be represented by one equivalent causal windowed function.The lemma gives a bound on the equivalent function’s window size.Source: E Proving Theorem 1 · p. 28

E.1 Wrong use of ordering

• If an architecture profile A lacks comp(τ) as a subsequence, error at most 1/3 requires W ≥ cτm^(1/ϵ) or max{B, P} ≥ cτm.The theorem applies for all sufficiently large m and a task family parameterized by τ and ϵ.Source: §E.1 · p. 28
• Layer-merging lemmas reduce the lower-bound analysis to models with the complementary profile, which can represent models whose architecture lacks comp(τ).The analysis then characterizes information dependence across layers and positions.Source: §E.1 · pp. 28–29
• The recursive-chain task applies a sequence of functions to successive input blocks, and a model with profile comp(τ) can solve it exactly under the stated resource conditions.The task input consists of a starting key, blocks x_i, and intervening zero vectors.Source: §E.1 · pp. 29–30
• The proof records hidden SSM states at separated positions; conditioned on that record, the next stage’s states can be expressed as functions of the next input block.The record C_i is the tuple of SSM states at the stage-specific positions.Source: §E.1 · pp. 30–31
• A probabilistic construction chooses the recursive functions so that each stage’s candidate list of keys stays bounded, supporting the wrong-order resource lower bound.The construction uses lists of size l_m in the stated framework.Source: §E.1 · pp. 31–33

Theorem 4

• O(1/D)-error models with α ⪯ A but τ not ⪯ A require W ≥ λτ(m + 1) or max{B, P} ≥ λτCτ log m, while a τ-configured model solves exactly.This holds for sufficiently large m for the constructed task family Gτ,m.Source: §E.2 · p. 34

Size-2 cover reduction

• If τ is not a subsequence of A, some profile in the size-2 cover C2(τ) is also absent, reducing hardness for τ to hard tasks for each cover profile.The reduction constructs a task for τ that is exactly solvable by a τ-configured model.Source: §E.2 · p. 34

Hard tasks for simple profiles

• For σ = S^r and σ = T^r, the constructed tasks are exactly solvable by σ-configured models, while models missing σ have success probability at most 2/3 when resources are below the stated bounds.The S^r task computes the inner product of two bit vectors; the T^r task uses a causal window to move and process the input.Source: §E.2 · p. 37

Hard tasks for mixed profiles

• For σ = T^rS^s, a model with that profile solves exactly by selecting one of K vectors and computing its inner product with w, whereas a model missing σ has failure probability 1/3.The lower bound uses a random index J and a list-decoding argument to limit successful recovery of the selected vector.Source: §E.2 · pp. 37–38
• For σ = S^rT^s, an exactly solving model exists, while any model missing σ with error at most 1/3 must have W ≥ λ(m + 1) or max{B, P} ≥ λb_m.Here b_m = ⌈C log2 m⌉.Source: §E.2 · p. 38

Random-chain analysis

• The technical analysis of the S^rT^s case uses a random-chain lemma whose later transition functions are permutations and whose guarantees hold simultaneously across a specified range of source sets.The lemma considers uniform source and query blocks and establishes a fiber-size condition for almost all first query blocks.Source: §E.2 · pp. 39–40

E.2 Correct ordering but insufficient layers

• The task is an equal mixture of tagged chain and reach branches, and the same table realization works uniformly over eligible source sets, fixed old contexts, and post-table physical replay maps.The chain iterates H_i updates from z0 and query blocks; the reach branch asks for u_J from a database of entries u_1, …, u_k0.Source: §E.2 · p. 40
• S^rT^s solves both the chain and reach branches exactly.On the chain branch, S-layers copy the source into scratch and T-layers compute successive H_i updates; on reach, S-layers are identities and T-layers shift the database until the requested entry is available.Source: §E.2 · p. 40
• Independence and entropy subadditivity give I(uJ; Γ, J) ≤ DB*/k0; the standard error-indicator/Fano estimate yields p ≤ (DB*/k0 + 1)/log2 K = o(1), hence reach-task success is o(1).Source: §E.2 · p. 41
• Competitors avoiding S^rT^s achieve total success at most 1/2 + o(1), so their error exceeds 1/3 for sufficiently large m.The bound also holds for input-independently randomized models by fixing and then averaging over the random seed.Source: §E.2 · p. 42
• A τ-profile model solves Gτ,m exactly with B=P=C log m and W=m+1; configurations A lacking τ need W≥λ(m+1) or max{B,P}≥λC log m for error ≤1/(4D).The lower bound applies to A ∈ Σ≤D when τ̸ ⪯A; the tasks are constructed for each target profile τ.Source: §E.2 · p. 46

E.3 Proof of Theorem 1

• The proof establishes Theorem 1 by combining completeness constructions with soundness bounds for architectures that do not contain τ.The proof states that the two cases are exhaustive and completes both soundness statements.Source: §E.3 · pp. 46–47
• A uniform tag J selects between the wrong-order and insufficient-profile task families, whose inputs are padded to the same sequence length.The construction uses the corresponding profile-τ component models to solve either task exactly.Source: §E.3 · p. 46
• The exhaustive cases force either W ≥ cτm^(1/ϵ) or max{B, P} ≥ cτm, or W ≥ λτ(m + 1) or max{B, P} ≥ λτCτ log m.These bounds apply to the wrong-order and insufficient-profile cases, respectively.Source: §E.3 · p. 47

F Omitted Details in Section 2.2

• The omitted details identify Exact Copy and Segmented Retrieval as tasks in the dominant-TS family and give constructions using less window and state than ST requires.The appendix describes these as the task-specific constructions and bounds supporting the main result.Source: F Omitted Details in Section 2.2 · p. 47

F.1 Hardness Results

• For a deterministic ST model computing a task satisfying Assumption 1, the cut theorem requires |S| ≥ |U| when τ ≤ L − W.The prefix family must be distinguishable by common suffixes, so equal cut states would contradict correct outputs.Source: §F.1 · p. 47
• Exact Copy requires log |S| ≥ Ω(h log |V|) for deterministic ST models, where h = min(l, (L − W)+).The bound is stated for |V| ≥ 5 and L ≥ cl.Source: §F.1 · p. 48
• Segmented Retrieval requires log |S| ≥ Ω(h log |N|) for deterministic ST models, where h = min(λ, (L − W)+).The theorem assumes |N| ≥ 2 and L ≥ 3λ + 1.Source: §F.1 · p. 50

Construction for Exact Copy

• A TSS model solves Exact Copy with window W ≤ O(l) and state complexity log |S| ≤ O(log |V| + log(cl)).Its Transformer retrieves x_(t−l), an SSM counts consecutive matches, and a second SSM latches success.Source: Construction for Exact Copy · pp. 51–52

Construction for Segmented Retrieval

• A TS model solves Segmented Retrieval with window W ≤ O(λ) and log |S| ≤ O(log |V|).The Transformer retrieves the token λ positions earlier, while the SSM retains the latest retrieved token in N at a boundary.Source: Construction for Segmented Retrieval · pp. 53, 55

F.3 SSMs with a Convolutional Front-End

• An Sκ-layer’s state update and output can use the previous state and the local input window x_(t−κ:t); Mamba blocks fit this abstraction, but the converse does not hold.The local window is not charged to recurrent-state complexity.Source: §F.3 · p. 55
• Theorem 1’s resource separation carries over with cτm/(κ+2) in place of cτm, but its constant-factor separation does not follow from this conversion.Theorems 6 and 7 likewise replace log2 |S| with log2 |S| + κP.Source: §F.3 · p. 56
• For a stack of DS convolutional recurrent layers, matching states and a common suffix imply identical outputs when the prefixes agree on the final DSκ positions.This cut-through-stack result provides the basis for adjusting the natural-task lower bounds.Source: §F.3 · p. 56
• For Exact Copy and Segmented Retrieval, the lower bound is nonzero exactly when task distance exceeds DSκ; otherwise, a small-state construction exists.The converse uses W = 1 and places the task computation in the final recurrent layer.Source: §F.3 · pp. 56–57