Title: SAS: Simulated Attention Score
Venue: NeurIPS 2025
Accepted: yes
Keywords: sequence modeling

language models

model architecture

Transformer
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.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\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\n\nQUESTIONS:\nIt 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?\n\nLIMITATIONS:\nI don't see any concerns on negative societal effects. The contriubtion represents a small change to the dynamics of the attention layer.", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 3, "CLARITY": 2, "IMPACT": 2, "ORIGINALITY": 2, "DATE": "2025-06-28"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nIn 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.\n\nSTRENGTHS AND WEAKNESSES:\n# 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.\n\nQUESTIONS:\nSee weaknesses\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 2, "IMPACT": 2, "ORIGINALITY": 3, "DATE": "2025-06-30"}

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

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

STRENGTHS AND WEAKNESSES:
Strengths: 

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

- Experimental results seem strong and consistently outperforms other attention structures in multiple benchmarks.


Weaknesses: 

- The motivation of the method is weak. Discussions on the rationale how the number of heads contribute to better performance are largely omitted.

- Performance gains don’t seem substantial. Different hyperparameters may change the results.

QUESTIONS:
- How does SAS compare with other baselines in terms of compute? A pareto front graph would display the value of the model better.

- 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
TITLE: Official Review
CLARITY: 3
IMPACT: 1
ORIGINALITY: 3
DATE: 2025-07-06