Title: How Does Sequence Modeling Architecture Influence Base Capabilities of Pre-trained Language Models? Exploring Key Architecture Design Principles to Avoid Base Capabilities Degradation
Venue: NeurIPS 2025
Accepted: yes
Keywords: Pre-trained Language Models

Base Capabilities

Sequence Modeling
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.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS AND WEAKNESSES:\n**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.\n\nQUESTIONS:\nSee Weaknesses\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-06-16"}

comments: 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.

STRENGTHS AND WEAKNESSES:
Strengths:
1. The topic is an important topic: how to evaluate architecture designs besides the current pretraining method.
2. The discovery is interesting. Top-chunk attention can achieve similar performances as full attention.

Weaknesses:
1. 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.
2. The novelty is limited. As the authors mentioned themselves, MOBA and NSA already proposed top chunk selection architectures. The implementation is not novel enough.
3. 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.

QUESTIONS:
See Weaknesses.

LIMITATIONS:
Yes
TITLE: Official Review
CLARITY: 2
IMPACT: 2
ORIGINALITY: 2
DATE: 2025-06-29

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS AND WEAKNESSES:\nThis 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.\n\nQUESTIONS:\nThe 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?\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 1, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-01"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS AND WEAKNESSES:\nThis 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.\n\nQUESTIONS:\n1. 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?\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 2, "IMPACT": 2, "ORIGINALITY": 1, "DATE": "2025-07-09"}