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

> Multi-hop retrieval attention (MURA) enables a single hybrid block to compute m-function composition with constant width, while conventional top-k access requires polynomial capacity.

- **Source:** [arXiv](https://arxiv.org/abs/2610.25250)
- **Published:** 2026-10-24
- **Permalink:** https://picx.dev/p/F0lNlB
- **Whiteboard:** https://picx.dev/p/F0lNlB/image

## Summary

# 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 (%)

| Model | m = 1 | m = 2 | m = 3 | m = 4 |
|---|---|---|---|---|
| MURA (32,) | 83.05 ± 29.35 | 44.47 ± 23.64 | 5.81 ± 3.49 | 3.68 ± 0.11 |
| MURA (16,1) | 100.00 ± 0.01 | 99.30 ± 1.22 | 30.31 ± 23.22 | 3.73 ± 0.10 |
| MURA (11,1,1) | 100.00 ± 0.00 | 100.00 ± 0.00 | 99.99 ± 0.01 | 27.81 ± 24.98 |
| MURA (8,1,1,1) | 100.00 ± 0.00 | 100.00 ± 0.00 | 100.00 ± 0.00 | 100.00 ± 0.00 |
| Linear attention | 99.97 ± 0.01 | 57.83 ± 10.75 | 14.97 ± 5.66 | 8.78 ± 0.62 |
| Sliding window (16) | 100.00 ± 0.00 | 46.83 ± 8.49 | 3.76 ± 0.23 | 3.72 ± 0.08 |
| Sliding window (32) | 99.99 ± 0.01 | 85.18 ± 4.30 | 59.86 ± 4.18 | 3.65 ± 0.02 |
| Gated DeltaNet | 82.07 ± 31.03 | 25.55 ± 16.14 | 13.89 ± 10.93 | 7.97 ± 2.48 |

**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.

---

_Markdown view of https://picx.dev/p/F0lNlB, served by PicX — AI-generated visual whiteboard summaries of research papers._