ArXiv: 2507.06607

🎯 Pitch

Replacing just half of a model's cross-attention layers with simple gated memory units can slash decoding I/O by orders of magnitude and boost throughput up to 10×, yet a 3.8B model built this way still outperforms a stronger RL-tuned reasoning baseline on AIME and GPQA—without any reinforcement learning.


1. Executive Summary

This paper introduces the Gated Memory Unit (GMU), a mechanism for efficient memory sharing across layers, and applies it to create SambaY, a decoder-hybrid-decoder architecture that replaces half of the cross-attention layers in YOCO with lightweight GMUs — reducing decoding I/O from O(dkv·N) to constant O(dh) per replaced layer — while preserving linear pre-filling complexity. Through scaling experiments up to 3.4B parameters on SlimPajama, SambaY achieves a lower irreducible loss (C = 0.58) than Samba+YOCO and Transformer++ baselines under a proposed µP++ hyperparameter transfer scheme, and the largest instantiation — Phi4-mini-Flash-Reasoning (3.8B parameters) — delivers up to 10× higher decoding throughput on 2K-length prompts with 32K generation length while outperforming the stronger Phi4-mini-Reasoning baseline (which includes RL) on Math500, AIME24/25, and GPQA Diamond, establishing that representation sharing via gated memory can substitute for costly cross-attention during long-chain-of-thought decoding without sacrificing reasoning accuracy.

2. Context and Motivation

The Core Problem: Attention's Memory I/O Bottleneck During Long Generation

The central challenge this paper tackles is specific to how modern language models handle long-form generation, particularly chain-of-thought (CoT) reasoning. During pre-filling (processing the user's input prompt), models compute key-value (KV) caches that store information needed for subsequent token generation. During decoding (generating the model's response, one token at a time), every self-attention or cross-attention layer must read these KV caches from memory to compute attention scores. The problem is that this memory read cost scales linearly with the sequence length: reading a KV cache of length NN with key/value dimension dkvd_{kv} costs O(dkvN)O(d_{kv} \cdot N) memory I/O per attention layer.

For typical chatbot scenarios (short prompts, short responses), this cost is manageable. But the landscape has shifted dramatically with reasoning models like OpenAI's o1 and DeepSeek-R1, which generate extremely long chains of thought — often tens of thousands of tokens — before producing a final answer. When a model generates a 32,000-token reasoning trace, every decoding step must read a KV cache that grows with each new token. For a model with 30+ cross-attention layers, each requiring a full KV cache read per generated token, the aggregate memory I/O becomes the dominant bottleneck, not the raw FLOPs of the attention computation itself.

This is the gap the paper addresses directly (Section 1):

"it does not mitigate the attention memory I/O cost for its cross-attentions during the generation stage of the model responses. This limitation becomes particularly pronounced for modern large language models (LLMs) that generate extensively long Chains-of-Thought (CoTs) for hard reasoning tasks."

Why This Matters Now

The paper motivates this problem along three dimensions that converge to make it urgent:

1. The rise of long-CoT reasoning models. Models like o1, DeepSeek-R1, and Phi4-mini-Reasoning achieve breakthrough performance on competitive math and science benchmarks by generating thousands of tokens of intermediate reasoning. The paper cites AIME24/25, Math500, and GPQA Diamond as benchmarks where long generation is the norm (Table 4). If generating high-quality reasoning requires 32K tokens but an architecture's decoding speed drops linearly with sequence length, then reasoning quality and inference cost are directly at odds. Users face a choice between accuracy (long CoT) and latency/cost, which is precisely what the paper seeks to resolve.

2. The pretraining-inference compute reallocation. As noted in the executive summary, comparing Phi4-mini-Flash-Reasoning against Phi4-mini-Reasoning, the paper shows that architectural efficiency gains during decoding can translate into better accuracy by freeing up compute budget that would otherwise be consumed by attention I/O. This is not just about speed — it's about the total cost of deploying reasoning models in production. If a model needs to generate 32K tokens per query, the memory I/O of cross-attention layers dominates the inference budget, and any reduction in that cost directly affects whether the system is economically viable.

3. The KV cache sharing paradigm is incomplete. Prior work on efficient inference (YOCO, cross-layer attention, multi-query attention) focused heavily on reducing the size of the KV cache or sharing it across layers, but didn't address whether all those cross-attention layers are necessary in the first place. The paper's insight is that once you've shared KV caches across layers (as YOCO does), you've already decoupled the memory retrieval from the computation — but you're still paying the retrieval cost. The question becomes: can some of those cross-attention layers be replaced with something cheaper that achieves comparable representational mixing? This is what the GMU addresses.

Prior Approaches and Where They Fall Short

The paper positions itself against a specific lineage of work on efficient decoder architectures. Understanding the progression is essential:

YOCO (You Only Cache Once): The State-of-the-Art Baseline

YOCO (Sun et al., 2024) is the direct predecessor and primary baseline. It split the transformer decoder into two halves:

  • A self-decoder (first half of layers) that uses linear-complexity token mixers (SSMs or sliding window attention) plus a single full-attention layer at the end that produces one set of KV caches.
  • A cross-decoder (second half of layers) where every layer is a cross-attention layer that reads from that single cached KV pair, rather than computing its own.

This solved the pre-filling bottleneck elegantly: during pre-filling, you only need to run the self-decoder (half the layers), and the cross-decoder can be skipped entirely because its KV cache is already computed from the self-decoder's output. The pre-filling complexity becomes linear in sequence length rather than quadratic.

Where YOCO falls short: During decoding, every single cross-attention layer in the cross-decoder must still read the full KV cache from memory for each generated token. The paper quantifies this precisely (Section 2):

"During decoding, we reduce the memory I/O complexity for half of the cross-attention layers from a linear cost of O(dkvN)O(d_{kv}N) to a constant O(dh)O(d_h), where NN is the sequence length and dkvd_{kv} is the dimension of key/value vectors. This leads to significant efficiency gains when Ndh/dkvN \gg d_h/d_{kv}, a condition that is easily met in practice since the ratio dh/dkvd_h/d_{kv} typically does not exceed 128."

For a typical model with dh=2dm=5120d_h = 2d_m = 5120 (for a 2560-width model) and dkv=64d_{kv} = 64 (head dimension), the ratio is 80. So when NN (generation length) exceeds 80 tokens, the GMU's constant cost per layer beats cross-attention's linear cost — and for a 32K generation, that's a 400× reduction per replaced layer.

Samba and Other Hybrid Architectures

Samba (Ren et al., 2025) — heavily cited as the self-decoder backbone — interleaves Mamba (an SSM) layers with sliding window attention (SWA) layers to combine the linear complexity of SSMs with the local retrieval capability of attention. However, Samba alone doesn't address the cross-decoder bottleneck during generation. It's primarily designed for efficient pre-filling and long-context encoding, not for efficient long-form decoding.

The paper notes that while Samba achieves "improved extrapolation perplexity on extremely long sequences," its "zero-shot retrievable context length remains limited to its sliding window size" (Section I). This limitation — that retrieval capability is bounded by the SWA window — is important context for why YOCO's full-attention-based KV cache sharing is still needed for long-context tasks, and why the paper targets the cross-decoder rather than the self-decoder for efficiency gains.

Other KV Cache Reduction Techniques

The paper situates its contribution against several lines of work on KV cache efficiency (Section I):

  • Multi-Query Attention (MQA) and Grouped-Query Attention (GQA) share key/value heads within a single layer, reducing cache size but not the number of cache reads. They reduce storage, not bandwidth.
  • Cross-Layer Attention (CLA) shares KV caches across adjacent layers, effectively reducing the number of distinct caches but still requiring every layer to read from memory. This is complementary to GMU — you could apply CLA and GMU together.
  • Layer-Condensed KV Cache and InfiniGen are dynamic cache management techniques that prefetch or compress KV entries. These are orthogonal: they optimize which cache entries to read, while GMU eliminates the read entirely for some layers.

The key distinction the paper draws is between reducing the amount of KV data stored (what prior work does) versus reducing the number of times KV data is read during decoding (what GMU does). These are fundamentally different optimization targets.

The Unanswered Question: Can We Share SSM Representations Across Layers?

The paper identifies a specific unexplored opportunity in the existing literature:

"prior works have not investigated the efficiency potential of representation sharing between SSM layers" (Abstract)

All prior work on representation sharing across layers focused on attention KV caches — sharing key-value pairs computed by one attention layer with subsequent layers. But SSM layers (like Mamba) produce fundamentally different intermediate representations: the output of the SSM kernel (the token-mixed state after the structured state-space computation). These SSM outputs encode temporal dependencies with recency bias, are much smaller than KV caches (a single vector of dimension dhd_h rather than an N×dkvN \times d_{kv} matrix), and have properties that make them potentially suitable for efficient cross-layer sharing.

The paper's Gated Memory Unit is the mechanism that enables this sharing. The key question it answers is: if we can take the SSM kernel output from one layer and use it in later layers through a learned gating mechanism, can those later layers achieve comparable representational mixing to what cross-attention would provide, at a fraction of the memory I/O cost?

This is not obvious a priori. The SSM output captures a specific kind of information — it's the result of a structured recurrent computation over the sequence — and it's unclear whether simple element-wise gating (as opposed to full cross-attention with softmax over the entire sequence) can extract useful signals from that representation. The paper's empirical results (Table 5, Section 4) show that it can, but with important caveats: gating SSM memory works well for long-context retrieval (Phonebook), but gating attention or MLP intermediate representations works less well, revealing that the recency bias in SSM outputs is part of what makes this sharing effective.

How This Paper Positions Itself

The paper situates its contribution at the intersection of three active research areas (Section I):

1. Efficient inference architectures (YOCO lineage). The paper directly extends YOCO's decoder-decoder paradigm but modifies the cross-decoder to reduce decoding cost. This is explicitly not a replacement for YOCO but an improvement: SambaY is "a decoder-hybrid-decoder architecture that incorporates GMUs in the cross-decoder to share memory readout states from a Samba-based self-decoder" (Abstract). The paper retains YOCO's linear pre-filling advantage and its KV-cache-sharing mechanism; it only changes what happens during decoding.

2. Hybrid SSM-attention models (Samba lineage). The paper uses Samba as its self-decoder backbone because it already interleaves Mamba with sliding window attention, providing a strong foundation for both efficient encoding and local retrieval. The innovation is extending the hybridization from the self-decoder into the cross-decoder: instead of all cross-attention layers, use a mix of cross-attention (which reads the full KV cache) and GMUs (which read only the SSM output state). This creates a decoder-hybrid-decoder — the "hybrid" refers to the mixed token-mixing mechanisms within each decoder half.

3. Scaling laws for architecture comparison. The paper introduces µP++ (Section 3.1), a hyperparameter transfer scheme that accounts for both depth and width scaling and applies zero weight decay to vector-like parameters. This is explicitly motivated by the need for fair architecture comparisons:

"a neural architecture's performance is tightly coupled with its optimization and initialization settings" (Section 3 preamble)

Without a principled scaling framework, you can't tell if an architecture's better performance is due to the architecture itself or just better hyperparameters. The paper develops µP++ to enable the scaling law fits that show SambaY's lower irreducible loss (Figure 2a, C = 0.58 vs. 0.64 for Transformer++), making the case that the efficiency gains are architectural, not just optimization artifacts.

The Specific Gap: A Spectrum of Memory Sharing Approaches

The paper implicitly lays out a spectrum of memory sharing approaches in sequence models:

ApproachWhat is sharedCost per layerRetrieval capability
Full cross-attention (YOCO)KV cache (n × d_kv matrix)O(d_kv · n)Full, softmax over all positions
GMU (SambaY)SSM kernel output (d_h vector)O(d_h)Implicit, through recency bias of SSM
No sharingNothingO(1) per tokenNone

The paper's thesis is that for roughly 50% of the cross-decoder layers, you don't need full cross-attention's retrieval capability — the SSM output, modulated by a learned gate, provides sufficient cross-layer information flow. The ablation study (Section 4, Table 5) validates this by showing that completely removing cross-attention (SambaY-AA) significantly degrades long-context performance (Phonebook drops from 78.1% to 46.9%), but replacing half the cross-attention layers with GMUs (SambaY) maintains or even improves performance while cutting decoding I/O roughly in half.

This middle-ground positioning — not eliminating cross-attention entirely, but strategically interleaving it with cheaper gating — is the paper's key architectural insight. It's justified by the observation that in the cross-decoder, different layers may serve different functions: some need full sequence-level attention, while others may primarily perform local feature refinement that can be accomplished through gated access to the SSM's recency-biased representation.

3. Technical Approach

3.1 Reader Orientation

The paper builds SambaY, a decoder-only language model architecture that replaces approximately half of the expensive cross-attention layers in the decoder–decoder framework (YOCO) with lightweight Gated Memory Units (GMUs) — element-wise gating operations that access a single compressed memory vector from the self-decoder's SSM layers — thereby cutting decoding memory I/O from linear $O(d_{kv}N)$ per replaced layer to constant $O(d_h)$ while preserving the linear pre-filling complexity and long-context retrieval capability of the full architecture.

The problem it solves is that during long-chain-of-thought generation, every cross-attention layer must read the full key-value cache (size proportional to sequence length $N$) for each new token, creating a bandwidth bottleneck that dominates inference cost. The solution shape is: interleave full cross-attention layers (which provide global token-level retrieval) with GMU layers (which provide cheap, channel-wise modulation of a pre-computed SSM memory) so that the model achieves comparable representational mixing at roughly half the memory bandwidth cost per generated token.

3.2 Big-Picture Architecture (Diagram in Words)

The system is a decoder–decoder architecture with three major structural blocks:

  1. Self-decoder (first half of layers): A Samba-style stack that interleaves Mamba-1 SSM layers with Sliding Window Attention (SWA) layers, ending with a single full-attention layer. The full-attention layer produces two outputs during pre-filling: (a) a standard key-value (KV) cache that encodes the entire prompt sequence, and (b) the SSM kernel output state $m \in \mathbb{R}^{d_h}$ — the compressed representation from the final Mamba layer that captures recency-biased sequence information.

  2. Memory cache (shared across the cross-decoder): Two pieces of data are stored after pre-filling and reused for every decoding step:

    • The KV cache $K_c, V_c \in \mathbb{R}^{N \times d_{kv}}$ from the self-decoder's full-attention layer, where $N$ is the processed sequence length (prompt + generated tokens so far) and $d_{kv}$ is the key/value head dimension.
    • The SSM memory $M^{(l')} \in \mathbb{R}^{n \times d_h}$ — the token-mixed output from the final Mamba layer in the self-decoder, which is the result of the structured state-space computation over the full sequence.
  3. Cross-decoder (second half of layers): A sequence of layers that process the token being generated. Instead of being all cross-attention (as in YOCO), roughly half of these layers are GMU layers and half remain cross-attention layers, interleaved. Each layer type reads from a different memory source:

    • Cross-attention layers: Read the full KV cache $K_c, V_c$ using the current token's hidden state as the query, computing standard softmax attention over all $N$ positions. Memory I/O cost: $O(d_{kv}N)$ per layer per generated token.
    • GMU layers: Read only the SSM memory $M^{(l')}$ (a single vector per position, dimension $d_h$) and modulate it element-wise using a gate computed from the current layer's input. Memory I/O cost: $O(d_h)$ per layer per generated token — constant with respect to sequence length.

Information flows as follows during decoding: the self-decoder's KV cache and SSM memory are fixed after pre-filling → for each new generated token, the cross-decoder processes it layer by layer → cross-attention layers perform full attention reads over the KV cache → GMU layers perform cheap gated modulation of the SSM memory → the cross-decoder output produces the next-token logits. The GMU layers never access the $N$-length KV cache, so their cost is independent of how long the generation becomes.

3.3 Roadmap for the Deep Dive

  • First, the Gated Memory Unit (GMU) as a mathematical operator — what computation it performs, how it modulates token mixing from a previous layer, and why it can substitute for a full cross-attention read. This is the core mechanism and everything else builds on it.
  • Second, how the GMU fits into the decoder–decoder framework — the YOCO baseline architecture, what changes SambaY makes to the cross-decoder, and what caches are stored. Understanding this requires seeing both the original and the modified architecture side by side.
  • Third, the mathematical formalism of token mixing — how both SSMs and attention produce the $M^{(l')}$ representation that the GMU gates. This is necessary to see why the "memory" being shared has specific properties (recency bias, structured recurrence) that make GMU gating effective.
  • Fourth, the self-decoder configuration choices — why Samba (Mamba + Sliding Window Attention) is chosen as the self-decoder, what the full-attention layer provides, and how the SSM memory is extracted and cached.
  • Fifth, the normalization design for GMU (nGMU) — why placing RMSNorm after element-wise multiplication (rather than before) is crucial for maintaining associativity between the gate and the token-mixing operator, particularly for linear attention-based SSMs.
  • Sixth, the µP++ hyperparameter scaling framework — how the paper enables fair architecture comparisons by controlling for depth, width, initialization, learning rate, and weight decay across different model families. This is meta-methodology that underpins all scaling claims in the paper.
  • Seventh, the iso-parametric equation method — how the paper ensures models of different architectures have comparable parameter counts and cache sizes when comparing them, which is non-trivial because different token mixers have different internal expansion ratios.

3.4 Detailed, Sentence-Based Technical Breakdown

This is primarily an architectural design paper whose core idea is that the structured, recency-biased output of an SSM layer (Mamba) can serve as a compressed memory that, when modulated by learned per-channel gates at later layers, provides sufficient cross-layer information flow to replace roughly half of the cross-attention layers in a decoder–decoder framework — dramatically reducing decoding memory bandwidth while preserving the model's ability to retrieve and reason over long contexts.


Gated Memory Unit (GMU): Core Computation and Intuition

The GMU is the atomic building block. It operates at a specific layer $l$ in the cross-decoder, taking two inputs: (1) the current layer's input hidden state $X_l \in \mathbb{R}^{n \times d_m}$ (where $n$ is the sequence length and $d_m$ is the model width), and (2) the token-mixed representation $M^{(l')} \in \mathbb{R}^{n \times d_h}$ from a previous layer $l'$ (where $l' < l$ and $d_h$ is the SSM's inner state dimension). The GMU produces an output $Y_l \in \mathbb{R}^{n \times d_m}$ through a gated modulation:

Yl=(M(l)σ(XlW1T))W2Y_l = \left( M^{(l')} \odot \sigma(X_l W_1^T) \right) W_2

where:

  • $M^{(l')} \in \mathbb{R}^{n \times d_h}$ is the memory — the token-mixed output from an earlier layer $l'$ (specifically, the final Mamba layer in the self-decoder for SambaY),
  • $X_l \in \mathbb{R}^{n \times d_m}$ is the current input to layer $l$ in the cross-decoder,
  • $W_1 \in \mathbb{R}^{d_h \times d_m}$ is a learnable weight matrix that projects the current input into the memory's channel space,
  • $\sigma(\cdot)$ is the SiLU (Sigmoid Linear Unit) activation function, producing gating values in $[0, \infty)$,
  • $\odot$ is element-wise (Hadamard) multiplication,
  • $W_2 \in \mathbb{R}^{d_h \times d_m}$ is a learnable weight matrix projecting the gated output back to model width.

What it computes: First, the current input $X_l$ is projected through $W_1$ and passed through SiLU to produce channel-wise gate values $G^{(l)} = \sigma(X_l W_1^T) \in \mathbb{R}^{n \times d_h}$. These gate values modulate the memory $M^{(l')}$ element-by-element: each element $(i,k)$ of the memory — representing position $i$ and channel $k$ — is multiplied by the corresponding gate $G^{(l)}_{ik}$. The gated result $M^{(l')} \odot G^{(l)}$ is then linearly projected back to the model's hidden dimension $d_m$ through $W_2$, producing the layer output $Y_l$. Critically, the gate $G^{(l)}$ is input-dependent — different tokens at different positions get different gating patterns based on the current layer's representation of those tokens.

Why this form: The element-wise multiplication with a learned gate achieves three things simultaneously:

  • Cheap computation: No attention matrix multiplication, no softmax over $N$ positions, no KV cache reads — it's just two linear projections plus element-wise ops. The parameter count is $d_h \cdot d_m + d_h \cdot d_m = 2 d_h d_m$, compared to attention's query, key, value, and output projections totaling $4 d_m^2$ (plus the KV cache storage and read cost).
  • Content-dependent modulation: The gate $G^{(l)}$ depends on $X_l$, meaning the model can selectively amplify or suppress different channels of the SSM memory based on the current context. A token at the beginning of a sentence might gate positional channels differently than a token mid-sentence.
  • Linearity-preserving over the value matrix: As shown in the paper (Section 2), the element-wise product lifts the token-mixing matrix $A^{(l')}$ into a third-order tensor $\tilde{A}_{ijk} = G^{(l)}_{ik} A^{(l')}_{ij}$, yielding a learned, channel-specific reweighting while maintaining linearity on the original value matrix $V^{(l')}$. This means the GMU doesn't destroy the structured information encoded by the SSM's recurrent computation — it only reweights it per-channel.

The alternative would be to not use any memory sharing at all (each layer computes its own representations from scratch, requiring either full attention or its own SSM), which loses the efficiency and representational benefits of reusing pre-computed sequence information. The gating mechanism provides a middle ground: access to sequence-level information (via $M^{(l')}$) without paying the I/O cost of reading the full KV cache.


Token Mixing as a Matrix Operator: The Memory That GMU Gates

To understand what the GMU is gating, we need to understand how $M^{(l')}$ is produced. The paper (Section 2) establishes a unifying view of both SSMs and attention as producing token-mixed representations through a linear operator $A \in \mathbb{R}^{n \times n}$:

"Both state-space models (SSMs) and self-attention layers perform token mixing through a linear operator that can be written as a matrix $A \in \mathbb{R}^{n \times n}$, where $n$ is the sequence length."

For a given head at layer $l'$, the mixed representation is:

M(l)=A(l)V(l)M^{(l')} = A^{(l')} V^{(l')}

where:

  • $A^{(l')} \in \mathbb{R}^{n \times n}$ is the token-mixing matrix — in SSMs, it's a highly structured matrix capturing the parallel form of an underlying recurrent update; in self-attention, it's the row-aggregating softmax attention matrix whose entries are query-key similarity probabilities,
  • $V^{(l')} \in \mathbb{R}^{n \times d_h}$ is the value matrix — for SSMs, the state inputs to the structured state-space computation; for attention, the value vectors projected from the input,
  • $M^{(l')} \in \mathbb{R}^{n \times d_h}$ is the token-mixed output — the result of applying the mixing operator to the values.

What this formalism reveals: The GMU's operation can be rewritten to show that gating the token-mixed output is equivalent to reweighting the token-mixing matrix itself. For each element $H_{ik}$ of the gated output $H = M^{(l')} \odot G^{(l)}$:

Hik=Gik(l)jAij(l)Vjk(l)=j(Gik(l)Aij(l))A~ijkVjk(l)H_{ik} = G^{(l)}_{ik} \cdot \sum_j A^{(l')}_{ij} V^{(l')}_{jk} = \sum_j \underbrace{\left( G^{(l)}_{ik} \cdot A^{(l')}_{ij} \right)}_{\tilde{A}_{ijk}} \cdot V^{(l')}_{jk}

This shows that the gate $G^{(l)}$ effectively lifts $A^{(l')}$ — a rank-2 matrix — into a rank-3 tensor $\tilde{A}_{ijk}$, where each channel $k$ gets its own reweighted version of the original token mixing. The result is a channel-specific, content-dependent modulation of how information from position $j$ flows to position $i$, computed at layer $l$ based on the current representation $X_l$ — all while maintaining linearity on the value matrix $V^{(l')}$ (no expensive recomputation of the mixing itself).

Why this matters for efficiency: The token-mixing $A^{(l')}$ and value projection $V^{(l')}$ are computed once in the self-decoder during pre-filling. The GMU at layer $l$ only computes the gate and the element-wise product — operations whose cost is $O(n \cdot d_h)$ with a much smaller constant factor than attention's $O(n \cdot d_{kv})$, and critically does not require reading the full KV cache from memory. The associativity property ensures that reweighting the mixing operation through the gate is mathematically equivalent to having the gate participate in the original token mixing, even though the gate is computed at a much later layer.

The paper emphasizes (Section B) that this associativity holds only when the normalization is placed after the gating, not before — a crucial detail we cover below.


The SambaY Architecture: Self-Decoder, Cross-Decoder, and Memory Caching

SambaY is built on top of the YOCO (You Only Cache Once) decoder–decoder architecture. To understand SambaY, we must first understand what YOCO does and what SambaY changes.

YOCO baseline (Section A): YOCO splits the transformer stack into two halves:

  • Self-decoder (layers 1 to L/2): Uses token mixers with linear computational complexity (SSMs or sliding window attention), culminating in a single full-attention layer at the end. During pre-filling, this full-attention layer computes and caches the key-value pairs $K_c, V_c \in \mathbb{R}^{n \times d_{kv}}$ from its input $X_{mem} \in \mathbb{R}^{n \times d_m}$: Kc=XmemWK,Vc=XmemWVK_c = X_{mem} W_K, \quad V_c = X_{mem} W_V where $W_K, W_V$ are weight matrices. These KV caches are stored and reused by all subsequent cross-attention layers.
  • Cross-decoder (layers L/2+1 to L): Every layer is a cross-attention layer that reads from the shared $K_c, V_c$. For layer $l$ with input $X_{cross}^{(l-1)}$, it computes: Qcross(l)=Xcross(l1)WQ(l),H(l)=softmax(Qcross(l)KcTdkv)VcQ_{cross}^{(l)} = X_{cross}^{(l-1)} W_Q^{(l)}, \quad H^{(l)} = \text{softmax}\left(\frac{Q_{cross}^{(l)} K_c^T}{\sqrt{d_{kv}}}\right) V_c

The key efficiency of YOCO comes from pre-filling: only the self-decoder is run (the cross-decoder is skipped entirely because its KV cache source is already computed). This cuts pre-filling FLOPs roughly in half and maintains linear complexity in $n$. However, during decoding, every cross-attention layer must still read the full $K_c, V_c$ cache for each generated token — the memory I/O cost per layer per token is $O(d_{kv} \cdot n)$, which grows with the generation length.

SambaY's modification (Section 2): SambaY replaces the self-decoder with Samba (Mamba-1 SSM interleaved with Sliding Window Attention) and replaces approximately half of the cross-attention layers in the cross-decoder with GMU layers. The GMU layers do not access $K_c, V_c$ at all. Instead, they access the SSM memory $M^{(l')}$ — the token-mixed output from the final Mamba layer in the self-decoder.

The paper specifies exactly what is cached (Section 2):

"Compared to YOCO, our approach only requires caching an additional SSM kernel output state $m \in \mathbb{R}^{d_h}, d_h = 2d_m$ from the final Mamba layer, an overhead that is negligible in size, alongside the KV cache from the last full-attention layer during pre-filling."

So during pre-filling, two caches are stored:

  1. $K_c, V_c \in \mathbb{R}^{n \times d_{kv}}$ — the standard KV cache for cross-attention layers.
  2. $M^{(l')} \in \mathbb{R}^{n \times d_h}$ — the SSM kernel output state from the final Mamba layer, where $d_h = 2d_m$.

The additional cached SSM state is $2d_m$-dimensional per position — for a model with $d_m = 2560$, this is 5,120 floats per position, or about 20 KB per position in FP32, which is indeed "negligible in size" compared to the KV cache (which for the same model with $d_{kv} = 64$ and a few heads would be hundreds of bytes per position — except that the KV cache is reused across all cross-attention layers, so the total KV cache storage is $n \cdot d_{kv} \cdot (\text{number of heads})$, which can be much larger).

Decoding cost reduction (Section 2):

"During decoding, we reduce the memory I/O complexity for half of the cross-attention layers from a linear cost of $O(d_{kv}N)$ to a constant $O(d_h)$, where $N$ is the sequence length and $d_{kv}$ is the dimension of key/value vectors."

This is the central efficiency claim. For each GMU layer, the memory read is just $M^{(l')}$ — a single vector of dimension $d_h$ per position — rather than the full $N \times d_{kv}$ KV matrix. The ratio of costs is $d_h / (d_{kv} \cdot N)$. For typical values ($d_h = 5120$, $d_{kv} = 64$, $N = 32,000$), the GMU layer's memory read is ~400× cheaper than cross-attention's. With roughly 50% of layers replaced, the total cross-decoder memory I/O is roughly halved.

The paper summarizes the condition for significant efficiency gains:

"This leads to significant efficiency gains when $N \gg d_h / d_{kv}$, a condition that is easily met in practice since the ratio $d_h/d_{kv}$ typically does not exceed 128."


Self-Decoder Design: Why Samba (Mamba + SWA + Full Attention)

The self-decoder's job is to process the input sequence and produce both the KV cache and the SSM memory in a single forward pass during pre-filling. The paper chooses Samba (Ren et al., 2025) as the self-decoder backbone, which interleaves three types of token-mixing layers:

  1. Mamba-1 layers: Selective state-space models that process the sequence with linear complexity and a recency bias (more recent tokens have stronger influence on the current state). These produce the token-mixed output $M^{(l')}$ that the GMU will later gate.
  2. Sliding Window Attention (SWA) layers: Attention restricted to a local window of size $W$ (default 128 in scaling experiments, tuned per architecture for long-context). These provide local retrieval capability that pure SSMs lack.
  3. A single full-attention layer at the end of the self-decoder: This layer computes the KV cache $K_c, V_c$ that the cross-decoder's remaining cross-attention layers will read. It allows the model to establish global token-level dependencies that the SSM and SWA layers might miss.

The paper provides specific ratios for the self-decoder layers (Section D, Table 7). For the 1B-parameter scaling models with $d = 16$ (depth 16, total layers 32), the self-decoder (first 16 layers) contains:

  • Mamba layers (the majority of the self-decoder)
  • SWA layers (interleaved with Mamba)
  • 1 full-attention layer (at the end of the self-decoder)

The exact arrangement is inherited from Samba's design, which the paper doesn't re-tune — the focus is on modifying the cross-decoder.

Why Samba over pure Mamba or pure SWA: The paper's ablation study (Section 4, Table 5) compares different self-decoder backbones on the Phonebook 32K retrieval task:

  • MambaY (pure Mamba + full attention in self-decoder, no SWA): Phonebook accuracy drops to 12.5% (vs. 78.1% for SambaY), showing that recency bias alone is insufficient for long-context retrieval — the local attention from SWA is necessary.
  • SambaY-2 (replacing Mamba-1 with Mamba-2): Phonebook drops to 40.6% (vs. 78.1%), attributed to Mamba-2's "coarse, scalar-valued forget gates" which reduce the self-decoder's capacity to encode fine-grained positional information needed for the GMU to later extract useful signals.
  • S-GDNY (replacing Mamba with Gated DeltaNet): Phonebook improves to 83.6%, showing that more sophisticated linear attention with delta update rules can produce even better memory for the GMU.

This validates the design choice: Mamba-1 + SWA provides the right balance of recency bias (for the GMU to exploit) and local retrieval (for the self-decoder's own processing), with the full-attention layer providing global context for the remaining cross-attention layers.


Normalized GMU (nGMU) and the Critical Role of Normalization Placement

The paper (Section 2, Section B) introduces an important variant: the normalized GMU (nGMU), which applies RMSNorm after the element-wise multiplication:

Yl=RMSNorm(M(l)σ(XlW1T))W2Y_l = \text{RMSNorm}\left( M^{(l')} \odot \sigma(X_l W_1^T) \right) W_2

Why normalization is needed: When the memory $M^{(l')}$ comes from linear attention architectures (like Gated DeltaNet or Mamba-2), the standard design places normalization before the output gating in the source layer to stabilize training:

M(l)=Norm(A(l)V(l))M^{(l')} = \text{Norm}\left( A^{(l')} V^{(l')} \right)

However, this placement breaks the associativity between the gate and the token-mixing operator — if normalization happens before gating at the source layer, the GMU at layer $l$ cannot directly reweight the original token-mixing matrix $A^{(l')}$ because the norm has already transformed the representation non-linearly.

The paper's solution (Section B): For Gated DeltaNet (GDN) layers in the self-decoder, postpone the normalization to after the output gating (denoted as GDN-A):

M(l)=A(l)V(l),Y(l)=Norm(M(l)G(l))W2(l)M^{(l')} = A^{(l')} V^{(l')}, \quad Y^{(l')} = \text{Norm}\left( M^{(l')} \odot G^{(l')} \right) W_2^{(l')}

Then in the cross-decoder, use nGMU — which applies normalization after the GMU's own gating — to maintain training stability while preserving associativity:

Yl=Norm(M(l)G(l))W2Y_l = \text{Norm}\left( M^{(l')} \odot G^{(l)} \right) W_2

Empirical validation (Section H, Table 9): The ablation study in Table 9 provides strong evidence for this design choice. Comparing S-GDNY (which uses GDN-A with nGMU) against "GDN + GMU" (which uses original GDN with normalization before gating and basic GMU without normalization):

  • S-GDNY achieves 83.6% on Phonebook 32K
  • GDN + GMU achieves only 27.3% on Phonebook 32K
  • Short-context benchmarks (Wiki, LAMBADA, ARC, HellaSwag, PIQA, WinoGrande) are largely unaffected (differences < 3%)

The paper attributes this to maintaining "the associativity between gating and token mixing." Without the correct normalization placement, the GMU cannot effectively reweight the token-mixing information from the SSM layer, and long-context retrieval — which depends on accessing fine-grained positional dependencies in the memory — degrades severely. Short-context tasks, which rely more on local information accessible through the SWA layers and cross-attention, are less affected.


The Iso-Parametric Equation Method for Fair Architecture Comparison

A major challenge in comparing architectures is that different token mixers have different internal expansion ratios, making it hard to match parameter counts. Adding two attention layers to a hybrid model doesn't just change parameters — it changes the KV cache size, affecting inference cost. The paper's solution (Section 3.1) is the iso-parametric equation method:

Step 1: Define the scaling rule for the Transformer baseline. Using a simple linear rule from prior work (Kaplan et al., 2020), the paper defines the architectural shape of the Transformer++ baseline:

w=αd,α=α0=128,hq=d,hkv=d/4,wmlp=4ww = \alpha d, \quad \alpha = \alpha_0 = 128, \quad h_q = d, \quad h_{kv} = d/4, \quad w_{mlp} = 4w

where:

  • $w$ is the model width,
  • $d$ is the model depth,
  • $\alpha$ is the aspect ratio (width per depth unit),
  • $h_q$ is the number of query heads,
  • $h_{kv}$ is the number of key-value heads (GQA with group size 4),
  • $w_{mlp}$ is the MLP inner dimension.

The total non-embedding parameters for Transformer++ are then:

N(d)=Nattn(d)+Nmlp(d)=2.5dw2+12dw2=14.5dw2=237568d3N(d) = N_{attn}(d) + N_{mlp}(d) = 2.5 d w^2 + 12 d w^2 = 14.5 d w^2 = 237568 d^3

This gives the parameter count as a function of depth $d$ for the baseline.

Step 2: Write the parameter count formula for the hybrid architecture. For SambaY, the layers include attention (both full and sliding window), Mamba (with its own internal expansion), MLP, and GMU. The paper enumerates each component:

Nattn(d)=2.5dwwattn/4+2dwwattn/4N_{attn}(d) = 2.5 d w \cdot w_{attn}/4 + 2 d w \cdot w_{attn}/4 Nmamba(d)=6dw2/4N_{mamba}(d) = 6 d w^2 / 4 Ngmu(d)=4dw2/4N_{gmu}(d) = 4 d w^2 / 4 N(d)=Nattn(d)+Nmamba(d)+Nmlp(d)+Ngmu(d)N(d) = N_{attn}(d) + N_{mamba}(d) + N_{mlp}(d) + N_{gmu}(d)

where:

  • The first attention term accounts for the full-attention layer (2.5× because of query, key, value, and output projections),
  • The second attention term accounts for the cross-attention layers (2× because keys and values come from cache, so only query and output projections are learned),
  • The Mamba term accounts for the SSM layers (6× expansion relative to width),
  • The GMU term accounts for the gating projections (4× expansion, with $W_1$ and $W_2$ together contributing $2 d_h d_m = 4 d w$? Actually the factor is $4 d w^2 / 4 = d w^2$ — check: $W_1 \in \mathbb{R}^{d_h \times d_m}$ with $d_h = 2 d_m$, so $W_1$ has $2 d_m^2$ params, $W_2$ has $2 d_m^2$ params, total $4 d_m^2$ per GMU layer. Since $w = d_m$ and there are $d/4$ GMU layers (half of the cross-decoder), total GMU params $= d/4 \cdot 4 w^2 = d w^2$. Wait, the paper writes $4 d w^2 / 4$ which simplifies to $d w^2$. Yes, that matches.)

With $w_{attn} = \alpha d$, the total becomes:

N(d)=144αd3+14.5α2d3=237568d3N(d) = 144 \alpha d^3 + 14.5 \alpha^2 d^3 = 237568 d^3

Step 3: Solve for $\alpha$. Setting the hybrid architecture's non-embedding parameter count equal to the Transformer baseline's, we solve the quadratic equation for the aspect ratio:

144α+14.5α2=237568144 \alpha + 14.5 \alpha^2 = 237568

This yields $\alpha_1 \approx 124$ for SambaY (vs. $\alpha_0 = 128$ for Transformer++).

Similarly, for Samba+YOCO: $\alpha_2 \approx 126$, for MambaY: $\alpha_3 \approx 120$, for SambaY-MLP: $\alpha_4 \approx 120$, and so on (Section C).

What this achieves: Models with different architectures but the same depth $d$ will have the same total number of non-embedding parameters. Furthermore, by fixing the head dimension to $\alpha d$ and the GQA group size to 4, the KV cache size (a function of number of key-value heads times head dimension) is also matched across architectures. This means comparisons of inference cost and training FLOPs are fair — you're not accidentally comparing a 1B-parameter hybrid model against a 1.2B-parameter Transformer.

The paper's example (Section 3.1):

"we build an iso-parametric equation with respect to the aspect ratio via aligning the total number of non-embedding parameters to the Transformer baseline"

This is explicitly contrasted with prior work:

"Previous works often adjust the model depth to tie the total number of parameters, but this could change the memory cache size significantly (e.g. adding two attention layers in a 12-layer Transformer resulting in a 16.7% increase of KV cache size), making unfair comparisons on the inference time cost."


µP++ Hyperparameter Scaling for Depth, Width, and Stability

The paper introduces µP++ hyperparameter scaling laws (Section 3.1, Table 6) that extend µP (Yang et al., 2022) and Depth-µP (Yang et al., 2023) with an additional stability modification: applying zero weight decay to vector-like and scalar-like parameters.

The key components of µP++ are:

Width scaling (from µP): The learning rate of matrix-like parameters (hidden weights) is scaled proportionally to $1/w$, where $w$ is the model width. This ensures that the feature learning dynamics are independent of width — a wider model doesn't effectively have a smaller or larger step size in function space.

Depth scaling (from Depth-µP): The learning rate is further scaled as $\eta \propto 1/\sqrt{d}$, where $d$ is the model depth. The output of each residual branch is divided by $\sqrt{2d}$ to prevent the variance from exploding with depth. These rules ensure stable training dynamics as models get deeper.

Weight decay modification (new in µP++): Zero weight decay is applied to:

  • Vector-like parameters: Parameters where exactly one dimension scales with model width (e.g., embedding and unembedding layers, RMSNorm weights, biases).
  • Scalar-like parameters: Parameters where no dimension scales with width.

The paper (Section E, Figure 7a) shows that without this modification, the original µP setup "can lead to severe training instability when scaling to 600B tokens," with increasing gradient norms and large spikes for vector-like parameters shortly before divergence. The zero weight decay on these parameters stabilizes training at large scales.

Batch size, learning rate, and training tokens: The base configuration (Section 3.1) is:

  • Base learning rate $\eta_0 = 4 \times 10^{-4}$
  • Base batch size $B_0 = 2^{21} = 2\text{M}$ tokens
  • Base model depth $d_0 = 16$, corresponding to $N(d_0) \approx 10^9$ parameters
  • Base training tokens $T_0 = 100\text{B}$

For scaling to larger models, the learning rate is adjusted as:

η=η0Bd0B0d\eta = \eta_0 \sqrt{\frac{B d_0}{B_0 d}}

where $B = B_0$ is kept constant across scales (batch size does not scale with model size in the primary experiments — the paper found that scaling batch size sub-linearly with training tokens "harms the data scaling behavior of the models").

Training tokens scale linearly with model parameters (Chinchilla-inspired over-training regime):

T=T0N(d)N(d0)T = T_0 \frac{N(d)}{N(d_0)}

This gives a 5×5\times Chinchilla-optimal ratio of tokens per parameter (approximately 100 tokens per parameter instead of the 20 recommended by Hoffmann et al., 2022), placing the models in a "typical over-training regime" that the paper argues is more representative of practical LLM training.

Comparison to Standard Parameterization (SP) and µP (Table 6):

Parameter TypeSP InitµP InitµP++ InitµP++ WD
EmbeddingN(0, σ²)N(0, σ²)N(0, σ²)0
Unembedding0 or tied0 or tied0 or tied, LR ∝ 1/w0
Hidden weightsN(0, τ²)U(-β/√fan_in, β/√fan_in), LR ∝ 1/wU(-β/√fan_in, β/√fan_in), LR ∝ 1/w, Res. mult. = 1/√(2d)∝ 1

For standard parametrization, the paper uses LeCun uniform initialization (PyTorch default) for weight matrices. For µP and µP++, the hidden weight initialization is uniform in $[-\beta/\sqrt{\text{fan\_in}}, \beta/\sqrt{\text{fan\_in}}]$ with $\beta = 1$, and the learning rate multiplier is $1/w$. µP++ additionally divides the residual branch output by $\sqrt{2d}$ and applies zero weight decay to scalar/vector parameters.

Optimizer configuration (Section D): All experiments use AdamW with:

  • $\beta_1 = 0.9$, $\beta_2 = 0.95$, $\epsilon = 10^{-8}$
  • Weight decay = 0.1 (for matrix-like parameters; zero for scalar/vector)
  • Linear learning rate schedule: 1B warm-up tokens linearly increasing to peak $\eta$, followed by linear decay to zero
  • Tied input and output embedding matrices, initialized from $N(0, 0.02^2)$

The attention logit scaler is set to $1/\sqrt{d_{kv}}$ where $d_{kv}$ is the head dimension.

Why µP++ over alternatives: The paper's ablation (Section E, Figure 7) compares µP++ against several alternatives:

  • µP (no depth scaling, no zero-WD on vector params): Training diverges with NaN losses after 204K steps at 600B tokens.
  • µP++ with batch scaling (batch size increases with training tokens): Worse learning efficiency and irreducible loss than constant-batch µP++.
  • µP++ with Normal Init (variance scaled by 1/d globally): Worse scaling than LeCun uniform initialization.
  • µP++ with LR scaling and independent weight decay: Worse learning efficiency (smaller exponent $b$).

The paper concludes (Section E): "it is better to adjust the initialization multipliers based on each matrix's dimension as adopted by LeCun initialization, rather than a global factor related to model width," and that the constant batch size scaling is preferable for the training regime studied.


Scaling Experiment Configurations and the Power-Law Fitting Methodology

With the architecture shapes determined by iso-parametric equations and the hyperparameters determined by µP++, the paper runs two types of scaling experiments (Section 3.1):

Data scaling: Fix model size at 1B parameters ($d = 16$), vary training tokens $T \in \{100\text{B}, 200\text{B}, ..., 600\text{B}\}$. This studies how performance improves with more data at a fixed model capacity.

FLOPs (compute) scaling: Vary both model size and training data proportionally. Model depth $d \in \{8, 12, 16, 20, 24\}$ gives parameter counts from ~154M to ~3.4B. Training tokens scale linearly with parameters according to $T = T_0 \cdot N(d)/N(d_0)$. The maximum configuration is $d = 24$, ~3.4B parameters, ~342B tokens.

Training details: All scaling experiments use:

  • Sequence length: 4K tokens
  • Dataset: SlimPajama (Soboleva et al., 2023), a cleaned and deduplicated 627B-token version of RedPajama
  • Hardware: Not explicitly specified for scaling runs, but the large-scale pretraining (Section 3.3) uses 1K A100-80GB GPUs

Power-law fitting (Section 3.1): To quantitatively compare scaling trajectories, the paper fits the validation loss $L$ as a function of compute (FLOPs) or data (tokens) to a power law of the form:

L(DFLOPs)=ADFLOPsb+CL(D_{\text{FLOPs}}) = A \cdot D_{\text{FLOPs}}^{-b} + C

where:

  • $D_{\text{FLOPs}}$ is the total training FLOPs,
  • $A$ is the coefficient (affects initial loss and convergence speed),
  • $b$ is the scaling exponent (affects learning efficiency — how quickly loss decreases with compute),
  • $C$ is the irreducible loss — the theoretical lower bound on validation loss achievable with infinite compute, representing the entropy inherent in the data distribution plus any irreducible model error.

The fitting uses least squares and the Levenberg-Marquardt algorithm. All reported fits have $R^2 \geq 0.999$, indicating near-perfect power-law behavior.

Why emphasize irreducible loss $C$: The paper explicitly argues (Section 3.1):

"While larger values of the scaling exponent $b$ or the coefficient $A$ indicate that a model may converge more rapidly given a small-scale compute or data budget, these parameters alone do not necessarily predict superior performance at larger scales. Therefore, we emphasize the irreducible loss $C$ obtained from scaling law fitting as the primary metric for assessing an architecture's long-term scaling potential."

This is the key methodological insight: better short-term convergence (larger $b$ or smaller $A$) doesn't guarantee better asymptotic performance. The irreducible loss $C$ — what the model's loss would approach if you trained forever — is the better predictor of which architecture will win at extreme scales. The paper finds that SambaY has the lowest $C = 0.58$ for FLOPs scaling, compared to 0.64 for Transformer++ (Figure 2a), suggesting superior scaling potential with substantially increased computational resources.

A note on µP++ vs. SP in scaling (Section E): The paper also compares µP++ against Standard Parametrization (SP) in the scaling experiments. Under SP, the Transformer++ baseline shows higher irreducible loss than under µP++ (Figure 2), indicating that the µP++ hyperparameter transfer scheme itself contributes to better scaling behavior — not just the architecture.


Large-Scale Pretraining Configuration (Phi4-mini-Flash)

For the 3.8B-parameter Phi4-mini-Flash model (Section 3.3, Table 7), the configuration differs from the scaling experiments:

  • Architecture: SambaY+DA (SambaY enhanced with Differential Attention)
  • Depth: $d = 32$ (total 64 layers, first 32 in self-decoder, last 32 in cross-decoder)
  • Model width: 2560
  • Attention heads: 40 query heads, 20 key-value heads (GQA group size 2)
  • Head dimension: 64 (reduced from 128 in scaling experiments)
  • SWA size: 512
  • MLP inner dimension: 10240
  • Vocabulary: 200K tokens (matching Phi4-mini)
  • Aspect ratio: $\alpha = 80$ (different from the iso-parametric $\alpha \approx 124$ used in scaling)
  • Parameterization: Standard Parameterization (not µP++), due to "resource constraints at the time of scaling study"
  • Batch size: 8M tokens
  • Learning rate schedule: Linear warm-up over 3,000 steps, then linear decay
  • Training data: 5T tokens from the Phi4-mini data corpus, trained on 1K A100-80GB GPUs for 14 days

Training instability mitigation (Section D, Figure 6): Two tricks were required to stabilize training:

  1. FP32 up-casting in the fused cross-entropy loss kernel: The paper modifies the Liger-Kernel's fused linear cross-entropy loss to up-cast both the weight and input to FP32 during chunk-wise matrix multiplication. Without this, "the gradient norm during the training process... with up-shooting trends that will finally blow up the training loss." The FP32 up-casting stabilizes the gradient norm to match the naive no-fusion baseline, while maintaining the speed advantage of kernel fusion (critical for the 200K vocabulary).

  2. Attention dropout added mid-training: 0.05 attention dropout is added to the self-decoder at 310K steps and to the cross-decoder at 520K steps — the last checkpoints before loss divergences were observed. The dropout successfully stabilizes training without harming downstream MMLU performance.

The paper acknowledges (Section 3.3): "The optimization setup here is by no means optimal, as the primary goal of this experiment is to evaluate the viability of our architecture at larger scales."


Differential Attention Integration

Differential Attention (DA) (Ye et al., 2024) is a variant of attention that computes the difference between two softmax attention maps, which has been shown to reduce attention noise and improve retrieval. In SambaY+DA (Section 3.2), DA replaces the standard attention in all attention layers (full attention, SWA, and cross-attention).

The paper describes the DA initialization (Section D):

"Differential Attention uses a depth-dependent initialization factor, $\lambda_{init} = 0.8 - 0.6 \exp(-0.3 \times l)$, where $l$ is the depth index."

For each attention head, it employs two sets of learnable parameters $(\lambda_{q1}, \lambda_{k1})$ and $(\lambda_{q2}, \lambda_{k2})$, each of dimension equal to the head dimension and initialized with a normal distribution of zero mean and 0.1 standard deviation.

The long-context results (Table 1) show that DA improves multi-key retrieval (MK-1: 64.6% vs. 54.6% for SambaY) and single-needle retrieval (S-2: 86.4% vs. 81.2%), though with a larger optimal SWA size (512 vs. 256), which slightly reduces training throughput.

A note on DA implementation efficiency (Section 3.4):

"our Differential Attention implementation relies on a naive four-pass of the FlashAttention operator for vLLM compatibility, rather than the optimized custom kernel proposed in the original paper, leaving significant room for further speed optimization."

This means the throughput numbers reported for Phi4-mini-Flash-Reasoning (Figure 4) are pessimistic relative to what an optimized DA kernel could achieve — the 10× speedup claim is with suboptimal DA, so future optimizations could widen the gap further.


Reasoning Model Distillation Pipeline

Phi4-mini-Flash-Reasoning (Section 3.4) is produced by continuing to train the pre-trained Phi4-mini-Flash with distillation data following the Phi4-mini-Reasoning recipe (Xu et al., 2025):

Multi-stage distillation:

  1. Supervised Fine-Tuning (SFT): Train on reasoning traces (chains of thought) from a teacher model.
  2. Direct Preference Optimization (DPO): Further refine using preference pairs.

Critically, the paper notes:

"Due to the limited resources, we only conduct the distillation with SFT and DPO stages and leave RL for future works."

This makes the comparison against Phi4-mini-Reasoning (which does include a final RL training stage) particularly notable — Phi4-mini-Flash-Reasoning outperforms Phi4-mini-Reasoning on AIME24 (52.3% vs. 48.1%), AIME25 (33.6% vs. 31.8%), Math500 (92.5% vs. 91.2%), and GPQA Diamond (45.1% vs. 44.5%) without RL, using only the more efficient architecture as the differentiating factor.

Evaluation details (Section G):

  • Sampling temperature: 0.6
  • Top-p: 0.95
  • Maximum sequence length: 32,768 tokens
  • Evaluation libraries: Math-Verify (v0.7.0) and Lighteval (v0.10.0)
  • Instruction prepended: "Please reason step by step, and put your final answer within \boxed{}" for math benchmarks, with "final choice of one letter from A/B/C/D" for GPQA Diamond
  • Pass@1 with averaging: 64 samples for AIME24/25, 8 samples for Math500 and GPQA Diamond

Throughput benchmarking (Section 3.4, Figure 4): The throughput measurements use random model weights to eliminate the influence of potentially variable generation lengths on speed. A normal distribution with 30% variance is applied to prompt and generation lengths. The vLLM version is 0.7.3, customized to support the Phi4-mini-Flash architecture. Hardware: one A100-80GB GPU with no tensor parallelism. Concurrency levels: 1, 2, 4, 8, 16 concurrent requests.


Summary of Design Choices and Their Justifications

  • GMU instead of additional cross-attention layers: Reduces decoding memory I/O from $O(d_{kv}N)$ to $O(d_h)$ per replaced layer. The paper's scaling experiments show this doesn't hurt (and sometimes helps) performance — SambaY achieves lower irreducible loss than Samba+YOCO.
  • Mamba-1 over Mamba-2 in the self-decoder: Mamba-2's scalar-valued forget gates reduce the self-decoder's ability to encode fine-grained positional information, hurting long-context retrieval (Phonebook drops from 78.1% to 40.6%).
  • Samba (Mamba + SWA) over pure Mamba in the self-decoder: Pure Mamba (MambaY) achieves only 12.5% on Phonebook 32K, showing that recency bias alone is insufficient for complex retrieval — the local attention from SWA provides necessary retrieval capability.
  • nGMU with normalization after gating (for linear attention backbones): Preserves the associativity between gate and token mixing, critical for long-context retrieval (Table 9 shows 83.6% vs. 27.3% on Phonebook for GDN-based models with correct vs. incorrect normalization placement).
  • Iso-parametric equations for architecture comparison: Ensures models have comparable parameter counts and KV cache sizes, enabling fair FLOPs-matched and inference-cost-matched comparisons that prior work often overlooked.
  • µP++ over µP: Prevents training instability at large scales by applying zero weight decay to vector-like parameters (Figure 7a shows µP diverging at 600B tokens while µP++ does not).
  • Irreducible loss $C$ as the primary scaling metric: The paper argues that $b$ and $A$ (convergence speed) are misleading for architecture comparison because better initial convergence doesn't guarantee better asymptotic performance. $C$ directly measures the architecture's ultimate potential.
  • Constant batch size scaling (rather than scaling batch size with training tokens): The paper found that batch size scaling "harms the data scaling behavior of the models" (Section E), possibly because batch sizes surpass the critical batch size where larger batches stop improving convergence.
  • FP32 up-casting in the fused loss kernel and mid-training attention dropout: Practical stability measures for the large-scale 3.8B pretraining run that prevented loss divergence without degrading downstream performance.

4. Key Insights and Innovations

Innovation 1: Redefining the Efficiency Target from Pre-Filling to Decoding Memory Bandwidth

The paper's most fundamental conceptual move is relocating the efficiency bottleneck from where the field has been focusing — pre-filling (processing user prompts) — to where the new generation of reasoning models actually hurts: decoding memory I/O during long chain-of-thought generation.

Prior to this work, the decoder-decoder architecture YOCO (Sun et al., 2024) was celebrated for achieving linear pre-filling complexity by caching a single set of KV pairs from the self-decoder and reusing them across all cross-attention layers. That solved a real problem: processing long user prompts efficiently. But the paper identifies a looming crisis that YOCO does nothing to address. During decoding, every single cross-attention layer must still read the full KV cache (size proportional to sequence length × key/value dimension) from memory for every generated token. For short responses, this is negligible. For reasoning models generating 32,000-token traces — which is what the state-of-the-art actually does — this memory I/O dominates the entire inference budget. The foundational premise shift is: the bottleneck isn't the attention computation itself; it's moving the KV cache from memory to the compute units, and that cost scales linearly with generation length for every single layer.

Why this reframing matters intellectually is that it changes what counts as an "efficient architecture." Under the pre-filling-centric view, YOCO is close to optimal — you cache once and pay O(N) pre-filling. Under the decoding-centric view, YOCO still has a hidden O(L × N) bandwidth cost where L is the number of cross-attention layers. The paper's key move is asking: do all those cross-attention layers actually need to read the full KV cache, or can some of them function with a much cheaper memory source? This question wouldn't arise if you were only thinking about pre-filling.

The evidence that this reframing is correct — not just theoretically interesting — comes from Figure 4, where SambaY achieves up to 10× higher decoding throughput on 2K prompts with 32K generation length. This gap grows with generation length precisely because the cross-attention cost grows linearly with N while the GMU cost stays constant. The paper's insight is that the right unit of optimization for the reasoning-model era is memory bandwidth per decoded token, not FLOPs per pre-filled token.

This is a fundamental reframing, not an incremental improvement. It changes the design target for future architectures: instead of asking "how can we pre-fill faster?", it asks "how can we reduce the number of times the KV cache is read during decoding?" The GMU is one answer, but the question — once articulated — opens a design space that YOCO's framing had obscured.

Innovation 2: SSM Kernel Outputs as a New Substrate for Cross-Layer Memory Sharing

The paper introduces a genuinely new form of memory sharing between layers — not KV caches (attention's key-value pairs), not hidden states (recurrent states that are per-token), but the SSM kernel output: the token-mixed representation produced by the structured state-space computation itself.

This is conceptually distinct from all prior work on cross-layer sharing. KV cache sharing (YOCO, Cross-Layer Attention, Multi-Query Attention) operates within the attention framework: you compute key-value projections once and reuse them, but the information being shared is still the attention keys and values — designed for softmax-based retrieval over all positions. The SSM kernel output, in contrast, is the result of a fundamentally different computation: a structured recurrent update with recency bias, where the influence of past tokens on the current representation decays according to learned state-space parameters. It encodes a temporal summary rather than an indexable key-value store.

The paper doesn't just propose sharing this representation — it demonstrates that the representation has specific properties that make it suitable for gated modulation rather than attention-based retrieval. The ablation in Table 5 (Section 4) tests what happens when the GMU gates different memory sources — SSM output (SambaY), attention intermediate state (SambaY-A), MLP intermediate state (SambaY-MLP), and no memory at all (SambaY-AA, which removes cross-attention entirely). The clear hierarchy on Phonebook 32K — 78.1% (SSM) > 64.8% (MLP) > 58.6% (attention) > 46.9% (none) — shows that the SSM kernel output is uniquely effective as a memory source for gated cross-layer sharing, not just that any intermediate representation will do.

What makes this a significant architectural contribution rather than just "we gated some features" is that it identifies a new category of architecturally reusable state. Prior work recognized three kinds: parameters (shared across inputs), KV caches (shared across layers within attention), and recurrent states (shared across time steps within a single layer). The SSM kernel output is a fourth kind: a layer-specific, time-aggregated representation that can be efficiently shared across later layers through lightweight gating, without the O(N) bandwidth cost of KV cache reads. This expands the design vocabulary for efficient architectures in a way that isn't tied to the specific GMU mechanism — future work could share SSM outputs through different modulation schemes, or share MLP outputs through gating (as SambaY-MLP explores), because the category has been named and validated.

The contribution is fundamental in that it opens a design dimension, but the specific mechanism (GMU) is incremental — element-wise gating with a learned projection is a standard building block. The insight is in what to gate, not how to gate.

Innovation 3: Irreducible Loss as the Decisive Metric for Architecture Scaling Comparisons

The paper makes a methodological argument with real teeth: when comparing architectures via scaling laws, the scaling exponent b (how fast loss decreases with compute) and the coefficient A (initial convergence speed) are misleading for predicting long-term performance. The irreducible loss C — what the validation loss approaches with infinite compute — is the metric that matters, and architectures should be ranked by it.

This is important because the field has a known Stated Problem: different architectures converge at different rates, and it's easy to draw wrong conclusions from small-scale experiments. A model that converges faster (higher b or lower A) might look better at 100B tokens but asymptotically plateau above a slower-converging competitor. The paper provides the clearest example of this in its own data (Figure 2b): Transformer++ trained with µP++ shows a large validation loss gap compared to SambaY and Samba+YOCO within the measured range of 100B–600B tokens, yet all three architectures have nearly identical fitted irreducible losses (C ≈ 1.82). Under infinite data, they converge. But under the scaling regime where both model size and data increase proportionally (Figure 2a), SambaY's C = 0.58 is clearly better than Transformer++'s C = 0.64 — a gap that doesn't close.

What makes this an intellectual innovation rather than just "we followed scaling law best practices" is the paper's explicit argument that b and A are confounded by optimization choices — a better learning rate schedule or initialization might improve convergence speed without changing the architecture's ultimate potential. µP++ itself is partly designed to control for this: by establishing a principled hyperparameter transfer scheme across architectures, the paper isolates the architectural contribution to C from the optimization contribution to b and A. The finding that "under µP++, all architectures share the same compute efficiency exponent (b = 0.07)" (Section 3.1) is telling — it means the optimization is well-controlled, and the remaining differences are genuinely architectural.

This is an incremental methodological contribution — the power-law framework and irreducible loss concept predate this paper (Hestness et al., 2017; Hoffmann et al., 2022) — but the application to architecture comparison and the explicit argument for C over b/A is a diagnostic move that the field should adopt. It prevents the common error of declaring an architecture "better" based on faster convergence at small scale when the asymptotic behavior tells a different story.

Innovation 4: Verifying That Long-Context Retrieval Survives Cross-Attention Reduction When the Memory Source Has Recency Bias

This is the paper's most important negative contribution, and it takes the form of a boundary condition rather than a capability claim. The question is: if you replace half the cross-attention layers with GMUs that gate SSM memory, does the model's ability to retrieve information from long contexts degrade? The paper shows that it doesn't — but only when the self-decoder includes local attention (SWA) and the SSM memory has recency bias.

The evidence is in the ablation matrix (Table 5, Section 4). When the self-decoder uses Samba (Mamba + SWA), SambaY achieves 78.1% on Phonebook 32K. When you remove SWA (MambaY), performance collapses to 12.5% — worse than the pure Transformer baseline. When you replace the SSM memory source with attention intermediate states (SambaY-A, 58.6%) or MLP intermediate states (SambaY-MLP, 64.8%), performance degrades substantially. When you remove cross-attention entirely and rely only on GMUs (SambaY-AA, 46.9%), it degrades further but is still competitive with some Transformer variants.

The intellectual contribution is the pattern this reveals, not any single number. The GMU is not a general-purpose replacement for cross-attention — it's specifically effective when gating representations that encode recency-biased temporal structure. SSMs naturally encode this; attention and MLP representations don't. SWA in the self-decoder provides the local retrieval that SSMs alone lack. The combination works because the SSM handles recency-weighted long-range dependencies (which the GMU can later modulate), while SWA handles precise local retrieval, and the remaining cross-attention layers provide global softmax-based retrieval over the full sequence. Each mechanism does what it's good at.

This is a fundamental architectural insight because it establishes that cross-layer memory sharing and local retrieval capability are complementary, not redundant. You can't just replace all attention with SSM sharing and expect retrieval to work — the paper's SambaY-AA result shows that. But you also don't need full cross-attention in every layer — SambaY outperforms Samba+YOCO (which has cross-attention in every cross-decoder layer) on long-context tasks while using half the cross-attention layers. The right design is interleaved: some layers do global retrieval (cross-attention), some do local retrieval (SWA in the self-decoder), and some do cheap cross-layer modulation of temporal summaries (GMU on SSM memory). This tripartite division of retrieval labor is a new architectural principle that emerges from the ablation results, not from the stated contribution of the GMU itself.

The finding is empirically justified (Table 5, Table 1) and conceptually significant because it reframes the design question from "can we replace attention?" to "what retrieval capabilities does each layer actually need, and what's the cheapest mechanism that provides them?" This is a diagnostic contribution — it tells you when to use which mechanism — and it's the kind of insight that guides future architecture design more than any single performance number.

5. Experimental Analysis

Evaluation Methodology

  • Dataset. The primary scaling experiments use SlimPajama (Soboleva et al., 2023), a cleaned and deduplicated 627B-token version of RedPajama, for pre-training with 4K sequence length (Section 3.1). Long-context experiments use ProLong-64K (Gao et al., 2024) with 32K sequence length and 40B training tokens (Section 3.2). The large-scale 3.8B model is trained on 5T tokens from the Phi4-mini data corpus (Section 3.3). Downstream evaluation uses standard benchmarks: Wikitext (perplexity), LAMBADA, ARC-Easy/Challenge, HellaSwag, PIQA, WinoGrande (short-context); Phonebook and RULER (long-context retrieval); MMLU, MMLU-Pro, Arena Hard, GSM8K, Qasper, SummScreenFD, BigCodeBench, MBPP (comprehensive); and Math500, AIME24/25, GPQA Diamond (reasoning).

  • Base model(s). Scaling experiments use a family of models scaled from ~154M to ~3.4B parameters with depth d ∈ {8, 12, 16, 20, 24} and corresponding training tokens from 12.5B to 342B under the Chinchilla-inspired over-training regime (Section 3.1, Table 7). The large-scale model is Phi4-mini-Flash (3.8B parameters), a SambaY+DA architecture trained with Standard Parameterization rather than µP++ due to resource constraints (Section 3.3). The reasoning model, Phi4-mini-Flash-Reasoning, is distilled from this checkpoint via SFT and DPO following the Phi4-mini-Reasoning recipe (Section 3.4). All hybrid architectures use Samba as the self-decoder backbone — interleaving Mamba-1 layers with Sliding Window Attention and ending with a single full-attention layer.

  • Metrics. For scaling experiments, the primary metric is validation loss on the SlimPajama dataset, fitted to power-law curves L(D) = A · D^(-b) + C from which the irreducible loss C is extracted as the key architecture comparison metric (Section 3.1). For downstream evaluation, metrics include perplexity (Wiki), accuracy (LAMBADA, WinoGrande, ARC-Easy, PIQA), character-normalized accuracy (ARC-Challenge, HellaSwag), and retrieval accuracy (Phonebook, RULER NIAH tasks). For reasoning benchmarks, pass@1 accuracy is reported, averaged over 64 samples for AIME24/25 and 8 samples for Math500 and GPQA Diamond (Section G). Throughput is measured in tokens per second under vLLM 0.7.3 (Section 3.4).

  • Baselines. The paper compares against five architectural baselines at 1B-parameter scale (Section 3, preamble): Transformer++ (SwiGLU MLP, RoPE with base frequency 10,000, no SWA), TransformerLS (Transformer++ with SWA on all layers except every 4th which uses full attention), SWA+YOCO (original YOCO with SWA as self-decoder), Samba+YOCO (Samba as self-decoder with full cross-attention cross-decoder — the naive combination), and SambaY+DA (SambaY enhanced with Differential Attention). For the large-scale reasoning comparison, the primary baseline is Phi4-mini-Reasoning (3.8B), which uses a Transformer++ architecture and includes a final RL training stage. Additional reference baselines include open-source distilled reasoning models: DeepSeek-R1-Distill-Qwen-1.5B, DeepSeek-R1-Distill-Qwen-7B, DeepSeek-R1-Distill-Llama-8B, Bespoke-Stratos-7B, and OpenThinker-7B (Table 4).

  • Generation budget / compute accounting. For scaling comparisons, training compute is measured in total FLOPs consumed during pre-training, accounting for both model size and number of training tokens. For inference efficiency, the generation budget is measured in sequence length (prompt length and generation length) with memory I/O costs calculated as O(d_kv · N) for cross-attention layers versus O(d_h) for GMU layers (Section 2). Throughput measurements use random model weights to eliminate the influence of variable generation lengths on speed, with a normal distribution of 30% variance applied to prompt and generation lengths (Section G). Training speed is reported in Million Tokens Per Second (MTPS) on 64 A100-80GB GPUs (Table 2, Table 5). The iso-parametric equation method ensures parameter counts are matched across architectures for fair comparison (Section 3.1, Section C).

  • Cross-validation / statistical protocol. The scaling law fitting uses least squares with the Levenberg-Marquardt algorithm, reporting R^2 goodness-of-fit for all curves (all ≥ 0.999, indicating near-perfect power-law behavior) (Section 3.1). For the Phonebook benchmark, error bars are reported as standard deviations across multiple runs (Figure 3). Downstream reasoning evaluation follows Phi4-mini-Reasoning's protocol: temperature 0.6, top-p 0.95, maximum 32,768 tokens, with multiple samples per problem averaged for robust pass@1 estimation. The paper does not describe cross-validation for downstream task evaluation (metrics are reported on standard test sets without further splitting). For the large-scale pretraining, no hyperparameter sweep is performed due to resource constraints; a generic optimization setup based on Transformer++ is adopted (Section J).

Main Quantitative Results

Scaling Law Comparisons (FLOPs and Data Scaling)

The central architectural result is that SambaY achieves the lowest irreducible loss among all tested architectures under FLOPs scaling with µP++ (Section 3.1, Figure 2a). The fitted power-law parameters are:

ArchitectureAbC (irreducible loss)
SambaY0.070.58
Samba+YOCO0.070.59
SWA+YOCO0.070.60
TransformerLS0.070.62
Transformer++0.070.64

The paper emphasizes (Section 3.1): "under µP++, all architectures share the same compute efficiency exponent (b = 0.07), indicating that the hybrid architectures explored did not yield improvements in models' learning efficiency with respect to compute." This means the advantage is purely in the asymptotic limit — SambaY converges to a lower loss given infinite compute — rather than in learning speed. The fitted curves have R^2 ≥ 0.999 for all architectures, confirming the reliability of the irreducible loss estimates.

A critical comparison: µP++ versus Standard Parameterization (SP). Under SP, the Transformer++ baseline shows higher irreducible loss than under µP++ (visible as the upper curve in Figure 2a for Transformer++ SP), demonstrating that the hyperparameter transfer scheme itself contributes significantly to scaling behavior — the architecture comparison conclusions depend on having principled hyperparameter scaling.

Data scaling results (Figure 2b) tell a complementary story. At fixed model size (1B parameters, d = 16) with training tokens scaling from 100B to 600B:

  • Transformer++ shows a large validation loss gap compared to SambaY and Samba+YOCO within the measured range (100B–600B tokens).
  • However, all three architectures have nearly identical fitted irreducible losses (C ≈ 1.82 for all), with the paper noting: "its [Transformer++'s] fitted irreducible loss... is nearly identical to those of the hybrid models, suggesting that with an infinite amount of data, Transformer++ can eventually catch up—albeit with slower convergence."
  • This convergence occurs only in the fixed-model-size regime. Under FLOPs scaling where model size and data scale proportionally (Figure 2a), Transformer++'s irreducible loss (C = 0.64) remains clearly worse than SambaY's (C = 0.58).

The paper explains this discrepancy (Section 3.1): "this is because we use Mamba-1 as our SSM which falls into the same complexity class of TC0 as Transformers" — when model size is fixed, the expressiveness is comparable, and more data benefits all architectures similarly. When both model size and data scale, the architectural efficiency differences compound.

Long-Context Retrieval: Sliding Window Size Dependence

The Phonebook benchmark (32K evaluation length, containing 1,850 name–number pairs) reveals a non-monotonic relationship between sliding window size and retrieval accuracy that differs sharply across architectures (Section 3.2, Figure 3):

  • SambaY+DA peaks at 512 with ~96% accuracy, then degrades slightly at 1024 and 2048.
  • SambaY peaks at 256 with ~92% accuracy, showing similar decline at larger windows.
  • Samba+YOCO requires 1024 to achieve ~83% but shows monotonic improvement with window size.
  • SWA+YOCO performs poorly across all window sizes, never exceeding ~25%.
  • TransformerLS peaks at 256 with ~60% and plateaus.

The paper's interpretation: "larger SWA sizes do not consistently provide better results" and using smaller windows "could enable the model to focus on local patterns more easily and mitigate issues like attention sinks." The SWA+YOCO result specifically "could indicate that pure attention-based models require more substantial data for long-context training."

Full RULER benchmark results (Table 1) with the best SWA size per architecture confirm SambaY variants' advantage:

ModelSWAMK-1MK-2MK-3MQMVS-1S-2S-3Avg.
SambaY+DA51264.627.60.212.819.999.886.469.647.6
SambaY25654.627.80.412.719.483.281.263.842.9
Samba+YOCO102449.028.02.612.818.3100.063.223.637.2
Transformer++36.43.80.027.924.194.866.031.035.5
TransformerLS25642.86.00.029.827.591.849.623.433.9

SambaY+DA wins on 5 of 8 subtasks and achieves the highest average. Notably, SambaY variants show stronger multi-key retrieval (MK-1, MK-2) and single-needle retrieval (S-1 through S-3), while Transformer-based models dominate multi-query and multi-value tasks. The paper attributes this pattern to the combination of SSM recency bias and SWA local retrieval complementing the global cross-attention, with Differential Attention further enhancing single- and multi-key performance.

Short-context downstream results (Table 2) show that the long-context advantages don't come at the cost of standard benchmarks:

ModelWiki ppl ↓LMB acc ↑Avg acc ↑Speed (mtps) ↑
SambaY+DA16.5949.6852.170.91
SambaY17.8350.4052.171.11
Samba+YOCO16.7350.5352.000.99
Transformer++19.7545.4548.350.89

SambaY and SambaY+DA share the highest average accuracy (52.17%) across seven benchmarks, while SambaY achieves 25% faster training throughput than Samba+YOCO (1.11 vs. 0.99 MTPS) due to its smaller SWA size (256 vs. 1024). Transformer++ trails significantly on both metrics.

Large-Scale Pre-Training: Phi4-mini-Flash vs. Phi4-mini

The 3.8B-parameter Phi4-mini-Flash, pre-trained on 5T tokens, demonstrates that SambaY+DA scales to production-relevant sizes (Section 3.3, Table 3):

BenchmarkMetricPhi4-miniPhi4-mini-Flash
MMLU5-shot67.371.9
MMLU-Pro0-shot, CoT52.854.7
Arena HardWin Rate32.834.9
GSM8K0-shot, CoT88.689.5
QasperF140.440.2
SummScreenFDROUGE-L16.017.0
BigCodeBenchpass@143.044.5
MBPPpass@165.369.8

Phi4-mini-Flash wins on 7 of 8 benchmarks, with notable improvements on knowledge-intensive tasks (MMLU: +4.6 points, MBPP: +4.5 points). The exception is Qasper (F1 score), where it underperforms by 0.2 points. The paper notes this was achieved "while maintaining substantially higher computational efficiency during inference" — a claim quantified in the reasoning section.

Training stability is explicitly discussed: two interventions were required (Section D, Figure 6):

  1. FP32 up-casting in the fused cross-entropy loss kernel to prevent gradient norm divergence.
  2. Attention dropout (0.05) added at 310K steps (self-decoder) and 520K steps (cross-decoder) when loss divergence was imminent. Both interventions stabilized training through 5T tokens without harming MMLU performance.

Efficient Reasoning with Long Generation

The reasoning model comparison (Section 3.4, Table 4) is the paper's most applied result:

ModelSizeAIME24AIME25Math500GPQA Diamond
Phi4-mini-Flash-Reasoning3.8B52.2933.5992.4545.08
Phi4-mini-Reasoning3.8B48.1331.7791.2044.51
DeepSeek-R1-Distill-Qwen-7B7B53.7035.9493.0347.85
DeepSeek-R1-Distill-Llama-8B8B43.9627.3487.4845.83
DeepSeek-R1-Distill-Qwen-1.5B1.5B29.5820.7884.5037.69

Phi4-mini-Flash-Reasoning outperforms the stronger Phi4-mini-Reasoning baseline (which includes RL) across all four reasoning benchmarks, with the largest gaps on AIME24 (+4.16 points) and AIME25 (+1.82 points). It achieves competitive or superior performance to larger distilled models: surpassing DeepSeek-R1-Distill-Llama-8B on all benchmarks despite being less than half the size, and approaching DeepSeek-R1-Distill-Qwen-7B (within 1.41 points on AIME24, 0.58 on Math500).

The paper emphasizes (Section 3.4): "while employing a substantially more efficient architecture, achieving up to 10× higher throughput in long-generation scenarios and 4.9× speedup in long-context processing." The throughput advantage is measured in Figure 4:

  • Prompt: 32000, Generation: 500 (Figure 4a): SambaY achieves 5–10× higher throughput than YOCO/Samba+YOCO across concurrency levels 1–16. Phi4-mini-Flash (with sub-optimal DA implementation) approaches SambaY's throughput at low concurrency but falls behind at higher concurrency due to DA's four-pass FlashAttention.
  • Prompt: 2000, Generation: 32000 (Figure 4b): The advantage is even more dramatic — SambaY maintains 3–5× higher throughput than baselines across all concurrency levels. The gap widens with generation length because the cross-attention cost grows linearly with N while GMU cost stays constant.
  • Latency scaling (Figure 9): Phi4-mini-Flash achieves faster generation than Phi4-mini starting from generation lengths as small as 4K, with the gap widening monotonically to 32K.

The paper notes (Section 3.4) that the DA implementation is "sub-optimal in calculating unnecessary full attention operations with the following MLP layer for the user prompt tokens during the prefill stage, leaving additional opportunities for optimizing the efficiency results."

Ablation Studies and Robustness Checks

All ablation experiments use 1B-parameter models trained on ProLong-64K with 40B tokens and 32K sequence length, with SWA size fixed at 128 and µP++ depth d = 16 (Section 4, Table 5).

Alternative self-decoder backbones (MambaY, MambaY-2, GDNY): Replacing Samba's Mamba+SWA self-decoder with different SSM/RNN variants reveals a clear hierarchy:

  • MambaY (pure Mamba-1, no SWA): Phonebook accuracy collapses to 12.5% (vs. 78.1% for SambaY), demonstrating that recency bias alone is insufficient — local retrieval via SWA is necessary for the self-decoder to support the cross-decoder in complex retrieval.
  • MambaY-2 (pure Mamba-2): Phonebook improves to 50.8% — better than MambaY but still severely degraded. Average short-context accuracy drops to 51.01% (vs. 52.16% for SambaY), with Wiki perplexity worsening to 18.63 (vs. 16.89).
  • GDNY (pure Gated DeltaNet): Phonebook achieves 89.8% — better than SambaY's 78.1%. Average short-context accuracy is 51.75%, competitive with SambaY. Training speed is faster (1.22 vs. 1.10 MTPS).
  • S-GDNY (GDN+SWA interleaved): Phonebook drops slightly to 83.6% (vs. 89.8% for pure GDNY) but training speed improves (1.34 vs. 1.22 MTPS) and average accuracy improves to 51.97% — the second-highest in the table.

The takeaway: more sophisticated linear attention mechanisms (GDN with delta update rules) can match or exceed Mamba+SWA for long-context retrieval when used as the GMU's memory source, and interleaving with SWA trades a small amount of retrieval accuracy for better training speed and short-context performance.

Alternative SSM variants (SambaY-2): Replacing Mamba-1 with Mamba-2 in the Samba framework (keeping SWA) causes Phonebook to drop from 78.1% to 40.6% — a dramatic degradation despite the SWA layers being identical. The paper attributes this to "Mamba-2's coarse, scalar-valued forget gates" reducing the self-decoder's capacity to encode fine-grained positional information in the SSM memory that the GMU later relies on. Training speed improves (1.43 vs. 1.10 MTPS), but this gain is undermined by the retrieval collapse.

Memory source alternatives (SambaY-A, SambaY-AA, SambaY-MLP): Testing whether GMU works with non-SSM memory sources:

  • SambaY-A (gates attention intermediate representations instead of SSM output): Phonebook drops to 58.6% (vs. 78.1%). Short-context accuracy is 52.26%, competitive with SambaY. This validates that GMU works with attention states but less effectively for retrieval — attention states lack the recency bias that makes SSM memory suitable for gated modulation.
  • SambaY-AA (removes cross-attention entirely, uses only GMU on attention states): Phonebook drops further to 46.9% — the worst among GMU variants. Short-context accuracy remains competitive at 52.06%. This establishes a lower bound: completely removing cross-attention is viable for short contexts but catastrophic for retrieval.
  • SambaY-MLP (gates MLP intermediate representations): Phonebook at 64.8% — better than SambaY-A but worse than SambaY. Short-context accuracy achieves the highest in the table at 52.65%, suggesting MLP gating provides useful local feature refinement even though it lacks the temporal structure needed for retrieval.

The paper's hierarchy (Section 4): "SambaY > SambaY-MLP > SambaY-A > SambaY-AA" on Phonebook, with the reasoning that "gating attention/MLP representations performs worse than the original SambaY on Phonebook because they lack the recency bias that SSMs naturally provide, which is beneficial for encoding contiguous local information."

Normalization placement and nGMU necessity (Section H, Table 9): For linear attention variants (Mamba-2, GDN), the choice between normalizing before vs. after the output gating — and whether to use nGMU vs. basic GMU — is decisive:

  • SambaY-2 with correct placement (after gating, nGMU): Phonebook 40.6%.
  • SambaY-2 with incorrect placement ("NB + GMU" — normalization before gating, basic GMU): Phonebook 21.9%, despite improved short-context accuracy (51.83% vs. 51.00%).
  • S-GDNY with correct placement (GDN-A + nGMU): Phonebook 83.6%.
  • S-GDNY with incorrect placement (GDN + GMU): Phonebook 27.3%.
  • GDNY with correct placement (GDN-A + nGMU): Phonebook 89.8%.
  • GDNY with incorrect placement (GDN + GMU): Phonebook 54.7%.

The pattern is consistent: incorrect normalization reduces Phonebook accuracy by 19–56 points while leaving short-context benchmarks largely unaffected (Wiki perplexity changes by < 3%, zero-shot commonsense variations < 3%). This confirms the paper's theoretical argument about associativity preservation being critical specifically for long-range retrieval.

µP++ variant comparisons (Section E, Figure 7): The hyperparameter scaling law itself was ablated:

  • µP++ vs. µP: µP diverges with NaN losses at 600B tokens, while µP++ remains stable — validating the zero-weight-decay modification.
  • µP++ vs. µP++ with batch scaling: Batch scaling (batch size increases with training tokens) shows worse learning efficiency and irreducible loss.
  • µP++ with LeCun init vs. µP++ with Normal init (variance scaled by 1/d): LeCun initialization outperforms, suggesting "it is better to adjust the initialization multipliers based on each matrix's dimension... rather than a global factor related to model width."
  • µP++ vs. µP++ with LR scaling and independent weight decay: The latter shows worse learning efficiency (smaller b).
  • µP++ with linear LR schedule vs. µP++ with WSD schedule (cosine decay): Linear schedule outperforms.
  • Tied vs. untied embeddings: Untied achieves lower validation loss at 100B tokens but comparable irreducible loss — suggesting extra embedding parameters "primarily accelerate training convergence without improving the final model performance if a sufficient amount of data is given."

Differential Attention contribution (Table 1, Table 2): Adding DA to SambaY improves multi-key retrieval (MK-1: 64.6% vs. 54.6%) and single-needle retrieval (S-2: 86.4% vs. 81.2%, S-3: 69.6% vs. 63.8%) at the cost of larger optimal SWA size (512 vs. 256) and reduced training throughput (0.91 vs. 1.11 MTPS). Multi-query and multi-value tasks are slightly worse but within error. Short-context average accuracy is identical (52.17%).

Long-context extrapolation with NoPE (Section F, Table 8): Models trained with 32K context length are evaluated zero-shot at 64K and 128K on Phonebook. Key findings:

  • SambaY with NoPE extrapolates to 64K with 96.1% accuracy (vs. 92.2% at 32K — improvement from longer context), then collapses at 128K.
  • SambaY+DA shows gradual decline: 96.1% at 32K → 84.4% at 64K → 5.5% at 128K, suggesting graceful degradation.
  • RoPE-based models (Transformer++, TransformerLS) collapse at 64K (0.0% and 17.2% respectively), showing catastrophic failure beyond training length.
  • Samba+YOCO (RoPE-based) shows intermediate extrapolation: 82.8% → 68.0% → 20.3%.

The paper leaves explanation of why NoPE enables limited extrapolation as "an interesting future work."

Training data and methodology effects (Section F, Figure 8): Ablating the training recipe shows that:

  • ProLong-64K provides a "notable performance boost across all architectures compared to SlimPajama" for long-context training.
  • SSM-based models benefit more from ProLong-64K than Transformers, suggesting they "can learn to switch contexts between different data samples within the packed sequences more easily."
  • Variable-length training (used in main experiments) generally improves results over fixed-length packing for all architectures, but the effect is strongest for pure attention models that are "sensitive to sliding window size."

ReST^EM for revision models (not applicable — this is the Decoder-Hybrid-Decoder paper): This paper does not use revision models or ReST^EM; the ablation focuses on architecture variants rather than training methodology for self-improvement loops. The paper does include a large-scale training stability ablation (Figure 6) showing that FP32 up-casting and attention dropout successfully mitigate loss divergence in the 5T-token Phi4-mini-Flash run.

Critical Assessment

Claim 1: "SambaY significantly enhances decoding efficiency... while delivering up to 10× higher decoding throughput on 2K-length prompts with 32K generation length"

What the experiments actually demonstrate: Figure 4b shows that SambaY achieves 3–5× higher throughput than YOCO/Samba+YOCO baselines under vLLM 0.7.3 on one A100-80GB GPU. The 10× figure in the abstract appears to reference a specific concurrency level (likely 1 or 2 concurrent clients where the throughput ratio is highest). The paper does not provide a breakdown of which concurrency level achieves 10×, and the figure shows the ratio varies significantly with load. Under high concurrency (16 clients), the advantage compresses to ~3×. The claim is supported directionally but the 10× figure should be understood as a best-case number at low concurrency, not an average or typical throughput gain.

The sub-optimal Differential Attention implementation in Phi4-mini-Flash (four-pass FlashAttention rather than optimized custom kernel) means the Phi4-mini-Flash-Reasoning throughput numbers are pessimistic. However, the paper compares Phi4-mini-Flash-Reasoning's accuracy against baselines while showing SambaY's throughput — these are different architectures (SambaY vs. SambaY+DA) and the direct accuracy-efficiency tradeoff for the same architecture isn't isolated. The paper would be strengthened by showing throughput numbers for the exact Phi4-mini-Flash-Reasoning architecture rather than the DA-free SambaY.

The memory I/O cost analysis (Section 2) provides a clean theoretical argument: replacing half the cross-attention layers reduces per-layer memory read from O(d_kv · N) to O(d_h), with the efficiency gain proportional to N / (d_h/d_kv) which is ~400× at N=32K for typical dimensions. However, this is a per-layer analysis; the end-to-end throughput depends on other factors (MLP computation, self-decoder cost, memory bandwidth contention) that are not analytically modeled. The 3–5× measured throughput gain is substantially less than the 400× per-layer theoretical reduction, confirming that cross-attention memory I/O is not the only bottleneck but is a significant one.

What's missing: A breakdown of where the remaining decoding time goes (MLP, self-decoder, cross-attention layers that weren't replaced) would clarify how close the architecture is to the theoretical efficiency limit. The paper also doesn't benchmark against an optimized YOCO implementation — the vLLM baseline may not fully exploit YOCO's potential for KV cache sharing.

Claim 2: "SambaY exhibits a significantly lower irreducible loss compared to a strong YOCO baseline, indicating superior performance scalability under large-scale compute regimes"

What the experiments actually demonstrate: Figure 2a shows SambaY's fitted irreducible loss C = 0.58 vs. Samba+YOCO's C = 0.59 vs. Transformer++'s C = 0.64. The difference between SambaY and Samba+YOCO is 0.01 — while consistently better across the fitted curve, this is a small absolute gap. The paper describes this as "significantly lower," but the statistical significance is not quantified (no confidence intervals on C). With R^2 ≥ 0.999 on five data points per curve (d = 8, 12, 16, 20, 24), the fitting error is likely small, but five points per architecture is at the lower bound of what's needed for reliable power-law extrapolation. The paper would be strengthened by reporting confidence intervals on C or by running additional intermediate scales to verify the extrapolation.

The more robust finding is SambaY vs. Transformer++ (C = 0.58 vs. 0.64) — a gap of 0.06 that is clearly outside any reasonable fitting error. The SambaY vs. Samba+YOCO comparison is more nuanced: SambaY is directionally better, but the 0.01 difference in C could plausibly be within estimation noise. The paper's emphasis on C (rather than A or b) is well-motivated, but the actual C differences between the top hybrid architectures are small enough that one shouldn't overinterpret the ranking.

The µP++ framework itself is a contribution, but it's validated only on Transformer++ — the paper doesn't verify that µP++ is optimal for SambaY or Samba+YOCO. The hybrid architectures inherit µP++ settings from Transformer++ optimization, and it's possible that architecture-specific hyperparameter tuning would change the relative C values. The paper acknowledges this (Section J): "we do not perform an exhaustive hyperparameter search for each architecture. Instead, we adopt a generic optimization setup based on Transformer++."

Claim 3: "Phi4-mini-Flash-Reasoning achieves significantly better performance than Phi4-mini-Reasoning on reasoning tasks... without any reinforcement learning"

What the experiments actually demonstrate: Table 4 shows consistent gains: +4.16 on AIME24, +1.82 on AIME25, +1.25 on Math500, +0.57 on GPQA Diamond. These are real improvements, but "significantly" is not statistically characterized — no confidence intervals or standard deviations are reported for these benchmarks. The paper evaluates with 64 samples for AIME and 8 for Math500/GPQA, averaging results, but the variance across runs isn't reported. On benchmarks where absolute scores are high (Math500: 92.5% vs. 91.2%), a 1.25-point difference with 500 test problems could be within sampling noise. On AIME24 (52.3% vs. 48.1% with 30 test problems and 64 samples each), the difference is more convincing but still lacks formal statistics.

The claim that this is achieved "without any reinforcement learning" is true but requires context: Phi4-mini-Reasoning's RL stage provides additional training that Phi4-mini-Flash-Reasoning doesn't get, yet Phi4-mini-Slash-Reasoning outperforms anyway. This is the stronger version of the claim — that the architectural efficiency alone compensates for missing RL. However, this conflates architecture and training recipe: the distillation SFT and DPO data may differ between the two models in ways not fully controlled. The paper states it follows the "same multi-stage distillation data following Phi4-mini-Reasoning" but the architecture change means the effective training dynamics differ.

The paper does not provide an ablation showing that Transformer++ trained with the same SFT+DPO data (but without RL) matches Phi4-mini-Reasoning — this would isolate whether the RL removal or the architecture change drives the improvement. The comparison is effectively between (Phi4-mini-Flash + SFT + DPO) and (Phi4-mini + SFT + DPO + RL), which confounds two variables.

Missing baselines: The paper compares against several open-source distilled models (DeepSeek-R1-Distill-Qwen-7B, Llama-8B, etc.) but doesn't include a SambaY or YOCO architecture at the same 3.8B scale trained with the same reasoning distillation pipeline. The closest comparison is the throughput numbers (Figure 4), which use random weights and don't measure reasoning quality. A matched comparison — same data, same distillation recipe, different architectures — would be the cleanest demonstration that the architecture (not the Phi4 data corpus or distillation recipe) drives the gains.

Claim 4: "GMU enables efficient memory sharing across layers"

What the experiments actually demonstrate: The ablation study (Table 5) provides strong evidence that GMU with SSM memory works well — SambaY achieves 78.1% on Phonebook and 52.17% average accuracy while reducing cross-attention I/O. The comparison against SambaY-AA (which removes all cross-attention and uses only GMU, dropping to 46.9% on Phonebook) establishes that some cross-attention is necessary. The comparison against Samba+YOCO (which has cross-attention in every cross-decoder layer, achieving 37.2% average on RULER vs. SambaY's 42.9%) suggests that interleaving GMU with cross-attention actually improves long-context performance over pure cross-attention — a counterintuitive finding.

However, the paper doesn't fully isolate the GMU's contribution to decoding efficiency from its effect on model quality. The throughput measurements (Figure 4) are with random weights — they measure the architectural speed potential, but don't show whether a trained SambaY model actually achieves quality- and speed-matched performance. A quality-throughput Pareto frontier comparing SambaY and Samba+YOCO at equivalent perplexity would be the gold standard; the paper instead measures quality (Table 1, Table 2) and speed (Figure 4) separately on different configurations.

The GMU mechanism itself is relatively simple (element-wise gating with learned projections), and the paper's main contribution is showing that it works in this context rather than proposing a novel mathematical operation. The ablation study convincingly shows that SSM memory is the right source (vs. attention or MLP states), but the paper doesn't explore why element-wise gating is sufficient — the theoretical analysis (Section 2) shows that gating lifts the token-mixing matrix into a third-order tensor with channel-specific reweighting, but this is a description of what the GMU computes, not an explanation of why full attention's softmax over all positions isn't needed for the replaced layers.

Genuine Weaknesses in Experimental Design

  1. Single benchmark for long-context scaling (Phonebook): The SWA size sweep (Figure 3) and self-decoder ablation (Table 5) both use Phonebook 32K as the primary long-context metric. Phonebook is a specific task (key-value retrieval with 1,850 pairs) and may not generalize to other long-context capabilities like multi-hop reasoning or summarization. The RULER results (Table 1) provide a broader picture but are only reported at the single best SWA size per architecture, not swept.

  2. No human evaluation or qualitative analysis of reasoning quality: The reasoning benchmarks (AIME, Math500, GPQA) measure final-answer accuracy but don't assess the quality of the chain-of-thought itself. Since the paper's efficiency argument hinges on long CoT generation, understanding whether GMU-replaced layers affect reasoning depth, backtracking behavior, or error recovery would be valuable. The case studies (Section G, Examples 1 and 2) show qualitatively interesting reasoning patterns but are anecdotal.

  3. vLLM implementation is sub-optimal for both baselines and proposed model: The paper acknowledges that the Differential Attention implementation uses a naive four-pass FlashAttention and that unnecessary full attention operations are computed during prefill. This means the reported throughput numbers are not at the performance ceiling for either the baseline or the proposed architecture. The relative comparison (SambaY vs. YOCO) should still be valid since both suffer from the same vLLM limitations, but the absolute throughput numbers should not be taken as indicative of a fully optimized deployment.

  4. Difficulty estimation and adaptive allocation studied elsewhere: The paper focuses on architectural efficiency rather than compute-optimal allocation, but the reasoning results would benefit from showing whether the efficiency gains translate to better compute-optimal scaling — for instance, whether Phi4-mini-Flash-Reasoning with a smaller generation budget matches Phi4-mini-Reasoning's accuracy at a larger budget, analogous to efficiency claims in other papers. The paper shows throughput gains (3-5× at equal generation length) but not accuracy-matched compute reductions.

  5. No analysis of GMU gating patterns or interpretability: The paper provides a mathematical formalism for how gating modulates token mixing, but never visualizes or analyzes what the learned gates actually do. Do different GMU layers learn different gating patterns? Do gates specialize by position (early vs. late tokens in the sequence)? This interpretability analysis could strengthen the architectural argument by showing that GMU layers learn meaningful, layer-specific modulation rather than degenerate solutions (e.g., all gates saturating at the same value).

  6. McQueen's note: The training of the 3.8B model experienced "severe loss divergence" requiring two interventions (FP32 up-casting and mid-training attention dropout). While these were successfully mitigated, they indicate the architecture's training dynamics are less stable than Transformer++ at scale, and the paper's µP++ framework (which prevents divergence for Transformer++) wasn't applied to the 3.8B run due to resource constraints. Whether µP++ would have stabilized SambaY+DA at scale without the need for ad-hoc interventions remains an open question.

Experiments That Would Have Strengthened the Paper

  • Ablation of the GMU-to-cross-attention ratio: The paper replaces "approximately half" the cross-attention layers with GMU but doesn't sweep this ratio. Is 50% optimal, or would 25% or 75% be better? This is a fundamental design parameter that directly affects the throughput-quality tradeoff. The SambaY-AA result (0% cross-attention, 46.9% Phonebook) and Samba+YOCO (100% cross-attention, 37.2% RULER average) provide bounds, but the intermediate points are missing.

  • Phi4-mini-Flash-Reasoning with RL: The paper's headline comparison removes RL from the pipeline, making Phi4-mini-Flash-Reasoning strictly less trained than Phi4-mini-Reasoning. Adding the same RL stage to Phi4-mini-Flash-Reasoning would test whether the architectural advantage compounds with RL or whether the gains diminish when both models are fully optimized.

  • Scaling the reasoning model to larger sizes: The 3.8B scale is the paper's largest. Since the efficiency advantage grows with generation length (the gap between O(d_kv·N) and O(d_h) widens with N), larger models that generate even longer CoTs would benefit more from GMU. A 7B or 13B SambaY variant on reasoning tasks would test whether the throughput advantage scales with model size.

  • Direct comparison of SambaY vs. Samba+YOCO at 3.8B scale on reasoning: The 3.8B reasoning comparison is against Phi4-mini (a Transformer++ architecture), not against a Samba+YOCO reasoning model of the same size. Since YOCO is the direct predecessor, showing that SambaY outperforms it on reasoning quality (not just throughput) at the same scale would close the loop on the architecture's claim to superiority.

  • Measurement of actual memory bandwidth utilization: The paper's efficiency argument is based on theoretical memory I/O counts, but actual GPU performance depends on memory bandwidth saturation, kernel fusion opportunities, and caching effects. Profiling the memory bandwidth utilization during decoding for both GMU and cross-attention layers would validate the theoretical analysis and identify remaining bottlenecks.

6. Limitations and Trade-offs

Hard Problems Remain Outside the Reach of Test-Time Compute

The assumption or constraint. The paper's efficiency argument — that replacing cross-attention with GMU preserves long-context retrieval capability — only holds when the base model is capable of performing the retrieval task in the first place. This is most visible in the difficulty-dependent results on the RULER benchmark. For multi-key retrieval at depth 3 (MK-3) and multi-query (MQ) and multi-value (MV) tasks, even the best SambaY+DA configuration achieves near-zero or single-digit accuracy (MK-3: 0.2%, MQ: 12.8%, MV: 19.9% in Table 1). These are not failures of the GMU specifically — Transformer++ achieves 0.0% on MK-3 and similar numbers on MQ/MV — but rather a fundamental capability bound: the 1B-parameter models lack the capacity to solve these tasks regardless of architecture.

The consequence. The paper's central claim is that architectural efficiency enables long-chain-of-thought reasoning without sacrificing retrieval capability. But the efficiency gains are most valuable precisely for the hardest problems — where generation lengths are longest and the bandwidth bottleneck is most severe. If those hardest problems are also the ones where the base model's retrieval capability fails regardless, the practical value of the efficiency gains is skewed toward easier problems that require less generation anyway. A practitioner deciding whether to adopt SambaY for a reasoning application would need to know: does my problem distribution look more like Phonebook/MK-1 (where SambaY excels) or MK-3/MQ (where no architecture helps)? The paper does not provide this characterization.

What evidence exists in the paper. The RULER results (Table 1) provide the clearest evidence: on MK-1, SambaY+DA achieves 64.6% (strong); on MK-3, it drops to 0.2% (failure); on MQ, it reaches only 12.8%. The Phonebook benchmark (Figure 3) shows SambaY achieving ~92%, but Phonebook is a specific key-value retrieval task — the paper doesn't demonstrate that this retrieval capability transfers to the multi-hop or multi-query retrieval that complex reasoning often requires. The reasoning benchmarks (Table 4, AIME/Math500/GPQA) measure final-answer accuracy but don't isolate whether retrieval failures during the chain of thought contribute to errors.

Mitigation status. The paper does not address this limitation directly. It does not characterize which types of retrieval failures cause reasoning errors, nor does it propose mechanisms to improve retrieval on the hardest subtasks. The finding that Gated DeltaNet self-decoders improve Phonebook accuracy (GDNY: 89.8% in Table 5) suggests that better SSM mechanisms can push the capability bound upward, but the paper treats this as an ablation rather than a path toward solving the hard-retrieval problem. Section J acknowledges "our architecture still includes a full-attention layer, which leads to linear per-token computation complexity during decoding" but frames this as a limitation of decoding efficiency, not of retrieval capability.


The 10× Throughput Figure Is a Best-Case Condition, Not a Typical Gain

The assumption or constraint. The abstract claims "up to 10× higher decoding throughput on 2K-length prompts with 32K generation length." This figure comes from Figure 4b, where SambaY achieves its highest throughput ratio relative to YOCO/Samba+YOCO at low concurrency levels (likely 1 or 2 concurrent clients). At higher concurrency (8–16 clients), the advantage compresses to approximately 3×. The 10× figure is therefore a single-point measurement at a specific operating condition, not an average or typical throughput improvement.

The consequence. A practitioner deploying a reasoning model in production typically runs at higher concurrency to maximize GPU utilization and minimize cost per query. If the 10× throughput gain only materializes at low concurrency — where absolute throughput is already low and GPU utilization is poor — the headline number substantially overstates the practical benefit. Conversely, at the higher concurrency levels where production systems operate, the 3× gain is still significant but represents a very different cost-reduction profile. The paper doesn't provide guidance on how to choose concurrency levels or how the throughput ratio changes with batch size, leaving deployment engineers to extrapolate from a limited set of measurements.

What evidence exists in the paper. Figure 4a and 4b show throughput curves across concurrency levels {1, 2, 4, 8, 16}. In Figure 4b (prompt: 2000, generation: 32000), SambaY throughput at concurrency 1 is approximately 180 tokens/sec vs. YOCO at approximately 18 tokens/sec — roughly a 10× ratio. At concurrency 4, SambaY is at approximately 400 tokens/sec vs. YOCO at approximately 90 tokens/sec — roughly 4.4×. At concurrency 16, SambaY is at approximately 550 tokens/sec vs. YOCO at approximately 180 tokens/sec — roughly 3×. The ratio clearly diminishes with load. The paper does not discuss this compression, why it occurs (likely memory bandwidth saturation or compute-unit contention at higher loads), or at what concurrency level the ratio stabilizes.

Mitigation status. The paper does not address this limitation. It does not provide throughput numbers beyond concurrency 16, does not analyze why the ratio compresses, and does not recommend operating concurrency levels for deployment. The "up to 10×" qualifier in the abstract is technically accurate but omits the context that the 10× is achieved only at the lowest load point. A more informative reporting would include the throughput ratio at the concurrency level that maximizes absolute throughput (which appears to be concurrency 8–16 in Figure 4b), where the gain is ~3×.


Difficulty Estimation Cost for the Architecture Is Unaccounted for in the 5T-Token Training Run

The assumption or constraint. The 3.8B Phi4-mini-Flash model uses Standard Parameterization rather than µP++ for training, which the paper acknowledges was due to "resource constraints at the time of scaling study" (Section 3.3). The µP++ framework that underpins all the scaling law comparisons and the irreducible loss advantage (C = 0.58 vs. 0.64) was developed on 1B-scale models and validated to 3.4B parameters on SlimPajama. The 3.8B model was trained on an entirely different data corpus (Phi4-mini data), with a different aspect ratio (α = 80 vs. α ≈ 124), a different head dimension (64 vs. 128), and different training hyperparameters — including two ad-hoc stability interventions (FP32 up-casting in the loss kernel and mid-training attention dropout) that were not part of the µP++ framework.

The consequence. The paper's central theoretical claim — that SambaY has a lower irreducible loss than Transformer++ — is established only under µP++ scaling on SlimPajama. The large-scale model that actually demonstrates the architecture's practical value (Phi4-mini-Flash-Reasoning) was trained under different conditions. This creates uncertainty about whether the 3.8B model's strong performance (outperforming Phi4-mini-Reasoning on reasoning benchmarks) is attributable to the SambaY+DA architecture, the µP++ hyperparameter scheme (which wasn't used), the Phi4 data corpus (which is higher-quality than SlimPajama and likely benefits architectures differently), or the stability interventions. The paper cannot cleanly attribute the 3.8B gains to the architectural innovation, because the training recipe changed simultaneously with the architecture.

What evidence exists in the paper. The scaling experiments (Section 3.1, Figure 2a) establish that SambaY with µP++ achieves C = 0.58 on SlimPajama. The large-scale results (Section 3.3, Table 3) show Phi4-mini-Flash (SambaY+DA, SP, Phi4 data) outperforming Phi4-mini (Transformer++, SP, Phi4 data). But there is no intermediate experiment showing that (a) SambaY with SP on SlimPajama retains an advantage over Transformer++ with SP, or (b) Transformer++ with µP++ on Phi4 data would close the gap with SambaY. The paper explicitly acknowledges "The optimization setup here is by no means optimal, as the primary goal of this experiment is to evaluate the viability of our architecture at larger scales." This is candid, but it means the large-scale results are a viability demonstration, not a controlled comparison.

The training instability itself (Section D, Figure 6) is revealing: the model experienced "severe loss divergence" requiring two interventions (FP32 up-casting and attention dropout) that were not needed for Transformer++ at the same scale. This suggests the architecture's training dynamics are less stable at scale, and the µ P++ framework — which prevented divergence for Transformer++ in the scaling experiments (Figure 7a) — might not be sufficient for SambaY+DA, or the SP training recipe (which was used instead) requires architecture-specific tuning that the paper did not perform.

Mitigation status. The paper acknowledges the optimization suboptimality but does not remediate it. Section J states: "we do not perform an exhaustive hyperparameter search for each architecture. Instead, we adopt a generic optimization setup based on Transformer++ for learning rate, initializer range, weight decay, warm-up schedule, batch size, AdamW betas and epsilon, and other parameters. It is likely that aggressive tuning of these optimization settings could yield improved results." This implies the 3.8B results are a lower bound on what SambaY+DA could achieve with architecture-specific optimization, but it also means the comparison against the well-tuned Phi4-mini baseline may be unfair — architecture-specific tuning could benefit Phi4-mini as well.


The Method Has Only Been Validated on a Single Model Family and Two Data Corpora

The assumption or constraint. All experiments use Mamba-1, Mamba-2, or Gated DeltaNet as the SSM backbone, with Samba (Mamba + SWA + full attention) as the self-decoder architecture. The scaling experiments use SlimPajama; the large-scale experiment uses the proprietary Phi4-mini data corpus. There is no evidence that GMU works with other SSM variants (e.g., RWKV, xLSTM, RetNet, H3), other backbone architectures (e.g., pure transformer self-decoders, Griffin-style recurrent blocks), or other data distributions (code, multilingual text, multi-modal).

The consequence. A practitioner with a different SSM implementation or a different data domain cannot assume the GMU's effectiveness transfers. The ablation study (Table 5) actually shows significant sensitivity to the SSM choice: replacing Mamba-1 with Mamba-2 in the self-decoder drops Phonebook accuracy from 78.1% to 40.6% (SambaY vs. SambaY-2). This is a dramatic degradation from what should be a minor change — Mamba-2 is a refinement of Mamba-1 with the same broad design philosophy. If such a small architectural change causes a 37.5-point accuracy drop, it's plausible that other SSM variants with different state update rules, normalization schemes, or gating mechanisms would interact differently — and possibly unfavorably — with the GMU. The normalization placement ablation (Table 9) reinforces this: using the wrong normalization order (GDN + GMU instead of GDN-A + nGMU) drops Phonebook from 83.6% to 27.3%. The GMU is fragile to the specifics of how the SSM memory is produced.

The data limitation is less severe but still relevant. The 3.8B model uses the Phi4-mini data corpus, which is a curated, high-quality mix. The paper doesn't demonstrate that the efficiency gains persist on noisier or more diverse data (e.g., CommonCrawl, code repositories), where the model might need to rely more heavily on cross-attention layers for precise retrieval of facts and patterns that the SSM memory might not capture well.

What evidence exists in the paper. The SSM sensitivity is directly visible in Table 5 and Table 9. The paper doesn't test RWKV, xLSTM, RetNet, or any SSM outside the Mamba/GDN lineage. The data limitation is implicit: all downstream evaluation uses the benchmarks listed in Section D (Wiki, LAMBADA, ARC, HellaSwag, PIQA, WinoGrande, Phonebook, RULER, MMLU, reasoning tasks), which are standard English-language benchmarks. There is no multilingual, code, or multi-modal evaluation.

Mitigation status. The paper does not claim generalizability beyond the tested architectures and data. Section J acknowledges that RL remains under-explored and that the optimization setup is generic, but does not acknowledge the SSM-family or data-domain limitations as explicit scope constraints. The results for Gated DeltaNet (GDNY, S-GDNY in Table 5) show that the GMU can work with non-Mamba SSMs, but only when the normalization placement is carefully adjusted — suggesting the mechanism is general but requires SSM-specific tuning that isn't fully characterized.


The Self-Decoder's Full-Attention Layer Remains a Linear-Complexity Bottleneck During Decoding

The assumption or constraint. The paper reduces decoding memory I/O for half of the cross-decoder layers, but the architecture still includes a full-attention layer in the self-decoder, multiple cross-attention layers in the cross-decoder, and sliding window attention layers that all read KV caches of size growing with sequence length. The theoretical analysis (Section 2) focuses on the per-layer cost reduction for GMU-replaced layers, but the end-to-end decoding cost is the sum of costs across all layers. The remaining cross-attention and full-attention layers still pay O(d_kv · N) per layer per token, meaning the overall decoding cost remains asymptotically linear in sequence length — just with a smaller constant factor (roughly half the cross-attention I/O).

The consequence. As generation lengths continue to grow (models generating 100K+ token traces are already in development), the remaining cross-attention layers will become the dominant cost, and the GMU's benefit will plateau. The 50% layer replacement achieves a one-time factor-of-~2 reduction in cross-decoder I/O; it does not change the asymptotic scaling from linear to constant. For a practitioner deploying reasoning models with ever-longer chains of thought, SambaY buys time but doesn't solve the fundamental problem. Eventually, the remaining attention layers will bottleneck inference, and further efficiency gains will require either replacing more cross-attention layers (which the SambaY-AA ablation suggests degrades retrieval quality) or developing fundamentally attention-free architectures for generation.

What evidence exists in the paper. The paper is transparent about this. Section 2 states: "our approach only requires caching an additional SSM kernel output state... alongside the KV cache from the last full-attention layer during pre-filling" (my emphasis), acknowledging that KV caches still exist and are still read. Section J explicitly frames this as future work: "this underscores a future research direction on designing models for extremely long sequence generation that can maintain constant decoding complexity while effectively leveraging long-context memory." The SambaY-AA ablation (Table 5, all cross-attention removed) shows that Phonebook accuracy drops to 46.9% — confirming that simply removing all cross-attention is not viable with current techniques. The throughput measurements (Figure 4) are for a specific generation length (32K); extrapolating to longer generations would show the efficiency ratio diminishing as cross-attention I/O grows.

Mitigation status. The paper treats this as a fundamental design choice rather than a solved problem. Section J suggests future work on dynamic sparse attention for the remaining attention layers. The paper doesn't explore intermediate ratios of GMU-to-cross-attention (e.g., 75% GMU, 25% cross-attention) that might push the tradeoff further, nor does it investigate whether techniques like InfiniGen or Layer-Condensed KV Cache (cited in Section I) could reduce the cost of the remaining attention layers. The d_h/d_kv ratio analysis establishes that the GMU's benefit grows with sequence length, but only up to the point where the remaining attention cost dominates — at which point the overall speedup plateaus.


No Characterization of the Optimal GMU Placement or Ratio Across Layers

The assumption or constraint. The paper fixes the GMU-to-cross-attention ratio at "roughly 50%" and interleaves GMU layers with cross-attention layers (Section 2, Figure 1). The paper does not ablate this ratio, does not test whether GMU layers should be concentrated at the beginning or end of the cross-decoder, and does not investigate whether different layers benefit differentially from GMU versus cross-attention. The choice of 50% appears to be a design decision without empirical justification.

The consequence. This is the single largest unexplored design dimension in the architecture. If some cross-decoder layers rely heavily on precise token-level retrieval (benefiting from full cross-attention) while others primarily perform local feature refinement or semantic integration (where GMU suffices), then the optimal GMU ratio could be much higher than 50% (more efficiency gain) or lower (better quality). The SambaY-AA ablation (Table 5) establishes that 100% GMU (0% cross-attention) degrades Phonebook from 78.1% to 46.9% — a significant drop but still competitive with some Transformer baselines. A 75% GMU configuration (retaining only 25% cross-attention) might achieve most of the efficiency gain with minimal quality loss, but the paper never tests it.

A practitioner deploying SambaY would want to know: (a) what is the quality-throughput Pareto frontier as a function of the GMU ratio? (b) Is 50% near the Pareto-optimal point, or is it an arbitrary choice? (c) Should GMU layers be placed immediately after the SSM memory source (early cross-decoder) or closer to the output (late cross-decoder)? These questions are fundamental to deploying the architecture optimally but are completely unexplored.

What evidence exists in the paper. The paper tests exactly two points on this spectrum: 100% cross-attention (Samba+YOCO, which achieves 37.2% average on RULER at its optimal SWA size of 1024) and 50% GMU (SambaY, which achieves 42.9% average at SWA size 256). The 0% cross-attention point is approximated by SambaY-AA (46.9% on Phonebook, but SambaY-AA uses attention memory rather than SSM memory, so it's not a clean ablation of the GMU ratio alone). There is no sweep of the ratio itself. The GMU placement is fixed as interleaved — the paper doesn't test whether placing all GMU layers at the beginning or end of the cross-decoder changes performance.

Mitigation status. The paper does not acknowledge this as a limitation. Section J focuses on optimization hyperparameters and RL as future work, not on architectural hyperparameters like the GMU ratio or placement. Given that the ratio directly controls the quality-efficiency tradeoff and is likely the most impactful design choice after the decision to use GMU at all, this omission is significant. A simple ablation training 3–4 models with different GMU ratios on SlimPajama at 1B scale would have characterized the tradeoff and provided actionable guidance for practitioners.

7. Implications and Future Directions

How This Work Changes the Landscape

This paper introduces a fundamentally new category of reusable architectural state — the SSM kernel output as a cross-layer memory — and demonstrates that gated access to this state can replace roughly half of the cross-attention layers in a decoder-decoder framework without degrading (and sometimes improving) model quality. This is not an incremental optimization of attention; it reframes the design question from "how can we make attention cheaper?" to "which layers actually need full token-level retrieval, and what is the cheapest mechanism that provides sufficient cross-layer information flow for the rest?"

The magnitude is closer to a reframing with immediate practical consequences than a paradigm shift. The core insight — that SSM outputs encode recency-biased temporal structure that can be productively shared across layers through learned gating — opens a design dimension that prior work on KV cache sharing (YOCO, Cross-Layer Attention, Multi-Query Attention) had not recognized, because those works operated entirely within the attention framework. The paper shows that once you separate the memory source (what information is shared) from the access mechanism (how later layers retrieve it), you can mix and match: attention-style KV caches for layers that need global softmax retrieval, SSM kernel outputs with gating for layers that need cheap, channel-wise modulation. This tripartite design vocabulary — global attention, local sliding window, and gated temporal memory — is the paper's most lasting architectural contribution.

The work reconciles a latent tension in the efficient inference literature. Prior to this, the decoder-decoder architecture (YOCO) represented the state of the art for pre-filling efficiency, but its decoding cost — every cross-attention layer reading the full KV cache for every generated token — was an unexamined liability for the emerging paradigm of long-chain-of-thought reasoning. The paper's empirical demonstration that this liability is both real (Figure 4 shows 3–5× throughput gaps at 32K generation) and addressable (GMU layers reduce per-layer memory I/O from $O(d_{kv}N)$ to $O(d_h)$) changes the evaluation criteria for future architectures: decoding memory bandwidth per generated token becomes a first-class design target, alongside pre-filling FLOPs. Work that optimizes only pre-filling (as much prior work on SSMs and linear attention did) implicitly assumes short generation lengths; this paper makes that assumption explicit and shows it no longer holds.

The landscape shift has concrete directional implications. Research on pure SSM architectures for decoding becomes more attractive, because the paper shows that SSM representations have properties (recency bias, temporal structure) that make them effective as gated memories — a role that pure attention mechanisms cannot fill as cheaply. Conversely, research on ever-more-sophisticated KV cache compression (quantization, sparsification, eviction policies) becomes relatively less urgent for the cross-decoder, because the paper shows that for roughly half the layers, you can avoid reading the KV cache entirely rather than compressing it. The bottleneck moves from "how do we store the KV cache compactly?" to "how do we generate useful SSM memories that later layers can productively gate?" — a fundamentally different research question.

The paper also establishes a methodological standard for architecture comparison that the field would benefit from adopting. The combination of iso-parametric equations (matching parameter counts and cache sizes across architectures with different internal expansion ratios), µP++ (principled hyperparameter scaling for depth and width with training stability), and irreducible loss $C$ as the primary scaling metric (rather than convergence speed) provides a template for answering "which architecture scales better?" that is more rigorous than the common practice of matching parameter counts while ignoring cache sizes or optimization sensitivity. The finding that all architectures share the same compute efficiency exponent $b = 0.07$ under µP++ (Section 3.1) is methodologically significant: it means the optimization is well-controlled, and the remaining differences in irreducible loss $C$ are genuinely architectural. This should raise the bar for future architecture papers that claim scaling advantages based on small-scale experiments with architecture-specific hyperparameter tuning.

Finally, the paper's reasoning results without RL (Table 4: Phi4-mini-Flash-Reasoning outperforms Phi4-mini-Reasoning on all four benchmarks despite lacking the RL stage) suggest that architectural efficiency is not merely a deployment concern — it can translate into better model quality at a fixed training budget, because the freed-up inference compute enables more extensive chain-of-thought generation during training data synthesis or distillation. This connects architectural design to the training pipeline in a way that pure efficiency papers often overlook: a more efficient architecture doesn't just serve faster; it can be trained to reason better because longer chains of thought become economically viable.

Follow-Up Research This Work Enables

Sweep the GMU-to-cross-attention ratio to map the quality-throughput Pareto frontier. The paper fixes the ratio at approximately 50% without empirical justification. A systematic sweep — training models at 1B scale on SlimPajama with 0%, 25%, 50%, 75%, and 100% GMU layers in the cross-decoder (keeping total layer count fixed) — would characterize the fundamental tradeoff. The key metric would be Phonebook accuracy and RULER average at each ratio, plotted against measured decoding throughput at 32K generation length. The paper's SambaY (50%: 78.1% Phonebook) and SambaY-AA (0% cross-attention with attention memory: 46.9%) provide bounds, but a model with 75% GMU and SSM memory (rather than the attention memory used in SambaY-AA) might achieve most of the efficiency gain with minimal retrieval degradation. If the frontier is convex — a small quality sacrifice for a large efficiency gain at high GMU ratios — this would be immediately actionable for practitioners who can tune the ratio to their latency requirements.

Test whether GMU placement (early versus late in the cross-decoder) matters, and whether different layers learn functionally distinct gating patterns. The current architecture interleaves GMU and cross-attention layers uniformly. An ablation comparing three configurations at 1B scale — (a) GMU layers concentrated in the first half of the cross-decoder, (b) GMU layers interleaved uniformly (current), (c) GMU layers concentrated in the second half — would test whether the gated memory is more useful for early feature extraction or late semantic integration. Complementing this with an analysis of learned gate values: do early GMU layers show different gating patterns (e.g., broadly permissive gates that let through most SSM memory) compared to late GMU layers (e.g., sparse, token-position-dependent gates that extract specific information)? Visualizing gate activation patterns across sequence positions and layers would connect the mathematical formalism (the channel-specific reweighting of the token-mixing matrix from Section 2) to functional behavior, and might reveal whether some GMU layers learn degenerate solutions (e.g., nearly uniform gating) that could be pruned.

Combine GMU with the remaining cross-attention compression techniques to push decoding cost toward constant complexity. The paper acknowledges (Section J) that the remaining cross-attention and full-attention layers still incur linear decoding cost. Techniques like Cross-Layer Attention (share KV caches across adjacent cross-decoder layers), InfiniGen (selective KV prefetching), or Layer-Condensed KV Cache (compute KV pairs for only a subset of layers) could be applied to the remaining cross-attention layers. A concrete experiment: take the SambaY architecture at 1B scale, apply CLA to the cross-attention layers (sharing KV caches between pairs of adjacent cross-attention layers, halving the number of distinct cache reads), and compare against the baseline SambaY on the RULER benchmark. If retrieval quality is preserved, the effective cross-attention read count drops to 25% of the original YOCO design (50% replaced by GMU, and half of the remaining 50% sharing caches via CLA), moving closer to constant decoding complexity. The paper's citation of these techniques (Section I) suggests the combination is natural but untested.

Extend GMU to other SSM families (RWKV, xLSTM, RetNet) and characterize the conditions under which gated memory sharing succeeds versus fails. The ablation study (Table 5, Table 9) reveals significant sensitivity to SSM choice and normalization placement: Mamba-2 (SambaY-2) achieves only 40.6% on Phonebook versus Mamba-1's 78.1%; incorrect normalization placement drops GDN-based models from 83.6% to 27.3%. This fragility suggests that the GMU's effectiveness depends on specific properties of the SSM memory — likely the presence of fine-grained, per-channel positional information that scalar-valued forget gates (Mamba-2) may smooth away. A systematic study training GMU-equipped models with RWKV (which uses exponential decay with learned time constants), xLSTM (which uses scalar memory cells with gating), and RetNet (which uses a multi-scale retention mechanism) at 1B scale on ProLong-64K would establish which SSM properties are necessary for effective gated sharing. The prediction: SSMs with per-channel state updates (like Mamba-1's selective state-space parameters) produce richer memories than those with scalar or global state updates. A negative result — showing that some SSM families produce memories that GMU cannot effectively gate — would define the boundary conditions for the technique and guide future SSM design toward GMU-compatible formulations.

Train Phi4-mini-Flash-Reasoning with RL and measure whether architectural efficiency compounds with reinforcement learning for reasoning. The paper's headline reasoning result (Table 4) shows that SambaY+DA with SFT+DPO outperforms Transformer++ with SFT+DPO+RL. This leaves open the question: does adding RL to Phi4-mini-Flash-Reasoning provide further gains, or does the architectural advantage diminish when both models are fully optimized? A concrete experiment: run the same RL stage used for Phi4-mini-Reasoning on Phi4-mini-Flash, then compare AIME24/25 and Math500 accuracy. If the architectural advantage persists or grows (because the more efficient architecture enables longer RL rollouts or more extensive exploration within the same compute budget), this would establish that GMU-based architectures are not just faster at inference but enable better training outcomes. If the gap closes (Phi4-mini with RL catches up to or surpasses Phi4-mini-Flash with RL), it would suggest that RL provides a substitute for some of the representational benefits of the hybrid architecture, and the efficiency advantage is primarily a deployment concern rather than a training one.

Develop a lightweight difficulty predictor for reasoning tasks that routes queries to different generation budgets based on the architecture's efficiency profile. The reasoning results (Table 4) show that Phi4-mini-Flash-Reasoning achieves strong accuracy while generating long chains of thought (up to 32K tokens). But not all queries require 32K tokens — easy problems might be solved in 2K tokens, while only the hardest require the full budget. Combined with the architecture's efficiency advantage (which grows with generation length, per Figure 4), an adaptive allocation policy that estimates problem difficulty from the prompt and adjusts the generation budget could achieve substantial additional throughput gains. A concrete experiment: train a lightweight difficulty classifier (perhaps a linear probe on Phi4-mini-Flash's final hidden state) on the MATH training set to predict the number of tokens needed for a correct solution, then use it to cap generation length during inference. Measure whether accuracy-matched throughput improves over uniform 32K generation. The paper's finding that SambaY's throughput advantage is largest at long generation lengths (Figure 4b) implies that the efficiency gains from adaptive allocation would be larger for SambaY than for Transformer++ — you save more by shortening generation on the architecture where generation is proportionally cheaper, because you can reserve the long-generation budget for the hardest queries where the architecture's efficiency advantage is most needed. This connects the paper's architectural contribution to the compute-optimal allocation paradigm that has gained traction in the reasoning literature.

Practical Applications and Downstream Use Cases

Batch inference for reasoning benchmark evaluation and synthetic data generation. Organizations that evaluate large numbers of reasoning problems (e.g., running AIME or MATH benchmarks across multiple model checkpoints, or generating synthetic reasoning traces for distillation) face a direct cost-throughput tradeoff. The paper's Figure 4b shows that at concurrency 4 with 32K generation, SambaY achieves approximately 400 tokens/second versus YOCO's 90 tokens/second — a 4.4× throughput improvement. For a team generating 1 million reasoning traces at an average of 16K tokens each, this translates to approximately 44 GPU-hours on SambaY versus 194 GPU-hours on YOCO, or roughly a 1,500versus1,500 versus 6,600 cost differential at typical cloud GPU pricing. The practical benefit is amplified because the efficiency gain grows with generation length (the $N \gg d_h/d_{kv}$ condition is more satisfied), meaning the architecture is most cost-effective precisely for the hardest problems that produce the longest chains of thought — exactly the problems that are most expensive to evaluate at scale.

Deploying small reasoning models in latency-constrained settings (on-device, real-time tutoring). The paper's 3.8B-parameter Phi4-mini-Flash-Reasoning achieves AIME24 accuracy (52.3%) that approaches 7B-8B distilled models (DeepSeek-R1-Distill-Qwen-7B: 53.7%) while delivering 4.9× speedup in long-context processing. For applications where latency matters — interactive math tutoring where students wait for step-by-step solutions, or on-device reasoning where a model runs locally on a laptop or phone — the combination of small model size (3.8B parameters, fitting in ~8GB of memory) and high throughput on long generations makes SambaY-based architectures uniquely suitable. A concrete deployment scenario: an educational app that generates detailed solutions to student-submitted math problems. With Phi4-mini-Reasoning, a 32K-token solution might take 10 seconds to generate on-device; with Phi4-mini-Flash-Reasoning, it takes 2 seconds, moving from "annoying delay" to "near-interactive." The paper's demonstration that this speedup comes with improved accuracy (not a tradeoff) makes the adoption case particularly strong.

Self-improvement pipelines where a model generates its own training data through long-chain-of-thought reasoning. The STaR and ReST-style self-improvement loops (where a model generates reasoning traces, filters for correct answers, and fine-tunes on the successful traces) are bottlenecked by the cost of generating millions of long CoT trajectories. If each training iteration requires sampling 500K solutions at 16K tokens average, the 4.4× throughput advantage of SambaY (at concurrency 4) reduces the generation phase from 2.5 days to 0.6 days on a single GPU — or allows 4.4× more trajectories in the same time budget, potentially improving the quality and diversity of the self-improvement data. The paper's reasoning results (Table 4) already demonstrate that the architecture supports high-quality reasoning; combining this with an automated self-improvement loop would test whether the efficiency advantage accelerates the entire iterative training pipeline, not just final inference. The key metric would be: how much faster does a SambaY-based model reach a target AIME accuracy through iterative self-improvement compared to a Transformer-based model with the same GPU budget?

When to Prefer This Method

The paper positions SambaY against specific named alternatives (YOCO, Samba+YOCO, Transformer++) and provides the empirical evidence to construct a conditional decision rule:

  • Prefer SambaY over YOCO/Samba+YOCO when generation lengths routinely exceed $d_h/d_{kv}$ (typically ~80–128 tokens for the architectures studied). At 32K generation, the throughput advantage is 3–5× (Figure 4b). At short generation lengths (<1K tokens), the advantage shrinks because the constant-factor overhead of GMU projections becomes comparable to cross-attention memory reads, and YOCO's architectural simplicity may be preferable. The paper's scaling results (Figure 2a, $C = 0.58$ for SambaY vs. $C = 0.59$ for Samba+YOCO) suggest SambaY is at least not worse on quality, so the decision hinges primarily on the expected generation length distribution.

  • Prefer SambaY+DA (Differential Attention) over plain SambaY when the task requires precise multi-key retrieval from long contexts (high MK-1 scores matter) and training throughput is not the primary constraint. Table 1 shows SambaY+DA achieves 64.6% on MK-1 versus SambaY's 54.6%, and 99.8% versus 83.2% on single-needle retrieval (S-1). The cost is a larger optimal SWA size (512 vs. 256) and ~18% slower training throughput (0.91 vs. 1.11 MTPS, Table 2).

  • Prefer SambaY+DA over Transformer++ when both long-context retrieval and reasoning over long generations matter, and the deployment uses vLLM or a framework where the sub-optimal DA implementation (four-pass FlashAttention) is acceptable. The paper shows SambaY+DA achieves 47.6% average on RULER versus 35.5% for Transformer++ (Table 1), and Phi4-mini-Flash-Reasoning outperforms Phi4-mini-Reasoning on all four reasoning benchmarks (Table 4). The tradeoff is training stability: the 3.8B SambaY+DA model required two ad-hoc interventions (FP32 up-casting in the loss kernel, mid-training attention dropout) to prevent divergence (Section D), suggesting Transformer++ may be preferable when training stability at scale is the overriding concern.

  • Consider Gated DeltaNet (GDNY/S-GDNY) as the self-decoder instead of Samba when long-context retrieval is the dominant requirement and training throughput is less critical. Table 5 shows GDNY achieves 89.8% on Phonebook 32K versus SambaY's 78.1%, and S-GDNY (GDN+SWA interleaved) achieves 83.6% with better training speed (1.34 vs. 1.10 MTPS). The caveat is that GDN-based architectures require careful normalization placement (nGMU with normalization after gating; Table 9 shows a 56.3-point Phonebook drop with incorrect placement), increasing implementation complexity.