Decoder-Hybrid-Decoder Architecture for Efficient Reasoning with Long Generation
Authors: Liliang Ren
Congcong Chen
Haoran Xu
Young Jin Kim
Adam Atkinson
Zheng Zhan
Jiankai Sun
Baolin Peng
Liyuan Liu
Shuohang Wang
Hao Cheng
Jianfeng Gao
Weizhu Chen
yelong shen
Venue: NeurIPS 2025 poster ()
Keywords: Hybrid Neural Architectures;Long Sequence Models;State Space Models;Linear Recurrent Neural Networks
Primary area: deep_learning
TL;DR: A hybrid architecture with linear pre-filling complexity and up-to10x higher throughput on decoding.

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.

Decoder-HHyrrd-Deeodd Architecture for Efficient Reasoning with Long Generation
Liliang Ren'
Congcong Chen'
Haoran Xu
Young Jin Kim Baolin Peng
Adam Atkinson'
Zheng Zhan'
Jiankai Sun
Liyuan Liu"
Shuohang Wang' Hao Cheng'
Jianfeng Gao'
Weizhu Chen Yelong Shen
Microsoft Stanford University
{ (liliangren, yeshe cmicrosftt.com
Abstract
Recent advances in language modeling have demonstrated the effectiveness of State Space Models (SSMs) for efficient sequence modeling. While hybrid ar- chitectures such as Samba and the decoder-decoder architecture, YOCO, have shown promising performance gains over Transformers, prior works have not in- vestigated the efficiency potential of representation sharing between SSM layers. this paper, we introduce the Gated Memory Unit (GMU), simple yet effec- tive mechanism for efficient memory sharing across layers. We apply it to create architecture that incorporates GMUs the cross-decoder share memory readout states from Samba-based self-decoder. SambaY significantly enhances decoding efficiency, preserves linear pre- filliinnng time complexity, and boosts long-conteds performance, all while eliminating the need for explicit positional encoding. Through extensive scaling experiments, we demonstrate that model exhibits significantly lower irreducible loss compared to strong YOCO baseline, indicating superior performance scalability under large-scale compute regimes. Our largest model enhanced with Differential Atten- tion, 1e144-mmin-Flas--Raaaaaee Bhi4-mini-RRasoniing on reasoning tasks such Math500, AIME24/25, and GPQA Diamond without any reinforcement learning, while delivering up to 10x higher decoding throughput on 2K-length prompts with 32K generation length under the vLLM inference framework. We release our training codebase on open- operesource data https: //gg/ttttbu/// commmi wwwroossfftArrssshoooo
YambaY, becoder-hybrid-decooe
achieves significantly better performance than
Introduction
State Space Models (SSMs) [GGR21, GGGR22, GD23, DG24], including linear attention [HDLL22, SDH+23, QHS 22, YWS 23, YWZ 24, YKH25) and modern Recurrent Neural Networks (RNNs) [BPS 24, BPL 22, PAA + 23, GOA + 24], have recently shown promising results for more efficient sequence modeling over Transformers IVSP+ 17]. While pure SSMs/RNNs offer computational advantages with their linear complexity, they exhibit theoretical expressiveness gap the in- context retrieval capability relative to Transformers [WDL224]. This gap can be bridged by hybrid architectures [LLB+ 24, DSF+ 24, RLL 25, WBR+ 24, DFD 25, Min25], even with the inclusion of as few as single full attention layer [WDL24]. Recently, the decoder-decoder architecture, YOCO [SDZ+24], achieves linear complexity for long context processing through storing the Key-Value (KV) pairs from a single self-attention layer and re-using them across all subsequent layers through cross-attentions. In practice, YOCO has delivered substantial efficiency gains when processing user prompts with long sequences, but challenges remain; it does not mitigate the attention memory
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

I/O cost for its cross- rross-attentions during the generation stage of the model responses. This limitation becomes particularly pronounced for modern large language models (LLMs) [Ope24, DA25] that generate extensively long Chains-of-Thonn
(CoTTs) [WWS+22] for hard reasoning tasks.
In this paper, we investigate the potential of representation sharing between SSM layers to enhance decoding efficiency. We introduce the Gated Memory Unit (GMU), a versatile, simple yet effective mechanism for efficient memory sharing across layers. Applying GMUs to the cross-decoder of YOCO, we create a novel model with our decoder-hy decoodr-hyybidd-eccode brid-decoder architecture named Y that uses Samba 255] for the self-decoder and replaces half of the cross- crossaattention layers with GMUs to share the inner representations of the final SSM layer in the celf-decoder. Since around 50% of expensive cross- attention layers are replaced with cheap element-wist gating, Samba Y significantly improves decoding efficiency and maintains linear pre- rrelfiliing time complexity, all while removing the need for explicit positional encoding such as RoPE [SI.P-2 21].
To enable a robust comparison of the scaling capabilities across different architectures, we first design principled μP++ hyperparameter transfer scheme that accounts for both depth and width scaling, as well as the application of weight decay to vector-like parameters. We then conduct extensive experiments up to 3.4B parameters/600B tokens verify the scaling behaviors of both our μP++ scaling laws and the SambaY architecture. Comparing to Samba+YOCO an architecture that naively combines Samba with YOCO, we show that SambaY has significantly lower irreducible loss [HNA + 17] on the validation set when scaling with the training FLOPs, indicating a better scaling potential with large-scale computes. We also conduct extensive experiments verify the long-context retrieval capabilities of our architecture. Our results reveal that Samba Y achieves superior performance on challenging long-context tasks like Phonebook and RULER [HSK+ 24] benchmark, even with a modest Sliding Window Attention (SWA) size of 256. To further explore the capabilities of hybrid models with a single set of full attention memory, we augment SambaY with Differential Attention [YDX 24], resulting in the Phi4-mini-Flash architecture. We pre-train our 3.8B-parrmetter model Phi4-mini-Flas
with tokens from the same Phi4-mini data corpus and further follow Phi4-mini-RR Phi4-miii-Reessoinng IXPA SPA+22 25] to conduct the multi-stage distillation with Supervised Fine-Tuning (SFT) and Direct Preference Optimization (DPO) to produce our reasoning model, Fhii-mini-Flash-Reeuunnnnn Our model achieves significantly better performance than the strong Mhi4-miniRRcasomee baseline on challenging reasoning benchmarks such as Math500, AIME24/25, and GPQA Diamond, while excluding any stage of Reinforcement Learning (RL) that is used by Phi4-mini-Reasoning Critically, our wwwimiii--laaa-aaaae throughput on 2K-length prompts with 32K generation length under the vLLMM [KLZ* 23] inference framework, showcasing its substantial and practical efficiency gains for the LLM reasoning paradigm of generating long Chain-of-Thoughts.
delivers up to 10x higher decoding
2 Decoder-HtbrrddDDeood Decdde Architecture
Inspired by the gating mechanism that broadly exists in Gated Linear Units [Sha20], Gated Attention Units [HDLL22] and SSMs [GD23, YWS+23. YKH25), introduce our Gated Memory Unit (GMU) together with its application on YOCO, which produces our final decoder-hybrid decoder architecture. We include a dedicated related works section in Appendix I, the limitation section in Appendix and provide the background introduction of YOCO in Appendix A.
Token mixing as a matrix operator. Both state-space models (SSMs) and self-attention layers perform token mixing through a linear operator that can be written as a matrix A E R=*x", where is the sequence length. In SSMs, A is highly structured matrix that captures the parallel form of an underlying recurrent update, whereas in self-attention is the row-aggregating attention matrix whose entries are the query-ke key softmax probabilities. For given head at layer 1", the mixed representation is
M (')
(1) = (0')) v(')
where ((') guxd, denotes either the SSM state inputs or the attention value vectors.
Gated Memory Unit (GMU). From an inter-layer perspective, we define "memory" as hidden representations passed from preceding layers. Specifically, at a given layer GMU operates on two inputs: the current layer's input hidden state, X₁ E Raadd. and the mixed representation M('), E nad of previous layer l' (where L' < <). The GMU then produces an output Y e =dd

---

Mamba
Self-
Decoder
Linear
SWA
SSM kernel
g(-)
Mamba
Linear
Output State
Attn
GMU
Cache
Linear
ar)-
Cross Attn
GMU
Linear
Cross-
Decoder
Cross Attn
Figure Our decoder-hybee decoder--ybrid-deccdeer architecture taking Samba [RLL 225] the self-decoder. Gated Memory Units (GMUs) are interleaved with the cross-attention layers in the cross-decoder to reduce the decoding complexity. in YOCO [SDZZ+24], the full attention layer only need to compute the KV cache during prefillnn with the self-decoder, leading to linear computation complexity for the prefill stage.
through gating mechanism modulated by learnable projections. Formally, the GMU can be expressed as:
= (M'()) O o(
where (.) the SiLU [HUD17] activation function, @ denotes element-wise multiplication, and W1. W2 E [ddi.dd are learnable weight matrices. One can also apply an RMSNorm [ZS19] layer after the element-wise multiplication, yielding the normalized GMU (nGMU). We show that this normalization is crucial when the memory originates from linear attention mechanisms (DG24, YKH25], detailed in Appendix B and Appendix H. Intuitively, the GMU enables dynamical and fine- line-grained recalibration the token mixing performed a previous layer, conditioned on the current layer's input across each of the memory channels. Specifically, for each element His in the gated output H M @ G()) Raxdd G()) (X,W have
H G Σ
V
E A
v( 5k
Σ
(1) ik'
Aijk
which shows that the gate ( effectively lifts the previous token mixing matrix A(") into a third- order tensor with elements A.j G( A yielding a learned, channel-specific reweighting of the original token-mixing operator while maintaining linearity on the original value matrix V(() at later layer 7. The concept of the GMU generalizes beyond token-mixed memories. For instance, it can gate intermediate outputs from previous MLP layers, enabling retrieval from projected representation of earlier layer's 's input. In all cases, GMUs reduce both parameter count and computational cost compared to standard SSM, attention, or MLP layers.
Model architecture. In Figure 1, we illustrate the SambaY architecture, using our decoder-hybrid- decoder architecture with Samba [RLL+25 as the self-decoder. We apply GMUs the cross- decoder of YOCO to replace half of its cross-attention layers. The GMUs share the representation from the last SSMs layers in the self-decoder that the pre-filling time complexity is still linear. Compared to YOCO, our approach only requires caching an additional SSM kernel output state m dR da 2dm 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 During decoding, we reduce the memory I/O complexity for half the cross-attention layers from a linear cost of ((d,, to constant (dd)), where N is the sequence length and der is the dimension of key/value vectors. This leads to significant efficiency gains when N > dy, dgvs a condition that is easily met in practice since the ratio ds/dg, typically does not exceed 128.

---

Experiments and Results
Motivated by the theoretical efficiency of our SambaY architecture, we aim to address the follow- ing research questions: Does the architecture scale effectively? Does it compromise long-context performance? Can it support reasoning over long generations? Given that a neural architecture s performance is tightly coupled with its optimization and initialization settings, we begin by estab- lishing a generic scaling setup to encourage fair comparison of scaling behavior across different architectures.
Baseline Architecture. Apart from our proposed SambaY architecture, we consider the following baselines in this section: Transformer++ [HDLL22] (which uses SwiGLU [Sha20] for MLP and RoPE [SLP+21] with the base frequency of 10,000), Samba+YOCO (which uses Samba as self- decoder for the original YOCO architecture). SWA+YOCO (the original YOCO with SWA layers as self-decoder), TransformerLS (interleaving SWA with full attention using layer ratio of 3:1), and SambaY+DA (which uses Differential Attention (DA) [YDX* 24] for all attention layers in SambaY), More architecture details are included in Appendix D. We standardize the sliding window size to 128 for all architectures with SWA while conducting extensive ablations on window size effects in Section 3.2. Following the studies in recent hybrid models [LLB* 24, RLL CLL+25], we omit explicit positional encodings (NoPE) for hybrid SSMs architectures, and demonstrate that NoPE enables zero-shot 2x retrievable context extrapolation in Appendix F.
3.1 Scaling Experiments on Open-Source Data
Architecture scaling setup. We use simple linear rule from the previous works on Transformer models 20, TJY+24] for scaling the architectural shape of our Transformer++ baseline, including model width wo, model depth d, number of attention query heads hy and the MLP inner dimension Wmlpp i.e.
ad, ag 128, hg = d, hgu d/4,
Wmip 4w,
where the Transformer- specific aspect ratio ag and the number of key-value heads h&s are computed based on the Llama-3-44 3-BB [Met24] architecture. The total number of non non-embedding parameters N (d) for the Transformer++ architecture can then be calculated as,
Nat (d) = 2.5dw2, Nmip(c) (d) = 12dw",
(d) Naan (d) + (mii(d) 14.5dw² 237568d
where Natn; Nmip means the total number of parameters for attention/MLP layers respectively.
Scaling transfer for hybrid architectures. Since different token mixers have their own inner dimension expansion ratio, it is hard to balance the number of parameters between hybrid models and Transformers to make fair comparisons. Previous works [DA24a, 25, YKH25) 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 12-layer Transformer resulting in 16. 7% increase of KV cache size), making unfair comparisons on the inference time cost. We propose a simple solution that (1) builds an coo-parametric equation with respect to the aspect ratio via aligning the total number of non-embedding parameters to the Transformer baseline, (2) solves the equation to obtain the specific aspect ratio for the hybrid architectures. We maintain consistent hyperparameter settings with a 4x MLP inner dimension expansion ratio and GQA [ALTd]* 23] group size of 4 for self-attention layers, matching our Transformer++ baseline. We also fix the head dimension to be ag = 128, and set the inner dimension of the attention layers to be Wain and so that the number of key-value heads lgg a valid integer. Specifically, for SambaY, we have
Naun (d) = 2.5dw Wann/ + 2dw Wann/44 /4, Nmamhu (d) 6dw2/44
(d) = 4dw2 /4,
(dd) Naun (d) + Nmamba (d) + Nmup (d) + Ngmu (d) 14 24 237568d3
where Nam: Nmamha, Nmlp: Ngmm means the total number of parameters for attention/Mam- ba/MI.P/GMU layers. Solving for a, we get on R. 124. For Samba+YOCO, we can similarly solve -srammrmet equation to obtain ag se 126, with more details in Appendix C.

---

Hyperparameter scaling with pP++. To account for both width and depth scaling of model architectures, we propose JuP++ hyperparameter scaling laws that integrate μP [YHB* + 22], Depth-uuP [YYZH23), and apply zero weight decay to vector-like or scalar-like parameters' for training stability, Since we use the AdamW optimizer [LH18], we apply batch-size-basee scaling with vy x VB. The learning rate is further scaled as 1/vd following Depth-uP.. For studying the FLOPs scaling behavior across model architectures, we adopt the Chinchilla scaling law [HBM+22] to scale the number of training tokens T linearly with the number of model parameters. Formally, we have
Bdo Byd
N(d) B = Bo: T =o N(d))
7 == 100
where the base learning rate is set as 7/0 x 10-4 and the base batch size Bo = 221 2M number of tokens. We also explore scaling the batch size sub-linearly with respect to the training tokens [MKAT18, SWW+ 24, LZH+ but find that it harms the data scaling behavior of the models, so we keep the batch size as a constant across scales. The base model depth is set as do 16 so that N (dg) Rs 10 number of parameters. The base training token count To is set to 100B, corresponding to 5x the Chinchilla-opiee optimal ratio of tokens per parameter (upproximately 20 based on [HBM 22]). in order to study scaling behaviors in a typical over over-training regime. We summarize the differences between Standard Parametrization, μP and μP++ in Appendix D, while providing large scale ablation studies Appendix E.
Scaling experiment setups. A common concern with SSMs is that they not theoretically more expressive than self-attention for in-context retrieval [WDL224]. This raises the question of whether the better performance of hybrid SSM models is owing to their fast convergence from the recency bias, while Transformers could eventually match their performance given more training tokens. With the scaling laws we established in the previous paragraphs, we can now examine this hypothesis systematically. We first study the data scaling behavior across architectures through fixing the model size IB parameters with d 16 and scaling the number of training tokens T from 100B to 600B. We also study the FLOPs scaling behaviors of the model architectures with up to 3.4B parameters and 342B tokens through varying the model depth d {8, 12, 16, 20, 24]. We use 4K training sequence length and the SlimPajama [SAKM* 23] dataset for all our scaling experiments. More experimental details are included in Appendix D.
Comparison of scaling behaviors. To quantitatively compare the scaling trajectories, we fit the validation loss L as a function of compute (FLOPs), denoted as DFLOPs. to a power law [HNA+ + 17, HBM of the form:
I(DFLOPS) A D + C
This model enables us to estimate the irreducible loss C which represents the lower bound of achievable loss for a given architecture or parameterization under infinite compute, and the scaling exponent b that reflects the learning efficiency with respect to compute. We fit the curves with least squares and the LMA algorithm ILEV44, Mar63]. A similar power law model is employed for data scaling experiments, where loss is modeled as a function of the number of training tokens Duokens-
In Figure 2, we present the results of both FLOPs scaling and data scaling experiments, showing validation loss on the SlimPajama dataset as a function of total training FLOPs or number of training tokens. We show both the original data points and the fitted power-law curves. The goodness of fit for each curve is assessed using the R² statistic and all plots have a fitness score R² 2 0.999, indicating near- cear-perfect fits. 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 's long-term scaling potential. As illustrated in Figure 2a, the SambaY architecture exhibits the lowest irreducible loss (C 0.58) for FLOPs scaling, suggesting that can attain a superior validation loss when scaled further with substantially increased computational resources. However, 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
Following the definition in JPP. parameters are vector-like when exactly one dimension scales with model width (e.g.. embedding and unembedding layers), and scalar-like like when dimension scales with width.
5

---

Tunalfmmer
Tandfrreee (S%
To sifern ISC Ghw 44.0m18t
Thansh
2.06
Traad's phw.kit
Saniae roco (Aorl (2 20037)224444
2.04
Semba-1020 (PP+))
Sambar-0C0 gate Mark 10 27.
2.4
Sansan
2.02
Sardoo lune
Loos
12.00
1.40
1.8
10
1:34
12a
30
2 30"
5% 30
6X
6x 10
(a) Compute scaling comparisons
(b) Data scaling comparisons
Figure 2: Validation Loss v.s. FLOPs (left) or Training Tokens (right) on the SlimPajama dataset. Besides the architecture comparisons, we also compare our pP++ based scaling with the Standard Parametrization (SP).
efficiency with respect to compute. We also observe that μP++ yields lower irreducible loss than SP under both data and compute scaling, indicating more favorable scaling potentials.
Notably in Figure 2b, the Transformer++ model trained with μPP+ exhibits large validation loss gap compared to SambaY and Samba+YOCO within the measured range of training tokens. However, its fitted irreducible loss (C 1.82) 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 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 can be because we use Mamba-l our SSM which falls into the same complexity class of T'O" as Transformers [MPS24]. Interestingly, this convergence no longer holds when both model size and data scale proportionally. As illustrated in Figure 2a, under the μP++ setting, Transformer+ exhibits the highest irreducible loss C = 0.64, indicating that hybrid architectures could offer superior scalability under limited data regimes.
3.2 Efficient Long Context Retrieval
Given the presence full-attention layers, we aim to de- termine the minimal sliding window size required for our hybrid models to retain effective long-context retrieval ca- pabilities-a essential property for supporting advanced reasoning behaviors that involve long generation with backtracking. Specifically,
pre-train 1.0B parameter models with μPP++ and d = 16 using TransformerL.S, SambaY, SambaY+Da and Samba+ YOCO architectures respectively on the ProLong-664 [GWYC24] dataset with 32K sequence length and 40B tokens, varying the win- dow size of their Sliding Window Attention (SWA) in the range {64, 128, .... 2048). We align the number of parameters between different architectures through build- ing the sso-parametric equations as Section 3. We adopt variable-length training, where short documents are packed together and self-attended within the same seg- ment. We evaluate the long-contexx retrieval capabilities of the models using difficult Phonebook benchmark [JBKM24] with 32K context length (containing 1. 1,850 name-number pairs). We choose this benchmark because it is realistic which minimizes the confounding influence of instruction-folloviii ability when evaluating long- context retrieval performance. We use a RoPE base of 640K for Transformerl Transformer++ and SWA+ + YOCO, following the lower bounds proposed in [XMW We also examine how the training corpus and methods affect the long context performance, with more details in Appendix F.
SITIIAAS HAY WWN-1000
Figure 3: Accuracy (with error bars) v.s. Sliding Window Size on Phonebook with 32K evaluation length.
retrieval task with minimal instructions,

---

Table 1: Retrieval accuracy on BoedleIn-A-Haysteee (NIAH) tasks with 32K context from the RULER [HSK 24] long context benchmark. MK: Multi-Key, MQ: Multi-Query, MV: Value, S: Single-needle. We use the best Sliding Window Attention (SWA) size found on the Phonebook benchmark for each hybrid architecture. Best results are in bold, second best underlined.
Model
SWA
MK-1
MK-2
MK-3
MQ
MV
S-1
S-2
S-3
Avg.
Transformer++
36.4
3.8
0.0
27.9
24.1
94.8
66.0
31.0
35.5
TransformerL.S
256
42.8
6.0
0.0
29.8
27.5
91.8
49.6
23.4
33.9
SWA+YOCO
128
24.2
6.8
0.2
10.2
14.7
81.2
32.6
48.4
27.3
Samba+YOCo
1024
49.0
28.0
2.6
12.8
18.3
100.0
63.2
23.6
37.2
SambaY Y
256
54.6
27.8
0.4
12.7
19.4
83.2
81.2
63.8
42.9
Samba YDDA
512
64.6
27.6
0.2
12.8
19.9
99.8
86.4
69.6
47.6
As shown in Figure 3, which plots accuracy against SWA size on Phonebook dataset (32K evaluation length), surprisingly, larger SWA sizes do not consistently provide better results. Since learned full attention involves both sliding window (local) patterns and non-local patterns like global retrieval or attention sinks, using small sliding window sizes, where models like SambaY and Samba Y+DA show strong performance, could enable the model to focus on local patterns more easily and mitigate issues like attention sinks [GPD+ 25]. Moreover, shorter sliding windows facilitate faster convergence, a crucial factor in long context training scenarios that are often characterized by limited high-quality data. The lower scores of SWA+YOCK which consistently underperform the Y variants, could indicate that pure inrention-based models require more substantial data for long-context training.
Table 2: Downstream short-context evaluation language modeling and common-sense reasoning tasks in zero-shot for 1B B-parameter models with the tuned sliding window size. The training speed is measured in MTPS (Million Tokens Per Second) with 64 A100-80GB GPUs. Best results in bold, second best underlined.
Speed mips" 1
Wiki.
LMB. acc 1
ARC-c
ARC-e
Hella.
PIQA
Wino.
Avg.
Model
SWA
ppl
t
aec 1
aec
acc f
acc
ace
Transformer++
0.89
19.75
45.45
27.56
54.17
43.86
68.77
50.28
48.35
TransformerLS
256
1.46
18.49
48.77
28.84
57.11
45.85
69.21
53.67
50.57
SWA+YOCO
128
1.24
18.01
49.80
28.58
57.79
46.48
70.46
51.85
50.69
Samba+YOCO
1024
0..9
16.73
50.53
28.50
60.02
48.85
71.55
52.57
52.00
SambaY
256
1.11
17.83
50.40
29.44
57.87
49.08
71.00
55,25
52.17
SambaY+DA
512
0.91
16.59
49.68
28.33
60.65
49.53
71.38
53.43
52.17
Using the optimal sliding window size from the Phonebook benchmark, we evaluate our architectures on both long-comten retrieval tasks (Table 1) and traditional downstream benchmarks (Table 2). Across both contexts, hybrid models with SSMs consistently outperform pure Transformer architec- tures. YambaY variants demonstrate notable advantages in long-context retrieval while maintaining strong performance on short-context tasks, despite using much smaller sliding window sizes than Yamba+YOCO. The addition of DA further enhances multi-key and single-needle retrieval capa- bilities, while Transformer-based models show specific strengths in multi-query and multi-value scenarios. wwt.fforner..SSSAAYYYYYYY
outperforms Transformer++ on short-context tasks but falls behind on the RULER benchmark, highlighting the trade-off on long-context performance caused by introducing SWA full attention models. Overall, our results suggest that GMUs facilitate efficient representation sharing across layers and enable strong performance with smaller SWA sizes.
3.3 Large-Scale Pre-training High-quality Data
We apply our SambaY+DD architecture to pre-train larger-scale prototype model named Phi4-mini- Flash with 3.8B parameters. It uses an SWA size of 512 and GQA of group size 2. Compared to the configuration described in Section 3. this model uses a different aspect ratio a 80 and an attention head dimension of 64. It is trained with standard parameterization rather than μP+++ due to resource constraints at the time of scaling study. We pre-train our model on 5T tokens from the data corpus used by Phi4-mini [MAA- 25] on IK A100-80GB GPUs for 14 days. During training, encounter severe loss divergence, which mitigate by introducing label smoothing of 0.1 and
7

---

Table Downstream evaluation performance of post-trained models. We use the completion split for BigCodeBench evaluation. Bold indicates the best result per row.
Benchmark
Metric
Phi4-mini
Phi4-mini-Flash
MMLJU [21
5-shot
67.3
MMLU-Pro [WMZ* 24]
0-shot, CoT
52.8
54.7
Arena Hard [LCF+24 24]
Win Rate
32.8
34.9
GSM8K [CKB*21]
O-shot, CoT
88.6
89.5
Qasper [DLB 21]
F1
40.4
40.2
SummScreenFD [CCWG22]
ROUGE-L
16.0
17.0
BigCodeBench [5V
pass
43.0
44.5
MBPP [AON*21]
pass
65.3
69.8
attention dropout of 0.05. The optimization setup here is by no means optimal, as the primary goal of this experiment is evaluate the viability of our architecture larger scales. Phi4-mini-Flash uses 200K token vocabulary size consistent with Phi4-mini. Additional training and architectural details, including the mitigation of stability issues, are provided Appendix D. Table 3 demonstrates that Phi4-mini-Flast outperforms the Phi4-mini baseline across a diverse range of tasks, with notable benchmarks like MMLU and coding tasks such as MBPP.
improvements on snowledge-intenseee The consistent performance advantage, winning on 7 out of 8 benchmarks, is particularly significant given that Phi4-mini-Floe Flash achieves these gains while maintaining substantially higher computational efficiency during inference.
3.4 Efficient Reasoning with Long Generation
Table 4: Pass@1 performance of models reasoning benchmarks measured with maximum generation length of 32K. We report Pass@ I accuracy averaged over 64 samples for AIME24/25 and 8 samples for Math500 and GPQA Diamond to ensure evaluation robustness. We also evaluate popular open-source distilled reasoning models [DA25, Bes, Ope] as reference baselines.
Model
AIME24
AIME25
Math500
GPQA Diamond
DeepSeek-R1--isttll--nee
29.58
20.78
84.50
37.69
DeepSeek-R1-Disiiee -iiee-sll---
53.70
35.94
93.03
47.85
DeepSeek-R1 -asas-RRR--Distill--aaaaa
43.96
27.34
87.48
45.83
-tspoke-Strato--B OpenThinker-7B
21.51
18.28
80.73
38.51
29.69
24.32
87.25
41.60
Phi4-mini-RReasoingg (3.8B)
48.13
31.77
91.20
44.51
Phil-minii-Flssh-Reaaonnng (3.8B)
52.29
33.59
92.45
45.08
Our Phi4-mini-Flash-RRassee the same multi-stage distillation data following resources, we only conduct the distillation with SFT and DPO stages and leave RI. for future works. As shown in Table and Figure 4, our reasoning model achieves significantly better performance than 1ii4-mini-Reasomee [HBK+21], and GPQA Diamond [RHS*23] 23], while employing a substantially more efficient archi- tecture, achieving up to 10x higher throughput in geneeeetattion scenarios and 4.9x speedup in long-context processing. In Figure 4, we evaluate the throughput of our vLLM implementation' using random model weights to eliminate the influence of potentially shorter generation lengths speed measurements, ensuring a fair comparison across different architectures. The same hyperparameter configurations Phi4-mini-Flash We observe that SambaY achieves the best throughput in both long-context and Long-generation settings across various numbers of concurrent clients, highlighting the significant practical efficiency gains enabled by our GMU modules. Notably, our Differential Attention implementation relies on naive four-pass of the FlashAttention [Dao23] operator for vLLM compatibility, rather than the optimized custom kernel proposed in the original paper, leaving significant room for further speed
model is continually trained from the Phi4-mini-Flash model with
[XPA+25]] Due the limited
(which has final RL training stage) AIME24/25 [Art], Math500
applied for the YOCO and Y based baseline architectures.
We customize the official vLLM framework with the version 0.7.3 to support our Phi4-mini- mini-Flash architec- ture.

---

optimization of Phi4-mini- Mhi--mini-Flash-a Reasoning to catch up the efficiency of Samba Y. More evaluation details and case studies on our model's general reasoning ability are provided in Appendix G.
14-
Piskem--weering
Samba+1OCo
YOCO-DA
Sumba
Piie min-reasoning
Sambs+ YOCO+DA
Phis mn fiash roasoning
12
Sambol
2000
Latency (5)
Latency (5)
10-
1500
4.9x
10x
500
0.2
0.4
0.6
0.8
1.0
1.2
1.4
1.5
1.8
0.0025 0.0050 0.0075 0.0100 0,0125 0.0100 0.0755 0,0200 0.0225 Throughput (queries/s)
Throughput (queries/s)
(a) Prompt: 32000, Generation: 500
(h) Prompt: 2000, Generation: 32000
Figure 4: Throughput and latency of text generation with various architectures under the vLLM inference framework (using one A100-80GB GPU and no Tensor Parallelism). A normal distribution with 30% variance was applied to prompt and generation lengths with averages of 32000/2000 and 500/32000 respectively, following the setting [2T].
4 Ablation Study
Mamba
Membe
SWA
SWA
SWA
GON-A
Mamba
Mamba
Mamba
Full Atti
Full Attn
Full Attn
MLP
Output
Full Attn
Inner
Inner
State
KV
KV Cache
KV Cache
Cache
nOMU
Cross Attn
GMU
Cross Atth
Cross Attn
OMU
GMU
nOMU
Cross Attn
Cross Attn
Cross Attn
[
SambaY-A
SambaY-AA
SambaY-MLP
GDNY
Figure 5: Major architectural variants explored in this section. For GDNY, we use Gated DeltaNet [YKH25] with normalization after output gate (GDN-A) for self-decoder, and apply normalized GMU (nGMU) in cross drcoser
As illustrated Figure 5, we systematically study the design choices in our decoder-hybrrd-decom architecture through the following architectural modifications of SambaY:
Samba Y-2 S- G-DNY substitutes Mamba layers with Mamba-2 or GDN-A (as detailed in
Appendix B) respectively in the self-decoder, Mamba Maanha//aabaa2/GDDNA.a respectively in the self-decoder except the full attention layer. We find it is crucial to also use nGMU for Mamba-2/GDN-A based models to achieve strong long context performance, with ablation studies Appendix
employs only
Samba Y-A A applies GMUs to gate intermediate representations from the last full attention layer the self-decoder rather than from Mamba.
Samba Y-AA entirely removes aross-attention in the cross-decoder and instead uses GMU to gate the intermediate representations from the middle full attention layer.
SambaY-MM Y-MLP uses GMUs to gate the intermediate representation from the linear projection branch of the SwiGLU right after the full attention layer.
All ablation models are trained with 1.0 parameters on the ProLong-644K dataset with 40B tokens and maximum of 32K sequence length with variable length training, using SWA size of 128 as in

---

the scaling experiments. We leverage μP++ with depth d = 16 and construct iso- parameter equations maintain parameter count equivalence across all variants, with more details in Appendix D. We aim to answer the following research questions given the ablation results in Table 5.
Table 5: Downstream evaluation on Phonebook 32K (PB-32k), language modeling and common- sense reasoning tasks in zero-shot for IB-parameter models with a sliding window size of 128. We measure the training speed in MTPS (Million Tokens Per Second) with 64 A100-80GB GPUs. The average accuracy excludes PB-32K due to its relatively high variability, with a standard deviation of around 5%. Best results in bold, second best underlined.
Speed
Wiki.
PB-32K
LMB. acc f
ARC-c
ARC-e
Hella.
PIQA
Wino.
Avg. 1
Model
mtps 1
ppl
acc_n' t
t
acc_n1 f
acc
f
SambaY
1.10
16.89
78.13
50.22
28.58
59.18
49.07
70.84
55.09
52.16
MambaY
0.94
17.29
12.50
50.24
28.84
59.64
48.27
71.44
52.80
51.87
Samba Y-2
1.43
17.17
40.63
48.96
28.84
59.18
48.01
70.18
50.83
51.00
MambaY-2 Y-2
1.38
18.63
50.78
49.58
28.24
58.75
48.29
70.13
51.07
51.01
S-GDNY
1.34
16.78
83.59
50.94
29.61
58.96
48.93
71.55
51.85
51.97
GDNY
1.22
16.92
89.84
50.38
28.84
60.61
48.01
71.27
51.38
51.75
SambaY-A
1.11
18.12
58.59
49.85
30.29
59.60
48.41
71.33
"44066
$2.26
SambuY-AA
1.25
17.03
46.88
49.93
28.50
59.05
48.69
72.25
53.91
52.06
SumbaY-MLP
1.15
18.70
64.84
50.16
30.38
60.69
48.46
71.44
54.78
52.65
Alternative architectures for self-decoder in SambaY? As shown in Table 5, while SambaY performs well on the PB-32K benchmark. replacing its Mamba layers with Mamba-2 leads a significant drop in accuracy. This may be attributed to Mamba-2's s coarse, scalar-valued forget gates, which can reduce the self-decoder's 's capacity to encode fine- fine-grained positional information. The weaker PB-32K performance of Mamba compared to SambaY underscores the importance of local retrieval ability provided by SWA; recency bias alone appears insufficient for the self-decoder to support the cross-decoder in completing complex retrieval tasks. While GDN-based models achieve impressive PB-32K accuracy due to their enhanced retrieval capabilities with delta update rules, interleaving GDN with short-range SWA notably accelerates training without significantly degrading performance on either short or long-context tasks.
Is GMU effective for other memories beyond SSMs? We examine Samba Y-A and Samba AA, which gate attention inner representations, and SambaY-MLP, which gates MLP intermediate representations. As shown in Table 5, these variants achieve respectable performance on downstream tasks, with Y-ML even surpassing the original SambaY average accuracy for short-context tasks. However, for the long-context task, PB-32K, we observe a clear hierarchy: Samba Y > SambaY- MLP> Y-A Samba Y-An. AA. This pattern indicates that GMUs remain effective with alternative memory sources, but their performance on retrieval tasks depends significantly on the memory source's inherent characteristics. Gating attention/ML_P representations performs worse than the original Yamba'Y on Phonebook because they lack the recency bias that SSMs naturally provide, which is beneficial for encoding contiguous local information. S-mbaY-AA, which completely removes cross-attention, shows significant degradation, highlighting the importance of having sufficient number of cross-attention layers for the successful retrievals from a large pool of multiple key-value pairs.
Conclusion
In this work, we introduce the Gated Memory Unit (GMU), simple yet effective mechanism for efficient memory sharing across layers in sequence models. Replacing expansive cross attention architecture with Samba as the
layers with GMUs, we propose SambaY, a beccoder-hbbriddecoooe self-decoder, which achieves significant improvements in both computation efficiency and long- context performance. Our extensive scaling experiments demonstrated that SambaY exhibits lower irreducible loss compared to strong baselines when fitted with power laws against training FLOPs, indicating superior scaling properties with increasing computational resources. Our largest model, Phi4-mini-Flash -elass-naasnning outperforms Phi4-mini- mini-RRessonng on challenging reasoning benchmarks while delivering substantially higher decoding throughput on long-context generations. Given that our architecture still retains full attention layer with linear decoding complexity, future work could explore dynamic sparse attention mechanisms to further improve efficiency on extremely long sequence generation, particularly in agentic application scenarios.
10

---

Acknowledgement
We want to thank Yutao Sun, Li Dong, Songlin Yang and Yang Liu for helpful discussions and insights. We also want to thank Yi Zhu for an early version of the vLLM implementation of YOCO.
References
[33 J. Ainslie, Lee- Thorp. Michiel de Jong, Yury Zemlyanskiy, Federico Lebr' on, and Sumit K. Sanghai. Gqa: Training generalized multi-query transformer models from multi-head checkpoints. Conference on Empirical Methods in Natural Language Processing, 2023.
[ANZ22] Ibrahim M Alabdulmohsin, Behnam Neyshabur, and Xiaohua Zhai, Revisiting neural scaling laws in language and vision. Advances in Neural Information Processing Systems, 35:22300-2231 2022.
[1] Jacob Austin, Augustus Odena, Maxwell Nye. Maarten Bosma, Henryk Michalewski, David Dohan, Ellen Jiang, Carrie Cai, Michael Terry, Quoe Le, and Charles Sutton. Program synthesis with large language models. arXiv preprint arXiv: 2/08.07732, 2021.
[Art]
Art Prohlem Solving. Aime problems and solutions. httpp////rrtfprooleemlllloeeeeee aaki/index. phpp Aro___oooeeeessss and_aandddttt Soluutoouioos Accessed: 2025-04-20.
com/
[BBC+25]
Johan Bjorck, Alon Benhaim, Vishrav Chaudhary, Furu Wei, and Xia Song. Scaling optimal LR across token horizons. In The Thirteenth International Conference on Learning Representations, 2025.
[BDK 24]
Yasaman Bahri, Ethan Dyer, Jared Kaplan, Jachoon Lee, and Utkarsh Sharma. Explaining neural scaling laws. Proceedings of the National Academy of Sciences, 121(27)e23 2024.
[Bes]
Bespoke-stratos. Bespoke- Sttrrkos-ttttts-7 Accessed: 2025-06-18
Bespoke Labs.
https://huggiintaaet http:///muggggagaa.c.cccc///ettt
[BKН16]
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E. Hinton. Layer normalization. arXiv preprint
arXiv: 064500 2016.
[BMN*24] William Brandon, Mayank Mishra, Aniruddha Nrusimha, Rameswar Panda, and Jonathan Ragan Kelly. Reducing transformer key-value cache size with cross-layer attention. arXiv preprint arXiv: 2405. /298/, 2024.
[BPC20]
Iz Beltagy, Matthew E. Peters, and Arman Cohan. Longformer: The long-document transformer. arXiv preprint arviv: Arxiv-2004.057 05500, 2020.
[BPL++2]]
Maximilian Beck, Korbinian Pöppel. Phillip Lippe, Richard Kurle, Patrick M. Blies, Günter Klambauer, Schastian Bäck, and Sepp Hochreiter. xlstm 7b: A recurrent Ilm for fast and efficient inference. arXiv preprint arXiv: 2503. I 13427, 2025.
[BPS* 24]
Maximilian Beck, Korbinian Poppel, M. Spanring, Andreas Auer, Oleksandra Prudnikova, Michael K Kopp. G. Klambauer, Johannes Brandstetter, and Sepp Hochreiter. xlstm: Extended long short-term memory. Neural Information Prucessing Systems, 2024.
[BZB
Yonatan Bisk, Rowan Zellers, Ronan Le Bras, Jianfeng Gao, and Yejin Choi. PIQA: reasoning about physical commonsense in natural language. In The Thirty-Fourth AAAI Conference on Artificial Intelligence, AAAI 2020, The Thirty-Second, Innovative Applications of Artificial Intelli- gence Conference, IAAI 2020, The Tenth AAAI Symposium on Educational Advances in Artificial Intelligence, FAAI 2020, New York, NY, USA, February 7-12, 2020, pages 7432-7439. AAAT Press, 2020.
|CCE* 18]
Peter Clark, Isaac Cowhey, Oren Etzioni, Tushar Khot, Ashish Sabharwal, Carissa Schoenick, and Oyvind Tafjord. Think you have solved question answering? try arc, the ai2 reasoning challenge. arXiv preprint arXiv: 1803.05457 2018
[CCWG22]
Mingda Chen, Zewei Chu, Sam Wiseman, and Kevin Gimpel. SummScreen: A dataset for ab- stractive screenplay summarization. In Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics ( (Volume 1: Long Papers). pages 8602-88088 8615, Dublin, Ireland, May 2022. Association for Computational Linguistics.
[CKB*211 Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton. Reiichiro Nakano, Christopher Hesse, and John Schulman. Training verifiers to solve math word problems. arXiv preprint arXiv: 2110.14168 2021.
[DA24a]
DeepSeck-AL AL. Deepseek-v2: A strong, economical, and efficient mixture-of of-uue--o--eperttt language model. arXiv preprint arXiv: $405.04433 2024.

---

[DA24b]
DeepSeek-AI. Deepseek-v3 technical report. arXiv preprint arXiv: 2412.19437, 2024.
[DA25]
DeepSeek-AL. Deepseek-rl: Incentivizing reasoning capability in llms via reinforcement learning. arXiv preprint arXiv: 2501.12948, 2025.
[Dao23]
Tri Dao. Flashattention-22 Faster attention with better parallelism and work partitioning. arXiv preprint arXiv: 2307.0869/ 2023.
[DF
Xin Dong, Yonggan Fu, Shizhe Diao, Wonmin Byeon, Zijia Chen, A. Mahabaleshwarkar, Shih- Yang Liu, Matthijs Van Keirsbilek. MinnHHung Chen, Yoshi Suhara, Yingyan Celine Lin, Jan Kautz, and Pavlo Molchanow. Hymba: A hybrid-head architecture for small language models. International Conference on Learning Representations, 2025.
[DG24]
Tri Dao and Albert Gu. Transformers are ssms: Generalized models and efficient algorithms through structured state space duality. In Forty-first International Conference on Machine Learning, ICML 2024, Vienna, Austria, July 21-27, 2024. OpenReview.net, 2024.
[DLB+21]
Pradeep Dasigi, Kyle Lo, Iz Beltagy, Arman Cohan, Noah A. Smith, and Matt Gardner. A dataset information-sseking questions and answers anchored research papers. Proceedings of the 2021 Conference of the North American Chapter of the Association for Computational Linguisties: Human Language Technologies, page 4599 4610, Online, June 2021. Association for Computational Linguistics.
[DSF+24] Soham De. Samuel L. Smith, Anushan Fernando, Aleksandar Botev, George Cristian- Al- bert Gu, Ruba Haroun, Leonard Berrada, Yutian Chen, Srivatsan Srinivasan, Guillaume Desjardins, Arnaud Doucet, David Budden, Yec Whye Tch, Razvan Pascanu, Nando De Freitas, and Caglar Gulcehre. Griffin: Mixing gated lincar recurrences with local attention for efficient language models. arXiv preprint arXiv: 2402. 19427, 2024.
[DY Haojie Duanmu, Zhihang Yuan, Xiuhong Li, Jiangfei Duan, Xingcheng Zhang, and Dahua Lin. Skvq: widing-window key and value cache quantization for large language models, arXiv preprint arXiv:2405.0m 262.:2990.06619, 2024.
[EUDI7] Stefan Elfwing. E. Uchibe, and K. Doya. function approximation in reinforcement learning. Neural Networks, 2017.
linear units for neural network
[FD Daniel Y Fu, Tri Dao, Khaled K Saab, Armin W Thomas, Atri Rudra, and Christopher Ré. Hungry hungry hippos: Towards language modeling with state space models. arXiv preprint arXiv:2212. 14052, 2022.
[GD23]
Albert Gu and Tri Dao. Mamba: Lincar-time sequence modeling with selective state spaces. arXiv preprint ar1iv:2312.0000 00752,, 2023.
[GGGR22)
Albert Gu, Ankit Gupta, Karan Goel, and Christopher Ré. On the parameterization and initializa- tion of diagonal state space models. ARXIVORG, 2022.
[GGR21]
Albert Gu, Karan Goel, and Christopher R'c. Efficiently modeling long sequences with structured state spaces. International Conference On Learning Representations, 2021.
[GOA 24] Daniel Goldstein, Fares Oheid, Fric Alcaide, Guangyu Song, and Fugene Cheah, Goldfinc:: High performance rwkw/transformer hybrid with linear pie-fill and extreme kv-cache compression. arXiv preprini arXiv: 2407.12077. 2024.
[GPD+25]
Xiangming Gu, Tianyu Pang, Chao Du, Qian Liu, Fengzhuo Zhang, Cunxiao Du, Ye Wang, and Min Lin. When attention sink emerges in language models: An empirical view. In The Thirteenth International Conference on Learning Representations, 2025.
[GSF+25]
Riccardo Grazzi, Julien Siems, Jörg K.H. Franke, Arber Zela, Frank Hutter, and Massimiliano Pontil. Unlocking state- mate-tracking in linear RNNs through negative eigenvalues. The Thirteenth International Conference on Learning Representations, 2025.
[GWYC24] Tianyu Gao, Alexander Wettig, Howard Yen, and Danqi Chen. How train long-context language models (effectively). arXiv preprint arXiv: 24 10.006660 2024.
[HBB+21] Dan Hendrycks, Collin Burns, Steven Basart, Andy Zou, Mantas Mazeika, Dawn Song, and Jacob Steinhardt. Measuring massive multitask language understanding, 9th International Conference on Learning Representations, ICLR 2021, Viriual Event, Austria, May 3-7, 2021. OpenReview.net, 2021.
[HBD* 19]
Ari Holtzman, Jan Buys, Li Du, Maxwell Forbes, and Yejin Choi. The curious case of neural text degeneration. International Conference on Learning Representations, 2019.
[HBK 21]
Dan Hendrycks, Collin Burns, Saurav Kadavath, Akul Arora, Steven Basart, Eric Tang, Dawn Song, and Jacob Steinhardt. Measuring mathematical problem solving with the math dataset. arXiv preprint 2003:2200,.33874, 2021.
12

---

[H
Jordan Hoffmann, Sebastian Borgeaud, A. Mensch, Elena Buchatskaya, Trevor Cai, Eliza Ruther- ford, Diego de Las Casas, Lisa Anne Hendrricss, Johannes Welbl, Aidan Clark, Tom Hennigan, Eric Noland, Katie Millican, George van den Driessche, Bogdan Damoc, Aurelia Guy, Simon Osindero. K. Simonyan, Erich Elsen, Jack W. Rac, Oriol Vinyals, and L. Sifre. Training compute-optimal large language models, ARXIVORG, 2022.
[HDLL22] Weizhe Hua, Zihang Dai, Hanxiao Liu, and Quoc V. Le. Transformer quality in linear time. International Conference On Machine Learning, 2022.
[HNA+17] Joel Hestness. Sharan Narang, Newsha Ardalani, Gregory Diamos, Heewoo Jun, Hassan Kianine- jad, Md. Mostofa Ali Patwary, Yang Yang. and Yanqi Zhou. Deep learning scaling predictable, empirically. arXiv preprint arXiv: 1712.00409, 2017.
[HSK*24] Cheng- Ping Hsich, Simeng Sun, Samuel Kriman, Shantanu Acharya, Dima Rekesh, Fei Jia, and Boris Ginsburg. Ruler: What's 's the real context size your long-context language models? arXiv preprint arXiv: 2404.06654 2024.
[HTH*24] Shengding Hu, Yuge Tu, Xu Han, Chaoqun He. Ganqu Cui, Xiang Long. Zhi Zheng, Yewei Fang. Yuxiang Huang, Weilin Zhao, Xinrong Zhang, Zheng Leng Thai, Kaihuo Zhang, Chongyi Wang, Yuan Yao, Chenyang Zhao, Jic Zhou, Jie Cai, Zhongwu Zhai, Ning Ding, Chao Jia, Guoyang Zeng, Dahai Li, Zhiyuan Liu, and Maosong Sun. Minicpm: Unveiling the potential of small language models with scalable training strategies. arXiv preprini arXiv: 2404.06395, 2024.
[4T Connor Holmes, Masahiro Tanaka, Michael Wyatt, Ammar Ahmad Awan, Jeff Rasley, Samyam Rajbhandari, Reza Yazdani Aminabadi, Heyang Qin, Arash Bakhtiari, Lev Kurilenko, and Yuxiong He. Deepspeed- arXiv preprint arXiv: 2401.08671, 2024.
High-throughuut text generation for Ilms via mii and deepspecadinference.
[JBKM24] Samy Jelassi, David Brandfonbrener, S. Kakade, and Fran Malach. Repeat after me: Transformers are better than state space models at copying, International Conference on Machine Learning. 2024.
[KKL20] Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. arXiv preprint arXiv:2001. 04455, 2020.
[KI]*233] Woosuk Kwon, Zhuohan Li, Siyuan Zhuang, Ying Sheng, Lianmin Zheng, Cody Hao Yu, Joseph E. Gonzalez, Haotong Zhang, and I. Stoica. Efficient memory management for large language model serving with pagedattention. Symposium on Operating Systems Principles, 2023.
[KMH+20] Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B. Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. arXiv preprint arXiv: 2001,08361, 2020.
[LBOM12] Yann A. LeCun, Léon Botton, Genevieve B. Orr, and Klaus-Robert Müller. Efficient BackProp. pages 9-48. Springer Berlin Heidelberg, 2012.
[LCF++24] Tianle Li, Wei-Lin Chiang, Evan Frick, Lisa Dunlap, Tianhao Wu, Banghua Zhu, Joseph E. Gonzalez, and Ion Stoica. From crowdsourced data to high-quality benchmarks: Arena-hard and benchbuilder pipeline. arXiv preprint arXiv: 11939, 2024.
[LEV44]
KENNETH LEVENBERG. A method for the solution of certain non-linear problems in least squares. Quarterly of Applied Mathematics, 242):164-168, 1944.
[LH18]
Tlya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In International Conference on Learning Representations, 2018
[LLB*22] Opher Licher, Barak Lenz, Hofit Bata, Gal Cohen, Jhonathan Osin, Itay Dalmedigos, Erez Safahi, Shaked Meirom, Yonatan Belinkov, Shai Shalev-Shward Omri Abend, Raz Alon, Tomer Asida, Amir Bergman, Roman Glozman, Michael Gokhman, Avashalom Manevich, Nir Ratner, Noam Rozen, Erez Shwartz. Mor Zusman, and Yoav Shoham. Jamba: hybrid transformer-mante language model. arxiv preprint arXiv: 2403.19887, 2024.
[LLSS24]
Wonheom Lee, Jungi Lee, Junghwan Seo, and Jaewoong Sim. InfiniGen: Efficient generative inference of large language models with dynamic KV cache management. 18th USENIX Symposium on Operating Systems Design and Implementation (OSDI pages 155-172, Santa Clara, CA. July 2024. USENIX Association.
[LZH+25] Houyi Li, Wenzhen Zheng, Jingcheng Hu, Qiufeng Wang. Hanshan Zhang, Zili Wang, Shijie Xuyang, Yuantao Fan, Shuigeng Zhou, Xiangyu Zhang, and Daxin Jiang. Predictable scale: Part i optimal hyperparameter scaling law in large language model pretraining. arXiv preprint arXiv: 2503.04715, 2025.
[MA+*25 Microsoft, AbdeIrahman Abouelenin, Atabak Ashfaq, Adam Atkinson, Hany Awadalla, Nguyen Bach, Jianmin Bao, Alon Benhaim. Martin Cai, Vishrav Chaudhary, Congcong Chen, Dong Chen, Dongdong Chen, Junkun Chen, Weizhu Chen, Yen-Chur Chen, Yi ling Chen, Qi Dai, Xiyang
13

---

Dai, Ruchao Fan, Mei Gao. Min Gao, Amit Garg, Abhishek Goswami. Junheng Hao, Amr Hendy, Yuxuan Hu, Xin Jin, Mahmoud Khademi. Dongwoo Kim. Young Jin Kim. Gina Lee, Jinyu Li, Yunsheng Li, Chen Liang, Xihui Lin, Zeqi Lin, Mengchen Liu, Yang Liu, Gilsimia Lopez, Chong Luo, Piyush Madan, Vadim Mazalov, Arindam Mitra, Ali Mousavi, Anh Nguyen, Jing Pan, Daniel Perez-Becker, Jacob Platin, Thomas Portet, Kai Qiu, Bo Ren, Liliang Ren, Sambuddha Roy, Ning Shang. Yclong Shen, Saksham Singhal, Subhojit Som, Xia Song. Tetyana Sych, Praneetha Vaddamanu, Shuohang Wang. Yiming Wang. Zhenghao Wang. Haibin Wu, Haoran Xu, Weijian Xu, Yifan Yang. Ziyi Yang, Donghan Yu, Ishmam Zabir, Jianwen Zhang, Lyna Zhang, Yunan Zhang, and Xiren Zhou. Phi-4-mini technical report: Compact yet powerful multimodal language models via mixture-of loras, 2025.
[Mar63]
Donald W. Marquardt. An algorithm for least-squares estimation of nonlinear parameters. Journal of the Society for Industrial and Applied Mathematics, 11(2):331-44) 1963.
[Met24]
MctaAI. Introducing meta Ilama 3: The most capable openly available Ilm to date, 2024, URL: https: //ai.aaaaaaaa meta. com/ com/blog/ blog/meta meta- 11ama- 3/.
[MFG24]
Tsendsuren Munkhdalai, Manaal Faruqui, and Siddharth Gopal. Leave no context behind: Efficient infinite context transformers with infini attentitttiio. arXiv preprint arXiv:2404.0714: 07714,, 101, 2024.
[Min25]
MiniMax. Minimax-01: Scaling foundation models with lightning attention. arXiv preprint arXiv: 2501.08313, 2025.
[MKAT18] Sam McCandlish, Jared Kaplan, Dario Amodei, and OpenAl Dota Team, An empirical model of large-batch training. arXiv preprint arXiv: 1812.06162, 2018.
[MLPA22]
Sadhika Malladi, Kaifeng Lyu, Abhishek Panigrahi, and Sanjeev Arora. the SDEs and scaling rules for adaptive gradient algorithms. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems, 2022.
[MPS24]
William Merrill, Jackson Petty, and Ashish Sabharwal. The illusion of state in state-space models. arXiv preprint arXiv: 2404.08817 2024.
[MXBS16] Stephen Merity, Caiming Xiong. James Bradbury, and Richard Socher. Pointer sentinel mixture models. arXiv preprint arXiv: 1609: 060400784, 2016.
[MYX+24] Xuezhe Ma, Xiaomeng Yang, Wenhan Xiong, Beidi Chen, Lili Yu, Hao Zhang, Jonathan May, Luke Zettlemoyer, Omer Levy, and Chunting Zhou, Megalodon: Efficient Ilm pretraining and inference with unlimited context length. Advances in Neural Information Processing Systems, 37:71831-71854, 2024.
[MZK* 22] Xuezhe Ma, Chunting Zhou, Xiang Kong, Junxian He, Liangke Gui, Graham Neubig, Jonathan May. and Luke Zeulemoyer. Mega: Moving average equipped gated attention. arXiv preprint arXiv:2209. 10655, 2022.
[0pe]
Open Thoughts OpenThinker- 7B. Accessed: 2025-06-18.
Openthinker-7b.
https://huggingtee htt.
[Ope24]
OpenAL. Openai ol system card. arXiv preprint arXiv: 2412.16720, 2024.
[PAA+23]
Bo Peng, Eric Alcaide, Quentin G. Anthony, Alon Albalak, Samuel Arcadinho, Stella Biderman, Huanqi Cao, Xin Cheng, Michael Chung, Mattco Grella, G. Kranthikiran, Xuming He, Haowen Hou, Przemyslaw Kazienko, Jan Kocoñ, Jiaming Kong, Bartlomicj Koptyra, Hayden Lau, Krishna Ipsit Mantri, Ferdinand Mom, Atsushi Saito, Xiangru Tang. Bolun Wang. J. S. Wind, Stansilaw Wozniak, Ruichong Zhang, Zhenyuan Zhang, Qihang Zhao, P. Zhou, Jian Zhu, and Rui Zhu. Rwkv: Reinventing rnns for the transformer era. Conference Empirical Methods in Natural Language Processing, 2023.
[2F]
Jonathan Pilault, Mahan Fathi, Orhan Firat, Chris Pal. Pierre-Luc Bacon, and Ross Goroshin. slock-state transformers. Advances in Neural Information Processing Systems. 36:7311-7329 2023.
[PKI. 16]
Denis Paperno, Germán Kruszewski, Angeliki Lazaridou, Q. N. Pham, R. Bernardi, Sandro Pezzelle, Marco Baroni, Gemma Boleda, and R. Fernández. The lambada dataset: Word prediction requiring a broad discourse context. Annual Meeting of the Association for Computational Linguistics, 2016.
IQHS* 22]
Zhen Qin, Xiaodong Han, Weixuan Sun, Dongxu Li, Lingpeng Kong, Nick Barnes, and Yiran Zhong. The devil in linear transformer. Conference on Empirical Methods in Natural Language Processing, 2022.
[RHS* + 23]
David Rein, Betty Li Hou, Asa Cooper Stickland, Jackson Petty, Richard Yuanzhe Pang. Julien Dirani, Julian Michael, and Samuel R. Bowman. Gpqa: A graduate-level google-proof q&a benchmark. arXiv preprint arXiv: 2311.12022, 2023.
14

---

[RLL 25] Liliang Ren, Yang Liu, Yadong Lu, Yelong Shen, Chen Liang, and Weizhu Chen. Samba: Simple hybrid state space models for efficient unlimited context language modeling. In The Thirteenth International Conference on Learning Representations, 2025.
[RLW*23] Liliang Ren, Yang Liu, Shuohang Wang. Yichong Xu, Chenguang Zhu, and ChangXiang Zhai. Sparse modular activation for efficient sequence modeling. NEURIPS, 2023.
[RWC* 19] Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. Language models are unsupervised multitask learners, arXiv preprint, 2019.
[SAKM*23] Daria Soboleva, Faisal Al-Khateeb, Robert Myers. Jacob R Steeves, Joel Hes- tness, and Nolan Dey. Slimpajama: A 627b token cleaned and dedupli- cated version of redpajama, 2023. URL: slimpajama a- a-622b-tooo worrb-token-oee cleaned- and and-deduplicated wwwwsve-oisiieeeeeeeeeeeeeeeeeee
net/blog/
[SBBC21] Keisuke Sakaguchi, Ronan Le Bras, Chandra Bhagavatula, and Yejin Choi. Winogrande: An adversarial winograd schema challenge at scale. Communications of the ACM, (499)999116.. 2021.
[SD Yutao Sun, Li Dong, Shaohan Huang, Shuming Ma. Yuqing Xia, Jilong Xue, Jianyong Wang. and Furu Wei. Retentive network: A successor transformer for large language models. arXiv preprint arXiv: 08621, 2023.
[SDZ*24] Yutao Sun, Li Dong. Yi Zhu, Shaohan Huang, Wenhui Wang, Shuming Ma, Quanlu Zhang. Jianyong Wang, and Furu Wei. You only cache once: Decoder decoder architectures for language models. Neural Information Processing Systems, 2024.
[Sha19]
Noam Shazeer. Fast transformer decoding: One write-head is all you need. arXiv preprint arXiv: 1911.02150. 2019.
[Sha20]
Noam Shazeer. Glu variants improve transformer. arXiv preprint arXiv: 2002.05202, 2020.
[SLP*21] Jianlin Su, Yu Lu, Shengfeng Pan, Ahmed Murtadha, Bo Wen, and Yunfeng Liu. Roformer:
Enhanced transformer with rotary position embedding. arXiv preprint arXiv: 2104.09864, 2021.
[5W Xian Shuai, Yiding Wang, Yimeng Wu, Xin Jiang, and Xiaozhe Ren. Scaling law for language models training considering batch size. arXiv preprint arXiv: 2412,01508 2024.
[T] Keyu Tian, Yi Jiang, Zehuan Yuan, Bingyue Peng, and Liwei Wang. Visual autoregressive modeling: Scalable image generation via next-scale prediction. Neural Information Processing Systems, 2024.
(TMS* + 23]
Hugo Touvron, Louis Martin, Kevin Stone, Peter Alhert, Amjad Almahairi, Yasmine Babaci, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, Dan Bikel, Lukas Blecher, Cristian Canton Ferrer, Moya Chen, Guillem Cucurull, David Esiobu. Jude Fernandes, Jeremy Fu, Wenyin Fu. Brian Fuller. Cynthia Gao, Vedanuj Goswami. Naman Goyal, Anthony Hartshorn, Saghar Hosseini. Rui Hou, Hakan Inan, Marcin Kardas, Viktor Kerkez. Madian Khabsa, Isabel Kloumann, Artem Korenev, Punit Singh Koura, Maric-Anne Lachaux, Thibaut Lavril, Jenya Lee, Diana Liskovich, Yinghai Lu, Yuning Mao, Xavier Martinet, Todor Mihaylov, Pushkar Mishra, Igor Molybog, Yixin Nie, Andrew Poulton, Jeremy Reizenstein, Rashi Rungta, Kalyan Saladi, Alan Schelten, Ruan Silva, Eric Michael Smith, Ranjan Subramanian, Xiaoqing Ellen Tan, Binh Tang. Ross Taylor, Adina Williams, Jian Xiang Kuan, Puxin Xu, Zheng Yan, Iliyan Zarov, Yuchen Zhang, Angela Fan, Melanie Kambadur, Sharan Narang, Aurelien Rodriguez, Robert Stojnic, Sergey Edunov, and Thomas Scialom. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint arXiv: 2307,09288, 2023.
[VSP+ 17] Ashish Vaswani, Noam M. Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. NIPS, 2017.
[WA24]
Xi Wang and Laurence Aitchison. How to set adamw's weight decay as you scale model and dataset size. arXiv preprint arXiv: 14055,3698, 2024.
[WBR+24] Roger Waleffe, Wonmin Byeon, Duncan Riach, Brandon Norick, Vijay Korthikanti, Tri Dao, Albert Gu, Ali Hatamizadeh, Sudhakar Singh, Deepak Narayanan, Garvit Kulshreshtha, Vartika Singh. Jared Casper, Jan Kautz, Mohammad Shoeybi, and Bryan Catanzaro. An empirical study of
language models. arXiv preprini arXiv: 2406.07887 2024.
[WDL24] Kaiyue Wen, Xingyu Dang. and Kaifeng Lyu. Rans are not transformers (yet): The key bottleneck on in-context retrieval. arXiv preprint arXiv: 2402.18510 2024.
[WLX*24] Mitchell Wortsman, Peter J Liu, Lechao Xiao, Katie E Evereu,, Alexander A Alemi, Ben Adlam, John D Co- Co-Reyes, Izzeddin Gur, Abhishek Kumar, Roman Novak, Jeffrey Pennington, Jascha Sohl-Dit Dickstein, Kelvin Xu, Jachoon Lee, Justin Gilmer, and Simon Kornblith. Scall-scale proxies for large-scale transformer training instabilities. In The Twelfth International Conference on Learning Representations, 2024.
15

---

[WMZ+24] Yubo Wang, Xueguang Ma, Ge Zhang, Yuansheng Ni. Abhranil Chandra. Shiguang Guo, Weiming Ren, Aaran Arulraj. Xuan He, Ziyan Jiang, Tianle Li, Max Ku, Kai Wang, Alex Zhuang. Rongqi Fan, Xiang Yue, and Wenhu Chen. Mmlu-pro: A more robust and challenging multi-task language understanding benchmark. arXiv preprint Xrviv: 2406.01574, 2024.
[WRHS22] Yuhuai Wu, Markus N. Rabe, DeLesley S. Hutchins, and Christian Szegedy. Memorizing trans- formers. International Conference On Learning Representations, 2022.
[WT24]
Haoyi Wu and Kewei Tu. Layer-condensed kv cache for efficient inference of large language models. In Proceedings of the 62nd Annual Meeting of the Association for Computational 2024.
Linguistics (Volume 1: Long Papers), pages
[WWS+22] Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, E. Chi, F. Xia, Quoc Le, and Denny Zhou. Chain-of-thought prompting elicits reasoning large language models. Neural Information Processing Systems, 2022.
[WWZ*24] Jialong Wu, Zhenglin Wang, Linhai Zhang, Yilong Lai, Yulan Hc, and Deyu Zhou. Scope: Opti- mizing vay-v cache compression in long-context generation. arXiv preprint arXiv:2412.13649 2024.
[XMW+24] 24] Mingyu Xu, Xin Men, Bingning Wang, Qingyu Zhang. Hongyu Lin, Xianpei Han, and Weipeng Chen. Base of rope bounds context length. In Amir Globersons, Lester Mackey, Danielle Belgrave, Angela Fan, Ulrich Paquet, Jakub M. Tomezak, and Cheng Zhang, editors, Advances in Neural Information Processing Systems 38: Annual Conference on Neural Information Processing Systems 2024, NeurlPS 2024, Vancouver, BC, Canada, December 10- 15, 2024, 2024.
[XPA+25]
Haoran Xu, Baolin Peng, Hany Awadalla, Dongdong Chen, Yen-Chun Chen, Mei Gao, Young Jin Kim, Yunsheng Li, Liliang Ren, Yelong Shen, Shuohang Wang, Weijian Xu, Jianfeng Gao, and Weizhu Chen. -hi--onnnn--easooinng: Exploring the limits of small reasoning language models in math. arXiv preprint arXiv: 2504.21233, 2025.
[XYH+20] 20] Ruibin Xiong, Yunchang Yang, Di He, Kai Zheng, Shuxin Zheng, Chen Xing, Huishuai Zhang, Yanyan Lan, Liwei Wang. and Yan Liu. On layer normalization in the transformer architecture. In Proceedings of the 37th International Conference on Machine Learning, ICML 13-18 July 2020, Virtual Event. volume 119 of Proceedings of Machine Learning Research, pages 10524-100 10533, PMLR, 2020.
[YC Yu Yan, Jiusheng Chen, Weizhen Qi. Nikhil Bhendawade, Yeyun Gong, Nan Duan, and Ruofei Zhang. El-attention: Memory efficient lossless attention for generation. In International Confer- ence on Machine Learning, pages 11648 10--1166 11658. PMLR, 2021.
[YDX 24] Tianzhu Ye, Li Dong, Yuqing Xia, Yutao Sun, Yi Zhu, Gao Huang, and Furu Wei. Differential transformer. arXiv preprint arXiv:2410.05255 2024.
[YHIB+22] Greg Yang. Edward J. Hu, Igor Babuschkin. Szymon Sidor, Xiaodong Liu, David Farhi, Nick Ryder, Jakub Pachocki, Weizhu Chen, and Jianfeng Gao. Tensor programs v: Tuning large neural networks via zcro-shot hyperparameter transfer. arXiv preprint arXiv: 2203.03466, 2022.
[YKH25] Songlin Yang. Jan Kautz, and Ali Hatamizadeh. Gated delta networks: Improving mamba2 with delta rule. In The Thirteenth International Conference on Learning Repreventations, 2025.
23] Songlin Yang, Bailin Wang, Yikang Shen, Rameswar Panda, and Yoon Kim. Gated linear attention transformers with hfrfiete-efficient training. arXiv preprint 2xXiv:3312.066:: 2023.
[YW7 Songlin Yang. Bailin Wang, Yu Zhang, Yikang Shen, and Yoon Kim. Parallelizing linear trans- formers with the delta rule over sequence length. Neural Information Processing Systems, 2024.
[YYZH23] Greg Yang. Dingli Yu, Chen Zhu, and Soufiane Hayou. Tensor programs vi: Feature learning in infinite-depth neural networks. Internationaa Conference on Learning Representations, 2023.
[7HBB + 19] Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. Hellaswag: Can machine really finish your sentence? Annual Meeting of the Association for Computational Linguistics, 2019.
[2LJ+22]
Simiao Zuo Xiaodong Liu, Jian Jiao, Denis Charles, Eren Manavoglu, Tuo Zhao, and Jianfeng Gao. Efficient long sequence modeling via state space augmented transformer. arXiv preprint arXiv: 2212.08136, 2022.
[ZS19]
Biao Zhang and Rico Sennrich. Root mean square layer normalization. Neural Information Processing Systems, 2019.
[5V Terry Yue Zhuo, Minh Chien Vu, Jenny Chim, Han Hu, Wenhao Yu, Ratnadira Widyasari, Imam Nur Bani Yusuf, Haolan Zhan, Junda He, Indraneil Paul, Simon Brunner, Chen Gong. James Hoang, Armel Randy Zebaze, Xiaoheng Hong. Wen- Ding Li, Jean Kaddour, Ming Xu, Zhihan Zhang, Prateek Yadav, and et al. Bigcodebench: Benchmarking code generation with diverse function calls and complex instructions. In The Thirieenth International Conference on Learning Representations, ICLR 2025, Singapore, April 24-28, 2025. OpenReview.net, 2025.
16

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification:
We explain method and summarize the contribution in introduction.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results be expected to generalize other settings.
is fine include aspirational goals as motivation as long is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification:
The limitation included conclusion section.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification. asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low images are taken in low lighting. Or speech-to-text system might be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach address problems of privacy and faimess.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question:
For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [NA]
17

---

Justification:
Our paper does not include theoretical results.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: This paper discloses the information needed to reproduce the main experimen- tal results.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm. (b) If the contribution is primarily
model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome deseribe the particular way they provide for reproducibility, In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
18

---

Answer: [Yes]
Justification: The datasets we used are partially open-source, and we will provide our code in the footnote of page one.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.. cc/
Sabllic/guides/CdeSSubiii oolooo for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).
The instructions should contain the exact command and environment needed to run reproduce the results. See the NeurlPS code and data submission guidelines (https: //n /nwps..cc////////////// ctips.cc/pullic/uut ggwgguuiss////// /Cooeeuuoiiseeeeeeee wwweeissssonniliieeoeooe
for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information possible in supplemental material (appended to the paper) recommended, but including URLs to data and code is permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: We have specified all the training and test details necessary to understand the results
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: We report average results of multiple runs in our experimental section with error bars.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results accompanied by error bars, confi- dence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
19

---

The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative emrr rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification:
We explain the computation resources in experiment section.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https: http///eerr.p..eeeee
Answer: [Yes]
Justification: Research is conducted in the paper conform with NeurlPS Code of Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. if there is a special consid- eration due to laws or regulations their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: Our paper is not highly related to societal impacts.
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
20

---

The conference expects that many papers will be foundational research and not tied particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility of MI..
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: Our paper poses no such risks. Our work does not release a new model.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Lieenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: CC-BY 4.0, and referenced the works that we used to implement our code.
Guidelines:
The answer NA means that the paper does not existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (c.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
21

---

If this information is not available online, the authors are encouraged to reach out to the asset's creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification: We will provide our code in the footnote of page one.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/mmoel as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: This paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included the main paper.
According the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: This paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such the institution conducting the review.
16. Declaration of LLM usage
22

---

Does the paper describe the usage of LLMs if it is important, original, or non-standard component of the core methods this research? Note that if the LLM used only for writing, editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, or originality of the research, declaration is not required.
Question:
Answer: [No].
Justification:
LLM is used only for writing, editing, formatting purposes.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs any important, original, non-standard components.
Please refer to our LLM policy (https htt////eerppp.eee /ttt:p:/eearrpp..ccooofffffeee / for what should or should not be described.
23

---

A Background
YOCO (You Only Cache Once) [SDZ 24] is an inference-efffcient decoder-decoder architecture with linear pre-filling complexity. It comprises a self-decoder, formed by the first half layers, which employ token mixers with linear computational complexity and include final full-attention layer that generates Key- Value (KV) caches during pre-filling. And the second half of the layers forms the cross-decoder, which are cross-attention layers that attend to the KV caches produced by the self-decoder's ' last full attention layer. Specifically, for an input sequence of hidden states X Xmem E Raxdmm of the final full-attention layer in the self-decoder, it pre-computes and caches sequence KV pairs during pre-filling:
Kg = mem K. Ve -xmmm Wv,
where Ke Ve E Rnn are the cached matrices and WK. Wv are weight matrices. Subsequently, at the decoding stage, every cross-attention layer I in the cross-decoder reuses this single set of KV cache. Given the input hidden state x((-1) to that layer, the cross-attention output is calculated by generating new query, Q and attending to the shared cache:
2 -x(1I) @
Ceross
Qcfose K-T ddk
H()) softmax
Ve.
During pre-filling, this approach (1) entirely avoids the computation of full attention and (2) requires inference through only the first half of the layers, substantially reducing the computational cost of processing prompts for both short and long contexts. Our SambaY architecture further improves its efficiency by modifying the cross-decoder: we replace half its memory 1/0-expensive cross-attention layers with lightweight Gated Memory Units (GMUs), thereby enhancing decoding efficiency during response generation.
B Additional Theoretical Analysis
Normalization placement in linear attention. In linear attention architectures [QHS+22, SDH*23, YWS+23], including Gated DeltaNet (GDN) [YKH25], a normalization operator is applied immediately after token mixing to stabilize training at layer l':
M() = Norm (("),
(()) = (M o ?
G(")= (W x(")),
where Norm is typically RMSNorm (ZS19) and x(')) the layer input and y(')) is the layer output. Placing Norm before the output gating, however, breaks the associativity between the gating matrix G(") and the token-mixing operator A("), the Gated Memory Unit (GMU)
longer directly re-weight the token mixing based the current layer input with X(), To resolve this issue, we propose to postpone the Norm after output gating for GDN (denoted as GDN-A), following the design of Mamba-2 [DG24]. Concretely, for layer l' we instead compute
M = A(')O), y(O)) = Norm(M o G w(())
and employ the normalized GMU (nGMU) at layer / >> l' in the cross-decoder to maintain training stability while allowing the gate to modulate 1l²) directly, i.e.
y Norm @ (()) W!O
This simple reordering preserves associative re-weighting and, as demonstrated empirically in Appendix H, substantially improves the long context performance when compared to the original normalization before gating design in GDN.
24

---

C Additional Aspect Ratio Calculations
When solving for the aspect ratio through iso-parametric equations, we rounded it up to an even integer to guarantee the activation of Tensor Cores". Based on the Samba+ architecture, we can derive the iso- -so-parametric equation through calculating the number of non-embedding parameters as follows,
Vatta (d) = 2.5dw Wattn / + 2dw Watt/ /, Nmamba (d) = 6dw² /4,
(d) Naun (d) + Nmamba (d) Nuup (d) = 208ar + 13.5a d = 237568d
Solving for a, we get or 126. For the Samba SambaY+Da +DA architecture, the aspect ratio is not changed because the number of extra parameters introduced by DA is negligible. For MambaY, we have
Natin (d) = 2dw Watn/ /
Nmamhu (d) = 6dw² /2,
Ngma = 4dw? /4,
N(d) = Nau (d) + Nmaanba(d) + Nmip(d) (d) + Ng (d) = 64ad + 16aadd = 237568d
Solving for a, we get 0rg 120. For
we have
Naum (d) 2.5dw Watn/ / + 2dw Wann// /4, N Naabba (d) 6dw///4 N Nenu (d) Sdw2 /4,
(d)) Natm(() + Mammba (d) + Nmtp(d) N Ng (d) 144ad s 15.5022 237568d
Solving for a, we get 044 2 120. For SambaY-ANe Y-Attn, we have
Namm (d) = 2.5dw Watn/ /4 + 2dw Watn/4
Nmumba (d) = 6dw² /4, Ngmu (d) = 2dw Waua /4,
(d) = Naan(d) (d) + Nmaabaa (d) + Nmtp(d) + Nggmu (d) 208ad + 13.5a & = 237568d
Solving for a, we get az se 126, which is the same as Camba+YOCon For Yambta-AAtt--aa we similarly have
Natn (d) = 2.5dw Wam//4, Nmamba (d) = 6dwp /4, N Ngmu (d) = 2dw Watn /2
(d)) = Nun + N Nmambaad (d) + (mtp(d) + Ngmu (d) 208ad + 13. = 237568d
Solving for a, we get Og zs 126. For the GDNY architecture, we use a fixed head dimension 256 and SwiGLU output gating with the Gated DeltaNet layers. The query-key projections with 0.75 expansion ratio are used, while the value, gating, and output projections are using a 1.5 expansion ratio with respect to the model width. We also allow negative eigenvalues for improved expressiveness of the transition matrices [GSF+25] 25). Specifically, we have
N Naun(d (d) 2dw Watn/ /4, NGDN (d) 6dw² /2,
Ngmmu (d) 3dw2 /4.
(dd = Nam (d) + NGDN(d) + -Nipp(( (d) + Ngmu (d) = 64ad + 15.75a 2d = 237568d
Solving for a, we get ar 120. As in SambaY, we can similarly solve for S-GDNY to get 0g 126 and for SWA+YOCO to get Org 130.
D Implementation Details
the AdamW optimizer with 3, = 0.9, 82 = 0.95. 10-8 and weight decay of 0.1. A learning rate schedule applied with 1B warm-up tokens linearly increasing to the peak learning rate y,, followed by linear decay to zero. We use LeCun uniform initialization (i.e. Py Py'oorch default initialization) [LBOMI2] for the weight matrices following [GD23] and [RLL* 25], and tie the input and output embedding matrices which are initialized from the normal distribution (0, 0.02²). The attention logits scaler is set to 1/ /dz where d is the head dimension. We summarize the key differences between ,P, pPP++ and Standard Parameterization (SP) in Table 6, with additional details as follows. For pP++, we scale the output logits and the learning rate of matrix-like parameters proportional to 1/w. The output of each layer is divided by 22d following Depth-uu..
Details on scaling comparisons. Except for the learning rate, we other
For SP, we don't 't apply any pP++ scaling laws, and since LeCun initialization already scales its initialization variance with respect to 1/diz as the same as proposed in μdP, where d is the fan-in
httppe//deveooper:: www.spppr.riiiiiiiii.aa nvidia. cewiiia..comballi www/bloogopttiiiziiiiiiiiio/oiii optimizing-iiiieie gpu- performance perforrennee-ensor-ccrrree
25

---

Table Key differences between jal, p and Standard Parameterization (SP). LR mult. denotes the per-parameter multiplier applied on top of the global learning-rate (1)), Res. mult. is the multiplier applied to the output of residual branches and WD denotes the weight decay. For μP and zero weight decay is also applied other scalar or vector-like parameters such RMSNorm 10 for untied embedding and 0.02 for tied embedding, and in both cases T = 0.02 and B = "fan_in in" means the input dimension weight matrices.
1 1/vd
weights. In this work,
Parameter
Scheme
LR mult.
Initialization
Res. mult.
Weight mult.
WD
SP
1 oc
N o²)
1
Embedding
JPP
N(0, of)
x
x 1 0
pP+
oc
x
SP
oc
or tied
x
x x
Unembedding
JPP
oc
0 or tied
x
1/w
pP+
o
Or tied
1/w
0
SP
< 1 1/tw
(O,
1
1
Hidden Weights
JP
u1
x
Tan in
pP+ x 1/w
21(
Vian
Vianj
1//2d
x
dimension of the weight matrix, instead normal initialization with standard deviation of 0.02 for weight matrices to rule out this confounding effect. We divide the initialization standard deviation of the output projection of each layer by 22d, following [RWC+ 19, GD23, RLL The detailed architecture and optimization setups for each of the scales are shown in Table 7. Following [GD23, YWZ 24, RLL* 25, YKH25), our downstream evaluations are conducted on the following benchmarks: Wikitext [MXBS16], LAMBADA (LMB) [PKL, + 16), Arc-Easy/Challenge (ARC-e/ARC-e) (CCE+18) 18), HellaSwag (Hella.) [ZHIB+ 19), WinoGrande (Wino.) [SBBC21] and PIQA [B]B+20], where measure character normalized accuracy (acc_n) for Arc-Challenge and HellaSwag.
More details on architecture and large-scale pre-training. We provide a comprehensive summary of the architectures explored in this work, along with the large-scale pre-training setup, in Table 7. In our architectures, Differential Attention uses a depth-dependent initialization factor, Aini 0.8 0.6 -0.3 0.3 I), where 1 is the depth index. For each attention head, it employs sets of learnable parameters, (Ag1, A&1) and (Ag2, As2). each of dimension equal to the head dimension and initialized with normal distribution of zero mean and 0.1 standard deviation. RMSNorm [ZS19] with learnable element-wise affine parameters is adopted for attention output normalization. For each of the intermediate layers, LayerNorm [BKH16] used with Pre-LN [XYH * 20] for architecture.
-lash
Scones
3.45
0.55
Artin DP self-
decoder 520K
Aecrag
BF16
(a) Fused FP32 LM Head with Cross- Cross-Entropy Loss
(b) Adding Attention Dropout
Figure 6: Adopted tricks for mitigating the large-scale pre-training instability of Phi4-mini-Flash
Mitigation of instability in large scale pre-training. During the pre-training stage of Phi4-mini- Flash, meet severe loss divergence, which is mitigated with the following two tricks: (1)
26

---

cp-cast the weight and the input to FP32 during chunk-wise matrix multiplication for our fused linear cross-entropy loss kernel which is modified from the Liger-Kend "ggee-Kernl'.. As shown in Figure 6a, with the FP32 up-casting, the gradient norm during the training process can be stabilized to closely track the naive no-fusion baseline, compared to the normal BF16 matrix multiplications with up-shooting trends that will finally blow up the training loss. Without fusion, the training speed is substantially slower because of our large 200K vocabulary size. (2) As shown in Figure 6b, we add 0.05 attention dropout at 310K steps for self-decoder and at 520K steps for cross-decoder, which are the last checkpoint steps before the loss divergences happen. As we can see from the figure, adding attention dropout doesn't it harm downstream performance MMLU [HBB' 21], but successfully stabilize the whole training process until the end.
Table 7: Model and training configurations for the architectures explored in this work. TransformerLS adopts the same architecture as Transformer++, with Sliding Window Attention (SWA) applied to all attention layers except every fourth layer, which uses full attention. MLP Size denotes the intermediate dimension of the MLP. i.e. the input dimension of the output projection. Phi4-mini and Phi4-mini-Flast are trained with a batch size of 8M tokens, using linear learning rate schedule with 3,000 warm-up steps. We allow the intermediate dimension of attention layers be larger than the model width, so that the head dimension can be power of 2. Variants enhanced with Differential Attention adopt the same architectural configurations as their respective baselines. All models use tied embeddings. The 3.8B- --raammeeeeerr YambaY and Samba+ YOCO models are randomly initialized for benchmarking under the vLLM inference framework. Except for 3.8B- 3.8B-par aB-parameter models which use vocabulary of 200K tokens, apply Llama-2 [TMS + 23] tokenizer with 32K vocabulary for all other models.
Architecture
Depth
Model Width
Query Heads
KV
Head
MLP
Non-Embed Params (M)
Params (M)
Learning Rate
Training Tokens (B)
Heads
Dim
Size
Transformer++
1024
128
4096
121.6
154.4
5.66c-04
12.5
12
1536
12
3
128
6144
410.5
459.7
4.62e-04
42.2
2048
16
4
128
8192
973.1
1038.6
4.00e-04
100.0
20
2560
20
5
128
10240
1900.5
1982.5
3.58c-04
195.3
24
3072
24
6
128
12288
3284.1
3382.4
3.27e-04
337.5
SambaY
8
128
3968
123.3
155.0
5.66e-04
12.7
1488
12
3
128
5952
416.1
463.7
4.62c-04
42.8
1984
16
4
128
7936
986.3
1049.8
4.00c-44
101.4
2480
20
5
128
9920
1926.5
2005.8
3.58e-04
198.0
24
2976
24
6
128
11904
3328.9
3424.2
3.27c-04
342.1
Samba+YOCO
1008
8
128
4032
123.2
155.4
5.66e-04
1512
12
128
6048
415.6
464.0
4.62e-04
42.7
16
128
8064
985.2
1049.7
4.00e-04
101.2
2520
20
128
1924.3
2004.9
3.58e-04
197.8
3024
24
6
128
12096
3325.1
3421.9
3.27e-04
341.7
SWA+YOCO
2080
16
4
128
8320
984.0
1050.6
4.00c-04
40.0
Mamba
1920
16
128
7680
975.2
1036.6
4.00c-04
40.0
GDNY
1920
16
128
7680
960.4
1021.9
4.00c-04
40.0
S-GDNNY
2016
16
128
7680
1001.0
1065.5
4.00c-44
40.0
MambaY-2
1920
16
128
7680
1036.6
4.00e-04
40.0
SambaY-2
1984
16
128
7936
986.3
1049.8
4.00c-04
40.0
SambaY-A
2016
16
128
8064
985.2
1049.7
4.00c-04
40.0
Sam
2016
16
4
128
8064
985.2
1049.7
4.00c-04
40.0
SambaY-MLP
16
1920
16
4
128
7680
985.0
1046.4
4.00e-04
40.0
Phi4-mini
3072
24
128
8192
3221.2
3835.8
5.00ee44
5000
Pih4-mini-Flass
2560
40
20
10240
3329.2
3841.4
5.00e-04
5000
SambaY
2560
20
10240
3329.2
3841.4
Samba+YOCh
2560
40
20
10240
3224.4
3736.5
E
Ablation Study on Hyper-parameter Scaling Laws
We conduct a comprehensive ablation study of our pP++ scaling laws to validate their scaling behavior. All experiments performed using Transformer++ trained with 4K sequence length
27

---

JP-a
2.00
LAW37 SOLLO
51501 (4N2,50, A1.C
EJR boul hoss 1911
1.00
EJ bualing meEs 30,Cm
10 doddion Losp
1.30
1.0
159
1.34
2x-30" Tisding Boons IBilooss
3:30
6x100
10
3w 1=10"
6x 10
Taining Tikons (588ons
(a) Tied Embedding
(b) Untied Emhedding
Figure 7: Validation Loss v.s. Training Tokens on the SlimPajama dataset for Transformer++ trained with tied (left) or untied (right) embedding layers. For the training on 600B tokens with jPP, the model encountered NaN losses after 204K gradient update steps. We report the last valid validation loss prior to divergence as its final performance.
on the SlimPajama dataset. To ensure that the linear learning rate schedule fully decays to zero, we train six models at different training token budgets: (100B, 200B, 600B) for each of the scaling curves. We examine the scaling performance under both tied and untied embedding setups. For the untied setting, we follow RWKV [PAA + 23] by applying normal initialization with zero mean and standard deviation of 10 10-4, The unembedding layer is initialized to zero, following the zero-out trick proposed in μP [YHB* 22]. As shown in Figure 7a, we observe that the original μP setup (which uses LeCun initialization and does not include Depth-uP or weight decay modifications as in pP++) can lead to severe training instability when scaling 600B tokens. Since observe increasing gradient norms with large spikes for vector-like parameters shortly before the model diverges, this highlights the importance of the μP++ strategy of applying zero weight decay to vector-like parameters to enhance training stability at large scales. We also explore batch size scaling with respect to training token size, following [SWW+2] 24, L i.e.
T B = Boy To
As in Figure 7a, μP++ (Batch Scaling) shows both worse learning efficiency and irreducible loss than uP++. Generally, we think the batch size mainly affects parallelism and the computation efficiency as long as the batch size is not too large, and the worse scaling behavior can be because (1) when scaling up, the batch size can surpass the critical batch size [MKAT18] 8), which leads to worse model performance, (2) other optimizer hyper-paramett parameters are not adjusted accordingly with batch size as in [MLPA22] and we leave it for future works to study the large batch size training with pP++. We also try using Normal Initialization with 0.02 standard deviation for the weight matrices, and scale the variance with respect to 1/d. However, μPP++ (Normal Init.) shows worse scaling than μP++, indicating that it is better to adjust the initialization multipliers based as adopted by LeCun initialization, rather than a global factor related to model width. We explore integrating the empirical scaling law of the learning rate 7) scaling with respect to training tokens T [BBC+25] to μP++, ie.,
each matrix' dimension
Bd Bod
To T
1o
and adjust weight decay to maintain the same regularization effect across different training tokens with the setup of Independent Weight Decay [WLX +21x,244,,40 i.e.
Ao 71
A
where À is the weight decay in AdamW [LH18] and A 0.1. We denote this scaling law as μP++ (LR scaling + Indep. WD). As in Figure 7b, while the irreducible loss is comparable, we observe a worse learning efficiency with smaller b compared to pP++. We think that future work is needed to
28

---

have an empirical study of the learning rate scaling with respect to dataset size under μPP+,, instead of transferring the empirical law directly to our theoretical laws. We also explore using the WSD [HTH+24] learning rate scheduler for jPP++ where we set the final decay period to be 2/7 of the total period following [DA24b]. Unfortunately, depicts worse scaling behavior than pP+++ with linear learning rate schedule, as shown in Figure 7b. Interestingly, when comparing the performance of juP++ with tied versus untied embeddings, we observe that μP++ with untied embeddings achieves significantly lower validation loss with 100B training tokens, but its irreducible loss remains comparable to that tied embeddings. This suggests that the additional parameters from untied embeddings primarily accelerate training convergence without improving the final model performance a sufficient amount of data is given.
F Additional Long-Conter Retrieval Experiments
Long-conter, extrapolation with NoPE. In Table 8, directly measure the retrieval accuracy at 32K, 64K and 128K context length on the Phonebook benchmark for IB parameter models trained with 32K sequence length in Section 3.2. We see SambaY and its variants with NoPE can extrapolate their retrieval ability by 2x zero-shot, while RoPE-based models (Transformer+ and TransformerLS) have substantial drop beyond 32K. We leave the explanations of why NoPE can enable limited extrapolations on retrieval tasks as an interesting future work.
Table 8: Long-context extrapolation accuracy (with standard deviations) on the Phonebook benchmark. The models are trained on the Prol.ong-6dd 64k dataset with 32K sequence length and 1B parameters.
Model
SWA Size
32K Acc. (%)
64K Acc. (%)
128K Acc. (%)
Transformer++
60.94 10.00
0.00 0.00
0.00 0.00
TransformerLSS
256
60.16 + 7.12
17.19 5.63
0.78 = 1.35
Samba+YOCO
1024
82.81 = 6.44
67.97 11.13
20.31 = 8.41
Samba'Y
256
92.19 = 1.56
96.09 2.59
0.00 0.00
SambaY Y+DA
512
96.09 = 1.35
84.38 3.12
5.47 2.59
1.0
1.
Tranfummrrs
Sambor
Tranaformarts (Tott Sees 204) Sanberr
0.0
0.6
6couma
Accunecy
0.4
Tamifome Sarite-Y6co
Sarbba Sembatiiu
0.0
128
512
1024
2042
0.0
125
236
512
1024
2046
(a) SlimPajama
(b) ProLong-644%
Figure 8: Accuracy (with error bars) v.s. Sliding Window Size on Phonebook with 32K evaluation length using 40B training tokens from SlimPajama (left) ProLong-64K (right). As ablation to Figure 3, variable-length training is not applied for both setting.
Ablation on training data and methodologies. Figure 8 illustrates how different model architec- tures perform on the Phonebook long-contex, task as the sliding window size increases, using either for pre-training with 32K sequence length and without variable-length training. Specifically, we concatenate the data samples with EOS tokens as separation to form 32K length training sequences. On SlimPajama, overall accuracy is modest, with Samba showing some initial promise at smaller window sizes (peaking at 128) before declining, while Yamba+YOCO performs best a moderate window size of 512. Transformer-basee models generally struggle to achieve competitive accuracy across window sizes. Notably, reducing RoPE base from 640K to
SlimPajama or ProLong-
29

---

10k for TransformerLS significantly harms the performance across window sizes. Switching to the ProLong- dataset leads to a notable performance boost across all architectures compared to SlimPajama, even without variable-length training. We can observe that SSM-based models enjoy larger boosts on accuracies than Transformer++. This indicates that SSM-based models can learn to switch contexts between different data samples within the packed sequences more easily than pure at- tention models. Notably, S achieves competitive accuracy using a smaller sliding window (512), matching the performance of Samba+ at larger window sizes. While Samba+ continues to benefit from increasing window sizes, reaching peak accuracy at 2048, SDm demonstrates greater efficieney by achieving strong results with smaller sliding window size. Given that variable-length training on ProLong- 1ogg-64K generally yields better results as in Figure 3, these fixed-lengle training results indicate that while ProLong-64k benefits long-context performance, the full potential, especially for pure attention models that are sensitive to sliding window size (e.g. Transformerl can be further unlocked by training methodologies that explicitly account for varying sequence lengths of each data sample. The different optimal sliding window sizes and performance trajectories underscore that both the pre-training dataset and the training methodology significantly influence how effectively the training context length can be utilized for long-context pre-training,
G Additional Details on Efficiency and Reasoning Results
[XPA 25], the eval- uation conducted with a sampling temperature of 0.6, top-p [HBD* 19] value of 0.95, and maximum sequence length of 32,768 tokens. We leverage the Math-Verify library (version 0.7.0) and Lighteval (version 0. 0.10.0) to enable efficient and robust evalu- ation reasoning tasks. We prepend the instruction: "Please reason step by step, and put your final an- swer within A \booee (O)." for the evaluation on AIME24/25 and MATH500 and "Please reason step by step, and put your final choice of one letter from A/B/C/D within [())) for the evaluation on GPQA Diamond.
Following
Lateey 50 Whote Genestion
5000
10000
19300
20038
25000
35000
Oakene)
Figure 9: Generation latencies at length of IK, 2K, 4K. 8K, 16K and 32K for prompt length of 2000. Given a certain generation length, we measure the average latency of all the requests in all the loads of 1, 2, 4, 8, 16 concurrent requests.
The data-points of latency and throughput in Figure 4 measured by varying the concurrnnyy levels of the requests sent to the vLLM server, ranging from 21, 4, 8, 16). We also compare the latency of our Phi4-mini-Fi architecture with Phi4-mini across various generation lengths in Figure 9. We can see that our model can achieve faster generation speeds starting from a generation length as small as 4K. Apart from the sub-optimal Differential Attention implementation, our vL.MM implementation is also 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 in Figure 4 and Figure 9.
We conduct case studies for Phi4-mini- -aii-mmiii-F-------- Reasoning in Example 1 and Example 2 for both mathematical reasoning and daily-life reasoning. Notably, while our model is only trained on distillation data with math problems,
model exhibits advanced reasoning techniques such as double checking (e.g. "Alternatively maybe I could have completed the square instead?". "reaking it down:"), "), reflection (e.g. "That means there are no real solutions, right? ", "Wait, maybe I miscounted."", "), and backtracking (e.g. "But the problem might be expecting complex solutions.". "aait, maybe I'm overcomplicating.) ") in both cases. Surprisingly, in Example 2, our model will even guess the user intention (e.g. "But wait, sometimes people might confuse the number of R's in double letter. ") finally provide the correct answer.
"nccpe://githuub.com http::///hhhttccccc////gaa/aaaaa Veriiy
30

---

Example 1: Mathematical Reasoning (Part 1)
Query: solve x+++3x+==1
Model's reasoning process:
Think
Okay, let's see. need to solve the equation squared plus 3x plus equals 1. Hmm, first, side that the equation equals to solve quadratic equations. me subtract from both sides. That would give me x² + 3x + 4- = Simplifying that, 4 minus I is3, so the equation becomes x² 3x + 3 Alright, now I have quadratic equation in standard form, which is ax² + bx + c = 0. Here, is 1, b is and is 3. To solve this, can the quadratic formula, which is equals negative b plus or minus the square root b squared minus four all over Let me write that down: x = [-b /(22_.. 4ac)] (2a).
maybe should get all the terms
That usually
makes
Plugging in the values from my equation, is 3, so negative h would be -3. Then, squared is 3 squared, which is Thenn, 4ac is times times 3, which is 12. So the discriminant, which is the part under the square root, would be 9 12. That equals -3, Wait, the discriminant is negative. That means there are no real solutions, right? Because you can't of a negative number in the real number system.
square root
But the problem might be expecting complex solutions. Let me check. If complex numbers are allowed, then the square root of -3 can be written times the square root of 3,where is the imaginary unit. So, substituting back into the quadratic formula, we have x equals [-3 v(-3) (2*1). Which simplifies to [-3 i/33/2.
Let me double-check my steps to make sure didn't 't make a mistake. Starting equation: x² 3x 1. Subtract 1: x2 + 3x + O. Coefficients a=1, b=3, c=3, Discriminant: 3201:332_222 13.4 =2= 3. Yep, that's right. So the solutions complex: (-3 # 133322 Alternatively, maybe could
me try that method completed the square instead??
verify. Starting with x² + 3x = Move the constant term to the other side: x² = To complete the square, take half of the coefficient of x, which is 3/2, square to get (9/4). Add that to both sides: x²+ 3x 9/4 -3 + 9/4. The left side hecomes (x 3/2/2. The right side: -3 is -124, -1/4 9/4 is -3/4. So + = Taking the square of both sides gives x 3/2 = = = /3/4) = = /3M2. Subtract from both sides: = -322 + (1/3)/2. Which is the same as (-3 # 1/33/2. So that matches the quadratic formula result. Therefore. both methods lead to the same complex solutions. Since the original equation has negative discriminant, there expressed terms complex numbers.
real solutions, only complex ones. The answer should be
Just be thorough, let me plug one of the solutions back into the original equation to verify. Let's take x [-3 /332/2 Plugging into x² + 3x + 4. First, compute x²: [(-3 + / Let's expand that. (-3 i2 (-3) 2*(-3)(( (i/3)22 =9 ((2(3))))))))))))))))) 6i.0-ii63331 (-1)333 6i/33/4 (3 31//3//2. Next, compute 3x: 3*[(-3 +1/3)/1+332(0 33/3)/2
/i = /3. Then divide by
Add x 3x 4: (3- (3 (-9 3i.33// 4. Combine the fractions: 3 -9 3//33/2
4= Which matches the right side of the original equation. So the solution checks out.
Similarly, if plug in the other solution x [-3 ii/3//2 the process would be the same, and the imaginary parts would cancel out well, resulting in 1. Therefore, the solutions think that's 's all. The key steps were moving recognizing it's negative, and proceeding with complex solutions. Both quadratic formula and completing the square worked, and verifying by substitution confirmed the answers are correct.
correct. So the final answer one side, calculating discriminant,
complex numbers: (-3 1 i
To be continued...
31

---

Example 1: Mathematical Reasoning (Part 2)
Model's Response (continued): To solve the equation x² + + 4
1. **Rarrange the equation** * to standard quadratic form:
x² + + - 0 = x² +x
**Identify coefficients * a 1, 3, 3.
3. **Calculate the discriminant* *#
A = b² 4ac 3 4(1) = 9 - 12 -3
Since the discriminant is negative, solutions are complex. 4. ** Apply the quadratic formula** *::
b VA
3
-3+
x
2a
2
**Verify via completing the square**;
3
2
3 4
x² 3x
2
Solving gives + v confirming the solutions.
v3
**Final :: The solutions are
+
2
2
Example 2: Letter Counting in Misspelled Words (Part 1)
Query:
how many r's "strawherrry"?
Model's reasoning process:
Think
Okay, let's see. The question asking many times the letter "r" appears in the word "strawberry".. Hmm, first I need to make sure spell the word correctly. The user wrote "strawberryy.. . Let me check that. Normally, the plural of "strawberry" is "strawberries" but maybe they made typo spelling mistake. Anyway, the word given here is "strawberry".. Let count the letters one by one.
Breaking it down: s-t-r-a-w-be Fr-a-w-b-e-r-r-yy- Let's write each letter separately:
S. T. R, A, W. B, E, R, R, R, Y.
Wait, so starting from the beginning: S (1). T (2), R (3), A (4), W (5), B (6), (7). R (8), R (9), R (10), Y (11). So total letters
11. Now, how many times does R appearhere?
Let me count again to sure. S-T-R-A-WY W-B-E-R-BE R--E-R-R-RR R-R------ R- So after E comes R. then another R, another R, and then Y. So that's three R's in row? Wait, no. Let's see: after E is the first R. then the second R. third R, and then Y. Wait, no, the word is "strawherrry". So after E, there three R's in row, right? Because it's written "berrry". So breaking down:
Looking at the letters: positions 3. 8. 9, 10. That's four
-t-r-a--w--e------ee So after E, there are three
before So that's three R's. Wait, but when listed them earlier, counted four R's. Let me check again.
Wait, maybe I miscounted, Let's write out each letter with their positions:
1. S 2. T 3. R A W 6. B 7. 8. 9. R 10. R 11.
positions 3, are R's. That's four R's. But wait, the word is spelled as "strawberry".
To be continued...
32

---

Example 2: Letter Counting in Misspelled Words (Part 2)
Model's reasoning process (continued):
Think
So after "straw" and "b", , it's "erry". So "e", then three R's, then "y". So that's three R's. But according the breakdown above, there are four R's. Wait, maybe the user added an extra R mistake? Let check the exact spelling again. The wrote "strawberrry" So "berr". So "b", then "c", then three R's s, then "y". that's three R's. But in my positional list, have four R's. Wait, maybe made an error in counting.
Wait, let's write the word step by step:
s (I) (2) (3) a (4) (5) b (6) (7) (8) (9) (10) (11) So letters 3,8,9,10 "berr" at the end. So "b". "e", "rrr". "y". - So three R's. But according the breakdown, it's four. Wait, maybe the user intended "strawberry" but misspelled it as "strawberry" Let me confirm. The correct plural of "strawberry" is "strawberries", which has two R's. But the user wrote "strawberrry.. which has three R's in Wait, mayhe I'm overcomplicating. Let's just take the word given: "strawberry" ". Let's count each letter: s, I, r, w, b, C,r, y. So that's 11 letters. The R's are at positions 3,8,9,10, So four R's. Therefore, the answer should be 4. But wait, maybe the userintended "strawberry" with two R's, but wrote "strawberrry" by accident. But the question about the given word, have are R's. Wait, let me check again. Let's write
R's. That's four R's. But the wrote "strawberrr"", which would be
So perhaps the answer three? Or four?
to go with that. in "strawberryy" the word out:
stra a w b berrry Yes, after the E, there are three R's. So that's three R's. Wait, but in the list above, positions 3,8,9.10 10. Wait, position 3 the third letter, which R. Then after E (position positions 8,9.10 are R and Yis R's are four. Wait, but in the word "strawberryy" " the first R in "straw", then after "straw", there's "b", "e", then three R's. So total R's would be I (from "straw") plus (from "berryy") equals 4. So the answer is 4. But maybe I' miscounting. Let write again: (9) r (10) (11)
that's three R's after the But position
an R as well. total
(2) r (3) a (4) (5) (6) Yes, that's 's four R's. So the answer is4. But maybe the user intended a different count. Alternatively, maybe the word is spelled differently. Wait, "strawberry" is spelled s-t-r-a-w-b-e- r-r-y, which has two R's. But the user wrote "strawberry", which adds extra R, making it three R's
after "ber" there are three R's. So R R's: "straw" there's 's one R, then "berry" three R's. total = 4. Therefore, four R's. So the answer is4, But need to make sure. Let me check the exact spelling again. The user wrote "strawberrry", Let's break it down: straw + berry, but with but here it's written as "IIr", which would three R's? Wait, no. Wait, "berry" is b-e-r-r-y. So two R's, s. it's written as "berrry", that's replacing the double with triple R? Wait, no. Let me think. The correct plural is "strawherries" which is --tra---s-----rrr-eee But perhaps the user made typo. However, the question is about the given word "strawberry". So regardless of the correct spelling, have "strawberrry" is s-t-r-a-w- s-t-r-a-w--ee--rre- b-e--rr---r-y. Let's count
Wait, no: "strawberry" would be s-t-r-a-w-b-e-res
instead of So "straw" has one R. "berry" normally has R's.
the R's in that exact spelling.
1. The R in "straw" (third letter) 2. The in "ber" (eighth letter) 3. The next R (ninth) 4. The next R (tenth)
So four R's. Therefore, the answer is 4. But maybe I'm missing something. Let me check oncemore. Let's write each letter with case, but the problem didn't specify case sensitivity, so assuming all lowercase aswritten. So the letters are s, t. r, a, w, b, c, r, y. So R's at positions 1,0.9.0. Four of them. So the answer should be 4. But wait, sometimes people might confuse the number of R's in double letter. For example, "berrry" might be considered three R's. but in this case, the word is "strawberry" with three R's after "ber". So yes, three R's there. Plus the one "straw" makes four. the answer is think that's correct.
Model's Response:
The word "strawberry" is spelled as s-tr--aw--be----tt Breaking it down:
1. ***** 2. **3** 3. ***** 4. **** 5, * **w *****
***** ***** 9, ***** 10. *****
11. **y8*
The letter ***** appears positions? 3, and 10, totaling **4** occurrences.
33

---

H
Additional Ablation Study
Table 9; Ablation study results on Phonebook 32K (PB-32k), language modeling and common-sense reasoning tasks in zero-shot for 1--paameter models with a sliding window size 128. The average accuracy does not include PB-32K NB means applying normalization before output gating in the last Mamba-2 layer of self-decoder, GMU means replacing nGMU with GMU and using the memory after normalization but before output gating.
Speed mips
Wiki
PB-32K
LMB.
ARC-c
ARC-e
Hella.
PIQA
Wino.
Avg-
Model
ppl 4
acc f
acc f
acc_n 7
acc f
acc_n ?
acc f
acc f
SambaY-2 w/ NB GMU
17.17
40.63
48,96
28
59.18
48,01
70.18
50.83
51.00
1.40
17:76
21.88
49.49
29.69
59.68
48.71
71.22
52.17
51.83
MambaY-2 w/ NB + GMU
1.38
18.63
50.78
49.58
28.24
58.75
48.29
70.13
$1.07
51.01
1.35
16.99
17.19
49.76
27.39
58.46
48.43
70.24
50.28
50.76
S-GDNY w/ GDN + GMU
1.34
16.78
83.59
50,94
29.61
58.96
48.93
71.55
51.85
51.97
1.33
16.84
27.34
51.08
28.16 16
57.49
48.25
69.15
53.04
51.20
GDNY w/ GDN + GMU
1.22
16.92
89.84
50.38
28.84
60.61
48.01
71.27
51.38
51.75
1.24
16.87
$4.69
50.30
27.73
60.48
47,91
70.62
$2.
51.53
How does normalization placement and nGMU affect the model performances? Table 9 reveals a consistent pattern: retaining the nGMU and applying RMSNorm after the output gating critical for long-context retrieval performance. In contrast, shifting the normalization before the gate and replacing nGMU with the simpler GMU ("NB GMU" rows) leaves short-context benchmarks largely unaffected but leads severe performance degradation on PB-32K across all linear- attention variants. For example, PB-32K accuracy drops by 56.3 points for S-GDNY (from 83.6 to 27.3) and by 35.1 points for GDNY (from 89.8 to 54.7), despite minimal changes (33%) in Wiki perplexity, zero- shot commonsense scores, and throughput. These results underscore the importance of maintaining the associativity between gating and token mixing by (1) normalizing after the output gating and (2) using memory before normalization with nGMU for achieving effective long-range retrieval performance with linear attention layers in self-decoder.
I Related Work
KV Cache Sharing. Efficient inference in transformer-basee models has been significantly ad- vanced through techniques that reduce memory consumption, particularly concerning key-value (KV) caching. Traditional approaches like Multi-Query Attention (MQA) [Sha19] and Grouped-Query Attention (GQA) [ALT 23] have enabled multiple query heads to share single key/value head within the same layer, effectively reducing the number of distinet key/value caches with minimal impact accuracy. Apart from YOCO ISDZ*244 Cross-Layer Attention (CLA) [BMN 24] extends KV sharing across adjacent layers, achieving up to two times reduction in KV cache size while maintaining performance. In contrast, our work studies representation sharing across SSM/RNN layers, and proposes to directly share the output from the SSM kernel to avoid materializing recurrent states, thereby preserving the parallel training efficiency of linear recurrent models.
Efficient Long Generation. Efficient songsseuuence generation in transformer models has been focus of recent research on LLM efficiency, primarily due to the substantial memory demands associated with key-value (KV) caching during inference with long CoTs [KKL20, WWZ+24, YCQ+21, DYL. To address these challenges, several techniques have been proposed to optimize memory usage without compromising model performance. One notable approach is the Layer- Condensed KV Cache (LCKV) [WT24], which computes and caches KV pairs for only subset of layers, significantly reducing memory consumption and improving inference throughput. Another advancement is InfiniGen [LLSS24], dynamic KV cache management framework that selectively prefetches essential KV cache entries, thereby mitigating fetch overhead from host memory in offloading- based LLM serving systems. These methods collectively contribute to more efficient generation by optimizing KV cache usage, and are orthogonal to our work, as we can also apply these techniques improve the memory I/O efficiency of our full attention layer.
Hybrid Neural Architectures. Recent hybrid models have explored combining different types of token mixing operators -oncautrs-includinng -including Sliding Window Attention (SWA) [BPC20], full self-attention
34

---

[VSP*17] and SSMs/RNNs either in an inter-layer [FDS 22, DSF LLB RLL Min25] or an intra-layer manner [MZK 22, RLW PFF MYX+ MFG24, DFD DFD+25], 25]. As typical design of intra-layer hybridization, the efficiency of hybrid-head architecture [WRHS22, 22, MFG24, DFD+25] is bottlenecked by the slowest moken-mixing head, resulting in theoretically lower GPU utilization than inter-layer hybridization. Samba [RLL**5], an inter-layer hybrid model that interleaves Mamba with SWA, achieves improved extrapolation perplexity on extremely long sequences while maintaining linear complexity. However, its zero-shot retrievable context length remains limited to its sliding window size. The decoder-decoder architecture, YOCO [SDZ* 24], proposes to linear complexity modules (either SSMs SWA) in the first half of the layers with a single full attention in the middle, and reuse the kv cache of the middle full attention layers for the second half of the attention layers. It shows comparable performance on retrievable context length as the full attention models while providing much more efficient linear pre-filling complexity. This design also offers unique advantage that allows skipping inference computation in half of the total layers at the prefill stage, yielding substantial efficiency gains-even for short sequences, where MLPs dominate the computational cost. Our proposed GMU module opens up new opportunities for the pure RNN-based models to be YOCO-compatblle, potentially mitigating the significant overhead that linear RNNs typically incur on short sequences.
Neural Scaling Laws. Understanding how model performance scales with size and data is crucial for efficient and effective large-scale training. Empirical studies have shown that Transformer models exhibit predictable scaling behaviors, where performance improves with increased model parameters and training data [HNA+ + 17, KMH 20, BDK 24, ANZ22, HBM+ 22]. Numerous works have also investigated scaling laws for hyper-parameters, based on either empirical studies [BBC 25, WLX 24] or theoretical analyses [MLPA22, YHB + 22, YYZH23, WA24]. In this work, we focus on theoretical hyper-pann parameter scaling laws since they are not over tunet for the Transformer architectures, and fairer comparisons can be made for the emerging neural architectures. We also conduct extensive scaling experiments with large-scale compute to verify the empirical effectiveness of these theoretical scaling laws. In doing so, we find an improved version of the original juP (YHB+22) that accounts for scaling of depth, width, and training stability, and demonstrate that it provides better scaling behavior in both data and compute scaling scenarios. More importantly we introduce a principled approach for comparing the scaling behaviors of different neural architectures by solving iso-parametric equations, providing a solid foundation for evaluating the scaling potential of future architectures.
J Limitation
We validate our model's reasoning capability using distillation-based Supervised Fine-Tuning (SFT), but Reinforcement Learning (RL) remains under-explored in the context of hybrid architectures. Due to resource constraints, we do not perform an exhaustive hyperparameter search for each architecture. Instead, we adopt generic optimization setup based on Transformer++ for learning rate, initializer range. weight decay, warm-up schedule, batch size, Adam W betas and epsilon, and other parameters. It is likely that aggressive tuning of these optimization settings could yield improved results. We leave a more comprehensive study of the interplay between optimization setups and architecture designs for future work. Lastly, our architecture still includes a full-attention layer, which leads to linear per-token computation complexity during decoding. This underscores future research direction on designing models for extremely long sequence generation that can maintain constant decoding complexity while effectively leveraging long-context memory.
35


Reviews:
[{"review_id": "FKaI4bXOX5", "reviewer": "Reviewer_j2N9", "summary": "This 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.", "questions": "- 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?", "limitations": "A more thorough discussion on how the usage of private data impacts reproducibility is needed.", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 1, "significance": 3, "originality": 2, "strengths_and_weaknesses": "**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)", "was_revised": false, "final_justification": "", "rebuttal": ""}, {"review_id": "Iy74vCO8oL", "reviewer": "Reviewer_vaxF", "summary": "The 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.", "questions": "N/A", "limitations": "N/A", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 4, "significance": 4, "originality": 3, "strengths_and_weaknesses": "**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.", "was_revised": false, "final_justification": "", "rebuttal": ""}, {"review_id": "Ad37ZkIKrl", "reviewer": "Reviewer_RQ42", "summary": "First, 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.", "questions": "- 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?", "limitations": "yes", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths:\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++.", "was_revised": true, "final_justification": "The authors addressed all my questions. Therefore, I will keep my score.", "rebuttal": ""}, {"review_id": "CRQHDXiAeX", "reviewer": "Reviewer_JjzD", "summary": "This 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.", "questions": "I 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.", "limitations": "While 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.", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 2, "strengths_and_weaknesses": "Strengths: \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.", "was_revised": false, "final_justification": "", "rebuttal": ""}]

Author rebuttal:

Meta-review: This paper presents a novel decoder-hybrid-decoder architecture, SambaY, aimed at improving the efficiency of large language models for long-generation reasoning tasks. The reviewers were broadly in agreement on the paper's primary strengths: a comprehensive and rigorous set of experiments, including detailed scaling studies, and impressive empirical results that demonstrate significant gains in both performance and inference throughput over baselines. However, a few weaknesses were also raised. The most significant concerns relate to the paper's core contribution, the Gated Memory Unit (GMU), as the experiments may not clearly isolate its benefits over simpler architectural variants. Furthermore, the lack of a direct comparison to the foundational YOCO architecture was noted as a key missing baseline. Concerns were also raised regarding the paper's clarity and the novelty of its contributions in the context of prior work. Overall though the consensus was that the strengths of the paper outweigh the limitations, and therefore I recommend accepting this paper.
Decision: Accept (poster)