Bridging the Divide: Reconsidering Softmax and Linear Attention

We would first like to express our appreciation for your time and insightful comments. Please find our response to your concerns in the following:

1. Novelty and contribution.

Thanks for your valuable comment. The novelty and contribution of our work can be summarized as follows:

• We identify an important property of attention mechanism, injectivity, which has not been explored in literature. With both theoretical proof and experimental verification, we show that injectivity significantly contributes to the performance gap between Softmax and linear attention.
• We thoroughly validate the crucial role of local modeling in attention mechanism through a series of experiments.
• We propose a novel subtraction normalization method to achieve injectivity in linear attention, significantly improving performance with no extra FLOPs.
• We present InLine attention module, a simple, effective and efficient alternative to the commonly adopted Softmax attention. Extensive experiments fully validate the superiority of our design over Softmax attention.

2. The locality problem and related work.

Thanks for highlighting this important related work [1]. While our work shares some similar analyses with [1], it has fundamental distinctions:

• We thoroughly validate the crucial role of local modeling in attention mechanism, while [1] mainly focuses on identifying differences. There could be many different behaviors between Softmax and linear attention, but a large number of them may not be the key disparities. Therefore, our study conducts a series of experiments in Fig. 3, 4 and Table. 2, 4 to fully validate that: (1) Softmax attention's effectiveness depends significantly on robust local modeling. (2) Varied local priors contribute substantially to the performance gap between Softmax and linear attention. In contrast, [1] observes that Softmax attention is more concentrated locally than linear attention (see its Fig. 2), but it does not verify whether this is a key factor leading to the performance gap between the two.

• The design proposed in our paper is more effective than [1]. To enhance locality, [1] proposes DiagAttention, which restricts each token's receptive field to its local block. This approach resembles the window attention in Table 4 (left), which shows little improvement (window=$14^2$, acc=80.4) compared to vanilla global linear attention (acc=80.2). In contrast, our design boosts accuracy significantly from 80.2 to 82.4, fully unleashing the power of global linear attention. Additionally, we offer a comparison between our InLine model and TransNormer proposed in [1] under Swin-T structure. Our InLine model achieves better results.

  Method	#Params	FLOPs	Acc.
  TransNormer	30M	4.5G	79.9
  InLine	30M	4.5G	82.4

• Thanks for pointing this out again. We will include discussion and give more credits to this work in the revised manuscript.

[1] The Devil in Linear Transformer.

3. Ablation on different kernel functions.

Thanks for the insightful comment. We conduct additional ablation studies on different kernel functions $\phi$ using InLine-Swin-T.

It is shown that our InLine attention can effectively work with different kernel functions, further validating the effectiveness of our method. The ReLU and Exponential functions achieve similar results, slightly outperforming the Identity function. In our paper, we use Identity kernel function as default. Additionally, we will give more credits to [2] in the revised version.

[2] Linear Log-Normal Attention with Unbiased Concentration

4. Experiments on language models.

Please refer to our general response.

5. Negative attention scores.

Thanks for the insightful question. We already provide a brief discussion on the non-negativity assurance in L260-L263 of our paper. Here, we offer detailed analyses and additional experiments to further clarify this interesting issue.

• InLine attention does not ensure non-negative attention scores. As you pointed out, the subtraction normalization could possibly produce negative values, and we practically find that there do exist many negative attention scores in our InLine models.

• These negative values do not impair the performance of InLine attention. We directly verify this with an ablation experiment, where we ensure all attention values to be non-negative by introducing additional normalization on $Q,K$. The results are provided below.

  Method	#Params	FLOPs	Acc.
  InLine-Swin-T	30M	4.5G	82.4
  InLine-Swin-T non-negative	30M	4.5G	81.7

  It can be seen that the additional non-negative assurance actually leads to notable performance drop. This suggests that negative values in InLine attention do not hinder performance and might even enhance the model.

• Non-negativity is crucial for vanilla linear attention. As depicted in Table. 3 of our paper, unlike our InLine attention, vanilla linear attention fails to converge without non-negativity assurance. We attribute this to extreme non-injectivity and semantic confusion problem. For example, with Identity kernel function, vanilla linear attention is unable to distinguish completely opposite semantics, assigning identical attention scores to $q$ and $−q$.

We would first like to express our appreciation for your time and insightful comments. Please find our response to your concerns in the following:

1. The source code.

• Firstly, we provide the Pytorch-style pseudo code of our InLine attention module below. It can be seen that the proposed method is very simple and easy to implement.

  # b: batch size;  n: sequence length;  d: head dimension
  
  # q: query,		shape: (b, num_heads, n, d)
  # k: key,		shape: (b, num_heads, n, d)
  # v: value,		shape: (b, num_heads, n, d)
  
  # Eq. 5 of our paper
  o = q @ (k.transpose(2, 3) @ v) - (q @ k.transpose(2, 3).sum(dim=2) - 1) * v.mean(dim=2)
  # Eq. 6 of our paper
  o = o + local_residual_aggre(v)
  
  # o: output,	shape: (b, num_heads, n, d)
  

• Secondly, we will provide the full source code to promote reproducibility if this manuscript is accepted.

2. Comparison with vision Mamba models.

Thanks for your valuable comment. Here, we provide additional comparison with vision Mamba models.

• We offer comparison in terms of params, FLOPs and accuracy in the table below.

  Model	#Params	FLOPs	Acc.
  Vim-S	26M	5.1G	80.3
  LocalVim-S	28M	4.8G	81.2
  PlainMamba-L2	25M	8.1G	81.6
  EfficientVMamba-B	33M	4.0G	81.8
  VMamba-T	31M	4.9G	82.5
  LocalVMamba-T	26M	5.7G	82.7
  InLine-CSwin-T	21M	4.3G	83.2
  Mamba2D-B	94M	-	83.0
  VMamba-B	89M	15.4G	83.9
  InLine-CSwin-B	73M	14.9G	84.5

  As depicted in the above table, our InLine model surpasses various vision Mamba designs without bells and whistles.

• Furthermore, as our InLine attention is extremely simple, it achieves much faster inference speed than vision Mamba models. Accuracy-runtime comparison is provided in the PDF in general response. The results demonstrate that our InLine model is 4.6x, 6.3x faster than LocalVMamba and PlainMamba models respectively, while achieving comparable performance. Compared to the highly optimized VMamba model, our model also achieves 1.4x speedup and 0.2 accuracy gain.

3. Is this algorithm adaptive to NLP or time series domain?

Our method can be applied to NLP and time series domain. Please refer to our general response for the detailed discussion.

Dear Reviewer 3cj2, thank you for your insightful review and for engaging with our work. We would like to know if there are any additional questions or concerns. We are eager to engage in further discussion and provide clarification to fully address them.

We would first like to express our appreciation for your time and insightful comments. Please find our response to your concerns in the following:

1. The performance of InLine models and comparison with related works.

Thanks for the valuable comment.

Firstly, under fair comparison, the improvement of our method is significant and consistent.

• We propose injective linear attention and local attention residual to address two core limitations of linear attention. The effectiveness of these two designs is validated in Table 3 and Table 4 of our paper. Table 3 shows that injective linear attention achieves significant accuracy gains of 2.5 and 9.8 compared to vanilla linear attention. Table 4 shows that local attention residual further improves model performance from 80.2 to 82.4. We kindly argue that these improvements are significant and remarkable.
• Our InLine attention module significantly improves performance across four representative architectures: DeiT, PVT, Swin and CSwin. As shown in Table 5 and Table 6 of our paper, simply applying our InLine attention module to these four models leads to obvious accuracy gains. For example, InLine-Swin-S outperforms Swin-B with 57% of the Params and 56% of the FLOPs. These results demonstrate the superiority of InLine attention as an effective alternative to the widely used Softmax attention.

Secondly, we offer additional comparison with more competitive related works like VVT.

• The primary focus of our paper is to demystify the limitations of linear attention, rather than achieving SOTA results. Therefore, we mainly focus on building our InLine models under the four simple baseline models. In contrast, VVT utilizes advance macro designs like ConvFFN to achieve more competitive results. For a fair comparison with VVT, we employ similar advance designs and provide the results below.

  Model	#Params	FLOPs	Acc.
  VVT-L	61.8M	10.8G	84.1
  InLine	51.1M	6.8G	84.2

  It is shown that our InLine model outperforms VVT with 20% fewer Params and 40% fewer FLOPs.

• Additionally, since our InLine model is extremely simple and effective, it delivers much better speed-accuracy trade-off than other competitive works. For example, as depicted in Fig. 1 of the PDF in general response, our simple InLine-Swin model shows 1.8x-2.5x faster inference speed than VVT, while achieving comparable accuracy. Speed is tested on a single RTX3090 GPU.

• We will give more credits to these related works and include comparison in the revised manuscript.

2. The analysis of local modeling capability.

Thanks for your insightful question. We offer clarification on the analysis of local modeling capability.

• As shown in Eq. 6 of our paper, the local attention residual proposed in our paper is $\sum_{j=1}^9 r_j V_j^{N(i)}$, where $N(i)$ is the 3×3 neighborhood of $V_i$ and $V_j^{N(i)}$ represents the value token in this neighborhood, $j=1,\cdots,9$. Therefore, the attention local attention residual is not a simply residual $rV_i$ since it contains local property. The local property of the attention residual is beneficial for InLine attention, rather than the residual property. We offer additional experiments to further validate this.

• Firstly, we replace the local residual with vanilla residual $rV_i$ and provide the results below. The window size is fixed as global, i.e. $56^2$.

  Method	Window	#Params	FLOPs	Acc.
  InLine w/o residual	$56^2$	30M	4.5G	80.2
  InLine + vanilla residual	$56^2$	30M	4.5G	80.2
  InLine + local residual	$56^2$	30M	4.5G	82.4

  It can be seen that adding vanilla residual can not benefit InLine model and it achieves the same result as no residual. This indicates that the residual property is not the core of the proposed attention residual.

• Secondly, we provide further ablation studies on $N(i)$. In our original design, $N(i)$ is the 3×3 neighborhood of $V_i$. Here, we define $D_kN(i)$ as the dilated 3×3 neighborhood of $V_i$ with dilation $k$, where the concept "dilation" is the same as in dilated convolution. Thus we have $D_1N(i)=N(i)$. As $k$ enlarges, the local property of $\sum_{j=1}^9 r_j V_j^{D_kN(i)}$ is weakened.

  Model	$k$	#Params	FLOPs	Acc.
  InLine-Swin-T	1	30M	4.5G	82.4
  InLine-Swin-T	2	30M	4.5G	81.7
  InLine-Swin-T	3	30M	4.5G	80.8
  InLine-Swin-T	4	30M	4.5G	80.6

  As depicted in the above table, the model performance decreases significantly as $k$ enlarges, proving that local property is the real factor benefiting InLine model.

• These results will be included in the revised version to make the analysis of local modeling capability more solid. Thanks again for your valuable question.

Dear Reviewer My8Q, thank you for your insightful review and for engaging with our work. We would like to know if there are any additional questions or concerns. We are eager to engage in further discussion and provide clarification to fully address them.

We would first like to express our appreciation for your time and insightful comments. Please find our response to your concerns in the following:

1. The locality problem and related work.

Thanks for pointing out this important related work [1]. Here, we offer clarification on the relationship between this study and our work.

Firstly, our work verifies the vital importance of the injective property of attention mechanism. This unique and valuable perspective helps us develop the extremely simple yet effective injective linear attention.

Secondly, while our work shares some similar analyses on the locality problem with [1], it has fundamental distinctions:

• We thoroughly validate the crucial role of local modeling in attention mechanism, while [1] mainly focuses on identifying differences. There could be many different behaviors between Softmax and linear attention, but a large number of them may not be the key disparities. Therefore, our study conducts a series of experiments in Fig. 3, 4 and Table. 2, 4 to fully validate that: (1) Softmax attention's effectiveness depends significantly on robust local modeling. (2) Varied local priors contribute substantially to the performance gap between Softmax and linear attention. In contrast, [1] observes that Softmax attention is more concentrated locally than linear attention (see its Fig. 2), but it does not verify whether this is a key factor leading to the performance gap between the two.

• The design proposed in our paper is more effective than [1]. To enhance locality, [1] proposes DiagAttention, which restricts each token's receptive field to its local block. This approach resembles the window attention in Table 4 (left), which shows little improvement (window=$14^2$, acc=80.4) compared to vanilla global linear attention (acc=80.2). In contrast, our design boosts accuracy significantly from 80.2 to 82.4, fully unleashing the power of global linear attention. Additionally, we offer a comparison between our InLine model and TransNormer proposed in [1] under Swin-T structure. Our InLine model achieves better results.

  Method	#Params	FLOPs	Acc.
  TransNormer	30M	4.5G	79.9
  InLine	30M	4.5G	82.4

• Thanks for pointing this out again. We will include discussion and give more credits to this work in the revised manuscript.

[1] Qin et al., 2022; The Devil in Linear Transformer.

2. The similarities and differences compared to FLatten.

Thanks for your valuable comment.

Firstly, removing the division operation can make FLatten injective but hinders performance. Without the division operation, FLatten can not ensure the attention weights sum up to 1, which is important for stabilizing the model. Therefore, when the division operation is removed from FLatten-Swin-T, the model experiences an accuracy drop of 0.2.

Secondly, we offer a detailed analysis of the similarities and differences with FLatten.

• Our work and FLatten are largely orthogonal. Our work analyzes the injectivity and locality of Softmax and linear attention, while FLatten discusses the focus ability and feature diversity of linear attention. Furthermore, we present subtraction normalization to make linear attention injective,

  $$\rm{InL_K}(Q_i)=\left[\phi(Q_i)^\top \phi(K_1), \cdots, \phi(Q_i)^\top \phi(K_N)\right]^\top-\frac{1}{N} \sum_{s=1}^N \phi(Q_i)^\top \phi(K_s)+\frac{1}{N},$$

  while FLatten introduces a specific mapping function $\phi=\frac{||x||}{||x^{**p}||}x^{**p}$ to sharpen attention distribution. These methods are orthogonal and can be employed together.

• The findings of our paper are more fundamental, and could be more beneficial for the community in designing effective linear attention patterns. As discussed in our paper (L145-L152), the focus ability identified in FLatten can be viewed as a special case of non-injectivity and confusion problem. And our studies on locality can also explain why dwconv improves FLatten.

• Our InLine attention is more effective and efficient than FLatten. The proposed InLine attention is extremely simple yet effective, yielding both faster speed and higher accuracy compared to FLatten. For example, a single InLine attention block is 1.8x faster than a FLatten block, and InLine-PVT-T achieves 0.4 accuracy gain with 1.2x speed up than FLatten-PVT-T. These results are tested on a single RTX3090 GPU.

3. Language modeling experiments.

Please refer to our general response.

Dear Reviewer TKLm, thank you for your insightful review and for engaging with our work. We would like to know if there are any additional questions or concerns. We are eager to engage in further discussion and provide clarification to fully address them.