Theory, Analysis, and Best Practices for Sigmoid Self-Attention

• Activation parity for vision tasks:
  • To paint a complete and thorough picture we highlight various activation alternatives in vision. Here we find that many activations perform equivalently. However when tried with language we observe a disparity (also observed inHua, Weizhe, et al, see Table 2).
• “Putting aside your FlashSigmoid hardware aware attention. Could you explain what is the real benefit of using sigmoid over softmax? Is there any real benefit in terms of training/performance? Could you also clearly explain the disadvantages of using sigmoid attention?“
  • Thank you for this thoughtful question about sigmoid versus softmax attention. While the benefits of sigmoid extend beyond our FlashSigmoid implementation, the fundamental advantage lies in sigmoid's local, element-wise nature versus softmax's row-wise operations. This locality enables superior parallelization, particularly valuable for multi-device scenarios where we eliminate one round of node-to-node communication. From a computational perspective, softmax attention requires two reductions along the key/value sequence length (one for softmax normalization, another for V vectors summation), while sigmoid attention needs only one. On modern hardware like the H100, which scales FP32 performance more than bandwidth and warp communication compared to A100, softmax's non-matrix-multiply reduction is less efficient than standard reductions. Importantly, our empirical results (detailed in section 5 and provided as summary tables below) demonstrate that these computational benefits come without compromising model performance, and in some cases even improve it. We believe this work opens valuable directions for exploring alternatives to the conventional softmax monopoly in transformer attention mechanisms.
  • SigmoidAttn is faster than SoftmaxAttn without Flash applied to either variant. Table 3 uses JAX with XLA without any Flash variant for either formulation. Even here we see a 20% speedup.
• Comparison to prior works:
  • We compare SigmoidAttn with [5] in Figure 9 (a) — this is $n^{-\alpha}$ without layerscale or qknorm. We also ablate multiplicative sequence length normalization in Appendix G.1.1.
  • We compare SigmoidAttn with [6] in Figure 9 (c) [this is ReLU$^2$] — parity is only achieved with QK Norm and LayerScale]. The authors of [6] already have results that highlight that going from softmax → ReLU$^2$ results in a performance degradation (see Table 2 in [6] where perplexity increases). In contrast we show parity in our language modeling results at a much larger scale.
• Hybrid-Norm for longer sequence lengths:
  • The core issue to mitigate to ensure a high performant SigmoidAttn is that of attention norms. While the local Lipschitz constant shows the worst case scenario for SigmoidAttn being lower than for SoftmaxAttn, this is not informative of the average-case behavior, and we do observe instabilities in practical training. Other reasons beyond the Lipschitz constant that can cause issues:
    • The Lipschitz constant is calculated on a per-layer basis and does not account for interaction in deep transformer networks with other components like normalization layers and MLPs.
    • As the sequence length increases there can be a large attention norm which is not properly tempered by the b = -ln (max seq len) approximation. A stronger normalization with hybrid-norm helps mitigate the instabilities at larger sequence lengths. Note that hybrid-norm is applied on the embedding dimension of models (typically much smaller than sequence length for modern models).
  • To help simplify the exposition we have attached the relevant runs comparing Sigmoid and Softmax Attention below in markdown. Here you can see that SigmoidAttn outperforms SoftmaxAttn on 5/9 tasks at 1B n=4096.
  • The number of parameters is neglible compared to the rest of the network: for a regular pre-norm transformer block using an embed dim d this increases total parameters by 1/(6d + 3) — so for d=2048 this is 0.0081% extra parameters. We also found that in some situations normalization without learnable affine terms works well (making the extra parameters 0). There is a slight tradeoff to throughput, however at sequence length 1024 we still observe a 8.36% causal attention throughput improvement.

Thank you for your thoughtful review of our paper on sigmoid attention. We appreciate your recognition of our theoretical and empirical contributions. We'd like to address your concerns:

• Novelty & Performance:
  • We build upon the proof of Yun et al. and extend contextual mapping to the sigmoid activation (this is much more involved than softmax which tends towards hardmax — elaborated more in the point below).
  • The gains reported by FlashSigmoid are not only due to it being a hardware-aware implementation, but rather by the fact that Sigmoid itself is a local operation (applied element-wise), as opposed to the non-locality of Softmax (applied row-wise) with direct implications on the parallelisability of the algorithm. Not all FLOPs are created equal.
  • We empirically analyze practical training of SigmoidAttn and highlight stability issues and ways to mitigate them.
  • We improve inference throughput of attention by 17% in non IO bound scenarios. Note that 17% inference speedup can result in 17,000 dollars per day in typical language model serving [4]. This can directly reduce energy requirements which has a real world impact. Even without FlashSigmoid, SigmoidAttn is faster than SoftmaxAttn (see JAX results for LM training in Table 3 where no Flash variant is used for either -- at sequence length 4096 we observe $\approx$ 20% speedup).
  • Empirically SigmoidAttn outperforms SoftmaxAttn on:
    • Language Modeling: 6/9 tasks on 1B n=2048 [truncated summary table below]
    • Language Modeling: 5/9 tasks at 1B n=4096 [truncated summary table below]
    • ASR: Sigmoid + RoPE (“Sigmoid” in Table 2 under “RoPE” row — note that this is SigmoidAttn + QKNorm + b = -ln(n) ) outperforms SoftmaxAtttn on all sequence generalization tasks evaluated on TEDLIUM-v3. Lower WER is better here. We provide a transparent view of various changes here to highlight what works and what does not.
    • And matches performance in all vision tasks (supervised and self-supervised learning).
• “The proof of UAP is very nice but is essentially a slight change of the proof given by Yun et al., which the authors do say themselves”
  • We do agree with the reviewer that the proof for UAP in SigmoidAttn largely follows the blueprint of that for SoftmaxAttn, which is why we did ensure to provide due credit to [Yun et al.] in the manuscript. At the same time, we believe it unfair to downplay the technical difficulty associated with our adaptation. Arguably, even for the original proof the main challenge consisted in identifying the proper way of combining selective shift layers to assemble a contextual mapping: SigmoidAttn modifies the action of a selective shift layer non-trivially (compare (7) with (8)), and forced us to re-think the proper way to combine them, as outlined in C2. On top of this, and regardless of the actual complexity of the proof, we believe this contribution positively impacts the quality of the paper, as it allows to highlight the theoretical equivalence of the expressive power of SigmoidAttn to SoftmaxAttn
• “In section 3.2 they show that the Lipshitz constant of sigmoid attention grows quadratically w.r.t a ball of radius R and that softmax grows exponentially.” + Q.2 :
  • We want to clarify this statement: we show that for any input sequence in a ball of radius R, the Lipschitz constant of sigmoid attention is at most $O(R^2)$. On the other hand, [Castin et al. 23] show that there exists an input sequence in a ball of radius R such that the Lipschitz constant of softmax attention grows at least like $R^2 \exp(R^2)$. This result is therefore to be understood as an adversarial case, where a special input sequence is used. It says nothing about the average case behavior that we observe in practice, which is likely much better behaved (otherwise, at the reviewer points out, it would be practically impossible to train transformers). These Lipchitz constants are also of independent interest to study the dynamics of attention [1, 2]. We have clarified this in the revised pdf.
• “This implies that sigmoid attention has much better Jacobian regularity and softmax has very bad Jacobian regularity.”
  • As stated above, this statement should rather become “This implies that sigmoid attention has much better Jacobian regularity in the worst case than softmax in the worst case”. In practice, we are likely far from these worst cases.
• “Also in figure 2 I can't really see any difference between softmax and sigmoid.”
  • The goal of this figure is to higlight that on various domains and tasks SigmoidAttn matches SoftmaxAttn, while providing throughput improvements (Figure 1). Faster training for equivalent performance can lead to substantial energy savings as we keep scaling our modern models. We have also improved the clarity of this figure in the manuscript to better highlight this.

Thank you for your detailed review and recognition of our paper's depth and potential impact. We appreciate your positive comments on the presentation and analysis. We'd like to address your concerns and questions:

• “the local Lipschitz constant of SigmoidAttn is much lower than the worst local Lipschitz constant of SoftmaxAttn. This suggests that sigmoid attention should be easier to train than softmax.”
  • We thank the reviewer for this insightful remark. First, we want to clarify what these theoretical results mean. Theorem 3.2 shows that the local lipschitz constant of sigmoid attention is $O(R^2)$ for any input sequence that has vectors of norm <= R. On the other hand, the results of [Castin et al. 23] show that there exists a sequence of vectors of norm <= R such that the local lipschitz constant of softmax attention grows like $R^2 \exp(R^2)$. This is a worst case result, and it is not informative of the average case, which might be closer to what happens in practice. Note that we never intend to imply that these results lead to more stable training for sigmoid attention in practice. However, these results regarding the regularity of sigmoid attention are of great interest to theoretically study the dynamics of attention (see e.g. [1, 2]). We have made these points clearer in the revised pdf.
• “This suggests that sigmoid attention should be easier to train than softmax. However, this is inconsistent with the experiments presented by the authors in Fig. 2, Fig. 3, and Fig. 4.”
  • Fig 2 demonstrates that SigmoidAttn matches SoftmaxAttn on various domains and tasks. We have improved the figure as you suggested, which should enhance the paper presentation.
  • Fig 3 and Fig 4 are to highlight that a naive application of SigmoidAttn results in instabilities, particularly when naively using SinCos or RoPE embeddings. Note that Figure 5 and Figure 6 highlight that this is completely resolved using (a) ALiBi in Figure 5, (b) RoPE with b=-10 in Figure 6. With Figure 6 (RoPE) the 1:1 comparison is between Sigmoid RoPE and Softmax RoPE. With the bias in place the norm of SigmoidAttn matches that of SoftmaxAttn (RoPE). We also see that Sigmoid slightly outperforms RoPE in terms of the train loss.
• “Table 2 shows that sigmoid attention causes unstable training, whereas softmax does not. I believe a theoretical analysis of training stability is necessary.”
  • Table 2 provides a thorough ablation of all the components of SigmoidAttn, some of which are indeed unstable. The goal of this table is to be as transparent as possible and ablate what does and does not work with SigmoidAttn.
  • The suggested approach [ QK Norm with b =-ln(n) ] is called just “Sigmoid” in the table [the default]. Sigmoid + RoPE (“Sigmoid” in the RoPE section) performs better sequence level generalization than SoftmaxAttn on the ASR task that uses TEDLIUM-v3 that evaluates on ever growing sequence lengths [see summary markdown below]. Lower word error rate (WER %) is better.
• "It would be helpful if the authors could make Fig. 2 clearer."
  • Thanks for the feedback! We have improved the plot in the updated manuscript with a wider line for SigmoidAttn and set the z-level of SoftmaxAttn to be in front as well as apply smoothing to the language modeling train loss. It is now easy to see that SigmoidAttn matches the train loss of SoftmaxAttn on a wide variety of domains and tasks (the main goal of the figure).
• "It would be better if the authors conducted experiments on large-scale language models (around 7B parameters)."
  • We appreciate the reviewer's suggestion to consider larger models. Our study's primary objective is a methodical comparison between sigmoid and softmax attention, necessitating precise control over all variables. Given that training a single 1B LLM requires about 11,300 A100 40G GPU hours and we train 5 (Tab 3) + 16 (Appendix Table 11) we believe this already presents a fair evaluation of SigmoidAttn vs. SoftmaxAttn [237,300 GPU hours total].

Lower WER% is better in table below:

[1]: Geshkovski, Borjan, et al. "The emergence of clusters in self-attention dynamics." Advances in Neural Information Processing Systems 36 (2024).

[2]: Geshkovski, Borjan, et al. "A mathematical perspective on transformers." arXiv preprint arXiv:2312.10794 (2023).

• “it looks like the proposed remedy ( bias b=-10 only partially helps to control the increase in the amplitude of |act * V} (graph 6)“
  • The correct comparison in Fig 6 is Sigmoid + RoPE b=-10 and Softmax + RoPE. In this case the norms are approximately overlapping. Softmax + ALiBi does have a lower norm, but that model has objectively different learning dynamics.
  • Sigmoid + RoPE + b=-10 does outperform (in terms of train loss) the Softmax + RoPE baseline in Fig 6. This is also the (practically) infinite data regime, so train loss \approx test loss.
• “ASR: sigmoid (RoPE with bias) performs badly on long audio ( Table 2)“
  • “Sigmoid” in the table refers to “SigmoidAttn uses LayerScale, QK norm, b=−log n, and no sequence normalization.” (lines 473-474). This model actually outperforms the softmax baseline (see summary ASR table below). We will further clarify this detail in the updated manuscript.
• “LLM: Sigmoid (ALibi) underperform vs Softmax (Alibi) . No results for Sigmoid (RoPE)”
  • Sigmoid (ALiBi) outperforms Softmax (ALiBi) on 5/9 tasks at sequence length 2048 (see summary table below).
  • There is a gap at sequence length 4096, but this is resolved via hybrid-norm. We also highlight this below for reference where Sigmoid (ALiBi) + H-Norm outperforms the Softmax (ALiBi) baseline on 5/8 tasks.
  • We have many more results (incl Sigmoid (RoPE)) in Appendix G.4 as well as in the ASR results.
• Relation to Yang et al., 2018; Ganea et al., 2019:
  • We highlight these two works to demonstrate the limited expressivity of Softmax [independent of attention]. The two works above demonstrate this relationship through rank bottlenecking in matrix factorization as the reviewer highlights. There are newer works such as [5] and [4 — empirical] that also highlight softmax as being problematic w.r.t. expressivity — we will include these references in the paper.

Higher % values are between in the two LM tables below:

Lower word error rate percent (WER%) is better in the table below:

[1] Patel, D., & Ahmad, A. (2023, February 9). The inference cost of search disruption – large language model cost analysis. semianalysis.com. https://www.semianalysis.com/p/the-inference-cost-of-search-disruption

[2] Xiao, Guangxuan, et al. "Efficient streaming language models with attention sinks." arXiv preprint arXiv:2309.17453 (2023).

[3] Sun, Mingjie, et al. "Massive activations in large language models." arXiv preprint arXiv:2402.17762 (2024).

[4] Team, Chameleon. "Chameleon: Mixed-modal early-fusion foundation models." arXiv preprint arXiv:2405.09818 (2024).

[5] Veličković, Petar, et al. "softmax is not enough (for sharp out-of-distribution)." arXiv preprint arXiv:2410.01104 (2024).

We thank you for the positive review and recognition of our paper’s strengths. We appreciate your insightful comments about our work and would like to address your questions and concerns:

• "Why should we switch from original softmax attention to sigmoid attention?"
  • As the scale of modern ML models grows so does the energy consumption. For example [1] suggests that the cost of inference is in the range of hundreds of thousands of dollars per day. Saving 17% of this can significantly reduce cost and energy, which becomes a more critical concern as model size increases. There are also associated inference speedups.
  • We also want to highlight that SigmoidAttn does outperform SoftmaxAttn on:
    • Language: 6/9 tasks on 1B n=2048 [see summary markdown below]
    • Language: 5/9 tasks at 1B n=4096 [see summary markdown below]
    • Sigmoid + RoPE performs better sequence level generalization using RoPE on the TEDLIUM-v3 ASR task evaluates the model on ever-growing sequence lengths [see summary markdown below]. This is “Sigmoid” in Table 2 (the default setting is QK Norm + b = -ln(n) that we propose). Here lower WER% is better.
  • Concurrent work has suggested that Softmax attention produces suboptimal representations [2, 3] and faces issues when working with multi-modal learning scenarios [4]. We hope to explore SigmoidAttn in these contexts in the future.
• “will it help with long context? ”
  • We do have preliminary evidence that it does well in this regard. See the summary ASR table below w/ TEDLIUMv3 that evalutes on growing sequence lengths (lower WER% is better).
• Computational complexity of reductions.
  • Mathematically there are 2 reductions in SoftmaxAttn along the seq length (once for the softmax, once for the sum of V vectors) - while SigmoidAttn only has one along the sum of V vectors. In the case of efficient kernel implementations, with SoftmaxAttn, the first reduction for the softmax happens over tensor-core formats while not being a tensor core matrix-multiply-accumulate operation (the second one with V can use the mma) - that leaves it to be fairly inefficient, especially when compared to a regular reduction. This gap is exaggerated on machines like the H100 which scale their fp32 performance more than their bandwidth and warp communication times as compared to the A100.
• UFA proof:
  • We agree with the reviewer that the UAP proof of sigmoid attention stems from an adaptation of the one for softmax attention, and indeed we highlight this in the manuscript. However we argue that, also in the original case, proving the architecture can represent a contextual mapping indeed represents the main challenge (the remainder in fact reuses established results from UAP of MLPs and approximation of continuous function via piecewise constant ones), and as such is not trivial. Moreover, the softmax operation lends itself well as a building block for building a contextual mapping, thanks to its ability to extract individual values (as it tends to hardmax); sigmoid provides additional challenges in this sense as it can only extract ranges of values (as it tends to heaviside). This makes the proof rather more convoluted, as testified by the lengthiness of C2.
• “I think the proof of Theorem 2 is missing the key assumption for reqularity is that ||Wq, Wk"|| and ||Wv|| should be limited."
  • The theorem does provide an explicit relationship between the Lipschitzness bound of sigmoid attention and the norms of its parameters ||Wk||, ||Wq||, ||Wv||. To clarify, we will further highlight that the bound on the Lipschitzness constant degrades as the norm of the weight matrices increases.

Lower WER% is better.

Higher % values in the LM tables below are better.