Title: Decoder-Hybrid-Decoder Architecture for Efficient Reasoning with Long Generation
Venue: NeurIPS 2025
Accepted: yes
Keywords: Hybrid Neural Architectures;Long Sequence Models;State Space Models;Linear Recurrent Neural Networks
Abstract: Recent advances in language modeling have demonstrated the effectiveness of State Space Models (SSMs) for efficient sequence modeling. While hybrid architectures such as Samba and the decoder-decoder architecture, YOCO, have shown promising performance gains over Transformers, prior works have not investigated the efficiency potential of representation sharing between SSM layers. In this paper, we introduce the Gated Memory Unit (GMU), a simple yet effective mechanism for efficient memory sharing across layers. We apply it to create SambaY, a decoder-hybrid-decoder architecture that incorporates GMUs in the cross-decoder to share memory readout states from a Samba-based self-decoder. SambaY significantly enhances decoding efficiency, preserves linear pre-filling time complexity, and boosts long-context performance, all while eliminating the need for explicit positional encoding. Through extensive scaling experiments, we demonstrate that our model exhibits a significantly lower irreducible loss compared to a strong YOCO baseline, indicating superior performance scalability under large-scale compute regimes. Our largest model enhanced with Differential Attention, Phi4-mini-Flash-Reasoning, achieves significantly better performance than Phi4-mini-Reasoning on reasoning tasks such as Math500, AIME24/25, and GPQA Diamond without any reinforcement learning, while delivering up to 10× higher decoding throughput on 2K-length prompts with 32K generation length under the vLLM inference framework. We release our training codebase on open-source data at https://github.com/microsoft/ArchScale.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper addresses the efficiency limitations of Transformer models when handling long sequences by proposing a hybrid architecture that integrates a state-space model (SSM) into a decoder-decoder framework. Specifically, the authors enhance the YOCO architecture by replacing its self-decoder with Samba, a hybrid SSM, and introduce a gated memory unit (GMU) to improve cross-decoding. The resulting model, SambaY, is evaluated across several tasks:\n\n- Scaling on NLP tasks: SambaY demonstrates superior FLOPs scaling compared to Transformer baselines. However, the GMU offers limited additional benefits.\n- Long-context retrieval: SambaY outperforms Transformer baselines, with the GMU and differential attention (DA) contributing positively. The GMU, however, only gives minor benefits on other downstream tasks (Wiki, ARC-c, etc.). \n- Large-scale pretraining: A scaled version of SambaY (Phi4-mini-Flash) trained on proprietary data outperforms Phi4-mini on downstream tasks.\n- Reasoning tasks: Phi4-mini-Flash matches Phi4-mini-Reasoning in performance while improving throughput.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths: \n- The paper presents a comprehensive experimental evaluation across multiple tasks and scales.\n- The integration of SSMs into a decoder-decoder architecture is well-motivated and practically relevant.\n\nWeaknesses: \n- **Missing baselines** As this work mainly builds on top of YOCO, it should be considered as a baseline in the experiments (compute and data scaling experiments, retrieval, and downstream tasks). Otherwise, we cannot see the effect when going from SWA to Mamba in the self-decoding. Moreover, Table 4 should also include the Phi-mini-Flash (without reasoning SFT). \n- **Mixed results of the gating mechanism (GMU)** The experiments fail to clearly demonstrate the benefits of GMU. In many benchmarks, just merging Samba with YOCO (Samba+YOCO) is on par with SambaY. Indeed, the performance of SambaY (which uses GMU) had to be pushed with differential attention (DA). In that regard, testing Samba+YOCO+DA may even give the best results.\n\nQUESTIONS:\nI would appreciate if the additional experiments (mentioned in Strengths and Weaknesses) could be included in the rebuttal. Moreover, I have the following comments/questions: \n- Line 164 states “This aligns with our expectation, as there is no theoretical expressiveness gap between Transformers and our hybrid models when the number of parameters is held constant.” This statement is very strong and requires more investigations (e.g., [Merill et al, ICML 2024]). \n- Line 288 states “Our extensive scaling experiments demonstrated that SambaY\nexhibits a lower irreducible loss compared to strong baselines, indicating superior scaling properties with increasing computational resources.” This statement is clearly conditioned on the modelling/fitting of the loss function. Hence, it should also be stated that way.\n\nLIMITATIONS:\nWhile the second point in the NeurIPS Paper Checklist states that the limitation is included in the conclusion section, they are instead (a bit hidden) in the appendix. It would be better to have limitations in the main paper.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 2, "DATE": "2025-06-20"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nFirst, this paper introduces the Gated Memory Unit (GMU), a neural network layer for accessing memories of other/previous layers.\nSecond, the GMU is then applied to a new Hybrid Decoder-Decoder architecture called SambaY, which uses Samba (a hybrid Mamba Attention architecture) for the self decoder and replaces half of the cross-attention layers with GMUs to access the inner representations of the SSM state of the Samb self decoder layer.\nFinally, the paper introduces Phi4-mini-Flash, a 3.8B model trained on 5T tokens, which replaces the attention layers in SambaY with Differential Attention (DA).\n\nThe post-trained version Phi4-mini-Flash-Resoning achieves comparable performance to the Phi4-mini baseline on reasoning tasks while delivering higher decoding throughput.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\n- Extensive and large-scale experiments on scaling behavior and good scaling performance\n- The scaling setup is explained in detail and the relations between width, depth, and hyperparameters are given as formulas.\n- Detailed ablation study on long context performance\n- Superior performance compared to Transformer++ baseline on Long-Context Benchmarks\n\nWeaknesses:\n- In the ablation study in Table 5, the Differential Attention used in the final Phi4-mini-Flash model is missing\n- The authors observe significantly better scaling behavior with their µP++ parameterization over SP. Since µP++ is a (new) combination from µP and Depth-µP it would be helpful to clarify the differences to SP in more detail, i.e. how does SP look like in this paper and what was modified going from SP to µP++.\n\nQUESTIONS:\n- How do you explain the drop in performance on PB-32k when moving from SambaY to SambaY-2?\nIt seems that without SWA Mamba2 is better, while with SWA Mamba1 is better. \nThis seems to contradict the explanation in L. 263. \n( Could it be that the reason is a less effective memory access to matrix states with GMU? (see next question) )\n- How do you apply the GMU to memory states which are matrices instead of vectors (i.e. Mamba2 has a matrix memory state)? Do you flatten the matrices?\n\n- L.279: „Gating attention/MLP representations performs worse … because they lack recency bias ..“: Why is the recency bias in the memory state necessary/helpful for GMU?\n\n- Is the fitting of the powerlaw in Figure 2 performed in logspace? \n- How does the SambaY or SambaY-DA architecture scale with number of parameters beyond 3.8B paramaters? Since the authors allocated a significant amount of compute to train the 3.8B model on 5T tokens, why were larger models not explored? Is there anything that prevents scaling to larger models?\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-01"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper introduces SSMs to the space of decoder-decoder architectures such as YOCO. It introduces a new Gated Memory Unit (GMU) to efficient memory sharing across layers. This, when introduced into the cross-decoder of the base YOCO architecture (modified with Samba SSMs), creates a new decoder-hybrid-decoder architecture, referred to as Samba. Since the model shares KV cache across the decoder and cross-decoder, it helps ensure that prefill is linear in complexity (since only last full attention computes the KV cache for rest of cross-decoder). The authors then follow this up with rigorous training experiments, including µP++ - which combines both regular µP and depth-µP, batch size scaling and data scaling to show scaling behaviors (through power law curve fits) of the models in comparison to transformers and its variants transformer++ - both with SP (standard trianing) and µP++. They follow this up with long context extension (& evaluation) to show that the model scales well for inference (better than transformers, just slighlty worse relative to pure Samba architectures). They finally follow this up by scaling to a new model, Phi4-mini-flash, including large scale pre-training + SFT for reasoning.\n\nSTRENGTHS AND WEAKNESSES:\n**Strengths**\n\n1. The architecture introduces GMU, and shows how the modifications help with modeling\n2. The paper throughly explores all axes (scaling model size, dataset size, pre-training haprams, long-context retrieval, downstream evaluation)\n3. The different architecture ablations are useful to understand the design choices of the model.\n4. The paper also shows inference scaling (through random weights in vLLM) to showcase how the linear-time complexity behavior they talk about.\n\n**Weaknesses**\n\n1. For some of the architecture ablations, it would have been good to see better downstream tasks for evaluation. Most of the tasks selected are very generic language modeling tasks with relatively high variance in scores - making the presented results a little tricky to parse.\n\nQUESTIONS:\nN/A\n\nLIMITATIONS:\nN/A", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 4, "IMPACT": 4, "ORIGINALITY": 3, "DATE": "2025-07-02"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis work leverages representation sharing between SSM layers in hybrid LLM architectures by gating states (memories) from earlier layers. To this end, the Authors propose the Gated Memory Unit (GMU). The Authors conduct extensive parameter and data scaling studies to compare the resulting architectures with competitive baselines. In a proof of concept based on the Phi4-mini family, they show performance gains in standard benchmarks and throughput increase in reasoning settings.\n\nSTRENGTHS AND WEAKNESSES:\n**Strengths**\n- The experiments are quite extensive and include parameter and data scaling studies for the considered architectures, with carefully scaled hyperparameters.\n- The proofs of concept experiments on proprietary data are quite convincing. The gains of Phi4-mini-Flash, which incorporates the proposed GMU as well as pre-existing DA and other improvements, over Phi4-mini are substantial. Throughput gains over Phi4-mini-Reasoning in reasoning tasks are also noteworthy.\n\n**Weaknesses**\n- I am partly concerned about the novelty of the contributions and their coherence as part of a singular narrative. In particular:\n    - GMUs, the µP++ hyper-parameter transfer scheme, and the Differential Transformer modification are given equal consideration (e.g., emphasis in the introduction); however, only the first one is novel to this paper. \n    - To my understanding, the main contribution is extending YOCO’s idea (memory sharing across layers) from attention layers to SSM layers in a hybrid model. While this is a meaningful extension, it does not constitute a fully original idea.\n- The main paper is not entirely self-contained: concepts such as YOCO, Cross-Decoder, SWA layers, etc., which the present approach builds on, should be described as part of the method or background and not taken for granted.\n- The range of downstream tasks seems somewhat limited for long context, which is one of the main foci of the paper: why are only a subset of RULER tasks included (e.g., NIAH but no variable state tracking)?\n\n**Minor**\n- Why is the architecture defined in section 3.1 called Transformer++?\n- The symbols in lines 99 and 111 should be explained to understand their meaning.\n- In the ablation study in section 4, you list the model variants without the specific question/ablation each one is meant to test.\n\n**Typos**\n- “Stat[e]” (line 1)\n\nQUESTIONS:\n- If YOCO reuses keys and values across layers, shouldn’t this also benefit decoding efficiency, as fewer values must be retrieved from GPU memory? (line 29)\n- What is the advantage of removing RoPE in practice? This does not affect efficiency, and you do not provide empirical evidence that it mitigates some of RoPE’s possible challenges (e.g., length generalisation).\n- Are you trying to claim that this architecture is superior for long context in general or just with respect to specific baselines? If the former, you would need to include state-of-the-art efficient transformer baselines, e.g., Quest sparse attention.\n- How are the learning rate and batch size chosen in line 114?\n- It might be just because of the image granularity, but how is C smaller for SambaY than Samba+YOCO in Figure 2b when the layer seems to display a consistent delta between the two lines?\n\nLIMITATIONS:\nA more thorough discussion on how the usage of private data impacts reproducibility is needed.", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 4, "CLARITY": 1, "IMPACT": 3, "ORIGINALITY": 2, "DATE": "2025-07-03"}