SambaY Architecture: Efficient Long-Context Modeling

• SambaY is a hybrid decoder architecture that integrates state space models with attention mechanisms for efficient long-context sequence modeling.
• It employs a Gated Memory Unit to selectively share memory between decoder layers, reducing computational cost and eliminating explicit positional encoding.
• Empirical results show SambaY achieves up to 10× higher decoding throughput and lower irreducible loss compared to YOCO baselines in long-context tasks.

SambaY is a decoder–hybrid–decoder neural architecture designed to advance efficient sequence modeling, especially for LLMs tasked with long-context reasoning and generation. It synthesizes state space modules (specifically, the Samba self-decoder), selective cross-decoder attention, and a newly introduced Gated Memory Unit (GMU) to enable memory-efficient token mixing and improved throughput. SambaY achieves linear pre-filling time complexity, obviates explicit positional encoding, and empirically demonstrates lower irreducible loss and superior long-context task performance compared to strong baselines while remaining amenable to large-scale compute (2507.06607).

1. Architectural Structure and Key Components

SambaY belongs to the family of decoder–hybrid–decoder architectures that combine state space models and attention mechanisms for efficient reasoning with long sequences. The architecture comprises two main parts:

• Self-decoder (lower block): Implements Samba, a variant of efficient state space models or linear recurrent units, to process input tokens in parallel with linear-time complexity. This module naturally captures temporal dependencies and positional information without explicit encoding schemes.
• Cross-decoder (upper block): Incorporates a mixture of full-attention and memory-sharing mechanisms. It reuses the key–value (KV) cache from a single full-attention layer—mirroring YOCO's strategy—thereby maintaining long-range contextual access.

A principal novel element in this design is the Gated Memory Unit (GMU), which facilitates memory state sharing from the self-decoder to the cross-decoder. The GMU enables selective, element-wise reweighting of Samba layer output, controlling how memory representations are leveraged in higher decoder layers and significantly reducing the computational/passthrough cost for long generations.

Mathematically, the GMU is expressed as: $$\mathbf{y}_l = \Bigl(\mathbf{m}_{l'} \odot \sigma\bigl(W_1\,\mathbf{x}_l\bigr)\Bigr) W_2,$$ where

• $\mathbf{x}_l \in \mathbb{R}^{d_m}$: Input at current layer $l$
• $\mathbf{m}_{l'} \in \mathbb{R}^{d_h}$: Memory state (typically final self-decoder output)
• $W_1, W_2$: Learnable weight matrices
• $\sigma(\cdot)$: Non-linear activation (e.g., SiLU) A normalized variant (nGMU) may be adopted to support stable training.

2. Efficiency Mechanisms and Positional Encoding Elimination

SambaY's efficiency gains stem primarily from two mechanisms:

1. Substitution of Cross-Attention: Roughly half the cross-attention layers in the cross-decoder are substituted with GMU blocks. Element-wise gating is substantially less computationally intensive than full cross-attention, especially at large sequence lengths.
2. Memory Cache Reuse: The architecture uses the prefill KV cache of a single attention layer across multiple steps, mirroring the compute strategy of YOCO but with improved cost reduction by integrating the GMU.

A notable by-product of the Samba-based self-decoder is the natural encoding of recency and position information due to recurrence. As a consequence, explicit positional encodings (e.g., RoPE) are eliminated, reducing model complexity without loss of performance on tasks requiring sequential understanding.

3. Comparative Analysis: SambaY vs. YOCO

While YOCO employs a decoder–decoder strategy that reuses a single self-attention cache to reduce per-token decoding cost, it still incurs significant overhead from multiple cross-attention computations during autogressive generation. SambaY, by alternately replacing such layers with GMUs, reduces this overhead and demonstrates lower irreducible loss under scale.

Empirically, SambaY matches YOCO in structural and parametric design but delivers up to 10× higher decoding throughput for generation over long sequences (e.g., 2K-token prompts with 32K-token generation) and outperforms YOCO in loss scaling curves. This efficiency is achieved without resorting to architectural overparameterization or increased model depth, highlighting the effectiveness of memory gating and representation sharing.

4. Differential Attention and Task-Specific Enhancement

SambaY can be further enriched with Differential Attention (DA), wherein each attention layer incorporates additional channel-wise reweighting through learnable multipliers, affording finer adaptation to channel-specific information. In the Phi4-mini-Flash–Reasoning model, which is based on the SambaY+DA stack, Differential Attention leads to improved performance in complex reasoning and retrieval domains, notably on Math500, AIME24/25, and GPQA Diamond. These gains were achieved without reinforcement learning; the enhancements stem from architectural design and supervised training.

5. Mathematical Formulation and Computational Complexity

The architecture’s efficiency is stated in terms of its cross-decoder layer cost. In classic attention, the cost per layer is

$O(d_{kv}\,N)$

for cross-attention, with $N$ as context length and $d_{kv}$ as the key–value dimension. In SambaY, the GMU-supplanted layers require only

$O(d_h)$

per token, where $d_h$ (typically $2d_m$) is small; for reasonable $d_h/d_{kv}$ (e.g., below 128), this is especially advantageous for long contexts. The memory state output from the final Samba layer is cached along with the lone attention KV cache, supporting efficient representation sharing throughout decoding.

6. Empirical Results and Scaling Experiments

Extensive scaling studies confirm that, across models up to 3.4B parameters and training regimes exceeding hundreds of billions of tokens, SambaY establishes a lower irreducible loss than both the YOCO and a naïve Samba+YOCO combination. Table summarization:

Metric	SambaY	YOCO Baseline
Irreducible loss (scaling curves)	Lower	Higher
Decoding throughput (2K/32K setting)	Up to 10× higher	—
Math500/AIME24/25/GPQA Diamond Tasks	Consistent increases	Baseline

The improvements extend both to raw inference throughput and to final evaluation accuracies ("Pass@1" on AIME, multi-parameter benchmarks), particularly in long-context and reasoning benchmarks, establishing the architecture's practical relevance for large-scale sequence modeling.

7. Availability and Open-Source Resources

The entire training and evaluation codebase related to SambaY, including implementation of the GMU, differential attention enhancements, and experiment protocols, is available as open source at https://github.com/microsoft/ArchScale. This repository enables direct replication of published results as well as further exploration and extension of memory-efficient long-context reasoning models.