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

Authors: Yam Eitan¹, Gilad Yehudai¹, Haggai Maron¹,²
¹Technion – Israel Institute of Technology, ²NVIDIA

---

Summary (Overview)

• Core contribution: The paper introduces Multi-hop Retrieval Attention (MURA), a novel hybrid block architecture that enables efficient sequence layers to perform multi-hop retrieval from a full-attention layer's shared KV cache, substantially improving expressive power without sacrificing asymptotic efficiency.

• Key negative result: The authors prove that both conventional hybrid blocks and those with top-κ KV-cache access require Ω(√n) capacity to compute even two-function composition over a domain of size n, revealing fundamental limitations of existing architectures for multi-hop reasoning.

• Key positive result: A single m-hop MURA block can compute m-function composition with constant width, constant depth, and only O(log N) precision—the first architecture with a single full-attention layer provably capable of such multi-step function composition.

• Expressivity hierarchy: Computing (m+1)-function composition with only m retrieval hops requires polynomial model size Ω(n^(1/(m+1))), establishing a strict expressivity hierarchy in retrieval depth, even when the total number of retrieved KV entries is matched.

• Empirical validation: On synthetic function-composition benchmarks with up to four retrieval hops, MURA achieves near-perfect accuracy on compositions matching its retrieval depth, substantially outperforming existing hybrid baselines including linear attention, sliding-window attention, and Gated DeltaNet.

---

Introduction and Theoretical Foundation

Background and Motivation

Transformers have become the dominant architecture for language modeling, but full attention poses significant computational and memory challenges:

• Full attention requires O(N²) computation in sequence length and maintains a KV cache that grows linearly with context length during autoregressive inference
• Efficient sequence layers (sparse attention, recurrent layers like Mamba and Gated DeltaNet) process sequences in O(N) time with fixed-size states

Hybrid architectures combine a few full-attention layers with many efficient layers, as adopted by prominent models including Kimi K3, Gemma 3, Nemotron 3 Ultra, and OLMo Hybrid.

The Research Gap

The key observation driving this work:

Recurrent layers within hybrid blocks typically maintain their own fixed-size hidden states, without direct access to the richer KV cache already maintained by the full-attention layer. Recent approaches (HySparse, HySparse2) allow top-κ retrieval from shared KV caches, but the theoretical benefits of such access remain poorly understood.

Theoretical Framework

The analysis centers on function composition, a well-studied task in sequence-model expressivity that captures multi-hop reasoning:

Definition (Function Composition Task): Given functions $f_1, \dots, f_m: [n] \to [n]$ and a query $x \in [n]$, the sequence is constructed as:

$$\mathbf{a}(f_1, \dots, f_m, x) = (f_1(1), \dots, f_1(n), \dots, f_m(1), \dots, f_m(n), x) \tag{6}$$

The task requires outputting at the final position:

$$F_{m,n}(\mathbf{a}(f_1, \dots, f_m, x)) = (f_1 \circ \dots \circ f_m)(x) \tag{7}$$

Communication Complexity Connection

The paper leverages one-way multiparty number-on-the-forehead (NOF) communication protocols. A classical result establishes that any one-way (m+1)-player NOF protocol computing m-function composition requires Ω(n^(1/m)) bits of communication.

---

Methodology

Architectural Components

Full-attention layer (Definition 2.1): Standard multi-head attention with causal masking:

$$\text{head}_j(\mathbf{X}) = \text{softmax}\left(\frac{\mathbf{Q}_j \mathbf{K}_j^\top}{\sqrt{d/h}} + \mathbf{M}\right) \mathbf{V}_j \tag{1}$$

Recurrent layer (Definition 2.2): With state transition and output maps:

$$\mathbf{h}_i = f(\mathbf{h}_{i-1}, \mathbf{x}_i), \quad \mathbf{y}_i = g(\mathbf{h}_i, \mathbf{x}_i) \tag{2}$$

Top-κ access recurrent layer (Definition 2.3): Extends recurrent layers with access to top-κ KV entries from a full-attention layer's cache:

$$\mathbf{C}_i = \left[\left[\mathbf{k}_t^j, \mathbf{v}_t^j\right]_{t \in \mathcal{N}_j^{(\kappa)}(i)}\right]_{j=1}^h \tag{3}$$ $$\mathbf{h}_i = f(\mathbf{h}_{i-1}, \mathbf{x}_i, \mathbf{C}_i), \quad \mathbf{y}_i = g(\mathbf{h}_i, \mathbf{x}_i, \mathbf{C}_i) \tag{4}$$

Top-κ hybrid block (Definition 2.4): Combines pre-attention recurrent layers, one full-attention layer, and post-attention top-κ access layers:

$$\mathsf{B}_\kappa(\mathbf{X}) = \left(\mathsf{R}^{\text{post}}_{L_{\text{post}}} \circ \dots \circ \mathsf{R}^{\text{post}}_1 \circ \mathsf{A} \circ \mathsf{R}^{\text{pre}}_{L_{\text{pre}}} \circ \dots \circ \mathsf{R}^{\text{pre}}_1\right)(\mathbf{X}) \tag{5}$$

MURA: Multi-Hop Retrieval Attention

The key innovation is recursive multi-hop neighborhood expansion. Given retrieval widths $\kappa_1, \dots, \kappa_m$, the multi-hop neighborhood is defined recursively:

$$\mathcal{N}_j^{(\kappa_1, \ldots, \kappa_r)}(i) = \bigcup_{t \in \mathcal{N}_j^{(\kappa_1, \ldots, \kappa_{r-1})}(i)} \mathcal{N}_j^{(\kappa_r)}(t), \qquad 2 \leq r \leq m \tag{10}$$

This allows retrieving chains of tokens connected by high-scoring attention links, rather than only immediate predecessors.

MURA sparse-attention layer (Definition 3.1): Attends only to retrieved multi-hop positions:

$$\text{head}_j(\mathbf{X}) = \text{softmax}\left(\frac{\mathbf{Q}_j \mathbf{K}_j^\top}{\sqrt{d/h}} + \mathbf{M}_{\boldsymbol{\kappa}_j}\right) \mathbf{V}_j \tag{11}$$

where the mask $M_{\kappa_j}$ allows attention only to positions in the multi-hop neighborhood.

MURA hybrid block (Definition 3.2):

$$\mathrm{B}_\kappa(\mathbf{X}) = \left(\mathrm{S}^{\text{post}}_{L_{\text{post}}} \circ \dots \circ \mathrm{S}^{\text{post}}_1 \circ \mathrm{A} \circ \mathrm{R}^{\text{pre}}_{L_{\text{pre}}} \circ \dots \circ \mathrm{R}^{\text{pre}}_1\right)(\mathbf{X}) \tag{12}$$

Efficiency Guarantees

For fixed retrieval depths and widths κ, each neighborhood contains a constant number of positions independent of sequence length N, preserving O(N) asymptotic computational and memory complexity.

---

Empirical Validation / Results

Key Theoretical Results

Theorem 2.7 (Limitation of top-κ blocks): Any top-κ access hybrid block computing m-function composition requires capacity:

$$C(\mathsf{B}_\kappa) = pd(L_{\text{post}} + \kappa + 1) = \Omega(\sqrt{n}) \tag{8, 9}$$

Theorem 3.3 (Sufficiency of MURA): For every fixed m ≥ 1, a single (1,...,1)-MURA block can compute m-function composition with:

• A single attention head
• A single MURA sparse attention layer
• No recurrent layers
• Width $d = O(m)$ independent of sequence length
• Precision $p = O(\log(N))$

Theorem 3.4 (Hierarchy): A κ-MURA block with m retrieval hops computing (m+1)-function composition requires:

$$C(\mathsf{B}_\kappa) = pd \, m \cdot \prod_{i=1}^m \kappa_i = \Omega(n^{1/(m+1)}) \tag{14}$$

Synthetic Experiments

Setup: 250,000 sequences encoding five functions $f_1, \dots, f_5: [n] \to [n]$ (n ∈ {25,...,30}). Models trained on causal token-prediction with multiple composition targets per token:

$$f_{i-1}(f_i(j)), (f_{i-2} \circ f_{i-1})(f_i(j)), \dots, (f_1 \circ \dots \circ f_{i-1})(f_i(j)) \tag{15}$$

MURA configurations with matched retrieval budgets (~32 tokens): (32,), (16,1), (11,1,1), (8,1,1,1) for 1–4 hops.

Results: Function-Composition Accuracy (%)

Key empirical findings:

• MURA achieves near-perfect accuracy on m-function composition when retrieval depth ≥ m
• The (8,1,1,1) configuration achieves 100% accuracy on all depths m ∈ {1,2,3,4}
• Existing baselines struggle with compositions beyond m = 1
• Sliding window (32), despite comparable retrieval budget, fails on deeper compositions

---

Theoretical and Practical Implications

Theoretical Significance

1. First architecture with single full-attention layer provably capable of multi-step function composition with constant width and logarithmic precision

2. Strict expressivity hierarchy in retrieval depth: Each additional retrieval hop enables higher-order composition with logarithmic capacity, while fewer hops require polynomial capacity—even with matched retrieval budgets, demonstrating that retrieval structure matters more than retrieval size

3. Stronger lower bounds than previous work: vs. Ye et al.'s Ω(N^(1/1024)) bound for hybrid architectures, this work establishes Ω(√n) with arbitrary recurrent depth and extension to top-κ access

Practical Implications

1. More efficient hybrid architectures: More expressive blocks could reduce the number of full-attention layers required, saving computation and KV-cache memory

2. Implementation feasibility: Multi-hop retrieval requires only an additional traversal of selected token indices, building naturally on existing FlashAttention-based implementations (as in HySparse)

3. Design principle: Retrieval depth emerges as a fundamental architectural principle for designing more expressive hybrid language models

---

Conclusion

Main Takeaways

• Existing hybrid blocks, whether conventional or with top-κ KV-cache access, have fundamental limitations in multi-hop reasoning (Ω(√n) capacity for two-function composition)

• MURA enables multi-hop retrieval from a shared KV cache while preserving the asymptotic computational and memory efficiency of existing hybrid architectures

• Strict expressivity hierarchy in retrieval depth: m hops suffice for m-function composition with constant width and logarithmic capacity, while (m+1)-function composition requires polynomial capacity

• Empirical validation confirms the theoretical hierarchy emerges through training, with MURA substantially outperforming existing hybrid baselines

Future Directions

• Scaling MURA to larger language models and real-world tasks
• Exploring whether the benefits of multi-hop retrieval transfer to long-context reasoning benchmarks
• Investigating adaptive retrieval depth and learned retrieval strategies
• Extending the communication-complexity framework to other architectural variations

The paper positions retrieval depth as a fundamental design principle, suggesting that enabling efficient layers to perform multi-hop retrieval from existing attention memory could be key to building more powerful and efficient hybrid language models.