How Does Sequence Modeling Architecture Influence Base Capabilities of Pre-trained Language Models? Exploring Key Architecture Design Principles to Avoid Base Capabilities Degradation
Authors: Xin Lu
Yanyan Zhao
Si Wei
Shijin Wang
Bing Qin
Ting Liu
Venue: NeurIPS 2025 poster ()
Keywords: Pre-trained Language Models
Base Capabilities
Sequence Modeling
Primary area: deep_learning

Abstract: Pre-trained language models represented by the Transformer have been proven to possess strong base capabilities, and the representative self-attention mechanism in the Transformer has become a classic in sequence modeling architectures. Different from the work of proposing sequence modeling architecture to improve the efficiency of attention mechanism, this work focuses on the impact of sequence modeling architectures on base capabilities. Specifically, our concern is: How exactly do sequence modeling architectures affect the base capabilities of pre-trained language models? In this work, we first point out that the mixed domain pre-training setting commonly adopted in existing architecture design works fails to adequately reveal the differences in base capabilities among various architectures. To address this, we propose a limited domain pre-training setting with out-of-distribution testing, which successfully uncovers significant differences in base capabilities among architectures at an early stage. Next, we analyze the base capabilities of stateful sequence modeling architectures, and find that they exhibit significant degradation in base capabilities compared to the Transformer. Then, through a series of architecture component analysis, we summarize a key architecture design principle: A sequence modeling architecture need possess full-sequence arbitrary selection capability to avoid degradation in base capabilities. Finally, we empirically validate this principle using an extremely simple Top-1 element selection architecture and further generalize it to a more practical Top-1 chunk selection architecture. Experimental results demonstrate our proposed sequence modeling architecture design principle and suggest that our work can serve as a valuable reference for future architecture improvements and novel designs.

How Does Sequence Modeling Architecture Influence Base Capabilities of Pre-trained Language Models? Exploring Key Architecture Design Principles to Avoid Base Capabilities Degradation
Xin Lut, Yanyan Zhaol.*, Si Wei Shijin Wang", Bing Qin', Ting Liu! Research Center for Social Computing and Interactive Robotics, Harbin Institute of Technology
2iFLYTK Co., Ltd {xlu, yyzhao, qinb, tliu)@ir.hit edu.cn
2 {siwei, sj com
Abstract
Pre-trained language models represented by the Transformer have been proven to possess strong base capabilities, and the representative self-attention mechanism the Transformer has become classic in sequence modeling architectures. Dif- ferent from the work of proposing sequence modeling architecture to improve the efficiency of attention mechanism, this work focuses on the impact of sequence modeling architectures on base capabilities. Specifically, exactly do sequence modeling architectures affect the base capabilities of pre- trained language models? In this work, we first point out that the mixed domain rre-tranning setting commonly adopted in existing architecture design works fails adequately reveal the differences in base capabilities among various architec- tures. To address this, we propose a limited domain pre-training setting with out-of-distribution testing, which successfully uncovers significant differences in base capabilities among architectures at an early stage. Next, we analyze the base capabilities of stateful sequence modeling architectures, and find that they exhibit significant degradation in base capabilities compared to the Transformer. Then, through series architecture component analysis, we summarize key architec- ture design principle: A sequence modeling architecture need possess full-sequence arbitrary selection capability avoid degradation base capabilities. Finally, we empirically validate this principle using an extremely simple Top- element selection architecture and further generalize it more practical Top-1 chunk selection architecture. Experimental results demonstrate our proposed sequence modeling architecture design principle and suggest that our work can serve as a valuable reference for future architecture improvements and novel designs.
concern is: How
pre-
Introduction
Recent research has discovered that pre-trained language models [35 12 6 30), represented by Transformer [47]. possess strong base capabilities and can achieve excellent performance in language modeling, few-shot learning, etc. Delving into the specific architecture design of Transformer, its self-attention mechanism is widely regarded as one of the key components behind its success and has since become classic in sequence modeling architectures.
* Email corresponding.
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

(a) Mixed Domain Pre- Training
(b) Limited Domain Pre- ree-Trinnng
Love on Moverved Docmen lest S400000
Looson Docmen (est
2.34
3.00
223
230
206
2.70
2,30
2:50
2.70
2.00
2.00
3.00
3.10
3.0
2.0
2.70
2.00
2.92
3.00
3.10
3.20
3,30
Pro Pre-Trannin Loss
Pre- Trunting Loss
Based
Mamba-1
Mamba-2
RWKY-5
RWKV-6
RWKV-7
DeltaNet
Gated DeltaNet
Translomee (W -230)
Transtermer
Transormer++
Figure Language modeling results of various sequence modeling architectures under two pre-training settings. (Model parameters 100M, pre-trained tokens 100B and sequence length-2k)
However, due to its quadratic time complexity, the self-attention mechanism has long been plagued by high computational costs in Long-sequeeen scenarios. To improve the efficiency of sequence modeling architectures, numerous novel stateful sequence modeling architectures have recently emerged, such as Mamba 16 I RWKV [31 32 33 Gated DeltaNet [5T]. These architectures primarily inherit
stateful modeling mechanisms of linear attention 23 and linear RNNs (28). offering advantages in time and space efficiency over self-attention. Additionally, they introduce new mechanisms like data-dependent decay and delta rule to enhance the expressive power of linearized modeling.
Although experimental results demonstrate that these stateful sequence modeling architectures can match even surpass Transformer in performance while maintaining significant efficiency advantages, some studies have noted their deficiencies in specialized capabilities such retrieval [4), copy (21], associative recall [1] and dynamic programming [50]. supported by empirical validation on synthetic datasets and tasks. While these studies focus on specific issues, they are highly insightful and directly raise our new questions about sequence modeling architectures:
Are the limitations of stateful sequence modeling architectures not only confined specialized capabilities but also present in base capabilities?
What architecture factors truly influence base capabilities?
What design principles can prevent the degradation of base capabilities?
To address these questions, we first point out that existing research sequence modeling architectures typically adopts the same mixed domain pre-training settings used in large model development. While beneficial for practical applications, these settings detrimental to architecture analysis, they turn base capability tests (e.g., language modeling) into in-distribution evaluations, failing to reveal differences in base capabilities during early pre-training stages. To address this, we propose a limited domain pre-training approach, employing nnt-of-distribution language modeling performance to measure base capabilities, successfully uncovering significant differences in base capabilities among architectures at an early stage.
Next, under this setting, analyze the base capabilities of stateful sequence modeling architectures like Mamba, RWKV, Gated DeltaNet. Our findings reveal that these architectures exhibit notable degradation in base capabilities compared to Transformer, confirming that stateful sequence modeling architectures suffer from deficiencies not only in specialized capabilities but also in base capabilities.
We then investigate the architecture factors that truly impact base capabilities. Through ablation studies on the Mamba family of architectures and analyses of common sequence modeling factors, we identify that mechanisms like data-dependent decay, convolution and position encoding only affect convergence speed rather than base capabilities. Conversely, we determine that full-sequence visibility, real relation calculation and non-uniform distribution are critical architecture factors

---

(8] Maxed Domain Ere Traiting
(B) Linzited Donaza Prs Traning
(c) Sow Shot Petfomance
3.00
30.0
$3.43
4345
1681
54.01
48.0
2055
2.50
1.00
2770
1.75
2.50
2.0
Pro Thatiig Loss
Pro Truzing Loss
Modee
Transfrmer++
Transformor* (W-256)
Mambu-1
Mamba-2
Figure 2: The Illustration include: (a) and (b) Language modeling test results of various sequence modeling architectures under two pre-training settings, (c) Few-shot learning performance results of these architectures. (Model parameters= 1.3B, pre-trained tokens=1008 and sequence length=2k)
influencing base capabilities. Based on this, we summarize the principle that "a sequence modeling architecture need to possess full-sequence arbitrary selection capability" as the key design rule avoid degradation in base capabilities.
Finally, we validate this principle using an extremely simple Top-1 Element Selection architecture, experimentally confirming its ability achieve base capabilities nearly on par with the Transformer. Furthermore, we extend this validation to the more practical Top-1 Chunk Selection architecture direct generalization of the Top-1 Element Selection architecture, and implement GPU Kernels to ensure its time efficiency. Experiments show that the Top-1 Chunk Selection architecture, which adheres to our design principle, outperforms stateful architectures in base capabilities for both short (2k) and long (100k) sequences while maintaining competitive time efficiency.
The main contributions of our work follows:
We propose a more indicative limited domain pre-training and -fddsiiiiitttnoonnn testing framework, successfully revealing the base capability degradation in stateful sequence modeling architectures during early pre-training stages. (Section 2)
We investigate the impact of various common sequence modeling architecture factors on base capabilities and summarize the principle that "full-sequence arbitrary selection capability" is critical to avoiding degradation in base capabilities. (Section 3
We validate the principle using extremely simple Top-1 Element Selection architecture and demonstrate its generalization to the accompanied by open-source GPU Kernel to ensure time efficiency. (Section and 5)
practical Top-1 Chunk Selection architecture,
Through our analysis and experiments, we prove the effectiveness the proposed architec- ture design principle, providing valuable insights for future architecture designs.
Base Capabilities Evaluation
This work focuses on how sequence modeling architectures influence the base capabilities of pre- trained language models, thus necessitating the design of an evaluation framework and the implemen- tation of assessments for different sequence modeling architectures.
2.1 Evaluation Scheme for Architecture Analysis
Existing sequence modeling architecture design works typically adopt the same Mixed Domain Pre-Traine Pre-Training setting as large language model development, where corpora from as many domains as possible are collected and mixed as pre-training data. While this setting benefits practical large language model applications, it is detrimental to architecture analysis. It turns base capability tests

---

like language modeling into in-distribution evaluations, failing to reveal differences in architecture base capabilities during early pre- pre-training stages. Moreover, poorly predicts the true usability of different architectures when applied to unknown
domain tasks.
We tested this setting. Specifically, we pretrained models by mixing all domains (cc, c4, arxiv, book, github, stack and wiki) from the SlimPajama dataset [42 and retained mixed-domain test set for evaluation. For models with x 110M parameters. we pretrained Transformer (6). Transformer++
(with Rotary Embedding 43 GeGLU 41 and RMSNorm (56]), Transformer++ (Window=256),
Based Mamba-1 [6 Mamba-2 II RWKV-5 32 RWKV-6 32, RWKV-7 [33], DeltaNet [40]
and Gated DeltaNet [51] with a sequence length of 2K and 100B tokens of pre-training The scatter
plots of pre-training loss versus mixed-domain test loss are shown in Figure [T). For models with s
1.3B parameters, we pretrained Transformer++, Transformer++ (Window=256), Mamba-1 [T6] and Mamba-2 [1 under similar settings, with results plotted in Figure (aa)
From Figure D a) and Figure it is evident that under the Mixed Domain Pre- Training setting, different sequence modeling architectures achieve similar test performance in language modeling when reaching comparable pre-training performance levels. Thus, this setting fails to reveal base capabilities difference among architectures during early pre-training and unsuitable for architecture design and analysis.
To address this issue, propose Limited Domain Pre- Pre-Training Training with testing framework, where models are pretrained on a restricted set of domains and evaluated on unseen domains.
(OOD)
Since the SlimPajama dataset has already undergone deduplication across domain subsets, we easily tested this setting. Specifically, we pretrained models on the cc and c4 domains of SlimPajama and evaluated them on the arxiv, github and stack domains. Other pre-training settings remained similar, and the scatter plots of pre-training loss versus OOD test loss are shown in Figures ( and Zb).
From Figure ([) and Figure 2b), significant differences emerge among sequence modeling architec- tures. At the same pre- pre-training performance level, their OOD test performance varies substantially. confirming that base capabilities difference exist across architectures.
Additionally, we evaluated models with s 1.3B parameter using the commonly adopted few-shot learning evaluation in prior work, which results in Figure ce. Similarly, no significant performance gaps were observed, suggesting that this evaluation method is also suboptimal for architecture design and analysis.
Based on these findings, we establish Limited Domain Pre-Training with OOD testing as our evaluation framework for base capabilities assessment. All subsequent tests in this work follow it.
2.2 Architecture-Indueed Degradation of Base Capabilities
We further analyzed Figure Tbb) and Figure bb). The results show that only the standard attention- based Transformer and Transformer++ achieve optimal base capabilities- OOD test performance is consistently the best at the same pre-training level, with no significant difference between them. In contrast, stateful sequence modeling architectures (e.g., Mamba. RWKV) exhibit varying degrees of base capabilities degradation, performing significantly worse than standard attention- based models under identical pre-training conditions.
These findings suggest that sequence modeling architectures directly induce base capabilities degrada- tion, independent of data or other factors. The standard self-attention mechanism likely contains key architecture factors critical to base capabilities, some of which may be missing in stateful sequence modeling architectures.
3 Sequence Modeling Architecture Analysis
Based on the previous results, we know that the sequence modeling architecture can directly determine the model's base capabilities, and certain key architecture design factors are likely to play a significant role in these base capabilities. Therefore, identifying which architecture design factors are truly critical could be of great importance for subsequent architecture improvements or the design of new architectures.

---

La M Ablaion
(3) 2anbaa Ablaiio
(e) Ptolion Encodex
1.50
L00 Tou Sen
3.59
3.54
3.15
2.30
250
131
245
225
Pro-Trainkng
Pro Twiiiinngggg. Leve
ProTninngg..
(5) Shiling Window Sizr
(c) Rehaiion Calcalation
(
Derruttion Unifeenity
Aich veer vid
4.3
4.33
0m Toa
Loo Toa
3.0
1.50
3.3
172
355
270
370
280
2000
3.0
210
3.3
3.0
231
233
1.3
1.50
275
255
315
535
355
575
3
415
Pee-Triiing
PosTuiiig
PeeTrohing
Figure 3: Analysis of the influence of various sequence modeling architecture components on base capabilities. (Model parameters 110M, pre-trained tokens- or 25B and sequence length-2k)
3.1
Non-Determinative Factors of Base Capabilities
3.1.1 Mamba Ablation
Since the classical stateful sequence modeling architecture, Mamba, has been observed to exhibit degradation in base capabilities, we conducted ablation studies on some of its key components to assess their actual impact on base capabilities.
First, we focused on the Data-Dependent Decay in Mamba. By replacing the multi-dimensional independent data-dependent decay (Full) in Mamba-1 and Mamba-2 with either shared data- dependent decay (Decay One) or no explicit decay (Decay None), we evaluated the effect of data- dependent decay on base capabilities, as shown in Figures 3 and The results indicate that data-dependent decay only accelerates pre-training convergence but does not positively impact base capabilities. Specifically, at the same pre-training level, the antoof-distribution test performance of the ablated models did not decrease and even improved slightly.
Next, we examined the Convolution in Mamba. We further removed the convolution from the Decay None architecture (Decay None w/o Conv) to assess its effect on base capabilities, as shown in Figures B a) and (b). The results show that while convolution significantly speeds up pre-training convergence, it still does not contribute positively to base capabilities.
Finally, we investigated the CroupNorm in Mamba-2. By removing GroupNorm from the full architecture (Full + w/o GN), we evaluated its impact on base capabilities, as shown Figure 3b). The results confirm that GroupNorm also does not enhance base capabilities.
3.1.2 Position Encoding
Previous results have shown that Transformer and Transformer++ exhibit nearly identical base capabilities, with their primary difference in sequence modeling lying in their position encoding schemes. Therefore, we analyzed the impact of position encoding base capabilities.
Specifically, we tested four common position encoding schemes: no position encoding, absolute position embedding, 344 and rotary position embedding [43), as illustrated in Figure c). The results indicate that no position encoding, absolute position embedding and rotary position embedding yield very similar base capabilities, differing only in convergence speed. In contrast, AliBi exhibits some degradation base capabilities. Based on these findings, we conclude that position encoding does not positively influence base capabilities.
5

---

3.2 Determinative Factors of Base Capabilities
3.2.1 Full-Sequence Visibility
The first factor we identified as critical to base capabilities is sliding window size. This observa- tion was prompted by the significant difference in base capabilities between Transformer+ and Transformer++ (Window=256), suggesting that sliding window size may be a key factor.
To investigate, we tested various sliding window sizes, as shown in Figure The results demon- strate that as the sliding window size increases, both pre-training convergence speed and base capa- bilities improve significantly. Conversely, reducing the sequence context leads to gradual degradation base capabilities. Thus, we conclude that Full-Sequence Visibility is an essential requirement for sequence modeling architectures.
3.2.2 Real Relation Calculation
In previous results, we observed that under mixed domain pre-training settings, different models exhibited nearly identical test performance, creating the illusion that base capabilities is architecture- independent. This reminded us of the Synthesizer [45) series of models, where one variant replaced the true query-key computed scores with trainable random constant scores yet did not exhibit significant degradation in language modeling tasks. We hypothesized that this might also be due to mixed domain pre-training and that real relation computation between queries and keys could be determinant of base capabilities.
To test this, we conducted experiments where we replaced keys with trainable random constants or replaced both queries and keys with trainable random constants, as shown in Figure (3)). The results reveal that models without real relation computation suffer substantial degradation in base capabilities. Therefore, we conclude that Real Relation Computation a necessary feature for sequence modeling architectures.
3.2.3 Non-Uniform Distribution
Unilomity
Another key factor we examined is the uniformity of attention distribution in self-attention mechanisms. We reasoned that if attention distribution were entirely uniform, the model would degenerate into a naive averaging structure over values, which exhibits poor base capabilities. However, if we start from this structure and gradually introduce non-uniformity, the model would evolve toward normal attention distribution, which has strong base capabilities. This led us to suspect that distribution uniformity directly impacts base capabilities.
59
Temperture
Figure 4: The relationship be- tween attention distribution en- tropy and temperature.
To control attention distribution uniformity, we employed two tech- niques: 1) Adjusting distribution uniformity via the temperature of the Softmax function. 2) Applying normalization to queries and keys to ensure consistent temperature interpretation, unaffected by the optimization of query and key transformation parameters. We tested a series of models with varying temperatures and plotted the relationship between temperature and attention distribution entropy in Figure 4 The results show that as temperature decreases, attention distribution entropy also decreases, contirming that the con-uniform.
distribution becomes
Subsequently, we evaluated the ofdiiiiiiiiittiinnn performance of these models, shown in Fig- ure f.. The results indicate that as temperature decreases (i.e., distribution becomes more non- uniform), base capabilities improves. Thus, we conclude that Non-Uniform Distribution necessary feature for sequence modeling architectures.
3.3 Key Architecture Design Principle: Full-Sequence Arbitrary Selection
Integrating the three key elements derived from the preceding analysis: Full-Sequence Visibility, Real Relation Calculation and Non- Non-Uniform Distribution, summarize them into unified expression:

---

(a) Top-1 Element Selection A rrchitecture
(b) Top- 1 Chunk Selection Kernel Design
Order Lop
Quary
XX
Outar Logp
Secuumc Viaiiiit
Nae Won-tr Strthedke
Output
Sce bydde
Loop
Figure 5: The Illustration include: (a) The overall architecture design of the Top-1 Element Selection architecture. (b) The kernel design for key component of the Top-1 Chunk Selection architecture.
Supporting full-sequence arbitrary selection in sequence modeling architecture is the key architecture design principle for preventing the degradation of base capabilities.
This principle directly impacts the model's 's base capabilities, and violating will lead to degradation.
Note that this is an architecture design principle rather than capability principle. The latter is often expressed as the model's retrieval capability. However, retrieval capability itself is not inherently tied to architecture design principles. For instance, stateful sequence modeling architectures like Mamba exhibit certain retrieval capabilities, yet they lack the architecture design of full-sequence arbitrary selection. is this aspect of architecture design that our work truly focuses on.
Top-1 Element Selection Architecture
In previous section, we proposed a key architecture design principle for avoiding degradation in base capabilities. However, this conclusion was derived inductively and has not yet been experi- mentally validated. To address this, we designed an extremely minimalist Top-1 Element Selection architecture, which directly adheres to this design principle while maintaining strong base capabilities.
4.1 Architecture Design
We posit that "sequence modeling architectures with full-sequence arbitrary selection" is key archi- tecture design principle for preventing degradation in base capabilities. The simplest implementation of this principle tion as the output, as illustrated in Figure 5(a). We refer to this architecture as the Top-1 Element additional operations: 1) applying normalization
directly select the element with the highest probability the attention distribu-
Selection architecture. In practice, it involves to queries and keys to enhance stability. 2) during training, attention scores are still computed, but the straight-through trick introduced to reconcile top-1 selection with gradient updates.
As shown in Figure a)), the Top-1 Element Selection architecture is remarkably simple yet satisfies the design principle. also incorporates three key elements: full-sequence visibility, real relation calculation and non-uniform distribution, making an excellent candidate for validating our analysis.
4.2 Out-of- Distri--Distiiiation Generalization Evaluation
We pre-trained models with 110M and 1.3B parameters, maintaining the same pre-training settings as in previous experiments, and evaluated their OOD performance. The results, shown in Figures ([a) and b(b), demonstrate that the Top-1 Element Selection architecture achieves OOD performance
7

---

(a) Models with Parameter:: 110M
Models with Parameter 1.3B
3.00
3.50
3.50
2.20
Loss-on OOD Test Set
3.00
3.10
2.10
2.30
2.00
2.30
2.90
Pro Trining Loss
Pro: Truning Loss
Transormert (W -210
Mamba-1
Mamba-2
Transformert+
Top-1 Element
Top-1 Chunk (Exort)
Top-1 Chink (Approx.)
Figure 6: The out-of-distribution language modeling results of baselines and Top-1 Element / Chunk Selection architectures. (Pre-trained tokens=100k 100B, sequence length=2k and chunk size=1288
nearly on par with Transformer++ at both scales. This confirms that the architecture does not suffer from base capability degradation and validates the correctness our proposed design principle.
4.3 Few-Shot Learning Evaluation
Although we previously argued that few-shot learning evaluation is not an ideal method for comparing base capabilities across architectures, remains necessary to verify that the new architecture does not exhibit sig- nificant degradation few-shot learning performance. Thus, we evaluated the few-shot learning performance of the 1.3B parameter model, as detailed Table The results show that the Top-1 Element Selection architec- ture achieves few-shot learning performance comparable to models like Transformer+++ Mamba-1 and Mamba-2, indicating its strong performance in this setting.
Model
SWDE
SQUAD (seef)
(scef)
(acef)
Original Sequence Modeling Archilectore
Transformer++
75.07
45.54
30.22
Stangful Seguence Madeling Architecture
Mamba-1
36.09
31.03
6,44
Mamha-2
37.98
35.25
17.60
Transformer++ (W=256)
12.96
41.72
1.63
Architecture Bused the Analyzed Principles
Top- Element
68.23
43.77
26.95
Top-I Chunk (Exact)
69.22
44,40
22.32
Beyond standard tasks, we also evaluated retrieval tasks, common benchmark in prior work. As shown in Ta- ble [ the Top-1 Element Selection architecture signif- icantly outperforms stateful sequence modeling archi- tectures like Mamba-1 I and Mamba-2 in retrieval tasks, further demonstrating effectiveness.
Top-1 Chunk (Approx.)
61.66
45.38
27.13
Table I: Results of retrieval tasks. (Model parameters 1.3B, pre-trained tokens= 100B, seq len=2k and chunk size= 128)
5 Top-1 Chunk Selection Architecture
In the previous section, we designed the Top-1 Element Selection architecture, successfully validating our proposed architecture design principles. However, this architecture serves as a proof-of-concept and lacks practical utility. To address this, we extended it to the Top-1 Chunk Selection architecture and implemented GPU kernels to optimize efficiency while maintaining strong base capabilities.
5.1
Architecture Design
5.1.1
Basic Design
The Top-1 Element Selection architecture has been verified to possess strong base capabilities but performs poorly in terms of pre-training convergence and time efficiency. To resolve this,

---

OBQA nft
BoolQ
SIQA
SCIQ
Hella.
PIQA
Score
Origisal Siqpinne Moddiing Avohlientve
Transtommrr+
25.21
33,60
48.48
25.9 94
60.43
33,88
39.56
77.60
48.42
74.00
68.99
53.91
65.57
50.43
Suntial Modiling
Mamba-1
24.30
34.60
52.82
27.22
$4.65
32.73
39.10
75,90
50.53
74.00
70.67
53.91
57.51
49.34
Mamba-2
$2.80
$0.04
28,92
33,68
75.70
71.44
60:44
50
Tranffrmmnt- (W-256)
73.70
63.00
Acchttatu Maree the Analyced Principtee
Top-l Flement
25.94
32,60
47.85
26.45
60.15
31.29
40.94
73.40
45.05
72.00
69.15
53.25
62.27
49.30
Top I Chunk (Exact)
25.96
32.80
50.55
27.22
60.43
51.77
40.07
77.40
48.75
75,00
69.64
53.51
66.30
50.72
Top-1 Chank (Appeos.)
25.24
33.00
49.33
26.02
5887
33.30
39.82
75.30
48.59
70.00
69.70
55.01
64.47
49.50
Table 2: The few-shot learning experimental results, where MMLU is 5-shot and the remaining tasks are 0-shot. (Model parameters==1.355 pre-trained tokens=1000 100B, seq length=2k and chunk size=128)
generalized it to the Top-1 Chunk Selection architecture. The key difference from the Top-1 Element Selection architecture is that the query selection target shifts from fine-grained kv elements to coarse- grained kv chunks, and the attention mechanism operates only within the selected chunks. This preserves the design principles while enabling efficiency optimizations.
More specifically, the kv sequence divided into mul- tiple kv chunks. Each query attends to one selected full kv chunk (remote chunk) and the nearest partial kv chunk (local chunk), performing attention operations on them. Similar the Top-1 Element Selection architec- ture, both queries and keys undergo normalization to enhance stability.
Models with Sequence Length = 100K
Mandoo
5.1.2 Exact and Approximate Variants
The selection of full kv chunks can be implemented in two ways, corresponding to two variants:
Top-1 Chunk (Exact): The query computes the full attention distribution with all keys, obtaining the prob- ability for each kv chunk. The chunk is selected based on these exact probabilities.
435
285
2.75
2.55
2.35
Figure 7: The OOD language modeling re- sults of baselines and ours. (Model param- eters Rs 135M, pre- rre-trained tokens= 100B, seq length=1000 100k and chunk size= 128)
Top-1 Chunk (Approx.): The exponential function in probability computation can be approximated by a first- order linear function. Under this approximation, the selection process simplifies computing the mean vector of each key chunk and performing a dot product between the query and the mean vector select the chunk. This significantly reduces computational overhead.
5.1.3 GPU Kernels Design
To achieve better time efficiency, we designed and implemented GPU kernels for the Top-1 Chunk Selection architecture. An overview is given in this section, with details in Appendix D
For exact chunk selection, approximate chunk selection and local chunk attention, directly implemented Triton kernels using simple, naive designs that already deliver good performance.
For remote chunk attention, we conducted specialized optimizations when implementing Triton Kernels, shown in F(gbre(5(b). Specifically, parallel processing performed along the kv dimension in complete chunks, meaning different kv chunks are processed independently in parallel. This partitioning is feasible because each query selects only one remote chunk, ensuring no overlap between queries associated with different kv chunks. Within each kv chunk, we load index via offset, then load the queries relevant to the current chunk by the index and compute attention in SRAM, storing results by the index. Efficient offset and index handling is achieved by custom CUDA kernels.

---

During our research, we observed that DeepSeek and Kimi released two sparse attention works NSA [53] and MoBA (27] in February 2025. Their sequence modeling approaches closely resemble our Top-1 Chunk (Approx.). Although we discovered this architecture independently, we no longer claim Top-1 Chunk (Approx.) as our primary contribution. Instead, our key contributions include Top-1 Element, Top-1 Chunk (Exact), analytical experiments and the released GPU kernels.
5.2 Out-of-Distribution Generalization and Few-Shot Learning Evaluation
We pre-trained models with = 110M and 1.3B parameters (chunk size 128), shown in Figures 6(a) and bb). Results show both Top-1 Chunk Selection variants achieve OOD performance close to Transformer++ and outperform Mamba-1/2. We also evaluated few-shot learning tasks, as detailed in Tables 2 [ Top-1 Chunk Selection exhibits results and conclusions similar to Top-1 Element Selection: strong performance on general tasks without degradation and advantages in retrieval tasks.
5.3 Long Sequence and Architecture Combination
In this section, we evaluate two aspects: first, base capabilities, time efficiency and long-sequence retrieval capability of T Chunk Selection architecture on long sequences (100k); second, per- formance of Top-1 Chunk Selection architecture when transferred other models. We designed an integrated setting where we ported Transformer+++ Transformer++ (W=256), Top-1 Chunk (Exact) and Top-1 Chunk (Approx.) to the RWKV-7 architecture, replacing only the sequence modeling while retaining the rest of RWKV-7. Pre-training data remained unchanged, and OOD test data were filtered to retain samples from arxiv and github with original lengths exceeding 100k. Since stack contained almost data with original lengths over 100k, we supplemented the test set with samples from AutoMathText (57] met this criterion.
The base capabilitties results are shown in Figure Transformer++ demonstrated the
Modal
Pre-Teiin
Speed
Infervnee
Speed
Orignnal Areistereer (ownsddedd)
strongest base capabilities, followed closely by Top-1 Chunk (Exact). Top-1 Chunk (Approx.) outperformed other stateful ar- chitectures in base capabilities, but due to its approximate nature, it exhibited signif- icant pretraining convergence degradation and base capabilities degradation compared to Transformer++ (similar issues may also affect DeepSeek NSA and Kimi MoBA, which employ analogous approaches).
KWRY
2:4H
0.71x
70.90
14.90
0.00
0.30
Aps'
Traaffrrazes
1.0%
1,00x
76.20
43.90
25.60
16.20
13.90
Truasfommree (Wa250)
1,77
1.50
0.40
0.50
0.40
Top- (Foor)
1lpp- Chark (Appeox.)
1.20
70.00
43.00
21.80
Table 3: Results of time efficiency and long se- quence retrieval. (Model 135M, pre-trained tokens=1000 100B, seq length=100k and chunk size= 128)
Time efficiency and long-sequence retrieval results are presented in Table 3 Top-1 1 Chunk (Approx.) achieved substantial speed improvements over Transformer++, reaching time efficiency levels compa- rable to Mamba-2. Top-1 Chunk (Exact) also showed speed gains, achieving time efficiency similar to RWKV-7, In oong-sequence retrieval, results on S-NIAH (20) shown Top- Chunk Selection delivered performance close to Transformer++, while other stateful models suffered severe degradation.
Limitations
This work primarily focuses on models pre-trained with language modeling objectives, without exploring other pre-training objectives. As a result, our conclusions are limited to language models, narrowing the current applicability of our findings. We plan to expand this research with additional experiments in the future.
7 Conclusion
This work investigates how sequence modeling architectures influence base capabilities of pre-trained language models. We first reveal the degradation of base capabilities in stateful sequence modeling architectures by limited domain pre-training. Subsequently, via architecture analysis, we identify "full- sequence arbitrary selection" as the key architecture design principle for preventing such degradation. Finally, we validate our analysis by proposing Top-1 Flement Selection and Top-I Chunk Selection architecture, which may provide valuable references and foundations for future research.
10

---

Acknowledgments
This work was supported by the New Generation Artificial Intelligence-Nationn Science and Tech- nology Major Project 2023ZD0121 1000, the National Natural Science Foundation of China (NSFC) via grant 62441614 and 62176078.
References
[1] Arora. S. Eyuboglu, A. Timalsina, I. Johnson, M. Poli, J. Zou, A. Rudra, and C. Re. Zoology: Measuring and improving recall in efficient language models, In The Twelfth International Conference Learning Representations, 2024.
[2] S. Arora, S. Eyuboglu, M. Zhang, A. Timalsina, S. Alberti, J. Zou, A. Rudra, and C. Rc. Simple linear attention language models balance the recall- -troouhptth tradeoff. In ICLR 2024 Workshop on Mathematical and Empirical Understanding of Foundation Models, 2024.
[3] S. Arora, B. Yang. S. Eyuboglu, A. Narayan, A. Hojel, I. Trummer, and C. Ré. Language models enable simple systems for generating structured views of heterogeneous data lakes, 2025.
[4] I. Beltagy, M. E. Peters, and A. Cohan. Longformer: The long-document transformer, 2020.
[5] Y. Bisk, R. Zellers, J. Gao, Y. Choi, et al. Piqa: Reasoning about physical commonsense in natural language. In Proceedings of the AAAI conference on artificial intelligence, volume 34, pages 7432-7439, 2020.
[6] T. Brown, B. Mann, N. Ryder, M. Subbiah, J. D. Kaplan, P. Dhariwal, A. Neelakantan, Shyam, G. Sastry, A. Askell, Agarwal, A. Herbert- Voss, G. Krueger, Henighan, R. Child, A. Ramesh, D. Ziegler, J. Wu, C. Winter, C. Hesse, M. Chen, E. Sigler, M. Litwin, S. Gray, B. Chess, Clark, C. Berner, S. McCandlish, A. Radford, I. Sutskever, and D. Amodei. Lan- guage models are few-shot learners. In H. Larochelle, M. Ranzato, R. Hadsell, M. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 1877-1901. Curran Associates, Inc., 2020.
[7] R. Child, S. Gray, A. Radford, and I. Sutskever. Generating long sequences with sparse transformers, 2019.
[8] C. Clark, K. Lee, M.-W. Chang, T. Kwiatkowski, M. Collins, and K. Toutanova. BoolQ: Explor- ing the surprising difficulty of natural yes/no questions. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long and Short Papers), pages 2924-2936, Minncapolis, Minnesota, June 2019, Association for Computational Linguistics.
[9] Clark, I. Cowhey, O. Etzioni, T. Khot, A. Sabharwal, C. Schoenick, and O. Tafjord. Think you have solved question answering? try arc, the ai2 reasoning challenge. ArXiv, abs/ 800.005457 2018.
[10] T. Dao. FlashAttention-2: Faster attention with better parallelism and work partitioning. 2023.
[11] Dao and A. Gu. Transformers are SSMs: Generalized models and efficient algorithms through structured state space duality. In R. Salakhutdinov, 7. Kolter, K. Heller, A. Weller, N. Oliver, Scarlett, and F. Berkenkamp, editors, Proceedings of the 41st International Conference on Machine Learning, volume 235 of Proceedings of Machine Learning Research, pages 10041-10071, PMLR, 21-27 Jul 2024.
[12) Devlin, M.-W. Chang, Lee, and K. Toutanova. BERT: Pre-training of deep bidirectional transformers for language understanding. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume I (Long and Short Papers), pages 4171-4186, Minneapolis, Minnesota. June 2019. Association for Computational Linguistics.
[13] Ding, S. Ma, Dong, X. Zhang, S. Huang, W. Wang, N. Zheng and F. Wei. Longnet: Scaling transformers to 1,000,000,000 tokens, 2023.

---

[14] J. Dong, B. Feng, D. Guessous, Y. Liang, and H. He. Flex attention: A programming model for generating optimized attention kernels, 2024.
[15] L. Gao, Tow, B. Abbasi, S. Biderman, S. Black, A. DiPofi, C. Foster, L. Golding, Hsu, A. Le Noac"h, h, H. Li, K. McDonell, N. Muennighoff, C. Ociepa, J. Phang, L. Reynolds, H. Schoelkopf, A. Skowron, L. Sutawika, E. Tang, A. Thite, B. Wang, K. Wang, and A. Zou. The language model evaluation harness, 07 2024.
[16] A. Gu and T. Dao. Mamba: Linear-time sequence modeling with selective state spaces. In First Conference on Language Modeling, 2024.
[17] A. Gu, K. Goel, and C. Re. Efficiently modeling long sequences with structured state spaces. In International Conference on Learning Representations.
[18] J. He, J. Qiu, A. Zeng, Z. Yang. J. Zhai, and J. Tang. Fastmoe: A fast mixture-of-expert training system. arXiv preprini arXiv:2 13262, 2021.
[19] D. Hendrycks, C. Burns, S. Basart, A. Zou, M. Mazeika, D. Song, and J. Steinhardt. Measuring massive multitask language understanding, 2021.
[20] C.-P. Hsich, S. Sun, S. Kriman, S. Acharya, D. Rekesh, F. Jia, and B. Ginsburg, RULER: What's s the real context size of your long-context language models? In First Conference on Language Modeling, 2024.
[21] S. Jelassi, D. Brandfonbrener, S. M. Kakade, and E. Malach. Repeat after Transformers better than state space models at copying. In R. Salakhutdinov, Z. Kolter, K. Heller, A. Weller, N. Oliver, J. Scarlett, and F. Berkenkamp, editors, Proceedings of the 41st International Conference on Machine Learning, volume 235 of Proceedings of Machine Learning Research, pages 111502-21521. PMLR, 21-27 Jul 2024.
[22] A. Q. Jiang, A. Sablayrolles, A. Mensch, C. Bamford, S. Chaplot, D. de las Casas, F. Bressand, G. Lengyel, G. Lample, I.. Saulnier, I. R. Lavaud, M.-A. Lachaux, P. Stock, T. L. Scao, T. Lavril, T. Wang, T. Lacroix, and W. E. Sayed. Mistral 7b, 2023.
[23] A. Katharopoulos, A. Vyas, N. Pappas, and F. Fleuret. Transformers are RNNs: Fast autoregres- sive transformers with linear attention. In H. D. III and A. Singh, editors, Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pages 5156-5165. PMLR, 13-18 Jul 2020.
[24] G. Lai, Q. Xie, H. Liu, Y. Yang, and E. Hovy. RACE: Large-scale ReAding comprehension dataset from examinations. In Proceedings of the 2017 Conference on Empirical Methods in Natural Language Processing, pages 785-794, Copenhagen, Denmark, Sept. 2017. Association for Computational Linguistics.
[25] H. Levesque, E. Davis, and L. Morgenstern. The winograd schema challenge. In Thirteenth international conference on the principles of knowledge representation and reasoning, 2012.
[26] C. Lockard, P. Shiralkar, and X. L. Dong. OpenCeres: When open information extraction meets the semi-structured web. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume I (Long and Short Papers), pages 3047-3056, Minneapolis, Minnesota, June 2019. Association for Computational Linguistics.
[27] E. Lu, Z. Jiang, J. Liu, Y. Du, T. Jiang, C. Hong, S. Liu, W. He, E. Yuan, Y. Wang, Z. Huang, H. Yuan, S. Xu, X. Xu, G. Lai, Y. Chen, H. Zheng, Yan, J. Su, Y. Wu, N. Y. Zhang, Z. Yang. X. Zhou, M. Zhang, and J. Qiu. Moba: Mixture of block attention for long-context Imss, 2025.
[28] F. Martin and C. Cundy. Parallelizing linear recurrent neural nets over sequence length. In International Conference on Learning Representations, 2018.
[29] T. Mihaylov, Clark, T. Khot, and A. Sabharwal. Can a suit of armor conduct electricity? a new dataset for open book question answering. In Proceedings of the 2018 Conference on Empirical Methods in Natural Language Processing, pages 2381-2391, Brussels, Belgium, Ocv.-Nov. 2018. Association for Computational Linguistics.
12

---

[30] OpenAl. Gpt-4 technical report, 2023.
[31] B. Peng, E. Alcaide, Q. Anthony, A. Albalak, S. Arcadinho, S. Biderman, H. Cao, X. Cheng, M. Chung, L. Derezynski, X. Du, M. Grella, K. Gv, X. He, H. Hou, P. Kazienko, J. Kocon, J. Kong, B. Koptyra, H. Lau, J. Lin, K. S. I. Mantri, F. Mom, A. Saito, G. Song, X. Tang, J. Wind, S. Wožniak, Z. Zhang. Q. Zhou, J. Zhu, and R.-J Zhu. RWKV: Reinventing RNNs for the transformer era. In H. Bouamor, J. Pino, and K. Bali, editors. Findings of the Association for Computatioonal Linguistics: EMNLP 2023, pages 1008-14077, Singapore, Dec. 2023. Association for Computational Linguistics.
[32] B. Peng, D. Goldstein, Q. Anthony, A. Albalak, E. Alcaide, S. Biderman, E. Cheah, X. Du, Ferdinan, H. Hou, P. Kazienko, K. K. GV, J. Kocoñ, B. Koptyra, S. Krishna, R. M. Jr., J. Lin, N. Muennighoff, F. Obeid, A. Saito, G. Song, H. Tu, C. Wirawan, S. Wožmiak, R. Zhang, B. Zhao, Q. Zhao, P. Zhou, J. Zhu, and R.-J. Zhu. Eagle and finch: Rwky with matrix-valued states and dynamic recurrence, 2024.
[33] B. Peng, R. Zhang, D. Goldstein, E. Alcaide, X. Du, H. Hou, Lin, J. Liu, J. Lu, W. Merrill, G. Song, K. Tan, S. Utpala, N. Wilce, J. S. Wind, T. Wu, D. Wuttke, and C. Zhou-Zheur Rwkv-7 "goose" with expressive dynamic state evolution, 2025.
[34] O. Press, N. Smith, and M. Lewis. Train short, test long: Attention with linear biases enables input length extrapolation. In International Conference on Learning Representations, 2022.
[35] A. Radford, K. Narasimhan, T. Salimans, I. Sutskever, et al. Improving language understanding by generative pre-training. 2018.
[36] P. Rajpurkar, R. Jia, and P. Liang. Know what you don't 't know: Unanswerable questions for SQuAD. In Proceedings of the 56th Anmual Meeting of the Association for Computational Lin- guistics (Volume 2: Short Papers), pages 784-789, Melbourne, Australia, July 2018. Association for Computational Linguistics.
[37] M. Roemmele, C. A. Bejan, and A. S. Gordon. Choice of plausible alternatives: An evaluation of commonsense causal reasoning. In 2011 AAAI Spring Symposium Series, 2011.
[38] K. Sakaguchi, R. L. Bras, C. Bhagavatula, and Y. Choi. Winogrande: An adversarial winograd schema challenge at scale. Communications of the ACM, 64(9)999-106, 2021.
[39] M. Sap, H. Rashkin, D. Chen, R. Le Bras, and Y. Choi. Social IQa: Commonsense reasoning about social interactions. In Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing (EMNLP-IJCNLP) IJCNLP),, pages 4463-4473, Hong Kong, China, Nov. 2019. Association for Computational Linguistics.
[40] I. Schlag. K. Irie, and J. Schmidhuber. Linear transformers are secretly fast weight programmers. In M. Meila and T. Zhang, editors, Proceedings the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pages 9355-9366 PMLR, 18-24 Jul 2021.
[41] N. Shazeer. Glu variants improve transformer, 2020.
[42] Z. Shen, Tao, L. Ma, W. Neiswanger, Z. Liu, H. Wang, B. Tan, Hestness, N. Vassilieva, D. Soboleva, and E. Xing. Slimpajama-dc: Understanding data combinations for Ilm training, 2024.
[43] Su, Y. Lu, S. Pan, A. Murtadha, B. Wen, and Y. Liu, Roformer: Enhanced transformer with rotary position embedding, 2023.
[44] Y. Sun, L. Dong, S. Huang, S. Ma, Y. Xia, J. Xue, J. Wang, and F. Wei. Retentive network: A successor to transformer for large language models, 2023.
[45] Y. Tay, D. Bahri, D. Metzler, D.-C. Juan, Z. Zhao, and C. Zheng. Synthesizer: Rethinking self-attention for transformer models. In M. Meila and T. Zhang, editors, Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pages 10192. PMLR, 18-24 Jul 2021.
13

---

[46] P. Tillet, H. T. Kung, and D. Cox. Triton: intermediate language and compiler for tiled neural network computations. Proceedings of the 3rd ACM SIGPLAN International Workshop on Machine Learning and Programming Languages, MAPI. 2019, page 10-19, New York, NY, USA, 2019. Association for Computing Machinery.
[47] A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, L. u. Kaiser, and I. Polosukhin. Attention is all you need. In I. Guyon, U. V. Luxburg. S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett, editors, Advances Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017.
[48] Welbl, N. Liu, and M. Gardner. Crowdsourcing multiple choice science questions, 2017.
[49] K. Wen, X. Dang, and K. Lyu. RNNs are transformers (yet): The key bottleneck on in-context retrieval. In The Thirteenth International Conference on Learning Representations, 2025.
[50] K. Yang. J. Ackermann, Z. He, G. Feng, B. Zhang, Feng. Q. Ye, He, and Wang. Do efficient transformers really save computation? In R. Salakhutdinov, Z. Kolter, K. Heller, A. Weller, N. Oliver, J. Scarlett, and F. Berkenkamp, editors, Proceedings of the 41st Inter- national Conference Machine Learning, volume 235 of Proceedings of Machine Learning Research, pages 55928 0228--5547. PMLR, 21-27 Jul 2024.
[51] S. Yang, J. Kautz, and A. Hatamizadeh. Gated delta networks: Improving mamba2 with delta rule. In The Thirteenth International Conference Learning Representations, 2025.
[52] S. Yang and Zhang. Fla: triton-based library for hardwarre-fffccient implementations of linear attention mechanism, Jan. 2024.
[53] J. Yuan, H. Gao, D. Dai, J. Luo, L. Zhao, Z. Zhang, Z. Xie, Y. X. Wei, L. Wang. Z. Xiao, Wang, C. Ruan, M. Zhang, Liang, and W. Zeng. Native sparse attention: Hardware-aligned and natively trainable sparse attention, 2025.
[54] M. Zaheer, G. Guruganesh, K. A. Dubey, J. Ainslie, C. Alberti, S. Ontanon, P. Pham, A. Ravula, Q. Wang, L. Yang, and A. Ahmed. Big bird: Transformers for longer sequences. In H. Larochelle, M. Ranzato, R. Hadsell, M. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 17283-17297. Curran Associates, Inc., 2020.
[55] R. Zellers, A. Holtzman, Y. Bisk, A. Farhadi, and Y. Choi. HellaSwag: Can a machine really finish your sentence? In Proceedings the 57th Annual Meeting of the Association for Computational Linguistics, pages 4791-4800, Florence, Italy, July 2019. Association for Computational Linguistics.
[56] B. Zhang and Sennrich. Root Mean Square Layer Normalization. In Advances in Neural Information Processing Systems 32, Vancouver, Canada, 2019,
[57] Y. Zhang, Y. Luo, Y. Yuan, and A. C.-C. Yao. Autonomous data selection with zero-shot generative classifiers for mathematical texts, 2025.
14

---

A Model Specifications and Pre-training Procedures
This work implemented and pre-trained multiple models with different architectures under various configurations to support experiments in different sections.
The first configuration involves models with approximately 110M (small scale) and 1.3B (large
scale, available only for certain architectures) parameters, with a sequence length of 2k. It includes
the following architectures: Transformer [6]. Transformer++, Transformer++ (W=256). Based [2]
Mamba-1 1 [6]. Mamba-2 [T]), RWKV-5 [32) RWKV-6 [2), RWKV-7 [33], DeltaNet [4)] Gated
DeltaNet Top-1 Element, Top-1 Chunk (Exact) and Top-1 Chunk (Approx.).
Transformer adopts the GPT 35 architecture, is based on FlashAttention-2"
2 [0]. and has a hidden dimension of 768, 12 layers and 12 attention heads for the small-scale version; while the large- scale version has a hidden dimension of 2048, 24 layers and 16 attention heads. Transformer++ introduces Rotary Embedding [43, GeGLU 4T and RMSNorm (56) over the Transformer, also built using FlashAttention-22, with identical model specifications as the Transformer. Transformer++ (W=256) incorporates window size of 256, retains the same model specification, and is also based on FlashAttention-22.
Based 22 is implemented according to the official open-source code with a hidden dimension of 768 and 15 layers for the small-scale version; among them, there are layers of window attention with 12 attention heads and a window size of 128, and layers of TaylorExp linear attention with 12 attention heads. Mamba-1 1 [16] is implemented based on the official open-source code" with a hidden dimension of 768. 24 layers and state size of 16 for the small-scale version; the large-scale version has a hidden dimension of 2048, 48 layers and the same state size of 16. Mamba-2 [ is implemented based on the official open-source code (URL is the same Mamba-1), with a hidden dimension of 768, 24 layers and state size of 128 for the small-scale version: the scage version has hidden dimension of 2048, 48 layers and the same state size of 128. RWKV-5/6/7 [32 33] are
implemented based on the official open-source codd with a hidden dimension of 768, 12 layers and 12 attention heads for the small-scale versions. DeltaNet 440 is based on Flash Linear Attention [2
with a hidden dimension of 768, 12 layers and 8 attention heads for the small-scale version. Gated DeltaNet [51 is also based on Flash Linear Attention, with a hidden dimension of 768, 12 layers, attention heads and a head dimension of 96 for the small-scale version.
Top-1 Element is implemented using PyTorct and the transformer library, matching the model specifications of the Transformer++. Top-1 Chunk (Exact/Approx: shares the same model specifica- tions as Transformer++, with chunk size of 128. Its sequence modeling part is implemented using Tritor [46). and its offset and index CUDA kernel is modified from the fastmoe 10
[18]
In addition, all of the above models the Mixtral 22 vocabulary, which contains 32,000 tokens.
Large-scale models involved in the training process were trained using the DeepSpeed ! open-source project. All models involving Rotary Embedding have a base of 10,000.
The second configuration involves models with approximately 135M parameters (small scale) and a sequence length of 100k. This mainly consists of architectures obtained by replacing the sequence modeling in RWK V-7. Specific models include: Transformer++, Transformer++ (W=256), Top-1 Chunk (Exact) and Top-1 Chunk (Approx.).
Transformer++ is based on FlashAttention-2,, while Transformer++ (W=256) uses FlexAtten- tior"" [2] [1]. The implementations of Top-1 Chunk (Exact) and Top-1 Chunk (Approx.) follow those described in the previous configuration. All four models have hidden dimension of 768, 12
" https://github..oo/ssaa spaces/mamba
"https//github.com/ffe
'htrps//pyytrrch.or org//
"hutp:://ittubb.c
'ttps://ggthhubbb
15

---

layers and 12 attention heads, with chunk size of 128 for the two Top-1 Chunk models. Additionally, RWKV-7 was implemented using Flash Linear Attention when the sequence length was set to 100k to improve time efficiency. Gradient checkpointing was applied to reduce GPU memory consumption across all models. All models use the Mixtral vocabulary with a size of 32,000 tokens. Models incorporating Rotary Embedding maintain base of 1, 1,000,000.
All models were pre-trained from scratch on the SlimPajama [42] dataset using language modeling objective across 100B tokens (mixed domain or limited domain). These settings apply to most of the main experiments and analysis experiments in this work (with the exception being the "distribution uniformity" analysis experiment, where some models were pre-trained on 25B tokens). More detailed test data partitions and other related settings are introduced in Section 2.1 and Section [3. Pre-trained models were trained using b precision and the Adam optimizer. Small-scale models used a learning rate of 2e-4, while large-scale models used a learning rate of Se-5. The optimizer's 8 and B3 values were set to 0.9 and 0.95, respectively. Learning rates for small-scale models were warmed up over the first 2,500 steps, while those for large-scale models were warmed up over the first 5,000 steps, followed by linear decay. The batch sizes for small-scale models were 256 (sequence length=2k) and (sequence length=100k), while that for large-scale models was 128 (sequence length=2k). All models were trained in distributed manner across 8 Nvidia Tesla A100 GPUs, totaling approximately 1,221 days of single-GPU equivalent training time.
B Evaluation Procedures
For language modeling evaluation, since all models in this work were pre-trained from scratch using the same vocabulary, we directly use the loss on the test set for evaluation.
For few-shot learning evaluation, general tasks include MMLU 19), OpenBookQA 22). ARC Easy &
Challenge 99, BoolQ [8]. RACE 24), SIQA (39], SCIQ 48 HellaSwag [55), COPA 337, PIQA (5]
WinoGrande 38 and Winograd 225 retrieval tasks include SWDE (26). SQUAD 3d and FDA 3
Among these, MMLU uses a 5-shot setting, while all others are evaluated a O-shot setting. To ensure
the stability and reproducibility of the evaluation results, we use the Im-evaluatioon-hares 13 15
open-source evaluation framework for evaluation.
For time efficiency evaluation (Table we compare the pre-training speed of different models with an input length of 100k tokens and a total batch size of 8. For inference speed comparison, we use an input length of 50k tokens and generate 5k tokens, with total batch size of 64.
For long- ung-sequence retrieval evaluation (Table 3). to ensure the stability and reproducibility of the results, we use the RULER 14 200 open-source evaluation framework for evaluation.
C Related Work
Our work is related to research on multiple topics, primarily including stateful sequence modeling architectures, sparse attention architectures and the analysis of sequence modeling architectures.
Stateful sequence modeling architectures have gradually evolved to address the inefficiency of standard self-attention [47] in handling long sequences. The core of these architectures lies in linear RNNs (28) and linear attention (23], with key characteristics being their ability to support both parallel and efficient training while also being expressible in the form of iterative state updates. Additionally, they offer certain advantages in terms of time and space efficiency compared to standard self-attention. Subsequent works have continuously introduced new structures or mechanisms to enhance the expressive power of stateful sequence modeling architectures. Early works incorporated data-indeppndent decays, such as S4 177. RetNet 44), RWKV-5 (32), among others. Later works added data-dependent decays, including Mamba-1 [16), Mamba-2 (T]) and RWKV-6 (32). The most recent developments have introduced the delta rule, exemplified by Gated DeltaNet [51] and RWKV-7 [33] Our work does not involve the development of new stateful sequence modeling architectures; instead, it primarily reveals existing deficiencies in base capabilities within established stateful sequence modeling architectures.
13 hhtps://githubbb "hhtps//ggihua..om//Heeutt http:/iiiinc.ooo///euteteeeeeeee htwarr///mmmmaaaaionnaaaaaaa "aapss//gtttb..co////iIDAAAU
16

---

Sparse attention architectures represent another approach to addressing the efficiency issues of standard self-attention. These works mainly achieve improved computational efficiency by reducing the number of terms involved in attention calculations. Early efforts were dominated by manually designed fixed sparsity patterns, including Sparse Transformer 77,, Longformer 4 BigBird (54), LongNet [13], and others. Recent works, such as NSA [33] and MoBA [27]. introduce dynamic block selection- based sparsification schemes that effectively balance computational efficiency and practical performance. The core of our work is to reveal how sequence modeling architectures influence base capabilittes. Although the improved architecture we ultimately propose also adopts a dynamic block selection-basee sparsification scheme, unlike previous works that focus on heuristic designs from an efficiency perspective and only evaluate in-domain performance, we begin our analysis from base capabilities, gradually building toward a well-foundee proposal for dynamic block selection. We further discover unique advantages of such architectures in out-of-distribution performance across domains and find that architectures like NSA and MoBA, which employ approximate selection strategies, may lead to limitations in out-of-distribution generalization capability. In addition, as mentioned in Section although discovered this architecture independently, no longer claim Top-1 Chunk (Approx.) as our primary contribution. Instead, our key contributions include Top-1 Element, Top-1 Chunk (Exact), analytical experiments and the released GPU kernels.
Works focusing on the analysis of sequence modeling architectures mainly concentrate on revealing specific capability deficits in stateful sequence modeling architectures, particularly in tasks such as retrieval [49]. copy (2I), associative recall (.. and dynamic programming [50]. These studies have successfully conducted experimental validations on synthetic datasets and synthetic tasks. Additionally, some of these works analyze theoretical differences in representational capabilities between RŇNS and Transformers. While our work also focuses potential defects in stateful sequence modeling architectures, we primarily investigate deficiencies in base capabilities rather than specific ones for example, whether such deficiencies can already be manifested in language modeling tasks. Moreover, our work does not delve into comparing strengths and weaknesses in model representation or retrieval capabilities but instead pays more attention to the architecture design itself. aiming to uncover truly effective components through changes in base capabilities and offering insights into architecture design.
D Algorithm Details
To achieve better time efficiency, we designed and implemented GPU kernels for the Top-1 1 Chunk Selection architecture.
Specifically, we designed Triton kernels for Top-1 I Chunk Selection architecture, with its forward pass detailed in Algorithm and its backward pass outlnned in Algorithm 2 Beyond these two core algorithms,, the implementation involves several other components. In both algorithms, F and I represent the offset and index, respectively. The offset F sequentially records the number of times each chunk is selected by querys, while the index I sequentially records the indices of all queries associated with each chunk. These can be obtained through naive Py implementation; we developed hardware-efficient CUDA kernels based on the logic of the naive implementation. Additionally, the chunk indices for both exact selection and approximate selection can also be acquired via naive PyTorch implementation. Similarly, we developed hardware-efficient Triton kernels following the logic the naive implementation. Additionally, we designed Triton kernels for the decoding stage, which can effectively accelerate the decoding process. The source code has been released at: https::/ hatg:////httbbb cow Lww/ixuzzzxx/AA
17

---

Algorithm Top-1 Chunk Selection Forward Pass
Require: Q, K, V Txxd
Fe
PERI I T chunk size Be, block size B, blocks K....., Kr. and VI..... Vr. of size Be x d each, divide
1: Divide K, V into T.
=
R.
T
F into Tc
T
scalars f1,..., ST..
B.
2: Initialize 0 (0)rxa RTxd M (-00)r e R L (0)r R R= (0)r E RR.
3:
for S T. do
4: Load Kj. v. from HBM to on-chip SRAM.
5: if = then
6:
On chip, Initialize fj-1 0.
7:
else
8:
Load / from HBM on-chip SRAM.
9:
end if
10: On chip, compute M ʃg JJ-1.
11:
Divide [f, into M,
blocks I,,... I, of size B, each.
12:
for 1 S < M, do
Load I, from HBM on-chip SRAM. Load Q; E RR by index I, E IR n. from Q in HBM to on-chip SRAM. On chip, compute S. Q: B.
14:
16:
On chip, compute mht rowmax Sig) E RRB.
17:
On chip, compute P., = exp(S.j mey) E RBB.BB..
18:
On chip, compute l, = rowsum P E R...
On chip, compute O = P, E RRB. Bd
O in HBM. 21: Write maj E IR by index I, R. to M in HBM.
Write O IR. Bd by index I, R.
22: Write lij E B, by index E BB. to L in HBM.
23:
end for
24:
end for
25: Divide K, V, O into T.
T
blocks Q Qr Ki,..., Kr. Vi,..., Vr, and
Be
O....., Or, of size B, x d each.
26: Divide M, into T.
blocks m my,.... my, and l,...,,y, of size B, each.
B
for < i do Load Q; K, Vi, 00, a.m.f. L, from HBM to on-chip SRAM.
28:
On chip, compute S, Q.K7 E B.B..
30:
On chip, compute m
(T) maz(m;, rowmax (S:)) E B.
On chip, compute P. = exp(S, - mm (T)) E B.B.B.
31:
32:
On chip, compute ( m.
rowsum P.) RBB
33:
On chip, compute ( diag( (, (T) (diag(e
(T)
)0 + P..... E BBxd
On chip, compute ri (T) = m; + togg(( ( (T)) E B.B
34:
Write O (T) E BRxd to O in HBM.
35:
36: Write E B. to R in HBM.
end for
38: Return O, R.
18

---

Algorithm 2 Top-1 Chunk Selection Backward Pass
Require: Q, V, 0, E RTxd RE B,.
IE IERT- RT-B, chunk size Be, block size
1: Divide K, V into T.
T Be
blocks KI...., Kr, and VI, V..... Vr, of size Bg x d each, divide
F into T
I
scalars f1,-..,
Be
2: Initialize dQ = (0)rxd E TRxx dK = (0)rxd Rxd dV = (0)rxd TRd
3: for S x Tg do
4: Load K, J. from HBM to on-chip SRAM.
5: Initialize dk, = R.x E Bexd dV, (0) n. R
6:
j jj=1 then
7:
On chip, Initialize ʃJ 0.
8:
else
9:
Load f-- from HBM to on-chip SRAM.
10: end if
11: On chip, compute M = f3 fj-1.
12:
Divide [f. into M,
B. blocks I....., .....,II,, IM. size B, each.
13:
for 1 < <is M, do
14:
Load I, from HBM to on-chip SRAM.
15:
Load Q, B. by index
B. from Q in HBM on-chip SRAM. Load O; E R B.xd by index I; E R B. from O in HBM to on-chip SRAM.
16:
Load dO; Rd by index E B. from dO in HBM on-chip SRAM.
18:
Load r; E E f by index I; E B. from R in HBM to on-chip SRAM.
19:
On chip, compute S, K BBB.B.
20:
On chip, compute Pú; = exp(S.) - r;) R..B.
21:
On chip, compute dV, + dV, + P,, do, Baxd
22:
On chip, compute dP;; ; = do,V7 E B.
23:
On chip, compute D; cowsum(d0) On chip, compute dS.j = Pij o (dP;) D.) B.B.B..
O,) E RRB
24:
25:
On chip, compute dK, + dK, ds [S]JQ.. Q E BBxd
26:
On chip, compute dQ, dS E R B.xd
27:
Write dQ; ER, Bd by index I; R B, dQ in HBM.
28:
end for
29: Write dk, RRBxdd dV, B. dK, dV in HBM.
30:
end for
31: Divide Q, V into T. B.
T
blocks Q Qr K Kr., VI.... Vr, and
O Or. of size Be x each.
Divide dQ, dk, dV, dO into To dOr, of size B. d each.
T
blocks dQ,..... dQr,, dk,..... dkr..
B.
dV,, .V.....,V,.. and dO,
33: Divide R into T.
T B.
blocks , rr. of size B, each.
34: for < SiST.4 do Load Q K, V., O,, dQ,, dK, dV, dO,, from HBM to on-chip SRAM.
35:
36:
On chip, compute S: = Q.K? E BB
On chip, compute P. exp(S: - r2) E R B.B
38:
On chip, compute dV; ; dV. WP. dO, RBB.xdd
39:
On chip, compute dP; dO, B..BB
40:
On chip, compute D; rowsunídO, 0,) E BBB
41:
On chip, compute dS; = (dP, E B.XB.
42:
On chip, compute dQ; <- dQ, dS, E R Boxd
43:
On chip, compute dK; + dK, + dS, Q E BB
Write dQ, E B. dk, B. dV, B B.xd dQ, dK. dV in HBM.
45: end for
Return dQ,dK,dl dV.
19

---

E More Results Under Mixed Domain Pre-Training Setting
Although our primary focus is on experiments under the limited domain pre-training setting, we also conducted preliminary experiments on the Top-1 Element Selection Architecture under mixed domain pre-training setting during the early stages of this work. The results, as shown Figures and and | indicate that the performance of the Top-1 Element Selection Architecture aligns with the conclusions drawn the paper, with no degradation base capabilities. This suggests that the usability of proposed architecture remains consistent across different pre-training settings.
Mixed Domain Pre-Training
Loss on Maxed Domain Teat Set (ID)
2.60
2.47
2.34
2.08
1.9
2.50
2.60
2.70
2,80
2,90
3.10
3.20
Pro-Training
Based
Mamba-1
Mamba-2
RWKY.
RWKV
RWKV.7
DeitnNet
Gated DeltaNet
Transformert+ (W-256)
Transformer
Transformert
Top-1
Figure Language modeling results of various sequence modeling architectures under mixed domain pre-training settings. (Model parameters=110M, pre-trained tokens 10kens=100B and sequence length-2k)
Mixed Domain Pre-Traine reearainnng
2.50
Loss on Mixed Domain (I)
2.35
1.30
1.75
1.60
Pre-Trinnng Loss
Transformer++
Transformert+ (W-256)
Mambu-1
Mumba-2
Top-1 Element
Figure Language modeling results of various sequence modeling architectures under mixed domain pre-training settings. (Model p 1.3B, pre- tre-trained tokens= 100B and sequence length=21
20

---

NeurIPS Paper Checklist
1. Claims
Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Question:
Answer: [Yes]
Justification:
In Abstract and Section 1.
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:
In 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]
21

---

Justification:
Not involving proof.
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 referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
Justification: In Appendix A and B.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a 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 a 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) the contribution is primarily a new algorithm, the paper should make clear how to reproduce that algorithm.
(b) the contribution is primarily a new model architecture, the paper should describe the architccture clearly and fully.
(c) If the contribution is new model (e.g., a large language model), then there should either 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 to describe the particular way they provide for reprodueibility. the case of closed-source models, may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers 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?
22

---

Answer: [Yes]
Justification: We have provided the most critical GPU kernel implementation code in the supplemental material, and will also make the source code publicly available.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips. cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to to reproduce the results. See the NeurlPS code and data submission guidelines (https
ssiinlolicy) 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, preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code 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: In Appendix and
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
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: We use default settings of open-source LLM evaluation frameworks to ensure the stability and reproducibility of the results.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are 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 capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
23

---

It should be clear whether the error bar is the standard deviation or the standard error of the mean.
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 figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures 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 to reproduce the experiments?
Question:
Answer: [Yes]
Justification:
In Appendix A.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU 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 as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported the paper (e.g., preliminary or failed experiments that didn't 't make it 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 http:///waragccccp//aa/a//aaaaaa
Answer: [Yes]
Justification: Yes.
Guidelines:
The answer NA means that the authors have 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., there a special consid- eration due to laws or regulations in 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: This work is a foundational research.
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 socictal 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.

---

The conference expects that many papers will be foundational research and not tied to 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 are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
We do not release models.
Guidelines:
The answer NA means that the paper poses no 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 to 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. Licenses for existing assets
Question: 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?
Answer: [Yes]
Justification: The corresponding references are provided for the existing assets used in this paper.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package 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 (e.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 Datasets.papersstt paperswi thcode. cww/ datasets package should be provided. For popular 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.
25

---

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: [NA]
Justification: N/A.
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: N/A.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing 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: N/A.
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
26

---

Question: 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.
Answer: [NA]
Justification: N/A.
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 for what should or should not be described.
27


Reviews:
[{"review_id": "DKsRZUsQFz", "reviewer": "Reviewer_jJs2", "summary": "This paper investigates non-transformer architectures developed in the past several years to determine which architectural design decisions they make drive their performance. It argues that limited domain ablation experiments - instead of the general purpose / mixed domain pretraining that's typically used - is important for ablation studies. They explore various design decisions made in Mamba and RWKV, as well as introduce and validate two simplifications of ideas common in non-transformer architectures and show that they work.", "questions": "1. Why does Section 4 mostly focus on Mamba and Section 5 mostly on RWKV?\n2. If one specifically wants to use a non-transformer architecture, what would you recommend they use? I think this is a question a lot of people have.\n3. How robust are these results to changing the training dataset?", "limitations": "Yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 2, "originality": 1, "strengths_and_weaknesses": "This is a good paper full of careful experiments. My main hesitency is around the fact that some of the experimental design choices seem poorly motivated in ways that could be the result of cherrypicking. To be clear I'm not saying I think the authors did, but that ad hoc design choices are non-desirable and one of the main reasons is that they can enable cherrypicking. The two examples that stand out the most are:\n1. Section 4 focuses on comparing to Mamba and Section 5 focuses on comparing to RWKV. It's not clear why this is. The sections are titled \"Top-1 Element Selection Architecture\" and \"Top-1 Chunk Selection Architecture\" which doesn't obviously connect to one vs the other. \n2. Very little space is spent justifying the decision to use a single data source to do the ablation studies. I would like to hear more about both the decision-making process and results on other datasets. For example, it could be that the specific dataset used is a major driver of the results. Alternatively, I could see people wondering if the results matter if they don't replicate on the types of training data people typically use.\nThese factors make it somewhat challenging to assess the significance of the results, as they make me wonder about the generality.\n\nThe idea of doing an ablation study of architecture design is not at all novel and entirely standard. Nobody should read the previous sentence as being a knock against this paper though.\n\nMany of the plots are unreadable in their current form due to small fonts and sizes. A particularly bad case is Figure 3, where at 250x scale I still can't read the key.", "was_revised": true, "final_justification": "I find the response to my questions about justifying their dataset decisions and what architecture they'd recommend wholly unsatisfying.", "rebuttal": "Thanks for your insightful review and valuable feedback!\n\nWe answer your questions below.\n\n---\n\n**Q1:** Section 4 focuses on comparing to Mamba and Section 5 focuses on comparing to RWKV. It's not clear why this is. The sections are titled \"Top-1 Element Selection Architecture\" and \"Top-1 Chunk Selection Architecture\" which doesn't obviously connect to one vs the other ... Why does Section 4 mostly focus on Mamba and Section 5 mostly on RWKV?\n\n**A1:** We apologize for any confusion caused—this may have been a misunderstanding, and we would like to clarify the following points:\n\nIn fact, most of the experiments in Section 4 (Top-1 Element Selection Architecture) and Section 5 (Top-1 Chunk Selection Architecture) are consistent, as both involve comparisons with the Mamba series of models. **Figures 6, Table 1, and Table 2 in the paper are shared between Sections 4 and 5**, presenting the results of comparing the three architectures proposed in these two sections with Mamba-1/2 and others. Due to an oversight in layout, these results were placed closer to Section 4, which may have misleadingly suggested they belonged exclusively to that section. We will improve this in subsequent versions.\n\nRegarding the experimental design related to RWKV-7, the situation is somewhat complex, and we would like to explain further:\n\n1. RWKV-7 was indeed not included as an independent baseline in Figures 6, Table 1, and Table 2 (i.e., the three architectures from Sections 4 and 5 were not compared with RWKV-7 in these figures/tables, though they were all compared with Mamba-1/2). This is because RWKV-7 incorporates many additional design elements beyond its sequence modeling architecture, which significantly improve its convergence speed. Since these enhancements are not attributable to improvements in the sequence modeling architecture itself, a direct comparison would be unfair and could obscure our assessment of the sequence modeling architectures' merits. Therefore, RWKV-7 was excluded from Figures 6, Table 1, and Table 2.\n\n2. To address this limitation, **we designed additional experiments in Figure 7 and Table 3**. The core idea here was to replace RWKV-7's original sequence modeling component with different sequence modeling architectures to enable a fair comparison with the original RWKV-7. Since the experiments in Figures 6, Table 1, and Table 2 had already demonstrated that the architecture in Section 4 was weaker in convergence compared to the two architectures in Section 5, we only transplanted the latter two architectures into RWKV-7 for these experiments—primarily to avoid unnecessary computational costs.\n\nTo summarize, we did not deliberately design Sections 4 and 5 to compare against different baselines. A more accurate description would be:\n\n**After comparing the architectures in Section 4 and Section 5 with Mamba-1/2, we determined that the two architectures in Section 5 were superior. We then transplanted the two architectures into RWKV-7 and compared them with the original RWKV-7.**\n\n---\n\n**Q2:** Very little space is spent justifying the decision to use a single data source to do the ablation studies. I would like to hear more about both the decision-making process and results on other datasets. For example, it could be that the specific dataset used is a major driver of the results. Alternatively, I could see people wondering if the results matter if they don't replicate on the types of training data people typically use. These factors make it somewhat challenging to assess the significance of the results, as they make me wonder about the generality ... How robust are these results to changing the training dataset?\n\n**A2:** We appreciate your concerns and would like to provide the following clarification:\n\nAll our architecture ablation analyses were conducted during the pre-training phase. Therefore, the \"training dataset\" you mentioned primarily refers to the pre-training dataset. Pre-training datasets have a distinctive characteristic compared to task-specific datasets: **most publicly available pre-training datasets likely share similar original data sources.** Tracing back to their origins, they are essentially crawled from the public internet, differing mainly in collection scale and data processing methods.\n\nFor instance, the SlimPajama dataset we employed consists mainly of data from CommonCrawl, C4, GitHub, Books, ArXiv, Wikipedia, and StackExchange. **The CommonCrawl and C4 subsets we used for pre-training are also core components of many public pre-training datasets** (either directly referenced or with highly overlapping sources), which was the primary reason we selected them for pre-training.\n\nGiven such intrinsic connections among different pre-training datasets, we chose not to switch between various pre-training datasets. Instead, we conducted relevant experiments using SlimPajama, which has already undergone cross-subset deduplication. **We believe that** selecting other pre-training datasets while performing similar cross-subset deduplication would yield largely comparable results in architecture ablation analyses.\n\n---\n\n**Q3:** Many of the plots are unreadable in their current form due to small fonts and sizes. A particularly bad case is Figure 3, where at 250x scale I still can't read the key.\n\n**A3:** We sincerely apologize for any inconvenience caused during the reading process. We fully appreciate your suggestions and will make improvements to the presentation of figures, tables, and text in the revised version.\n\n---\n\n**Q4:** If one specifically wants to use a non-transformer architecture, what would you recommend they use? I think this is a question a lot of people have.\n\n**A4:** While we hold our work in high regard, it has yet to undergo large-scale validation and thus may not currently represent the most robust option.\n\nRegarding non-Transformer architectures, our primary recommendation remains **Mamba-2**, as it demonstrates superior developmental momentum and maturity. Additionally, we would suggest considering **RWKV-7**, which incorporates numerous unique design elements that significantly enhance convergence (though many are engineering tricks unrelated to sequence modeling, their practical effectiveness is remarkable). Its block design—featuring separated MHA and FFN components similar to Transformers—makes it notably more amenable to customization or reintegration of recent advancements from Transformer architectures.\n\n---"}, {"review_id": "NmKkeoJHbR", "reviewer": "Reviewer_zo9Q", "summary": "The authors study the language modeling performance of popular architectures, including Transformer++, Mamba, RWKV, and DeltaNet, on an out-of-distribution (OOD) held-out set compared to an in-distribution (ID) held-out set. While some architectures converge faster and achieve a lower ID loss in the same number of gradient steps, none surpass the Transformer's ability to generalize (as measured by the OOD performance relative to the ID one). The authors conducted ablation studies on Mamba's decay, convolution, normalization, and the Transformer’s positional encodings, revealing that these design choices impact convergence speed but not generalization. Instead, the authors identify three mechanisms as critical for generalization: full sequence visibility, direct token-to-token relationship computation, and a non-uniform attention distribution. To validate this, they evaluated \"Top-1 Element\" and \"Top-1 Chunk\" Transformers, which restrict attention to the most significant token or chunk. The performance of these models aligned with that of the Transformer++, supporting the authors' claims about what is required for generalization.", "questions": "The experimental design is sound, and I believe the observations are valuable to the machine learning community. However, the lack of clarity and technical detail are holding the paper back.\n\n- Can you explicitly define base capabilities and specific capabilities?\n- Can you provide the exhaustive list of hyperparameters and the pseudo-code for the kernel?", "limitations": "Yes.", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 1, "significance": 3, "originality": 3, "strengths_and_weaknesses": "This work is timely and relevant considering the recent surge in popularity of alternative architectures and mechanisms to the original Transformer. The authors acknowledge the similarities between their “Top-1 Chunk” attention and Native Sparse Attention, prompting them to shift their focus on extensive ablations and alternative formulations to the Top-1 Chunk attention.\n\nThe pre-training methodology is good. The authors pre-trained a broad range of moderately sized models (~110M) on a large dataset (100B tokens from SlimPajama) with a decent sequence length (2k), providing a robust foundation for their observations. The inclusion of larger-scale experiments (1.3B models, 100k context length) and few-shot learning results is appreciated.\n\nDespite its strengths, the paper is limited by two significant weaknesses:\n- The central concepts of \"base capabilities\" and \"specific capabilities\" are never explicitly defined and are confusing. It seems that “base capability\" refers to a model's generalization performance as measured by the ratio of its OOD loss to its ID loss. Note that the OOD set is constructed with (English?) data from domains not included in the train set, such as arXiv and GitHub. The performance is measured on the same task as used for the pre-training. In opposition, \"specific capabilities\" seems to refer to the performance on downstream tasks. The paper would be substantially improved by explicitly defining these concepts early, and by potentially using more direct phrasing, such as \"generalization on the pre-training task\".\n- The current manuscript lacks the technical details required for reproducibility, including the depth and width of the different models, the optimizer, the scheduler, and all other hyperparameters. The pseudo-code for the custom kernel is also missing.\n\nMinor suggestions:\n- Line 33 and 55, I suggest removing “etc” as it is redundant with “such as”.\n- Line 64, I believe there is a missing “to” after “needs”.\n- In the figures, I suggest adding the identity line to all figures as the x and y-axes are not the same and change between plots.", "was_revised": true, "final_justification": "I am pleased with the clarifications. Please disregard my comment on the reproducibility, as the information was available at the time of submission (the appendix was included in the supplementary material). Given the authors’ commitment to including the clarifications and pseudo-code in the paper, I have increased my score from 4 to 5.", "rebuttal": "Thanks for your insightful review and valuable feedback!\n\nWe answer your questions below.\n\n---\n\n**Q1:** The central concepts of \"base capabilities\" and \"specific capabilities\" are never explicitly defined and are confusing. It seems that “base capability\" refers to a model's generalization performance as measured by the ratio of its OOD loss to its ID loss. Note that the OOD set is constructed with (English?) data from domains not included in the train set, such as arXiv and GitHub. The performance is measured on the same task as used for the pre-training. In opposition, \"specific capabilities\" seems to refer to the performance on downstream tasks. The paper would be substantially improved by explicitly defining these concepts early, and by potentially using more direct phrasing, such as \"generalization on the pre-training task\" ... Can you explicitly define base capabilities and specific capabilities?\n\n**A1:** We sincerely apologize for the confusion. You are absolutely correct—**we indeed focus on the out-of-distribution (OOD) generalization capability of language modeling as the base capability in our paper**, and we acknowledge the lack of explicit clarification in the earlier sections. We will revise this according to your suggestion.\n\n**Regarding base capabilities,** prior work on stateful sequence modeling primarily evaluated them through in-distribution (ID) generalization capability of language modeling or few-shot learning capability. However, we observed that these metrics fail to adequately distinguish between different sequence modeling architectures. Thus, we propose using OOD generalization of language modeling to measure base capabilities, which has yielded meaningful insights.\n\n**As for specialized capabilities,** we define them as skills such as retrieval, copy, associative recall, and dynamic programming. These are specific capabilities deliberately designed in prior stateful sequence modeling work for architecture analysis, distinct from conventional focuses like language modeling or few-shot learning. Hence, we categorize them as specialized capabilities.\n\n---\n\n**Q2:** The current manuscript lacks the technical details required for reproducibility, including the depth and width of the different models, the optimizer, the scheduler, and all other hyperparameters. The pseudo-code for the custom kernel is also missing ... Can you provide the exhaustive list of hyperparameters and the pseudo-code for the kernel?\n\n**A2:** We sincerely apologize for any confusion caused. **The technical details (including the depth and width of the different models, the optimizer, the scheduler, and all other hyperparameters) were included in the appendix.** However, the appendix was not embedded in the main paper but instead uploaded as a separate PDF file under **Supplementary Material (the official download link appears at the top of this page)**. Due to the tight deadline for the main paper submission (the Supplementary Material deadline was one week later), we made this arrangement. We will incorporate the appendix into the main paper if given the opportunity in future revisions.\n\nRegarding the kernel implementation, we acknowledge that we only provided the core design diagram without including pseudo-code, and we will follow your suggestion to add it (we attempted several pseudo-code display solutions, such as introducing relevant LaTeX packages or inserting formulas in Markdown code blocks, but **OpenReview doesn't fully support these features**. Therefore, we can only update this in the PDF version later). \n\nMeanwhile, **we are committed to releasing the kernel source code**, which has already been uploaded as **Supplementary Material (the official download link appears at the top of this page)**. We believe the executable source code will help compensate for any insufficient description of kernel details in the paper.\n\n---\n\n**Q3:** Line 33 and 55, I suggest removing “etc” as it is redundant with “such as” ... Line 64, I believe there is a missing “to” after “needs” ... In the figures, I suggest adding the identity line to all figures as the x and y-axes are not the same and change between plots.\n\n**A3:** Thank you very much for your valuable suggestions regarding the improvement of expressions in our paper. We will carefully revise the manuscript according to your comments.\n\n---"}, {"review_id": "FDImfkyhvx", "reviewer": "Reviewer_EQyK", "summary": "The paper studies the base capabilities of LLMs through OOD tests. And the paper discovers that full selection is a key property and designs the top-chunk selection architecture.", "questions": "See Weaknesses.", "limitations": "Yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 2, "significance": 2, "originality": 2, "strengths_and_weaknesses": "Strengths:\n1. The topic is an important topic: how to evaluate architecture designs besides the current pretraining method.\n2. The discovery is interesting. Top-chunk attention can achieve similar performances as full attention.\n\nWeaknesses:\n1. It's unclear what the definition of base capabilities is. The author seems to evaluate it by OOD tests. What is base capabilities? Why can OOD test evaluate the base capabilities? These two questions need to be answered and emphasized clearly in the main part of the paper. Currently both are vauge.\n2. The novelty is limited. As the authors mentioned themselves, MOBA and NSA already proposed top chunk selection architectures. The implementation is not novel enough.\n3. The connection between full-selection and the top-element training seems to be not strong enough. From my understanding, full-selection is an inference concept. So the authors need to inference with only one element. But the authors seem to training only backward through certain chunk/element but inference, from my understanding, uses the full sequence.", "was_revised": true, "final_justification": "The rebuttal is clear. The authors justify my questions. But due to the novelty limitation, I only increase my score by 1.", "rebuttal": "Thanks for your insightful review and valuable feedback!\n\nWe answer your questions below.\n\n---\n\n**Q1:** It's unclear what the definition of base capabilities is. The author seems to evaluate it by OOD tests. What is base capabilities? Why can OOD test evaluate the base capabilities? These two questions need to be answered and emphasized clearly in the main part of the paper. Currently both are vauge.\n\n**A1:** We sincerely apologize for any confusion caused and would like to provide the following clarification:\n\nPre-trained models can achieve performance improvements in areas such as few-shot learning through training solely on language modeling tasks. **This suggests that pre-trained models possess certain base capabilities before being specifically adapted to downstream tasks.**\n\nRegarding base capabilities, there exist various measurement approaches. Previous work on stateful sequence modeling primarily measured them through **in-distribution (ID) generalization capability of language modeling** or **few-shot learning capability.** However, we found these metrics insufficient for effectively distinguishing between different sequence modeling architectures. Therefore, we proposed using **out-of-distribution (OOD) generalization capability of language modeling** as a measure of base capabilities, which proves to be a more discriminative approach and has led to several meaningful findings.\n\n**To summarize,** we employ OOD generalization of language modeling as our measure of base capabilities. We argue that this approach is valid because it inherits aspects of previous language modeling-based measurement methods while **offering better discriminative power for comparing different sequence modeling architectures.** However, as you pointed out, we did lack specific explanations about this in the earlier sections of our paper. We will carefully follow your suggestion to make corresponding improvements.\n\n---\n\n**Q2:** The novelty is limited. As the authors mentioned themselves, MOBA and NSA already proposed top chunk selection architectures. The implementation is not novel enough.\n\n**A2:** As you mentioned, we acknowledge in the paper that our independently proposed **Top-1 Chunk (Approx.)** shares certain similarities with MOBA and NSA, and we explicitly state that due to the timing of submission, we do not position it as our primary contribution. The main contributions we claim in the paper include **Top-1 Element, Top-1 Chunk (Exact), architectural analysis conclusions, and released GPU kernels**, which better demonstrate the uniqueness of our work.\n\nFurthermore, regarding the Top-1 Chunk architecture, **we have made two additional contributions that go beyond the scope of MOBA and NSA:**\n\n1. Rather than designing this architecture directly from an efficiency perspective, we derived it through step-by-step analysis of architecture factors affecting capability, thereby revealing its deeper advantages.\n\n2. In Figure 7, we demonstrate that Top-1 Chunk still exhibits capability limitations due to approximate selection, which to some extent suggests that MOBA and NSA might share similar capability limitations. This provides an important starting point for future research.\n\nTherefore, **we believe our work still possesses significant unique aspects** that should compensate for the lack of primacy regarding the Top-1 Chunk selection architecture.\n\n---\n\n**Q3:** The connection between full-selection and the top-element training seems to be not strong enough. From my understanding, full-selection is an inference concept. So the authors need to inference with only one element. But the authors seem to training only backward through certain chunk/element but inference, from my understanding, uses the full sequence.\n\n**A3:** We appreciate your feedback and would like to clarify some points that may have caused confusion:\n\nRegarding your mention of **\"full-selection,\"** we understand this likely refers to our concept of **\"Full-Sequence Arbitrary Selection\"** in the paper. We suspect this terminology may have been misleading. To clarify, this does not imply physically selecting all sequence elements for weighted summation. Rather, it denotes the architecture's inherent capability to freely select any individual element from the sequence. As implemented in our Top-1 Element architecture, this mechanism actually chooses only the element with the maximum dot-product while disregarding others (i.e., no weighted summation occurs).\n\nWe should emphasize that the Top-1 Element architecture is not designed as a computationally efficient framework for either training or inference. It serves as a transitional architecture, with computational costs comparable to standard Multi-Head Attention, primarily for analysis purposes. **Although both training and inference involve selecting just one element from the full sequence, the architecture still requires maximum inner-product search across the full sequence.** Thus, the computations are performed over the entire sequence, only a single element is ultimately selected.\n\nThese explanations are offered based on our understanding of your concerns. However, if we have misinterpreted your original intention, we would be grateful for any additional clarification you could provide!\n\n---"}, {"review_id": "daNPb2AlSR", "reviewer": "Reviewer_z1L8", "summary": "This paper explores the design principles of sequence modeling architectures and their performance across various tasks, particularly in enhancing fundamental capabilities and handling long sequences effectively. The paper first introduces a key design principle: \"sequence modeling architectures must possess the capability of arbitrary selection across the entire sequence,\" aimed at preventing the degradation of foundational abilities.\n\nThey then provide a detailed description of two methods for implementing full key-value (KV) block selection: the Exact and Approximate (Approx.) variants.  This simplifies the process of calculating the dot product between the average vector of each key block and the query, thereby significantly reducing computational overhead.\n\nIn addition, the paper discusses differences in model performance under various pre-training setups. The study finds that although mixed-domain pre-training is beneficial for practical applications of large-scale language models, it transforms fundamental capability testing into in-distribution evaluation. This approach fails to reveal differences in foundational capabilities among different architectures during the early pre-training stages.\n\nFinally, the paper presents several experimental results showing that both Top-1 element selection and Top-1 block selection methods outperform state-of-the-art architectures in fundamental capabilities on both short sequences (2k) and long sequences (100k), while maintaining competitive time efficiency.", "questions": "See Weaknesses", "limitations": "Yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "**Strengths**\n\n+ Despite the unclear practical implications of OOD loss, especially since the authors avoid discussing its connection with retrieval capabilities, the paper provides a highly insightful perspective on revealing the differences in capabilities among various models and components. It identifies that supporting full-sequence arbitrary selection in sequence modeling architecture is the key design principle for preventing degradation of base capabilities.\n\n+ The authors propose a toy architecture to validate their viewpoint and implement Triton and CUDA operators, achieving efficient training on contexts up to 100k in length with minimal performance loss.\n\n**Weaknesses**\n\n+ The exact **nature of what OOD loss or base capabilities reflect about a model's abilities remains unclear**. While the authors mention in the introduction that strong base capabilities can lead to robust language modeling and few-shot learning abilities, the results in Figure 2 show that OOD Loss and few-shot performance **do not seem strongly correlated**.\n\n+ Further questions regarding base capabilities: **Does the author consider retrieval capabilities as part of base capabilities**? Different data types have varying impacts on different architectures. Linear attention models are generally thought to require additional context-dependent data to aid in learning retrieval capabilities. Would adding such data to both transformers and Mamba models reduce the performance gap between them?\n\n+ Although this point does not affect the validity of the article, it raises concern: Typically, pre-training lengths are 4k or 8k, but the maximum length used by the authors when testing full-sequence visibility was 2k. It is uncertain whether conclusions drawn from longer sequences would remain valid.\n\n+ Additionally, during the testing of full-sequence visibility, the **terminology \"window size\"** used by the authors could be misleading. Clarification is needed on whether it refers to the sample length during training or if a sliding window attention mechanism was employed.", "was_revised": true, "final_justification": "Yes", "rebuttal": "Thanks for your insightful review and valuable feedback!\n\nWe answer your questions below.\n\n---\n\n**Q1:** The exact nature of what OOD loss or base capabilities reflect about a model's abilities remains unclear. While the authors mention in the introduction that strong base capabilities can lead to robust language modeling and few-shot learning abilities, the results in Figure 2 show that OOD Loss and few-shot performance do not seem strongly correlated.\n\n**A1:** We appreciate your feedback and would like to clarify some potentially confusing expressions in our manuscript:\n\nIn our work, **\"base capabilities\"** refers to an attribute of pre-trained models that influences their ultimate performance when adapted to downstream tasks.\n\nWhile base capabilities can be measured through various approaches, prior work on stateful sequence modeling primarily evaluated it via **in-distribution (ID) generalization capability of language modeling** or **few-shot learning capability**. However, we found these metrics insufficient for revealing architecture differences between sequence modeling approaches. Consequently, we proposed using **out-of-distribution (OOD) generalization capability of language modeling** as a more discriminative measurement, which led to several meaningful discoveries.\n\nThus, your observation that **\"the results in Figure 2 show that OOD Loss and few-shot performance do not seem strongly correlated\" actually aligns with our expectations**. Precisely because few-shot performance proved inadequate for distinguishing base capabilities differences between architectures, we introduced OOD loss as an alternative metric. This doesn't reflect an intrinsic contradiction in base capabilities, but rather demonstrates that different measurement approaches reveal distinct levels of capability differences. **The OOD loss captures more profound architecture distinctions, which explains its weaker correlation with the metrics reflecting more superficial differences.**\n\n---\n\n**Q2:** Further questions regarding base capabilities: Does the author consider retrieval capabilities as part of base capabilities? Different data types have varying impacts on different architectures. Linear attention models are generally thought to require additional context-dependent data to aid in learning retrieval capabilities. Would adding such data to both transformers and Mamba models reduce the performance gap between them?\n\n**A2:** Whether to consider **\"retrieval capability\"** as part of the base capabilities may be a matter of definition. **In our paper,** we did not define it as a base capability but rather treated it as a specialized skill, though defining it as base capability would not be incorrect.\n\nOur decision not to regard **retrieval capability** as being on par with **language modeling OOD generalization** as a base capability stems primarily from the observation that **it appears more easily satisfied than the latter**. The key evidence lies in our 100k sequence-length experiments, where Top-1 Chunk (Approx.) and Top-1 Chunk (Exact) demonstrated nearly identical retrieval performance (Table 3), yet Top-1 Chunk (Approx.) still showed significant degradation in language modeling OOD generalization (Figure 7). Therefore, we prefer using language modeling OOD generalization as our measure of base capabilities, which also has the advantage of being more closely aligned with the model's original pre-training objective.\n\nYou raised an insightful point about **different data types affect architectures differently** - this indeed presents an alternative perspective where data could be treated as an architecture-dependent dynamic variable. However, **much current research still tends to view Transformer as the ideal reference point for capabilities, considering it a weakness when other efficient sequence modeling architectures require specific data to complete certain capabilities**. Our work follows this conventional approach, aiming to design architectures that match Transformer's data efficiency without capability degradation. Thus, regarding your question about whether adding retrieval-oriented auxiliary training data to both Transformer and Mamba would narrow their retrieval performance gap: while this might help, we would still interpret it as reflecting Mamba's inherent limitations. **That said**, we fully acknowledge that your perspective of treating data as an architecture-dependent dynamic variable is more profound, and we hope to explore this direction in future work.\n\n---\n\n**Q3:** Although this point does not affect the validity of the article, it raises concern: Typically, pre-training lengths are 4k or 8k, but the maximum length used by the authors when testing full-sequence visibility was 2k. It is uncertain whether conclusions drawn from longer sequences would remain valid.\n\n**A3:** We understand your concerns and would like to provide the following clarification:\n\nThe majority of our pre-training experiments were conducted with 2k sequences, including the relevant experiments that established the principle of full-sequence arbitrary selection. However, **we additionally performed pre-training experiments with 100k sequence lengths (Figure 7)**. The out-of-distribution experimental results remain consistent with the principles summarized from the 2k sequence experiments, with no violations observed. **We believe this supplementary experiment helps address the limitations of using only 2k sequence lengths to some extent.**\n\n---\n\n**Q4:** Additionally, during the testing of full-sequence visibility, the terminology \"window size\" used by the authors could be misleading. Clarification is needed on whether it refers to the sample length during training or if a sliding window attention mechanism was employed.\n\n**A4:** We sincerely apologize for any confusion caused. Please allow us to clarify:\n\nThe term **\"window size\"** used in testing full-sequence visibility refers to the **sliding window size**, while the sample length during training remains unchanged at 2k. We will revise the relevant expressions according to your comments to prevent any potential misunderstandings.\n\n---"}]

Author rebuttal:

Meta-review: This paper studies how neural architecture choices affect the base capabilities of LLMs. The core claim is that mixed-domain pretraining often hides architectural differences early, and that limited domain pre-training setting with out-of-distribution testing is a more discriminative way for probing. With this setup, the authors find that stateful architectures (lie Mamba/RWKV) show notable degradation relative to Transformers. They also propose a design principle:  "a sequence modeling architecture need possess full-sequence arbitrary selection capability".
They validate this by introducing two intentionally simple Top-1 Transformers, which basically restricts attention to the most significant token or chunk. The Top-1 Transformers can match Transformer++ on OOD generalization at small scale (≈110M params, 2k context) and remain competitive at larger scale (≈1.3B, 100k context).

Strengths:
1. The topic is important and the conclusion is insightful for LLM pretrain community
2. The FSAS principle is simple, testable, and of high practical relevance for ongoing non-Transformer efforts.
3. Experiment design is good, it controlled training budgets, early-stage OOD evaluations, and scaling checks.

Weaknesses:
1. One major concern is presentation. I feel difficult to understand the paper at my first try. Especially when the authors tried to introduce the top-1 models. Many sections begin with the lack of motivations.
2. The main conclusions rely on a limited-domain pretraining + OOD evaluation protocol. Although authors claimed that mixed-domain pretraining creating the illusion of base capabilities. It still raised the question that how generalizable the conclusion of this setting can be applied to real-world LLMs.

Rationale for Accept:
Even if majority community members continue to favor mixed-domain pretraining, the presented framework and results are still valuable to architecture ablation studies. Also, two reviewers increased their scores following clarifications and added results.
Decision: Accept (poster)