SAS: Simulated Attention Score

TL;DR

This work investigates how the number of attention heads and the per-head hidden size influence Transformer performance. It introduces Simulated Attention Score (SAS), which expands head and feature dimensions through non-linear projections to imitate a larger model while preserving parameter count, and extends this idea to Q/K embeddings. To manage costs, SAS pairs with Parameter-Efficient Attention Aggregation (PEAA), aggregating expanded-head outputs efficiently. Across several datasets, model sizes, and pretraining scales, SAS consistently achieves lower perplexities and stronger downstream performance than standard MHA and other variants, demonstrating that moderately sized models can approximate the attention capacity of larger models with meaningful practical benefits.

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.

Figures (7)

• 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.
• Figure 2: The performance of different methods on the Arxiv and Books3 dataset, with training lengths of 512 and 1024.
• Figure 3: The performance on training length 2048.
• Figure 4: The performance of different methods on the Books3 dataset, with model sizes of 350M and 2.7 B.
• Figure 5: The Ablation study of SAS (SAS-HeadOnly and SAS-FeatureOnly) on the Books3 dataset, with training length 1024.