Rethinking Softmax: Self-Attention with Polynomial Activations

Despite the interesting angle of their work I do have a number of criticisms.

1. The rigour with which the theories have been developed:

   1. Theorem 4.1 specifically Eq 4.2 which is proved in appendix 4.1.1, how does L719 work, I do not believe the factor of 2 be present, e.g. if $F_{i 1} = 1$ and all other values are 0 the equation gives leads to $1 - 0 = 2(1 + 0 + \ldots + 0) = 2$ which is obviously wrong.
   2. In Theorem 4.2, the authors have assumed X (hidden states), Q and K projections are i.i.d and therefore $A = XQ$ is has 0 covariance with $B^{T} = KX$, I would argue this is an extremely strong assumption except in the trivial cases and so is not applicable to Transformers. I therefore question the utility of the theory developed on this basis. In the general case $\text{Cov}(A, B) = Q K^\top$ when X is drawn from $\mathcal{N}(0, 1)$.
   3. In Theorem 4.2, the authors comes to final terms which contain both $N$ (the sequence length) and $D$ (the feature dimension) e.g. Eq A.41 and then claim the quantity grows with $\mathcal{O}\left(\sqrt{N}\right)$. This will only be the case when $N \gg D$. Similar arguments can be made for all of Theorem 4.2, Corollary 4.3, Theorem 4.4 and Corollary 4.5. Noting that for ViT-Base $D = 768$ and $N = 14^2 = 196$ (https://huggingface.co/google/vit-base-patch16-224/blob/main/config.json). Therefore the much larger condition cannot be satisfied. As a side note A.1.2 introduces a parameter $d$ (not $D$) which is not defined, perhaps this is the dimension of a single attention head but it is not clear.

2. The linkage between the derived theory and the experiments.

   1. L201-207 say that corollaries 4.3 and 4.5 show that scaled polynomial activation functions can perform similarly to softmax. It is not made clear why this is the case unless simply because the relationship between the activation function use and norm both grow with $\sqrt{N}$.
   2. Figure 1 and L240-245. If you argue that polynomial activation functions should be scaled by $\frac{1}{\sqrt{N}}$, but then say as the sequence length decreases the scale applied should decrease. Should this not be the opposite i.e. the scale applied should increase? $N = 256, 64, 16, 8 \to \frac{1}{\sqrt{N}} = 0.125, 0.0625, 0.25, 0.353. Perhaps I have misunderstood the argument, but regardless it is not clear.

3. Significance of the attention map visualisations?

   1. Figure 6 and Figure 7 show pairs of attention map visualisations for specific heads for softmax and polynomial activations. The arguments made here about "contrast" and "difference" seem obvious given the function is so different. E.g. softmax maps cannot be negative. Moreover, there is no pre-ordained behaviour of a specific attention head and so I do not understand what point the authors are trying to make other than the activation functions are different.

4. Natural language processing sequence length. Transformers for NLP make use of causal masking and so the effective sequence length for each query is different (incremented by 1) and this has not been considered in the experiment, to match the theory shouldn't the activation scaling evolve as the effective sequence length increases?

• Does not mention what to do with dynamic sequence length tasks (e.g. language generation). Do you dynamically scale during both training and inference dependent on current sequence length (or do you take e.g. maximum during training)?

• Seems to mainly focus on ViT, lacking LM experiments (even at small scale). As seen, ViT is a lot more flexible and even gets decent results without doing the attention scaling. The NLP experiments are very small, and lack details. Experiments are minimal.

• Lacking details for the NLP experiments. No mention of model size, hyperparameters overview/recipe and information about the tasks themselves. Does refer to a work (https://ojs.aaai.org/index.php/AAAI/article/view/17664), but does not mention what they use from that work.

• Does not cite https://arxiv.org/abs/2409.04431, which this works seems to have a lot of overlap with. The focus of this work is different, but especially section 3.2 (Regularity of Sigmoid Attention) has a lot of overlap.

• Typos

  • “accuracr” (A.2.1) -> accuracy

1. Motivation: Softmax stabilizes the gradient but also causes gradient vanishing (the gradient is less than 1 and saturates to 0, see [1, 2]). Therefore, i disagree with the argument in the paper contribution stating that: "we theoretically show that softmax has a regularization effect on attention and argue this plays a crucial role in its success."

In addition, in their theory, there is only the analysis of softmax and polynomial activation (used in attention) but I have not found the theoretical analysis that shows that this regularized effect correlates with the performance of transformer models.

In the experimental results, I also fail to see the same correlation. For example, in Figures 2 and 3, the attention norm and Jacobian norm of models using (x^3/1) are constantly lower than that of models using (x/16). However, in Table 1, the model with (x^3/1) has significantly worse performance than the model with (1/16).

The paper also states: "We advise the reader that this paper diverges from the usual pursuit of creating cutting-edge transformer architectures for achieving state-of-the-art results on benchmark datasets. Instead, our focus is on critically examining softmax attention to determine whether its effectiveness is a result of true interpretability or a more nuanced regularization mechanism". As far as I understand, the paper aims to examine the reason why behind the use of softmax activation in attention, less about aiming to improve it. However, I think the paper fails to argue why regularization is the source of softmax attention success and their proposed polynomial activation also fails to significantly improve over softmax.

2. Not convincing experiment results Again, as I pointed out above, the experiments do not convince me how the norm-regularized effect correlates with the performance of transformer models. In addition, in Table 2, ViT-base and ViT-small (x/14) have lower performance than (x^3) and that seems to correlate with norm analysis in Figures 13 and 14 (where Vit-models x/14 have higher attention and Jacobian norm than those of ViT-models x^3). However, also in Table 2, Deit-base and -small (x/14) outperform the rest. However, no norm analysis of this model is shown in the experiment.

Most of the improvements in performance made by polynomial activation are only marginal

LRA benchmark is not an NLP benchmark, 3 out of 5 tasks are not NLP, including Listops, Image Classification and pathfinder. The authors should not make this mistake, and they should understand the benchmark that they are testing on.

3. Small mistakes in the proof: from the line 665-669, the author makes a replicate typo in $e^{2a_{11}}$ and $e^{2a_{NN}}$.

Reference [1]. Oliver Paul Richter and Roger Wattenhofer. Normalized attention without probability cage, 2022. [2]. Escaping the gradient vanishing: Periodic alternatives of softmax in attention mechanism

1. The proposed polynomial activation method does not seem to achieve the promising results observed in vision tasks. Do the authors have any insights into why this might be the case?
2. The paper only adopts $x$ and $x^3$ as the polynomial activation functions. Is there a particular reason why other values were not considered?
3. The paper lacks a discussion on the computational overhead. Given that the proposed method introduces a new form of self-attention, it would be beneficial to include an analysis of time complexity or provide a comparison of actual running times between the softmax-based and polynomial-based methods.
4. On line 255, $\frac{1}{8}x^3$ should be $\frac{1}{16}x^3$. The same mistake in the caption of Table 1.