Improving the Expressive Power of Hybrid Sequence Models via Multi-Hop Retrieval

Yam Eitan; Gilad Yehudai; Haggai Maron

cs.LGOct 18, 2026

Improving the Expressive Power of Hybrid Sequence Models via Multi-Hop Retrieval

Authors: Yam Eitan, Gilad Yehudai, Haggai Maron

Organizations: Technion – Israel Institute of Technology · NVIDIA

Abstract

Hybrid architectures offer a promising approach to efficient language modeling by combining a small number of full-attention layers with a larger number of efficient sequence layers, such as sparse attention and recurrent layers. While full-attention layers maintain a KV cache that grows linearly with sequence length, efficient layers typically rely on separate, fixed-size states rather than accessing this richer existing memory. Recent work has begun exploring ways to allow efficient layers access to the KV cache, but the theoretical benefits of such access, and the extent to which existing approaches can realize them remain poorly understood. We investigate these questions by analyzing the expressive power of individual hybrid blocks on function composition, a well-studied task in sequence-model expressivity associated with multi-hop reasoning. We first establish that both conventional hybrid blocks and those with top- $κ$ KV-cache access require $Ω(\sqrt{n})$ capacity to compute even two-function composition over a domain of size $n$ . To overcome this limitation, we introduce the Multi-hop Retrieval Attention (MuRA) hybrid block, a simple architecture that enables multi-hop retrieval from a shared KV cache while preserving asymptotic computational and memory efficiency. We prove that a single fixed size $m$ -hop MuRA block can compute $m$ -function composition, whereas computing $(m+1)$ -function composition with the same retrieval depth requires polynomial model size, establishing a strict expressivity hierarchy in retrieval depth. Synthetic experiments support this hierarchy, showing that increasing retrieval depth enables learning progressively deeper compositions, substantially outperforming existing hybrid architectures. Together, our results identify multi-hop retrieval as a promising mechanism for improving the expressive power of efficient hybrid architectures.

Figures & tables

Panel of Figure 1 — Figure 1: MuRA versus top- $\kappa$ retrieval. Both panels illustrate the same input for $3$ -function composition. (a) The construction of Theorem 3.3 performs three successive lookups to compute $f_{1}(f_{2}(f_{3}(x)))$ . (b) An illustrative attention-score-based top- $\kappa$ selection retrieves several tokens directly from $x$ , but fails to reach the target token through the required retrieval chain.

Model	$m=1$	$m=2$	$m=3$	$m=4$
MuRA $(32,)$	$83.05\pm 29.35$	$44.47\pm 23.64$	$5.81\pm 3.49$	$3.68\pm 0.11$
MuRA $(16,1)$	$100.00\pm 0.01$	$99.30\pm 1.22$	$30.31\pm 23.22$	$3.73\pm 0.10$
MuRA $(11,1,1)$	$100.00\pm 0.00$	$100.00\pm 0.00$	$99.99\pm 0.01$	$27.81\pm 24.98$
MuRA $(8,1,1,1)$	$100.00\pm 0.00$	$100.00\pm 0.00$	$100.00\pm 0.00$	$100.00\pm 0.00$
Linear attention	$99.97\pm 0.01$	$57.83\pm 10.75$	$14.97\pm 5.66$	$8.78\pm 0.62$
Sliding window (16)	$100.00\pm 0.00$	$46.83\pm 8.49$	$3.76\pm 0.23$	$3.72\pm 0.08$

Table 1: Function-composition accuracy (%), reported as mean

\pm

standard deviation over three seeds. Blue , green , and purple indicate the first, second, and third highest mean accuracies in each column, respectively.

Multi-Hop Retrieval

See all 16 papers

Oct 13, 2026

The Effect of Layer Ordering in Hybrid Architectures

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.

John Cooper, Mingchen Ma, Jiaqi Mo +4

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

Hybrid Transformer Architectures Hybrid Attention Transformer Attention

Ask PDF

Mar 9, 2026

Expressivity-Efficiency Tradeoffs for Hybrid Sequence Models

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.

John Cooper, Ilias Diakonikolas, Mingchen Ma +1

Department of Computer Sciences, University of Wisconsin-Madison

Hybrid Transformer-State Space Model Architectures Sequence Modeling Transformer Expressivity

Ask PDF

Jun 13, 2026

Rethinking the Role of Efficient Attention in Hybrid Architectures

Modern language models increasingly adopt hybrid architectures that combine full attention with efficient attention modules, such as sliding-window attention (SWA) and recurrent sequence mixers. However, how these efficient modules shape model capabilities remains poorly understood. To address this gap, we conduct a systematic analysis across hybrid architectures from three perspectives: scaling behavior, mechanism analysis, and architecture design. First, from a scaling perspective, we find that efficient-attention design primarily affects how fast long-context capability emerges, while different hybrids eventually converge to comparable long-context performance under sufficient training. Second, mechanistically, we show that long-range retrieval is mainly carried by full attention, whereas efficient attention shapes its optimization trajectory. This explains a counter-intuitive phenomenon we call Large-Window Laziness: larger SWA windows can delay the formation of retrieval heads in full-attention layers. Third, guided by this mechanism, we show that applying NoPE to only the full-attention layers of a small-window SWA hybrid substantially improves long-context performance with negligible impact on short-context performance.

Ziqing Qiao, Yinuo Xu, Chaojun Xiao +6

1Tsinghua University · 2OpenBMB

Attention Head Analysis Long-Context Modeling Efficient Attention Sliding-Window Attention Attention Mechanisms Hybrid Attention

Ask PDF

Abstract

Multi-Hop Retrieval

Explore similar work

The Effect of Layer Ordering in Hybrid Architectures

Expressivity-Efficiency Tradeoffs for Hybrid Sequence Models

Rethinking the Role of Efficient Attention in Hybrid Architectures