Dear Reviewer DaRN,

Thank you very much for your comment. We will address your concerns below.

Q1: The discussion of the time cost

A: 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.**

The time cost with different model sizes and micro batch size 1, with H100

The time cost with different model sizes and micro batch size 4

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

The discussion of the next step of the model architecture

• 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.
• The high-quality data is limited. The current Large Language Model almost uses all the web data.
• 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.

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?

A2: 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

• With the same parameter number budget, the SAS could have the best performance. This has already been discussed in the paper experiment part.
• With the same/comparable time, the SAS has the potential to have better performance.

The following is the perplexity performance with 6.7B and 10.6B.

• With the same parameter cost, the SAS is definitely better than other methods. We have proved this in the paper.
• 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.

Adding Head

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

Adding Layer or Width or Feature Dimensionality.

• Such an operation will lead to a higher model parameter cost. Our setting is using the same parameter number to achieve better performance

The discussion of previous works

• Mixture-of-Experts (MoE) could achieve better performance than da ense model with the same activated parameters, while the MoE needs more parameter cost.
• 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.
• Similarly, this work targets to improve the performance with the same/comparable parameter number.

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?

A3: **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.

Why average pooling (use 125M as an example):

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

Let us calculate the parameter cost for upscaling the full H outputs back

• Assume we simulate triple attention heads. Therefore, the additional parameter cost is $3D^2$.
• The FFN takes $8D^2$ parameter cost.
• The original Attention takes $4D^2$ cost.
• Therefore, upscaling the full H outputs back will take an addtional 25% parameter cost.
• Our experiment setting uses the same parameter cost to improve the performance. Therefore, we cannot afford such an expensive additional parameter cost.

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?

A4: Yes, the understanding is correct. Finally, there are $H_{out}$ simulated attention heads.

Q5:Why upscale the head dimension before the feature dimension? Does the opposite approach not work as well?

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

The perplexity result on the Books dataset with length 512 and H100

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

• SAS_FeatureFirst:12x(64x96+96x96)+96x(12x36+36x36)=350208;
• SAS: 64x(12x36+36x36)+36x(64x96x96x96)=663552.

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.

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

If the response addresses the concerns, could you consider improving the score?

Dear Reviewer DaRN,

We further conducted an experiment of SAS with 2.7B, which has a comparable time cost to baseline MHA 10.6B.

The result on the Pile dataset with length 512 and H100

• In this work, we aim to achieve better performance with the same/comparable parameter cost. Therefore, we have confidence to claim that the SAS is definitely better than all other methods with the same parameter cost
• With the same time cost, the SAS has the potential to achieve comparable performance. For example, SAS 2.7 achieves comparable performance with MHA 10.6B.

Dear Reviewer DaRN,

We have addressed the concerns in the above response, including the time cost, performance with the comparable time cost (2.7B SAS has comparable performance with 10.6B MHA), the discussion of mapping H', process feature dimension first and so on. If there are any questions, please let us know before the discussion period ends (August 8 AoE). After the detailed explanation, if possible, could you consider improving the score?

Again, thank you for your attention, and wish you a good day.

Dear Reviewer DaRN,

Thank you for the message. We will answer your question below.

Q1: Performance of MHA on Books, for instance, shows a significant improvement from 45K to 50K. could you clarify the results to support the fully converged quality numbers?

A1: We provide the large-scale training results with larger training steps to support that the SAS is still better than MHA. Also, we have to point out that as long as we further add training data/increase training steps (even with trillion tokens), the loss will gradually decrease (e.g., loss drops from 1.4 to 1.3 with training tokens increasing from 12T to 15T), which is supported by recent work Kimi K2 Figure 3 [1].

The loss on FinWeb-Edu with training step 100K, 50K tokens, and cosine learning rate scheduler

At the end of the training, the loss reduction is not significant compared to previous steps. For example, from 80K to 90K, the loss of MHA reduces by 0.04, while from 90K to 100K, the loss reduces only by 0.02. Also, even if we double the training steps, the SAS is still better than MHA. We can be sure that with the same training steps, the SAS is always better than MHA. We have to point out that as long as adding training data/training steps, the loss will gradually decrease (e.g., loss drops from 1.4 to 1.3 with training tokens increasing from 12T to 15T), as supported by Kimi K2.

As long as the training length, training tokens and training steps, and others are the same, with the same parameter cost, we have confidence to say that SAS is better than MHA.

Q2: Can you make the same argument for other datasets where there is a clear advantage to increased model size?

A2: The main target of this work is to increase performance with the same Parameter Cost.

• We have confidence to say that with the same parameter cost, we can be sure that SAS is better than MHA.
• With the same wall-clocking time, SAS has the potential to achieve comparable performance than MHA, and this depends on the model size, datasets, and so on.

The following question of Reviewer DaRN may be: if we only promise that SAS achieves better performance than SAS with the same parameter cost but not wall-clocking time, then why should we believe SAS is a good work?

• Computational Cost becomes cheaper. In the past several years, the GPUs have become more and more powerful, and we can train larger models at a larger cost.
• As we know, the Transformer is only able to solve the $NC^0$ problem [2]. Therefore, we need a more powerful model to further solve problems that cannot be solved by Transformer, even if the computational cost is higher. Recently, there have also been other researchers agreed, so that they have developed models with $O(T^3)$ time cost, such as Deltaformer [2], Simplex Attention [3], and so on.
• Large Cost may be the future.
  • We increase more costs on the parameter so that we have the Mixture-of-Experts.
  • We increase more costs on the model, so we have a Transformer from LSTM.
  • We scale the model up with more cost, including data, model size, and so on, so we have a large language model.

We hope that the above insight could address the concerns. If there are any other questions, please let us know.

Reference:

[1] Team, K., Bai, Y., Bao, Y., Chen, G., Chen, J., Chen, N., ... & Zhang, H. (2025). Kimi K2: Open Agentic Intelligence. arXiv preprint arXiv:2507.20534.

[2] Zhong, S., Xu, M., Ao, T., & Shi, G. (2025). Understanding Transformer from the Perspective of Associative Memory. arXiv preprint arXiv:2505.19488.

[3] Roy, A., Chou, T., Duvvuri, S. S., Chen, S., Yu, J., Wang, X., ... & Anil, R. (2025). Fast and Simplex: 2-Simplicial Attention in Triton. arXiv preprint arXiv:2507.02754.

Dear Reviewer DaRN,

Thank you very much for the message. We will answer your question.

Before the explanation, we would like to reclaim our argument:

• With the same parameter cost, the SAS is definitely better than MHA, as long as others are the same, including training step, batch size, learning rate, and so on.
• With the same time cost, the SAS only has the potential (but not promising) to have comparable performance, as MHA with the same time cost either has more training data or a larger model size.

Q1: SAS at equivalent training steps outperforms MHA, but the wall clock time is longer. demonstration that for the same architecture and the same training time, even with smaller training step

A1: Our main claim is: SAS at equivalent training steps outperforms MHA. But we never promise that SAS outperforms MHA with the same blocking time, but SAS only has the potential to have comparable performance with MHA with a comparable time cost.

We have to notice that, if we directly use previous results by reducing the training step, we also change the other hyperparameters, such as training data size. For example, when MHA is trained with 50K steps with 10.6B, if with the same time cost, the SAS with 10.6B is trained with 25K steps, which only.use half dataset. Therefore, forcing such comparison (same model size and time cost) is actually unfair, as too many other hyperparameters are changed, including learning rate, training dataset size and so on. On the other hand, compare MHA 10.6B and SAS 2.7B for same time cost is relatively acceptable , as the dataset size is not.changed.

The training loss on Pile with 50K training steps, 10.6B model size, with the each column result are reported with the same time cost

Q2: whether the larger model is not helpful for that dataset

A2: SAS still significantly improves the performance with the increase of model size

The perplexity on Pile

The model scale-up follows the suggestion of the GPT-3 paper.

Dear Reviewer DaRN,

Our Core Claim is: SAS at equivalent training steps outperforms MHA, with the same parameter cost. All experiments are designed for such a claim.

• This is readily proved in the paper experiment.

However, we never claim that SAS will definitely outperform MHA with the same training time, but SAS only has the potential to achieve comparable performance with MHA with the same time cost, depending on the dataset, model size, and so on.

• 2.7B SAS achieves comparable performance with 6.7B MHA and 10.6B MHA.
• However, this is not always promised, depending on the dataset, model size, and so on.

Therefore, we both agree: 1) SAS at equivalent training steps outperforms MHA, with the same parameter cost; 2) SAS has the potential to achieve comparable performance with MHA with the same time cost (2.7B SAS and 10.6B MHA on Pile dataset), depending on model size, dataset, and so on.

If there are any questions, please let us know

Dear Reviewer DaRN,

We thank Reviewer DaRN for acknowledging our clarifications on the contributions. We are open to further modifications if any.

Dear Reviewer wQTE,

Thank you very much for your comment. We will address your concerns below.

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. A1: 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. How 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.

Q2: We do not know the configuration of the larger models (350M and 2.7B). Relationship between the configuration data and Equations 1-12.

A2: A: We further provide the configuration in the following.

Why do we have such a configuration? Our configuration follows the GPT-3 paper [1] and TPA paper [2].

• From 125M to 2.7B, the layer number, hidden size, learning rate, and so on follow the GPT-3 paper.
• 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.

For the hidden size (use 125M model as an example):

• For MHA:
  • hidden size: 768. This means that the hidden size is 768 for token embedding.
  • Attention, head number. The 125M has 12 attention heads and 64 hidden size per head, so that 12 x 64=768.
  • 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.
• For SAS
  • 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].
  • 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.

Q3: The time cost

A3: We evaluate the time cost via H100.

The time cost with different model sizes and micro batch size 1, with H100

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

Q4: The training lengths of 512 and 1024 are too short. It remains unclear whether SAS maintains its advantages in long-context settings

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

The validation perplexity with 125M model size and 16384 training length on the Pile, with H100/200

Q5: thorough double-checking of all reported metrics is necessary to ensure accuracy.

A5: Thank you very much for your notice. We have carefully checked the accuracy and the highlight.

Reference:

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

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

If the response addresses the concerns, could you consider improving the score?

Dear Reviewer wQTE,

Hope this finds you well.

We have addressed the concerns in the above response, including the discussion of Figure 1, model configuration, time cost, the performance of length 16384, and so on. As the discussion period will be closed soon, if possible, could you consider improving the score?

Dear Reviewer rUv9,

Thank you very much for the support. We will address the concerns below.

Q1: The disucssion of computational cost

A1: We further test the time cost on H100 GPUs. The time cost with different model sizes and micro batch size 1, with H100

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

Q2: The SAS with larger training (e.g. larger model size)

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

The validation perplexity with 6.7B and 10.6B model size and 512 training length on the Pile, with H100/200

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

Q3: Adding another non-linear layer before averaging them

A3: The following is the results. The initial embedding is MHA. And the following is the operation (use 125M as example):SAS-NonLinearBack:

• After SAS attetnion, we will have tensor with shape [36 64], while 36 is the head number and 64 is the hidde size.
• For SAS, we will reshape it to [3 12 64] and then mean pooling to have [1 12 64].
• 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]).

The perplexity result on the Books dataset with length 512 and H100

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

Q4: On L200, 'Linear_2' is probably meant by the second 'Linear_1'.

A4: Thank you very much for the notice. We have fixed it in the revision.

If the response addresses the concerns, could you consider improving the score?

Dear Reviewer rUv9，

Thank you very much for your response. We will answer your question below.

Q1: The performance of MHA and SAS with comparable time cost

A1: We have discussed it in Response to All Reviewers. We directly copy the results in the following.

• With the same model size, the SAS is better than MHA..
• With the same time cost and smaller model size, SAS 2.7B has comparable performance with MHA 6.7B and MHA 10.6B.

If there are any other questions, please let us know. If we have addressed the concerns, could you consider improving the score?

Again, thank you very much for your attention, and wish you a good day.

Dear Reviewer yD2Q,

Thank you very much for your comment. We will address your concerns below.

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

A1: 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

According 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. Therefore, more attention heads and larger hidden sizes per head may further improve the performance.

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

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

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

• Different Length. We validate the SAS in Section 4.1, with lengths from 512 to 2048.
• 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.
• • Different Model Size. We validate the SAS with sizes 125M, 350M and 2.7B in Section 4.1 and Section 4.2.
• Different Head Number. We validate the SAS with different head number in Section 4.3 with different head number.
• Different Kernel Size. We validate the SAS with different kernel sizes, from 1 to 7.
• Different Seed. We validate the SAS with different random seed in Appendix D.
• Different initial attention type. We try different attention methods as the initial query, key and value embedding. We present the results in Appendix G.

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.

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

A3: We test the time cost on H100 GPUs.

The time cost (ms) with different model sizes and micro batch size 1, with H100

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

Q4: Are the baselines state-of-the-art?

A4: The baseline is the state-of-the-art. The baselines contain the most recent works, including MLA [1] and TPA[2].

Q5: How different hyperparameters or other transformer techniques may affect the performance is not clear.

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

The following is the performance with 6.7B and 10.6B.

The validation perplexity with 125M model size and 16384 training length on the Pile, with H100/200

If the response addresses the concerns, could you consider improving the score?

Reference:

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

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

Dear Reviewer yD2Q,

Hope this finds you well.

We have addressed the concerns in the above response, including time cost, the effect of attention head number, different hyperparameters, and so on. As the discussion period will be closed soon, if possible, could you consider improving the score?

Dear Reviewer yD2Q,

We have addressed the concerns in the above response, including time cost, the effect of attention head number, different hyperparameters, and so on. If there are any questions, please let us know before the discussion period ends (August 8 AoE). After the detailed explanation, if possible, could you consider improving the score?

Again, thank you for your attention, and wish you a good day.

Dear Reviewer yD2Q,

Thank you very much for your response. We will answer your question below.

Q1: The Theoretical Analysis of maximum attention head number

A1: The meaning of increasing the head number to improve expressiveness: output relation pattern independently. The answer for the maximum number of heads: $O(D^2)$.

The Softmax Attention with Fully Parameterized Bilinear Attention [1,2].

$O_i = \sum_{c=1}^{H}\left(\sum_{j=1}^i\phi(\frac{x_{i}W_{c}x_{j}}{\sqrt{D}})x_jU_c\right)$

where:
$W_c \in \mathbb{R}^{D \times D}$ = learnable bilinear transformation for head $c$, $\phi$ = the kernel function such as softmax, $U_c \in \mathbb{R}^{D \times D}$ = value projection matrix, $H$ = number of attention heads, $D$ = hidden dimension per head.

Each head computes a weighted sum of values, with weights determined by $x_i W_c x_j^T$. For $m$-dimensional $x_i$ and $n$-dimensional $x_j$, one potential relation pattern is related to $x_i^mx_j^n$. Therefore, at most $D^2$ relations or interactions).

The bilinear term $x_i W_c x_j^T=\sum_{m=1}^D \sum_{n=1}^D x_i^{(m)} W_c^{(m,n)} x_j^{(n)}$ can be expressed as $(x_i \otimes x_j) \cdot \text{vec}(W_c)$, where $\text{vec}(W_c) \in \mathbb{R}^{D^2}$. With $H$ heads, the model can represent up to $D^2$ distinct relations [6]. We hope that the output of each head is a single relation but not mixed relations, so that we can utilize the relations independently (e.g., $\phi$ processes the relation independently).

Example

• Assume $Rank(W_c)=1$, and there is only one non-zero value in $W_c$. And there are $H=D^2=4$.
• Input：

$$x_i = [a b] x_j =[c d]$$

• $W_c$:

$$W_1= [1 0] [0 0] W_2= [0 1] [0 0] W_3= [0 0] [1 0] W_4= [0 0] [0 1]$$

The first head $W_1$ will gives the result of relation $ac$, the $W_2$, $W_3$, $W_4$ 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., $\phi$ processes the relation independently), compared to only one head with $W_c=[ 1~ 1$ \ \ $1 ~1]$ that outputs the mix of $ac+ad+bc+bd$. With a large $D$, we can have more combinations. With a large $H$, we can capture more independent relations. And previous [5, 6, 7] suggest that the expressiveness is related to O(HD), where $H$ is the attention head and $D$ is the hidden size per head.

If we increase $D$ (e.g., hidden size per head), we increase the maximum number of relations. If we increase the $H$, we increase the number of independent relations that are actually captured.

And this is why GQA, MQA, MLA, and TPA have the potential to increase performance, though they are low-rank: they either increase the attention head number (GQA, MQA, MLA, and TPA) or increase the hidden size per head (MLA), with the same parameter cost.

If the response addresses the concerns, could you consider improving the score?

Reference:

[1] Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., ... & Polosukhin, I. (2017). Attention is all you need. Advances in neural information processing systems, 30.

[2] Shazeer, N. (2019). Fast transformer decoding: One write-head is all you need. arXiv preprint arXiv:1911.02150.

[3] Shazeer, N., Lan, Z., Cheng, Y., Ding, N., & Hou, L. (2020). Talking-heads attention. arXiv preprint arXiv:2003.02436.

[4] Ainslie, J., Lee-Thorp, J., De Jong, M., Zemlyanskiy, Y., Lebrón, F., & Sanghai, S. (2023). Gqa: Training generalized multi-query transformer models from multi-head checkpoints. arXiv preprint arXiv:2305.13245.

[5] 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.

[6] Hu, J., Li, H., Zhang, Y., Wang, Z., Zhou, S., Zhang, X., ... & Jiang, D. (2024). Multi-matrix factorization attention. arXiv preprint arXiv:2412.19255.

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