Theory, Analysis, and Best Practices for Sigmoid Self-Attention
Authors: Jason Ramapuram
Federico Danieli
Eeshan Gunesh Dhekane
Floris Weers
Dan Busbridge
Pierre Ablin
Tatiana Likhomanenko
Jagrit Digani
Zijin Gu
Amitis Shidani
Russell Webb
Venue: ICLR 2025 Poster (2025)
Keywords: Sigmoid Attention
Pointwise Attention
Universal Function Approximators
Primary area: other topics in machine learning (i.e., none of the above)
TL;DR: Revisiting sigmoid attention in transformers, proving universality, identifying training factors, introducing FlashSigmoid, and demonstrating competitive performance across domains, establishing it as a softmax alternative.

Abstract: Attention is a key part of the transformer architecture. It is a sequence-to-sequence mapping that transforms each sequence element into a weighted sum of values. The weights are typically obtained as the softmax of dot products between keys and queries. Recent work has explored alternatives to softmax attention in transformers, such as ReLU and sigmoid activations. In this work, we revisit sigmoid attention and conduct an in-depth theoretical and empirical analysis. Theoretically, we prove that transformers with sigmoid attention are universal function approximators and benefit from improved regularity compared to softmax attention. Through detailed empirical analysis, we identify stabilization of large initial attention norms during the early stages of training as a crucial factor for the successful training of models with sigmoid attention, outperforming prior attempts. We also introduce FLASHSIGMOID, a hardware-aware and memory-efficient implementation of sigmoid attention yielding a 17% inference kernel speed-up over FLASHATTENTION2 on H100 GPUs. Experiments across language, vision, and speech show that properly normalized sigmoid attention matches the strong performance of softmax attention on a wide range of domains and scales, which previous attempts at sigmoid attention were unable to fully achieve. Our work unifies prior art and establishes best practices for sigmoid attention as a drop-in softmax replacement in transformers.

Published as a conference paper at ICLR 2025
THEORY, ANALYSIS, AND BEST PRACTICES FOR SIGMOID SELF-ATTENTE
Jason Ramapuram; Federico Danieli, Eeshan Dhekane, Floris Weers, Dan Busbridge, Pierre Ablin," Tatiana Likhomanenko"" Jagrit Digani, Zijin Gu, Amitis Shidani, Russ Webb
Apple
ABSTRACT
Attention is key part of the transformer architecture. is sequence-fo-sequence mapping that transforms each sequence element into a weighted sum of values. The weights are typically obtained the softmax of dot products between keys and queries. Recent work has explored alternatives to softmax attention transformers, such as ReLU and sigmoid activations. In this work, revisit sigmoid attention and conduct an in-depth theoretical and empirical analysis. Theoretically, we prove that transformers with sigmoid attention are universal function approximators and benefit from improved regularity compared to softmax attention. Through detailed empirical analysis, we identify stabilization of large initial attention norms during the early stages of training as a crucial factor for the successful training of models with sigmoid attention, outperforming prior attempts. We also introduce FLASH- SIGMOID, hardware-aware and memory-efficient implementation of sigmoid attention yielding 17% inference kernel speed-up over FLASHATTENTION on H100 GPUs! Experiments across language, vision, and speech show that properly normalized sigmoid attention matches the strong performance of softmax attention wide range of domains and scales, which previous attempts at sigmoid atten- tion were unable to fully achieve. Our work unifies prior art and establishes best practices for sigmoid attention drop-in softmax replacement in transformers.
INTRODUCTION
The success of modern machine learning can be largely attributed to the attention mechanism (Bahdanau et al., 2015; Vaswani et al., 2017). Attention uses sequence-to-sequeeee (seq-to-seq) map to build context-award aware token representations. Classically, attention relies on the softmax function (SoftmaxAttn) to recover token representations data-dependent convex combinations of values.
Despite its widespread use and effcctiveness, softmax in SoftmaxAttn is not without limitations. For instance, the softmax function can sometimes lead concentration of attention just few features (Yang et al., 2018; Ganea et al., 2019), potentially neglecting other informative aspects of the input data. Moreover, applying Softmax requires performing row-wise reduction along the length of the input sequence, which in the case of efficient attention kernels (Dao et al., 2022: Dao, 2023), slows down computations. In this work, we relax this constraint by substituting the row-wise softmax operation with an element-wise sigmoid nonlinearity. We highlight that the central problem with naive sigmoid attention (SigmoidAttu) is that of large initial attention norms and propose solutions to alleviate it. Our contributions are as follows:
(1) We prove SigmoidAttn a universal function approximator on seq-to-seq tasks (Sec. 3.1).
(2) We analyze SigmoidAttn's regularity and provide its worst-case Jacobian bound (Sec. 3.2).
33) We extend FLASHATTENTION2 (Dao et al., 2022; Dao, 2023) with the sigmoid kernel, reducing kernel inference wall-clock time by up 17% and real world inference by up to 8% (Sec. 4).
(4) We show that SigmoidAttn matches SoftmaxAtta in various tasks and domains (Sec. 5).
Primary contributor. For detailed breakdown author contributions see Appendix H. Correspondence to: Jason Ramapuram <jramapuram @ apple.com> Code is available https: / wwwb.cco/aapppppcccccccc wtt. e /m1 imoi d- atteetttion

---

Published as a conference paper at ICLR 2025
2 SIGMOID ATTENTION
Let E Rund be the input sequence of vectors, where each vector has dimension d. We define three learnable weight matrices W, E RdxdØt Wa E Rdxdd and W, E Rdxé, which are used to compute the queries Q E Ruxdde keys K E RBxxd and values V E RRax as follows:
Q XW
XWk, and V = =Ww.
Self-attention (Bahdanau et al., 2015; Vaswani et al., 2017) can be compactly written as SoftmaxAttn(XX X) Softmax(QK" /vx)))
(2)
where the Softmax function normalizes each row of the input matrix. We replace the Softmax with
Sig) [QKK" (QQKK/VdxxV) with o: :u + sigmoid( b) = (1 + (
(3)
Here, is applied element-wise the input matrix in (3). The activation function has hyper- parameter E R. In App. E, we discuss an intuitive way to choose the order-optimal bias term, resulting in b = - (n). This choice of allows make sense of SigmoidAttn for any sequence length. Indeed, letting (y1,... ,....,y)))))) SigmoidAttn( X) be the output sequence, we have
(W D W W y exp( (W Wx,,)) +
/ exp((W Wa, Wax) Woxdu(x),
W Pj
(4)
where /4 n دド= day is the empirical measure corresponding to X. Notably, (4) still makes sense in the infinite length limit, where the measure μ is not sum of Diracs. Wortsman et al. (2023a) do bias, and propose n1 normalization for various attention activations, such sigmoid and ReLU, but leave the reason as an open question. Our variable bias has a similar effect the large n limit, and posit that recovering a finite output limit as n increases the why works in practice.
A multi-head version of (3) is obtained by combining the outputs of several Sigmoid as follows:
[Sigmoid Attn, ( Sigmoid (X)) W.:
(5)
for a learnable output weight matrix W. E Igdd, xd where h denotes the number of heads.
3 THEORETICAL PROPERTIES OF SIGMOID ATTENTION
We analyze SigmoidAtun Attn, with two objectives: (1) showing that transformer architecture remains universal function approximator when SigmoidAttn replaces SoftmaxAttn, and (2) recovering measure of regularity of SigmoidAttn by computing its Lipschitz constant.
3.1 ARE TRANSFORMERS WITH SIGMOID ATTENTION UNIVERSAL APPROXIMATORS?
Yun et (2020) demonstrate that classical transformers can approximate continuous sequence- to-sequence functions to arbitrary precision, property known as the Universal Approximation Property (UAP). UAP highly desirable as it provides proof of an architecture's generalizability and representation capability. As Sigmoid Attn modifies the transformer architecture, it is crucial to theoretically guarantee that this modification does not impact the representation capability and that UAP is retained. We provide this guarantee with the following theorem.
Attn
Theorem 3.1 (UAP for SigmoidAttn). We denote with Th.durr the class of transformer networks obtainable by combining an arbitrary number of SigmoidAttn layers (each of h heads of dimension followed by FFN layers of hidden dimension For any given continuous, germutationeeuuvarre function f C n Raad with compact support Q and for any arbitrarily small error e, there exists transformer network g E T..1. such that
( [l((X) g x()) S.,
for 1 1 < x.
(6)
2

---

Published as a conference paper at ICLR 2025
Theorem 3.1 is the exact counterpart of (Yun et al., 2020, Thm. 2), which shows UAP for classical transformers. Our proof largely follows the same path, an outline of the original proof provided in App. C. Here, present overview of the main adaptations required to prove Thm. 3.1 for Sigmoid Attn, with further details in App. C.I and C.2.
Sigmoid Attention layers can implement contextual mappings: A key step in proving Thm. 3.1 is showing that, even with Attn, a sequence of transformer blocks can implement a Contextual Mapping (Yun et al., 2020, Def. 3.1). A contextual mapping characterizes function that maps each input sequence element to an output uniquely dependent on the whole sequence. This property allows transformer to capture and store global context within each token, even if each layer only performs pairwise comparisons. Subsequent layers can then use this global information to map individual tokens to the correct output, ultimately approximating any arbitrary sequence-to-sen to-quuence function.
In Yun et al. (2020), the contextual mapping is assembled by modifying individual transformer blocks: each block is tuned react specific input token. By stacking a sequence of these blocks, a transformer can be turned into accumulator, mapping given input token sequence to a unique global index. This outcome is achieved via a selective shift layer (Yun et al., 2020, App. B.5):
max Xx,1 - ming Xx,1 if b< x <
V (x;b,
otherwise,
and can be approximated using classic attention. Although SigmoidAttn cannot directly approxi- mate (7), our accumulator definition relies on an equivalent selective shift operation:
(E X if b < X < D'
V (X;;b,00))
otherwise,
which can be approximated by Sigmoid Attn (described in App. C.I). In App. C.2.4, we show that (8)
shares similar properties with (7), allowing us to use the original proof framework Yun al. (2020)
and demonstrate that UAP holds in our case as well.
Our proof is largely equivalent to that in Yun et al. (2020), with two relevant differences: to approximate (8), we require SigmoidAttn with at least four heads and shifts included in both query and key definitions. In contrast, SoftmaxAttn requires at least two heads to approximate (7), with shifts only in the query definition. However, this is primarily theoretical requirement for the proof and does not affect performance. Notably, the total number of parameters required by both architectures for the approximation follows the same tight scaling of Yun al. (2020).
3.2 REGULARITY OF SIGMOID ATTENTION
As with any layer in a neural network, the regularity of Sigmoid Attn is important to study, as it gives insights into the robustness of the corresponding network and the ease of optimizing it. The most standard way to quantify the regularity of a layer function & is to compute its Lipschitz constant over a set X, that a constant C' such that for all X, E X', it holds (Øx(X) ((Y)) S C Y||. where || is the standard Frobenius norm. The local Lipschitz constant is the spectral norm of the Jacobian of á X. The two are related: the Lipschitz constant of é over is the greatest local Lipschitz constant for all X E X. We turn to the theorem giving the regularity of SigmoidAtins
Theorem 3.2. Define A {(W WE i.j E (l....,,) n.} C R the set of attention weights, and the scaled activation norms ao n X suPue, (a(u) and d'a the Jacobian of Sttn at X = (I1..... Inn) has spectral norm of at most:
supae A (u)). Then,
Wolle Joo + 20 W W& l1
a 113
The proof is found in App. D. In SigmoidAttn, if we assume that the attention weights (Waz; Waæz)) are all bounded by a constant μ this is true, e.g., if the activations are bounded we get de I exp(μ) and d'c I схр(μμ thanks to the choice of b = - log(n). The bound in Thm. 3.2 depends only on the average squared-norm of the input sequence I, while classical results for the study of attention all rely on the largest value of ||0;11 (Kim et al., 2021: Castin et al., 2023). This is another consequence of the simplicity of sigmoid attention and is due to the removal of the
3

---

Published as a conference paper at ICLR 2025
normalizing constant in SoftmaxAttn. Our result implies that if all a; are within ball of radius R then the Lipschitz constant of Sttn grows at most like R², but it is stronger since we can apply this to unbounded distributions zü it matters only that the second moment bounded. This result contrasts sharply with the bounds obtained for Softmax Castin et al. (2023, Thm. 3.4.) show that there exists sequence X (#11--... Fnn) with | R for all such that the spectral norm of the Jacobian of Attn at X is at least cft exp(cIt) for some constant c >> O. On the other hand, bound scales in Rº:: this means that the local Lipschitz constant of SigmoidAttn is much lower than the worst local Lipschitz constant of SoftmaxAttn. Note that this result does not inform us of the practical average case Lipschitz constant, which likely be much lower for both Softmax and Sigmoid attention. Upper bounds on the Lipschitz constant of SigmoidAtln are of particular interest to study the dynamies of attention, as done, e.g., in (Geshkovski et al., 2024; 2023)
4 FLASHSIGMOID: HARDWARE-AWAR AWARE IMPLEMENTATION
FlashAttention2 (Full)
FlachAttention2 (Full)
HlshSSggmoi (Full)
FlashSigmoid (Full)
FlastAtttetion2 (Causal)
FlashMttention22 (Causal)
Kernel GPU Time (ms)
FlashSigmoid (Causal)
FlashSigmaid (Causall
bi
Kernel GPU Time
anco
30000
2000
1000
1880
1000
21000
1083
7000
Tokens
Tokens
(a) Inference mode kemelss on H100.
(b) Training mode kernels on H100.
Figure/ 1: Average kernel speed-up for FLASHSIGMOIN FLASHATTENTION2 for sequence lengths 64-7k. Inference is 17.39% faster for self-attention and 18.76% for causal attention. Training is 6.53% faster for self-attention and 9.46% for causal attention.
Memory speed has not kept pace with recent gains in computation speed (Choquette, 2023; Jouppi et al., 2017; Hanmun et al., 2023). Consequently, attention computations on modern architectures have been 10-bound by memory accesses (Ivanov et al., 2021). FLASHATTENTION (Dao et al., 2022) and PLASHATTENTION2 (Dao, 2023) address these shortcomings by optimizing GPU memory hierarchy utilization to accelerate attention computations. Motivated by the speed boost provided by these approaches, we develop FLASHSIGMOID, a hardware-aware implementation of SigmoidAttn. Like previous works, FLASHSIGMOID employs three core ideas:
Tiling: Divide and Conquer Approach to Attention: Similar to FLASUATTENTION and
P FLASHISGGMOI processes input parts in parallel compute attention outputs in blocks, efficiently combining partial results to generate the final attention output.
Kernel Fusion: Like FLASHATTENTION and FLASHATTENTION2. FLASHSIGMOID implements
the computational steps of both forward and backward passes of SigmoidAttn as single GPU kernels, minimizing memory accesses and improving memory efficiency by avoiding materialization of intermediate activations on High- Bandwidth Memory (HBM).
Activation Recomputation: The backward pass of sigmoid attention requires the sigmoid acti- vation matrix, which, if materialized on GPU HBM, results in slower implementation and memory inefficiencies. FLASHSIGMOID addresses this by retaining only query, key, and value tensors for re-computation of the sigmoid activation matrix during the backward pass. Despite increased FLOPs, this approach proves faster in wall-clock time as well as more memory- -fffciett than the alterantive approach of materializing and retaining the attention matrix.
The forward and backward pass algorithms of FLASHSIGMOID can be found in App. F.I. Here, we highlight key differences between FLASHSIGMOID and The point-wise wise nature of SigmoidAttu results in a faster and more memory-efficient implementation by removing the need to compute the sofumax normalization and materialize HBM. A reduction
4

---

Published as a conference paper at ICLR 2025
in the number of kernel dispatches also speeds up FLASHSIGMOID. Further, FLASHSIGMOID does not require accumulation and tracking of intermediate variables (row-sum and maximum of blocks) in the forward and backward passes which saves computation cost and reduces register pressure. We use sigmoid (x= = 0.5 - (1 + tanh (0.5 x)) to optimize the sigmoid computation on GPU. The speed up in FIGSSID compared to ATTENTION arises from optimizing hardware bottlenecks; theoretically, SignoidAttu is slower than SoftmaxAttn (App. G.6).
To measure the performance improvements of FLASHIIGGOID, we compare the timings of the kernels in its forward and backward passes against those of FLASIIATTENTION. The details of this benchmarking on H100 and A100 GPUs can be found in App. F.2. Measuring GPU computation time, we observe 17.39% speed-up during inference and 6.53% speed-up during training for attention over randomly initialized data on H100 GPU (Fig. 1). In practice, these gains may affected by other bottlenecks, such movement of tensors between CPU or GPU memory, computations in other layers, and communication overhead in distributed training and inference. However, we demonstrate that FLASHSIGMOID speeds up training by 44% and inference by ~8% in realistic end-to-end setup. The details wall-clock time improvements with FLASHSIGMOID are in App. F3. We also note that practical machine learning workflows are dominated by inference rather than training.
5
EXPERIMENTS
BYOL ViT-B/16
SimCLR ViT-B/16
MAE ViT-L/16
0.7
8.8
0.9
Train Loss
0.5
6.8
0.8
0.3
4.8
0.6
0.1
2.8
0.5
-0.0
0.8
0.4
0
200
400
600
0
100
200
300
0
200
400
Epoch
Epoch
Epoch
Supervised ViT-B/16
255M ASR Transformer
1B & 7B LLM
7.1
60.0
2.8
Sigmoid
Train Loss
5.9
50.8
2.5
Softmax
4.7
41.5
2.2
IB
3.5
32.2
1.9
7B H-Norm
2.3
1.6
100
200
300
100000
200000
300000
100000
200000
300000
Epoch
Steps
Steps
Figure 2: Train losses comparing SigmoidAtta with SoftmaxAtin.
To empirically validate Attn, evaluate across several domains: supervised image classi- fication using vision transformers (Dosovitskiy et al., 2021), self-supervised image representation learning with SimCLR (Chen et al., 2020; Zhai al., 2023a), Bootstrap Your Own Latent (BYOL) (Grill et al., 2020; Busbridge et al., 2023) and Masked AutoEncoders (MAE) (He et al., 2022) as well as automatic speech recognition (ASR) (Synnaeve et al., 2020; Gulati et al., 2020b) and reto-segressvve language modeling (LM) (Brown et al., 2020). We also validate sequence length generalization on TED-LIUM v3 (Hernandez al., 2018) for ASR and in small scale synthetie exper- iments in App. G.5.4. Across all these domains and algorithms. we demonstrate that SigmoidAttn matches the performance of SoftmaxAttn (Fig. 2 and 22), while offering training and inference speed-ups highlighted in Sec. 4. Empirically we make the following observations:
(1) SigmoidAttn is effective for vision tasks without a bias (except MAE), but relies on Layer- Scale (Touvron al., 2021) to match the performance of the baseline Softmax Attn (Fig. 10-a)
5

---

Published as a conference paper at ICLR 2025
in hyper-parameter free manner?? All results presented for SoftmaxAttn also fairly add LayerScale unless specified.
(2) LM and ASR sensitive the initial norm /VIlQk//ddyy VII. Modulation is required via (a) relative positional embeddings like ALiBi (Press et al., 2022), which reduces the initial attention norm by shifting logit mass
under Attn, (b) appropriate initialization of b to achieve the same effect enabling usage of any positional embedding, (c) using hybrid- norm (Appendices G.3.3, G.4 and G.7) at the expense of an extra normalization layer.
5.1 ABLATIONS
sctrras ROD
sottraa rope ai
Train NLL
109
artt
10.4
10
10-1
10
0
10000
20300
30000
45000
50003
60000
10000
20000
30093
40000
50035
50000
Steps
Steps
Figure SigmoidAttn with SinCos.
Figure 4: SigmoidAtba with RoPE.
aetimac oohmer al
softeer Sgra
Train NILL
dg rpp
acr ||
000
30"
30a
107
10'
10-
10
10--
10t
10
10
30000
20320
30003
40200
50002
52000
30000
20200
30032
45000
50025
50000
Steps
Stops
Figure 5: SigmoidAttn with ALiBi.
Figure SigmoidAtta with RoPE, b = -00
We begin with ablations to dissect the benefits of each of our introduced components. To gain intuition about SigmoidAttn, we developed research-friendly auto-regressive (AR) LM training framework to measure all components of attention and validate the effects of LayerScale, LayerNorm applied to and K (QK norm), different positional embedding techniques, and initialization values for b.
Mitigating Large Attention Norms We train a single layer AR transformer block (E=3072, D. FF==12288) on the realnews split of C4 (Raffel al., 2020). We train for 216 steps using a batch size of 6 and max sequence length of 4096 using single cycle cosine learning rate (LR) schedule without weight decay. SigmoidAttn initially underperformed SoftmaxAttn when using absolute sinusoidal (SinCos) (Fig. 3) or relative (Fig. 4) positional embeddings (PE), which attribute
A Appendix G.2.2 demonstrates that supervised vision tasks using SigmoidAttr Attn without LayerScale can match haseline Softmax Attn performance by relying learnable scalar bias and temperature: (b,t) B.

---

Published as a conference paper at ICLR 2025
to high initial attention Frobenius norms, ||o |[<()k A corresponding evolution of the attention distribution and sparsity can be seen in Appendix Fig. 30 and Fig. 31 on a symthetic task. To address these larger attention norms, we propose: (a) using ALiBi (Press et al., 2022) whose relative bias moves initial attention logit mass to the zero region under the sigmoid activation, producing equivalent train negative log-likelihoods (Fig. 5) (b) set the attention logit bias b to negative offset proportional to the sequence length, b x - In n (see App. G.1.2 for an ablation on 6). This enables the usage of other PE techniques like RoPE (Su et al., 2024) (Fig. 6).
No HyhridNrm
With HybridNarm
[A (XWg) Jank
Emptricall
Sigroid Bourd 2th. 1.3
Sigrwon Tipptiidd b hio)
107
10
10
32
125
Sequenze Length
Sequence Length
Figure 7: Regularity analysis comparing SigmoidAttn vs. Softmax (10x trials per nt). SoftmaxAttn theoretical bound is off scale and thus omitted.
Empirical Analysis of Attention Regularity To validate our theoretical analysis (Section 3.2), we measure Jacobian norms of SigmoidAttn and Softmax across sequence lengths (Figure 7). Us- ing autograd, we compute exact Jacobian norms for both mechanisms, with and without HybridNorm (Appendices G.3.3 and G.7), comparing them to theoretical bounds maxfmm bound omitted as it exceeds scale). Both variants show empirical norms (solid lines) well below their theoretical bounds (dashed lines). With our proposed bias initialization (b - In(n)). SigmoidAttn achieves lower norms than SoftmaxAttn in both settings, suggesting improved regularity. This aligns with its strong task performance (Section 5). Additionally, HybridNorm (Figure 7, right) reduces norms for both mechanisms compared to baseline (left), highlighting normalization's 's role in attention stability longer sequences.
LayerScale To validate the need for LayerScale, we follow Wortsman et al. (2023b) to quantify the impact stability. All models are trained with RoPE with b x In,, using AdamW (Loshchilov & Hutter, 2017) on the realnews split of C4 with (81, B2) = (0.9, 0.95). e = 10 8 wd = O, batch size 24, maximum token sequence length of 512 from the T5 tokenizer (Raffel et al., 2020), cosine LR schedule of 214 steps including a linear warmup of 2nD steps. Models have Raeads = K. Plaayers = 2 x K, dmodel = 64 x and dieed-forwand = 256 x K for a scaling value λ E {1,,, 2, 4, 8, 16) leading to models with 2,.2, 4.9, 15.0, 67.0, 440.0) M trainable non-embedding parameters. Following Wortsman et al. (2023b), we sweep learning rates 7) E {3 x 10-1, x 10 10-3, 3 x 10 x 10- x 10-2, 1 x 10-1,3 x 10-1). LR sensitivity is defined as By (ni(()(A(r)) fo) {*] where ((A(T)) is the loss achieved by the learning algorithm A with LR 17. lo is the loss at initialization, and E is the loss achieved by the best LR. LayerScale is initialized at 10-4. Unlike vision tasks, where LayerScale improves performance (Fig. 10-a), in LM, we observe that Softmax slightly benefits from LayerScale, while the performance of iigmoidAtth remains largely unaffected.
Stability with QK Norm To explore the stability of SoftmaxAttn vs. SigmoidAttn we repeat the analysis of Wortsman et al. (2023b), as described in the LayerScale analysis, to investigate the impact of QK norm (Dehghani et al., 2023). For language modeling, both SigmoidAttn and maltmaxxAii exhibit sensitivity to learning rate changes without QK norm. However, incorporating QK norm significantly stabilizes performance (Fig. 9). In vision tasks, SigmoidAttn demonstrates robustness with and without QK norm (Fig. 10-a) and without the need for normalization from Wortsman et al. (2023a)²
Multi-query attention (MQA) In Fig. 10-b we explore MQA (Shazeer, 2019) for vision using only one head for {K, V}. We find that both SigmoidAttn and Softmax perform equally well with or without multiple heads even at the small scale of ViT-B/16.
A Activation Function Ablations As in Wortsman et al. (2023a), various activation functions, when combined with LayerScale and QK norm, perform equally well for vision tasks (Fig. 10-c). However, for sequence-critical tasks like ASR, activation functions such as ReLU pose instabilities
We ablate multiplicative sequence length scaling in more detail in App- G.I.I.
7

---

Published as a conference paper at ICLR 2025
1S42. layet scelowFllo
140, Tayer TORe
FICEE Time
Lpmal 6Pan
ichad
4050
6050
30
APC EX02
1P Leuningg 7802
Leatting 2012
Learning
Non-cmbed Panaro
Noz-crbed Purura (M)
Non-crited Pazaro 34)
New-mmbed Parara (44)
Figure 8: LR sensitivity LayerScale ablation.
Figure 9: LR sensitivity QK norm ablation.
Tost TopTTo-18
Settraa, el alayereeat eKKeem
Sigrord d, el speetude #OKNomn
Laaeriick
Tiali, (Lousfeic (QKSan
MBL Softras
RalAl -taNetickk-Okherr
Mil Sipreit
GeU.U. slggesede sOCN
MGA Softnas
MCA Signedl
Rel1P, Jayadaaks OWNen
65
100
150
20
300
150
150
200
250
30.3
100
150
250
30
Epoch
Epach (b)
Epoch (c)
(2)
Figure 10: ImageNetlk VIT-B/16 classification. (a) SigmoidAttn is robust without QK norm (+Lay- erScale, -QKNorm). Removing I LayerScale reduces accuracy by 1.0% (-LayerScale, +/-QKNorm). n normalization (Wortsman et al., 2023a) underperforms without LayerScale. (b) SigmoidAttn multi-query attention (MQA) (Shazeer, 2019) with one head matches multi-head attention (MHA). (c) Sigmoid with LayerScale and QK norm performs comparably to other activations, except TanH. ReLU (Hua et al., 2022) underperforms without LayerScale and QK norm.
and underperform. In the same figure, we also compare to the ReLU proposal from Hua et al. (2022) and find that it underperforms without LayerScale and QK norm.
5.2 SUPERVISED IMAGE CLASSIFICATION
Vision transformers (Dosovitskiy et al., 2021) extend transformers (Vaswani et al., 2017) to treat K K image grids disparate tokens. All tokens are refined through sequential layers of self- attention, pooled using a CLS token or global average pooling layer, and optimized using the negative log likelihood, In p We train ViT-B/16 models using 2222xx24x3 images for 300 epochs using the recipe provided in App. G.2.4. We use the same set of training hyper-parameters for both Softmax Attn and SigmoidAAtn, changing only the activation function between trials. The train negative log ligeliiolooo is reported in Fig. 2 and the test top- 1% is reported in Fig. 22. We find that SigmoidAtta matches both the training dynamics and the evaluation performance of SoftmaxAttn.
5.3 SULE-SUPERVISED IMAGE REPRESENTATION LEARNING
Self Self-supervised representation learning (SSL) exploits vast quantities of unlabeled data to learn semantic representations based on inductive biases such as augmentation invariance (SimCLR Chen et al. (2020), BYO (Grill et al., 2020) or reconstruction from compressed representations (MAE (He et al., 2022)). We employ vision transformer training recipes from Zhai et al. (2023a) and Busbridge al. (2023) (App. G.2.4) for SimCLR and BYOL. As with supervised learning, we use the same set of training function between trials. Figure reports the train losses, and Fig. 22 highlights the linear probe
for both SoftmaxAttn and SigmoidAttn, changing only the activation
8

---

Published as a conference paper at ICLR 2025
Table 1: Word eπrоr rate (%) on LibriSpeech test sets and TED-LIUM (Hernandez et al., 2018) ("TED", joint validation and test sets split according to duration) for transformer (255M params) with either SottmaxAttn or SigmoidAttn (LayerScale and QK norm are used with b log n) trained on LibriSpeech 960h data (mean duration is 10-15s) Hlyper-parameter are in App. G.4.
ATTN
PE
TEST-CLEAN
TEST-OTHER
TED 0-10s
TED 10-20s
TED 20-30s
TED 30s+
SOFTMAX
2.3
5.7
12.4
10.5
11.9
SIGMOID
2.4
5.5
12.4
10.3
12.3
9.7
LAVERSCALE
CAPE
2.5
6.1
13.6
11.5
13.4
8.9
SIGMOID (b-
10, LEARNABLE)
2.3
5.5
12.1
10.5
13.0
9.3
SIGMOID
LEARNABLE)
QK NORM
UNSTABLE, GRADENNT NORM AND LOSS SPIKES
SOFTMAX
2.2
5.5
12.7
10.6
12,8
SIGMOID
2.3
5.4
12.3
10.1
12.3
8.6
SIGMOID (b-
LEARNABLE)
ROPE
2.2
5.2
12.4
10.5
12.3
21.8
+
2.7
6.6
14.1
12.0
14.5
14.9
SIGMOID (b = -5 IN Q. LEARNABLE)
UNSTABLE, GRADIENT NORM AND LOSS SPIKES
3OFTMAX
2.2
5.4
12.3
10.7
12.1
8.6
SIGMOID
2.3
5.1
12.3
10.5
12.6
SIGMOID 10, LEARNABLE)
+
2.6
6.6
13.9
11.9
14.2
SIGMOID -5 LEARNABLE)
2.2
5.2
12.1
10.4
12.0
and finetuned test top-1%. Despite the diverse training objectives in SSL, SigmoidAttn matches SoftmaxAttn while improving training and inference throughput (Sec. 4).
5.4 AUTOMATIC SPEECH RECOGNITION (ASR)
We benchmark ASR using LibriSpeech data (Panayotov et al., 2015) on 100h and 960h settings of paired speech and text transcriptions. Our PyTTrcch implementations of encoder-based vanilla transformer (Synnaeve et al., 2020) and conformer (Gulati et al., 2020a) are trained with Connectionist Temporal Classification (CTC) (Graves et al., 2006) w/ BF16 mixed precision, w/o QK norm and w/o LayerScale. After extensively tuning SoftmaxAttn baselines, we switch to Sigmoid.Att per (3) without any other changes. We investigate the effects of post/pre-LayerNone model depth, optimizer type, small data regime, and connection to local attention, with details in App. G.4.
Our main findings are: i) CAPE (Likhomanenko et al., 2021) PE is the most unstable for AegmoidAtta, ii) post-LayerNorm models with Softmax Atin are hard to match with stable Sigmoid iii) w/o QK norm SigmoidAttn is unstable and significant spikes happen in both gradient norms and training loss; iv) LayerScale is needed for generalization; v) learnable bias
- 10 gives no loss and gradient norms spikes while matching the SoftmaxAttn (which does not benefit from the improved throughput of FLASHISIGMOID); vi) adding learnable bias, b -5, to Q instead of the attention logits also solves the initial large attention norms for CAPE and ALiBi but not for RoPE; vii) b log gives rare (2-5 times) marginal gradient norms spikes with smooth loss while matching SoftmaxAuIn.
Table 1 shows the main result for mae-LayerNomm transformers with CAPE, RoPE, and ALiBi, where Sigmoid Attn uses LayerScale, QK norm, ablated with learnable bias (one per layer) in attention or Q with or without sequence normalization. SigmoidA Attn is stabilized with bias while matching SoftmaxAttn, and = log n works well. In most cases, bias allows generalization longer sequences without sequence normalization, except for RoPE where it helps for longer sequences but hurts overall performance.
log n. and no sequence normalization. The bias is
5.5 AUTOREGRESSIVE LARGE LANGUAGE MODELING
We train all models using the Llama2 recipe (Touvron et al., 2023) (with ALiBi instead RoPE) and the RedPajama (Computer, 2023) dataset in JAX without FLASHATTENTION using the AXLeam framework* (App. G.3 for detailed pyper--arameters). Initial experiments at 85M parameters es- tablished basic stability requirements (App. G.3.1), with attention bias b - log(n) (n 4096) providing effective results. At IB n= 2048 scale with b = - log(n), SigmoidALur matches the train NLL and evaluation results of SoftmaxAttn (Tab. 2 top row) while improving throughput by 1.12x.

---

Published as a conference paper at ICLR 2025
Table LLM English evaluation. All models use ALiBi. Detailed ablations in Appendix G.3.3.
SEQ.
ARC
ARC CHAI..
HELLA-
WING- GRANDE
LAMBADA
TRIVIAQA (I-SHOT)
WEBQS (I-SHOT)
STGP
MODEL
SIZE
PIQA
SCIQ
AVG
(S)
SOFTMAX
1B
2K
62.2
26.8
42.4
59.0
72.3
88.1
58.4
19.9
15.4
49.4
0.38
SIGMOID
1B
2K
62.8
28.8
42.5
59.7
70.3
88.6
59.7
19.1
13.8
49.5
0.34
62.6
27.7
58.6
88.2
58.6
49.2
SIGMOI
1B
4K
60.5
27.3
41.3
57.8
70.5
87.0
57.6
18.9
12.6
48.2
0.67
SOFT (H-NORM)
1B
4K
61.7
26.8
43.4
59.4
70.6
88.6
60.8
20.5
12.9
49.4
SIGM. (HI NORM)
1B
4K
63.5
28.1
43.5
60.7
70.8
88.9
$9.0
20.9
16.0
50.2
SOFT (H-NORM)
7B
4x
71.2
39.9
53.2
65.5
75.6
91.8
67.2
37.7
21.8
59.0
3.85
(II NORM)
72.7
40.5
53.5
66.2
92.5
66.5
59.6
At the 1B n = 4096 scale, using just b - logg()) we observe a 1.25 x speedup: however, slight instabilities prevent SigmoidAttu from matching the strong performance of SoftmaxAttn (Tab. second row). We address these issues through hybrid-norm, an extra normalization layer applied on the output of the attention operation, x + norm(or (amm(r(QQ" (QK"/ /dgg)) V), more details in App. G.3 and G.7). With hybrid-norm, SigmoidAtta matches the train NLL and slightly outperforms SoftmaxAttn English evaluation results (50.2% vs. 49.4% - Tab. third row). In App. G.3.3 we ablate various design choices IB scale, including norm structures, position embedding techniques, and attention bias configurations.
At 7B 4096, SigmoidAtta with hybrid-norm demonstrates compelling advantages compared to SoftmaxAtta with hybrid-norm' matches the train NLL of SoftmaxAttn (Figure 2), while delivering 1.13x speedup (Tab. 2 bottom row). The model shows marginal improvements on challenging tasks, including both reasoning (ARC-Challenge: 40.5% vs 39.9%) and knowledge retrieval (TriviaQA: 39.5% vs 37.7%), with better average performance across all benchmarks (59.6% vs 59.0%). These results establish SigmoidAttn with hybrid-norm large-scale language modeling, offering both computational and performance benefits.
efficient alternative for
6 RELATED WORK
Recent studies in supervised image classification (Wightman et al., 2021) and self-supervised learning (SSL), including approaches like SigLIP (Zhai al., 2023b), shifting large-scale machine learning training from output conditional categorical distributions, traditionally parameterized by softmax functions, to richer pointwise Bernoulii conditionals parameterized by sigmoid functions. In this study, our focus shifts to refining the model's internal mechanics, specifically by substituting the softmax component of the attention mechanism with a pointwise sigmoid function.
Previous work has explored the replacing softmax with the ReLU activation in both practical (Shen al., 2023; Hron al., 2020) and theoretical settings (Bai et al., 2023; Fu al., 2023). Other works explores using the ReLU activation (Hua et al., 2022), exploring purely linear attention (Katharopoulos al., 2020; Lu et al., 2021; Koohpayegani & Pirsiavash, 2024) or cosine-similarity based attention (Luo et al., 2018: Liu et al., 2022). Our work builds upon these explorations, particularly Wortsman al. (2023a), which replaces softmax with various activation functions scaled by n n where corresponds the sequence length and hyper-parameter However, we find that their formulation does not match expected performance without proper b initialization and the of LayerScale (Fig. 10-a, App. G.I.1).
CONCLUSION
In this work, present comprehensive theoretical and empirical study of sigmoid attention an alternative to softmax attention in transformers. We prove that transformers with sigmoid attention are universal function approximators with improved regularity, and identify LayerScale and prevention of large initial attention norms as key factors for successful training. We introduce FLASHSIGMOID. memory-efficit variant providing 17% inference kernel speed-up. Extensive experiments across language, vision, and speech demonstrate that properly normalized sigmoid attention matches softmax attention performance on various tasks and scales. Our findings establish sigmoid attention viable alternative, unifying prior work and establishing best practices for its application in transformers.
'One-to-one trained 7B Sigmoid and 7B Softmax weights are available.
10

---

Published as a conference paper at ICLR 2025
8 ACKNOWLEDGEMENTS
We thank Zakaria Aldeneh, Samy Bengio, Navdeep Jaitly, David Koski, Pau Rodriguez Lopez, Hadi Pouransari, and Skyler Seto for their helpful feedback and critical discussions throughout the process of writing this paper; Okan Akalin, Hassan Babaie, Michael Brooks, Brian Gamp, Denise Hui, Mubarak Seyed Ibrahim, Li Li, Rajat Phull, Evan Samanas, Guillaume Seguin, and the wider Apple infrastructure team for assistance with developing and running scalable, fault tolerant code. Names are in alphabetical order by last name within group.
REFERENCES
Dzmitry Bahdanau, Kyunghyun Cho, and Yoshua Bengio. Neural machine translation by jointly learning to align and translate. In Yoshua Bengio and Yann LeCun (eds.), 3rd International Conference on Learning Representations, ICLR 2015, San Diego CA. USA, May 7-9, 2015, Conference Track Proceedings, 2015. URL http::p:// http:://aaxiv.v.ro wwwp/agg/////ooooo/////////eaaae 1tt//1199..0..
Yu Bai, Fan Chen, Huan Wang, Caiming Xiong, and Song Mei. Transformers as statisti- cians: Provable in-context learning with in-context algorithm selection. In Alice Oh, Tris- tan Naumann, Amir Globerson, Kate Saenko, Moritz Hardt, and Sergey Levine (eds.), Ad- vances in Neural Information Processing Systems 36: Annual Conference on Neural Infor- mation Processing Systems 2023, NeurlPS 2023, New Orleans, LA, USA, December 10 / html.
16, 2023, 2023. URL http::/ http:::/pappaaee /wtppaaaprssse/aa/eeeeeeee nips sccpapper_filieeeeeee oaa.aaaccS3...00a22222aaaaaaaaan
Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhari- wal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agar- wal, Ariel Herbert- Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeffrey Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCan- dlish, Alec Radford, llya Sutskever, and Dario Amodei. Language models are few-shot learners. In Hugo Larochelle, Marc' Aurelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and Hsuan-Tiees Lin (eds.), Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurIPS 2020, December 6-12, 2020, virtual, 2020. URL https: //t:: http:///ooceeeeeecooooo ngs neurips htcpp..cccee... /tte//aper/22/ 2eaper/2020/ //aah// 1457c0d cbcdddbfc4496970 96744499441411 Bbfbaacllblccccccc 42 242 besstracte html.
Dan Busbridge, Jason Ramapuram, Pierre Ablin, Tatiana Likhomanenko, Eeshan Gunesh Dhekane, Xavier Suau Cuadros, and Russell Webb. How scale your EMA. In Alice Oh, Tris- tan Naumann, Amir Globerson, Kate Saenko, Moritz Hardt, and Sergey Levine (eds.), Ad- vances in Neural Information Processing Systems 36: Annual Conference on Neural Infor- mation Processing Systems 2023, NeurlPS 2023, New Orleans, IA, USA, December 10 16, 2023, 2023. URL http:: http:::://ppp.aa//eeee pwppppppppsss..ee pimpers..ips. cc secc/appr__fille wtw.s//ppaa///223//33 e7 47681dd6felee (1w.ee665523333 at www.aactttooofeennneecmcccc html.
555bdc9- Abstract
Mathilde Caron, Hugo Touvron, Ishan Misra, Hervé Jégou, Julien Mairal, Piotr Bojanowski, and Armand Joulin. Emerging properties in self-supervised vision transformers. In 2021 IEEE/CVF International Conference on Computer Vision, ICCV 2021, Montreal, QC. Canada, October 10-17, 2021, pp. 9630-9640. IEEE, 2021. doi: 10. 1119///CC48922.222.0000 10090.. URL https: //doi. www.orrg/100 1109 ICCV 2021. 00951.
Valérie Castin, Pierre Ablin, and Gabriel Peyré. Understanding the regularity of self-attention with optimal transport. arXiv preprint arXiv:2312. 14820, 2023.
Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey E. Hinton. A simple framework for contrastive learning of visual representations. In Proceedings of the 37th International Conference Machine Learning. ICML 2020, 13-18 July 2020, Virtual Event, volume 119 of Proceedings of Machine Learning Research, pp. 1597-1607. PMLR, 2020. URL http::ttpppcceeee.eaeeee pprocediings. ml m arr.press/vll 9/chn200 html.

---

Published as a conference paper at ICLR 2025
Xinlei Chen, Saining Xie, and Kaiming He. An empirical study of training self-supervised vision transformers. In 2021 IEEE/CVE International Conference on Computer Vision, ICCV 2021. Mon- treal, QC, Canada, October 10-17, 2021, pp. 2021 .009500 URL. ht.pp:: / /ooi. org/ 1 c 11 09/ ICCV4
IEEE, 2021. doi: 10.1 109110//CC489922. 00950.
Jack Choquette. NVIDIA hopper H100 GPU: scaling performance. IEEE Micro, 43(3):9-17, 2023. doi: 10.1109/MM..023332565999 URL https: http://ooo.org..////// 1109 /MM. 2023. 3256796.
Jack Choquette, Wishwesh Gandhi, Olivier Giroux, Nick Stam, and Ronny Krashinsky. NVIDIA A100 tensor core GPU: performance and innovation. IEEE Micro, 41(2):29-33 2021. doi: 10.1109//MM.222.3366..33 URL https: /ttp:////oooooo/./////// wwt///.gg/....oo...ooooo......./ 1101111001110 /MM 2021 2021.361394... 3061394.
Krzysztof Marcin Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamás Sarlós, Peter Hawkins, Jared Quincy Davis, Afroz Mohiuddin, Lukasz Kaiser, David Ben- jamin Belanger, Lucy J. Colwell, and Adrian Weller. Rethinking attention with performers. In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net, 2021. URL https: attps:///penree/// /wpp/r/eeeeeeeeeee//e/eeee net forum?id=
Ua UWWOO...
Together Computer. Redpajama: An open source recipe to reproduce Ilama training dataset, April 2023. URL httpt:////// thu.. com/ www.ggeeerrommooooo wwwp:theeeoooeeeecctctcoce /RedPa jama-Data.
Ekin Dogus Cubuk, Barret Zoph, Jonathon Shlens, and Quoc Le. Randaugment: Practical automated data augmentation with a reduced search space. In Hugo Larochelle, Marc Aurelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and Esuan-Tien Lin (eds.), Advances in Neural Information Pro- cessing Systems 33: Anmual Conference on Neural Information Processing Systems 2020, NeurlPS 2020, December 12. 2020, virtual, 2020. URL https: htattps://proceeeeeee paper /8 (088866332 63e f0ccb1 10c:b10440003111b761111/ 4bt
http://pwceeeennneeeecee neur ips. cc/ html.
Tri Dao. Flashattention Faster attention with better parallelism and work partitioning. CoRR, abs/2307.08691,, 2023. doi: 1044850/AAXXV22270006699 URL 48550/arxiv 2307 08661.
Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, and Christopher Ré, Flashattention: Fast and memory-efficiiet exact attention with io-awareness. In Sanmi Koyejo, S. Mo- hamed, A. Agarwal, Danielle Belgrave, K. Cho, and A. Oh (eds.), Advances in Neural Information Processing Systems 35: Annual Conference on Neural Information Process- ing Systems 2022, NeurlPS 2022, New Orleans, LA, USA, November 28 December 9, 2022, 2022. URL http: http//aaprrsree//eee/eeee/eeeeee nips. wwtppppr_.illes////eccepppseeeee www.feelsspppe/////// h/aah// 67d57c32e (dteeddaaa33000 (et.eebbbbd3664440 Abstract cunfferencehtml..
Mostafa Dehghani, Josip Djolonga, Basil Mustafa, Piotr Padlewski, Jonathan Heek, Justin Gilmer, Andreas Peter Steiner, Mathilde Caron, Robert Geirhos, Ibrahim Alabdulmohsin, Rodolphe Jenatton, Lucas Beyer, Michael Tschannen, Anurag Arnab, Xiao Wang, Carlos Riquelme Ruiz, Matthias Minderer, Joan Puigcerver, Utku Evci, Manoj Kumar, Sjoerd van Steenkiste, Gamaleldin Fathy Elsayed, Aravindh Mahendran, Fisher Yu, Avital Oliver, Fantine Huot, Jas- mijn Bastings, Mark Collier, Alexey A. Gritsenko, Vighnesh Birodkar, Cristina Nader Vascon- celos, Yi Tay, Thomas Mensink, Alexander Kolesnikov, Filipp Pavetic, Dustin Tran, Thomas Kipf, Mario Lucic, Xiaohua 7hai, Daniel Keysers, Jeremiah J. Harmsen, and Neil Houlsby. Scaling vision transformers to 22 billion parameters. Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett (eds.), International Conference on Machine Learning, ICML 2023, 23-29 July 2023, Honolulu, Hawaii, USA, vol- ume 202 of Proceedings of Machine Learning Research, pp- https:
Andreas Krause, Emma Brunskill,
PMLR, 2023. URL httpps:/prroceeiingg/ http//ptcceenngg..cee www.mmrrrr.rsse. press v202/ dehghani22 html.
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hierarchical image database. In 2009 IEEE Computer Society Conference on Computer Vision and Pattern Recognition (CVPR 2009), 20-25 June 2009, Miami. Florida, USA, pp. 248 248-255. IEBE Computer Society, 2009. doi: 10. 2000/CCVRR.209.50066333 URL https: /ooi org/10. 11 2009. 5206848.
12

---

Published as a conference paper at ICLR 2025
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recognition at scale. In 9th International Conference on Learning Reprexentations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net, 2021. URL https: https://openreve ew net/ forum? ?id= YicbFaNTTy.
Hengyu Fu, Tianyu Guo, Yu Bai, and Song Mei. What can a single attention layer learn? A study through the random features lens. In Alice Oh, Tristan Naumann, Amir Globerson, Kate Saenko, Moritz Hardt, and Sergey Levine (eds.), Advances in Neural Information Processing Systems 36: Annual Conference on Neural Information Processing Systems 2023, NeurIPS 2023, New Orleans, LA, USA, December 10 16, 2023, 2023. URL http:/ http::///ppprs...e wttp:ssssppppccppppp.e..pp cccpaper__a les http//pappp22/////ppppp 274 1
html.
Octavian Ganca, Sylvain Gelly, Gary Bécigneul, and Aliaksci Severyn. Breaking the softmax bot- tleneck via learnable monotonic pointwise non-linearities. In Kamalika Chaudhuri and Ruslan Salakhutdinov (eds.), Proceedings of the 36th International Conference on Machine Learning, ICML 2019, 9-15 June 2019, Long Beach, California, USA, volume 97 of Proceedings of Ma- chine Learning Research, pp. 2073-2082. PMLR, 2019. URL http:///p/// /ttp:////cccccegg../m/o .mm.innng............mmmmm.m...m mlr. press/v97 ganealnngggg. 9a. html.
Borjan Geshkovski, Cyril Letrouit, Yury Polyanskiy, and Philippe Rigollet. A mathematical perspec- tive on transformers. arXiv preprint arXiv:2312. 10794, 2023.
Borjan Geshkovski, Cyril Letrouit, Yury Polyanskiy, and Philippe Rigollet. The emergence of clusters in self-attention dynamics. Advances in Neural Information Processing Systems, 36, 2024.
Google. Praxis.
http:ttub..commgogge////r/aattta 2024.
Alex Graves, Santiago Fernández, Faustino J. Gomez, and Jürgen Schmidhuber. Connectionist temporal classification: labelling unsegmented sequence data with recurrent neural networks. In William W. Cohen and Andrew W. Moore (eds.), Machine Learning. Proceedings of the Twenty- Third International Conference (ICML 2006), Piusburgh, Pennsylvania, USA, June 25-29, 2006, volume 148 of ACM International Conference Proceeding Series, pp. 369-376. ACM, 2006. doi: 10.1 143844. 143891. URL https: doi. org/ 10. 1145 /11//3444444441 114389.14499.
Jean-Bastian Grill, Florian Strub, Florent Altché, Corentin Tallec, Pierre H. Richemond, Elena Buchatskaya, Carl Doersch, Bernardo Ávila Pires, Zhaohan Guo, Mohammad Gheshlaghi Azar, Bilal Piot, Koray Kavukeuoglu, Rémi Munos, and Michal Valko. Bootstrap your latent new approach to self-supervised learning. In Hugo Larochelle, Marc' Aurelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and Hsuan-Tien Lin (eds.), Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurIPS 2020, December 12, 2020, virtual, 2020. URL https: /ttLLtttt///pocceedtett 2020 hat0//hashaaae
nwurips. wttp:/s...r....c..ppppeeeeeeeeee 700777112221111 2at.a44446.000411781919222223389 Abstrract caaosttracthtm.l..... html.
Anmol Gulati, James Qin, Chung- Cheng Chiu, Niki Parmar, Yu Zhang, Jiahui Yu, Wei Han, Shibo Wang, Zhengdong Zhang, Yonghui Wu, and Ruoming Pang. Conformer: Convolution-augmented transformer for speech recognition. In Helen Meng, Bo Xu, and Thomas Fang Zheng (eds.), Interspeech 2020, 21 Annual Conference of the International Speech Communication Association, Virtual Event, Shanghai, China, 25-29 October 2020, pp. 5036-5040 ISCA, 2020a. doi: 10. 2145/INTERSPPECH.200 2000-033. URL https: //tp:////i..//... org 100 21437 IntIreseee...... 2020-3015.
Anmol Gulati, James Qin, Chung-Cheng Chiu, Niki Parmar, Yu Zhang, Jiahui Yu, Wei Han, Shibo Wang, Zhengdong Zhang, Yonghui Wu, et al. Conformer: cavoouutio-augmennt transformer for speech recognition. In Proc. Interspeech, 2020b.
Awni Hannun, Jagrit Digani, Angelos Katharopoulos, and Ronan Collobert. MLX: Efficient and flex- ible machine learning on apple silicon, 2023. httpttts: //e:://giitttte wwm/m//mmmmmmo.m -emmmmmeelxporeeeeemm
13

---

Published as a conference paper at ICLR 2025
Kaiming He, Xinlei Chen, Saining Xie, Yanghao Li, Piotr Dollár, and Ross B. Girshick. Masked autoencoders are scalable vision learners. In IEEE/CVE Conference on Computer Vision and Pattern Recognition, CVPR 2022, New Orleans, LA, USA, June 18-24, 2022, pp. 15979-15988. IEEE, 2022. doi: 10.1 P99CCCCP5268..22222.00. CVPR52688 2022. 01 01553.
URL ht.ps:: /doi http//oiooooooo/oooooooooooooo.t .org 1109/
François Hernandez, Vincent Nguyen, Sahar Ghannay, Natalia Tomashenko, and Yannick Esteve. Ted-lium 3: Twice much data and corpus repartition for experiments on speaker adaptation. In Speech and Computer: 20th International Conference, SPECOM 2018, Leipzig, Germany, September 18-22, 2018, Proceedings 20, pp. 198-208. Springer, 2018.
Jiri Hron. Yasaman Bahri, Jascha Sohl-Dickstein, and Roman Novak. Infinite attention: NNGP and NTK for deep attention networks. In Proceedings of the 37th International Conference on Machine Learning, ICML 2020, 13-18 July 2020, Virtual Event, volume 119 of Proceedings of Machine Learning Research, pp. 4376-4386. PMLR, 2020. URL http: http//poooddnnsss///ee mlr molr.press/ve 9 htt//hon2000 html.
Weizhe Hua, Zihang Dai, Hanxiao Liu, and Quoc V. Le. Transformer quality in linear time. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvári, Gang Niu, and Sivan Sabato (eds.), International Conference on Machine Learning, ICML. 2022, 17-23 July 2022, Baltimore, Maryland, USA, volume 162 of Proceedings of Machine Learning Research, pp. 9099-9117. PMLR, 2022. URL httpt::tt//......... /tte::/pprrreee http://occeeddeeeeooe ngs ml press v1 62/16/2// /ww.auaaaaa... html.
Niall P. Hurley and Scott Rickard. Comparing measures of sparsity, 2009.
Andrei Ivanov, Nikoli Dryden, Tal Ben-Nun, Shigang Li, and Torsten Hoefler. Data movement is all you need: A case study on optimizing transformers. In Alex Smola, Alex Dimakis, and Ion Stoica (eds.), Proceedings of Machine Learning and Systems 2021, MLSSy 2021, virtual, April 5-9, 2021. misys.org, 2021. URL https: httttpp:://rrceeeee http://poocceeengeeeeee myssy..... org/ wwrgg/pppe//2/// (aaah//cee1111ffbbfffefeeeaelal - html.
Charles H. Jones. Generalized hockey stick identities and n-dimensional blockwalking. 1994. URL
https: //ppi/ppppppp semanttttiii. ieteentcscchoo eschol ar. hwwpcchhlaa.rrr/eeeeoooooooo Corp//rrusstii TD: -2088017 7.
Norman Jouppi, Cliff Young, Nishant Patil, David A. Patterson, Gaurav Agrawal, Raminder Bajwa, Sarah Bates, Suresh Bhatia, Nan Boden, Al Borchers, Rick Boyle, Pierre-luc Cantin, Clifford Chao, Chris Clark, Jeremy Coriell, Mike Daley, Matt Dau, Jeffrey Dean, Ben Gelb, Tara Vazir Ghaemmaghami, Rajendra Gottipati, William Gulland, Robert Hagmann, C. Richard Ho, Doug Hogberg, John Hu, Robert Hundt, Dan Hurt, Julian Ibarz, Aaron Jaffey, Alek Jaworski, Alexander Kaplan, Harshit Khaitan, Daniel Killebrew. Andy Koch, Naveen Kumar, Steve Lacy, James Laudon, James Law, Diemthu Le, Chris Leary, Zhuyuan Liu, Kyle Lucke, Alan Lundin, Gordon MacKean, Adriana Maggiore, Maire Mahony, Kieran Miller, Rahul Nagarajan, Ravi Narayanaswami, Ray Ni, Kathy Nix, Thomas Norrie, Mark Omernick, Narayana Penukonda, Andy Phelps, Jonathan Ross, Matt Ross, Amir Salek, Emad Samadiani, Chris Severn, Gregory Sizikov, Matthew Snelham, Jed Souter, Dan Steinberg, Andy Swing, Mercedes Tan, Gregory Thorson, Bo Tian, Horia Toma, Erick Tuttle, Vijay Vasudevan, Richard Walter, Walter Wang, Eric Wilcox, and Doe Hyun Yoon. In-datacenter performance analysis of tensor processing unit. In Proceedings of the 44th Anmual International Symposium on Computer Architecture, ISCA 2017, Toronto, ON, Canada, June 24-28, 2017, pp. 1-12. ACM, 2017. doi: 10.1 10.1145/3079855030000466 URL httph:tt: /tts://ddoo.roo orr. org.orgggooooo.. 10. 4 30 788686666666.. 6. 3080246.
Angelos Katharopoulos, Apoory Vyas, Nikolaos Pappas, and François Fleuret. Transformers are rnns: Fast autoregressive transformers with linear attention. In Proceedings the 37th International Conference on Machine Learning, ICML 2020, 13-18 July 2020, Virtual Event, volume 119 of Proceedings of Machine Learning Research, pp. 5156-5165. PMLR, 2020. URL http: html.
//proceedings.mt mlr. press / vtt//vvl1l///ee/eee www.://aaaarroooooeess..oee
Hyunjik Kim, George Papamakarios, and Andriy Mnih. The lipschitz constant of self-attention. In International Conference on Machine Learning, pp. 5562-5571. PMLR, 2021.
Soroush Abbasi Koohpayegani and Hamed Pirsiavash. Sima: Simple softmax- free attention for vision transformers. In IEFF/CVF Winter Conference Applications of Computer Vision, WACV 2024,
14

---

Published as a conference paper at ICLR 2025
Waikoloa, HI. USA, January 3-8, 2024, pp. 2024.00259. URL https::ps///. /oi. org/
IEEE, 2024, doi: 10.1 10.1109/WAVV5700.. 00259.
Minchul Lee, Kijong Han, and Myeong Cheol Shin. Littlebird: Efficient faster & longer transformer for question answering. In Proceedings of the 2022 Conference on Empirical Methods in Natural
Language Processing, pp. 5261-5277, 2022.
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. ArXiv e-prints, pp.
arXiv-1607, 2016.
Tatiana Likhomanenko, Qiantong Xu, Gabriel Synnaeve, Ronan Collobert, and Alex Rogozh- nikov. CAPE: encoding relative positions with continuous augmented positional embeddings. In Mare' Aurelio Ranzato, Alina Beygelzimer, Yann N. Dauphin, Percy Liang, and Jennifer Wortman Vaughan (eds.), Advances in Neural Information Processing Systems 34: Annual Conference on Neural Information Processing Systems 2021, NeurlPS 2021, December 6-14, 2021, virtual, pp. 16079-16092, 2021. URL http:: / http:s/ee////ccpeceeeeeeee www.:iinnnnnnsnnunee.ee neurips pirripp....
wwwpeepp...pppppereeeeeeeee wwtpcppppp/////./// 2021/hashh / html.
Ze Liu, Han Hu, Yutong Lin, Zhuliang Yao, Zhenda Xie, Yixuan Wei, Jia Ning, Yue Cao, Zheng Zhang. Li Dong, Furu Wei, and Baining Guo. Swin transformer V2: scaling up capacity and reso- lution. In JEEE/CVF Conference on Computer Vision and Pattern Recognition, CVPR 2022, New Orleans, LA, USA, June 18-24, 2022, pp. 11999-12009, JEEE, 2022. doi: 10.1 10.1199/29822888 2022.01170. URL https::/p//// : /doi. www.door...ooooooooooo...oooooeo 10. 1109/ 22.009///PPP CB////VVR52888 2022. 01070.
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. arXiv:1711.05101 2017.
arXiv preprint
Jiachen Lu, Jinghan Yao, Junge Zhang, Xiatian Zhu, Hang Xu, Weiguo Gao, Chunjing Xu, Tao Xiang, and Li Zhang, SOFT: softmax-fre transformer with linear complexity. In Marc' Aurelio Ranzato, Alina Beygelzimer, Yann N. Dauphin, Percy Liang, and Jennifer Wortman Vaughan (eds.), Advances in Neural Information Processing Systems 34: Annual Conference on Neu- ral Information Processing Systems 2021, NeurIPS 2021, December 6-14, 2021, virtual, pp. 21297-21309, 2021. URL https: pttttps://ppooceediinn neurips. htt.aaaa....2.2...aeeeeeeeee....
Abstract html.
Chunjie Luo, Jianfeng Zhan, Xiaohe Xue, Lei Wang, Rui Ren, and Qiang Yang. Cosine normalization: Using cosine similarity instead of dot product in neural networks. In Vera Kurková, Yannis Manolopoulos, Barbara Hammer, Lazaros S. Iliadis, and Ilias Maglogiannis (eds.), Artificaa Neural Networks and Machine Learning ICANN 2018- 27th International Conference on Artificial Neural Networks, Rhodes, Greece, October 4-7, 2018, Proceedings, Part I, volume 11139 of Lecture Notes doi: 00.1007/978-3-00--0444 _8. URL https: /tw.//ooi..o////// org. / 10. 1007 978 10708883-30--000 -00--3000---8---0 -_338
in Computer Science, pp. 382-391. Springer, 2018.
OpenAI. GPT-4 technical report. CORR, wbtpaaa3.//33/////// 2023. doi: 10.45550AARXIIV2203..0777 URL pestps://doi.oor
200.00.85500///// /tt..////aa.11220....07777.. 230..
Vassil Panayotov, Guoguo Chen, Daniel Povey, and Sanjeev Khudanpur. Librispeech: An ASR corpus based on public domain audio books. 2015 IEEE International Conference on Acoustics, Speech and Signal Processing, ICASSP 2015, South Brisbane, Queensland, Australia, April 19- 24, 2015, pp. 5206-5210. IEEE, 2015. doi: 10. 1091CASSP2220157719996 URL https: //doi. org/ / 10. 1 1109/ ICASSP 2015. 7178964.
Daniel S. Park, William Chan, Yu Zhang, Chung-Cheng Chiu, Barret Zoph, Ekin D. Cubuk, and Quoc V. Le. Specaugment: A simple data augmentation method for automatic speech recognition. In Gernot Kubin and Zdravko Kacic (eds.), Interspeech 2019, 20th Annual Conference of the International Speech Communication Association, Graz. Austria, 15-19 September 2019, pp. 2613- 2617. ISCA, 2019. doi: (0.1457/INTRRSSEECH2000 21 21437/ / Interseecrrreeeeeeeeeeeeeeeeeece 2 2680.
URL https: org/10,
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Köpf, Edward Z. Yang, Zachary Vetito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy,
15

---

Published as a conference paper at ICLR 2025
Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. Pytorch: An imperative style, high-performanc deep learning library. In Hanna M. Wallach, Hugo Larochelle, Alina Beygelzimer, Florence d' Alché-Buc, Emily B. Fox, and Roman Gamett (eds.), Advances in Neural Information Processing Systems 32: Annual Conference on Neural Information Pro- cessing Systems 2019, NeurlPS 2019, December 8- 8-14, 2019, Vancouver, BC, Canada, pp. 8024-8035, 2019. URL https: https::/proceeeeiiieeeeee http:///poccceeeeeeeeeeeeeececee neurips httpa//sss.iipspspppa///////// bdbca288 weeeefeeeeeee www.222290000222000022
Abstract html.
Ofir Press, Noah A. Smith, and Mike Lewis. Train short, test long: Attention with linear bi- ases enables input length extrapolation. Representationss ICLR 2022, Virtual Event, April 25-29, 2022. OpenReview.net, 2022. URL nttps://openzeviie www.:/neeetttiieeeeeeeee forumiiddR&s@PpsSSS
The Tenth International Conference on Learning
net/
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J. Liu. Exploring the limits of transfer learning with unified text-to- text transformer. J. Mach. Learn. Res., 21:140:1-14000 2020. URL http://jmlr:
papers/v21/20-e1 httml...
Noam Shazeer. Fast transformer decoding: One write-head all you need. CoRR, abs/1911.022150 2019. URL
111...1..11 02150.
Kai Shen, Junliang Guo, Xu Tan, Siliang Tang, Rui Wang, and Jiang Bian. A study on relu and softmax in transformer. CoRR, abs/2302.0646/ 2023. doi: 1048550/RXIV.2300.066 https: //opi: /wiiooiooo./////../ /10 / www...
URL
Jianlin Su, Murtadha H. M. Ahmed, Yu Lu, Shengfeng Pan, Wen Bo, and Yunfeng Liu. Roformer: Enhanced transformer with rotary position embedding. Neurocomputing, 568: 167:12, 2024. doi: 10.1016/..EECCM.... 127666 URL https: httpt://dooiioo http///....../1..1//.//. http:///....11/.eeeeeee...eeee neucom. 2023. 127063.
Gabriel Synnaeve, Qiantong Xu, Jacob Kahn, Tatiana Likhomanenko, Edouard Grave, Vineel Pratap, Anuroop Sriram, Vitaliy Liptchinsky, and Ronan Collobert. End-to-end ASR: from supervised to semi-supervised learning with modern architectures. In /CML 2020 Workshop on Self-superviision in Audio and Speech, 2020. URL https: https::/openeeviieeee /ttp///eeeeeeeeeee/e/eeee net/ forum?
id=0SVxDDc360z.
Hugo Touvron, Matthieu Cord, Alexandre Sablayrolles, Gabriel Synnaeve, and Hervé Jégou. Going deeper with image transformers. In Proceedings of the IEEE/CVF international conference on computer vision, pp. 32-42, 2021.
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint urXiv:2307.09288. 2023.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Isabelle Guyon, Ulrike von Luxburg, Samy Bengio, Hanna M. Wallach, Rob Fergus, S. V. N. Vishwanathan, and Roman Garnett (eds.), Advances in Neural Information Processing Systems 30: Annual Conference on Neural Information Processing Systems 2017, December 4-9, 2017, Long Beach, CA, USA, pp. 5998-6008 2017. URL https: http://poooeeeecccccoccoccc neur neurips cc/ pppper/203 /hash/ 3f5ee22354 7dee91 fbd053c 44aa Aast. html.
Ross Wightman, Hugo Touvron, and Hervé Jégou. Resnet strikes back: An improved training procedure in timm. CoRR, abs/2110.00476, 2021. URL htthht:://// /tttts:///rxiv.. wwwp///..///////e/////.//...e/// a2s/2110. 00476.
Mitchell Wortsman, Jachoon Lee, Justin Gilmer, and Simon Kornblith. Replacing softmax with ReLU vision transformers. arXiv preprint Sttt:2200.08588 2023a.
Mitchell Wortsman, Peter J. Liu, Lechao Xiao, Katie Everett, Alex Alemi, Ben Adlam, John D. Co- Reyes, Izzeddin Gur, Abhishek Kumar, Roman Novak, Jeffrey Pennington, Jascha Sohl-Dickstein, Kelvin Xu, Jaehoon Lee, Justin Gilmer, and Simon Kornblith. Small-scale proxies for large-scale transformer training instabilities. CoRR, htbs//309.14322 2023b. doi: 10.48550/ARR(VV230999 rxiv 2309. 14322.
14322. URL https: htt///oooooo/////////
.org/ 48550
16

---

Published as a conference paper at ICLR 2025
xai-org. Grok-1. https: httppp/s/gggtha....c////aaaaa/// Date].
2024. Accessed: [Insert Access
Ruibin Xiong, Yunchang Yang, Di He, Kai Zheng, Shuxin Zheng, Chen Xing, Huishuai Zhang, Yanyan Lan, Liwei Wang, and Tieyan Liu. On layer normalization in the transformer architecture. In International Conference on Machine Learning, pp. 10524-10533. PMLR, 2020.
Zhilin Yang. Zihang Dai, Ruslan Salakhutdinov, and William W. Cohen. Breaking the softmax bottleneck: A high-rank RNN language model. In 6th International Conference on Learning Representutions, ICLR 2018, Vancouver, BC, Canada, April 30 May 3, 2018, Conference Track Proceedings. OpenReview.net, 2018. URL hwwph:tp////ot/ https////ppeeeeee..cmmmm /t net forum?id= Hkw2SG-C2.
Chulhee Yun, Srinadh Bhojanapalli, Ankit Singh Rawat, Sashank J. Reddi, and Sanjiv Kumar. Are transformers universal approximators of sequence-to-sequenec functions?, 2020.
Shuangfei Zhai, Tatiana Likhomanenko, Etai Littwin, Dan Busbridge, Jason Ramapuram, Yizhe Zhang, Jiatao Gu, and Joshua M. Susskind. Stabilizing transformer training by preventing attention entropy collapse. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett (eds.), International Conference on Machine Learning, ICML 2023, 23-29 July 2023, Honolulu, Hawaii, USA, volume 202 of Proceedings of Machine Learning Research, pp. 10770-40803. PMLR,, 2023a. URL https: http:s::pprcceeeddnngg mlr press/ htmll
Xiaohua Zhai, Basil Mustafa, Alexander Kolesnikov, and Lucas Beyer. Sigmoid loss for language image pre-training. CoRR, abs/2303. 153443 2023b. doi: 10.48350/ARXIV233355555 URL
https://dbi.org// 48500 arXiv 33333.11333
Biao Zhang and Rico Sennrich. Root Mean Square Layer Normalization. In Advances in Neural Information Processing Systems 32, Vancouver, Canada, 2019, URL https: / /penreviee. htth//eeeerremccee/p////...o.o.
17

---

Published as a conference paper at ICLR 2025
Appendices
Limitations
20
Broader Impact
20
C Universal Approximation Property for Sigmoid Attention
21
C.I Proof of Step (3): Sigmoid Transformers can Approximate Modified Sigmoid Trans- formers
21
C.2 Proof Step (2-b): Modified Sigmoid Transformers can Implement Contextual Mappings
22
C.2.1
Basic Building Blocks of Contextual Mapping
23
C.2.2 Result of Applying Sequence of Selective Shifts
24
C.2.3
Result of Applying One Last Global Shift Layer
26
C.2.4 A Sequence of Selective Shifts Followed by Global Shift Produces Con- textual Mapping
27
Lipschitzness of Sigmoid Attention
31
E
The Bias Term of Sigmoid Attention
32
F Details of FLASHSIGMOID
34
F.1 Details of FLASHSIGMOLD Algorithm
34
F2
Benchmarking of FLASHSIGMOID Kernels
37
F3
Speed Boosts of FLASUSIGMOID in Realistic Settings
................................
38
F4
FLASHSIGMOID with ALiBi
40
F.5
Directions for Future Work on FLASHSIGMOID
41
G Experiments
42
G.1 Extra Ablations
42
G.1.1
The Effect of Multiplicative Sequence Length Normalization
42
G.1.2
Attention Bias Stability Ablation
G.2
Vision
43
G.2.1
Test ImageNet Top- 1%
43
G.2.2
LayerScale Free Sigmoid Attention
44
G.2.3
Sigmoid Attention Attention Relaxations
44
G.2.4
Hyper-Parameters.
45
G.3
Language Model
45
G.3.1
Hyper-Parameterr
45
G.3.2 Gradient Norm
46
G.3.3
Norm structure
47
G.4
Automatic Speech Recognition
47
18

---

Published as a conference paper at ICLR 2025
G.4.1
Training Details
47
G.4.2
Results and Ablations
49
G.5
Simple Experiments
51
G.5.1
k-Summation Problem Definition
................................
G.5.2
Comparison to Softmax
................................
52
G.5.3
Attention Evolution
52
G.5.4
Pair Repeat Problem
53
G.6
Computational Complexity of Sigmoid and Softmax.
54
G.7
Practitioner's ' Guide
54
Contributions
55
19

---

Published as a conference paper at ICLR 2025
A LIMITATIONS
While our work demonstrates that SigmoidAttn can serve as viable drop-in replacement for SoftmaxAttn in many domains and scales, there are a few key limitations to note:
(1) In large-scale (1B parameter, 4096 context length) language modeling, observed some gradient norm spikes and a slight performance gap between SigmoidAttn and SoftmaxAttn (Table 2). While runs smaller context lengths (IB parameter, n=2048) were stable and matched Soft madt t&n performance, we required the use of hybrid-norm to stabilize n=4096 sequence length models. Hybrid-norm does incur slight extra performance penalty which we quantify in Appendix G.7.
(2) Our theoretical analysis proves that transformers with Sigmoid Attn are universal function approximators and have improved regularity compared to SoftmaxAttn However, the bounds we derive, while tighter than those for SoftmaxAttn, may not be maximally tight. There could be room for further theoretical refinements.
(3) We focused our empirical evaluation on standard benchmarks in language, vision, and speech domains. Performance on more niche emerging applications remains be validated.
(4) In automatic speech recognition experiments, we observed that Sigmoid Attn can be sen- sitive to the choice of positional embeddings and may require careful initialization of the attention bias term ensure stable training. Specifically, found that the CAPE positional embedding was the most unstable for Sigmoid Attn. Further work is needed to develop ro- bust initialization schemes that work well across different positional embeddings. Moreover we found that w/o QK norm or with post-LayerNorm SigmoidAttn is unstable and can underperforms SoftmaxAtta thus further investigation needed.
(5) FLASHSIGMOID demonstrates promising inference and training speed-ups by exploiting SigmoidAttan 's simpler kernel structure compared to Softmax However, realizing these gains at scale in distributed training setups may require additional engineering to optimize communication bottlenecks.
Despite these limitations, we believe this work establishes strong foundation for Sigmoid unifying prior art and demonstrating its potential as drop-in SoftmaxAttn replacement. We hope our theoretical grounding and empirical results motivate further research into this simple yet effective architectural variation.
B BROADER IMPACT
The development of efficient and theoretically grounded attention mechanisms has the potential for significant positive impact across a range of applications. By establishing Sigmoid Attn viable alternative to SoftmaxAttn, our work expands the toolkit of architectural choices available to researchers and practitioners. Positive impacts of this work may include:
(1) Improved computational efficiency: FLASHSIGMOID'S 's faster kernel implementation could lead to more efficient training and inference for sumption and enabling deployment on resource-constrainee devices. This could democratize access to powerful models.
models, reducing energy con-
(2) Theoretical understanding: Our universal approximation results and tighter bounds on the regularity of SigmoidAttn contribute deeper theoretical understanding of this key component. A stronger theoretical foundation can guide principled model design and architectural search.
(3) Application-specifee benefits: Across language, vision, and speech domains, SigmoidAttn' performance could translate into improved user experiences, such as more natural language interactions, enhanced image understanding, and robust speech recognition. These advance- ments could have positive societal impacts, such as improved accessibility tools and more effective educational technologies.
However, as with any foundational machine learning advance, there are also risks of negative impacts that must be considered and mitigated:
20

---

Published as a conference paper at ICLR 2025
(1) Faimess and bias considerations: As with any machine learning model, it is important to carefully evaluate SigmoidAttn based models for fairesss and potential biases when applied to sensitive
cases. The unique properties of SigmoidAttn may have unexpected interactions with data biases. Researchers and practitioners should follow best practices for auditing and mitigating unwanted biases to ensure equitable outcomes.
(2) Environmental impact: While FLASHSIGMOID is more computationally efficient than FLASHATTENTION, the overall trend of scaling up attention-based models has significant energy costs. Further efficiency improvements and the use of renewable energy sources are important to mitigate environmental harms.
We believe that the benefits of SigmoidAttn outweigh the risks, but it is crucial for the research community to actively consider and address these potential negative impacts. By doing so, we can work towards future where the efficiency and expressivity of Sigmoid Attn are used for societal benefit.
C UNIVERSAL APPROXIMATION PROPERTY FOR SIGMOID ATTENTION
This section is dedicated to the proof for the Universal Approximation Property for attention equipped with sigmoid nonlinearity. The proof follows closely the one provided in Yun et al. (2020, Sec. 3), of which we inherit much of the notation, and we encourage the interested reader to refer to the original source for more comprehensive understanding of its details. Here we first provide context by outlining the main steps in the original proof, before proceeding to adapt key components to the SigmoidAttn case.
The proof aims at showing that a transformer network can approximate to arbitrary accuracy any continuous, permutation-equivarraanl function with compact support. The proof is constructive in nature, in that explicitly defines the architecture (and particularly, the sequence of self-attention and wred-forward layers) that can approximate given target function. To do so, it proceeds in steps (see Yun et al. (2020, Sec. 3.2)):
(1) prove that any continuous function with compact support can be approximated to arbitrary accuracy by a piecewise constant function
(2) prove that an aptly-conttructed modified transformer network, (where the softmax nonlin- earity is substituted with hardmax nonlinearity), can exactly represent such piecewise constant function. This step is further divided into three sub-steps (see Yun et al. (2020, Sec. 4)):
(a) prove that series of feed-forward layers can quantize any input to a specific discretiza- tion grid in the compact domain
(b) prove that series of self-attention layers can implement contextual mapping (see Yun et al. (2020, Def. 3.)))
(c) prove that series of feed-forward layers can map the output of the contextual mapping
to the desired output of the target piecewise-constan constant approximation
prove that a (classical) transformer network can approximate such modified transformer network to arbitrary accuracy
Fortunately, some of the steps outlined above do not rely on a specific nonlinear function being used within the attention mechanism, and can be directly reused in our proof, virtually unchanged. Notice however that Steps (2-b) and (3) are directly impacted by modifications to the attention layer, and hence require adaptation in our case. This is the focus of the next sections.
C. PROOF OF STEP (3): SIGMOID TRANSFORMERS CAN APPROXIMATE MODIFIED SIGMOID TRANSFORMERS
In Yun al. (2020), to implement contextual mappings. the authors rely on modified version of transformers, for the sake of simplifying the analysis. In their modified version, the (row-wise) softmax operation is substituted with (row-wise) hardmax operation. This substitution is valid
21

---

Published as a conference paper at ICLR 2025
because classical transformer can still be made arbitrarily close to such modified transformer, in light of the fact that
softmax (AX) A hardmax (X).
(10)
In our proof, we follow a similar strategy to define our modified sigmoid transformer (and in particular, its self-attention mechanism). We have that
(AX)) A-o H
(11)
where (x) (1 e ")-
1 is the (elementwise) sigmoid function, while
1
0
((x)
1 x
(12)
2
x <00 0
denotes the (elementwise) Heaviside step function. This allows to define modified sigmoid self-attention layer, as follows.
Definition C.1 (Modified sigmoid modified sigmoid alff-attention layer with shifts and single one mee-imensional head is defined as (X; q, by. k. by: U, 0), where
layer), Given an input X Rdxx, the action of
X
v/( x; q, by: k. bk, v, 0) o o H
(k² (6"x-bb) -xb)))
(13)
with q, k, u E Rd representing the query, key, and value vectors, by. by R." the corresponding query and key bias vectors, while E IRd denotes the output vector.
Analogously to (10), (11) guarantees that sigmoid attention can approximate modified sigmoid attention by simply increasing the magnitude of its inner parameters.
Here and in the following, the length of the input sequence is denoted as n. while d represents the dimensionality of the tokens. Notice that are considering the input tensor X Rdxn, (as opposed to E IR x)) to better align out notation with the one used in Yun et al. (2020).
C.2 PROOF OF STEP (2-B): MODIFIED SIGMOID TRANSFORMERS CAN IMPLEMENT CONTEXTUAL MAPPINGS
The core of the proof consists in showing how, by opportunely combining the operations in (13), one can build an architecture capable of implementing a contextual mapping. For completeness, we report next the definition of such a map (see also Yun et al. (2020, Def. 3.1)).
Definition C.2 (Contextual mapping). A map q L + R" from finite set C [d is said to be contextual mapping if both the following conditions hold:
q
X
Vi + j and AX E LL
iii)
q:
9j
Vi,j and VX, X' E L, with X & X'
where (X) denotes the i-th component of (X))
Namely, a contextual mapping is such that it transforms each token in an input sequence to a value depending uniquely on the whole sequence. By satisfying this property, we can ensure that any element of the quantization of the input domain (achieved by Step (2-a)) can be mapped to unique identifying value (depending on the whole input) via a sequence of modified sigmoid self-attention layers. It is then up to the MLP (in Step (2-c)) correctly map this value to the corresponding output value in the piece-wisc wise constant approximation.
In particular, after defining uniform discretization (characterized by the parameter 8) of the unitary hypercube [0, 1]d C Rdd. namely
G (g: 9: E {0, 28, 5), Vi = 11.. d}.
(14)
we consider as input tensor X (composed of columns X [w.]]>.> such that
X I = (X Go Vi 1 I...n, and + æ. Vi + Rd
(15)
22

---

Published as a conference paper at ICLR 2025
that is, 2D tensor whose columns element of the discretization G5. and that all differ from each other (at least for one element). We want to build a contextual mapping acting on L, by stacking layers parameterized according to Def. C.I. In App. C.2.1 we define the basic building blocks of our architecture; in App. C.2.2 we describe how to stack them, and the effect the architecture has on a given input; finally, in App. C.2.4 we prove that this architecture indeed implements contextual mapping.
C.2.1 BASIC BUILDING BLOCKS OF CONTEXTUAL MAPPING
The strategy we follow to assemble a contextual mapping consists in sequentially looking at each column of the input, progressively updating and storing information regarding its content in a uniquely identifiable manner, and finally broadcasting this information back to every element in the sequence. The difficulty lies in the fact that each of these updates must be carried on while relying solely on applications of the modified SigmoidAtin Attn layer in Def. C.I. In the following, we describe how we can tweak its parameters achieve exactly this.
From d-dimensional quantized vectors to scalars As a first simplification, we can get rid of the d-dimension in the X tensor by mapping each of its columns to corresponding identifying scalar, uniquely defined by the specific column components. This step is also performed in Yun et al. (2020, App. B.5), and can be achieved rather straightforwardly. by defining
e = q = u [1,6- 2......,, 5 d+1TT.
.....,
(16)
Notice in fact that, since each column æ; belongs to G5. it can equivalently be written in the form 2; = 5 [ido.s. id........iid. idg-1.s'", where idj,, E {0, 1.2..... 8-1 I 1} represents the (indexed) coordinate of the discretization along the j-th dimension. Scalar-muliiplying by u in (16). then, turns this tuple of indices into single one, in bijective fashion
This allows us to equivalently consider a single vector u E R", rather than the whole tensor X E Rdxn in the remainder of our analysis. Analogously, choosing o constraints the effect of the layer application to impact only the first row of the tensor: the goal is then to store in this row the result of the target contextual mapping q in Def. C.2. To slim our notation, in the following we often refer to u X as the vector E TR with components 4,
== (1.0..... 00" in (13)
In light of the simplification above, can rewrite (13) more compactly, as follows:
5 (x: q k u. en: by b&) eat H ((1 b,) @ (L bb))
(17)
Notice that, since the elements of both X and u are always non-negative, so are those of I, too. Moreover, since we are interested in connutation-equivrrinn functions with respect to the columns of X, without loss of generality we can consider the elements of = u² X to be ordered: 0 < 4 < lj. Vi j.
Selective shift operation for sigmoid attention Since we aim to recover contextual map by sequentially updating the elements of I, we proceed by designing a modification (17) which affects only a certain selected element at time. This is were our second simplification comes into play, and this time it pertains the roles of the bias vectors by and by- Since 2 O, these vectors have the effect of tweaking the sign of the inner arguments of the Heaviside function in (17), hence directly impacting when its application outputs 0 or 1. By aptly selecting the values of b, and b,,. then, we can explicitly decide when specific layer triggers update, which elements are affected by the update, and what elements to consider to compute the update itself.
More in detail, take b, 1b, and b, 18g. for some scalars by: by and with being the all-ons vector. Plugging this into (17), we have
(X;; by bb) =
== =-= u eo, b, -1b by = 16k)
Emdiee 4
if < by if 1 > bk
(18)
= = H ((1 - 11bb) o (1- 1bx)) = en Σ Endiibo 4
'For example, consider d and the column defined as æ; = [38, 106, 28)". that is. column identified
by the triplet of indices [3, 10, 2). Multiplying by would then give the scalar ux" m, (3 + 10N + 2N2)., where N N=o 8-1, which is uniquely identified by the single index (3 + 10N + 2N2).
23

---

Published as a conference paper at ICLR 2025
notice how by determines what elements of I compose the update (as it impacts the indices considered in the sum), while by defines the elements impacted by the update itself 7, If we opportunely combine four modified sigmoid self-attention heads (xx:; by. bx), we recover, for given index i= -0...0-
1b (X; by = 0, bk (i - 2) 6)
16 (X; by O, by (i 6)
(x) = 2
(X; by bk (it 2)
+ (X;ba (i 2),
-))
H (1 (1 (i-
H1 (1 (1+
=X ccel
(22)
- H ((1 (1+
(i+ ! 8)
+H ((1 (i+
(i 6)
Exdisia 4
if i6
- (X) X1.5 + c
otherwise
1. (X) XK.i
where c = c(6, d, n) is multiplicative constant which will be chosen later.
The operator assembled in (22) defines the basic layer of the architecture that we use in our proof. Notice ()(0 has the effect of modifying only the column æj which has index 1j u æj i6 (if at all present in the input X). This layer covers a similar role to the selective shift operation introduced in Yun et (2020, App. B.5), but it has been adapted account for the presence of sigmoid nonlinearity: notice this required us to use 4-headed attention, while in Yun et al. (2020) 2-headed version is sufficient.
C.2.2 RESULT OF APPLYING A SEQUENCE OF SELECTIVE SHIFTS
Ultimately want to show how. by stacking sequence of selective shift layers (22) for increasing
i O.... 8-d and one additional global shift, we can build an architecture capable of representing
7 This can be hetter seen by considering independently the effects of the two paramcters bR. by on the modified sigmoid attention matrix H ((1 1b,) @ (L 1b²)). We have in fact, with b, = 0.
1 < by 1 > bk
0 ... 0
(1 (1 1bx))
:
=
:
(19)
This shows how, by modifying by, one can decide which columns will receive an update: namely, all those with index by- By combining such operators with by (i _1) 5 and (i+ 1) 8, then recover
1 18
0 0
H ( (L (1-1 H (I (1 (1- (1-1
(20)
which allows us to limit update to only one specific column: the one with index 1, 15. The parameter by acts analogously, but varies the output of the Heaviside function as we move down the rows,
rather than the columns. The same operator as in (20), but with b (i 4) 8 gives us in fact:
1 i8
0
L < i8
(1 H (1-1
(21)
14 1 15
Finally, (22) be recovered by combining (20) and (21): this has the effect of removing the -1's in (21).
24

---

Published as a conference paper at ICLR 2025
contextual mapping. As preliminary step, in this section provide explicit formula for the result of applying such an architecture. Once again, we are proceeding analogously to Yun et al. (2020, App. B.5.1).
After the first selective shift application Consider a quantized input sequence X E LL as defined in (15), with its columns ordered according to their scalar indices u X. The sequence of selective shift layers ()((), y .... initially has no effect the input itself, and it leaves it unchanged until we
hit the layer corresponding the index of the first column in the input, g((), where 1 u² æ1 id. At this point, following (22), the first column of the input is modified into
X 21 + CEO Σ 4= ceo
Σ 4 -4-4 11
(23)
while the other columns are still left untouched. In the following, we compactly refer to the quantities Σ lx
8g:
8 = [11. : sn
4
11
4k
In
(24)
According to (23), the index 1 of column is then analogously mapped to
u
i (x)
u xx + 081 c81
(25)
Notice that, by choosing 1, we can ensure
<> 1 =
i
1/
1/
Vi,
(26)
1=1
k=1
and particularly & 12, implying that at the next (effective) application of the selective shift operation, this term, too, will contribute to the update.
Subsequent selective shift applications Following similar considerations, the next effective update will be applied by the layer (( with u" æ2 id. At this point, the second column index is updated as follows:
la u² x22
1== 1 (X)
u² x2 +c
Σ
4 le+
(27)
=12 +
Σ 4. -2 12
K
+ 1/ +
cs1
= 1/ + C82 +
where 7, is also included in light of (26), and we used the definitions (24) and (25). Continuing to apply (+)) (X), for increasing i, and unrolling the recursion, we recover
L la +
Σ 4k -1-4 11 12 - -bai, + 12 la + csa + ( + 81) 6881
k=1
4 4 +
4 1 12 - 1/ 44 + 7, + i2 +
4 + 0S4
2 (53 s2) e552 + 251) + c's1
(28)
1 15 +
Σ 44-444--
4
+ 4 + 12 + 1 +11
1s + CSs + + c 83 83 + + $1) + c² (33 + 239 + 331) + c + 381) + d'
25

---

Published as a conference paper at ICLR 2025
which eventually allows to write the general formula"
k
1 csj
8k-1+1
j j= 1...n. n.
(29)
C.2.3 RESULT OF APPLYING ONE LAST Global Shift LAYER
After the last selective shift layer, the original input X has been mapped modified one X whereby each column 3y is characterized by the index i, u² át given in (29). Remember our goal is to recover contextual mapping, but notice that these I indices are nor uniquely defined by the input"; in other words, they do not satisfy property (2) in Def. C.2. The only exception to this is the last index i,, as (loosely speaking) it has "seen" all the previous updates and indeed in App. C.2.4 we prove this rigorously, under some assumption on the yet-undefined coefficient c(ó, d, n).
A straightforward way to recover one-to-one mapping for the whole sequence, then, is to update every index in via quantity directly depending on l,, This is precisely what the last global shift layer (X) aims to accomplish. This last layer is also defined starting from the simplified modified sigmoid attention (18), by picking by- = 0 and by (c(6, d, n)" + 2)) ó: if, for any input, we can guarantee that
i S d. n) inn and in > then the application of the global shift layer would result in
(30)
(X)) X + catl 1 b
c"
1
6, by =0)
2
(32)
- 1 X
X1 i. + ccttlin
(X)
X s.jj
The global shift (32) is the last layer we need to define our candidate contextual mapping. Collecting the results from this section together, our architecture defined by sequentially composing the selective shift layers with the global shift one,
(X) = T (
(2 (X).
(33)
After being scalar-multiplied by u. this results in sequence
(X) u² T (X) = i 1
(34)
which we aim to prove is a contextual mapping. This is shown in the next section.
FFrom (28), we can notice that, for given lt., the coefficients (k) each of the á¹ terms, are first given by list of ones, a
a. appearing in front of the various for (k)
(+) 1, then list of increasing numbers a
d=2 = cumsum(a"" (a then list of triangular numbers a i(i + 1)/2 = a= = cumsum (a 2) and = cumsum(a"" (a.) j-1). .). The result of iterated applications of cumsum, starting from all-one vector. be compactly described via the binomial coefficient: have in fact
01. = jcumsum / 1,,, = i (+
j-1 (
The actual formula (29) can be recovered after few algebraic steps, by rearranging the summation indices.
"To convince ourselves of this, suffices to look the formula for (25): two sequences with different
elements &f 1', but such that I li and sı = of (that is, with Σn= Σn (t) would map to the same 7 K-
As in footnote 7, this is also better seen by considering the resulting modified sigmoid attention matrix. n))" n)" + 1) 6. in fact, if condition (30) is verified, this matrix is given by With by 0 and b, (c(6, d,
H (i- +1))) =
T, <n
(31)
L
26

---

Published as a conference paper at ICLR 2025
C.2.4 A SEQUENCE OF SELECTIVE SHIFTS FOLLOWED GLOBAL SHIFT PRODUCES CONTEXTUAL MAPPING
To complete the proof, it remains to show that the recovered sequence (34) represents a contextual mapping and, in particular, that it is (i) one-to-one in L, and that (ii) all of its elements are distinct for different inputs. To do so, we need few preparatory lemmas. The first few are needed to show that each of the basic components of (34) is indeed a one-to-one map.
Lemma C.3. The map 8 in (24) is one-to-one.
Proof. The target map can be compactly represented linear operator S:
1 Σ L - (1 (1-1- St
(35)
k=1
which is invertible", denoting that - s is bijective.
Lemma C.4. The map + in in (29) is one-to-one, under the condition
c(o, d, n) (n 1)(8~
(36)
Proof. Consider two vectors of column indices 1, differing for least one element. We have by definition (29) that
n-
12- 2
I, - P, = (1, 1,) + c(sn x/ +
(8k-i41
(37)
1=0
k=
By absurd, assume I, - 1/,- = 0 even though : #
v.
We have then that must hold
n-2
n-2
(1n - 1n) = c(sn - ('n) Σ
142
(sk-it1 564-1)
i=0
k=
(38)
2
k
= c
(sn -
(8k- (sk-it1 - skp--11
1=0
k=
Notice that, for c(O, d, n) large enough, the right-hand side does not have enough granularity to counter the left-hand side: in fact, since in E {0, 8, 26..... 8-d+1 6), the left-hand side can attain values
H, - I, e (0, 26 106 d41 8)}
(39)
while the former, in light of the presence of the c(ō, d, n) factor, can only attain values (0, +to5, +-2c0.... Picking > 8-4 - 1, then, ensures that equality between the two sides of (38) can only be achieved if they are both 0. In this case, we need to impose
n-2
R-2
c(5/, -
8p)
5k sk-iit1 - 5 5k4))1)
k=
(40)
2
8/ 3n =
(sk sb-111)
1=0
k=
Similarly, notice !2 Vi,
|s: 5/1
Σ (1e
(2) - (4, 4))
Σ k= 1kky
(4+
(2)
< (n 1) -8), (41)
Indeed inverse can be explicitly recovered by directly applying Sherman-Morriso formula.
12This is a direct consequence of the definition of operator S in (35): since it has I's everywhere but its diagonal, its co-norm is simply 1.
27

---

Published as a conference paper at ICLR 2025
implying that s/1 8n E {0, =6, +26, .... =(n -)(6-~) - 1)0}. Again, by picking c(Ō, d, n) > (n () 1) we ensure that the right-hand side does not have enough granularity, and hence c(o, (( > (n 1) 11
s/2
Sn = 0,
(42)
implying
n-2 2
n 2
Σ
5/ sh-itt) = 0
(8k-i+1
k=
2
($1441
-8x+1) =c
n-
1
($t i+1 - sh-ie1)
(43)
k=
n-
Σ
(sk-it1 shonei)
Following a similar reasoning as the one applied above shows us that picking
c d, n) (n 1)²(0-1
1h-1 8/411) O,
(44)
and requires to satisfy
2-2
k
sh-et1)
8k
8k-1+1
=0
n-2
n-2
n-2
(
8k) =
Σ
2-2
Σ
(Sk-1+11 .b-s11))
(8k
(45)
1=2
2-2
R-
k( 8k) =c
(8h-141
1=
kmi
Once again, then, by choosing
c(6,d,n) (n - (n - 1)² (8
2
1)
Σ k(sp - 5/)) = 0.
(46)
2
This reasoning can be repeated recursively: at each step i of the recursion, by imposing a stricter and stricter bound c(o, d, n) we gain more and more conditions that the quantity s' needs to satisfy:
(o (n (1
Σ (SA-it1 - - 1 = 0. (47)
k=
Notice that, every time we increase s/1-411' h ===i. i... 2: if we were collect all these conditions within single linear system, the system would have an triangrri structure, and hence be non-singular This implies that for the set of independent conditions on s s' hold (we have n 1 in (47), plus one more in (42)), the only possibility is that = s'. Because of Lemma C.3, though, this also implies 7 = Z': we have finally reached contradiction, and proven that indeed 1 + I, is one-to-one, under opportune condition on c(6, d, n). Such condition can be promptly recovered by (47):
O... - 2, these conditions involve one less term 8k-i+1
max v O....
2
Σ
n22
max 1m....
n 11-1
(i+1
7
nー1
(48)
Substituting this (47), we recover that it suffices to impose
c(6, d, n) (n 1)(6
(49)
n-
17This is consequence of some useful properties of the binomial coefficient, namely the Hockey stick identity Jones (1994), and the symmetry of (t) with respect to i.
28

---

Published as a conference paper at ICLR 2025
The next few lemmas needed bound the elements in the 1 sequence, which in turn used to prove property (it) in Def. C.2.
Lemma C.5. Ij in (29) an increasing sequence.
Proof.
This can be proven directly: we have in fact, by definition (29),
j-2
2
1 = 17 + c8j +
i41
1=0
J-3
j 3
>1j-1 + cS
222
sk-it1
j-2
j - 2
combine sums
(11-)j-)(11c))- -) +
.42
8j 8--1-1 > 0
i
1=0
j-2
(77) 21, 0 20 (b-14-1)(1-))) -) c
aj-1-i20
j-2
22
c) = (1-- -
Σ
4
1--1
20
1=0
k=
= (1-1--)(11 +c²
(j-1) Σ
4/
4
>0 0
k=1
- (1 - c)l +
+ (c - 1)1 1- (j -2) Σ
1 +
1 > 0
k=1
k=
(c- c 1)ly + (c 1) + - Σ
4
c² Σ
1 >
0
k=1
k=j+1
(50)
Already with e 1, all the coefficients are positive (and at least one is non-zero), implying that the condition above is always satisfied and that indeed I is an increasing sequence.
Lemma C.6. Under constraint (36), each term [g, j >> 1 in (29) is bounded from below by
i, > 8,
and each term i,, j < n is bounded from above by
i
Proof. We start by proving the lower bound. By definition (29), we have
2-2
j-2
k i
i, 1 + cs,
8k-it1 = + t + é
8k-it1- (51)
Since by assumption 1ú is ordered sequence without repetitions, for j we necessarily have
4 > 4 2 0, and hence I, 2 8. All the other terms in (51) are non-negative, so we can safely claim that
i d8
vj>1,
(52)
which confirms the lower bound.
29

---

Published as a conference paper at ICLR 2025
For the upper bound, we start again from the definition of Iy:
j-2
j-2
1 + 08
+2
sk-it1
1=0
j-2
j 1-1
< <00- - -16 + c(n 1) -d
1)6 +
81
(53)
i+1)
1=
$ 11- (
1
c
1-c
where we used relationship (48) and collected all terms within the sum. Notice that, for given a > we have that
S ac'
(54)
1 -
provided that c <> _ In fact,
1 1- d+1
J+1 acd + acit1
1-c S ac
0
1-c c
(55)
a
1
0
20
a-1 a a-1
a
which is always satisfied, After substituting (54) in (53), this allows us to write
i, < - -)
ا óc.
(56)
To prove that I2 < 5ct+1, then, it remains to show that
e 2 a(n - 1)(6-- 1-1
AI < < n.
(57)
2
Substituting condition (36) in the inequality above, we are left with proving
(14))
2
[
max n-
2=
(58)
The outcome depends on the parity of n. For n odd, we have
2
2 2 a
(59)
1=2
to satisfy which it suffices to pick = 2. This requires having c c --T = 2, which is automatically satisfied. For even, on the other hand, the binomial coefficients simplify to
n 1-11 Z
n -
n 2 a.
2
a
(60)
To satisfy this, need pick ا=_ which requires - 2/ however, this too is automatically satisfied by (36) provided "
4. This completes the proof.
Lemma C.7. Under the constraint (36), condition (30) holds.
Proof. We remind that condition (30) is necessary for the correct "functioning" of the global shift layer, and it composes of two parts. The first part requires that i, < c= Vj < n. Thanks to Lemma C.5, it suffices to show that i < c"5, but this is already granted by the upper bound in Lemma C.6. Analogously, for the second part, we need to show that i, > c².:: for this too we can use the lower bound Lemma C.6.
We finally have all the ingredients prove the main theorem of this section:
30

---

Published as a conference paper at ICLR 2025
Theorem C.8. The map in (34), given by
X - xX) = u² (x)
represents a contextual mapping.
Proof. As defined in Def. C.2, contextual mapping must satisfy two conditions. The first one is that
X * (X),,
Vi + j
and
v E.
(61)
This is directly proven by considering Lemma C.5: since I, (strictly) increasing sequence, all its elements are already distinct. The action of the last global shift layer merely translates all these elements by same quantity, but they remain distinct nonetheless.
The second condition for a contextual mapping is given by
(X) # (X'),
Vi,j and vx, L, with X *
(62)
We prove that this holds for (34) by directly considering the difference between two components i.j for different inputs:
q: (x()) x - 25 (X') = (x
catt (hn-2;) =0 = 1 iz ca (h2
(63)
Notice that, due Lemma C.4, have In + 0 and particularly, 1T, -1,1
1%1] 2 8. On the other hand,
in light of the bounds in Lemma C.6, we have that the left-hand side 1 7,1 < 6 Consequently, the two sides can never cancel each other out, and the proof is complete.
D LIPSCHITZNESS OF SIGMOID ATTENTION
In the following, we report the proof for the recovering the Lipschitzness constant associated with SigmoidAttn, as stated in Thm. 3.2.
Letting A = W Wx. and calling Øij o((W Waxz)) and d'y = o' d'WW a, Wazz)), we find that the Jacobian of & in the direction (81,...,5,) for the sample o; is given by:
Jac, = Σ da
A"
AT
6, + Σ
(64)
j=1
We see that this Jacobian is the sum of terms. To control its norm, we can control each norm individually.
The first term, (Ejn dajzzzz A²) 8, of the form U,6; with U, a matrix. Its squared-norm therefore:
Σ
|U,5 < max |
(65)
Hence, its squared spectral norm is bounded by max, |1.112.
We now let o' be bound on n x |a'l: We have:
|U 2 10 Alla
(66)
< d A|ll
|1,511
(67)
j=1
S d || A A|1,81112
(68)
We that if the points 2; have norm < R, then the Jacobian grows at most like R² because is "quadratic" x. However, we see that the quadratic term likely to be mitigated by the of (a..)) term that goes to 0 if aij is large.
31

---

Published as a conference paper at ICLR 2025
The second term, Σj= (d A + dsyIp) 6j. is the sum of two terms. Here, too, we the triangular inequality to control their norm individually. We get:
0,f =110 1
(69)
S |
(70)
where o; RP is the i-th column of oij. and E Raxp, and by summing, letting doe an upper bound on n x (x)]:
Σ
GIE - 108
(71)
So that doo upper bounds the spectral norm of the last term
For the final term, Σ/=1 daaeaa Adj. define = SA We get:
Σ da Aoj Σ (
(72)
Hence, letting M the matrix of entries Mj- = (2i, 85), we see that the previous term is simply x² M so that get the upper bound on the norm of the term:
Σ
a M
12
S 1
(73)
and ||M = Ey(o)) 8,)² s # |> | B1 . giving overall:
Σ MFI d'
|| || ||
(74)
Notice how this quantity matches the one in (68). Finally, summing all together gives:
|Jac|la S 20'2 [ + ooo
(75)
which completes the proof.
Remark: The previous upper bound might not be tight. Indeed, intuitively, if the æ; are large, then the term of should be exponentially small (provided, of course, that Wy and W&zy are not orthogonal), which would even remove the dependency the variance in the sigmoid attention.
E THE BIAS TERM OF SIGMOID ATTENTION
One of the differences between SigmoidAttn and SoftmaxAttn is the normalization constant. In SigmoidAttn, one way to emulate the effect of normalization constant (which links all the elements of the input together and defines distribution over them), include bias term in the definition as proposed in (3).
For input vector RR. the output of the sigmoid with bias b is exp(2;) (z) exp(2:) + exp(-b)
Contrary to the softmax, this output cannot always sum to one because there is no normalization. We therefore seek value for that approximately normalizes (z), i.e., such that Σn od (z); 1. We have
Proposition E.1. Let E R", and take m, M such that for all i, it holds m S 2, S M. Then, the equation Σn (2); = with variable b has single solution by with
-lo(nn 1) M < b* S log((n 1)
32

---

Published as a conference paper at ICLR 2025
Proof. The function o b + Σi= o (z); is smooth and monotonically increasing, and we have (o - 1) - M) - 1 and @0 lo - 1) m) 2 1. This shows the existence of b" as well as the advertised bound on b*.
This suggests using a b of the order of - log(n): in practice we use b = log(n).
We can also look for bias term b, which helps to approximate the softmax function by the sigmoid function.
We assume that softmax provides us with the true distribution p", where p Eine e The goal is to find the bias term b such that sigmoid function with weights over all elements denoted by p. where pi = o. (x):, approximates p". Note that, as mentioned before, p is not necessarily a distribution, i.e. Di=1 pi is not always equal to onc.
In technical terms, we aim to estimate the normalizing factor Z - --E== eac. The existing approaches for estimating Z is compute-xppensiii for high dimensions and requires resampling methods. Also, the optimal value of b would depend on the exact values of z, which is unknown beforehand. Therefore, we propose a more intuitive way to estimate the order of bias but possibly with larger disparity. To distribute the independent masses in SigmoidAttn, we assume that each element has uniform weight for the model apriori, which means that none of the elements of the input vector = has any known importance over the others. In the simplest case when softmax is uniform distribution, we ideally want to have the same order of values for sigmoid as of softmax, which should be Therefore, can write down the following:
Vi
Pi
1 e (2 12
pt
(76)
Requiring that p p" in the case where all the z are 0 gives exp(-b)- = n - i.c. b b= - log(n) for large n. In the case that all the 2; are bounded, | S . < so for some constant M, then b ~ - (M + log(n)) in most cases we do not know M. When the sequence length n is large enough, the constant M loses its importance while in short sequence length, it impacts distributing the weights over elements more. To resolve this issue, we assume that z; are sampled from a standard Gaussian distribution, i.e. 24 ~ NN o2) where = 1. Note that this assumption comes from the fact that 2; in our problem is one of the elements of QK // dyg. which is the sum of dyp random variables. Using Central Limit Theorem, we can assume that 2; is sampled from Gaussian distribution. The idea is to estimate M, such that with high probability, |2:| S M. i.. IP (12gl > M) S e for a desired e. Therefore, we have M
Ideally, we would like to have 1 + -
max M, log(n)). However,
1
o²
(( M)
|
5
Se
(77)
M
M
where the inequality is resulted from Chebychev's inequality. Setting = 1, we have M = VIe. Therefore, the order-optimal value would be b max( log//))o(()) and for long sequence length, b ~ - log. For example, if we want 90% accuracy in our estimation, M RJ 30 = 3. which means b max (3, log(n)). ). Note that this approximation also follows the intuition that as grows, we expect the SigmoidAttn without bias term overestimate the mass on each point, so we need to normalize the mass according to n at each point as well.
On another side, one may be more interested in the gradients of p" and p with respect to 2; to behave similarly. We show that b = - -log(n) still good choice in this scenario. Let us derive the derivative of SigmoidAttn and SoftmaxxAtnn with respect to the input. We note that for any i, both functions can be written as --++2_, where Z_ is the share of normalization factor except element i of z. For SttnmaxAttn, Z_, = Σjei 6 and for SigmoidAttn, Z_ = e Now, we have
e";
e
e Z
(78)
dz, e Z=i
(e + 2_1)
Therefore, we have the following
0p:
pt (1 pi)
(79)
02,
Opi
p:(1 - -p.).
(80)
0z
33

---

Published as a conference paper at ICLR 2025
We can see that if p. pl. then dp: approximates the order of gradients well. In fact, this is the only valid choice even though we have a quadratic term.
Opt So, the previous choice of bias term - (og
de, D5,
Op: 0z, = Qp:
(1 - p.) p; ps)
(81)
0z,
(p:- (p (1- pi))
=
(82)
Which means either p; = p or p; = 1 p.. The first one provides us with b = - log(n) while the second one cannot happen since the nominator of is dependent on while the nominator of 1 pt is independent of Zi.
F DETAILS OF FLASHSIGMOID
This appendix provides details of the FLASHSIGMOID algorithm. We begin by discussing the implementation details of FLASHSIGMOID, which we build as an extension of FLASHATTENTION2, followed by a benchmark of the performance of the involved kernels. We show that the kernels of FLASHSIGMOID provide considerable performance boost in model inference over those of ATTSTTTTNN and a modest performance boost for model training. Further, we demonstrate that the kernel speed boosts also reflect in a considerable performance gain in realistic end-to-end experiments, with an example training vision transformers (Dosovitskiy et al., 2021) on the ImageNet dataset (Deng et al., 2009). Finally, we also provide kernel benchmarking details of FLASHSIGMOID implementation by taking into account ALiBi slopes (Press et al., 2022), which is one of the important components of SigmoidAttn seen in the main text of the paper.
F.1 DETAILS OF FLASHSIGMOOD ALGORITIIM
Softmax vs. Sigmoid Attention: In this subsection, we discuss the implementation details of FLASHSIGMOID algorithm, which is hardware-awe implementation of SigmoidAttn approach. We begin with the expressions of the forward and backward passes of softmax and sigmoid attention mechanisms. Let Q, K, and V represent the query, key, and value tensors. Then, the desired forward and backward pass expressions are reported in Tab. 3. The application of sigmoid and
SOFTMAX
SIGMOID
FORWARD
BACKWARD
FORWARD
BACKWARD
S K- Vd
S Q KT
dV = PT dO
dV PT dO dP d V ds P o P) dP dQ dS K
Vd
P SOFTMAX
dP d0 VT ds P O (dP ROWSUM (dO 0 0))
p= 0 (S)
0 P.V
0 - P.V
dQ dS ds K
dK = Vd ds Q
dK = Vd.dS .Q
Table 3: Description of the forward and backward passes of softmax and sigmoid attention. With C, we denote Hadamard (element-wise) multiplication.
softmax activation functions, as highlighted in orange color in Tab. 3, is the only implementation difference in the forward passes. Similarly, the expressions for the gradients of the preactivation (dS). as highlighted in purple color in the table above, is the only implementation difference in the backward passes. In light of this, we implement the FLASHSIGMOID algorithm an extension of the FIASSHATENTION2 (Dao, 2023) algorithm, which highly optimized hardware-aware implementation of SoftmaxAttn.
Flash Attention in Brief: As pointed at in the main text, the FLASILATTENTION (Dao et al., 2022) and FLASIATTENNION22 (Dao, 2023) algorithms provide hardware- implementations of exact attention mechanism by optimizing for bottlenecks of modern accelerators (Choquette et al., 2021: Choquette, 2023). These GPUs possess massive amounts (e.g.. 80GB) of Bigh-Bandwidth Memory (HBM), which stores large tensors but slow in moving the data to the accelerators. On the other hand, they have smaller amounts (e.g. 20 MB) of SRAM, which is often more than
34

---

Published as a conference paper at ICLR 2025
Algorithm 1 FLASHSIGMOID Forward Pass
1: procedure FORWARD( Q, K, V, Br, Be):
2:
3: inputs: Matrices Q, K, V E Raxd are on HBM of the GPU.
4: inputs: Integers B, and Be are the block size for queries and key-values respectively.
5:
6:
outputs: Matrix O E RRd on HBM of the GPU.
7:
# No need output logsumexp vector L E RD on HBM.
8:
109
9:
Divide Q into T.-
blocks: Q1 Qr. with Q E RB.
10:
Divide K into Tg == n blocks: K1," Kr. with K; E RBxdd
B.
11:
Divide V into T. blocks: V,,",, Vr. with V; E B.xd
12:
Divide O into T, blocks: O1, ,0,, with O; E R.Bxdd
for i 1,- 1,.... TT do
14:
Load block Q: from HBM to SRAM of the GPU,
On chip, initialize O; with zeros: O f 0 0B.xd
16:
# No allocation of either row-sum & E R Br or row-max m. E RR B. on chip.
17:
for j 1.. Te do Load blocks Kj, V, from HBM to SRAM of the GPU.
18:
19:
On chip, evaluate arevactivations: Sūj Q; K. //a
E RBxBB.
20:
On chip, evaluate sigmoid attention: P.j < (S.)).
21:
On chip, update output block: O; < O - P.j. V,-
22:
No need to update and track & and m; vectors.
23:
end for
24:
Store
O, from chip to HBM as the i-th block of O matrix. # No post-processing of O L; blocks on chip. No movement of L; block from chip to HBM.
25:
26:
27:
end for
28:
return matrix O.
29:
end procedure
order magnitude faster for carrying out actual computations using the registers/tensor cores of the GPU. This trade-off between memory size and computation speed across hierarchies results in the attention mechanism computation being bottlenecked by memory accesses between the HBM and the SRAM (Ivanov et al., 2021). Consequently, flash algorithms optimize for memory accesses across the hierarchy of GPU memory types in order to accelerate computation of attention mechanism and its gradients. FLASHSIGMOID no exception to this approach.
Algorithm 1 describes the forward pass and Alg. 2 describes the backward pass of the FLASHSIGMOID algorithm. We highlight in orange color the steps in the forward pass of FLASHSIGMOID that differ from those in FLASHATTENTION by virtue of sigmoid activation. Similarly, we highlight in purple color the differences in the backward pass. Finally, we highlight in blue color the salient points of FLASHSIGMOID that further help minimize bottlenecking factors on modern accelerators.
In FLASHATTENTION and FLASIIATENTTION, the attention mechanism computed by splitting the attention matrix into blocks. Since softmax activation requires row-wise reduction to compute its normalization factor (i.e., the denominator), one needs to properly compute and track such factor across blocks. Moreover, in FLASHATTENTION this normalization factor is stored after being computed in the forward pass, to have easily accessible to further speed-up the backward pass. By contrast, substituting sigmoid to softmax eliminates the need to allocate and move across the GPU memory hierarchy the tensors related to the normalization factor (i.c., moving the logsumexp tensor E R² on HBM in the forward and backward passes). In addition, applying softmax in stable manner requires tracking the row-max variable m; on chip, which instead is not needed for sigmoid activation. This further helps reducing some on-chip operations and lowering register pressure in FLASHSIGMOID.
Fewer Tensor Allocations, Fewer Memory Accesses, Fast-
Moving on to the backward pass (described in Alg. 2), FLASHATTENTION2 requires computing rowsum (dO O), which is needed to backpropagate the gradients of softmax attention outputs
35

---

Published as a conference paper at ICLR 2025
Algorithm 2 FLASHSSIGMOD Backward Pass
1: procedure BACKWARD( Q, K. V. dO, Br. Be):
2:
3: inputs: Matrices Q, K, V, dO E Raxd are on HBM of the GPU.
4: inputs: Integers B, and Be are the block size for queries and key-values respectively.
5: # No need of logsumexp vector L R" to be saved for the backward pass.
6:
7:
outputs: Matrices dQ, dK, dv RRad on HBM of the GPU.
8:
9:
Divide Q into T.-
blocks: Q1 Qr, with Q E BR
10:
Divide K into T. == x blocks: Ki,".. Kr. with K, E BBxdd [ B.
11:
Divide V into T. blocks: V,, Vr. with V; E BR
12:
Divide O into T, blocks: O,.... Or. with V, BRxdd
13:
Divide dO into T2 := B, blocks: dO,,... dOO, with dO, E BBxd
14:
Allocate dQ on HBM and divide into T, blocks: dQ dQr. with dQ; E B,
Allocate dK on HBM and divide into T. blocks: dK, 1, dKr. with dK; E B.xd
16:
Allocate dV on HBM and divide into T. blocks: dV1, dVr. with dV, E RB.xd
17:
# No need to compute rowsum (dO ④ 0) as sigmoid and gradients are pointwise. for 1. Te do
18:
19:
Load blocks K, V, from HBM to SRAM of the GPU. On chip, initialize d, dV, with zeros: d + 0B.xd dV, + 0Bxxd
20:
21:
for = 1..... T, do
22:
Load blocks Q; dO,, dQ, from HBM to SRAM of the GPU.
23:
# No need of movement of blocks rowsum (dO @ O)2 and logsumexp L,.
24:
On chip, evaluate pre-activations: Sij f Q: K /Vd R B.
25:
On chip, evaluate sigmoid attention: P.j + (S
26:
On chip, update gradient of values: dV; +- dV; + P.j dO,. j.
27:
On chip, compute gradients of attention matrix: dP - dO, V. E B..BB
28:
On chip, compute gradients of pre-activations: dS. + P.j (
o
29:
Load query gradient block dQ, from HBM to SRAM, and then on to chip.
30:
Update query gradient block on chip: dQ; + dQ: + vdl v ds.; K,
31:
Store query gradient block dQ; from chip back to HBM.
On chip. update key gradient block: dkj- + dK, vd ds. Q, end for
33:
34:
Store ,Kj,, dV, from chip to HBM as the j-th th blocks of dK, dV matrices respectively. end for
35:
36:
return matrices dQ, dK, dV.
37:
end procedure
to the preactivations. However, since sigmoid activation is applied element-wisc. its gradients also backpropagate across sigmoid element-wise, eliminating the need of the row-sum variable and the movement of its blocks across the memory hierarchy. Another optimization of FLASHÄTTENTION and FLASIIATTENTION2 consists of partially re-computing the forward pass of attention mechanism in the backward pass avoid bottlenecks and speed-up the implementation. To keep the backward pass implementation fast, they require the logsumexp variable to be available and transferred between HBM and SRAM the backward pass. FLASHSIGMOID, being an element- wise activation, eliminates the need of this variable from the backward pass, and consequently, from the entire algorithm. Finally, a major component in our implementation is the usage of GPU-based implementation of the tanh activation. Sigmoid activation is related to Tanh activation via the following relation: a (x) = 0.5 (1 + tanh (0.5- z)). We utilize the fast GPU-implementation of Tanh activation, which trades off some precision for better speed, in order to compute sigmoid activation in both the forward and the backward pass. This provides a considerable speed- in both the forward and backward passes of FLASHSIGMOID, while maintaining parity in performance with naïve implementation of sigmoid attention. Based on these points of modification, we extend FLASHATTENTION2 to obtain FLASHSIGMOID, hardware-aware implementation of SigmoidAtin.
36

---

Published as a conference paper at ICLR 2025
F.2 BENCHMARKING OF FLASHSIGMOID KERNELS
Benchmarking Setup:
Having seen the details of the FLASHSIGMOID algorithm, we next consider the benchmarking of its kernels. For this, we create a small model PyTorch (Paszke et al., 2019) that inputs query, key, and value tensors (all of shape |batch, tokens, heads, features]) and passes these through a number of attention layers. Mimicking the design of vision transformers (ViTB- 16/224) (Dosovitskiy et al., 2021), we set the number of heads and per-head features as 12 and 64, respectively, We set batch size 32, and consider a 10-layer architecture. Then, for the number of tokens sampled from wide range of [64, 78k], compute the forward and backward passes of this model. For these computations, we measure the kernel GPU time using Py TyTorch's 's profiler. We carry experiments both H100 (Choquette, 2023) and A100 (Choquette al., 2021) GPUs.
FlashAttention2 FlashSigmoid (Full)
(Full)
FlashSigmaid (Full)
Kernel GPU Time (ms)
--.--
FiashSigmoid (Causal)
FlashAttention2 (Causal)
Kernel GPU Time (ms)
FlashSigmoid (Causal)
FlashAttention2 (Causal)
1250
1003
1200
200cs
3083
1983
0500
xix
2000
38203
4500
0838
1900
Tokens
Tokens
(a) Inference mode kernels on H100.
(b) Training mode kernels on H100.
Figure 11: On average, for sequence lengths between 664, 78k], the inference mode kernel of FLASHSIGMOID 17.39% faster than PLASIIATTNNTION2 for self-attention and 18.76% for causal attention. The training mode kernels of FLASHSIGMOID are 6.53% faster than FLASHATTENTION for self-attentic and 9.46% for causal attention. Note that inference involves only the forward pass of the model and training involves both the forward and the backward pass of the model.
X00
FlashAttention2 (Full)
FlachAttentian2 (Full)
FlashSigmoid (Full)
1991
FlashSigmoid (Full)
2003
FlashAttention2 (Causal)
FlashAttention2 (Causal)
Kernel GPU Time (ms)
Kernel GPU Time (ms) )
FlashSigmoid (Causal)
FlashSigmoid (Causal)
2009
1883
Tokens
Tokens
(a) Inference mode kernels on A100.
(b) Training mode kernels on A100.
Figure 12: On average, for sequence lengths between 664, 78k], the inference mode kernel of FLASHSIGMOID is 14.33% faster than FLASHATTENTION2 for self-attention and 16.92% for causal attention. The training mode kernels of FLASHSIGMOID are 6.02% faster than FLASHATTENTION2 for self-attention and 5.27% for causal attention. Note that inference involves only the forward pass of the model and training involves both the forward and the backward pass of the model.
Figures 11 and 12 show the GPU time comparisons of kernels in inference mode and training mode of FLASHSIGMOID and FLASHATTENTIONS respectively. We observe that we obtain large average speed-boost for inference and modest average speed-boom for training. Note that
Results:
37

---

Published as a conference paper at ICLR 2025
the speed-ups in all the subsequent figures are obtained by averaging the performances for tokens sampled in the range of [64, 78k].
Next, also show the performance of individual flash kernels of FLASHSIGMOID and FLASHATTENTION2. Note that inference mode involves only the forward pas of the model, while training mode involves both the forward and the backward pass of the model. The forward pass of both these approaches involves one kernel, which we term flash. flash_fwd__ kernel, and the backward pass of both these approaches is made up of three kernels, which we term bwd_ dv, swd_dor_doo and swd_convert_de In code, the real names of these kernels are as follows.
Details of Individual Kernels:
fwd == flash_fwd kernel
pwd_dq_dd dk dv flash flash_bwd_ www_bbpdd_____e_eeeeeddebb l___ddd___voog htt__ooopsseee_peee saap__ssk___paallee kernel
(83)
bwd_dot wwd__dot_d____ o flash_bwd, ______dd_dog lidd_dot_do___e cooo__do___kee o_kenel
bwd__wome bwd_convert dq := flash_bwd_____
_ kernel
Here, we first provide a brief description of the tasks performed by each of these kernels; for a detailed explanation, refer the reader to FLASHATTENTION2 (Dao, 2023) paper and code. The fwd kemel computes the full forward pass of the model as shown in Tab. 3. The bulk of computations of the backward pass happen in the bwd_dq, pwd_dq_dk__g dv kernel, which performs re-computation of attention matrix and reduction of key and value gradient tensors (dK, dV). Again, the exact steps carried out in the backward pass can be checked from Tab. The bwd_convert of query gradient tensor (dQ). Finally, note that the bwd _od_ddtt____ do o kernel in FLASHATTENTIONZ performs the task of computing the rowsum dO O) tensor along with clearing of the accumulators of query gradients (dQ). Although FLASHSIGMOID does not require this row-sum tensor, the clearing of accumulators of query gradients still needed. For this reason, bwd_dot__o____ kernel also appears in the profiling of FLASHSIGMOID.
kernel performs the reduction
Performance of Individual Kernels: Figures 13 and 14 show the performance comparison of each flash kernel in FLASHSIGMOID with the corresponding kernel in FLASIATTENTION2 when tested on an H100 GPU and an A100 GPU respectively. We observe that on both the H100 and A100 GPU architectures, the fwd kernel of FLASHSIGMOID is significantly faster than that of FLASHATTENTION2 and the bwd _d__kk_g______k dk dv kernel of FLASHSIGMOID has a modest average speed boost over FLASHATTENTION
is significantly faster on A 100 GPUs Note that even though the wwv__dot_do___ kernel of FLASHSIGMONT appears to be slower on average H100 GPUs, the kernel time of bwd_ www__do_ddooo (~ 5ms) is negligible compared to that of the main bwd pwd__g__d _g__d_________l ___kk_o dv kernel (~~ 5000ms). Thus, the combined backward pass kernel in FLASHIIGGOID time does not suffer from this slowdown. Finally, note that for FLASHSIGMOID IGOOID and PLASHATTENTION? have identical performance. This is expected, since the task of this kernel to reduce the gradient of the queries dQ, which is common step in both the approaches and is not modified in FLASHIIGGMOI.
The bwd d___ddt_dd______ kernel in FLASHH
_dq
Rasoimamtoo Futl witttaa fuscitartoo Kauap Itantog diceaaas
INNIOuttr TVD
1000 090 Trme 1y01
T0saro
Takoro
fwd: 17.39% faster for self-attention and 18.76% for causal.
bwd dq ldddd__d____dddopo dv: 3.29% faster for
bwd dot_ do__o oo:
bwd convert dq: 0.03% faster for self-attention, 0.02% slower for causal.
2. 2.24% slower for self-attention and 2.17% for causal.
self-attention and
6.97% for causal.
Figure 13: FLASHSIGMOID and FLASHATTENTION kernel comparison H100 GPUs.
F.3 SPEED BOOSTS OF FLASHSIGMOID REALISTIC SETTINGS
In this section, we demonstrate how the performance boosts measured in App. F.2 for the individual kernels of FIASHSIGGOOI contributes to speeding-up realistic runs with end-to-end training.
38

---

Published as a conference paper at ICLR 2025
Rustii to Rusiily Clansh
Bund amo 1ul
Rludd yeu Cluaal
Packiigg oe Siautt
Corce GEU Time
Tocou
Tolors
Takena
fwd: 14.33% faster for self-attention and 16.92% for causal.
bwd _q__k__vv _d__v: 3.50% faster for
bwd_dot_ed _o__do_o: 7.95% faster for self-attention and 8.00% for causal.
bwd_ codd_conver__oe 0.01% faster for self-attention, 0.03% slower for causal.
self-attention and
1.39% for causal.
Figure 14: FLASHSIGMOID and FLASHATTENTION2 kernel comparison on A100 GPUs.
Setup: As a target experiment, we consider training a vision transformer (Dosovitskiy et al., 2021) on the ImageNet dataset (Deng al., 2009). We create vision transformer model variants- one with FLASHATTENTION attention and the other with FLASHSIGMOID attention. We carry out the training of these models with a distributed data-parallel (DDP) setup using PyTorch (Paszke et al.. 2019). We perform two sets of experiments- i. the first performs DDP training on four nodes of H100 GPUs with eight GPUs per node and EFA/RDMA interconnect for the nodes, and ii. the second performs DDP training on four nodes of A100 GPUs with eight GPUs per node. In each set of experiments, we use three different image sizes (64 x 64, 90 x 90, and 100 x 100), along with patch size of to result in different number of tokens for the underlying attention mechanism in the vision transformer model (64 x 64 4096, 90 x 90- 8100, and 100 x 100 10000 tokens). For each of these configurations, than 80%. These considcrations are in order to minimize, if not eliminate, other confounders that can unfairly affect estimation speed-ups in realistic runs. For instance, a low GPU utilization would lead to a larger number of updates, which in turn would incur unnecessary delays, variations, and slow-downg due across-nodes communications.
select batch sizes so that the GPU memory utilization would be greater
The results of the runs on H100 nodes and A100 nodes are shown in Tab. 4 and 5 respectively. There, we show how the kernel GPU times for forward and backward passes vary according the number of tokens considered, and include the wall-clock time of the end-to-end runs as explained above. We observe that the kernel speed-up reflects significantly in the speed-up of inference of the models (during testing) and modestly in the training of the models. We observe ~ 8% speed-up in wall-clock time of inference and 4% speed-up in wall-clock time of training.
Results:
TOKINS
KERKEL GPU TIME COMPARISON
Fuu RUN WALL-CLock TiMe COMPARSSON
KERNEES
Ft. SLMMMD (6S)
Mons
(LASBATTENTION2(0)
FLASSSSGMOttD (5)
4096
PWD
4.98:-0.1
4.17=0.01 ( 16.31%)
INPERENCE
11.17=0.18
10.08=0.18 ( 4.43%)
19.58860.00 .6
18.12:00.44 T.45%)
1563,80=11.0
0.21.08-22227 2.07%)
20. 4610 05
16.70.0000
22%)
INHERSACE
28.31
25, 25.0020011 (-8 06%)
07.63100.11
78.70.00 12 (.5%)
TRAINING
4282,70.22:14
1129.35.54.14 5656)
HWD
31.1745.00
25.49:20.00 (- 18. (0%))
INPERSNCE
38.71=0.19
35.37 (-8 8.00%)
P40 F BWD
117. 13
109.87:00:2 6.62%)
THAINING
0990.72-33.21
0701.41100.77 (- 3.00%)
Table 4: FLASHSIGMOID FLASHATTENTION2 H100 nodes. The kernel GPU time for both the approaches is reported in milliseconds and wall-clock times is reported in seconds per epoch.
Connection of Wall-Clock Time Speed-Up and Kernel Speed-Up: From Tab. 4 and 5, is clear that the speed-up in kernels is larger than that in the wall-clock times of the full runs. In fact, the speed-up in kernels is the upper bound for the speed-up that we would see in wall-clock times. To see why, let us denote by Tam and Ta the total kernel GPU time for softmax attention and sigmoid attention respectively. Then, the kernel speed-up is given by Skemel := 1- To However, in a full run, the total wall clock time also incorporates the time required to load data, time taken by other layers of the underlying models, time required to communicate gradients and other data across GPUs and across nodes, and so on. For our corresponding sigmoid and softmax runs, these extra factors are designed to add, upon expectation, in the same extra time Thus, the wall-clock time speed-up of full run
39

---

Published as a conference paper at ICLR 2025
KRRORI GPU COMPARISON
RUN WartaChoo COMPARSION
KERNELS
PLASHATIENTIONN (NS)
PLASISICHOED (965)
Mone
FLASHAAEETTOOO (S)
PLASUSScMMM (5)
PWD
6.32210.00 02
7.84 03 6.78%)
INFERENCE
18 18.0549:22
18.74 (-1 1.05%))
PAb + IWD
31 11.81::0000 08
10.11140.0000 (-2.10%)
TRAINING
2785. 03:3:33
2769, 44255 (-0.021))
$100
PWD
83.85220009
27. 27.9220.07 (-17.04%)
INFERENCE
47 37.35:20.30
64.06:20.11 (-6.96%)
1WD | BWD
128:19:0 33
119.04=0.12 7.18%5)
TRAINING
7010, 64:14
7254.84t 8.62%)
10000
PWD
51.1720:07
12.1910.06 (-16.6661))))) 96%)
61.61109.39
59.30.00.18 (-7.82%)
194. 194.5000 14
180.3920.15 (-7.17%)
TRAININO
10455 61.88.85
10052.042 (-3.86%)
Table 5: FLASHSIGMOID vs. FLASHATTENTION2 on A100 nodes. The kernel GPU time for both the approaches reported in milliseconds and wall-clock times is reported in seconds per epoch.
with end-to-end training is Swall-clock == +on+7
Since have faster sigmoid kernels, have
To < Tsm- which in turn shows that Swall = To < Te Skemme- This explains the speed boost trends in kernel time versus full run wall-clock time for each setting in Tab. 4 and 5. However, in particular, if model performs attention mechanism over large number of tokens, the attention mechanism, and hence the corresponding kernel time, starts to dominate the other computations the network. In that case, Mathematically, if To Tam > we have: To + T 2 To. + R Tsm- Thus, Skemel
that the wall-clock time speed-boost is closer the kernel speed-boost.
Swall-clock.
thereby making Swall-clock/ Skemel +
FlashAttention22 + ALIBI (Full)
700
FlashAttention2 + ALBII (Full)
Washhiigmols + AUBII (Rull)
FlashSigmoid + (Full)
FlashAttention2 (Causal)
FlashAttention2 (Causal) FlashSigmaid AUBI (Causal)
Kernel GPU GPU Time (ms)
FlashSigmoid ALIBI (Causal)
Kernel GPU Time (ms)
5000
+000
159
1910
2.10
Tokens
Tokens
(a) Inference mode kernels on H100.
(b) Training mode kernels on H100.
Figure 15: On average, for sequence lengths between (64, 78k], the inference mode kernel of FLASHSIGMOID is 17.04% faster than FLASHATTENTION2 for self-attention and 10.87% for causal attention. The training mode kernels of FLASHSIGMOID for self-attention and 4.7 for causal attention. Note that inference involves only the forward pass of the model and training involves both the forward and the backward pass of the model.
sequeha
8.91% faster than FLASHATTENTION
Significance of Wall-Clock Speed-Up of Inference: Although FLASHSIGMOID provides only modest gains during training, the speed-up in inference is significant (> 15% for underlying kerness and 5 - 10% during inference of full runs). We posit that this speed-up in inference is extremely critical well. Contemporary large-scale models, once trained, spend huge portion of the rest their lifetime in inference mode (OpenAI, 2023). Thus, significant performance boosts in inference mode have immense potential for saving resources deployment of large models for inference.
F.4 FLASHSIGMOID WITH ALIBI
is evident from the main text of the paper that improved positional embeddings, like ALiBi (Press et al., 2022), can be crucial for certain tasks and data modalities. Thus, we also provide FLASH- SIGMOID implementation that incorporates ALiBi. We compare the FLASHISIGMOID with ALiBi implementation with the FLASIIATTENTION2 with ALiBi implementation (Dao, 2023). Figures 15 and 16 show the kernel GPU time for the forward and backward pass kernels of FLASHSIGMOID with ALiBi implementation versus FLASHATTENTION2 with ALiBi implementation. Again, observe that FLASHSIGMOID kernels for inference have significant speed-up in wall-clock time over those in FLASHÄTTENTION2 and the kernels for training also have modest wall-clock improvements.
40

---

Published as a conference paper at ICLR 2025
FlashAttention2 ALIBI (Full)
FlachAttentian2 + ALIBI (Full)
1283
FlashSigmoid + ALiBi (Full)
FlashSigmoid ALiBI (Full
Kernel GPU Time (ms)
2003
FlashSigmoid ALIBi (Causal)
FlashAttention2 ALIBi (Causal)
Kernel GPU Time (ms) ms)
FlashAttention2 (Causal)
FlashSigmoid ALiBi (Causa))
13300
3800
<dice
68000
Tokens
Tokens
(a) Inference mode kernels on A100.
(b) Training mode kernels on A
Figure 16: On average, for sequence lengths between [64, 78k]. the inference mode kernel of sequehar FLASHSIGMOID is 12.28% faster than FLASIIATTENTION2 for self-attention and 5.30% for causal attention. The training mode kernels of FLASHSIGMOID are 14.6 14.64% faster than FLASHATTENTION2 for self- aeffaattentonn and 6.80% for causal attention. Note that inference involves only the forward pass of the model and training involves both the forward and the backward pass of the model.
F.5 DIRECTIONS FOR FUTURE WORK ON FLASHSIGMOID
In this section, we discussed FLASHSIGMOID, a hardware-aware implementation of the SigmoidAtan Attn algorithm. Then, we demonstrated via kernel benchmarking and realistic setting runs that FLASH- SIGMOID provides significant gains in inference as well as modest gains in training of models with attention mechanism. this subsection we further discuss additional avenues for improving the implementation of FLASHSIGMOID. and point out some interesting directions for future work.
Optimization of Block Shapes for Different Input and GPU Settings: As stated before, our FLASHSIGMOID implementation builds on FTASHATTENTION2 by adding functionality for forward and backward pass of sigmoid attention place of the standard softmax attention. In particular, results discussed far, inherit directly from FLASHATTENTION the details of optimal block shapes, grid shapes, and other kernel launch parameters, and keep them unchanged in our implementation. For instance, this is the case for the block sizes B,,, B. in Alg. 1 and 2, which are identical in FLASHATTENTION' and FLASHSIGMOID. This choice is dictated by the need to ensure fair comparison between the two implementations, and allows to demonstrate the speed-up of sigmoid attention by minimizing confounders associated with parallel computations different GPU architectures for different input shapes.
for all FLASHSIGMON
Although FLASHSIGMOID kernels lead to speed-ups in inference and training for both H100 and A100 GPUs, observe that the kernel timing speed-ups A100 are not uniform across sequence lengths: for a small subset of these, our kernel provides significantly lower speed-up compared to the overall trend for other sequence lengths. Ideally, the implementation not assume any information on the token count in input, and it is then desirable to have uniform speed- ups across all input lengths. Here, we show that this is achievable by simply updating the block shape information in FLASHSIGMON implementation of ATTENTIINN templated according to block shapes, grid shapes, and other kernel launch parameters. Note that FLASHATTENTION? provides various tailored implementations, optimized for different input shapes (e.g.. different ranges of feature dimension per head), input types (e.g., causal attention vs. self-attention, ALiBi vs. ALiBi in attention, etc.), and GPU types (e.g., A100 vs. H100 via checking shared memory size on GPUs). This is achieved by opportunely selecting the kernel template parameters defining block shapes. grid shapes, and other kernel launch parameters for parallel computation on GPUs. In our case, we create a variant of FLASHSIGMOID, denoted by FLASHSIGMOID where we update the block sizes for query and key tensors from (B,, B.) = (128, 128) of FLASHSIGMOID to (B,, B.) = (128, 64) of FLASHSIGMOID' only for input setting (template with features per head being 64).
attention mechanisms should
to values that are different than those in TTASSHTTENNION
The
Experimentation and Results: For this variant, we perform kernel benchmarking as described in App. F.2, and report the corresponding results in Fig. 17. Comparing the plots for kernel timing
41

---

Published as a conference paper at ICLR 2025
X0X
FlashAttention2 (Full)
FlachAttention2 (Full)
FlashSigmold" (Full)
FlashSigmoid (Full)
2003
FlashAttention2 (Causal)
FlashAtenttin22 (Causal)
Kernel GPU Time (ms)
Kernel GPU Time (ms)
giashSigmmod" (Full)
FlashSigmoid (Full)
2003
1400
0.0
104
309
3800
3883
Tokens
Tokens
(a) Inference mode kernels on A100.
(b) Training mode kernels on A100.
Figure 17: On average, for sequence lengths between [64, 78k], the inference mode kernel of FLASHSIGMOID' is 14.82% faster than FLASHATTENTION2 for self-attention and 18.02% for causal attention. The training mode kernels of FLASHSIGMOID" are 6.18% faster than FLASHATTENTION2 for self-attention and 76% for causal attention. Note that inference involves only the forward pass of the model and training involves both the forward and the backward pass of the model.
with FLASHSIGMOID plots from Fig. 12, observe that FLASHSIGMOID' not only provides more uniform inference and training kernel speed-up on all sequence lengths, but also improves the average of these speed-ups across all lengths. To further bolster our observations, Tab. shows the inference mode and training mode kernel speed-ups for subset of sequence lengths under consideration. This experiment indicates that it is possible obtain higher and more uniform speed-ups in kernel timings across a wide range of tokens by investigating optimal block shape, grid shape, and other kernel launch parameters for each input setting and GPU type. We leave this optimization for future work.
KERNEL GPU TIME COMPARISON
KERNELS
FLASHATTENTION2 (MS)
FLASHSIGMOID (MS)
FLASHSIGMO1 (MS)
4096
FWD
8.32-10.02
7.84:10.03 (-5.9%%)
7.2640.02 (-13.213333
31.81#0.08
31.11=0.08 (-2.19%)
30.62
33.65:20.09
27.9210.07 (-17.04%)
28.5410.07 (- 5.5%
8100
FWD BWD
128.18:10.13
119.04:10.12 (- 7.13%)
119.85:10.13
8.81%)
10000
FWD
51.1740.07
42.49:10.6 (-16.96%)
43.53:10.09 (- 15.11..8 3.32%)
194.54:10.14
180.59:10.15 (-7.177 17%%
181.9710017
6.87%)
131.19+0.12
(-6. (-6.55-63 -6.53%
116.75+00.00
40%)
16384
PWD BWD
494.6540.25
482.0840.22 23 ( 54.%)
474.5240.25 (-4.48%)
Table 6: FLASHATTENTION2 vs. FLASHSIGMOID vs. FLASHSIGMOID' on A100 nodes. The kernel GPU time for all three approaches are reported in milliseconds. We observe that FLASHSIGMOID" provides better and more uniform speed-ups across all example tokens.
G EXPERIMENTS
G.1 EXTRA ABLATIONS
G.1.1 THE EFFECT OF MULTIPLICATIVE SEQUENCE LENGTII NORMALIZATION
Wortsman al. (2023a) notes that models trained with sigmoid ReLU attention require scaling by the sequence length, -2 "a(QQ" / We ablate this by comparing the scaled solution to the one we propose in App. E. We also generalize the variant proposed (Wortsman et al., 2023a) variadic sequence lengths such that works with auto-regressive (AR) training, for example for
42

---

Published as a conference paper at ICLR 2025
Figure 20: n -05 normalization.
Figure 18: b
In
Figure 19: n- normalization.
3:
1
1
1
1
0
0 0
0.5
0.5
1 0.33
O
o (QK")
(dx))V
(84)
0.33 aa
0.33 -
1
Cousal Mask M
We repeat the experiment from Fig. 5, using ALiBi positional embeddings for all trials. We apply {1, 0.5) AR normalization proposed in (84). While there is observable difference in terms /K/Q/K//// / dyk ViI, we find that the train NLL is slightly worse for both normalized variants (Fig. 19 and 20) in comparison to the b = n variant in Fig. 18.
of the attention norm, ||o(Q?
G.1.2 ATTENTION BIAS STABILITY ABLATION
To validate the stabilizing effects of attention bias we repeat the experiment from Fig. and keeping
all of the same hyper-parameters, while enabling QK norm and LayerScale (initialized at 10-*). We
train with a range of constant bias offsets, b {-15, 10, -6, -4 1} and visualize the results
below in Fig. 21. We observe systematic increase in stability (and lower SigmoidAttn NLL) for
Kum
Figure 21: Attention bias ablation.
values less than -1 up till -10, after which the -15 plot shows an over-regularizing effect with decreased performance.
G.2
VISION
G.2.1 TEST IMAGENETIK Top-1%
Fig. 22 reports the test linear probe results for the ViT-B/16 BYOL (Grill et al., 2020; Busbridge al., 2023), ViT-B/16 SimCLR (Chen al., 2020; Zhai al., 2023a) and the finetuned performance
43

---

Published as a conference paper at ICLR 2025
BYOI, VIT-R/16
SimCLR VIE-//6
MAE VIT-1116 Finctune
VIT-R/16
Text Top. 19
Sottraa
Epoch
Epoch
Epoch
Epoch
Figure 22: ImageNetlk test top-1% for SoftmaxAttn vs. SigmoidAttn using models from Fig. 2.
for the ViT-L/16 MAE (He et al., 2022) and the test top- % results for for ViT-B/16 supervised model (Dosovitskiy et al., 2021). Across these wide range of SSL and supervised learning tasks, trained with contrastive (SimCLR), EMA distillation (BYOL) and reconstructive objectives (MAE), find that Sitn not only matches the training dynamics (Fig. 2), but also the linear probe and finetuned performance of the baseline SottmaxAttn.
G.2.2 LAYERSCALE FREE SIGMOID ATTENTION
Train
Test
Figure 23: A competitive SigmoidAttn ViT-B/16 model can be learned with- out LayerScale or QK norm using large initial learnable scalar tempera- ture L = 10 and bias b = -10 (sim- ilar to SigLIP (Zhai et al., 2023b)): a (e'(QKT / + b)V, {b, t) TR.
1-10, b
b-10, h-110
NLL
ReLU
Sigmoid
This regularizes the model, move the temperature to a learnable regime. The t 10, -10 curve makes no progress in train NLL or test top- for ~25 epochs (near max LR), but ultimately outperforms baselines.
it must
Tog-1 %
Epects
Epocts
While Fig. 23 demonstrates the possibility of learning Sigmoid Attn without LayerScale, it involves task specific tuning of {t, 6}. We also explored gating attention from learning (through simple multiply by zero) for ~25 epochs and were able to match the = 10, -10 training curves from above. However, opted for the LayerScale method due to its simplicity.
G.2.3 SIGMOID ATTENTION VS. ATTENTION RELAXATIONS
Figure 24: Supervised ViT-B/16 ImageNetlk classification. We contrast SigmoidAttn and SoftmaxAttn against (a) linear attention with no activation: QKT / and (b) fast atten- tion via positive orthogonal random features, used in Performer (Choromanski et al., 2021). SigmoidAttn, like SoftmaxAttn, differs from attention relaxations like Performer which uses low-rank representations of the attention ma- trix. Sigmoid maintains performance parity with SottmaxAttn, while outperforming other efficient attention variants.
Supervised VIT-B/16
80.0
Test Top- 1%
40.0
20.0
Sigmoid
Performer
Linear
0.0
0
50
100
150
200
250
300
Epoch
44

---

Published as a conference paper at ICLR 2025
G.2.4 HYPER-PARAMETE PARAMETERS
Table 7: SigmoidAttn SimCLR and BYOL ViT-B/16 hyperparameters.
Parameler
SimCLR
BYOL
Attention bias
None
None
LayerScale Init
10
10
QK Norm
Yes
Yes
Pos Embed
SinCos
Learnablc
Freeze Patchar
Yes MocoV3 (Chen al., 2021)
No tronc_normal (. (.02)
Weight
Normalization
LayerNorm
LayerNorm
schedule
Single Cycle Cosine
Single Cycle Creine 40 Epuchs
LR warmup
10 Eguchs
Min LR
1 x 10 300 Epochs
1 x 10-s
Training duration
600 Epochs
Optimizer
Adam Lincar (0.9, 0.95)
Adam
Optimizer scaling rule
Lincar
Adam (f1. (2)
(0.9, 0.95)
Base LR
x 10
x 10
Base hatch size
256
256
Total batch
4096
4096
Base teacher momentum
0.996
Weight decuy
0.1
0.3
Weight decay skip bias
Yes bC16
Yes
Numarical precision
bC16
Stochastic depth
0.0
0.2
Angmentation stack
SimCLR (Chen al., 2020) 0.5 (Chen et al., 2021)
DINO multicrp (Caron al., 2021)
Color Jitter Scaling
Table 8: SigmoidAAtn Supervised ViT-B/16 and MAE ViT-L/16 hyperparameters.
Parameter
Supervised
MAE
Attention bias
None
In
LayerScale Init
10"
10
QK Norm
Yes
Yes
Pos Embed
Learnable
Learnable
Archilecture
VYT-B/16
ViT-L/16
Mask Ratio
0.75
Freeze Patcher
No
No
Weight init
mocre__ (. (.02)
Leuuc__oonee (.02)
Normalization
LayerNorm
LayerNorm
I.R schedule
Single Cycle Cosine
Single Cycle Cosine 40 Epochs 0.0 400 Epochs
LR warmup
20 Epocch
Min LR
x 10
Training duration
300 Epochs
Optimizer
Adam
Adam W Linear (0.9. 0.95)
Optimizer scaling rule
Linear (0.9. 0.95)
Base Adam (2),5))
Base LR Base batch size
1 x 10
1.5 x 10
256
256
Total batch size
4096
4096
Weight decay
0.3
0.05
Weight decay skip bias
Yes
Yes
Numerical precision
b
bf16
Stochastic depth
0.28
0.0
Augmentation stack
RandAug (Cubuk et al., 2020)
RRC HFLIP
G.3 LANGUAGE MODEL
G.3.1 PARPEE-PR
Tab. 9 shows the hyper- parameters for the final comparison. MuP-simple (Wortsman et al., 2023b) is used, where the peak learning rate is set to le-2. Weight decay is decoupled, following Loshchilov & Hutter (2017). In addition, to confirm that applying QK-Norm does not hurt the bascline, we show training parity with and without QK-Norm in App. G.3.I.
45

---

Published as a conference paper at ICLR 2025
Table 9: Training details for the Llama-style IB and 7B LM training.
Pammeter
1B Value
78 Value
Parans
18
7B 4096
Comest Length
2048
Total Tokens
300B
IT
Bonch sioe
4M tokess
4M tokens
1R Schedile
Cusine
Cosine
LR Warnup Sieps
5000
5000
Peak IR
Ie-3
Io-2
Finul ER
10P5 of pask
10% of peak
Optimizer
Adlan
Adlam
Optimizer momentum
0.9, D.95
0.9, 1:95
Weight dreoy
le-4
lo-4
Goudient dippiin
1.0
1.0
Position mooding
Al.:ii
A1i8i
Q/K Norm
Applied
Applied
Norm
ReMSNeem (hang Senneich, 2019)
RMSNomm (hhng Sennrich, 2019)
Norm stracture
Pre-norm
Hehrid-oom
Num byers
24
92
Num heads
Ilidden dim
2018
4096
375
DJ Norm
350
2.0
325
Gradierk Nom
Vain.4
3.00
2.75
2.25
0.5
2.00
50000
100000
150000
200000
250000
300000
50000
350000
250000
300000
Steps
Figure 25: Left: IB SoftmaxAttn LLM training with and without QK Norm converging to the same loss. Right: 85M and 1B LLM training using Sigmoid (n 4096) without hybrid-norm. Smooth training loss curves, but gradient norm shows spikes. Hybrid norm mitigates gradient spikes.
4.00
Sig
5.75
350
3.5
Tean RL
325
3.0
3.00
2.75
2.50
2.0
$9000
100080
150000
200000
250003
300000
Figure 26: Left: 85M training using SigmoidAttn and SoftmaxAttn (n= 4096) without hybrid- norm. Training loss matches. Right: 1B training using SigmoidAtta (n 4096). Higher sequence length with a larger model shows a slightly different loss curve. Hybrid norm resolves the issues enabling stable training up 7B (Table 2 and Fig. 2).
G.3.2 GRADIENT NORM
While Sttn based LM using aforementation pareme-arseeters has smooth loss curve, do see more gradient norm fluctuations. See App. G.3.1, where spikes larger than 0.5 are not visible in the Softmax Attn equivalent.
46

---

Published as a conference paper at ICLR 2025
Table 10: Different norm structure ablations for Aitn with IB language-moddelngg
SEQ.
ATTN.
Pos:
ARC
ARC CHALL.
HELLA- SWAG
WINO-
LAMBADA OPENAT
TRIVIAQA (I-SHOT)
WEBQS (I-SHOT)
MODEL
NORM
PIQA
SCIQ
AVG
ExcoD.
SOFT.
28
No
ALIB:
PRE
62.2
26.8
42.4
59.0
72.3
88.1
58.4
19.9
15.4
49.4
SOFT.
2K
No
ROPE
PRE
64.5
30.4
43.9
61.0
71.9
88.7
59.3
21.1
15.0
50.6
SIGM.
2K
YES
ALIB]
PRE
62.8
28.8
42.5
$9.7
70.3
88.6
59.7
19.1
13.8
49.5
YES
ROPE
17.3
12.8
SIGM.
2K
No
ROPE
PRE
59.3
26.5
39.4
55.1
69.4
88.2
59.2
11.7
7.5
46.0
SOPT.
2K
No
ALIB]
HYBRID 64.9
30.5
43.3
61.9
71.6
88.4
60.9
23.6
12.8
50.9
SIGM.
2K
YES
ALIB1
HYBRID
60.5
26.9
42.2
59.2
70.8
89.6
$7.9
17.7
13.4
48.7
SIGM.
No
ALIBI
HYBRID
62.8
28.2
42.0
59.1
70.3
88.7
59.8
18.6
15.1
49.4
29.3
20.4
15.6
SOFT.
4K
No
ALIB]
PRE
62.6
27.7
42.4
58.6
71.1
88.2
58.6
18.9
14.7
49.2
SIGM.
4x
YES
ALIB1
PRE
60.5
27.3
41.3
57.8
70.5
87.0
$7.6
18.9
12.6
48.2
SOFT.
4K
No
ROPE
HYBRID
64.1
27.2
43.3
61.4
71.2
88.5
60.0
21.4
15.3
50.3
SIGM.
4x
No
ROPE
ILYBRID
63.3
27.1
43.4
61.3
70.4
88.2
57.5
20.5
14.8
49.6
60.7
70.8
20.9
SIGM.
4K
No
ALIB]
HYBRID
62.4
28.9
43.5
60.8
71.3
89.6
59.2
20.2
14.3
50.0
G.3.3
NORM STRUCTURE
Due to the slight performance difference observed at 4096 context length when using SigmoidAtta versus Softmax and marginally lower downstream results, evaluated various norm struc- tures to address potential instabilities (see Tab. 10). Some of these structures replace the required attention bias (in this case, column 'Attn. Bias' is 'No'). All use QK-norm with RMSNorm (Zhang & Sennrich, 2019), without LayerScale. We examined pre-norm and hybrid-nore (where do both pre-norm and normalization of the output of the attention layer following Xiong al. (2020)): norm (QKQ dykk) V). Post-norm, which normalizes the combined residual data stream, norm norm CammooQQKK' dgg V), is omitted from our analysis did not train stably for SigmoidAttn.
G.4 AUTOMATIC SPEECH RECOGNITION
G.4.1 TRAINING DETAILS
All acoustic models are fed 80 channel log-mel filterbanks with a 25ms sliding window strided by 10ms.
The transformer-Be based encoder model has 255M parameters: ID convolution kernel 7 and stride followed by CAPE positional embedding if it is used and 36 transformer blocks with pre-LayerNomm, an embedding dimension of 768, 4 heads, 3072 units in the MLP layers. The model is trained with CTC loss and character vocabulary, including apostrophe (*). In additional experiments, we vary the depth to 12 and 24 layers, and change pre-LayerNorm to post-LayerNorm.
We implemented our own conformer-based encoder model, also trained with CTC loss and character vocabulary. The conformer model has 104M parameters and consists of ID convolution of kernel 7 and stride followed by 16 conformer blocks with embedding dimension of 512, 4 heads, 2048 units in the MLP layers. Variational noise is not used and RoPE is used as a relative positional embedding instead of relative sinusoidal positional embedding.
For all models, SpecAugment (Park et al., 2019) is used for augmentation with 2 frequency masks (max width 30) and time masks (max width 50, ratio 0.1). All models are trained with dynamic batching and mixed precision with BF16. Models are trained with different configurations of optimizers and hyperparameters to have diverse coverage of use-cases. We first optimize every configuration for SoftmaxAttn and then change only attention to the introduced configuration of SitnoidAttn while all other parameters are kept the same. Detailed configurations are shown in Table 11. We train models until the greedy WER stops improving on the validation sets (dev-clean. dev-other) and report final test sets (test-clean, test-other) greedy WER without integration of any external language model.
47

---

Published as a conference paper at ICLR 2025
Table 11: Training details for the ASR models on LibriSpeech 100h (LS-100) and LibriSpeech 960h (LS-960) for transformers and conformers.
Parameter
Transformer LS-960
Conformer LS-960
Transformer LS-100
Transformer LS-100
Params
255M
104M
255M/ 170M 85M
255M
LayerNorm
post
post
Dropout
0.1
0.1
0.3
0.3
Layer drop
0.1
0.0
Training steps
400k
400k
400k
500k
Batch size
3.56h
4.44h
1.1h
1.1h
LR schedule
step-wise
step-wise
step-wise
step-wise
Apecmnngentt start
Ok
10k
Ok
Ok
LR Warmup Steps
64k
64k
64k
Peak LR
lc-3
2e-3
0.03
LR start decay
250k
250k
200k
330k
LR decay step
50k
50k
30k
50k
Optimizer
Adam
Adam
Adagrad
Adagrad
Optimizer momentum
0.9, 0.999
0.9,
Weight decay
le-6
le-6
Gradient clipping
1.0
0.5
1.0
1.0
Position eacoding
CAPE ALiBi RoPE
RoPE
CAPE
CAPE ALiBi / RoPE
Q/K Norm SoftmaxAttu
Not Applied
Not Applied
Not Applied
Not Applied
Q/K Norm SigmoidAttu
Applicd
Applied
Not Applied
Applicd
Num layers
36
16
36/ 24 12
36
Num heads
Table 12: Word error rate (%) LibriSpeech dev/test sets and TED-LIUM v3 (Hernandez al., 2018) ("TED", joint validation and test sets with split according to audio duration) for pre-LayerNorm transformer (255M/ 170M / 85M params) with CAPE and with either SoftmaxAttn SigmoidAttn (w/ LayerScale, w/o QK norm, w/ b = log n)) trained on LibriSpeech 100h data (average duration 10-15s). Hyper-parameters can be found in Table 11.
ATTN
A LAYERS
DEV-CLEAN
TEST-CLEAN
DEV-OTHER
TEST-OTHHRR
TED 0-10s
TEn 10-20s
TEn 20-30s
TED 30s+
SOFTMAX
36
6.7
7.1
20.0
20.4
26.4
22.4
23.3
21.8
SIGMUID
36
7.0
7.3
20.3
20.5
26.2
23.4
23.6
21.8
8
36
6.8
20.3
SOFTMAX
24
6.4
6.8
20.2
20.5
25.4
22.1
23.3
21.8
SIGMOID
24
7.1
7.3
21.0
21.3
26.6
23.3
24.0
22.0
SOFTMAX
12
8.2
8.7
25.0
25.4
29.0
25.6
27.1
27.4
8.3
25.2
26.3
8.7
24.7
For the bias term b = - loog n in SigmoidAttn, we do not use max sequence length as language model experiments. Instead, for every audio sample we its own duration a bias terms resulting into non-trainable bias vector for the minibatch. For experiments with sequence normalization, we also use the max sequence length in the minibatch but rather the ground truth sample duration to properly normalize encoder attention.
To evaluate behaviour for length generalization we use TED-LIUM v3 dataset Hernandez al. (2018) as its validation and test sets have longer audio duration than LibriSpeech: LibriSpeech has in average 10-15s duration, while TED-LIUM there are audio longer than 30s (the max duration of LibriSpeech). To perform evaluation on TED-LIUM v3, we combine together validation and test sets of TED-LIUM (we don't use them for training and hyper-parameters search and just perform final evaluation) and split them into 4 datasets according to the duration: 0- 10-20s, 20-30s, and 30s+.
For positional embeddings we use not only CAPE, but change it to AliBi or RoPE. As ALiBi was originally introduced for the decoder only models and there is official adoption of it yet for the encoder models (without causal masking), we follow the best practices found https: / //iclr-blogpost github. 1 2C24 /blog//l////// awiaa-miimmmm of nonsymmetric ALiBi with different slopes instead of symmetric version used by (Lee et al., 2022).
MSee discussion in http::////gghhh..ccooooooooooo// http:p////.tiooooooooee/ereeeeer with inaar biasss/
ssues/5.
48

---

Published as a conference paper at ICLR 2025
Table 13: Word error rate (%) on LibriSpeech dev/test sets for post-LayerNorm transformer (255M)
with either moftmaxAttn (w/o QK norm) or SigmoidAttn (by default w/ LayerScale, w/ QK norm, w/ b log n) trained on LibriSpeech 100h data. Syper-parameters can be found in Table 11.
ATTN
DEV-CLEAN
TEST-CLEAN
DEV-OTHER
TEST-OTHER
SOFTMAX
CAPE
6.4
6.5
18.4
18.2
6.1
6.3
18.2
18.1
SIGMOID
8.0
8.4
22.7
22.7
QK NORM
7.9
27.6
LAYERSCALE
UNSTABLE, GRADIENT NORM AND LOSS SPIKES
QK NORM- LAYERSCALE
6.5
6.9
19.9
20.1
SIGMOID (b 10, LEARNABLE)
8.7
9.4
23,5
24.0
ROPE
6.9
SIGMOID
6.8
7.1
20.8
20.8
SIGMOID (b 10, LEARNABLE)
SOFTMAX
ALIBI
6.4
6.9
18.3
18.3
SIGMOID
6.9
7.2
21.1
SIGMOID (b = 10, LEARNABLE)
6.8
7.1
20.4
20.5
Table 14: Word error rate (%) on LibriSpeech dev/test sets and TED-LIUM v3 (Hernandez et al., 2018) ("TED", joint validation and test sets with split according to audio duration) for conformer (104M) with RoPE and with either SoftmaxAAtn or SigmoidAttn (w/ LayerScale, w/ QK norm, w/ b = log n) trained on LibriSpeech 960h data (average duration is 10- 10-15s). Hyper-parameters can be found in Table 11.
ATTN SOITMAX SIGMOID
DEV CLEAN
TEST CLEAN
DEV OTHER
TEST OTHER
TED 0 10s
TED 10 20s
TED 20 30s
TED 305+
2.2
2,5
5.4
5.6
13.0
11.1
13.2
7.1
2.3
2.5
5.8
13.5
10.3
13.3
10.2
100
2.4
2.7
5.8
12.9
14.1
G.4.2 RESULTS AND ABLATIONS
Initial investigation on post-LayerNorm and pre-LayerNorm transformers on both LibriSpeech 100h and 960h revealed that SigmoidAtin without any bias is unstable resulting in huge and frequent gradient norm and training loss spikes throughout the training which in turn result in spikes of validation and test WER, see Figure 27. Neither LayerScale nor QK norm were able to stabilize the training, though we did not observe any model divergence.
Further experiments with bias term in the SigmoidAttn definition for Post-LayerNorm transformers on LibriSpeech 100h reveal that training is now stable (only few marginal spikes in gradient norm occur, while train loss is smooth all the time). However, both LayerScale and QK norm restrict model capacity thus not matching Softmax Moreover, some combination of them is needed for the stable training, though w/o both of them we got the best performance for Sigmoid Attn (still behind SoftmaxAttn), see Table 13. We believe, further adaptation and deeper investigation is needed for SigmoidAttn and Layerromo,,,, though recent advances in machine learning do not use post-LayerNNorm models due to high training instability even for SoftmaxAttn.
Switching to pre-LayerNorm transformers and varying the depth of the models lead to stable training with SigmoidAttn and bias term b - log with few (2-5 times) spikes in the gradient norm and smooth loss. In this case, SigmoidAttn matches results for Softmax Attn and they both generalize to TED-LIUM data similarly, see Table 12. If the bias term is removed, SigmoidA can still match SoftmaxAttn but large spikes in gradient norm and loss can occur.
Finally, experiment with conformer model, in Table 14. Again, we found that bias term b = - log n stabilizes training. The learnable = -10 works though we see significant gradient norm spikes while the train loss remains smooth. Besides, b log generalizes well to longer sequences while learnable = -10 fails to do so with RoPE for conformer. Overall, SigmoidAttn able to match SoftmaxAttn having stable training with b log
In experiments with different variants of bias term for SigmoidAttn, the bias = log n is found to be the most stable (only few marginal gradient norm spikes are observed with the train loss
49

---

Published as a conference paper at ICLR 2025
LayseLLayeNomm
Norer
080
sggadd
Train Loss
12.5
dev-ather WER
Steps
Steps
Figure 27: ASR Transformer model (255M) training with Post-LayerNorm (left) and pre-LayerNorm (right) on LibriSpeech 960h with AigmoidAttn (w/ bias term, b = O, w/o QK norm, w/ LayerScale) or with max Huge gradient norms and training loss spikes are observed for SigmoidAttn which can result in worse final model performance hence models for SigmoidAtts are unstable.
alabi
rope
Thain Loss
cape
Figure 28: ASR Transformer model (255M) training with pre-LayerNorm on LibriSpeech 960h with SigmoidAtta (w/ bias term, b = - -llog n, w/ QK norm, w/ LayerScale) and different positional embeddings CAPE, RoPE, ALiBi. The bias b is able to stabilize SigmoidAttn training: smooth training loss and only marginal rare spikes in gradient norms are observed.
60-
40-
50000 100000 150000 100000 75000 30000
Steps
being smooth) and it provides similar performance as SoftmaxAttn in most settings. The source of instability is coming from the larger attention output norms (80k for CAPE, 40k for RoPE and 20k for AliBi while being 200 for SoftmaxAttm) This happens due to high attention weight of every token which can be biased towards zero with a bias term in SigmoidAttn definition. Preliminary results to connect this to the local attention property needed at the beginning of the training for stable training failed, local attention did not converge well at all (it is deactivated after some initial training).
50

---

Published as a conference paper at ICLR 2025
Table 15: Word error rate (%) LibriSpeech dev/test sets and TED-LIUM v3 (Hernandez et al., 2018) ("TED", joint either SoftmaxAttn or SigmoidAttn (LayerScale and QK norm are used with b log n) trained on LibriSpeech 960h data (mean duration is 10-15s). Hlyper-parameten are in App. G.4. H-Norm corresponds to hybrid-norm, no LS-Attn corresponds to removing the LayerScale from the attention outputs, and no LS corresponds to removing the LayerScale from both the attention and MLP outputs.
joint validation and test sets split according to duration) for transformer (255M params) with
ATTN
PE
DOV-CLIAN
TEST-CLEAN
DES-OTIER
TEST-OOUHE
TED 0-10s
TED
10-206
TED
10-305
TED 3054
SUFTMAX
2.2
2.3
5.7
12.4
10.5
11.9
SHAMODD
CARE
2.4
5.2
10.3
12.3
lng[mitasses
2.3
10.6
2/0
2.3
10.1
12.3
S1GWORD, Ing/minasees m]
SIGHODD (R (-KOEM)), so QK-N0EN, LS-ATTX
ROPE
23
33
23
22
5.0
5.0
531
5.0
13.8
12.3
15.1
10.2
12.1
12.3
10.5
10.4
aRIWOR KORM). LSS-ATT
233
5.0
5.1
13.0
10.2
17.4
11.8
ALIRI
10.5
SIGMODD, log(maxssrakk
2.0
23
5.2
5.2
12.3
10.5
11.9
10.2
To fully benefit from the improved throughput of FLASHSIGMOID, for the bias term - log Attn we experimented with configuration when the maximum audio duration in the minibatch is used resulting into non-trainable bias scalar which changes between minibatches as we use dynamic batching. Comparison between the bias vector with per sample own duration normalization and the bias scalar maximum duration in the minibatch is shown in Table 15: final model performance is similar and stability is same (only 2-3 minor spikes in CAPE for gradient norms are observed). Thus, per batch maximum audio duration can be used with final configuration.
log n as the
We also experimented with hybrid-norm (see Table 16) to check if is able to stabilize the attention magnitudes as well as gradient norms. We did ablation with configuration similar to Table 16 with the following changes: LayerScale after attention is replaced to LayerNorm, only RoPE is used for positional embedding; we either keep or remove QK-norm and we either keep or remove LayerScale MIP part of transformer block.
First, for all variants we observe that training loss and gradient norms are smooth without any spikes during training while we see abnormally large attention activations compared to prior experiments. Second, while observe that QK-norm or its removal behave similarly, the LayerScale on top of MLP output is necessary to get performance on par with Softmax or with Sigmoid Attn with bias term.
G.5 SIMPLE EXPERIMENTS
G.5.1 K-SUMMATION PROBLEM DEFINITION
Here we look at synthetic, simple task in order to investigate the behavior of softmax and sigmoid attention activations. The problem chosen is to minimize the MSE loss of IR - target function. In the first half of each input are samples from N 1) distribution, and the second half is a k-hot binary vector indicating which values in the first half sum.
The results presented here are for the = 40 problem with various values for k. Where transformer used, the transformer is a single layer to aid visualization. In all cases (unless noted otherwise), the optimizer is Adam with constant learning rate of 0.001, and the training data is continuously generated to preclude over-fitting.
A few examples for 10 (not drawn from (0, 1)) are shown below. Inputs in the second half of
the input are show in orange only as a visual aid.
1
2
3
4
5
0
0
5 0
1
1
2
3
4
5
0
0
1
4 6
8
1
2
0
1
0
4 3
2
0
2
2
2
1
1
0
1
0
+ 4
51

---

Published as a conference paper at ICLR 2025
G.5.2 COMPARISON TO SOFTMAX
In Figure 29, the performance of three architectures on the k-summation problem as increases. The sigmoid activated transformer has similar scaling to the softmax activation.
MLP RaUU
0.07
transformer softmax(A)
transformer -igmoid(A-parrm(-44)
0.06
thaa loss (MSE)
0.05
0.04
0.03
0.02
0.01
2
8
Figure 29: Final loss is shown after training convergence k-summation problem complexity increases. The ReLU MLP has hidden layers (900, 300) for 307k parameters, while the transformer has an embedding dimension of 120, 8 heads, and an MLP ratio of 4, giving 187k parameters. The SigmoidAtun is applied after learned offset initialized -4, A+param (-4).
G.5.3 ATTENTION EVOLUTION
In Figures 30 and 31, forty samples are used to monitor the single head, single layer post-activation attention matrix as training progresses. Figure 30, the distribution of values is visualized over time; note the sigmoid attention is more variable but reaches comparable values at convergence. The main difference at convergence that the sigmoid has fewer high magnitude values than softmax indicating more distributed attention.
9.10
0.04
004
190044 prvatee
0.00
0.01
002
0.0
02
0.4
66
0.8
80
02
10
Taaiing Tozzion
Softmax
Sigmoid
Figure 30: The cost-activation attention evolves during training on the l: = 1, n = 40 summation problem. The model has head to simplify the visualization. Forty repeated test samples are used.
In Figure 31, metrics on the post-activation attention matrices used and show comparable behavior the first half training. In the second half training, the SigmoidAtta can be seen to reduce in norm and in sparsity. (see following discussion of Figure 32 for further insights).
In Figure 32, we see post-activation attention values for eight samples at training progresses. The most notable difference between the activations is, that the end of training, the Sigmoid is
52

---

Published as a conference paper at ICLR 2025
THITs sertrep
1.0
Exna actmes
0.6
0.5
0.4
Norm
Hoyer Sparsity
Figure 31: Metrics on the post-activation attention evolve during training on the n 40 summation problem. The model has one head to simplify the visualization. Quartiles and mean from 40 repeated test samples are shown. On the right, the Hoyer Sparsity (Hurley Rickard, 2009) is used
measure the change in sparsity training progresses: Hoyer
V
Σ (V7 1)-1
VE, 5
less sparse in the (N, 1) self-attention in the upper-left quadrant. We can see that softmax tends to produce sparser values (as designed to) while sigmoid controls the magnitude and location of peak attention independently, leading less sparse attention at the end of training.
25
55
1
sample indices
2
3
4
5
0.0
0.1
0.2
0.3
0.4
0.5
0.6
0.7
0.8
0.9
1.0
training fraction
Figure 32: For samples, the post-activation attentions is visualized as training progresses on the k 1, n2 40 summation problem. The model has one head to simplify the visualization. The attention shown in pairs for each sample with softmax attention is in black and sigmoid blue. A 2 x 2 block structure is evident in both cases, resulting from each halve of the input containing different information.
G.5.4 PAIR REPEAT PROBLEM
We define synthetic task of identifying if the first two symbols in sequence repeat. The symbols, 8; below come from fixed vocabulary of size and the repeat location (when present) is uniformly distributed in the sequence.
f(s,, $1, $2. .... SN) {f-
{1,
if 1/ (50, 81) (8n, 1n41),
otherwise
53

---

Published as a conference paper at ICLR 2025
A simple two layer transformer is trained on this problem. The model has embedding dimension of 160, MLP ratio of 4, QK norm, and layers with eight heads. The results for different model architectures shown in Figure Figure 33. The maximum input length is 22, K 9, shorter lengths are padding with value K, and the training set only contains lengths 14 and 15. cosine learning rate schedule with 5% linear warmup and maximum learning rate of 1e-3 is used with the Adam optimizer.
In this result, distribution cases. From shorter runs, we estimate that the softmax network would fit the training with 4 4- more data. Our conjecture is that the two layer transformer more easily learns the pair finding task with sigmoid because softmax is biased to focus on single values, though it is unclear why multiple heads are not able compensate for this proposed cause in the softmax case.
the sigmoid activation has higher data efficiency and similar fall-off in the out of
MLP RelU
1.0
transfomeer softmax(A)
0.9
accuracy
0.7
0.6
0.5
problem length
Figure 33: Validation accuracy for out of distribution sequence lengths is shows after 5M samples of training; trained lengths are shown with vertical lines. Quartiles and means are shown from six trials. The MLP has two hidden layers, ReLU activation, and similar number of parameters. The sigmoid transformer has a learned offset initialized to -4.
G.6 COMPUTATIONAL COMPLEXITY OF SIGMOID AND SOFTMAX.
Table 16: Forward floating operations per token per attention head. 7exx and dhead are the context length and head dimension respectively. A measures the compute difference between sigmoid and softmax. accounts for causal (c (nax 1)/2n values from the IB LLM results are Nexx = 2048, dread = 64. Sigmoid and softmax share the same number of floating operations (softmax: max-subtraction (2), exponentiation, summation, division; Sigmoid: bias-add, hign-flip exponentiation, addition, division). Remaining differences are due implementation details, and are subleading (~~ 1%) compared to other attention operations like computing attention logits L (shown below). This analysis precludes hardware aware improvements (Section 4).
1/2), standard (c = attention. Typical
= QK"
Softmax (L)
a (L b)
Expression
eeex dnead
cenax
5 cena
G.7 PRACTITIONER' GUIDE
In Table 17 summarize recommended settings for practitioners who aim to SigmoidAtta for training in their respective domains learning scenarios. While each setting has fully enumerated hyper-parameters listed in the Appendix, highlight some sane SigmoidAttn choices below.
Stabilizing larger models beyond sequence length 2048: We propose a non-learned scalar bias to mitigate large attention norms with SigmoidAttn (Appendix E), but observe instabilities
54

---

Published as a conference paper at ICLR 2025
Table 17: Simplified recipe for different domains and tasks with Sigmoid Attention. S/C/A/R refers to using any of SinCos, CAPE, ALiBi or RoPE positional encoding methods.
DOMAIN
OSIECTIVE
MODEL SIZE
Pos EMBED
QK NORM
LAYERSCALE
SIGMOID BIAS
NORM STRATEGY
SUPERVISED
PRE-NORM
LEARNABLE
PREE-NORM
SINCOS
PRR-NORM
MAE
304M
LEARNABLE
YES
YES
YES
PRE NORM
ASR
SUPERVISED (CTC) (CTC)
100M 250M
S/C/A/R
YES
YES
YES
PRE NORM
100M-250M
ROPE
No
HVDRID-NORA
AR LANGUAGE
NEXI-TOKEN (c-2x SEQ LEN)
1B
ALIBI
YES
No
YES
Par: NORM
NEXT-TDKEN (c-2x sRp LEN)
1R
AI.81
YES
No
No
HYBRID-NORM
NEXT TOKEN (>2K SEQ LBN)
1B
ALIDI
YES
No
YES No
TIYRRI NORM
NEXT-TOKEN (>2K
ALIDI
YES
No
YES /oo
HYBRID-
at sequence length n 4096 for autoregressive language modeling (Section 5.5). Hybrid-norm (without learnable affine parameters) resolves these instabilities (Table 10). Hybrid-norm differs
Post-norm
Hybrid-norm + (a(m / V)
norm x + o (QK") [dgk)))
from post-norm, which normalizes the combined residual data stream and attention block output. Hybrid-norm is used in models such as Grok-1 (xai-org, 2024) and frameworks such Praxis (Google, 2024) under the normalization strategy "primer_hybrid" When both SoftmaxAthn and observe similar kernel speedup times highlighted in Section 4. However, with LayerNorm (Lei Ba et al., 2016) only for SigmoidAttn, a token length of 10,000 needed to achieve a performance gain of 5.04% for full self-attention and token length of 1024 is needed to achieve performance gain of 8.36% for causal ceff-attention on H H100 GPUs. For LayerNorm (with and without affine terms), we summarize approximate regimes for positive throughput gains in Table 18.
BigmoidAttn use hybrid-norm,
FLASHSIGNOID
FLASHATTENTIONZ COMPARISON
ATTENTION TYPE
A
11000
AFFINE PROJECTION
No AFFINE PROJECTION
AFFINE PROJECTION
No AFINE PROJECTION
FULL
16384 (5.22% 1)
12544 (5.08% t)
10000 (4.82% ?)
10000 (5.04% f)
12544 (8.18%
5184 (4.14% f)
2048 (7.65% t)
1024 (8.36% f)
Table 18: FLASHSIGMOID along with LayerNorm vs. FLASHATTENTION2 A100 GPUs. Based on benchmarking on set randomly sampled tokens from the range [64, 60000], we report the token T after which FLASHSIGMOID with normalization consistently outperforms ATTEAATTNN,, along with the total CUDA time speed-up averaged over subsequent tokens (T > T*).
H CONTRIBUTIONS
All authors contributed to writing this paper, designing the experiments, discussing results each stage of the project.
Preliminary work Preliminary viability of SigmoidAttn done Jason Ramapuram.
Universal Function Approximation Proof of UFA (Section 3.1 and App. C) sculpted by Federico Danieli.
Lipschitzness of Sigmoid Attention Lipschitzness analysis (Section 3.2 and App. D) molded by Pierre Ablin. Empirical regularity analysis (Figure 7) done by Jason Ramapuram.
FlashSigmoid Implementation and analysis driven by Eeshan Dhekane in collaboration with Jagrit Digani (Section 4 and App. F).
55

---

Published as a conference paper at ICLR 2025
Bias Analysis Theoretical grounding for bias (Appendix E) done by Amitis Shidani in discussion with Pierre Ablin.
Language Modeling Results All large scale language model pretraining and evaluation (Section 5.5 and App. G.3) driven by Floris Weers.
Stability Analysis QK norm (Figure 9), LayerScale (Figure 8) and bias (Figure 21) ablations crafted by Dan Busbridge using Attention Simulator. Attention Simulator written by Jason Ramapuram and used to validate norm growth (Figures 3 to 6, 19 and 20).
ASR Results All ASR experiments (Section 5.4) and ablations (Appendix G.4.2) are conducted by Tatiana Likhomanenko in discussions with Jason Ramapuram and Zijin Gu. Baseline ASR models code written by Zijin Gu and Tatiana Likhomanenko. Baseline models are optimized by Zijin Gu be close state-of-the-an results.
Vision Results All vision experiments (Sections 5.2 and 5.3) and ablations (Figures 10, 23 and 24 and App. G.2.1) conducted and written by Jason Ramapuram.
Simple Experiments Simple experiments to compare Sigmoid to SoftmaxAtin, including visualizing attention evolution and simple sequence length generalization analysis (Appendix G.5) conducted by Russ Webb.
56


Reviews:
[{"review_id": "04RbabL75U", "reviewer": "Reviewer_6HF1", "summary": "The paper provides theoretical and empirical analysis for sigmoid attention. The theoretical contribution is the proof that transformers with sigmoid attention are universal function approximators. Authors also empirically linked the training instability of original sigmoid attention to large  attention norm, and proposed  a few options to improve it (switch to Alibi, or adding attention bias. They also did ablation study of the effect of layerScale and QK normalization on training stability. Since sigmoid attention is theoretically more expensive than regular soft-max, authors proposed new Flash-sigmoid  - efficient implementation for GPU. Combining multiple techniques they demonstrated that sigmoid attention can potentialy replace the  classical softmax attention on multiple tasks without loss in accuracy or speed.", "questions": "Very solid experimental study, but I am still not convinced with paper conclusion: \"Our findings establish sigmoid attention as a viable alternative to [soft-max]. \n1) is it worth to  switch from original softmax attention to sigmoid attention? \n2) will training Sigmoid-Attention be more stable than original Softmax Attention?", "limitations": "", "rating": 6, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "The paper is very detailed analysis of sigmoid attention vs softmax attention. The sigmoid attention is not very original, but authors did a solid job by exploring its theoretcial foundations and execution solid ablation study to decide if sigmoid attention is viable replacement for  soft-max attention . \n \nThe first theoretical contribution - the proof that sigmoid attention can be used as universal approximation to continuous function- is executed well, but has limited novelty, and not very interesting . For me the most interesting theoretical  contribution was the analysis of regularity for sigmoid attention. Btw, I think  the proof of Theorem 2 is missing the key assumption for reqularity is that  ||Wq, Wk\"|| and ||Wv|| should be limited. So we have to  make special measure during training to control or normalize  them. \n\nThe experimental part (Section 5) is very good. Most interesting parts \n1.  observation that training of sigmoid-attention based transformers is very sensitive to initial norm of |sigma (QK)\"|v| . \n2. demonstarion that combination of layer-scale with  QK norm can help to  stabilize training in some cases\n3. very solid ablation study (CV, ASR, LM) \n \nI also liked that authors spent time and effort to rewrite Flash-attention to support Flash-sigmoid and its description in Appendix F. This clearly helped to speed-up whole model (Step-time in Table 3) \n\nThe paper is well writen and it is easy to read.", "weaknesses": "The main weakness of the paper that it doesn't answer main question  \"Why should we switch from original softmax attention to sigmoid attention\"?: \n- will Sigmoid-Attention  be more stable than original Softmax Attention? \n- will it help with long context? \nOn the positive  side, thanks to flash-sigmoid we can see some speed-up for LLM inference (Table 3) \n\nMore details: \nThe paper started with explanation that original \"softmax in SoftmaxAttn is not without limitations. For instance, the softmax function can sometimes lead to a concentration of attention on just a few features (Yang et al., 2018; Ganea et al., 2019), potentially neglecting other informative aspects of the input data.\" As far as I remember Yang's 2018 paper was about soft-max attention work as low-rank factorization, and Ganea's 2019 was about last soft-max layer.  How this is related to sigmoid-attention? \nThe paper says that the second reason to explore sigmoid-attention was \" applying SoftmaxAttn requires performing a row-wise reduction along the length of the input sequence, which in the case of efficient attention kernels (Dao et al., 2022; Dao, 2023), slows down computations.\" You will still need gather operation along previous sequence to sum  V vectors scaled with attention weights over previous keys, so it's not clear how using  sigmoid attention can do noticable difference in this case (see e.g.  \"Online normalizer calculation for softmax\" by M. Milakov, 2018 for detailed analysis)\n \nThe first theoretical contribution -- the proof that sigmoid attention can be used as universal approximation to continuous function--  has limited novelty. This is mostly  re-write of original Yun's 2020 proof.  The main  addition is  the proof that sigmloid-based transformers can do contextual mapping (Appendix C2)  . The observation in C1 that sigmoid can approxiamte Heaviside step function when lambda --> infty is  obvious.  \n\nThe analysis of computational complexity is based on computing flops. But the real complexity should take into consideration memory access (read/write ) instead. \n\nA few comments related to the experimental part: \n-  it looks like the proposed remedy ( bias b=-10  only partially helps to control the  increase in the amplitude of |act * V} (graph 6) \n- ASR: sigmoid (RoPE with bias) performs badly on long audio ( Table 2) \n- LLM: Sigmoid (ALibi) underperform vs Softmax (Alibi) . No results for Sigmoid (RoPE)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "atJ9cjLKfG", "reviewer": "Reviewer_sBpo", "summary": "This paper revisits sigmoid activations in attention mechanism and point out that the main problem with sigmoid attention is that of large initial attention norms. Then they propose solutions with in-depth theoretical and empirical analysis. Also, authors introduce FLASHSIGMOID, a hardware-aware and memory-efficient implementation yielding 17% inference kernel speed-up over FLASHATTENTION2. Furthermore, authors demonstrate that performance across language, vision, and speech are comparable with softmax attention.", "questions": "1. It would be helpful if the authors could make Fig. 2 clearer.\n\n2. It would be better if the authors conducted experiments on large-scale language models (around 7B parameters).", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 2, "strengths": "**Originality:** This paper provides in-depth mathematical analysis of sigmoid attention specially in Universal Approximation Property and Regularity. Also, authors identify stabilization of large initial attention norms during  the early stages of training as a crucial factor for sigmoid attention.\n\n**Quality:** This paper provides detailed mathematical proofs and very extensive experiments and ablation study on sigmoid attention including supervised image classification, self-supervised image representation learning as well as automatic speech recognition and auto-regresive language modeling. The performance is comparable with softmax attention. Furthermore, I appreciated the implementation of FLASHSIGMOID for efficiency and speedup\n\n**Clarity:** The paper is written very well and easy to follow. \n\n**Significance:** Challenging softmax and understanding how activation works in attention mechanism is appreciated. Softmax attention is computationally cost and authors revisit sigmoid activation function to speedup inference time. I believe this paper will inspire the community to rethink attention mechanism.", "weaknesses": "1. In Sec 3.2, the authors state that the Lipschitz constant provides insight into the robustness of the network and the ease of optimizing it. They then present a theorem stating that, in $\\mathbb{R}^2$, 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. However, this is inconsistent with the experiments presented by the authors in Fig. 2, Fig. 3, and Fig. 4.\n\n2. Table 2 shows that sigmoid attention causes unstable training, whereas softmax does not. I believe a theoretical analysis of training stability is necessary.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ZYla5NcYv7", "reviewer": "Reviewer_q17K", "summary": "This paper gives an in-depth analysis of sigmoid attention for use on transformer architectures. It shows theoretically that sigmoid attention based transformers have the universal approximation property and that sigmoid attention has a better local Lipshitz constant than softmax based attention, namely that the former has Lipshitz constant that scales quadratically and that the latter has Lipshitz constant that scales exponentially. The authors conduct a variety of ablations and experiments giving an in depth analysis of when and where sigmoid attention should be used. Furthermore, they develop a hardware aware based implementation on sigmoid attention termed FlashSigmoid that has a much faster inference than the well known FlashAttention2.", "questions": "1). 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?\n\n2). I am very interested in your regularity results in section 3.2. I thank you for providing an in depth proof in appendix D. I found this result rather paradoxical though. Your proof shows that the Lipshitz constant of sigmoid attention scales quadratically, see line 159, and it is known that the Lipshitz constant of softmax scales exponentially, see line 162. This implies that sigmoid attention has much better Jacobian regularity and softmax has very bad Jacobian regularity. Wouldn't we see this in training in that it should therefore be very difficult to train transformers with softmax activation as the backpropagated weights would explode as they grow exponentially? However, we don't see this in practice. Most transformers use softmax as an activation and have no real problem with training. For example, most ViTs use softmax and their models can converge. Yet looking at figure 2 for example, I don't see this. Could you comment on this? I wonder if other architectural mechanisms like skip connections is in some sense negating the effect of the exponential scaling of the Lipshitz constant of softmax attention and that is why we still see such transformers being able to train well. Are the other parts of the transformer architecture somehow mitigating the bad Lipshitz scaling that softmax attention has?\n\n3). I noticed in various experiments, see figure 2, table 2, table 3, that sigmoid attention is comparable to softmax attention. However, since sigmoid attention has a Lipshitz constant that has much better regularity than that of softmax (quadratic over exponential) shouldn't we see sigmoid attention doing much better than softmax?\n\n4). You mention in your abstract, see line 023 to 024, that previous attempts at sigmoid attention have not been able to achieve results on par with softmax. However, in the main paper I don't see you comparing to any of those previous attempts. Could you cite those previous attempts for me so I can go and check those papers? Did you compare your methods with those papers? I think this would help your case as it would show just how your methodology does in comparison to those previous attempts you mention in the abstract. The paper is very dense so if this is somewhere in the appendix and I have missed it then I do apologize.\n\n5). In section 5.5 on Autoregressive large language modelling, you mention in line 508 that a slight disparity between sigmoid and softmax is observed at 1B scale with sequence length 4096 which you are able to address using a hybrid-norm by adding an additional normalization layer. I checked appendix G.3 but was not able to follow the explanation. Could you explain why you need extra normalization and in what way it helps the efficiency of sigmoid attention? As the Jacobian of sigmoid attention has quadratic regularity I would expect that much less normalization should be needed for a transformer architecture employing such an attention layer and rather you would need it on the transformer using the softmax based attention. Also, what is the parameter increase by adding an extra normalization layer and does this add any practical disadvantages to sigmoid attention for such tasks?", "limitations": "", "rating": 5, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "**Originality:** The paper gives a good in depth analysis of the use of the sigmoid function as an activation for attention and provides an original perspective of how it should be used. I applaud the authors for taking the time to give such a detailed analysis.\n\n**Quality:** The quality of the paper is good clearly showing that the authors have taken the time to give the reader a clear understanding of sigmoid based attention mechanisms in transformer architectures. Furthermore, they undertake a variety of experiments giving the reader insights into how sigmoid attention performs on practical tasks.\n\n**Clarity:** The paper is well written and easy to follow with lots of details provided in the appendix. \n\n**Significance:** I believe the real significance of this paper is in understanding how sigmoid attention compares to softmax attention and when and where it should be used, which I believe is something that is not found in current literature. Furthermore, the authors do develop a hardware aware implementation of sigmoid attention that does much better than the current state of the art for softmax attention.", "weaknesses": "**Novelty:** The main issue I have with the paper is its novelty. Although the authors give an in depth analysis of sigmoid attention I don't see it adding value to the community as their results often show that it only does comparable to softmax. The main novelty of the paper I feel is that they develop FlashSigmoid but I feel this is more of an engineering feat and is not enough for a paper to be accepted into ICLR. I applaud the authors for their in depth analysis and their various ablations but I am still questioning whether there is real quality in the paper. Please see my questions below. 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. Furthermore, the authors provide worst case Jacobian bounds which I found very interesting. 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. This suggest that the gradients of softmax attention should explode causing issues with model convergence, yet I don't see this in their experiments. Softmax often performs on par with sigmoid. Thus I don't see how this theoretical analysis transfers into a useful statement for the reader to gain insight into transformers using sigmoid attention over softmax. I think it would be helpful if the authors could tell me in a few sentences what is novel about sigmoid attention apart from FlashSigmoid?\n\n**Experiments:** I found some of the experimental results in the main paper a bit confusing as the authors tend to over do it with ablation after ablation. It would have been much better for the authors to put most of the ablations in the appendix and keep the main experimental results in the main paper. I felt this clouds the other experiments. For example, it might be fine to just put 6 of the graphs in the section 5.1 in the paper and then put the other 6 in the appendix. Similarly, in figure 7 and 8 I feel it would be enough to just put one of them, say layer scale in the main paper, and put qk_norm results in the appendix. This would free up space for you to really talk about in what way you see sigmoid attention yielding a better attention mechanism for the community to use.\n\nAlso in figure 2 I can't really see any difference between softmax and sigmoid. In figure 11 the authors compare sigmoid with layer scale and QK norm with other activations and I notice some of the other activations like GeLU and ReLU do on par with sigmoid on the vision tasks. Does this mean for vision tasks activations like ReLU and GeLU are just as good as softmax so therefore sigmoid is not offering any benefit? \n\nMy overall feeling is that this paper would make a great journal paper as it gives a very in depth analysis of just one activation and its use for attention based mechanisms. However, I will be willing to change my mind if the authors could answer my questions below.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)