SAS: Simulated Attention Score
Authors: Chuanyang Zheng
Jiankai Sun
Yihang Gao
Yuehao Wang
Peihao Wang
Jing Xiong
Liliang Ren
Hao Cheng
Janardhan Kulkarni
yelong shen
Zhangyang Wang
Mac Schwager
Anderson Schneider
Xiaodong Liu
Jianfeng Gao
Venue: NeurIPS 2025 poster ()
Keywords: sequence modeling
language models
model architecture
Transformer
Primary area: deep_learning
TL;DR: We propose SAS to simulate larger attention head numbe and hidden size per head for better performance, keeping the original model size.

Abstract: The attention mechanism is a core component of the Transformer architecture.
Various methods have been developed to compute attention scores, including multi-head attention (MHA), multi-query attention, group-query attention and so on. We further analyze the MHA and observe that its performance improves as the number of attention heads increases, provided the hidden size per head remains sufficiently large. Therefore, increasing both the head count and hidden size per head with minimal parameter overhead can lead to significant performance gains at a low cost.
Motivated by this insight, we introduce Simulated Attention Score (SAS), which **maintains a compact model size while simulating a larger number of attention heads and hidden feature dimension per head.** This is achieved by projecting a low-dimensional head representation into a higher-dimensional space, effectively increasing attention capacity without increasing parameter count. Beyond the head representations, we further extend the simulation approach to feature dimension of the key and query embeddings, enhancing expressiveness by mimicking the behavior of a larger model while preserving the original model size.
**To control the parameter cost, we also propose Parameter-Efficient Attention Aggregation (PEAA).**
Comprehensive experiments on a variety of datasets and tasks demonstrate the effectiveness of the proposed SAS method, achieving significant improvements over different attention variants.

SAS: Simulated Attention Score
Chuanyang Zheng'
Jiankai Sun
Yihang Gao Yuehao Wang
Peihao Wang
Jing Xiong
Liliang Ren'
Hao Cheng'
Janardhan Kulkarni
Yelong Shen'
Zhangyang Wang
Mac Schwager
Anderson Schneider'
Xiaodong Liu
Jianfeng Gao
Morgan Stanley
Stannfodd
Microsoft Research ANUS *UT Austin
HKK
Abstract
The attention mechanism is core component of the Transformer architecture. Various methods have been developed to compute attention scores, including multi- head attention (MHA), multi-query attention, group-query attention and on. We further analyze the MHA and observe that its performance improves as the number of attention heads increases, provided the hidden size per head remains sufficiently large. Therefore, increasing both the head count and hidden size per head with minimal parameter overhead can lead to significant performance gains at low cost. Motivated by this insight, we introduce Simulated Attention Score (SAS), which maintains a compact model size while simulating larger number of attention heads and hidden feature dimension per head. This achieved by projecting a low-dimensional head representation into higher-dimensions space, effectively increasing attention capacity without increasing parameter count. Beyond the head representations, further extend the simulation approach to feature dimension the key and query embeddings, enhancing expressiveness by mimicking the behavior of larger model while preserving the original model size. To control the parameter cost, we also propose Carameter-Ffffiicee (PEAA). Comprehensive experiments on a variety of datasets and tasks demonstrate the effectiveness of the proposed SAS method, achieving significant improvements over different attention variants.
Attention Aggregation
Introduction
The Transformer architecture [77] has emerged as the predominant framework modern machine learning, demonstrating remarkable success across diverse domains. Transformer-bas models consistently achieve state-of-the-an performance numerous natural language processing tasks, including machine translation [5, 89], question answering [96, 27, 4], and commonsense reason- ing [71, 106, 99, 92, 56]. The fundamental components of the Transformer comprise attention mech- anisms and feed-forward networks, with attention score computation serving as critical element. Beyond the original multi-head attention (MHA) formulation [77], researchers have developed various efficient sparse attention variants, such sliding window approaches (e.g.. Streaming LLMs [80]), linear transformers (e.g., Performer [14]), and sparse attention mechanisms (e.g., Reformer and Sparse Sinkhorn transformer [73, 21)).
Modern architectures employ diverse strategies for attention score computation. The standard MHA approach [77] maintains identical head counts across query, key, and value embeddings. To optimize memory efficiency, subsequent innovations have introduced parameter sharing schemes: Multi Multi-Query Attention (MQA) [61] and Group- Queuy- Attention (GQA) [3] share key and value projections across attention heads. The Multi-Latent Attention (MLA) approach [40] adopts low-rank compression
"Contact Email: cyzhengme @ gmail.ccmm
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

for key and value projections, caching only latent representations. More recently, Tensor Product Attention (TPA) [100] has introduced tensor decomposition techniques for compact representation of queries, keys, and values. We further analyze the MHA in Figure 1, and we find that the MHA performance gradually increases when the attention head number increases, as long as the hidden size per head is not too small. Hence, by increasing the number of attention heads and the hidden size per head at a minimal cost in terms of additional parameters, we can expect a substantial improvement in performance.
Building on these observations and insights, we propose Simulated Attention Score (SAS) to enhance Transformer performance through simulated attention head and feature expansion. Traditional query embeddings have dimensionality [B, T, H. D), where B represents the batch size, T is the sequence length, H denotes the head count, and D is the feature dimension per head. Our approach applies linear projections with nonlinear activations along the head dimension (H) to effectively increase the number of attention heads to 1l with H > H. Similarly, we also simulate more feature dimensions (D) for query/key embedding for the attention score computation. The proposed method effectively and efficiently expands both the head and feature dimensions, resulting in a query representation of shape [B, T, H, D). where the expanded dimensions H and ò exceed the original dimensions H and D, respectively. This expansion simulates greater number of attention heads and higher hidden feature dimensionality per head, thereby enhancing model expressiveness while preserving a compact overall model size. Additionally, we propose the parameter-efffciee attention aggregation control parameter size and computation. To summarize, our main contributions are the following:
1. By analyzing previous attention variants, we empirically observe that the number of attention heads plays critical role in the Transformer performance, while more attention heads can lead better performance, as long as the hidden size per head is not too small.
2. Based on the above observation, we propose SAS use a moderately sized model to simulate the attention mechanism of a larger model, effectively increasing attention heads and hidden dimension per head. To control the parameter cost, we also propose Errameter-Efficientt Attention Aggregation (PEAA).
3. We conduct extensive experiments on various datasets, tasks, and hyperparameter settings to validate the effectiivness of the proposed method, under different training lengths, model sizes, and datasets. Furthermore, we pretrain the model with SAS on the large-scale datasets and evaluate its performance on a wide range of downstream tasks, further showing the effectiveness of the SAS method.
2 Related Works
Attention Mechanism in Transformer. The Transformer architecture, leveraging Multi-Heas Attention (MHA) [77] for effective long-range dependency modeling [84, 85], suffers from high memory footprint and bandwidth demands during training. To alleviate these issues, several efficient attention mechanisms have been proposed. Multi-Query Attention (MQA) [61] reduces memory usage by sharing key and value heads among multiple queries. Grouped-Query Attention (GQA) [3] addresses MQA's Krawbacks by grouping queries with shared key and value heads. Multi- Latent Attention (MLA) [40] optimizes inference through low-rank key and value representations, surpassing MHA with a smaller KV cache. Multi-matrix Factorization Attention (MFA) [33] builds upon MQA by factorizing the query matrix. Tensor Product Attention (TPA) [100] achieves performance improvements and KV cache reduction by employing contextual tensor decomposition for query, key, and value representations.
Linear Attention and RNN. Linear recurrent language models have emerged as a compelling alternative due to their training and inference efficiency. The evolution of these models has progressed from early linear RNNs with data- iita-independent decay, exemplified by S4 [26], S5 [65]. LRU [48], RWKV4/5 [51, 52], and RetNet [70], towards more sophisticated architectures incorporating data- dependent decay, such as HGRN [55], Mamba [25, 19]. and GSA [101]. This transition highlights the effectiveness of gating and forgetting mechanisms (termed "selective mechanisms" in Mamba), principle inspired by classical gated RNNs. A key distinction in modern forget gates, unlike those in I. LSTMs [24], is their independence from the previous hidden state, relying solely on input data and enabling efficient parallelization [55]. Further advancements include Longhorn [41] and

---

TTT [69], which reformulate state space learning as a gradient- based online optimization, and Gated DeltaNet [91], which enhances a linear RNN's expressiveness through gating while preserving training efficiency. [63] provides a comprehensive comparison of expressivity between Attention- based architectures, Linear attention mechanisms, and Recurrent Neural Networks.
Non-Linear Mapping. The non-linear mappings typically have a higher expressiveness than linear mappings. Neural networks are well-suited for implementing such non-linear transformations. The most basic neural network architecture for this purpose the multilayer perceptron (MLP) [53], which is composed of linear transformations and non non-linear activation functions. In certain data structures and configurations, the MLP can replaced with convolutional neural network (CNN) [36, 35] to better capture spatial or local patterns. Common activation functions include ReI.U[2], Sigmoid [29]. Tanh, Leaky ReLU [88], among others.
3
Method
3.1
The Number of Attention Heads Plays an Important Role
To investigate the impact of varying head counts on model performance, we conduct experiments with MHA models using different numbers of
Arxiv Dataset (Length 512)
11
MHd-Heasd]
Validation perples perplexity
MHA-Heas2
attention heads, ranging from 1 head up to 96 heads, alongside a hidden feature of 768 dimen- sions. The model with 12 heads contains 125
10-
MHAAHead4
9-
MHA Hesd12
MHAAHesss2 24
Mith-teeed 48
million parameters. The results of these exper- iments are illustrated in Figure 1. When only single attention head (1 head) used, the model achieves perplexity score of 6.08. However, as the number of attention heads incrementally increases, the model's performance improves, reaching an optimal perplexity of 5.82 when the head count matches the original configuration of 12 heads. Beyond this point, further increasing the number of attention heads leads to a gradual degradation in performance, as evidenced by the rising perplexity scores. Therefore, more attention heads and larger hidden sizes per head may further improve the performance.
MHA
MIHA-Hes090
10K 15K 20K 30K 35K 40K 45K 50K Validation Step
Figure 1: The performance of MHA with different head numbers on the Arxiv dataset, with a 125M model (hidden size is 768) and training length 512.
3.2 Simulated Attention Score
Inspired by the above observations, we propose Simulated Attention Score (SAS), a method that simulates the behavior of a larger Transformer model with more attention heads and larger hidden size per head, using smaller and more efficient model. Our approach maintains a moderate model size while significantly enhancing representational expressiveness. The implementation is presented in Appendix L.
Initial query, key and value embeddings. Theoretically, our method is compatible with various existing attention mechanisms, including MQA, GQA, MHA, and TPA (e.g. the result presented in Appendix J). to construct the initial query and key tensors. Among these, MHA offers relatively stable performance, and it is widely used in previous works [77, 34]. Therefore, we adopt MHA as the foundation for initializing the query and key embeddings in our design.
Expanding feature dimensions enhances model performance. In addition to simulating more attention heads, we extend our method to include feature dimension expansion, as the feature dimension also plays critical role in shaping attention representations. We apply non-linear mapping to both the query and key tensors to simulate a larger feature space, As result, SAS enhances attention modeling through dual expansion, which increases both the number of heads and the feature dimension. Considering computational efficiency, we apply feature dimension expansion only to the key and query tensors, leaving the value tensor with its unchanged feature dimension.
3

---

3.3 Simulating Attention Heads via Non-linear Mappings
In this section, we present the step-by-step procedure for simulating a greater number of attention heads using non-linear mapping. Our goal is to enrich the model's 's expressiveness by mimicking the behavior of wider multi-heud attention mechanism, while keeping the actual number of parameters and computational cost moderate.
Given input query tensor Q RBxIxHxD, where B denotes the batch size, T is the sequence length, H represents the number of attention heads, and D is the feature dimension per head. This tensor structure follows the standard formulation in multi-head attention, where each head processes a D-dimensional feature vector independently.
We first reshape the query tensor to combine the batch size, sequence length, and feature dimensions, isolating the head dimension for further processing:
Reshapee Q, (T-T. D) RIB-T-DXXH
Our goal is to simulate a larger number of attention heads to enhance the model's expressiveness. To achieve this, we apply transformations along the head dimension. This reshaping step is crucial, allows us to treat the head dimension as the primary axis for non-linear mapping while collapsing the other dimensions into a single batch-like dimension.
The transformation applied to the head dimension consists of a (two-layer) multilayer perceptron (MLP):
Q Linear 2n (Qo),
Q2 = (QI),(Q:;;
Q Linear? Qs (Q2) + Q1
where Linear Qn RR + R and Linear, R + R are general linear transformation layers. Here, H denotes the original number of attention heads, and H represents the expanded (simulated) head dimension. To improve training stability and mitigate the risk of vanishing gradients introduced by the transformation, parameter cost introduced by this transformation is negligible, as both the original head dimension H and the expanded head dimension / are significantly smaller than the feature dimension D. Conse- quently, in auto-regressive Transformer models, the majority of the parameters are concentrated in the processing of feature vectors rather than in the introduced transformations along head dimensions.
incorporate residual (skip) connection from Q to Q. The additional
based transformations, convolution is another widely adopted approach for dimensional expansion, particularly in computer vision tasks. In our setting, we treat the head dimension as the channel and the feature dimension structured ID representation, analogous to multi-channel feature maps. This interpretation allows us to apply convolutional operations effectively. This perspective enables us exploit local patterns across attention heads and features, introducing an alternative mechanism for enhancing representational capacity with feature-aware processing of heads. We begin by reshaping the query tensor as follows:
In addition
Qo = Reshape Q (B H, E ((B-))xHDD
This reshaping step results in stack of B T multi-channel features, each of shape H X D, corresponding to individual time steps across the batch. We can then apply convolutional layers with ReLU non-linearity to expand the representational space of these features, enabling more expressive modeling of attention head interactions. The simulation of an expanded head dimension is performed as follows:
Q Conv Qa (Qo),
(5)
Q2 = Re)U(Q:),
(6)
Q Conv (Q2) + Q
(7)
where Conv HxD + p and Conv? : B + A are standard convolution
layers. Here, H denotes the original number of attention heads (interpreted as channels) and A the expanded (simulated) head dimension. Similarly, introduce the ReLU nonlinear activation function

---

and the residual connection to enhance representational capacity and improve training stability. The channel expansion is efficiently implemented using convolutional operators. Moreover, convolution naturally facilitates feature sharing across different heads, promoting effective communication both spatially and temporally (i.e., across channels). Note that the convolution operation with a kernel size of one is equivalent to the MLP applied along the channel dimension. Therefore, in our implementation, we adopt convolutional operators for head dimension expansion due to their efficiency and flexibility.
Head expansion for query tensor Q, key tensor K, and value tensor V. In the Transformer architecture, each attention head operates on a distinct query, key, and value tensor. Accordingly, the head expansion technique described above is applied to the query (Q), key (K), and value tensors (V), using the same expanded head dimension H. The additional parameters and computational overhead introduced by this head expansion mechanism are minimal, as both the original and expanded head dimensions (H and H) are significantly smaller than the feature dimension D. As a result, the overall parameter size and computational cost remain comparable the original Transformer computation. A detailed comparison of the parameter sizes between the vanilla Transformer and our proposed head-simulatn simulation method provided in the experimental section.
3.4 Simulating Features via Non-linear Mappings
In addition to the number of attention heads, the feature dimension also plays critical role in determining the capacity of Transformer models. Building upon the concept of attention head simulation, we extend the idea to simulate larger feature dimensions, aiming to enhance model expressiveness while preserving moderate model size and computational cost. After performing head dimension simulation, we apply a similar technique to the feature dimension. Taking the query tensor as example, first reshape the tensor to isolate the feature dimension for processing:
Qo Reshape Q (B-T
T H, D)),
(8)
Q1 Linear Q/ 'QQo)
(9)
Q2 = ReLU(Q:),
(10)
Q Linear Q (Q2) + Q1:
(11)
where Linear 01
TR. : D and Linear RDO - D are standard linear transformation layers.
Here, D denotes the original feature dimension, and D represents the expanded (simulated) feature dimension. This procedure mirrors the head dimension simulation described earlier, with the key difference being that the transformation is now applied along the feature dimension instead of the head dimension. As before, we incorporate a ReLU non-linear activation and a residual connection to improve both representation capacity and training stability.
Feature dimension expansion for the query tensor Q and key tensor K, excluding the value tensor V. We apply feature dimension expansion (simulation) to the query tensor Q and key tensor K, which form the core components of the attention mechanism. However, we do not apply this expansion to the value tensor V. This because the value tensor typically determines the output token representation, which is subsequently passed through feedforward MLP layer. Expanding the value dimension would directly increase the token embedding size, thereby leading to a significant increase in the parameter size and computational cost of the feedforward layer. This undesirable for maintaining model efficiency. Therefore, to balance expressiveness and efficiency, we restrict feature dimension expansion only Q and K.
3.5 Parameter-Effiiienn Attention Aggregation (PEAA)
To control the parameter cost, we utilize Parfmeter-Efficceett Attention Aggregation (PEAA). Denote the expanded key, query, and value tensors at i-th head by K,, Q., and V,, respectively. The output from the i-th head after the attention is obtained by
h = V Softmax (KTO.)
5

---

Then, the attention features are aggregated by
8/H
Output R/H (b ". b Wo
Σ
(12)
for i 1,2. / HH, where we H/H as positive integer, and W ((H-D)x(H-D) is responsible for a linear projection. In general, the attention feature aggregation is conducted by concatenating the output h, among all expanded heads followed by a linear projection. However, due to the expanded (simulated) heads, it will introduce additional parameters and computation overhead. Here, we adopt the PEAA method that aggregates the attention features by averaging over groups of heads, as shown in Equation (12).
4 Experiment
Baselines. We evaluate the proposed SAS against range of established baselines, including MHA [77], MQA [61], GQA [3], MLA [40] and TPA [100].
Datasets. Our analysis involves training language models on the Arxiv and Books3 datasets, which are frequently used benchmarks for evaluating model performance [54, 12, 37, 20]. Also, we train the model on the large-scale dataset FinWeb-Edu [45].
Experiment settings. Initially, we compare SAS with other baselines at training lengths 512, and 1024, with model size 125M decoder-only Transformers [7], whose configuration is shown in Appendix C. Subsequently, we evaluate the performance of larger model sizes with 350M and 2.7 B. We also train on large-scale datasets and evaluate downstream tasks.
Anxiv Dataset (Length 512)
Book Dataset (Length 512)
MHA
60
Validation perpiexity
MGA
perplexiy
MBA
G0A
MLA
50
M.A
TPA
SAS
30
5k
10K
15K
20K
25K 30K 35K 40K
45K
50K
5k
10K
15K
20K 25K 30K 35K 40K 45K 50K
Validation Step
Validation Step
Anxiv Dataset (Length 1024)
Books3 Dataset (Length 1024)
7.5
Validation Merplexdyy..
45
MHH
perplexity
7.0
MQA
GQA
40
GOA
6.5
MLA
M.A
335
TPA
6.0
30
25
5k 10K 15K 20K 25K 30K 35K 40K 45K 50K
10K 15K 20K 25K 30K 35K 40K 45K 50K 5k
Validation Step
Validation Step
Figure 2: The performance of different methods on the Arxiv and Books3 dataset, with training lengths of 512 and 1024.
4.1 Compare SAS with Baselines
SAS achieves better performance on different training lengths and datasets. As shown in Figure 2, compared with baselines, SAS achieves the best performance on different training lengths. For example, on the Arxiv dataset with training lengths 512 and 1024, the SAS achieves the lowest perplexity (ppl) 5.65 and 4.76, while the MHA achieves 5.82 and 4.93. Similarly, on Books3 dataset with training lengths 512 and 1024, SAS achieves 28.37 and 22. ppl, while MHA achieves 29.80 and 23.67 ppl. These results suggest that SAS is robust and effective approach, performing better across different training lengths and datasets.

---

SAS saves training token number. MHA achieves 29.86 ppl at the final training step (50K iterations) on the Books3 dataset and with training length 512. SAS already achieves 30.49 ppl at training step 30K. Therefore, SAS could save about 40% training tokens. Similarly, MHA achieves 23.67 ppl at the final training step, with Books3 dataset and training length 1024, SAS achieves 23.85 ppl at training step 30K. We have similar observation on the Arxiv dataset. In summary, the proposed SAS can achieve certain performance (low ppl) with fewer training steps, thus saving training tokens and computational costs during training.
SAS maintains better performance on longer training length. As shown in Fig- ure 3, SAS achieves the best performance (the lowest 18.34 ppl), while the standard MHA achieves 19.72 ppl. In comparison, other baselines such as GQA, MQA, MLA, and TPA achieve 16. 19.59, 19.87 and 20.25 ppl, respectively. These results show
Books3 Dataset (Length 2048)
MHA
Validation perp cxts
35
MCA
00A
30-
TIA
5A5
25-
that SAS consistently excel even with a longer training length of 2048, indicating its potential to perform well with longer training lengths.
20:
10K
15K
20K 25K 30K 35K 40K 45K 50K
Validation Step
Figure 3: The performance on training length 2048.
4.2
The Performance for Larger Models
350M Model: Books3 Dataset (Length 512)
2,78 Model: Books3 Dataset (Length 512)
MHA
MHH
Validation perplexity
Validation perplexity
50-
50
GQA
GOA
MLA
M.A
TIA
TPA
40
30
5k
10K
15K 20K 25K 30K 35K 40K 45K
50K
5k 10K 15K 20K 25K 30K 35K 40K 45K 50K Validation Step
Validation Step
Figure 4: The performance of different methods on the Books3 dataset, with model sizes of 350M and 2.7 B.
SAS achieves better performance for larger model sizes. As shown in Figure 4, when the model size is 350M, SAS achieves 24.91 ppl the final training step (50K steps), while MHA achieves 26.05 ppl. GQA, MQA, MLA, and TPA achieve 26.50, 26.29, 26.48, and 26.58 ppl. When we consider the model of size 2.7B, SAS achieves 21.62 ppl at the final training step. while MHA achieves 22.18 ppl. GQA, MQA, MILA, and TPA archive 22.29, 22.34, 22.52, and 22.25 ppl, respectively. These results demonstrate SAS's ability to leverage larger model sizes to achieve better performance.
SAS still achieves better performance across all validation steps, on larger model sizes. Across validation sizes ranging from 5K to the final 50K, SAS consistently outperforms competing methods, demonstrating not only superior final performance but also sustained improvements at every stage of training. Notably, SAS achieves this while exhibiting faster convergence, significantly reducing loss earlier in training compared to alternatives, highlighting its superiority in both training efficiency and ultimate model accuracy.
4.3 Ablation Study
The Effeet of Head Simulation. As shown Figure 5, the head simulation is employed to increase the attention head number. The SAS with only head simulation and 12 attention heads achieves 23.1 ppl, which is better than MHA with 12 attention heads, achieving 23.67 ppl. This suggests that SAS head simulation further improves the expressiveness, even with the same attention head number. Also, the proposed SAS without head simulation (FeatureOnly) achieves 22.79 ppl, with 36 attention heads, which achieves lower performance than SAS with both head simulation, feature simulation and 36
7

---

SAS Head-Only: Books3 Dataset (Length 1024)
SAS Feature-Only: Books3 Dataset (Length 1024)
45
45
MHA
Validation perplexity
SAS Headdnly Heod12
Validation perplexity
FestursOnly
40
SAS Heddd
40
SAS
35-
SAS
35
30
30
5k
10K
15K
20K 25K, 30K 35K 40K 45K 50K
10K
15K
20K 25K 30K 35K 40K 45K 50K
Validation Step
Validation Step
Figure 5: The Ablation study of SAS (SAS-HeadOnly and SAS-FcatureOnly) on the Books3 dataset, with training length 1024.
attention heads, achieving 22.14 ppl. This suggests that the head simulation is important to improve performance.
The Effect of Feature Simulation. As shown in Figure 5, the feature simulation is used to increase the query and key feature dimension. Empirical results demonstrate its effectiveness: on the Books3 dataset with a training length of 1024, the SAS- AAS-FeatureOnly achieves 22.79 ppl, which is better than MHA with 23.67 ppl. Furthermore, SAS without feature simulation (HeadOnly) achieves a higher perplexity of 22.43, which is worse than SAS with 22.14 ppl. This suggests that the feature simulation is crucial to improve the performance.
Ablation SAS Linear: Books3 Dataset (Length 1024)
Ablation Head Number: Books3 Dataset (Length 1024)
Jgl dettion perplexity
45
MHA
SAS-Liear
Vallddttton perplexity
45
NHA
SAS-Head12
40-
5A5-twad24
SAS-twwi35
35
30
25
10K 15K 20K 25K 30K 35K
50K
10K
15K
20K 25K 30K 35K
45K
Validation Step
Validation Step
Figure 6: The ablation study of SAS-Lineat and different attention head number of SAS, with training length 1024 and Books3 dataset.
The Effect of Non-Linear Mapping. SAS adopts ReLU as the nonlinear activation function in the mapping of simulation. As shown in Figure 6, with non-linear mapping, SAS achieves 22. 14 perplexity on the Books3 dataset with training length 1024. SAS without non-linear mapping achieves 22.43 ppl, which is higher than SAS with non-linear mapping and lower than MHA. This suggests the advantages of introducing nonlinear mappings in simulation.
The Effect of Attention Head Number. As shown in Figure 6, when the attention number is 12, SAS achieves 22.96 ppl, outperforming MHA. This implies that SAS can improve the expressiveness of query, key and value embeddings, leading to better performance even with the same attention head number. Moreover, increasing the attention head number from to 36 leads to further performance gains, with SAS achieving perplexity of 22.14. This demonstrates that SAS benefits from additional attention heads, which enhance the expressiveness of query, key and value embeddings and result in better performance compared to MHA.
The Effect of Kernel Size. As shown in Appendix E with a 125M model, kernel size of for the head convolution yields significant performance gains. With only less than about 0.4% additional parameters, which is 12x x ((12 x 36 + 36 x 36) x 3 + (64 x 96 + 96 x 96) x 2) = 0.43.M, SAS is able to significantly improve the performance, reducing the perplexity from 23.67 to 22.50 MHA. Moreover, as the model size increases, the additional parameter cost ratio becomes even smaller, such as less than 0.1 % for a 2.7B model. The results also show that as the kernel size increases from 7, the perplexity gradually decreases, with the performances of kernel sizes 5 and 7 being comparable.
8

---

These findings suggest that SAS can improve performance even with a small kernel size of 1, and the additional parameter cost ratio may gradually decrease as the model size becomes large.
4.4 The Performance on Large-Scale Pretrain Dataset
Table I: Main language modeling results against different methods. All models are trained on the same subset of the FineWeb- Edu dataset [45] with the GPT-2 tokenizer and 100K training steps.
Model
ARC-E
ARC-C
Hellaswag
PIQA
ScIQ
SocialIQA
Winograde acc
Avg
Metric
acc
acc
acc
acc
acc
acc
350M params/ / 50B tokens
MHA
56.57
29.01
45.45
69.10
76.40
40.94
52.96
52.92
MQA
58.12
30.55
45.71
69.10
78.70
39.25
$1.46
53.27
GQA
57.62
30.38
46.04
68.99
76.20
39.56
53.28
53.15
MI.A
57.15
28.92
44.77
67.95
75.90
38.89
53.67
52.46
TPA
59.09
32.08
46.51
69.53
76.20
39.82
53.12
53.76
SAS
60.44
31.66
47.79
70.67
80.70
40.07
54.14
55.07
350M params/ / IOB tokens
MHA
54.29
27.99
39.94
66.38
74.00
37.67
$3.59
50.55
MQA
54.45
28.75
40.51
66.21
73.10
37.92
$1.30
50.32
GQA
53.03
27.73
40.34
66.59
74.20
38.95
51.22
50.29
MLA
53.28
27.65
39.51
66.
76.00
38.08
$2.09
50.39
TPA
52.44
27.99
41.27
67.57
72.60
38.33
53.35
50.51
SAS
55.93
29.61
43.04
68.82
75.90
38.43
53.20
52.13
760M params
JOB tokens
MILA
56.57
29.01
45.45
67.25
77.20
38.54
$2.96
52.43
MQA
55.85
29.95
43.63
67.85
76.20
39.30
53.35
52.32
GQA
54.21
29.35
44.40
68.34
77.70
38.38
52.17
$2.08
MLA
57.15
29.10
44.77
67.95
75.90
38.89
$3.67
52.49
58.12
30.89
44.71
69.37
77.90
39.30
$2.80
$3.30
SAS
59.43
31.91
45.84
69.53
78.30
39.61
55.25
54.27
1.5B params 10B tokens
57.87
31.40
45.51
68.82
75.90
38.18
$2.33
$2.86
MQA
55.85
31.74
46.40
69.53
77.30
38.13
54.70
53.38
GQA
57.62
30.20
46.20
69.48
76.30
38.95
53.59
53.19
MLA
57.24
29.95
44.90
68.50
75.20
39.76
53.43
52.71
TPA
59.81
31.23
46.84
68.99
75.60
39.46
$5.09
$3.86
SAS
60.44
34.39
48.66
70.08
81.40
39.92
54.93
55.69
Downstream Evaluation. We evaluate zero-shot and two-shot performance standard benchmarks, including ARC [18], HellaSwag [93], PIQA [6]. ScIQ [78], SocialIQA [60] and WinoGrande [59], using the Am-evallation--aarta ScIQ, we report accuracy norm; for other tasks, we report standard accuracy. We display the zero-shot evaluation results of models here in Tables
codebase [22]. For ARC-E, ARC-C, HellaSwag, PIQA, and
SAS behaves well from small data size (e.g.. 10B) to larger data size (e.g. 50B). When the training tokens 10B, the 350M model achieves the best performance most benchmarks, such as ARC-E, Hellawag and so on. Also, SAS achieves the best average performance 52.13. When the training tokens number increases from 10B 50B, SAS still achieves the best performance most benchmarks, and SAS achieves the best performance with an average score 55.07. This demonstrates that SAS's performance superiority holds across range of data sizes, from small to large.
SAS works well from small model size (c.g. 350M) to larger model size (c.g. 1.5B). SAS consistently performs competitively. When the model size 350M, SAS excels, topping most benchmarks (e.g. ARC-E, Hellaswag and Winograde), with average score of 52.13. TPA also performs well with an average score of 50.51, but SAS shows broader dominance. When the model size becomes 760M or 1.5B, SAS still achieves better performance than other methods, with average score of 54.27 and 55.69. This demonstrates that SAS's performance henefits persist with larger model sizes.
5 Conclusion
In this work, we find that the attention head number and hidden size per head play crucial roles the Transformer model's performance. Therefore, we propose SAS to simulate both the attention head

---

number and feature dimension per head to improve the performance. To control the parameter cost, we also propose the PEFA aggregate attention outputs from each head. We conduct extensive ex- periments to validate our methods, including different lengths, different datasets, and different model sizes. Additionally, we scale up SAS and evaluate on downstream tasks, proving the effectiveness of our method. We believe that this paper provides an insight into using moderately sized models to mimic larger model attention score with enhanced expressiveness.
References
[1] Muhammad Adnan, Akhil Arunkumar, Gaurav Jain, Prashant J Nair, Ilya Soloveychik, and Purushotham Kamath. Keyformer: Ky cache reduction through key tokens selection for efficient generative inference. Proceedings of Machine Learning and Systems, 6:114-127, 2024.
[2] Abien Fred Agarap. Deep learning using rectified linear units (relu). arXiv preprint arXiv:/803.08:: 2018.
[3] Joshua Ainslie, James Lee-Thorp. Michiel De Jong, Yury Zemlyanskiy, Federico Lebrón, and Sumit Sanghai. Gqa: Training generalized multi-query transformer models from multi-head checkpoints. arXiv preprint arXiv:2305. 13245, 2023.
[4] Joshua Ainslie, Tao Lei, Michiel de Jong, Santiago Ontanon, Siddhartha Brahma, Yury Zemlyanskiy, David C Uthus, Mandy Guo, James Lee-Thorp, Yi Tay, et al. Colt5: Faster long-range transformers with conditional computation. Proceedings of the 2023 Conference Empirical Methods in Natural Language Processing, pages 5085-5100, 2023.
[5] Dzmitry Bahdanau, Kyunghyun Cho, and Yoshua Bengio. Neural machine translation by jointly learning to align and translate. arXiv preprint arXiv: 1409. 0473, 2014.
[6] Yonatan Bisk, Rowan Zellers, Jianfeng Gao, Yejin Choi, al. Piqa: Reasoning about physical commonsense in natural language. In Proceedings of the AAAI conference on artificial intelligence, volume 34, pages 7432-7439, 2020.
[7] Tom Brown, Benjamin Mann, Nick Ryder, Melanic Subbiah, Jared D Kaplan, Prafulla Dhari- wal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877- 1901, 2020.
[8] Sébastien Bubeck, Varun Chadrasekaran, Ronen Eldan, Johannes Gehrke, Eric Horvitz, Eie Kamar, Peter Lee, Yin Tat Lee, Yuanzhi Li, Scott Lundberg, et al. Sparks of artificial general intelligence: Early experiments with gpt-4, 2023.
[9] Kerim Büyükakyüüz Olora: Orthonormal low-rank adaptation of large language models. arXiv preprint arXiv:2406.01777 2024.
[10] Zefan Cai, Yichi Zhang, Bofei Gao. Yuliang Liu, Tianyu Liu, Keming Lu, Wayne Xiong, Yue Dong, Baobao Chang, Junjie Hu, al. Pyramidkv: Dynamic ky cache compression based on pyramidal information funneling. arXiv preprint 24t.::20000020099 2024.
[11] Yukang Chen, Shengju Qian, Haotian Tang, Xin Lai, Zhijian Liu, Song Han, and Jiaya Jia. Longlora: Efficient fine-tuning of long-context large language models. In The Twelfth International Conference on Learning Representations, 2024.
[12] Ta-Chung Chi, Ting-Han Fan, Peter J Ramadge, and Alexander Rudnicky, Kerple: Kernelized relative positional embedding for length extrapolation. Advances in Neural Information Processing Systems, 35:8386 8399, 2022.
[13] Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever. Generating long sequences with sparse transformers. arXiv preprint arXiv: 10509, 2019.
[14] Krzysztof Marcin Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter Hawkins. Jared Quincy Davis, Afroz Mohiuddin, Lukasz Kaiser, David Benjamin Belanger, Lucy Colwell, and Adrian Weller. Rethinking attention with performers. In International Conference Learning Representations, 2021.
10

---

[15] Krzysztof Marcin Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, An- dreea Gane, Tamas Sarlos, Peter Hawkins, Jared Quincy Davis, Afroz Mohiuddin, Lukasz Kaiser, et al. Rethinking attention with performers. In International Conference on Learning Representations, 2019.
[16] Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, et al. Palm: Scaling language modeling with pathways. Journal of Machine Learning Research, 24(240):1- 113, 2023.
[17] Christopher Clark, Kenton Lee, Ming Wei Chang, Tom Kwiatkowski, Michael Collins, and Kristina Toutanova. Boolq: Exploring the surprising difficulty of natural yes/no questions. 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 2924-292222 2019.
[18] Peter Clark, Isaac Cowhey, Oren Etzioni, Tushar Khot, Ashish Sabharwal, Carissa Schoenick, and Oyvind Tafjord. Think you have solved question answering? try are, the ai2 reasoning challenge. arXiv preprint 10X3.0880000555 2018.
[19] Tri Dao and Albert Gu. Transformers are ssms: Generalized models and efficient algorithms through structured state space duality. In Forty-first International Conference on Machine Learning.
[20] Yiran Ding, Li Lyna Zhang, Chengruidong Zhang, Yuanyuan Xu, Ning Shang, Jiahang Xu, Fan Yang, and Mao Yang. Longrope: Extending Ilm context window beyond 2 million tokens. In Forty-first International Conference on Machine Learning.
[21] Zafeirios Fountas, Martin A Benfeghoul, Adnan Oomerjee, Fenia Christopoulou, Gerasimos Lampouras, Haitham Bou-Ammar, and Jun Wang. Human-like episodic memory for infinite context llms. arXiv preprint arXiv:2407.09455 2024.
[22] Leo Gao, Jonathan Tow, Baber Abbasi, Stella Biderman, Sid Black, Anthony DiPofi, Charles Foster, Laurence Golding, Jeffrey Hsu, Alain Le Noac"h, Haonan Li, Kyle McDonell, Niklas Muennighoff, Chris Ociepa, Jason Phang, Laria Reynolds, Hailey Schoelkopf, Aviya Skowron, Lintang Sutawika, Eric Tang, Anish Thite, Ben Wang, Kevin Wang, and Andy Zou. The language model evaluation harness, 07 2024.
[23] Xavier Glorot and Yoshua Bengio. Understanding the difficulty of training deep feedforward neural networks. In Proceedings of the thirteenth international conference on artificial intelligence and statistics, pages 249-256. JMLR Workshop and Conference Proceedings, 2010.
[24] Alex Graves and Alex Graves. Long short-term memory. Supervised sequence labelling with recurrent neural networks, pages 37-45, 2012.
[25] Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces. In First Conference on Language Modeling.
[26] Albert Gu, Karan Goel, and Christopher Re. Efficiently modeling long sequences with
structured state spaces. In International Conference on Learning Representations.
[27] Mandy Guo, Joshua Ainslie, David C Uthus, Santiago Ontanon, Jianmo Ni, Yun-Hsuan Sung, and Yinfei Yang. Longt5: Efficient text-to-text transformer for long sequences. In Findings of the Association for Computational Linguistics: NAACL 2022, pages 724-736, 2022.
[28] Insu Han, Rajesh Jayaram, Amin Karbasi, Vahab Mirrokni, David Woodruff, and Amir Zandieh. Hyperattention: Long-context attention in near-linear time. In The Twelfth International Conference on Learning Representations, 2024.
[29] Jun Han and Claudio Moraga. The influence of the sigmoid function parameters on the speed of backpropagation learning. In International workshop on artificial neural networks, pages 195-201. Springer, 1995.
11

---

[30] Zhenyu He, Guhao Feng, Shengjie Luo, Kai Yang, Liwei Wang, Jingjing Xu, Zhi Zhang, Hongxia Yang, and Di He. Two stones hit one bird: Bilevel positional encoding for better length extrapolation. In International Conference on Machine Learning, pages 17858-17876. PMLR, 2024.
[31] Coleman Hooper, Sehoon Kim, Hiva Mohammadzadeh, Michael W Mahoney, Sophia Shao, Kurt Keutzer, and Amir Gholami. Kvquant: Towards 10 million context length Ilm inference with kv cache quantization. Advances in Neural Information Processing Systems, 37:1270- 1303, 2024.
[32] Edward Hu, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, Weizhu Chen, et al. Lora: Low-rank adaptation of large language models. In International Conference Learning Representations, 2022.
[33] Jingcheng Hu, Houyi Li, Yinmin Zhang, Zili Wang, Shuigeng Zhou, Xiangyu Zhang, Heung-Yene Yeung Shum, and Daxin Jiang. Multi-matrix factorization attention. arXiv preprint arXiv:2412. 19255, 2024.
[34] Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. arXiv preprint arXiv:2001.08361, 2020.
[35] Yann LeCun, Yoshua Bengio, and Geoffrey Hinton. Deep learning. nature, 521(7553):436 2015.
[36] Yann LeCun, Léon Bottou, Yoshua Bengio, and Patrick Haffner. Gradient-based learning applied to document recognition. Proceedings of the IEEE, 86(11):2278-232) 1998.
[37] Shanda Li, Chong You, Guru Guruganesh, Joshua Ainslie, Santiago Ontanon, Manzil Zaheer, Sumit Sanghai, Yiming Yang. Sanjiv Kumar, and Srinadh Bhojanapalli. Functional interpola- tion for relative positions improves long context transformers. In The Twelfth International Conference on Learning Representations.
[38] Vladislav Lialin, Sherin Muckatira, Namrata Shivagunde, and Anna Rumshisky. Relora: High-rank training through low-rank updates. In The Twelfth International Conference on Learning Represenlations, 2024.
[39] Yan-Shuo Liang and Wu-Jun Li. Inflora: Interference-free low-rank adaptation for continual learning. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 23638 22647,, 2024.
[40] Aixin Liu, Bei Feng, Bing Xue, Bingxuan Wang, Bochao Wu, Chengda Lu, Chenggang Zhao, Chengqi Deng, Chenyu Zhang, Chong Ruan, et al. Deepseek-v3 technical report. arXiv preprint 2xXi:24412119437 2024.
[41] Bo Liu, Rui Wang, Lemeng Wu, Yihao Feng, Peter Stone, and Qiang Liu. Longhorn: State space models are amortized online learners. arXiv preprint 2xXiv:2407.14207 2024.
[42] Zirui Liu, Jiayi Yuan, Hongye Jin, Shaochen Zhong, Zhaozhuo Xu, Vladimir Braverman, Beidi Chen, and Xia Hu. Kivi: A tuming-free free asymmetric 2bit quantization for kv cache. In Forty-first International Conference Machine Learning, 2024.
[43] Ilya Loshchilov and Frank Hutter. Sgdr: Stochastic gradient descent with warm restarts. arXiv preprint arXiv: 1608.0383, 2016.
[44] Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In International Conference on Learning Representations, 2019.
[45] Anton Lozhkov, Loubna Ben Allal, Leandro von Werra, and Thomas Wolf. Fineweb-edu: the finest collection of educational content, 2024.
[46] Sadhika Malladi, Alexander Wettig, Dingli Yu, Danqi Chen, and Sanjeev Arora. A kemel-- based view of language model fine-tuning. International Conference on Machine Learning, pages 236 10-223464. PMLR, 2023.
12

---

[47] Todor Mihaylov, Peter Clark, Tushar Khot, and Ashish 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 138111399, 2018
[48] Antonio Orvieto, Samuel L Smith, Albert Gu, Anushan Fernando, Caglar Gulcehre, Razvan Pascanu, and Soham De. Resurrecting recurrent neural networks for long sequences. In International Conference on Machine Learning, pages 26670-26698. PMLR, 2023.
[49] A Paszke. Pytorch: An imperative style, hegh-performance deep learning library. arXiv preprini arXiv: wwxXiv1912.017000 2019.
[50] Guilherme Penedo, Hynek Kydlíček, Anton Lozhkov, Margaret Mitchell, Colin A Raffel, Leandro Von Werra, Thomas Wolf, et al. The fineweb datasets: Decanting the web for the finest text data scale. Advances in Neural Information Processing Systems, 37:3081 1-30849, 2024.
[51] Bo Peng, Eric Alcaide, Quentin Gregory Anthony, Alon Albalak, Samuel Arcadinho. Stella Biderman, Huanqi Cao, Xin Cheng, Michael Nguyen Chung, Leon Derezynski, et al. Rwkv: Reinventing rnns for the transformer era. In The 2023 Conference Empirical Methods in Natural Language Processing.
[52] Bo Peng, Daniel Goldstein, Quentin Gregory Anthony, Alon Albalak, Eric Alcaide, Stella Biderman, Eugene Cheah, Teddy Ferdinan, Kranthi Kiran GV, Haowen Hou, et al. Eagle and finch: Rwkv with matrix-valued states and dynamic recurrence. In First Conference on Language Modeling.
[53] Marius-Constantin Popescu, Valentina E Balas, Liliana Perescu-Popescu, and Nikos Mas- torakis. Multilayer perceptron and neural networks. WSEAS Transactions Circuits and Systems, ((7):579-588, 2009.
[54] Ofir Press, Noah Smith, and Mike Lewis. Train short, test long: enables input length extrapolation. In International Conference on Learning Representations.
Attention with linear biases
[55] Zhen Qin, Songlin Yang, and Yiran Zhong. Hierarchically gated recurrent neural network for sequence modeling. Advances in Neural Information Processing Systems, 36:33202-3322 2023.
[56] Jianing Qiu, Lin Li, Jiankai Sun, Hao Wei, Zhe Xu, Kyle Lam, and Wu Yuan. Emerging cyber attack risks of medical ai agents. arXiv preprint 2rXiv:2504403759, 2025.
[57] Alee Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models unsupervised multitask learners. OpenAl blog, 1(8):9, 2019.
[58] Anian Ruoss, Grégoire Delétang, Tim Genewein, Jordi Grau-Moya. Róbert Csordás, Mehdi Bennani. Shane Legg, and Jocl Veness. Randomized positional encodings boost length generalization of transformers. In Proceedings of the 61 Annual Meeting of the Association for Computational Linguistics (Volume 2: Short Papers), pages 1 1889-1903 1903, 2023.
[59] Keisuke Sakaguchi, Ronan Le Bras, Chandra Bhagavatula, and Yejin Choi. Winogrande: An adversarial winograd schema challenge at scale. Communications of the ACM, 64(9):99-106,
2021.
[60] Maarten Sap, Hannah Rashkin, Derek Chen, Ronan Le Bras, and Yejin 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 Natural Language Processing GEMNLP-IJCNLPP pages 4463-4473, 2019.
[61] Noam Shazeer. Fast transformer decoding: One write-head is all you need. arXiv preprint /rXiv:/91..02/50/ 2019.
[62] Zhenmei Shi, Jiefeng Chen, Kunyang Li, Jayaram Raghuram, Xi Wu, Yingyu Liang, and Somesh Jha. The trade-off between universality and label efficiency of representations from contrastive learning. arXiv preprint arXiv:2303.00 2023.
13

---

[63] Jerome Sieber, Carmen Amo Alonso, Alexandre Didier, Melanie N. Zeilinger, and Antonio Orvieto. Understanding the differences in foundation modelssattention, state space models, and recurrent neural networks. Advances in Neural Information Processing Systems, 37, 2024.
[64] Prajwal Singhania, Siddharth Singh, Shwai He, Soheil Feizi, and Abhinav Bhatele. Loki: Low-rank keys for efficient sparse attention. arXiv preprini arXiv:2406.02553 2024.
[65] Jimmy TH Smith, Andrew Warrington, and Scott Linderman. Simplified state space layers for sequence modeling. In The Eleventh International Conference on Learning Representations.
[66] Jianlin Su, Murtadha Ahmed, Yu Lu, Shengfeng Pan, Wen Bo, and Yunfeng Liu. Roformer: Enhanced transformer with rotary position embedding. Neurocomputing, 568: 127063, 2024.
[67] Hanshi Sun, Li-Wen Chang, Wenlei Bao, Size Zheng, Ningxin Zheng. Xin Liu, Harry Dong, Yuejie Chi, and Beidi Chen. Shadowkv: Kv cache in shadows for high-throughput long-context 2024.
Ilm inference. arXiv preprint arXiv:2440..144440
[68] Jiankai Sun, Chuanyang Zheng, Enze Xie, Zhengying Liu, Ruihang Chu, Jianing Qiu, Jiaqi Xu, Mingyu Ding, Hongyang Li, Mengzhe Geng, al. A survey of reasoning with foundation models: Concepts, methodologies, and outlook. ACM Computing Surveys, 2023.
[69] Yu Sun, Xinhao Li, Karan Dalal, Jiarui Xu, Arjun Vikram, Genghan Zhang, Yann Dubois, Xinlei Chen, Xiaolong Wang, Sanmi Koyejo, et al. Learning to (learn at test time): Runs with expressive hidden states. arXiv preprint arXiv:2407.04620 2024.
[70] Yutao Sun, Li Dong, Shaohan Huang, Shuming Ma, Yuqing Xia, Jilong Xue, Jianyong Wang, and Furu Wei. Retentive network: A successor to transformer for large language models. arXiv preprint ar1iv:3307.0866 2023.
[71] Alon Talmor, Jonathan Herzig, Nicholas Lourie, and Jonathan Berant. Commonsenseqa: A question answering challenge targeting commonsense knowledge. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume (Long and Short Papers), pages 4149 4158, 2019.
[72] Jiaming Tang, Yilong Zhao, Kan Zhu, Guangxuan Xiao, Baris Kasikei, and Song Han. Quest: Query- sparsity for efficient long-context 11m inference. arXiv preprint 1r77444406.10774 2024.
[73] Yi Tay, Dara Bahri, Liu Yang, Donald Metzler, and Da-Cheng Juan. Sparse sinkhorn attention. In International Conference on Machine Learning, pages 1438-9447. PMIR, 2020.
[74] Gemma Team, Thomas Mesnard, Cassidy Hardin, Robert Dadashi, Surya Bhupatiraju, Shreya Pathak, Laurent Sifre, Morgane Rivière, Mihir Sanjay Kale, Juliette Love, et al. Gemma: Open models based on gemini research and technology. arXiv preprint arXiv:2403.08S. 088.99 2024.
[75] Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timo- thée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, al. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302.133) 2023.
[76] Yao-Hung Hubert Tsai, Shaojie Bai, Makoto Yamada, Louis-Philippee Morency, and Ruslan Salakhutdinov. Transformer dissection: unified understanding of transformer's attention via the lens of kernel. arXiv preprint arXiv: 1775, 2019.
[77] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. Advances in neural information processing systems, 30, 2017.
[78] Johannes Welbl, Nelson F Liu, and Matt Gardner. Crowdsourcing multiple choice science questions. Proceedings of the 3rd Workshop on Noisy User- eeserageeeetted Text, pages 94-106, 2017.
[79] Guangxuan Xiao, Ji Lin, Mickael Seznec, Hao Wa, Julien Demouth, and Song Han. Smoothquant: Accurate and efficient post-training quantization for large language models. In International Conference on Machine Learning, pages 38087-38099. PMLR, 2023.
14

---

[80] Guangxuan Xiao, Yuandong Tian, Beidi Chen, Song Han, and Mike Lewis. Efficient streaming language models with attention sinks, arXiv preprint arXiv:2309. 17453, 2023.
[81] Jing Xiong, Qiujiang Chen, Fanghua Ye, Zhongwei Wan, Chuanyang Zheng, Chenyang Zhao, Hui Shen, Alexander Hanbo Li, Chaofan Tao, Haochen Tan, al. A1: Asynchronous test-time scaling via conformal prediction. arXiv preprint arXiv:2509. 15148, 2025.
[82] Jing Xiong, Chengming Li, Min Yang, Xiping Hu, and Bin Hu. Expression syntax information bottleneck for math word problems. Proceedings of the 45th International ACM SIGIR Conference on Research and Development in Information Retrieval, pages
2022.
[83] Jing Xiong, Zixuan Li, Chuanyang Zheng, Zhijiang Guo, Yichun Yin, Enze Xie, Zhicheng Yang, Qingxing Cao, Haiming Wang, Xiongwei Han, et al. Dq-lore: Dual queries with low rank upproximation re-ranking for in-context learning. arXiv preprint arXiv:231.00255, 2023.
[84] Jing Xiong, Jianghan Shen, Fanghua Ye, Chaofan Tao, Zhongwei Wan, Jianqiao Lu, Xun Wu. Chuanyang Zheng. Zhijiang Guo, Lingpeng Kong, et al. Uncomp: Uncertainty- aware long-context compressor for efficient large language model inference. arXiv preprini arXiv:2410.03090 2024.
[85] Jing Xiong, Jianghan Shen, Chuanyang Zheng, Zhongwei Wan, Chenyang Zhao, Chiwun Yang, Fanghua Ye, Hongxia Yang, Lingpeng Kong, and Ngai Wong. Parallelcomp: Parallel long-context compressor for length extrapolation. arXiv preprint arXiv:2502.14317 2025.
[86] Jing Xiong, Jianhao Shen, Ye Yuan, Haiming Wang, Yichun Yin, Zhengying Liu, Lin Li, Zhijiang Guo, Qingxing Cao, Yinya Huang, et Trigo: Benchmarking formal mathematical proof reduction for generative language models. arXiv preprint arXiv:2310.10180 2023.
[87] Jing Xiong, Zhongwei Wan, Xiping Hu, Min Yang, and Chengming Li. Self-consistent reasoning for solving math word problems. arXiv preprint 10.:22200.1337.. 2022.
[88] Bing Xu, Naiyan Wang, Tianqi Chen, and Mu Li. Empirical evaluation of rectified activations in convolutional network. arXiv preprint 1505,
2015.
[89] Ting Xu, Zhichao Huang, Jiankai Sun, Shanbo Cheng, and Wai Lam. Seqpo-simt: Sequential policy optimization for simultaneous machine translation. arXiv preprint arXiv:2505.20622. 2025.
[90] Kai Yang, Jan Ackermann, Zhenyu Hc, Guhao Feng, Bohang Zhang, Yunzhen Feng, Qiwei Ye, Di He, and Liwei Wang. Do efficient transformers really save computation? International Conference on Machine Learning, 2024.
[91] Songlin Yang, Jan Kautz, and Ali Hatamizadeh. Gated delta networks: Improving mamba2 with delta rule. arXiv preprint (xXiv:24/2.064666 2024.
[92] Ziyu Ye, Jiacheng Chen, Jonathan Light, Yifei Wang, Jiankai Sun, Mac Schwager, Philip Torr, Guohao Li, Yuxin Chen, Kaiyu Yang, et al. Reasoning in reasoning: A hierarchical framework for better and faster neural theorem proving. In The 4th Workshop on Mathematical Reasoning and Al at NeurlPS'24.
[93] Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. Hellaswag: Can a machine really finish your sentence? In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, pages 4791-4800 2019.
[94] Yuchen Zeng and Kangwook Lee. The expressive power of low-rank adaptation. In The Twelfth International Conference on Learning Representations.
[95] Hengyu Zhang. Sinklora: Enhanced efficiency and chat capabilities for long-context large language models. arXiv preprint arXiv:2406.05668, 2024.
[96] Jingqing Zhang, Yao Zhao, Mohammad Saleh, and Peter Liu. Pegasus: Pre-training with extracted gap-sentences for abstractive summarization. In International conffeereencnnnfeeeen on machine learning, pages 11328-11333 11339. PMLR, 2020.
15

---

[97] Qingru Zhang. Minshuo Chen, Alexander Bukharin, Nikos Karampatziakis, Pengcheng He, Yu Cheng, Weizhu Chen, and Tuo Zhao. Adalora: Adaptive budget allocation for parameter- efficient fine-tuning. arXiv preprint arXiv:2303. 105/2, 2023.
[98] Ruiqi Zhang, Spencer Frei, and Peter L Bartlett. Trained transformers learn linear models in-context. Journal of Machine Learning Research, 25(49):1-55, 2024.
[99] Wencheng Zhang, Shiqin Qiao, Lingjie Luo, Yinfeng Li, Chuanyang Zheng, Qian Xu, Meng Li, Yong Gui, Yijun He, Jianing Qiu, et al. Synapseroute: An auto-route switching framework on dual-state large language model. arXiv preprint arXiv:2507.02883 2025.
[100] Yifan Zhang, Yifeng Liu, Huizhuo Yuan, Zhen Qin, Yang Yuan, Quanquan Gu, and Andrew Chi-Chih Yao. Tensor product attention is all you need. arXiv preprint arXiv:2501.0644 2025.
[101] Yu Zhang, Songlin Yang. Rui-Jie Zhu, Yue Zhang, Leyang Cui, Yiqiao Wang. Bolun Wang, Freda Shi, Bailin Wang, Wei Bi, et al. Gated slot attention for efficient linear-time sequence modeling. Advances in Neural Information Processing Systems, 37:1 17.116870-11689, 2024.
[102] Hongbo Zhao, Bolin Ni, Junsong Fan, Yuxi Wang, Yuntao Chen, Gaofeng Meng, and Zhaoxi- ang Zhang. Continual forgetting for pre-trained vision models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 22642, 2024.
[103] Chuanyang Zheng, Yihang Gao, Han Shi, Minbin Huang, Jingyao Li, Jing Xiong, Xiaozhe Ren, Michael Ng, Xin Jiang, Zhenguo Li, et al. Dape: Data-adaptive positional encoding for length extrapolation. Advances in Neural Information Processing Systems, 37:26659-26700, 2024.
[104] Chuanyang Zheng, Yihang Gao, Han Shi, Jing Xiong, Jiankai Sun, Jingyao Li, Minbin Huang, Xiaozhe Ren, Michael Ng. Xin Jiang, et al. Dape v2: Process attention score as feature map for length extrapolation. arXiv preprint 2rXiv:2410.0443 2024.
[105] Chuanyang Zheng. Zhengying Liu, Enze Xie, Zhenguo Li, and Yu Li. Progressive-hing prompting improves reasoning in large language models. In Al for Math Workshop@ ICML 2024.
[106] Chuanyang Zheng, Haiming Wang, Enze Xie, Zhengying Liu, Jiankai Sun, Huajian Xin, Jianhao Shen, Zhenguo Li, and Yu Li. Lyra: Orchestrating dual correction in automated theorem proving. arXiv preprint arXiv:2309. 15806, 2023.
16

---

A Limitation
Compared to the baseline method, SAS may have higher computational costs due to the additional nonlinear mappings and increased attention heads and feature dimensions. Specifically, SAS employs a greater number of simulated attention heads, each with an increased hidden dimension, leading to larger overall parameter count and more extensive matrix operations during both forward and backward passes.
B Broader Impacts
In this work, we introduce and develop novel, more powerful approach for advancing language modeling capabilities. The positive broader impact of this research lies in its potential find the way for the creation of even more effective models in the future, thereby contributing to progress in natural language processing and AI. However, alongside these benefits, it also crucial consider the potential negative broader impacts, emphasizing the importance of responsible usage, ethical considerations, and proactive measures to mitigate any unintended consequences associated with the deployment of such technology.
C Model Configuration
For the experiment Arxiv and Books3 dataset, the 125M and 350M model configurations as follows.
Table 2: The model configuration of different models follows the setting of TPA [100].
Model
Hidden Size
Head
Hidden Size Per Head
Group Size
125M parameters
MHA
768
12
64
MQA
768
23
64
23
GQA
768
22
64
11
MLA
768
13
64
TPA
768
36
64
SAS
768
36
96
350M parameters
MHA
1024
16
64
MQA
1024
31
64
31
GQA
1024
30
64
15
MLA
1024
23
64
TPA
1024
48
64
SAS
1024
48
96
For the experiment on the Fineweb-Edi dataset, the experiment setup is as follows: We follow the nanoGPT training configuration. In particular, (0.9, 0.95), a weight decay of 0.1, and gradient clipping at 1.0. We follow the same setting as nanoGPT that the learning rate is managed by cosine annealing scheduler [43]. For the model setting, we mostly follow the setting of TPA [100].
the AdamW [44] optimizer with (81, 82)
D Experiment Statistical Significance
In this section, we have shown the mean and standard deviation of different methods, with three random seeds. According to the above table, we can find that our proposed method consistently achieves the best performance. All methods show the expected trend of decreasing perplexity with increased training steps. However, SAS maintains a consistent advantage throughout the training process, with the performance gap remaining relatively stable. And our method is significantly better than all other methods, with p value 0.05. Therefore, we could have confidence to say that our SAS is significantly better than other proposed methods.
17

---

Table The mean and standard variance of perplexity metrics of different methods on three random seeds, with training length 512 and the Arxiv dataset.
Model
5K
10K
15K
20K
25K
30K
35K
40K
45K
50K
MHA
Mean
9.1802
8.0348
7.4585
7.0573
6.7825
6.4648
6.161
6.0311
5.956
5.8628
Std
0.146
0.0593
0.0644
0.0532
0.0161
0.0629
0.1094
0.083
0.04200
MQA
Mcan
9.1758
8.0427
7.4561
7.0629
6.7943
6.4823
6.182
6.0511
5.9797
5.89211
Std
0.0802
0.0747
0.0452
0.0453
0.028
0.0655
0.0979
0.0844
0.0248
0.0212
GQA
Mean
9.2267
8.0921
7.5054
7.1177
6.845
6.534
6.2292
6.1036
6.0298
5.9385
Std
0.1368
0.0428
0.0394
0.0538
0.0329
0.0776
0.0869
0.097
0.0244
0.0257
MLA
Mcan
9.1931
8.1181
7.533
7.1373
6.8651
6.5456
6.2372
6.1044
6.0331
5.9416
Std
0,1596
0.0888
0.0796
0.0325
0.0078
0.0516
0,1235
0.0727
0.0483
0.0382
TPA
Mean
9.059
7.9606
7.3845
7.0129
6.7562
6.4541
6.1567
6.0415
5.9695
5.8811
Std 0.1232
0.0583
0.0409
0.06
0.0226
0.0722
0.1037
0.0922
0.0403
0.0363
SAS
Mean
8,6952
7.6776
7.1403
6.7776
6.5234
6.2342
5,9493
5.8368
5.7662
5.6821
Std
0.1546
0.0286
0.0459
0.0623
0.0228
0.0703
1038
0.0854
0.0486
0.0425
E The SAS Performance with Varying Kernel Sizes
Books3 Dataset (Length 1024)
45
MHA
Validation perplexity
SAS Kernell
40
SAS Kernel3
SAS Kemel5
35
SAS Kemel7
25
5k
10K 15K
20K 25K 30K 35K 40K 45K 50K
Validation Step
Figure 7: The performance of SAS with different kernels, with the Books3 dataset and training length 1024.
The Training Cost
Table 4: The Time cost of different methods with training length 512 and batch size
Model
125M
350M
2.7B
6.7B
10.6B
MHA
36.20
72.87
117.82
186.58
217.01
GQA
38.87
71.66
133.76
205.69
239.16
MQA
39.37
72.01
132.62
202.61
238.85
MLA
51.68
104.95
180.03
254.43
637.78(Zero3 to avoid OOM;H200)
TPA
44.19
86.56
151.71
235.85
258.41
SAS
68.07
138.53
215.13
350.70
402.15
18

---

G Large Mode Size 6.7B and 10.6B performance
Table 5: The perplexity of 6.7B and 10B on the Pile dataset.
Model
5K
10K
15K
20K
25K
30K
35K
40K
45K
50K
6.7B
MHA
23.45
17.25
14.51
13.21
12.07
11.21
10.52
10.28
9.91
9.73
GQA
21.62
16.35
13.93
12.74
11.69
10.88
10.24
10.01
9.65
9.49
MQA
21.56
16.32
13.93
12.73
11.70
10.90
10.24
10.02
9.66
9.49
MLA
22.22
16.67
14.26
13.05
11.99
11.19
10.53
10.29
9.91
9.74
SAS
20.44
15.32
13.17
12.15
11.21
10.51
9.93
9.73
9.40
9.25
10.6B
MHA
23.67
17.34
14.59
13.18
12.00
11.15
10.46
10.20
9.82
9.64
GQA
21.86
16.45
13.92
12.70
11.61
10.81
10.17
9.94
9.58
9.41
MQA
22.13
16.55
14.03
12.80
11.70
10.88
10.21
9.98
9.61
9.43
MLA
22.40
16.67
14.22
12.99
11.90
11.11
10.42
10.18
9.83
9.66
SAS
20.53
15.12
12.92
11.90
10.98
10.29
9.71
9.51
9.18
9.04
H Large Training Length performance
Table
The perplexity of length 16384 on the Pile dataset
Model
5K
10K
15K
20K
25K
30K
35K
40K
45K
50K
MHA
18.31
14.50
12.85
11.96
11.45
11.11
10.44
10.37
9.95
9.94
GQA
17.74
14.40
12.86
11.96
11.47
11.14
10.45
10.39
9.99
9.98
MQA
18.02
14.46
12.82
11.96
11.44
11.10
10.46
10.38
9.96
9.99
MLA
17.74
14.40
12.86
11.96
11.47
11.14
10.45
10.39
9.99
9.98
TPA
18.49
14.54
12.96
12.08
11.59
11.26
10.62
10.58
10.16
10.18
SAS
14.93
12.45
11.23
10.56
10.16
9.90
9.32
9.28
8.95
8.97
I Theoretical Analysis
The Softmax Attention with Fully Parameterized
O E ( ( )2,U,) where: We DDD learnable bilinear transformation for head c, the kernel function such as softmax, U. E DDDx value projection matrix, H number of attention heads, D = hidden dimension per head. Each head computes weighted sum of values, with weights determined by x. w.W, For m-dimensional z; and n-dimensional 2j. one potential relation pattern is related to u.""."" Therefore, at most D2 relations or interactions).
Bilinear Attention
VD
The bilinear term v; W = -EP ( (17) w!(".)) ny
(n) can be expressed as (ee
where vec(W.) D= With H heads, the model represent up to D2 distinct relations [6]. We hope that the output of each head a single relation but not mixed relations, so that we can utilize the relations independently (c.g.. processes the relation independently). Assume Ronk and there is only one non-zero value in W.. And there are H D2 4. I, [a 6] x = df. W
0 0) W3- = [0 [1 0] W4 [0 0] [0 1]. The first head W1 will gives the result of relation ac, the W Ws. W. will give the result of ad, bc, bd. Therefore, this (ac, ad, bc, bd) suggests different attention patterns and relations. And they could be adjusted independently and not mixed (e.g., & processes the relation independently), compared to only one head with W
W, = [1 0]
1 1] that outputs the mix of ac + ad + be bd. With a large D. we can have more combinations. With large H, we can capture more independent relations. If we increase D (e.g., hidden size), we increase the maximum number of relations. If we increase the we increase the number of independent relations that actually captured.
19

---

J
SAS with Different Attention Types
Table The perplexity of SAS with different methods, with training length 512 and Books3 dataset.
Model
5K
10K
15K
20K
25K
30K
35K
40K
45K
50K
MIHA
60.0206
48.1325
42 22.9901
39,9132
36.6749
34.7711
32.1238
30.1071
30.019
29. 80,8
MHA-SAS
55.0498
45.1116
40.5295
37.6670
34.6701
32.9337
30.4943
28.6240
28.5052
28.3746
MQA
59.8441
48.4041
43.2021
40.128
36.9926
35.0613
32.4715
30.48
30.3322
30.1535
MQA-SAS
55.5273
45.7848
41.2885
38.4768
35,3813
33.6372
31.1418
29.2193
29.1556
29.0052
GQA
60.1109
48.2306
43.01
39,8758
36.5799
34.6976
32.0481
30.0991
29, 29,9904
29.8059
GQA-SAS
55.1054
45.5153
41.0740
38.1912
35.2073
33.3761
29.0752
28.9995
28.7853
MLA
60,4833
48.8735
43.6705
40,486
37.1407
35,2178
32.4767
30,4295
30.3255
30,1669
MLA-SAS
57.6595
46.8414
41,9994
38,9989
35,8767
34,0236
31.5449
29,5964
29.5283
29.3722
TPA
58.7126
47.9701
43.1154
40.1808
37,0918
35,3126
32,6225
30.6120
30.5273
30,3846
TPA-SAS
55.5189
45.2898
40.8298
37.9846
35.0361
33.2587
28.9701
28.8586
28.7323
present the results of SAS with varying attention types and model size of 125M, in Table 7. The MHA-SAS has triple attention heads of MHA with kernel size 5, while other SAS variants have double attention heads for their baseline with kernel size as other variants have larger attention head numbers. We observe that SAS demonstrates superior performance relative its corresponding standard attention types. Therefore, the SAS compatible with most current attention methods to further enhance the performance.
K
The Experiment Result on Large Model and Dataset
Table 8: The performance of different models with 1.5B model and 50B tokens. The performance of MHA, MQA, GQA, MLA and TPA comes from TPA paper [100]
Method
ARC-E
ARC-C
BoolQ
HellaSw.
OBQA
PIQA
W.G.
MMLU
SciQ
Avg.
MHA
64.81
35.41
61.90
54.32
37.20
72.74
55.80
25.44
82.80
54.49
MQA
64.10
36.01
62.26
54.38
39.00
72.58
56.43
23.70
81,90
$4.48
GQA
63.68
35.92
60.46
54.17
38.40
73.56
56.27
24.77
81.70
MI.A
64.14
35.92
60.12
53.60
39.20
72.25
55.17
24.71
81.60
$4.08
TPA
66.71
36.52
61.38
$4.03
40.40
72.52
56.83
24.49
82.20
55.01
SAS
66.84
37.88
61.41
56.05
41.60
72.52
57.85
26.11
82.60
55.85
Implementation
In this section, present the implementation of the proposed SAS module in PyTorch for research purposes, which is consistent with the intended use [49].
import torch import torch nn as nn import torch. nn. functional as F
class ResidualCNN (nn. Module)
def __ini self. input, output. kernel_ size, padding
): super in
self ReLU = nn nn. ReLU()
self, process = nn. annConvlee imput, output, kernel_size 1, padding)
def forward selff hidden_states,
output self ReLU(hidden states) output self, process ( output)) output output hidden states
20

---

return output
class ResidualMLP Module):
def --init_ self input. output
super (peer.))_iit_1 ):
self ReLU = mn ReL(()
self process
Linear (imput. output)
def (ell, hidden_states
output self Self.Re..U (eld(eiide__stttess output self. process (output)
output output hidden states hidden_states return output
class SAS( Module): def __init__ ori_head target_head ori_feature feature .kernel_size)> ...
SAS attention bias module.
Args:
ori_head: the original head number target_head: the target head number ori fri_ueeture: the original feature size target_feature: the target feature size kernel_size: the kernel size u*.
super(SAS, self)...init__.
self. ori oead self. target_head target_head
self ori_feature-oriffeauu self. taff.tttretefeaure target_featuur feature
self. kernelleeel__ie = kernel_size _ize
padding (kernel_size
//
self. sas sas_g na. Sequential
nn. C..Coovv ori head, self, target_ head pernel_size. 1, padding),
ResidualCNN target_head, self target_head, kernel_size padding),
nn. Linear self. ori ori_feature, self taagatarreauurreo
ResidualMLP (eeff target_feature self. target_feature) self sas_k nn. Sequential
Convid self ori ori_head, self target_head, kernel_size, 1, padding).
ResidualCNN target self target_heade kernel_size, padding).
Linear self fritufeauuee self taa.ata_faaturaee feature), ResidualMI.P self. target_feature, self. target_
self. sas_y = nn. Sequential
nn. Convid self ori_head, self target_head, kernel_size, 1, padding).
ResidualCNN (self head, self terget_head. hernel_size. padding),
self outpput_denseeneeeeeeetteee dense=an. dinear(sell ori_head=self ori_feature self. ori rri_headsself ori_feature)
def forward query, key, value):
21

---

Args:
query: query embedding, shape [bsz, seq_len, aeg_lenn,um_heaaee hidden_size.per_hent
key: key embedding,
shape [bsz,
sidden_size_per_hes
value: value embedding, ahape [baz, nem_lenn,ummheeeee sidden__ize.per_
Returns:
attention: attention output shape [bsz, seq_len, htw..u.odddiiiddd..ooo..r.......
uen
B, т, H,D = query size() query = query reshape (B T, self ori_head, self ori ori_fcature
key key reshaee T, self ori__eea self. ori ori_feature) value value. reshape (B T, ori head, self.ori_ feature)
#**e*e*e Attention Score query = self. salf. sas_q(query) reshape (B, T, self target_head, self. target_feature)
key sellf.ss_k(key) reshape self target_head, self.
target_feature)
value = val..ssssvvvvlvlla ress
reshape T, self. target_head, -1)
attention= scaled dot_ product_attention (query. transpose
key. transpose (I. 2). value transpose (I. 2). is causalatTrrre
# Attention Aggregation attention attention. transpose 2). contiguous () view (B, T,
self target_head self ori_head, self self. ori ori_feature)
attention = self self. outputuupuuuee _ettsddeeeeeeeeeeeeeeeeeeeeeeee_ attention)
attention attention mean dim=-2)
return attention
22

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: We have made the main claims in the Abstract and introduction, and also highlight them.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made 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 be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification:
We have discussed the Limitation in the Appendix A.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section 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 on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or speech-to-text system might not 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 to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
the paper. The authors should use their best
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: [Yes]
23

---

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

---

Answer: [Yes]
Justification: We have provided the dataset Section Experiment. And the code is shown in the Appendix L.
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:
/eddeS/Seuub ssi ssionnosssiiiiiioieo sesionnyocy) 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: We have provided the experiment setting details in Section Experiment and Appendix C.
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 a level detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, appendix, 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 have discussed the experiment statistical significance in Appendix D.
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).
25

---

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 or figures symmetric error bars that would yield results that are out range (e.g. negative enror rates).
If error bars reported in tables plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification:
We have discussed the training cost in Appendix K.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount compute required for each of the individual experimental runs well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary failed experiments that didn't make into the paper).
9. Code ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics https / htt////nuripppp cttppulic//thhhcccieeeeeeeeee
Answer: [Yes]
Justification:
We have reviewed the NeurlPS code of Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is special consid- eration due to laws or regulations their jurisdiction..
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification: We have discuss the both both potential positive societal impacts and negative societal impacts of the work performed Appendix
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
26

---

The conference expects that many papers will be foundational research and not tied particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: This work provides architecture but does not release well well-trained model.
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 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: We have properly credited them by citing the paper, such in the Related Work and Appendix L.
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 the asset used and, possible, include URI
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (c.g.. website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms use in package should be provided. For popular datasets, paperswi twwode. com datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
27

---

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's 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:
The paper does not release new assets.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moddl as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include anonymized zip file.
14. Crowdsourcing 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 as details about compensation (if any)?
Answer: [NA]
Justification: This work does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurIPS Code of Bthies, 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 equivalent approval/review based the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: This work does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending on 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 we expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, not include any information that would break anonymity (if applicable), such the institution conducting the review.
28

---

16. Declaration of LLM usage
Question: Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology. scientific rigorousness, originality of the research, declaration is not required.
Answer: [NA]
Justification: This work core method development in this research does not involve LLMs as any important, original, or non-standard components.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, or non-standard components.
Please refer to our LLM policy (https: : /ttp///erssp///// tteaurpp.cccCooffeeeeeeeeee.eeee /eeeneere/eeee22//////// 20LM/LLM) for what should or should be described.
29


Reviews:
[{"review_id": "QjLGUSeZwN", "reviewer": "Reviewer_yD2Q", "summary": "This paper suggests simulated attention score as a means to increase the expressivity of the model so as to improve model performance. Given the prior findings that a larger feature dimension as well as the number of heads contribute to better generalization, this method circumvents the bounded number of heads via expanding the dimensions with linear projections. The increased computational costs are compensated by PEAA which applies a grouping for feature aggregation across multiple heads. Results suggest that SAS outperforms baselines, such as MHA, consistently.", "questions": "- How does SAS compare with other baselines in terms of compute? A pareto front graph would display the value of the model better.\n\n- Are the baselines state-of-the-art? Judging by the nature of the proposed method, results are largely experimental, which may require well-tuned models for comparison purposes. How different hyperparameters or other transformer techniques may affect the performance is not clear.", "limitations": "Yes", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 1, "originality": 3, "strengths_and_weaknesses": "Strengths: \n\n- This paper has a good clarity, in that the method is built on top of popular MHA and is simple in nature, which is also well presented throughout the paper. I appreciate that the author offers the raw implementation in Appendix H, which helps reproducibility early on.\n\n- Experimental results seem strong and consistently outperforms other attention structures in multiple benchmarks.\n\n\nWeaknesses: \n\n- The motivation of the method is weak. Discussions on the rationale how the number of heads contribute to better performance are largely omitted.\n\n- Performance gains don’t seem substantial. Different hyperparameters may change the results.", "was_revised": true, "final_justification": "Thank you for the clarifications and I understood much better and by seeing other feedbacks to other reviewers, I'm more confident to raise my score.", "rebuttal": "Dear Reviewer yD2Q,\n\nThank you very much for your comment. We will address your concerns below.\n\n**Q1: The motivation of the method is weak. Discussions on the rationale how the number of heads contribute to better performance are largely omitted.**\n\nA1: **We have discussed the effect of the number of attention heads in Figure 1 and Section 4.3 Ablation Study: The Effect of Attention Head Number. For convenience, we copy the result in the following**\n\nAccording to Figure 1: When only a single attention head (1 head) is used, the model achieves a perplexity score of 6.08. However, as the number of attention heads incrementally increases, the model’s performance improves, reaching an optimal perplexity of 5.82 when the head count matches the original configuration of 12 heads. Beyond this point, further increasing the number of attention heads leads to a gradual degradation in performance, as evidenced by the rising perplexity scores.\nTherefore, more attention heads and larger hidden sizes per head may further improve the performance.\n\nIn Figure 6: when the attention number is 12, SAS achieves 22.96 ppl, outperforming MHA.\nThis implies that SAS can improve the expressiveness of query, key, and value embeddings, leading to better performance even with the same attention heads number. Moreover, increasing the attention head number from 12 to 36 leads to further performance gains, with SAS achieving a perplexity of 22.14. This demonstrates that SAS benefits from additional attention heads, which enhance the expressiveness of query, key and value embeddings and result in better performance compared to MHA.\n\n\n\n\n\n**Q2: Performance gains don’t seem substantial. Different hyperparameters may change the results.**\n\nA2: **We validate the SAS with different lengths, different datasets, different model sizes, different head numbers, different kernel sizes and so on. And the SAS consistently improves the performance.**\n* **Different Length**. We validate the SAS in Section 4.1, with lengths from 512 to 2048.\n* **Different Dataset**. We validate the SAS with the Arxiv and Books dataset. And we also use the FineWeb-Edu dataset for large-scale training and downstream tasks. Also, we use different downstream evaluation datasets, including ARC-E, ARC-C, Hellaswag, PIQA, ScIQ, SocialIQA and Winograde.\n* * **Different Model Size**. We validate the SAS with sizes 125M, 350M and 2.7B in Section 4.1 and Section 4.2.\n* **Different Head Number**. We validate the SAS with different head number in Section 4.3 with different head number.\n* **Different Kernel Size**. We validate the SAS with different kernel sizes, from 1 to 7.\n* **Different Seed**. We validate the SAS with different random seed in Appendix D.\n* **Different initial attention type**. We try different attention methods as the initial query, key and value embedding. We present the results in Appendix G.\n\n**Therefore, the SAS consistently achieves performance gains, with different lengths, different datasets, different model sizes, different kernel sizes, different seeds, different initial attention types, and so on.**\n\n**Q3: How does SAS compare with other baselines in terms of compute? A pareto front graph would display the value of the model better.**\n\nA3: We test the time cost on H100 GPUs.\n\n*The time cost (ms) with different model sizes and micro batch size 1, with H100*\n| Method | 125M  | 350M   | 2.7B    | 6.7B|10.6B|\n|--------|--------|---------|----------|----------|----------|\n| MHA    | 36.20  | 72.87   | 117.82   |186.58|217.01|\n| GQA    | 38.87  | 71.66   | 133.76   |205.69| 239.16|\n| MQA    | 39.37  | 72.01   | 132.62   |202.61| 238.85|\n| MLA    | 51.68  | 104.95  | 180.03   |254.43|637.78(Zero3 to avoid OOM;H200)\n| TPA    | 44.19  | 86.56   | 151.71   |235.85|258.41|\n| SAS    | 68.07  | 138.53  | 215.13   |350.70|402.15|\n\n* The computional cost cost of MHA:$O(12Th^2d^2)+O(T^2hd)$, while $h$ is the attention head number, $d$ is the hidden size per head and $T$ is the sequence length.\n* The computional cost of SAS: $O(11T h^2 d^2+Th'^2d'^2)+O(T^2h'd')$, while $h'$ is the simulated attention head number, $di$ is the simualted hidden size per head and $T$ is the sequence length.\n* Usually, the $h'/h$ is between 1 to 3, and the $d'/d$ is between 1 to 1.5 (less than 1.2 for larger models to control cost)\n\n\n\n**Q4: Are the baselines state-of-the-art?**\n\nA4: **The baseline is the state-of-the-art. The baselines contain the most recent works, including MLA [1] and TPA[2].**\n\n**Q5: How different hyperparameters or other transformer techniques may affect the performance is not clear.**\n\nA5: We have proved in Q2 that SAS consistently achieves performance gains, with different lengths, different datasets, different model sizes, different kernel sizes, different seeds, different initial attention types and so on. We further validate the model siz, the larger model size and training length.\n\n*The following is the performance with 6.7B and 10.6B.*\n| Model                     |5K|10K|15K|20K|25K|30K|35K| 40K|45K|50K|\n|-----|-----|-----|-----|-----|-----|-----|-----|-----|-----|-----|\n|6.7B|\n| MHA                 | 23.45 | 17.25 | 14.51 | 13.21 | 12.07 | 11.21 | 10.52 | 10.28 | 9.91 | 9.73|\n| MQA |21.56 | 16.32 | 13.93 | 12.73 | 11.70 | 10.90 | 10.24 | 10.02 | 9.66 | 9.49|\n|GQA|21.62 | 16.35 | 13.93 | 12.74 | 11.69 | 10.88 | 10.24 | 10.01 | 9.65 | 9.49|\n|MLA| 22.22 | 16.67 | 14.26 | 13.05 | 11.99 | 11.19 | 10.53 | 10.29 | 9.91 | 9.74|\n| SAS         | **20.44** | **15.32** | **13.17** | **12.15** | **11.21** | **10.51** | **9.93** | **9.73** | **9.40** | **9.25**|\n|10.6B|\n| MHA                 | 23.67 | 17.34 | 14.59 | 13.18 | 12.00 | 11.15 | 10.46 | 10.20 | 9.82 | 9.64|\n| MQA |22.13 | 16.55 | 14.03 | 12.80 | 11.70 | 10.88 | 10.21 | 9.98 | 9.61 | 9.43|\n| GQA            | 21.86 | 16.45 | 13.92 | 12.70 | 11.61 | 10.81 | 10.17 | 9.94 | 9.58 | 9.41|\n|MLA|22.40 | 16.67 | 14.22 | 12.99 | 11.90 | 11.11 | 10.42 | 10.18 | 9.83 | 9.66|\n| SAS         | **20.53** | **15.12** | **12.92** | **11.90** | **10.98** | **10.29** | **9.71** | **9.51** | **9.18** | **9.04**|\n\n*The validation perplexity with 125M model size and 16384 training length on the Pile, with H100/200*\n| Model                     | 5K|10K|15K|20K|25K|30K|35K| 40K|45K|50K|\n|-----|-----|-----|-----|-----|-----|-----|-----|-----|-----|-----|\n| MHA                 | 18.31 | 14.50 | 12.85 | 11.96 | 11.45 | 11.11 | 10.44 | 10.37 | 9.95 | 9.94|\nMQA |18.02 | 14.46 | 12.82 | 11.96 | 11.44 | 11.10 | 10.46 | 10.38 | 9.96 | 9.99|\n| GQA            | 17.74 | 14.40 | 12.86 | 11.96 | 11.47 | 11.14 | 10.45 | 10.39 | 9.99 | 9.98|\n|MLA|18.23 | 14.61 | 12.93 | 11.97 | 11.45 | 11.08 | 10.40 | 10.33 | 9.90 | 9.93|\n| SAS         | **14.93** | **12.45** | **11.23** | **10.56** | **10.16** | **9.90** | **9.32** | **9.28** | **8.95** | **8.97I\n\n**If the response addresses the concerns, could you consider improving the score?**\n\nReference:\n\n[1] Liu, A., Feng, B., Wang, B., Wang, B., Liu, B., Zhao, C., ... & Xu, Z. (2024). Deepseek-v2: A strong, economical, and efficient mixture-of-experts language model. arXiv preprint arXiv:2405.04434.\n\n[2] Zhang, Y., Liu, Y., Yuan, H., Qin, Z., Yuan, Y., Gu, Q., & Yao, A. C. (2025). Tensor product attention is all you need. arXiv preprint arXiv:2501.06425."}, {"review_id": "enG9yjUCri", "reviewer": "Reviewer_rUv9", "summary": "The authors propose a novel method of enhancing the performance of transformer language models with small extra cost to the model size, which they call SAS.\n\nThe number of attention heads and latent dimensions are 'expanded' by a non-linear transformation, before applying the attention mechanism. After that, the expanded heads are aggregated down to the normal number of heads. This 'simulates' having more attention heads and latent dimensions without increasing the number of parameters.\n\nSAS shows better perplexity on decoder-only transformers compared to MHA, MQA, GQA, MLA, and TPA, and better performance on downstream datasets.", "questions": "Did the authors conduct an experiment with larger models, preferably >10B parameters?\n\nAt the PEAA step, the authors just average out multiple 'virtual' heads together. In my view, this averaging may cause information loss contained in the virtual heads. Adding another non-linear layer to the $h_i$s, before averaging them may prevent this information loss. Did the authors attempt something like this?", "limitations": "Yes", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "**Strengths**\n\nSAS seems to slightly improve the perplexity and downstream performance on 125M - 1.5B parameter models.\n\nThe ablation study proves that artifically expanding the number of heads and latent dimensions helps model performance.\n\n**Weaknesses**\n\nWhile SAS does not increase the number of parameters by much, the increased computational cost may not be worth the slight improvement of model quality. The paper does not state how much the computational cost (e.g. latency or FLOPS) increased, or how it compares to the computational cost of other baselines.\n\nAs seen on Figures 2 and 4, it seems that the larger the models gets, the smaller the improvement in perplexity; the gap between the baselines and SAS in the 2.7B model seems to be much smaller than the gap in the 125M model. If this trend continues, it seems possible that SAS might not be very effective at sizes over 10B.", "was_revised": true, "final_justification": "The authors have provided a comparison against MHA at a fixed test time cost. SAS performs on par with MHA with much fewer parameters. This is significant and impressive. So I have raised my score.", "rebuttal": "Dear Reviewer rUv9,\n\nThank you very much for the support. We will address the concerns below.\n\n**Q1: The disucssion of computational cost**\n\nA1: We further test the time cost on H100 GPUs.\n*The time cost with different model sizes and micro batch size 1, with H100*\n| Method | 125M  | 350M   | 2.7B    | 6.7B|10.6B|\n|--------|--------|---------|----------|----------|----------|\n| MHA    | 36.20  | 72.87   | 117.82   |186.58|217.01|\n| GQA    | 38.87  | 71.66   | 133.76   |205.69| 239.16|\n| MQA    | 39.37  | 72.01   | 132.62   |202.61| 238.85|\n| MLA    | 51.68  | 104.95  | 180.03   |254.43|637.78(Zero3 to avoid OOM;H200)\n| TPA    | 44.19  | 86.56   | 151.71   |235.85|258.41|\n| SAS    | 68.07  | 138.53  | 215.13   |350.70|402.15|\n\n* The computional cost cost of MHA:$O(12Th^2d^2)+O(T^2hd)$, while $h$ is the attention head number, $d$ is the hidden size per head and $T$ is the sequence length.\n* The computional cost of SAS: $O(11T h^2 d^2+Th'^2d'^2)+O(T^2h'd')$, while $h'$ is the simulated attention head number, $di$ is the simualted hidden size per head and $T$ is the sequence length.\n* Usually, the $h'/h$ is between 1 to 3, and the $d'/d$ is between 1 to 1.5 (less than 1.2 for larger models to control cost)\n\n**Q2: The SAS with larger training (e.g. larger model size)**\n\nA2: The following is the result of SAS with model size 6.7B and 13B on the Pile dataset, which contains both the Arxiv and Books datasets. \n\n*The validation perplexity with 6.7B and 10.6B model size and 512 training length on the Pile, with H100/200*\n| Model                     | 5K|10K|15K|20K|25K|30K|35K| 40K|45K|50K|\n|-----|-----|-----|-----|-----|-----|-----|-----|-----|-----|-----|\n|6.7B|\n| MHA                 | 23.45 | 17.25 | 14.51 | 13.21 | 12.07 | 11.21 | 10.52 | 10.28 | 9.91 | 9.73|\n| MQA |21.56 | 16.32 | 13.93 | 12.73 | 11.70 | 10.90 | 10.24 | 10.02 | 9.66 | 9.49|\n|GQA|21.62 | 16.35 | 13.93 | 12.74 | 11.69 | 10.88 | 10.24 | 10.01 | 9.65 | 9.49|\n|MLA|22.22 | 16.67 | 14.26 | 13.05 | 11.99 | 11.19 | 10.53 | 10.29 | 9.91 | 9.74|\n| SAS         | **20.44** | **15.32** | **13.17** | **12.15** | **11.21** | **10.51** | **9.93** | **9.73** | **9.40** | **9.25**|\n|10.6B|\n| MHA                 | 23.67 | 17.34 | 14.59 | 13.18 | 12.00 | 11.15 | 10.46 | 10.20 | 9.82 | 9.64|\n| MQA |22.13 | 16.55 | 14.03 | 12.80 | 11.70 | 10.88 | 10.21 | 9.98 | 9.61 | 9.43|\n| GQA            | 21.86 | 16.45 | 13.92 | 12.70 | 11.61 | 10.81 | 10.17 | 9.94 | 9.58 | 9.41|\n|MLA|22.40 | 16.67 | 14.22 | 12.99 | 11.90 | 11.11 | 10.42 | 10.18 | 9.83 | 9.66|\n| SAS         | **20.53** | **15.12** | **12.92** | **11.90** | **10.98** | **10.29** | **9.71** | **9.51** | **9.18** | **9.04**|\n* From model size 125M to 10.6B, the SAS always achieves the best performance. And the improvement of SAS in 6.7B and 10.6B is still large enough.\n* The TPA meets overflow performance, and the loss scale is reduced to a minimum of 1 at step 3600 with fp16, and the loss becomes nan with bf16, with a larger model. \n\n\n\n\n\n\n\n\n**Q3: Adding another non-linear layer before averaging them**\n\nA3: The following is the results. The initial embedding is MHA. And the following is the operation (use 125M as example):**SAS-NonLinearBack**:\n* After SAS attetnion, we will have tensor with shape [36 64], while 36 is the head number and 64 is the hidde size.\n* For SAS, we will reshape it to [3 12 64] and then mean pooling to have [1 12 64].\n* For **SAS-NonLinearBack**, we will have a two-layer-gelu (MLP-GeLU-MLP) on the [36 64] map to back to [12 64] (first MLP [36 64] to [36 64]; second MLP [36 64] to [12 64]).\n\n*The perplexity result on the Books dataset with length 512 and H100*\n\n| Model                     | 5K|10K|15K|20K|25K|30K|35K| 40K|45K|50K|\n|-----|-----|-----|-----|-----|-----|-----|-----|-----|-----|-----|\n| MHA                 | 59.23 | 46.31 | 44.68 | 39.33 | 36.57 | 35.44 | 33.19 | 30.35 | 30.29 | 28.68 |\n|SAS-NonLinearBack|54.85 | 43.94 | 42.62 | 37.62 | 34.93 | 34.05 | 31.90 | 29.15 | 29.10 | 27.65 |\n| SAS         | **54.34** | **43.56** | **42.36** | **37.31** | **34.73** | **33.84** | **31.70** | **29.13** | **29.07** | **27.52** |\n\nAccording to the above results, the mean pooling achieves better performance. This may be caused by the mean pooling is easier for the network to be optimized.\n\n**Q4: On L200, 'Linear_2' is probably meant by the second 'Linear_1'.**\n\nA4: Thank you very much for the notice. We have fixed it in the revision.\n\n**If the response addresses the concerns, could you consider improving the score?**"}, {"review_id": "12LOm9cwx1", "reviewer": "Reviewer_wQTE", "summary": "In this paper, the authors introduce Simulated Attention Score to improve Transformer performance by simulating a larger number of attention heads and hidden feature dimensions without significantly increasing the model's parameter count.\nSAS achieves this through non-linear mappings that project low-dimensional head and feature representations into higher-dimensional spaces.\nThe authors also propose Parameter-Efficient Attention Aggregation (PEAA) to manage parameter costs.\nExperiments demonstrate that the proposed SAS outperforms existing attention mechanisms, achieving lower perplexity and faster convergence.", "questions": "See weaknesses", "limitations": "Yes", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 2, "significance": 2, "originality": 3, "strengths_and_weaknesses": "# Strengths\n\n1. The main idea of using more attention heads and larger dimensions is simple and reasonable.\n\n2. The experimental results show that the proposed SAS achieves promising improvements.\n\n\n# Weaknesses\n\n1. The writing of this paper is not clear and confuses me. For example:\n    - In Section 3.1, the discussion of using only one attention head (when the base model has 12) or 24 heads (when the original is 12) is confusing and requires clarification.\n    - In Sec. 4, we do not know the configuration of the larger models (350M and 2.7B). Although we know the configuration of 125M models, the authors do not explain the relationship between the configuration data and Equations 1-12. I don’t fully understand the dimension sizes for each part. A diagram might help clarify this.\n\n2. SAS introduces additional layers and residual connections, yet there is no discussion of its impact on inference speed, particularly for long-context scenarios. This is a critical consideration for practical deployment.\n\n3. The training lengths of 512 and 1024 are too short. It remains unclear whether SAS maintains its advantages in long-context settings (e.g., >=16K tokens).\n\n4. A thorough double-checking of all reported metrics is necessary to ensure accuracy. For example, in Table 1, the 350M SAS model does not consistently outperform baselines (e.g., in SocialIQA and Winogrande), yet the authors highlight it in bold as the best performer.", "was_revised": true, "final_justification": "The authors provided very detailed supplementary experiments that effectively addressed all my concerns. Therefore, I have decided to raise my score to borderline accept to reflect this.", "rebuttal": "Dear Reviewer wQTE,\n\nThank you very much for your comment. We will address your concerns below.\n\n**Q1: In Section 3.1, the discussion of using only one attention head (when the base model has 12) or 24 heads (when the original is 12) is confusing and requires clarification.**\nA1: According to Figure 1, we can find that with the attention head number gradually increasing and the hidden size per head gradually decreasing, the ppl first decrease and then increase. Then, we assume that if the hidden size is large enough, the performance may gradually increase as the attention head number increases.\nHow can we make the attention head and hidden size per head large? We then propose to use non-linear mapping to make them larger, which is the simulated attention score.\n\n**Q2: We do not know the configuration of the larger models (350M and 2.7B). Relationship between the configuration data and Equations 1-12.**\n\nA2: A:  We further provide the configuration in the following.\n| **Model** | Hidden Size | Head | Hidden Size Per Head | Group Size |\n|-----------|-------------|------|----------------------|------------|\n| *125M parameters* |  |  |  |  |\n| MHA       | 768         | 12   | 64                   | -          |\n| MQA       | 768         | 23   | 64                   | 23         |\n| GQA       | 768         | 22   | 64                   | 11         |\n| MLA       | 768         | 13   | 64                   | -          |\n| TPA       | 768         | 36   | 64                   | -          |\n| SAS       | 768         | 36   | 96                   | -          |\n| *350M parameters* |  |  |  |  |\n| MHA       | 1024         | 16   | 64                   | -          |\n| MQA       | 1024        | 31   | 64                   | 31         |\n| GQA       | 1024         | 30   | 64                   | 15         |\n| MLA       | 1024         | 23   | 64                   | -          |\n| TPA       | 1024         | 48   | 64                   | -          |\n| SAS       | 1024         | 48   | 96                   | -          |\n| *2.7B parameters* |  |  |  |  |\n| MHA       | 2560         | 32   | 80                   | -          |\n| MQA       | 2560        | 63   | 80                   |  63        |\n| GQA       | 2560         | 62   | 80                   | 31         |\n| MLA       | 2560         | 54   | 80                   | -          |\n| TPA       | 2560         | 96   | 80                   | -          |\n| SAS       | 2560         | 96   | 96                   | -          |\n\nWhy do we have such a configuration? Our configuration follows the GPT-3 paper [1] and TPA paper [2].\n* From 125M to 2.7B, the layer number, hidden size, learning rate, and so on follow the GPT-3 paper.\n* We simulate triple attention heads (e.g., 36 simulated heads for 125M with original 12 heads) because the previous model TPA uses the largest attention head number, which is almost 36 heads for 1the 25M model.\n\nFor the hidden size (use 125M model as an example):\n* For MHA:\n   * hidden size: 768. This means that the hidden size is 768 for token embedding.\n   * Attention, head number. The 125M has 12 attention heads and 64 hidden size per head, so that 12 x 64=768.\n   * Given a token with embedding 768, we will have query, key, and value, all with shape 768. And then we will reshape them to [12 64] so that we will have 12 attention heads and 64 hidden states per head.\n* For SAS\n   * We will use non-linear mapping to increase the attention head number, usually the triple attention head number, so that we have query/key with shape [36 64].\n   * Also, we will increase the hidden size per head so that we may have [36, 80] for query and key. Therefore, with a very small parameter cost, we simulate more attention heads and hidden size per head.\n\n**Q3: The time cost**\n\nA3: We evaluate the time cost via H100.\n\n*The time cost with different model sizes and micro batch size 1, with H100*\n| Method | 125M  | 350M   | 2.7B    | 6.7B|10.6B|\n|--------|--------|---------|----------|----------|----------|\n| MHA    | 36.20  | 72.87   | 117.82   |186.58|217.01|\n| GQA    | 38.87  | 71.66   | 133.76   |205.69| 239.16|\n| MQA    | 39.37  | 72.01   | 132.62   |202.61| 238.85|\n| MLA    | 51.68  | 104.95  | 180.03   |254.43|637.78(Zero3 to avoid OOM;H200)\n| TPA    | 44.19  | 86.56   | 151.71   |235.85|258.41|\n| SAS    | 68.07  | 138.53  | 215.13   |350.70|402.15|\n\n* The computional cost cost of MHA:$O(12Th^2d^2)+O(T^2hd)$, while $h$ is the attention head number, $d$ is the hidden size per head and $T$ is the sequence length.\n* The computional cost of SAS: $O(11T h^2 d^2+Th'^2d'^2)+O(T^2h'd')$, while $h'$ is the simulated attention head number, $di$ is the simualted hidden size per head and $T$ is the sequence length.\n* Usually, the $h'/h$ is between 1 to 3, and the $d'/d$ is between 1 to 1.5 (less than 1.2 for larger models to control cost)\n\n**Q4: The training lengths of 512 and 1024 are too short. It remains unclear whether SAS maintains its advantages in long-context settings**\n\nA4: We have provided the SAS performance with length 2048 in Section 4.1. And we also conducted an experiment with SAS with a length of 16384. The following are the results.\n\n*The validation perplexity with 125M model size and 16384 training length on the Pile, with H100/200*\n| Model                     | 5K|10K|15K|20K|25K|30K|35K| 40K|45K|50K|\n|-----|-----|-----|-----|-----|-----|-----|-----|-----|-----|-----|\n| MHA                 | 18.31 | 14.50 | 12.85 | 11.96 | 11.45 | 11.11 | 10.44 | 10.37 | 9.95 | 9.94|\nMQA |18.02 | 14.46 | 12.82 | 11.96 | 11.44 | 11.10 | 10.46 | 10.38 | 9.96 | 9.99|\n| GQA            | 17.74 | 14.40 | 12.86 | 11.96 | 11.47 | 11.14 | 10.45 | 10.39 | 9.99 | 9.98|\n|MLA|18.23 | 14.61 | 12.93 | 11.97 | 11.45 | 11.08 | 10.40 | 10.33 | 9.90 | 9.93|\n| SAS         | **14.93** | **12.45** | **11.23** | **10.56** | **10.16** | **9.90** | **9.32** | **9.28** | **8.95** | **8.97**|\n\n**Q5: thorough double-checking of all reported metrics is necessary to ensure accuracy.**\n\nA5: Thank you very much for your notice. We have carefully checked the accuracy and the highlight.\n\nReference:\n\n[1] Brown, T., Mann, B., Ryder, N., Subbiah, M., Kaplan, J. D., Dhariwal, P., ... & Amodei, D. (2020). Language models are few-shot learners. Advances in neural information processing systems, 33, 1877-1901.\n\n[2] Zhang, Y., Liu, Y., Yuan, H., Qin, Z., Yuan, Y., Gu, Q., & Yao, A. C. (2025). Tensor product attention is all you need. arXiv preprint arXiv:2501.06425.\n\n**If the response addresses the concerns, could you consider improving the score?**"}, {"review_id": "aSzQ2MZADy", "reviewer": "Reviewer_DaRN", "summary": "This paper proposes Simulated Attention Score, a technique that upscales the number of attention heads by applying a 2-layer residual 1-d convolution over the existing K, Q, and V tensors that increases the head dimension. It also uses a 2-layer residual MLP to increase the dimensionality of the K and Q tensors for the purposes of computing the attention weights.\n\nThe introduction motivates this idea of upscaling the number of attention heads through an additional feedforward layer with non-linearity.\n\nThe presentation is confusing. The authors in 3.3 begin with a detailed presentation of an MLP-based dimension upscaling for the heads. They then switch at line 162 to present a new convolution-based expansion of heads. I’m not sure, and the text gives little clue about how these two approaches relate, but I think (based on the code in Appendix G) the first MLP-based version is not used, and the second convolution-based one is used. Perhaps the authors can clarify this, and if it’s correct, can remove the initial section that’s not used.\n\nAs I understand the technique, the idea is as follows:\n* Apply a couple of convolutional layers to the HxD block corresponding to the query tensor, producing an H’xD output with additional heads. \n* Then apply a 2-layer MLP to this producing an H’xD’ ouptut with each new head having an upscaled dimension. This is the new query tensor\n* Do the same to the key tensor\n* Likewise for the value tensor, but without the upscaling of the feature dimension -- just increase the head dimension\n* Compute attention based on these new tensors\n* Mean-pool blocks of H’/H attention values to produce the final H-dimensional attention output\nOtherwise, the transformer operates in the standard manner.", "questions": "It appears from Appendix F that the training time of SAS is significantly higher than all other methods, except in the case of the larger 2.7B model. This raises several questions:\nWhy is this true? Why is SAS slower to train for smaller model sizes? And what changes in the 2.7B model?\nHow do these numbers change for batch size >1 – setting batchsize 1 doesn’t seem consistent with best practices for pre-training.\nWhat happens to the comparisons at equivalent wallclock training time? That is, if we increase the number of MHA heads until the training time on 350M is equivalent to SAS, how does the performance compare? Or, maybe a more general way to ask this: given a higher budget of clock training time, can you argue that SAS is the best way to use this additional budget, compared to other alternatives including adding heads or layers or width or feature dimensionality to previous techniques?\n\nSome other questions:\nWhy did you adopt the mean pooling to aggregate the attention scores? Is upscaling the full H’ outputs back to the original dimensionality of the next transformer MLP too expensive? Was this decision the result of experiments showing this was a necessary approach?\nCould you please clarify the number of weights in the convolution layers? Is it H_{in} * H_{out} for the first layer, and H_{out} * H_{out} for the second one?\nWhy upscale the head dimension before the feature dimension? Does the opposite approach not work as well?", "limitations": "I don't see any concerns on negative societal effects. The contriubtion represents a small change to the dynamics of the attention layer.", "rating": 3, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 2, "significance": 2, "originality": 2, "strengths_and_weaknesses": "Strengths:\n* The approach is straightforward and easy to implement\n* The performance analysis shows significant improvements in perplexity, though I have a number of questions below\n* There is some evidence of increased performance on downstream tasks\nWeaknesses:\n* There are a number of decisions that are not fully motivated\n* The wallclock time for training is significantly longer for SAS than the competitors. The baselines are not scaled appropriately.\n* There is little analysis of total number of parameters, though this is likely to be favorable for SAS\n* The originality of the approach is not high. The decision about how to upscale the heads and hidden attention dimensions is interesting (convolution in 1 case, MLP in the other), but not studied or fully motivated\n* Clarity of presentation is adequate but could be improved\n* Significance of findings misses a number of questions that are necessary to substantiate the improvements", "was_revised": true, "final_justification": "I think the analysis is thorough and the authors did a great job through the rebuttal. I'd love to see this paper appear. However, I still have a couple of points that cause me to retain my original rating. SAS appears to introduce an increase in training overhead per step in order to improve the test metrics as a function of number of training steps. As to whether this is a win overall, the question remains open:\n\n* For the benchmarks given by the authors, the reduction in perplexity seems to be equivalent to non-SAS with about 20K additional training steps -- in general, as SAS seems to add 2x-4x overheard per training step, this is generally not a favorable trade-off\n* As the authors point out, if the final converged quality is better under SAS, the additional overhead may well be worth it. Unfortunately, I don't think any experiments demonstrate this, as they don't seem to have reached convergence.\n* The authors give one example in which SAS performs better than non-SAS at equivalent wall-clock time: it is the 2.7B SAS model compared to the larger (10B) MHA model. This is a good example, but the larger models for this training setup don't seem to provide much benefit over the smaller ones. That is, one can always use a model that is too big in order to provide an example with very high wall-clock time and low quality, and this may be the situation.", "rebuttal": "Dear Reviewer DaRN,\n\nThank you very much for your comment. We will address your concerns below.\n\n**Q1: The discussion of the time cost**\n\nA: Thank you very much for the notice. We recheck and retest the time cost, and find that some 2.7B results of baseline may be recorded with batch size 4 (the GPT-NeoX default training batch size), or the GPU cluster is also occupied by other programs during the speed test. **For the updated test, we use a local standalone machine that has 80GB-HBM3 H100 GPUs**. The following are the updated results.**\n\n*The time cost with different model sizes and micro batch size 1, with H100*\n| Method | 125M  | 350M   | 2.7B    | 6.7B|10.6B|\n|--------|--------|---------|----------|----------|----------|\n| MHA    | 36.20  | 72.87   | 117.82   |186.58|217.01|\n| GQA    | 38.87  | 71.66   | 133.76   |205.69| 239.16|\n| MQA    | 39.37  | 72.01   | 132.62   |202.61| 238.85|\n| MLA    | 51.68  | 104.95  | 180.03   |254.43|637.78(Zero3 to avoid OOM;H200)\n| TPA    | 44.19  | 86.56   | 151.71   |235.85|258.41|\n| SAS    | 68.07  | 138.53  | 215.13   |350.70|402.15|\n\n\n*The time cost with different model sizes and micro batch size 4*\n| Method | 125M  | 350M   | 2.7B    | \n|--------|--------|---------|----------|\n| MHA    | 38.01  | 77.34   | 186.03   |\n| GQA    | 45.26  | 78.59   | 245.86   |\n| MQA    | 44.07  | 76.23   | 243.67   |\n| MLA    | 53.40  | 108.65  | 258.06   |\n| TPA    | 56.37  | 116.04  | 318.31   |\n| SAS    | 89.70  | 200.47  | 531.17   |\n\n* The computional cost cost of MHA:$O(12Th^2d^2)+O(T^2hd)$, while $h$ is the attention head number, $d$ is the hidden size per head and $T$ is the sequence length.\n* The computional cost of SAS: $O(11T h^2 d^2+Th'^2d'^2)+O(T^2h'd')$, while $h'$ is the simulated attention head number, $di$ is the simualted hidden size per head and $T$ is the sequence length.\n* Usually, the $h'/h$ is between 1 to 3, and the $d'/d$ is between 1 to 1.5 (less than 1.2 for larger models to control cost)\n\n\n\n\n\n\n\n**The discussion of the next step of the model architecture**\n* **The computational cost becomes cheaper and cheaper**. With the development of GPUs, the computational cost becomes cheaper and cheaper. This is the basic reason that Large Language Models could be true.\n* **The high-quality data is limited**. The current Large Language Model almost uses all the web data.\n* **Therefore, we have to develop more powerful models with higher expressiveness**. Compared to linear models (e.g. LSTM, RNN), the Transformer has higher cost but performance. And now, most STOA models usually use the Transformer. There are also recent works that propose the model with cost $T^3$ (e.g., DeltaFormer, Simplex Attention), where the $T$ is the sequence, while the Transformer has cost $T^2$. **More powerful with higher cost may be the next-generation architecture.**\n\n**Q2: Or, maybe a more general way to ask this: given a higher budget of clock training time, can you argue that SAS is the best way to use this additional budget, compared to other alternatives including adding heads or layers or width or feature dimensionality to previous techniques?**\n\n\n\nA2: **We can be sure that using the same parameter number budget, the SAS could achieve the BEST Performance, compared to previous techniques. With the comparable budget of clock training time, the SAS has the potential to have better performance**\n* **With the same parameter number budget, the SAS could have the best performance**. This has already been discussed in the paper experiment part.\n* **With the same/comparable time, the SAS has the potential to have better performance.**\n\n*The following is the perplexity performance with 6.7B and 10.6B.*\n| Model                     |5K|10K|15K|20K|25K|30K|35K| 40K|45K|50K|\n|-----|-----|-----|-----|-----|-----|-----|-----|-----|-----|-----|\n|6.7B|\n| MHA                 | 23.45 | 17.25 | 14.51 | 13.21 | 12.07 | 11.21 | 10.52 | 10.28 | 9.91 | 9.73|\n| MQA |21.56 | 16.32 | 13.93 | 12.73 | 11.70 | 10.90 | 10.24 | 10.02 | 9.66 | 9.49|\n|GQA|21.62 | 16.35 | 13.93 | 12.74 | 11.69 | 10.88 | 10.24 | 10.01 | 9.65 | 9.49|\n|MLA|22.22 | 16.67 | 14.26 | 13.05 | 11.99 | 11.19 | 10.53 | 10.29 | 9.91 | 9.74|\n| SAS         | **20.44** | **15.32** | **13.17** | **12.15** | **11.21** | **10.51** | **9.93** | **9.73** | **9.40** | **9.25**|\n|10.6B|\n| MHA                 | 23.67 | 17.34 | 14.59 | 13.18 | 12.00 | 11.15 | 10.46 | 10.20 | 9.82 | 9.64|\n| MQA |22.13 | 16.55 | 14.03 | 12.80 | 11.70 | 10.88 | 10.21 | 9.98 | 9.61 | 9.43|\n| GQA            | 21.86 | 16.45 | 13.92 | 12.70 | 11.61 | 10.81 | 10.17 | 9.94 | 9.58 | 9.41|\n|MLA|22.40 | 16.67 | 14.22 | 12.99 | 11.90 | 11.11 | 10.42 | 10.18 | 9.83 | 9.66|\n| SAS         | **20.53** | **15.12** | **12.92** | **11.90** | **10.98** | **10.29** | **9.71** | **9.51** | **9.18** | **9.04**|\n\n\n* With the same parameter cost, the SAS is definitely better than other methods. We have proved this in the paper.\n* With the comparable clocking time, the SAS has the potential to have better performance. For example, the SAS with 6.7B is significantly better than MHA, GQA, MQA, MLA with 10.6B.\n\n**Adding Head**\n* We have discussed the effect of attention head in Figure 1. With MHA and the same parameter cost and hidden size, the SAS is always better than MHA, whatever the attention head number is 1 or 96.\n\n**Adding Layer or Width or Feature Dimensionality**.\n* Such an operation will lead to a higher model parameter cost. **Our setting is using the same parameter number to achieve better performance**\n\n\nThe discussion of previous works\n* **Mixture-of-Experts (MoE) could achieve better performance than da ense model with the same activated parameters, while the MoE needs more parameter cost.**\n* **Transformer usually is better than a linear model with the same parameter cost, while Transformer needs a quadratic cost of sequence length and more time cost.**\n* **Similarly, this work targets to improve the performance with the same/comparable parameter number**.\n\n**Q3: Why did you adopt the mean pooling to aggregate the attention scores? Is upscaling the full H’ outputs back to the original dimensionality of the next transformer MLP too expensive?**\n\nA3: **The Attention output feature dimension should still be the hidden size (e.g., 768 for 125M). Also, we have to control the number so that we can use mean pooling. Ifwe directly map the H' back to the original dimension, the parameter may increase by 25%. Our setting is using the same parameter number to achieve better performance so that we do not directly map the full H' back. \n\nWhy average pooling (use 125M as an example):\n* Originally, the attention output will have shape [12 64] so that we could reshape it to [768] and then multiply a matrix with shape [768 768].\n* For SAS, the attention output will have [36 64]. Therefore, we have to first reshape it to [3 12 64] and then average along the first dimension to have [1 12 64]. Then, we reshape it to have [768] and multiply with the matrix [768 768].\n\n**Let us calculate the parameter cost for upscaling the full H outputs back**\n* Assume we simulate triple attention heads. Therefore, the additional parameter cost is $3D^2$. \n* The FFN takes $8D^2$ parameter cost.\n* The original Attention takes $4D^2$ cost.\n* Therefore, **upscaling the full H outputs back will take an addtional 25% parameter cost.**\n* **Our experiment setting uses the same parameter cost to improve the performance. Therefore, we cannot afford such an expensive additional parameter cost.**\n\n\n**Q4: Could you please clarify the number of weights in the convolution layers? Is it H_{in} * H_{out} for the first layer, and H_{out} * H_{out} for the second one?**\n\nA4: Yes, the understanding is correct. Finally, there are $H_{out}$ simulated attention heads.\n\n**Q5:Why upscale the head dimension before the feature dimension? Does the opposite approach not work as well?**\n\nA: Thank you very much for the suggestion. We first realize that more attention is necessary, and then extend it to feature expansion. And this is why we first process the attention head, and then the feature expansion. The following is the result of processing the feature dimension first and then the attention head.\n\n*The perplexity result on the Books dataset with length 512 and H100*\n\n| Model                     | 5K|10K|15K|20K|25K|30K|35K| 40K|45K|50K|\n|-----|-----|-----|-----|-----|-----|-----|-----|-----|-----|-----|\n| MHA                 | 59.23 | 46.31 | 44.68 | 39.33 | 36.57 | 35.44 | 33.19 | 30.35 | 30.29 | 28.68 |\n| SAS_FeatureFirst            | **53.98** | **43.44** | **42.35** | **37.27** | 34.86 | 33.94 | 31.80 | 29.15 | **29.06** | 27.63 |\n| SAS         | 54.34 | 43.56 | 42.36 | 37.31 | **34.73** | **33.84** | **31.70** | **29.13** | 29.07 | **27.52** |\n\nAccording to the above result, the SAS_FeatureFirst   achieves better performance at the early stage, while the SAS finally achieves better performance. This may be caused by that the SAS (head first) utilizes more computation cost. \n*  SAS_FeatureFirst:12x(64x96+96x96)+96x(12x36+36x36)=350208; \n* SAS: 64x(12x36+36x36)+36x(64x96x96x96)=663552.\n\n**Q6: or example, equation 5 vs 7 both use Q_h for Conv_1 and Conv_2 respectively, even though the dimensionalities of the operators are different.**\n\nA6: The $Q_h$ means the operation is for Query Head. The Conv1 is the first operation (e.g., mapping head from 12 to 36), and the Conv2 is the second operation (e.g., from head number 36 to 36). The $Q_h^{Conv1}$ is the first convolution operation for Query Head, so the $Q_h^{Conv2}$ is the second convolution operation for Query Head.\n\n**If the response addresses the concerns, could you consider improving the score?**"}]

Author rebuttal:

Meta-review: The paper develops a novel concept of simulated attention head and simulated hidden size per head. This concept should translate to performance lifts with smaller models and make them as powerful as larger models. This is a very interesting idea and there was consensus about this. However, the major concern was the overhead of training such models might outweigh the advantage of training  smaller models. This trade off seems to be not clear. During post rebuttal discussion the following points emerged.

<em>For the benchmarks given by the authors, the reduction in perplexity seems to be equivalent to non-SAS with about 20K additional training steps -- in general, as SAS seems to add 2x-4x overheard per training step, this is generally not a favorable trade-off</em>

<em>The authors give one example in which SAS performs better than non-SAS at equivalent wall-clock time: it is the 2.7B SAS model compared to the larger (10B) MHA model. This is a good example, but the larger models for this training setup don't seem to provide much benefit over the smaller ones. That is, one can always use a model that is too big in order to provide an example with very high wall-clock time and low quality, and this may be the situation</em> 

Given these concerns, it is thus not absolutely clear on how much impact the paper will have on the practice of large transformers and hence it is difficult to make a case of strong acceptance. Nevertheless, 
the paper already brings fresh ideas which should be of interest to Neurips community. If accepted,  the author(s) should  modify the manuscript with all the additional experiments discussed in rebuttal and incorporate appropriate commentary in the manuscript to address the concerns listed above.
Decision: Accept (poster)