Decoder-Hybrid-Decoder Architecture for Efficient Reasoning with Long Generation

We would like to thank the reviewer for the positive review of our work and the constructive feedback on our manuscript. In the following, we address the remaining concerns to hopefully motivate a clear acceptance score.

More baselines: Thanks for the suggestions. Due to the limited time and computation resources, we are unable to draw the expensive scaling curves (which needs more than 12K H100 hours for the FLOPs scaling experiment) during the rebuttal time. However, we do conduct the long context retrieval experiments as follows.

We pre-train 1.0B parameter SWA+YOCO models with µP++ and d = 16 on ProLong-64k dataset with 40B tokens. The SWA+YOCO architecture uses all SWA layers for self-decoder and all cross-attention layers for cross-decoder. We align the parameters through solving the iso-parameteric equation to obtain the aspect ratio $\alpha = 130$. A 640K RoPE base is also adopted. The following table shows the PB-32K evaluation results of six models trained with various sliding window sizes. We can see that SWA+YOCO achieves the best result with a SWA size of 128, while its performance is consistently worse than Samba+YOCO or SambaY.

We also measure the NIAH tasks performance in the RULER benchmarks as follows. We can see that SWA+YOCO with its best 128 SWA size provides the worst average accuracy across all models.

We further measure the downstream short-context performance on language modeling and common-sense reasoning tasks in zero-shot, and we can see that SWA+YOCO is comparable to TransformerLS but has worse performance than Samba+YOCO or SambaY.

GMU Performance: Thank you for the detailed analysis. While performance on some short-context benchmarks may appear comparable, this perspective overlooks the two crucial advantages that the Gated Memory Unit (GMU) provides to our SambaY architecture: (1) The primary motivation for the GMU is to reduce the severe memory I/O bottleneck present in standard cross-attention during decoding. This is where SambaY excels. As shown clearly in Figure 4, SambaY delivers substantially higher throughput than Samba+YOCO. (2) On tasks specifically designed to test long-context retrieval, SambaY demonstrates a clear performance advantage over Samba+YOCO. In the challenging Needle-in-a-Haystack (NIAH) benchmark from RULER (Table 1), SambaY achieves an average accuracy of 42.9, significantly outperforming Samba+YOCO's score of 37.2.

Limitations: Thanks for the suggestion. Due to the page limitations, we will add a reference at the beginning of section 3 to guide the reader for the limitation section in the appendix.

(to be continued)

Dear reviewer JjzD,

As we approach the end of the discussion period, we would greatly appreciate your input on the paper. We hope our responses and additional results address your concerns and welcome any further questions or suggestions.

Thank you for your time and effort.

We would like to thank the reviewer for the very positive review of our work and the interesting questions on our manuscript. We address these questions in the following paragraphs.

Ablation Studies with Differential Attention: Since our primary focus is on ablating the architectural design of the self-decoder and cross-decoder components in the SambaY architecture, we did not include Differential Attention (DA) in the main ablation results. DA is a simple, orthogonal extension to SambaY and does not directly impact the architectural comparisons we aim to study. However, for completeness, we do include results with DA in Figure 3 and Table 2 for ablating the design of Phi4-mini-Flash.

Differences between SP and µP++: Thanks for the suggestion! We summarize the key differences between μP, μP++, and Standard Parameterization (SP) in the following table. LR mult. denotes the per-parameter multiplier applied on top of the global learning rate (η), Res. mult. is the multiplier applied to the output of residual branches, and WD denotes the weight decay. For μP++, η ∝ 1/√d and zero weight decay is also applied to other scalar or vector-like parameters such as RMSNorm weights. In this work, σ = 10⁻⁴ for untied embedding and σ = 0.02 for tied embedding; in both cases, τ = 0.02 and β = 1. “fan_in” refers to the input dimension of weight matrices.

Q1: How do you explain the drop in performance on PB-32k when moving from SambaY to SambaY-2? It seems that without SWA Mamba2 is better, while with SWA Mamba1 is better. This seems to contradict the explanation in L. 263. ( Could it be that the reason is a less effective memory access to matrix states with GMU? (see next question) )

A1: Thank you for the insightful question! We believe the performance difference is best understood from a capacity trade-off perspective. Mamba-1 provides strong recency bias through fine-grained memory decay, enabling effective local representation and positional encoding. Mamba-2, on the other hand, uses a larger matrix-valued recurrent state, which enhances memorization capacity but weakens recency bias. When combined with SWA, which already offers strong local memorization and retrieval, the positional representation power of Mamba-1 becomes more valuable. In this setting, replacing Mamba-1 with Mamba-2 (as in SambaY-2) leads to redundant memorization and a net loss in positional expressiveness, hurting performance on tasks like PB-32k that benefit from precise local memorization. Conversely, in the absence of SWA, the increased memorization capacity of Mamba-2 in self-decoder becomes advantageous, enhancing the long-context retrieval of complex information (such as phone numbers) for the cross-decoder, where local memory support from SWA is missing. We will include this discussion in the next version of our paper.

(to be continued)

We would like to thank the reviewer for the very positive review of our work. We address the remaining concerns in the following paragraphs.

Downstream Evaluation: Thanks for your suggestions! Our primary objective in the ablation studies was to test the efficiency of different memory-gating strategies for long-context retrieval. To this end, we deliberately chose the Phonebook (PB-32K) benchmark. This is not a generic language modeling task but is a synthetic, highly-controlled test designed to provide a precise signal on a model's ability to perform multi-key retrieval. As shown in Table 5, the results on this task are not noisy but they reveal a clear and significant performance hierarchy (SambaY > SambaY-MLP > SambaY-A > SambaY-AA), which directly validates our core design choices.

Paper Formatting: Thanks for pointing it out! We have double checked the style guide, which does not exclude the usage of the current formatting since we maintain the internal consistency.

We would like to thank the reviewer for the constructive feedback on our manuscript. In the following, we address the remaining concerns to hopefully encourage a positive evaluation.

Novelty and coherence of the contributions: We think that our work actually presents a cohesive narrative centered on our main contribution: the Gated Memory Unit (GMU), a novel and versatile architectural component that enables our efficient SambaY architecture. The GMU is our primary contribution, designed as a simple yet effective mechanism for sharing representations across layers. While inspired by the previous works on KV cache sharing, our work introduces a completely new way to achieve memory sharing for State Space Models (SSMs). As we show in L271-283, the GMU is a generalizable concept, capable of gating memory from attention and MLP layers, making it a distinct contribution beyond the specific context of YOCO (L75-78). We do not claim novelty for the µP++ scheme or Differential Attention (DA). Their roles are distinct and supportive: µP++ is a methodological tool for ensuring a rigorous and fair comparison between architectures. This was necessary to scientifically validate our claim that SambaY exhibits superior scalability with a lower irreducible loss (L42-43, L154-156). DA is an extension that demonstrates the power and versatility of our SambaY framework. By integrating it, we create a state-of-the-art model that proves SambaY's practical value and efficiency gains on challenging benchmarks (L54-55, L234-237). Therefore, our contributions are unified: we introduce a novel architecture (SambaY) enabled by a new component (GMU), use rigorous methodology (µP++) to prove its effectiveness, and showcase its full potential with a powerful extension (DA).

More backgrounds: Thanks for pointing this out. We will add the following backgrounds in Appendix for more self-contained presentations:

Our work builds upon the YOCO (You Only Cache Once), an efficient decoder-decoder architecture. It consists of a self-decoder, which is the first half of the layers that produces KV caches in linear timem and a cross-decoder, which is the second half of the layers that re-uses the KV caches from the first half. For an input sequence of hidden states $H_{mem} \in \mathbb{R}^{N \times d}$ from a single final attention layer in the self-decoder, it pre-computes and caches a single Key-Value (KV) pair:

where $K_c, V_c \in \mathbb{R}^{N \times d_{kv}}$ are the cached matrices and $W_K, W_V$ are weight matrices. Subsequently, every cross-attention layer $l$ in the cross-decoder reuses this single KV cache. Given the input hidden state $H_{cross}^{(l)}$ to that layer, the cross-attention output is calculated by generating a new Query, $Q_{cross}^{(l)}$, and attending to the shared cache:

This approach avoids doing full attention computation at the prefill stage, reducing the prefill time complexity to linear.

More long-context evaluation: Our selection of NIAH tasks was intentionally focused to evaluate the core hypothesis of Section 3.2 that SambaY can do efficient long-context retrieval. We also add additional evaluation on the remaining RULER tasks as follows. As we can see, all the models show near-zero Variable tracking (VT) performance because of the small scales of training (1B parameters with 40B tokens). We can also see that SambaY still has a better average accuracy than baselines.

Minor & Typos:

1. Why ...Transformer++?

This is because it improves over the original Transformer with RoPE and GLU, following the notation in [1].

2. The symbols in lines ... understand their meaning.

N_{attn, mamba, mlp, gmu} is the total number of parameters for attention/Mamba/MLP/GMU layers respectively.

(to be continued)