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

## Abstract

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 1 1 1 Code is available at https://github.com/apple/ml-sigmoid-attention . . 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.

## 1 Introduction

The success of modern machine learning can be largely attributed to the attention mechanism (Bahdanau et al., ; Vaswani et al., ). Attention uses a sequence-to-sequence (seq-to-seq) map to build context-aware token representations.
Classically, attention relies on the softmax function ($\mathrm{SoftmaxAttn}$) to recover token representations as data-dependent convex combinations of values.

Despite its widespread use and effectiveness, softmax in $\mathrm{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., ; Ganea et al., ), potentially neglecting other informative aspects of the input data. Moreover, applying $\mathrm{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., ; Dao, ), 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 naïve sigmoid attention ($\mathrm{SigmoidAttn}$) is that of large initial attention norms and propose solutions to alleviate it. Our contributions are as follows:

- (1)
We prove $\mathrm{SigmoidAttn}$ is a universal function approximator on seq-to-seq tasks (Sec.˜3.1).
- (2)
We analyze $\mathrm{SigmoidAttn}$’s regularity and provide its worst-case Jacobian bound (Sec.˜3.2).
- (3)
We extend FlashAttention2 (Dao et al., ; Dao, ) with the sigmoid kernel, reducing kernel inference wall-clock time by up to 17% and real world inference by up to 8% (Sec.˜4).
- (4)
We show that $\mathrm{SigmoidAttn}$ matches $\mathrm{SoftmaxAttn}$ in various tasks and domains (Sec.˜5).

## 2 Sigmoid Attention

Let $\bm{X}\in\mathbb{R}^{n\times d}$ be the input sequence of $n$ vectors, where each vector has dimension $d$. We define three learnable weight matrices $\bm{W}_{q}\in\mathbb{R}^{d\times d_{qk}}$, $\bm{W}_{k}\in\mathbb{R}^{d\times d_{qk}}$, and $\bm{W}_{v}\in\mathbb{R}^{d\times d_{v}}$, which are used to compute the queries $\bm{Q}\in\mathbb{R}^{n\times d_{qk}}$, keys $\bm{K}\in\mathbb{R}^{n\times d_{qk}}$, and values $\bm{V}\in\mathbb{R}^{n\times d_{v}}$ as follows:

$$ $\bm{Q}=\bm{X}\bm{W}_{q},\quad\bm{K}=\bm{X}\bm{W}_{k},\quad\text{and}\quad\bm{V}=\bm{X}\bm{W}_{v}.$ (1) $$

Self-attention (Bahdanau et al., ; Vaswani et al., ) can be compactly written as

$$ $\mathrm{SoftmaxAttn}(\bm{X})=\operatorname*{Softmax}(\bm{Q}\bm{K}^{T}/\sqrt{d_{qk}})\bm{V},$ (2) $$

where the $\operatorname*{Softmax}$ function normalizes each row of the input matrix.
We replace the $\operatorname*{Softmax}$ with

Here, $\sigma$ is applied element-wise to the input matrix in eq.˜3.
The activation function $\sigma$ has a hyper-parameter $b\in\mathbb{R}$. In App.˜E, we discuss an intuitive way to choose the order-optimal bias term, resulting in $b=-\log(n)$.
This choice of $b$ allows us to make sense of $\mathrm{SigmoidAttn}$ for any sequence length.
Indeed, letting $(\bm{y}_{1},\dots,\bm{y}_{n})=\mathrm{SigmoidAttn}(\bm{X})$ be the output sequence, we have

$$ $\bm{y}_{i}=\sum_{j=1}^{n}\frac{\exp(\langle\bm{W}_{q}\bm{x}_{i},\bm{W}_{k}\bm{x}_{j}\rangle)}{\exp(\langle\bm{W}_{q}\bm{x}_{i},\bm{W}_{k}\bm{x}_{j}\rangle)+n}\bm{W}_{v}\bm{x}_{j}\xrightarrow[n\to+\infty]{}\int\exp(\langle\bm{W}_{q}\bm{x}_{i},\bm{W}_{k}\bm{x}\rangle)\bm{W}_{v}\bm{x}d\mu(\bm{x}),$ (4) $$

where
$\mu=\frac{1}{n}\sum_{j=1}^{n}\delta_{\bm{x}_{j}}$ is the empirical measure corresponding to $\bm{X}$.
Notably, eq.˜4 still makes sense in the infinite length limit, where the measure $\mu$ is not a sum of Diracs. Wortsman et al. () do not use a bias, and propose a $n^{-1}$ normalization for various attention activations, such as sigmoid and ReLU, but leave the reason as an open question.
Our variable bias has a similar effect in the large $n$ limit, and we posit that recovering a finite output limit as $n$ increases is the why it works in practice.

A multi-head version of eq.˜3 is obtained by combining the outputs of several $\mathrm{SigmoidAttn}$, as follows:

$$ $\left[\mathrm{SigmoidAttn}_{1}({\bm{X}}),\dots,\mathrm{SigmoidAttn}_{h}({\bm{X}})\right]{\bm{W}}_{o},$ (5) $$

for a learnable output weight matrix ${\bm{W}}_{o}\in\mathbb{R}^{hd_{v}\times d}$, where $h$ denotes the number of *heads*.

## 3 Theoretical Properties of Sigmoid Attention

We analyze $\mathrm{SigmoidAttn}$, with two objectives: (1) showing that a transformer architecture remains a universal function approximator when $\mathrm{SigmoidAttn}$ replaces $\mathrm{SoftmaxAttn}$, and (2) recovering a measure of regularity of $\mathrm{SigmoidAttn}$ by computing its Lipschitz constant.

### 3.1 Are Transformers with Sigmoid Attention Universal Approximators?

Yun et al. () demonstrate that classical transformers can approximate continuous sequence-to-sequence functions to arbitrary precision, a property known as the *Universal Approximation Property* (UAP). UAP is highly desirable as it provides proof of an architecture’s generalizability and representation capability.
As $\mathrm{SigmoidAttn}$ 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.

###### Theorem 3.1 (UAP for $\mathrm{SigmoidAttn}$ ) .

We denote with $\mathcal{T}^{h,d_{v},r}_{\sigma}$ the class of transformer networks obtainable by combining an arbitrary number of $\mathrm{SigmoidAttn}$ layers (each of $h$ heads of dimension $d_{v}$) followed by FFN layers of hidden dimension $r$.
For any given continuous, permutation-equivariant function $f:\Omega\subset\mathbb{R}^{n\times d}\to\mathbb{R}^{n\times d}$ with compact support $\Omega$, and for any arbitrarily small error $\varepsilon$, there exists a transformer network $g\in\mathcal{T}_{\sigma}^{4,1,4}$ such that

$$ $\left(\int_{\Omega}\|f({\bm{X}})-g({\bm{X}})\|^{p}_{p}d{\bm{X}}\right)\leq\varepsilon,\qquad\text{for}\quad 1\leq p<\infty.$ (6) $$

Theorem˜3.1 is the exact counterpart of (Yun et al., , 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, we present an overview of the main adaptations required to prove Thm.˜3.1 for $\mathrm{SigmoidAttn}$, with further details in Sec.˜C.1 and C.2.

###### Theorem 3.1 (UAP for SigmoidAttn \mathrm{SigmoidAttn} ) .

##### Sigmoid Attention layers can implement contextual mappings:

A key step in proving Thm.˜3.1 is showing that, even with $\mathrm{SigmoidAttn}$, a sequence of transformer blocks can implement a *Contextual Mapping* (Yun et al., , Def. 3.1). A contextual mapping characterizes a function that maps each input sequence element to an output *uniquely* dependent on the *whole* sequence. This property allows a 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-sequence function.

In Yun et al. (), the contextual mapping is assembled by modifying individual transformer blocks: each block is tuned to react to a specific input token. By stacking a sequence of these blocks, a transformer can be turned into an accumulator, mapping a given input token sequence to a unique global index. This outcome is achieved via a *selective shift layer* (Yun et al., , App. B.5):

$$ $\Psi({\bm{X}};b,b^{\prime})_{i,1}\coloneqq\begin{cases}\max_{k}{\bm{X}}_{k,1}-\min_{k}{\bm{X}}_{k,1}&\text{if}\quad b<{\bm{X}}_{i,1}<b^{\prime}\\ 0&\text{otherwise},\end{cases}$ (7) $$

and can be approximated using classic attention.
Although $\mathrm{SigmoidAttn}$ cannot directly approximate eq.˜7, our accumulator definition relies on an equivalent selective shift operation:

$$ $\Psi_{\sigma}({\bm{X}};b,b^{\prime})_{i,1}\coloneqq\begin{cases}\sum_{k:{\bm{X}}_{k,1}>b^{\prime}}{\bm{X}}_{k,1}&\text{if}\quad b<{\bm{X}}_{i,1}<b^{\prime}\\ 0&\text{otherwise},\end{cases}$ (8) $$

which can be approximated by $\mathrm{SigmoidAttn}$ (described in Sec.˜C.1). In Sec.˜C.2.4, we show that eq.˜8 shares similar properties with eq.˜7, allowing us to use the original proof framework in Yun et al. () and demonstrate that UAP holds in our case as well.

Our proof is largely equivalent to that in Yun et al. (), with two relevant differences: to approximate eq.˜8, we require $\mathrm{SigmoidAttn}$ with at least four heads and shifts included in both query and key definitions. In contrast, $\mathrm{SoftmaxAttn}$ requires at least two heads to approximate eq.˜7, with shifts only in the query definition. However, this is primarily a 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 et al. ().

### 3.2 Regularity of Sigmoid Attention

As with any layer in a neural network, the regularity of $\mathrm{SigmoidAttn}$ 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 $\phi$ is to compute its *Lipschitz constant* over a set $\mathcal{X}$, that is a constant $C>0$ such that for all $\bm{X},\bm{Y}\in\mathcal{X}$, it holds $\|\phi(\bm{X})-\phi(\bm{Y})\|\leq C\|\bm{X}-\bm{Y}\|$, where $\|\cdot\|$ is the standard Frobenius norm.
The *local* Lipschitz constant is the spectral norm of the Jacobian of $\phi$ at $\bm{X}$.
The two are related: the Lipschitz constant of $\phi$ over $\mathcal{X}$ is the greatest local Lipschitz constant for all $\bm{X}\in\mathcal{X}$.
We turn to the theorem giving the regularity of $\mathrm{SigmoidAttn}$:

###### Theorem 3.2 .

Define $A=\{\langle\bm{W}_{q}\bm{x}_{i}\bm{W}_{k}\bm{x}_{j}\rangle|,\kern 5.0pti,j\in\{1,\dots,n\}\}\subset\mathbb{R}$ the set of attention weights, and the scaled activation norms $\sigma_{\infty}=n\times\sup_{u\in A}|\sigma(u)|$ and $\sigma^{\prime}_{\infty}=n\times\sup_{u\in A}|\sigma^{\prime}(u)|$.
Then, the Jacobian of $\mathrm{SigmoidAttn}$ at $\bm{X}=(\bm{x}_{1},\dots,\bm{x}_{n})$ has a spectral norm of at most:

$$ $\|\bm{W}_{v}\|_{2}\left(\sigma_{\infty}+2\sigma^{\prime}_{\infty}\|\bm{W}_{q}^{T}\bm{W}_{k}\|_{2}\left(\frac{1}{n}\sum_{i=1}^{n}\|\bm{x}_{i}\|_{2}^{2}\right)\right).$ (9) $$

The proof is found in App.˜D.
In $\mathrm{SigmoidAttn}$, if we assume that the attention weights $\langle\bm{W}_{q}\bm{x}_{i},\bm{W}_{k}\bm{x}_{j}\rangle$ are all bounded by a constant $\mu$ — this is true, e.g., if the activations are bounded — we get $\sigma_{\infty}\leq\exp(\mu)$ and $\sigma^{\prime}_{\infty}\leq\exp(\mu)$ 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 $\bm{x}_{i}$, while classical results for the study of attention all rely on the largest value of $\|\bm{x}_{i}\|^{2}_{2}$ (Kim et al., ; Castin et al., ).
This is another consequence of the simplicity of sigmoid attention and is due to the removal of the normalizing constant in $\mathrm{SoftmaxAttn}$.
Our result implies that if all $\bm{x}_{i}$ are within a ball of radius $R$ then the Lipschitz constant of $\mathrm{SigmoidAttn}$ grows at most like $R^{2}$, but it is stronger since we can apply this to unbounded distributions $\bm{x}_{i}$; it matters only that the second moment is bounded.
This result contrasts sharply with the bounds obtained for $\mathrm{SoftmaxAttn}$: Castin et al. (, Thm. 3.4.) show that there exists a sequence $\bm{X}=(\bm{x}_{1},\dots,\bm{x}_{n})$ with $\|\bm{x}_{i}\|_{2}\leq R$ for all $i$ such that the spectral norm of the Jacobian of $\mathrm{Attn}$ at $\bm{X}$ is at least $cR^{2}\exp(cR^{2})$ for some constant $c>0$.
On the other hand, our bound scales in $R^{2}$: this means that the local Lipschitz constant of $\mathrm{SigmoidAttn}$ is much lower than the worst local Lipschitz constant of $\mathrm{SoftmaxAttn}$.

###### Theorem 3.2 .

### 3.3 Computational Complexity of Sigmoid and Softmax.

**Table 1: Forward floating operations per token per attention head. $n_{\text{ctx}}$ and $d_{\text{head}}$ are the context length and head dimension respectively. $\Delta$ measures the compute difference between sigmoid and softmax as a multiple of the floating operations for computing the attention logits. $c$ accounts for causal ($c=(n_{\text{ctx}}+1)/2n_{\text{ctx}}\sim 1/2$), or standard ($c=1$) attention. Typical values are taken from the 1B LLM results ($n_{\text{ctx}}=2048$, $d_{\text{head}}=64$). The difference in floating operations between Sigmoid and Softmax attention mechanisms is subleading ($\sim 1\%$) compared to other operations in the attention mechanism like computing attention logits $\bm{L}$ (shown below), and the attention matrix $\times$ values operation. This analysis precludes hardware aware improvements (Section˜4).**
|  | $\bm{L}=\bm{Q}\bm{K}^{T}$ | $\operatorname*{Softmax}\left(\bm{L}\right)$ | $\sigma\left(\bm{L}+\bm{b}\right)$ | $\Delta$ |
| --- | --- | --- | --- | --- |
| Expression | $2\,c\,n_{\text{ctx}}\,d_{\text{head}}$ | $3\,c\,n_{\text{ctx}}$ | $5\,c\,n_{\text{ctx}}$ | $1/d_{\text{head}}$ |
| Value | $262144\,c$ | $6144\,c$ | $10240\,c$ | $1/64$ |

## 4 FlashSigmoid : Hardware-Aware Implementation

Figure: (a) Inference mode kernels on H100.
Refer to caption: x1.png

Memory speed has not kept pace with recent gains in computation speed (Choquette, ; Jouppi et al., ; Hannun et al., ).
Consequently, attention computations on modern architectures have been IO-bound by memory accesses (Ivanov et al., ).
FlashAttention (Dao et al., ) and FlashAttention2 (Dao, ) 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 $\mathrm{SigmoidAttn}$.
Like previous works, FlashSigmoid employs three core ideas:

##### Tiling: Divide and Conquer Approach to Attention:

Similar to FlashAttention and FlashAttention2, FlashSigmoid processes input parts in parallel to 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 $\mathrm{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 activation 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-efficient than the alterantive approach of materializing and retaining the attention matrix.

The forward and backward pass algorithms of FlashSigmoid can be found in Sec.˜F.1.
Here, we highlight key differences between FlashSigmoid and FlashAttention/FlashAttention2.
The point-wise nature of $\mathrm{SigmoidAttn}$ results in a faster and more memory-efficient implementation by removing the need to compute the softmax normalization and materialize it to HBM.
A reduction 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 $\textrm{sigmoid}\left(x\right)=0.5\cdot\left(1+\textrm{tanh}\left(0.5\cdot x\right)\right)$ to optimize the sigmoid computation on GPU.
The speed up in FlashSigmoid compared to FlashAttention arises from optimizing hardware bottlenecks; theoretically, $\mathrm{SigmoidAttn}$ is slower than $\mathrm{SoftmaxAttn}$ (Sec.˜3.3).

To measure the performance improvements of FlashSigmoid, we compare the timings of the kernels in its forward and backward passes against those of FlashAttention2.
The details of this benchmarking on H100 and A100 GPUs can be found in Sec.˜F.2.
Measuring GPU computation time, we observe a $17.39\%$ speed-up during inference and a $6.53\%$ speed-up during training for attention over randomly initialized data on H100 GPU (Fig.˜1).
In practice, these gains may be affected by other bottlenecks, such as 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 $\sim$4% and inference by $\sim$8% in a realistic end-to-end setup.
The details of wall-clock time improvements with FlashSigmoid are in Sec.˜F.3.
We also note that practical machine learning workflows are dominated by inference rather than training.

## 5 Experiments

Figure: Figure 2: Train losses comparing $\mathrm{SigmoidAttn}$ with $\mathrm{SoftmaxAttn}$.
Refer to caption: x3.png

To empirically validate $\mathrm{SigmoidAttn}$, we evaluate across several domains: supervised image classification using vision transformers (Dosovitskiy et al., ), self-supervised image representation learning with SimCLR (Chen et al., ; Zhai et al., ), Bootstrap Your Own Latent (BYOL) (Grill et al., ; Busbridge et al., ) and Masked AutoEncoders (MAE) (He et al., ) as well as automatic speech recognition (ASR) (Synnaeve et al., ; Gulati et al., ) and auto-regressive language modeling (LM) (Brown et al., ). We also validate sequence length generalization on TED-LIUM v3 (Hernandez et al., ) for ASR and in small scale synthetic experiments in Sec.˜G.5.4.
Across all these domains and algorithms, we demonstrate that $\mathrm{SigmoidAttn}$ matches the performance of $\mathrm{SoftmaxAttn}$ (Fig.˜2 and 21), while offering training and inference speed-ups as highlighted in Sec.˜4. Empirically we make the following observations:

- (1)
$\mathrm{SigmoidAttn}$ is effective for vision tasks without a bias (except MAE), but relies on LayerScale to match the performance of the baseline $\mathrm{SoftmaxAttn}$ (Fig.˜9-a) in a hyper-parameter free manner.(^2^22Section G.2.2 demonstrates that supervised vision tasks using $\mathrm{SigmoidAttn}$ without LayerScale can match baseline $\mathrm{SoftmaxAttn}$ performance by relying on *learnable* scalar bias and temperature: $\{b,t\}\in\mathbb{R}$.) All results presented for $\mathrm{SoftmaxAttn}$ also fairly add LayerScale unless specified.
- (2)
LM and ASR are sensitive to the initial norm $||\sigma(\bm{Q}\bm{K}^{T}/\sqrt{d_{qk}})\bm{V}||$. Modulation is required via (a) relative positional embeddings like ALiBi (Press et al., ), which reduces the initial attention norm by shifting logit mass to the zero regime under $\mathrm{SigmoidAttn}$, (b) appropriate initialization of $b$ to achieve the same effect – enabling usage of any positional embedding.

Figure: Figure 3: $\mathrm{SigmoidAttn}$ with SinCos.
Refer to caption: figures/attn_norm_seed1000001_softmax_rope_vs_softmax_alibi_vs_sigmoid_sincos.png

### 5.1 Ablations

We begin with ablations to dissect the benefits of each of our introduced components. To gain intuition about $\mathrm{SigmoidAttn}$, we developed a research-friendly auto-regressive (AR) LM training framework to measure all components of attention and validate the effects of LayerScale, LayerNorm applied to Q and K (QK norm), different positional embedding techniques, and initialization values for $b$.

Figure: Figure 7: LR sensitivity LayerScale ablation.
Refer to caption: x4.png

Figure: Figure 9: ImageNet1k ViT-B/16 classification. (a) $\mathrm{SigmoidAttn}$ is robust without QK norm (+LayerScale, -QKNorm). Removing LayerScale reduces accuracy by 1.0% (-LayerScale, +/-QKNorm). $n^{-\alpha}$ normalization (Wortsman et al., ) underperforms without LayerScale. (b) $\mathrm{SigmoidAttn}$ multi-query attention (MQA) (Shazeer, ) with one head matches multi-head attention (MHA). (c) Sigmoid with LayerScale and QK norm performs comparably to other activations, except TanH. ReLU^2 (Hua et al., ) underperforms without LayerScale and QK norm.
Refer to caption: x6.png

##### Mitigating Large Attention Norms

We train a single layer AR transformer block (E=3072, D_FF=12288) on the realnews split of C4 (Raffel et al., ). We train for $2^{16}$ steps using a batch size of 6 and max sequence length of 4096 using a single cycle cosine learning rate (LR) schedule without weight decay. $\mathrm{SigmoidAttn}$ initially underperformed $\mathrm{SoftmaxAttn}$ when using absolute sinusoidal (SinCos) (Fig.˜6) or relative (Fig.˜6) positional embeddings (PE), which we attribute to high initial attention Frobenius norms, $\lVert\sigma(\bm{Q}\bm{K}^{T}/\sqrt{d})\bm{V}\rVert$. A corresponding evolution of the attention distribution and sparsity can be seen in Appendix Fig.˜31 and Fig.˜32 on a synthetic task.
To address these larger attention norms, we propose: (a) using ALiBi (Press et al., ) whose relative bias moves initial attention logit mass to the zero region under the sigmoid activation, producing equivalent train negative log-likelihoods (Fig.˜6); or (b) set the attention logit bias $b$ to a negative offset proportional to the sequence length, $b\propto-\ln n$ (see Sec.˜G.1.2 for an ablation on $b$). This enables the usage of other PE techniques like RoPE (Su et al., ) (Fig.˜6).

##### LayerScale

To validate the need for LayerScale, we follow Wortsman et al. () to quantify the impact on stability.
All models are trained with RoPE with $b\propto-\ln n$, using AdamW (Loshchilov & Hutter, ) on the
realnews split of C4
with $(\beta_{1},\beta_{2})=(0.9,0.95)$, ${\epsilon}=10^{-8}$, $wd=0$,
batch size 24, maximum token sequence length of 512 from the T5 tokenizer (Raffel et al., ), cosine LR schedule of $2^{14}$ steps including a linear warmup of $2^{10}$ steps.
Models have
$n_{\text{heads}}=\kappa$,
$n_{\text{layers}}=2\times\kappa$,
$d_{\text{model}}=64\times\kappa$ and
$d_{\text{feed-forward}}=256\times\kappa$
for a scaling value $\kappa\in\{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. (),
we sweep learning rates
$\eta\in\{3\times 10^{-4},1\times 10^{-3},3\times 10^{-3},1\times 10^{-2},3\times 10^{-2},1\times 10^{-1},3\times 10^{-1}\}$.
LR sensitivity is defined as
$\mathbb{E}_{\eta\in[a,b]}\left[\min(\ell(\mathcal{A}(\eta)),\ell_{0})-\ell^{*}\right]$
where $\ell(\mathcal{A}(\eta))$ is the loss achieved by the learning algorithm $\mathcal{A}$ with LR $\eta$,
$\ell_{0}$ is the loss at initialization, and
$\ell^{*}$ is the loss achieved by the best LR.
LayerScale is initialized at $10^{-4}$.
Unlike vision tasks, where LayerScale *improves performance* (Fig.˜9-a), in LM, we observe that $\mathrm{SoftmaxAttn}$ slightly benefits from LayerScale, while the performance of $\mathrm{SigmoidAttn}$ remains largely unaffected.

##### Stability with QK Norm

Theorem˜3.2 indicates that the Jacobian of $\mathrm{SigmoidAttn}$ has favorable properties compared to $\mathrm{SoftmaxAttn}$. We explore this by repeating the analysis of Wortsman et al. (), as described in the LayerScale analysis, to investigate the impact of QK norm (Dehghani et al., ). For language modeling, both $\mathrm{SigmoidAttn}$ and $\mathrm{SoftmaxAttn}$ exhibit sensitivity to learning rate changes without QK norm. However, incorporating QK norm significantly stabilizes performance (Fig.˜8). In vision tasks, $\mathrm{SigmoidAttn}$ demonstrates robustness with and without QK norm (Fig.˜9-a) and without the need for $n^{-\alpha}$ normalization from Wortsman et al. ().(^3^33We ablate multiplicative sequence length scaling in more detail in Sec. G.1.1.)

##### Multi-query attention (MQA)

In Fig.˜9-b we explore MQA (Shazeer, ) for vision using only one head for $\{\bm{K},\bm{V}\}$. We find that both $\mathrm{SigmoidAttn}$ and $\mathrm{SoftmaxAttn}$ perform equally well with or without multiple heads even at the small scale of ViT-B/16.

##### Activation Function Ablations

As in Wortsman et al. (), various activation functions, when combined with LayerScale and QK norm, perform equally well for vision tasks (Fig.˜9-c). However, for sequence-critical tasks like ASR, activation functions such as ReLU pose instabilities and underperform. In the same figure, we also compare to the ReLU^2 proposal from Hua et al. () and find that it underperforms without LayerScale and QK norm.

### 5.2 Supervised Image Classification

Vision transformers (Dosovitskiy et al., ) extend transformers (Vaswani et al., ) to treat $K\times K$ image grids as 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, $\ln p(\bm{y}|\bm{x})$. We train ViT-B/16 models using $\mathbb{R}^{224\times 224\times 3}$ images for 300 epochs using the recipe provided in Sec.˜G.2.4. We use the same set of training hyper-parameters for both $\mathrm{SoftmaxAttn}$ and $\mathrm{SigmoidAttn}$, changing only the activation function between trials. The train negative log-likelihood is reported in Fig.˜2 and the test top-1% is reported in Fig.˜21. We find that $\mathrm{SigmoidAttn}$ matches both the training dynamics and the evaluation performance of $\mathrm{SoftmaxAttn}$.

### 5.3 Self-Supervised Image Representation Learning

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. (), BYOL (Grill et al., )) or reconstruction from compressed representations (MAE (He et al., )). We employ vision transformer training recipes from Zhai et al. () and Busbridge et al. () (Sec.˜G.2.4) for SimCLR and BYOL. As with supervised learning, we use the same set of training hyper-parameters for both $\mathrm{SoftmaxAttn}$ and $\mathrm{SigmoidAttn}$, changing only the activation function between trials. Figure˜2 reports the train losses, and Fig.˜21 highlights the linear probe and finetuned test top-1%. Despite the diverse training objectives in SSL, $\mathrm{SigmoidAttn}$ matches $\mathrm{SoftmaxAttn}$ while improving training and inference throughput (Sec.˜4).

### 5.4 Automatic Speech Recognition (ASR)

**Table 2: Word error rate (%) on LibriSpeech test sets and TED-LIUM v3 (Hernandez et al., ) (“TED”, joint validation and test sets split according to duration) for transformer (255M params) with either $\mathrm{SoftmaxAttn}$ or $\mathrm{SigmoidAttn}$ (LayerScale and QK norm are used with $b=-\log n$) trained on LibriSpeech 960h data (mean duration is 10-15s). Hyper-parameters are in Sec.˜G.4.**
| attn | PE | test-clean | test-other | ted 0-10s | ted 10-20s | ted 20-30s | ted 30s+ |
| --- | --- | --- | --- | --- | --- | --- | --- |
| softmax | CAPE | 2.3 | 5.7 | 12.4 | 10.5 | 11.9 | 9.1 |
| sigmoid | 2.4 | 5.5 | 12.4 | 10.3 | 12.3 | 9.7 |  |
| - QK norm | unstable, gradient norm and loss spikes |  |  |  |  |  |  |
| - LayerScale | 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 ($b=-5$ in $Q$, learnable) | 2.3 | 5.4 | 12.2 | 10.8 | 12.4 | 9.9 |  |
| - QK norm | unstable, gradient norm and loss spikes |  |  |  |  |  |  |
| softmax | RoPE | 2.2 | 5.5 | 12.7 | 10.6 | 12.8 | 9.5 |
| sigmoid | 2.3 | 5.4 | 12.3 | 10.1 | 12.3 | 8.6 |  |
| sigmoid ($b=-10$, learnable) | 2.2 | 5.2 | 12.4 | 10.5 | 12.3 | 21.8 |  |
| + $\alpha=1$ | 2.7 | 6.6 | 14.1 | 12.0 | 14.5 | 14.9 |  |
| sigmoid ($b=-5$ in $Q$, learnable) | unstable, gradient norm and loss spikes |  |  |  |  |  |  |
| softmax | ALiBi | 2.2 | 5.4 | 12.3 | 10.7 | 12.1 | 8.6 |
| sigmoid | 2.3 | 5.1 | 12.3 | 10.5 | 12.6 | 9.1 |  |
| sigmoid ($b=-10$, learnable) | 2.2 | 5.2 | 12.4 | 10.4 | 11.7 | 9.1 |  |
| + $\alpha=1$ | 2.6 | 6.6 | 13.9 | 11.9 | 14.2 | 8.6 |  |
| sigmoid ($b=-5$ in $Q$, learnable) | 2.2 | 5.2 | 12.1 | 10.4 | 12.0 | 8.2 |  |

We benchmark ASR using LibriSpeech data (Panayotov et al., ) on 100h and 960h settings of paired speech and text transcriptions. Our PyTorch implementations of encoder-based vanilla transformer (Synnaeve et al., ) and conformer (Gulati et al., ) are trained with Connectionist Temporal Classification (CTC) (Graves et al., ) w/ BF16 mixed precision, w/o QK norm and w/o LayerScale. After extensively tuning $\mathrm{SoftmaxAttn}$ baselines, we switch to $\mathrm{SigmoidAttn}$ per eq.˜3 without any other changes. We investigate the effects of post/pre-LayerNorm, model depth, optimizer type, small data regime, and connection to local attention, with details in Sec.˜G.4.

Our main findings are: i) CAPE (Likhomanenko et al., ) PE is the most unstable for $\mathrm{SigmoidAttn}$; ii) post-LayerNorm models with $\mathrm{SoftmaxAttn}$ are hard to match with stable $\mathrm{SigmoidAttn}$; iii) w/o QK norm $\mathrm{SigmoidAttn}$ is unstable and significant spikes happen in both gradient norms and training loss; iv) LayerScale is needed for generalization; v) learnable bias $b=-10$ gives no loss and gradient norms spikes while matching the $\mathrm{SoftmaxAttn}$ (which does not benefit from the improved throughput of FlashSigmoid); vi) adding a 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 n$ gives rare (2-5 times) marginal gradient norms spikes with smooth loss while matching $\mathrm{SoftmaxAttn}$.

Table˜2 shows the main result for pre-LayerNorm transformers with CAPE, RoPE, and ALiBi, where $\mathrm{SigmoidAttn}$ uses LayerScale, QK norm, $b=-\log n$, and no sequence normalization. The bias is ablated with learnable bias (one per layer) in attention or $Q$ with or without sequence normalization. $\mathrm{SigmoidAttn}$ is stabilized with bias while matching $\mathrm{SoftmaxAttn}$, and $b=-\log n$ works well. In most cases, bias allows generalization to longer sequences without sequence normalization, except for RoPE where it helps for longer sequences but hurts overall performance.

### 5.5 Autoregressive Large Language Modeling

**Table 3: 1B LLM English evaluation.**
| Model | Seq.<br>Len. | ARC<br>Easy | ARC<br>Challenge | Hella-<br>swag | Piqa | Sciq | Wino-<br>grande | Lambada<br>OpenAI | TriviaQA<br>(1-shot) | WebQS<br>(1-shot) | AVG | Step<br>time (s) |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Softmax (ALiBi) | 2k | 62.2 | 26.8 | 42.4 | 59.0 | 72.3 | 88.1 | 58.4 | 19.9 | 15.4 | 49.4 | 0.38 |
| Sigmoid (ALiBi) | 2k | 62.8 | 28.8 | 42.5 | 59.7 | 70.3 | 88.6 | 59.7 | 19.1 | 13.8 | 49.5 | 0.34 |
| Softmax (RoPE) | 4k | 63.3 | 29.3 | 43.3 | 58.1 | 71.3 | 86.9 | 58.8 | 20.4 | 15.6 | 49.7 | 0.84 |
| Softmax (ALiBi) | 4k | 62.6 | 27.7 | 42.4 | 58.6 | 71.1 | 88.2 | 58.6 | 18.9 | 14.7 | 49.2 | 0.84 |
| Sigmoid (ALiBi) | 4k | 60.5 | 27.3 | 41.3 | 57.8 | 70.5 | 87.0 | 57.6 | 18.9 | 12.6 | 48.2 | 0.67 |

We initially iterated at the 85M scale, as it served as a proxy for larger scale training. Our findings show that: i) attention bias is required for stability, which can be learnable, but setting it to $-\log(n)$, where $n$ is the maximum training sequence length of 4096, works well and is faster; ii) RoPE is more challenging to stabilize; iii) the final setting exhibits smooth loss curves, but still shows gradient norm fluctuations. We then turn our attention to validating $\mathrm{SigmoidAttn}$ at scale.

We train a 1B language model using the Llama2 (Touvron et al., ) recipe with ALiBi instead of RoPE positional embedding, and the RedPajama (Computer, ) dataset (see Sec.˜G.3). At sequence length 4096, $\mathrm{SigmoidAttn}$ achieves a 1.23$\mathbf{\times}$ step-time improvement over $\mathrm{SoftmaxAttn}$ in JAX without FlashAttention (Tab.˜3). All LLMs are trained using the AXLearn framework, which include the recipe and $\mathrm{SigmoidAttn}$ implementation.(^4^44https://github.com/apple/axlearn)

$\mathrm{SoftmaxAttn}$ and $\mathrm{SigmoidAttn}$ have matching train and validation NLL at 85M (Fig.˜27) and at 1B scale when using 2048 sequence length (Fig.˜2). However, a slight disparity is observed at 1B scale when using 4096 sequence length, which we leave for future investigation (more details in Sec.˜G.3).

## 6 Related Work

Recent studies in supervised image classification (Wightman et al., ) and self-supervised learning (SSL), including approaches like SigLIP (Zhai et al., ),
are shifting large-scale machine learning training from output conditional categorical distributions, traditionally parameterized by softmax functions, to richer pointwise Bernoulli 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 et al., ; Hron et al., ) and theoretical settings (Bai et al., ; Fu et al., ). Other works explores using the ReLU^2 activation (Hua et al., ), exploring purely linear attention (Katharopoulos et al., ; Lu et al., ; Koohpayegani & Pirsiavash, ) or cosine-similarity based attention (Luo et al., ; Liu et al., ).
Our work builds upon these explorations, particularly Wortsman et al. (), which replaces softmax with various activation functions scaled by $n^{-\alpha}$, where $n$ corresponds to the sequence length and $\alpha$, a hyper-parameter. However, we find that their formulation does not match expected performance without proper $b$ initialization and the use of LayerScale (Fig.˜9-a, Sec.˜G.1.1).

## 7 Conclusion

In this work, we present a comprehensive theoretical and empirical study of sigmoid attention as 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, a memory-efficient variant providing a 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 as a viable alternative, unifying prior work and establishing best practices for its application in transformers.

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

## Appendix A Limitations

While our work demonstrates that $\mathrm{SigmoidAttn}$ can serve as a viable drop-in replacement for $\mathrm{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, we observed some gradient norm spikes and a slight performance gap between $\mathrm{SigmoidAttn}$ and $\mathrm{SoftmaxAttn}$ (Table˜3). While runs at smaller context lengths (1B parameter, n=2048) were stable and matched $\mathrm{SoftmaxAttn}$ performance, further research is needed to fully close the performance gap and ensure training stability for very large language models using $\mathrm{SigmoidAttn}$.
- (2)
Our theoretical analysis proves that transformers with $\mathrm{SigmoidAttn}$ are universal function approximators and have improved regularity compared to $\mathrm{SoftmaxAttn}$. However, the bounds we derive, while tighter than those for $\mathrm{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 or emerging applications remains to be validated.
- (4)
In automatic speech recognition experiments, we observed that $\mathrm{SigmoidAttn}$ can be sensitive to the choice of positional embeddings and may require careful initialization of the attention bias term to ensure stable training. Specifically, we found that the CAPE positional embedding was the most unstable for $\mathrm{SigmoidAttn}$. Further work is needed to develop robust initialization schemes that work well across different positional embeddings. Moreover we found that w/o QK norm or with post-LayerNorm $\mathrm{SigmoidAttn}$ is unstable and can underperforms $\mathrm{SoftmaxAttn}$, thus further investigation is needed.
- (5)
FlashSigmoid demonstrates promising inference and training speed-ups by exploiting $\mathrm{SigmoidAttn}$’s simpler kernel structure compared to $\mathrm{SoftmaxAttn}$. 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 a strong foundation for $\mathrm{SigmoidAttn}$, unifying prior art and demonstrating its potential as a drop-in $\mathrm{SoftmaxAttn}$ replacement. We hope our theoretical grounding and empirical results motivate further research into this simple yet effective architectural variation.

## Appendix 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 $\mathrm{SigmoidAttn}$ as a viable alternative to $\mathrm{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 faster kernel implementation could lead to more efficient training and inference for attention-based models, reducing energy consumption and enabling deployment on resource-constrained devices. This could democratize access to powerful models.
- (2)
Theoretical understanding: Our universal approximation results and tighter bounds on the regularity of $\mathrm{SigmoidAttn}$ contribute to a deeper theoretical understanding of this key component. A stronger theoretical foundation can guide principled model design and architectural search.
- (3)
Application-specific benefits: Across language, vision, and speech domains, $\mathrm{SigmoidAttn}$’s performance could translate into improved user experiences, such as more natural language interactions, enhanced image understanding, and robust speech recognition. These advancements 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:

- (1)
Fairness and bias considerations: As with any machine learning model, it is important to carefully evaluate $\mathrm{SigmoidAttn}$ based models for fairness and potential biases when applied to sensitive use cases. The unique properties of $\mathrm{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 $\mathrm{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 a future where the efficiency and expressivity of $\mathrm{SigmoidAttn}$ are used for societal benefit.

## Appendix 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. (, Sec. 3), of which we inherit much of the notation, and we encourage the interested reader to refer to the original source for a more comprehensive understanding of its details. Here we first provide context by outlining the main steps in the original proof, before proceeding to adapt its key components to the $\mathrm{SigmoidAttn}$ case.

The proof aims at showing that a transformer network can approximate to arbitrary accuracy any continuous, permutation-equivariant function with compact support. The proof is constructive in nature, in that it explicitly defines the architecture (and particularly, the sequence of self-attention and feed-forward layers) that can approximate a given target function. To do so, it proceeds in steps (see Yun et al. (, 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-constructed *modified* transformer network, (where the softmax nonlinearity is substituted with a hardmax nonlinearity), can exactly represent such piecewise constant function. This step is further divided into three sub-steps (see Yun et al. (, Sec. 4)):
(a)
prove that a series of feed-forward layers can quantize any input to a specific discretization grid in the compact domain
(b)
prove that a series of self-attention layers can implement a *contextual mapping* (see Yun et al. (, Def. 3.1))
(c)
prove that a series of feed-forward layers can map the output of the contextual mapping to the desired output of the target piecewise-constant approximation
- (3)
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)(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.1 Proof of Step ˜ (3) : Sigmoid Transformers can Approximate Modified Sigmoid Transformers

In Yun et al. (), to implement contextual mappings, the authors rely on a *modified* version of transformers, for the sake of simplifying the analysis. In their modified version, the (row-wise) softmax operation is substituted with a (row-wise) hardmax operation. This substitution is valid because a classical transformer can still be made arbitrarily close to such modified transformer, in light of the fact that

$$ $\text{softmax}(\lambda{\bm{X}})\xrightarrow[]{\lambda\to\infty}\text{hardmax}({\bm{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

$$ $\sigma(\lambda{\bm{X}})\xrightarrow{\lambda\to\infty}H({\bm{X}}),$ (11) $$

where $\sigma(x)=(1+e^{-x})^{-1}$ is the (elementwise) sigmoid function, while

$$ $H(x)=\begin{cases}1&x>0\\ \frac{1}{2}&x=0\\ 0&x<0\end{cases}$ (12) $$

denotes the (elementwise) Heaviside step function. This allows us to define our modified sigmoid self-attention layer, as follows.

###### Definition C.1 (Modified sigmoid self-attention layer) .

Given an input ${\bm{X}}\in\mathbb{R}^{d\times n}$, the action of a modified sigmoid self-attention layer with shifts and a single one-dimensional head is defined as ${\bm{X}}\mapsto{\bm{X}}+\psi({\bm{X}};{\bm{q}},{\bm{b}}_{q},{\bm{k}},{\bm{b}}_{k},{\bm{v}},{\bm{o}})$, where

$$ $\psi({\bm{X}};{\bm{q}},{\bm{b}}_{q},{\bm{k}},{\bm{b}}_{k},{\bm{v}},{\bm{o}})={\bm{o}}\left({\bm{v}}^{T}{\bm{X}}\right)H\left(\left({\bm{q}}^{T}{\bm{X}}-{\bm{b}}_{q}^{T}\right)^{T}\left({\bm{k}}^{T}{\bm{X}}-{\bm{b}}_{k}^{T}\right)\right)$ (13) $$

with ${\bm{q}},{\bm{k}},{\bm{v}}\in\mathbb{R}^{d}$ representing the query, key, and value vectors, ${\bm{b}}_{q},{\bm{b}}_{k}\in\mathbb{R}^{n}$ the corresponding query and key bias vectors, while ${\bm{o}}\in\mathbb{R}^{d}$ denotes the output vector.

Analogously to eq.˜10, eq.˜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 we are considering the input tensor ${\bm{X}}\in\mathbb{R}^{d\times n}$, (as opposed to $\in\mathbb{R}^{n\times d}$) to better align out notation with the one used in Yun et al. ().

###### Definition C.1 (Modified sigmoid self-attention layer) .

### C.2 Proof of Step ˜ (2)(2-b) : Modified Sigmoid Transformers can Implement Contextual Mappings

The core of the proof consists in showing how, by opportunely combining the operations in eq.˜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. (, Def. 3.1)).

###### Definition C.2 (Contextual mapping) .

A map ${\bm{q}}:\mathbb{L}\to\mathbb{R}^{n}$ from a finite set $\mathbb{L}\subset\mathbb{R}^{d\times n}$ is said to be a *contextual mapping* if both the following conditions hold:

- (i)
$q_{i}({\bm{X}})\neq q_{j}({\bm{X}})$, $\forall i\neq j$ and $\forall{\bm{X}}\in\mathbb{L}$
- (ii)
$q_{i}({\bm{X}})\neq q_{j}({\bm{X}}^{\prime})$, $\forall i,j$ and $\forall{\bm{X}},{\bm{X}}^{\prime}\in\mathbb{L}$, with ${\bm{X}}\neq{\bm{X}}^{\prime}$

where $q_{i}({\bm{X}})$ denotes the $i$-th component of ${\bm{q}}({\bm{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)(2-a)) can be mapped to a 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)(2-c)) to correctly map this value to the corresponding output value in the piece-wise constant approximation.

In particular, after defining a uniform discretization (characterized by the parameter $\delta$) of the unitary hypercube $[0,1]^{d}\subset\mathbb{R}^{d}$, namely

$$ $\mathbb{G}_{\delta}\coloneqq\{{\bm{g}}:g_{i}\in\{0,\delta,2\delta,\dots,1-\delta\},\quad\forall i=1\dots d\},$ (14) $$

we consider as input a tensor ${\bm{X}}$ (composed of columns ${\bm{X}}=[{\bm{x}}_{i}]_{i=1}^{n}$) such that

$$ ${\bm{X}}\in\mathbb{L}\coloneqq\{{\bm{X}}:{\bm{x}}_{i}\in\mathbb{G}_{\delta}\;\forall i=1\dots n,\quad\text{and}\quad{\bm{x}}_{i}\neq{\bm{x}}_{j}\;\forall i\neq j\}\subset\mathbb{R}^{d\times n},$ (15) $$

that is, a 2D tensor whose columns are element of the discretization $\mathbb{G}_{\delta}$, and that all differ from each other (at least for one element).
We want to build a contextual mapping acting on $\mathbb{L}$, by stacking layers parameterized according to Def.˜C.1. In Sec.˜C.2.1 we define the basic building blocks of our architecture; in Sec.˜C.2.2 we describe how to stack them, and the effect the architecture has on a given input; finally, in Sec.˜C.2.4 we prove that this architecture indeed implements a contextual mapping.

###### Definition C.2 (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 $\mathrm{SigmoidAttn}$ layer in Def.˜C.1. In the following, we describe how we can tweak its parameters to achieve exactly this.

##### From d d -dimensional quantized vectors to scalars

As a first simplification, we can get rid of the $d$-dimension in the ${\bm{X}}$ tensor by mapping each of its columns to a corresponding identifying scalar, uniquely defined by the specific column components. This step is also performed in Yun et al. (, App. B.5), and can be achieved rather straightforwardly, by defining

$$ ${\bm{v}}\equiv{\bm{q}}\equiv{\bm{k}}\equiv{\bm{u}}\coloneqq[1,\delta^{-1},\delta^{-2},\dots,\delta^{-d+1}]^{T}.$ (16) $$

Notice in fact that, since each column ${\bm{x}}_{i}$ belongs to $\mathbb{G}_{\delta}$, it can equivalently be written in the form ${\bm{x}}_{i}=\delta\cdot[\text{id}_{0,i},\text{id}_{1,i},\dots,\text{id}_{d-1,i}]^{T}$, where $\text{id}_{j,i}\in\{0,1,2,\dots,\delta^{-1}-1\}$ represents the (indexed) coordinate of the discretization along the $j$-th dimension. Scalar-multiplying ${\bm{X}}$ by ${\bm{u}}$ in eq.˜16, then, turns this tuple of indices into a single one, in a bijective fashion(^5^55For example, consider $d=3$ and the column defined as ${\bm{x}}_{i}=[3\delta,10\delta,2\delta]^{T}$, that is, the column identified by the *triplet* of indices $[3,10,2]$. Multiplying by ${\bm{u}}$ would then give the scalar ${\bm{u}}^{T}{\bm{x}}_{i}=(3+10N+2N^{2})\delta$, where $N=\delta^{-1}$, which is uniquely identified by the *single* index $(3+10N+2N^{2})$.).

This allows us to equivalently consider a single vector ${\bm{u}}^{T}{\bm{X}}\in\mathbb{R}^{n}$, rather than the whole tensor ${\bm{X}}\in\mathbb{R}^{d\times n}$ in the remainder of our analysis. Analogously, choosing ${\bm{o}}\equiv{\bm{e}}_{0}\coloneqq[1,0,\dots,0]^{T}$ in eq.˜13 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 ${\bm{q}}$ in Def.˜C.2.
To slim our notation, in the following we often refer to ${\bm{u}}^{T}{\bm{X}}$ as the vector ${\bm{l}}\in\mathbb{R}^{n}$, with components $l_{i}$.

In light of the simplification above, we can rewrite eq.˜13 more compactly, as follows:

$$ $\psi({\bm{X}};{\bm{q}}={\bm{k}}={\bm{v}}\equiv{\bm{u}},{\bm{o}}\equiv{\bm{e}}_{0};{\bm{b}}_{q},{\bm{b}}_{k})={\bm{e}}_{0}{\bm{l}}^{T}H\left(\left({\bm{l}}-{\bm{b}}_{q}\right)\otimes\left({\bm{l}}-{\bm{b}}_{k}\right)\right)$ (17) $$

Notice that, since the elements of both ${\bm{X}}$ and ${\bm{u}}$ are always non-negative, so are those of ${\bm{l}}$, too. Moreover, since we are interested in permutation-equivariant functions with respect to the columns of ${\bm{X}}$, without loss of generality we can consider the elements of ${\bm{l}}={\bm{u}}^{T}{\bm{X}}$ to be ordered: $0\leq l_{i}<l_{j}$, $\forall i<j$.

##### Selective shift operation for sigmoid attention

Since we aim to recover a contextual map by sequentially updating the elements of ${\bm{l}}$, we proceed by designing a modification of eq.˜17 which affects only a certain selected element at a time. This is were our second simplification comes into play, and this time it pertains the roles of the bias vectors ${\bm{b}}_{q}$ and ${\bm{b}}_{k}$. Since ${\bm{l}}\geq 0$, these vectors have the effect of tweaking the sign of the inner arguments of the Heaviside function in eq.˜17, hence directly impacting when its application outputs 0 or $1$. By aptly selecting the values of ${\bm{b}}_{k}$ and ${\bm{b}}_{q}$, then, we can explicitly decide when a specific layer triggers an update, which elements are affected by the update, and what elements to consider to compute the update itself.

More in detail, take ${\bm{b}}_{q}={\bm{1}}b_{q}$ and ${\bm{b}}_{v}={\bm{1}}b_{v}$, for some scalars $b_{q},b_{v}$, and with ${\bm{1}}$ being the all-one vector. Plugging this into eq.˜17, we have

$$ $\begin{split}\tilde{\psi}({\bm{X}};b_{q},b_{k})&\coloneqq\psi({\bm{X}};{\bm{q}}={\bm{k}}={\bm{v}}\equiv{\bm{u}},{\bm{o}}\equiv{\bm{e}}_{0},{\bm{b}}_{q}={\bm{1}}b_{q},{\bm{b}}_{k}={\bm{1}}b_{k})\\ &={\bm{e}}_{0}{\bm{l}}^{T}H\left(\left({\bm{l}}-{\bm{1}}b_{q}\right)\otimes\left({\bm{l}}-{\bm{1}}b_{k}\right)\right)={\bm{e}}_{0}\begin{cases}\sum_{i:l_{i}<b_{v}}l_{i}&\text{ if }l_{j}<b_{k}\\ \sum_{i:l_{i}>b_{v}}l_{i}&\text{ if }l_{j}>b_{k}\end{cases};\end{split}$ (18) $$

notice how $b_{q}$ determines what elements of ${\bm{l}}$ compose the update (as it impacts the indices considered in the sum), while $b_{k}$ defines the elements impacted by the update itself
(^6^66This can be better seen by considering independently the effects of the two parameters $b_{k}$, $b_{q}$ on the modified sigmoid attention matrix $H\left(\left({\bm{l}}-{\bm{1}}b_{q}\right)\otimes\left({\bm{l}}-{\bm{1}}b_{k}\right)\right)$. We have in fact, with $b_{q}=0$, $\begin{array}[]{cc}&\begin{array}[]{rl}l_{j}<b_{k}&l_{j}>b_{k}\end{array}\\ \begin{array}[]{r}H\left({\bm{l}}\otimes\left({\bm{l}}-{\bm{1}}b_{k}\right)\right)\end{array}=&\left[\begin{array}[]{ccc|ccc}0&\cdots&0&1&\cdots&1\[-4.30554pt] \vdots&\ddots&\vdots&\vdots&\ddots&\vdots\\ 0&\cdots&0&1&\cdots&1\\ \end{array}\right].\\ \end{array}$ (19) This shows how, by modifying $b_{k}$, one can decide which columns will receive an update: namely, all those with index $l_{j}>b_{k}$. By combining two such operators with $b_{k}=\left(i-\frac{1}{2}\right)\delta$ and $b_{k}=\left(i+\frac{1}{2}\right)\delta$, we then recover $\begin{array}[]{cc}&l_{j}=i\delta\\ \begin{array}[]{r}H\left({\bm{l}}\otimes\left({\bm{l}}-{\bm{1}}\left(i-\frac{1}{2}\right)\delta\right)\right)\\ -H\left({\bm{l}}\otimes\left({\bm{l}}-{\bm{1}}\left(i+\frac{1}{2}\right)\delta\right)\right)\end{array}=&\left[\begin{array}[]{ccc|c|ccc}0&\cdots&0&1&0&\cdots&0\[-4.30554pt] \vdots&\ddots&\vdots&\vdots&\vdots&\ddots&\vdots\\ 0&\cdots&0&1&0&\cdots&0\\ \end{array}\right],\end{array}$ (20) which allows us to limit the update to only one specific column: the one with index $l_{j}=i\delta$. The parameter $b_{q}$ 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 eq. 20, but with $b_{q}=\left(i+\frac{1}{2}\right)\delta$ gives us in fact: $\begin{array}[]{ccc}&l_{j}=i\delta&\\ \begin{array}[]{r}H\left(\left({\bm{l}}-{\bm{1}}\left(i+\frac{1}{2}\right)\delta\right)\otimes\left({\bm{l}}-{\bm{1}}\left(i-\frac{1}{2}\right)\delta\right)\right)\\ -H\left(\left({\bm{l}}-{\bm{1}}\left(i+\frac{1}{2}\right)\delta\right)\otimes\left({\bm{l}}-{\bm{1}}\left(i+\frac{1}{2}\right)\delta\right)\right)\\ \end{array}=&\left[\begin{array}[]{ccc|c|ccc}0&\cdots&0&-1&0&\cdots&0\[-4.30554pt] \vdots&\ddots&\vdots&\vdots&\vdots&\ddots&\vdots\\ 0&\cdots&0&-1&0&\cdots&0\\ \hline\cr 0&\cdots&0&1&0&\cdots&0\[-4.30554pt] \vdots&\ddots&\vdots&\vdots&\vdots&\ddots&\vdots\\ 0&\cdots&0&1&0&\cdots&0\\ \end{array}\right]&\begin{array}[]{c}l_{j}<i\delta\\ \\ \\ l_{j}>i\delta\\ \end{array}\end{array}.$ (21) Finally, eq. 22 can be recovered by combining eqs. 20 and 21: this has the effect of removing the $-1$’s in eqs. 45 and 21.).
If we opportunely combine *four* modified sigmoid self-attention heads $\tilde{\psi}({\bm{X}};b_{q},b_{k})$, we recover, for a given index $i=0\dots\delta^{-d}-1$,

$$ $\begin{split}\Psi^{(i)}({\bm{X}})\coloneqq&{\bm{X}}+\frac{1}{2}c\left(\begin{array}[]{l}\tilde{\psi}\left({\bm{X}};b_{q}=0,b_{k}=\left(i-\frac{1}{2}\right)\delta\right)\\ -\tilde{\psi}\left({\bm{X}};b_{q}=0,b_{k}=\left(i+\frac{1}{2}\right)\delta\right)\\ -\tilde{\psi}\left({\bm{X}};b_{q}=b_{k}=\left(i+\frac{1}{2}\right)\delta\right)\\ +\tilde{\psi}\left({\bm{X}};b_{q}=\left(i+\frac{1}{2}\right),b_{k}=\left(i-\frac{1}{2}\right)\delta\right)\\ \end{array}\right)\\ =&{\bm{X}}+\frac{1}{2}c{\bm{e}}_{0}{\bm{l}}^{T}\left(\begin{array}[]{l}H\left({\bm{l}}\otimes\left({\bm{l}}-\left(i-\frac{1}{2}\right)\delta\right)\right)\\ -H\left({\bm{l}}\otimes\left({\bm{l}}-\left(i+\frac{1}{2}\right)\delta\right)\right)\\ -H\left(\left({\bm{l}}-\left(i+\frac{1}{2}\right)\delta\right)\otimes\left({\bm{l}}-\left(i+\frac{1}{2}\right)\delta\right)\right)\\ +H\left(\left({\bm{l}}-\left(i+\frac{1}{2}\right)\delta\right)\otimes\left({\bm{l}}-\left(i-\frac{1}{2}\right)\delta\right)\right)\\ \end{array}\right)\\ \Longrightarrow&\Psi^{(i)}_{1,j}({\bm{X}})={\bm{X}}_{1,j}+c\begin{cases}\sum_{k:l_{k}>i\delta}l_{k}&\text{ if }l_{j}=i\delta\\ 0&\text{ otherwise}\end{cases}\\ \Longrightarrow&\Psi^{(i)}_{k>1,j}({\bm{X}})={\bm{X}}_{k,j},\end{split}$ (22) $$

where $c\equiv c(\delta,d,n)$ is a multiplicative constant which will be chosen later.

The operator assembled in eq.˜22 defines the basic layer of the architecture that we use in our proof. Notice $\Psi^{(i)}({\bm{X}})$ has the effect of modifying only the column ${\bm{x}}_{j}$ which has index $l_{j}={\bm{u}}^{T}{\bm{x}}_{j}=i\delta$ (if at all present in the input ${\bm{X}}$). This layer covers a similar role to the *selective shift operation* introduced in Yun et al. (, App. B.5), but it has been adapted to account for the presence of a sigmoid nonlinearity: notice this required us to use $4$-headed attention, while in Yun et al. () a $2$-headed version is sufficient.

#### C.2.2 Result of Applying a Sequence of Selective Shifts

Ultimately we want to show how, by stacking a sequence of selective shift layers eq.˜22 for increasing $i=0\dots\delta^{-d}-1$ and one additional global shift, we can build an architecture capable of representing a contextual mapping. As a preliminary step, in this section we provide an explicit formula for the result of applying such an architecture. Once again, we are proceeding analogously to Yun et al. (, App. B.5.1).

##### After the first selective shift application

Consider a quantized input sequence ${\bm{X}}\in\mathbb{L}$ as defined in eq.˜15, with its columns ordered according to their scalar indices ${\bm{l}}={\bm{u}}^{T}{\bm{X}}$. The sequence of selective shift layers $\Psi^{(0)},\Psi^{(1)},\dots$ initially has no effect on the input itself, and it leaves it unchanged until we hit the layer corresponding to the index of the first column in the input, $\Psi^{(\hat{i})}$, where $l_{1}={\bm{u}}^{T}{\bm{x}}_{1}=\hat{i}\delta$. At this point, following eq.˜22, the first column of the input is modified into

$$ ${\bm{x}}_{1}\quad\mapsto\quad\Psi^{(\hat{i})}_{|,1}({\bm{X}})={\bm{x}}_{1}+c{\bm{e}}_{0}\sum_{k:l_{k}>l_{1}}l_{k}={\bm{x}}_{1}+c{\bm{e}}_{0}\left(\sum_{k=1}^{n}l_{k}-l_{1}\right)$ (23) $$

while the other columns are still left untouched. In the following, we compactly refer to the quantities $\sum_{k=1}^{n}l_{k}-l_{i}$ as $s_{i}$:

$$ ${\bm{s}}=[s_{1},s_{2},\dots,s_{n}]^{T}\coloneqq\left[\sum_{k=1}^{n}l_{k}-l_{1},\sum_{k=1}^{n}l_{k}-l_{2},\dots,\sum_{k=1}^{n}l_{k}-l_{n}\right]^{T}.$ (24) $$

According to eq.˜23, the index $l_{1}$ of column ${\bm{x}}_{1}$ is then analogously mapped to

$$ $l_{1}={\bm{u}}^{T}{\bm{x}}_{1}\quad\mapsto\quad\tilde{l}_{1}\coloneqq{\bm{u}}^{T}\Psi^{(\hat{i})}_{|,1}({\bm{X}})={\bm{u}}^{T}{\bm{x}}_{1}+cs_{1}=l_{1}+cs_{1}.$ (25) $$

Notice that, by choosing $c>1$, we can ensure

$$ $c>1\quad\Longrightarrow\quad\tilde{l}_{1}>\cancel{l_{1}}+\sum_{k=1}^{n}l_{k}-\cancel{l_{1}}>\sum_{k=1}^{n}>l_{i}\quad\forall i,$ (26) $$

and particularly $\tilde{l}_{1}>l_{2}$, 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 $\Psi^{(\hat{i})}$ with $l_{2}={\bm{u}}^{T}{\bm{x}}_{2}=\hat{i}\delta$. At this point, the second column index is updated as follows:

$$ $\begin{split}l_{2}={\bm{u}}^{T}{\bm{x}}_{2}\qquad\mapsto\qquad\tilde{l}_{2}\coloneqq&{\bm{u}}^{T}\Psi^{(\hat{i})}_{|,2}({\bm{X}})={\bm{u}}^{T}{\bm{x}}_{2}+c\left(\sum_{k:l_{k}>l_{2}}l_{k}+\tilde{l}_{1}\right)\\ =&l_{2}+c\left(\sum_{k=1}^{n}l_{k}-l_{2}-\cancel{l_{1}}+\cancel{l_{1}}+cs_{1}\right)=l_{2}+cs_{2}+c^{2}s_{1}\end{split}$ (27) $$

where $\tilde{l}_{1}$ is also included in light of eq.˜26, and we used the definitions eqs.˜25 and 24. Continuing to apply $\Psi^{(i)}({\bm{X}})$, for increasing $i$, and unrolling the recursion, we recover

$$ $\begin{split}\tilde{l}_{3}&=l_{3}+c\left(\sum_{k=1}^{n}l_{k}-l_{1}-l_{2}-l_{3}+\tilde{l}_{1}+\tilde{l}_{2}\right)=l_{3}+cs_{3}+c^{2}(s_{2}+s_{1})+c^{3}s_{1}\\ \tilde{l}_{4}&=l_{4}+c\left(\sum_{k=1}^{n}l_{k}-l_{1}-l_{2}-l_{3}-l_{4}+\tilde{l}_{1}+\tilde{l}_{2}+\tilde{l}_{3}\right)\\ &=l_{4}+cs_{4}+c^{2}(s_{3}+s_{2}+s_{1})+c^{3}(s_{2}+2s_{1})+c^{4}s_{1}\\ \tilde{l}_{5}&=l_{5}+c\left(\sum_{k=1}^{n}l_{k}-l_{1}-l_{2}-l_{3}-l_{4}-l_{5}+\tilde{l}_{1}+\tilde{l}_{2}+\tilde{l}_{3}+\tilde{l}_{4}\right)\\ &=l_{5}+cs_{5}+c^{2}(s_{4}+s_{3}+s_{2}+s_{1})+c^{3}(s_{3}+2s_{2}+3s_{1})+c^{4}(s_{2}+3s_{1})+c^{5}s_{1}\\ &\vdots\end{split}$ (28) $$

which eventually allows us to write the general formula
(^7^77From eq. 28, we can notice that, for a given $\tilde{l}_{k}$, the coefficients $a_{i,j}^{(k)}$ appearing in front of the various $s_{k-i}$ for each of the $c^{j}$ terms, are first given by a list of ones, $a_{i,1}^{(k)}=1$, then a list of increasing numbers $a_{i,2}^{(k)}=i\Longrightarrow a_{-,2}^{(k)}=\text{cumsum}(a_{-,1}^{(k)})$, then a list of triangular numbers $a_{i,3}^{(k)}=i(i+1)/2\Longrightarrow a_{-,3}^{(k)}=\text{cumsum}(a_{-,2}^{(k)})$, and so on: $a_{-,j}^{(k)}=\text{cumsum}(a_{-,j-1}^{(k)})$. The result of iterated applications of cumsum, starting from an all-one vector, can be compactly described via the binomial coefficient: we have in fact $a_{i,j}=[\text{cumsum}^{j}([1,1,\dots])]_{i}=\binom{i+j-2}{j-1}.$ The actual formula eq. 29 can be recovered after a few algebraic steps, by rearranging the summation indices.)

$$ $\begin{split}\tilde{l}_{j}&\coloneqq l_{j}+cs_{j}+\sum_{i=0}^{j-2}c^{i+2}\sum_{k=i}^{j-2}\binom{k}{i}s_{k-i+1},\qquad\qquad j=1\dots n.\end{split}$ (29) $$

#### C.2.3 Result of Applying One Last Global Shift Layer

After the last selective shift layer, the original input ${\bm{X}}$ has been mapped to a modified one $\tilde{{\bm{X}}}$ whereby each column $\tilde{{\bm{x}}}_{j}$ is characterized by the index $\tilde{l}_{j}={\bm{u}}^{T}\tilde{{\bm{x}}}_{j}$ given in eq.˜29. Remember our goal is to recover a contextual mapping, but notice that these $\tilde{l}_{j}$ indices are *not* uniquely defined by the input(^8^88To convince ourselves of this, it suffices to look at the formula for eq. 25: two sequences with different elements ${\bm{l}}\neq{\bm{l}}^{\prime}$, but such that $l_{1}=l_{1}^{\prime}$ and $s_{1}=s_{1}^{\prime}$ (that is, with $\sum_{i=1}^{n}l_{i}=\sum_{i=1}^{n}l_{i}^{\prime}$) would map to the same $\tilde{l}_{1}=\tilde{l}_{1}^{\prime}$.); in other words, they do not satisfy property (2) in Def.˜C.2. The only exception to this is the last index $\tilde{l}_{n}$, as (loosely speaking) it has “seen” all the previous updates - and indeed in Sec.˜C.2.4 we prove this rigorously, under some assumption on the yet-undefined coefficient $c(\delta,d,n)$.

A straightforward way to recover a one-to-one mapping for the whole sequence, then, is to update every index $\tilde{l}_{j}$ via a quantity directly depending on $\tilde{l}_{n}$. This is precisely what the last *global shift* layer $\bar{\Psi}({\bm{X}})$ aims to accomplish. This last layer is also defined starting from the simplified modified sigmoid attention eq.˜18, by picking $b_{k}=0$ and $b_{q}=\left(c(\delta,d,n)^{n}+\frac{1}{2}\right)\delta$: if, for any input, we can guarantee that

$$ $\tilde{l}_{j}\leq c(\delta,d,n)^{n}\delta\quad j<n\qquad\text{and}\qquad\tilde{l}_{n}>c(\delta,d,n)^{n}\delta,$ (30) $$

then the application of the global shift layer would result in(^9^99As in footnote 6, this is also better seen by considering the resulting modified sigmoid attention matrix. With $b_{k}=0$ and $b_{q}=\left(c(\delta,d,n)^{n}+\frac{1}{2}\right)\delta$, in fact, if condition eq. 30 is verified, this matrix is given by $\begin{array}[]{rcl}\begin{array}[]{c}H\left(\left(\tilde{{\bm{l}}}-{\bm{1}}\left(c^{n}+\frac{1}{2}\right)\delta\right)\otimes\tilde{{\bm{l}}}\right)\end{array}=&\left[\begin{array}[]{ccc}0&\cdots&0\[-4.30554pt] \vdots&\ddots&\vdots\\ 0&\cdots&0\\ \hline\cr 1&\cdots&1\\ \end{array}\right]&\begin{array}[]{l}\\ \\ \tilde{l}_{j},\,j<n\\ \\ \tilde{l}_{n}\\ \end{array}\end{array}.$ (31)):

$$ $\begin{split}\bar{\Psi}(\tilde{{\bm{X}}})\coloneqq&\tilde{{\bm{X}}}+c^{n+1}\tilde{\psi}\left(\tilde{{\bm{X}}};b_{q}=\left(c^{n}+\frac{1}{2}\right)\delta,b_{k}=0\right)\\ \Longrightarrow&\bar{\Psi}_{1,j}(\tilde{{\bm{X}}})=\tilde{{\bm{X}}}_{1,j}+c^{n+1}\tilde{l}_{n}\\ \Longrightarrow&\bar{\Psi}_{k>1,j}(\tilde{{\bm{X}}})=\tilde{{\bm{X}}}_{k,j}.\end{split}$ (32) $$

The global shift eq.˜32 is the last layer we need to define our candidate contextual mapping. Collecting the results from this section together, our architecture is defined by sequentially composing the selective shift layers with the global shift one,

$$ $\Psi({\bm{X}})\coloneqq\bar{\Psi}\circ\Psi^{(\delta^{-d}-1)}\circ\cdots\circ\Psi^{(2)}\circ\Psi^{(1)}({\bm{X}}).$ (33) $$

After being scalar-multiplied by ${\bm{u}}$, this results in a sequence

$$ ${\bm{q}}({\bm{X}})\coloneqq{\bm{u}}^{T}\Psi({\bm{X}})=\tilde{{\bm{l}}}+c^{n+1}{\bm{1}}\tilde{l}_{n}$ (34) $$

which we aim to prove is a contextual mapping. This is shown in the next section.

#### C.2.4 A Sequence of Selective Shifts Followed by a Global Shift Produces a Contextual Mapping

To complete the proof, it remains to show that the recovered sequence eq.˜34 represents a contextual mapping and, in particular, that it is *(i)* one-to-one in $\mathbb{L}$, and that *(ii)* all of its elements are distinct for different inputs. To do so, we need a few preparatory lemmas. The first few are needed to show that each of the basic components of eq.˜34 is indeed a one-to-one map.

###### Lemma C.3 .

The map ${\bm{l}}\mapsto{\bm{s}}$ in eq.˜24 is one-to-one.

###### Proof.

The target map can be compactly represented as a linear operator $S$:

$$ ${\bm{l}}\mapsto{\bm{s}}\coloneqq{\bm{1}}\sum_{k=1}^{n}l_{k}-{\bm{l}}=({\bm{1}}\otimes{\bm{1}}-I){\bm{l}}\eqqcolon S{\bm{l}}$ (35) $$

which is invertible(^10^1010Indeed its inverse can be explicitly recovered by directly applying Sherman-Morrison formula.), denoting that ${\bm{l}}\mapsto{\bm{s}}$ is bijective.
∎

###### Lemma C.4 .

The map ${\bm{l}}\mapsto\tilde{l}_{n}$ in eq.˜29 is one-to-one, under the condition

$$ $c(\delta,d,n)>(n-1)(\delta^{-d}-1)\binom{n-1}{\left\lceil\frac{n-1}{2}\right\rceil}.$ (36) $$

###### Proof.

Consider two vectors of column indices ${\bm{l}},{\bm{l}}^{\prime}$ differing for at least one element. We have by definition eq.˜29 that

$$ $\tilde{l}_{n}-\tilde{l}_{n}^{\prime}=(l_{n}-l_{n}^{\prime})+c(s_{n}-s_{n}^{\prime})+\sum_{i=0}^{n-2}c^{i+2}\sum_{k=i}^{n-2}\binom{k}{i}(s_{k-i+1}-s_{k-i+1}^{\prime})$ (37) $$

By absurd, assume $\tilde{l}_{n}-\tilde{l}_{n}^{\prime}=0$ even though $\exists i:l_{i}\neq l_{i}^{\prime}$. We have then that it must hold

$$ $\begin{split}(l_{n}^{\prime}-l_{n})&=c(s_{n}-s_{n}^{\prime})+\sum_{i=0}^{n-2}c^{i+2}\sum_{k=i}^{n-2}\binom{k}{i}(s_{k-i+1}-s_{k-i+1}^{\prime})\\ &=c\left((s_{n}-s_{n}^{\prime})+\sum_{i=0}^{n-2}c^{i+1}\sum_{k=i}^{n-2}\binom{k}{i}(s_{k-i+1}-s_{k-i+1}^{\prime})\right)\end{split}$ (38) $$

Notice that, for $c(\delta,d,n)$ large enough, the right-hand side does not have enough *granularity* to counter the left-hand side: in fact, since $l_{n}\in\{0,\delta,2\delta,\dots,\delta^{-d+1}-\delta\}$, the left-hand side can attain values

$$ $l_{n}^{\prime}-l_{n}\in\{0,\pm\delta,\pm 2\delta,\dots,\pm(\delta^{-d+1}-\delta)\}$ (39) $$

while the former, in light of the presence of the $c(\delta,d,n)$ factor, can only attain values $\in\{0,\pm c\delta,\pm 2c\delta,\dots\}$. Picking $c>\delta^{-d}-1$, then, ensures that equality between the two sides of eq.˜38 can only be achieved if they are both 0. In this case, we need to impose

$$ $\begin{split}&c(s_{n}^{\prime}-s_{n})=\sum_{i=0}^{n-2}c^{i+1}\sum_{k=i}^{n-2}\binom{k}{i}(s_{k-i+1}-s_{k-i+1}^{\prime})\\ \quad\Longleftrightarrow\quad&s_{n}^{\prime}-s_{n}=c\left(\sum_{i=0}^{n-2}c^{i}\sum_{k=i}^{n-2}\binom{k}{i}(s_{k-i+1}-s_{k-i+1}^{\prime})\right).\end{split}$ (40) $$

Similarly, notice that(^11^1111This is a direct consequence of the definition of operator $S$ in eq. 35: since it has $1$’s everywhere but on its diagonal, its $\infty$-norm is simply $n-1$.), $\forall i$,

$$ $|s_{i}-s_{i}^{\prime}|=\left|\sum_{k=1}^{n}(l_{k}-l_{k}^{\prime})-(l_{i}-l_{i}^{\prime})\right|=\left|\sum_{k=1,k\neq i}^{n}(l_{k}-l_{k}^{\prime})\right|<(n-1)(\delta^{-d+1}-\delta),$ (41) $$

implying that $s_{n}^{\prime}-s_{n}\in\{0,\pm\delta,\pm 2\delta,\dots,\pm(n-1)(\delta^{-d}-1)\delta\}$. Again, by picking $c(\delta,d,n)>(n-1)(\delta^{-d}-1)$ we ensure that the right-hand side does not have enough granularity, and hence

$$ $c(\delta,d,n)>(n-1)(\delta^{-d}-1)\qquad\Longrightarrow\qquad s_{n}^{\prime}-s_{n}=0,$ (42) $$

implying

$$ $\begin{split}&c\left(\sum_{i=0}^{n-2}c^{i}\sum_{k=i}^{n-2}\binom{k}{i}(s_{k-i+1}-s_{k-i+1}^{\prime})\right)=0\\ \Longleftrightarrow\quad&\sum_{k=0}^{n-2}\binom{k}{0}(s_{k+1}^{\prime}-s_{k+1})=c\left(\sum_{i=1}^{n-2}c^{i-1}\sum_{k=i}^{n-2}\binom{k}{i}(s_{k-i+1}-s_{k-i+1}^{\prime})\right)\\ \Longleftrightarrow\quad&\sum_{k=0}^{n-2}(s_{k+1}^{\prime}-s_{k+1})=c\left(\sum_{i=1}^{n-2}c^{i-1}\sum_{k=i}^{n-2}\binom{k}{i}(s_{k-i+1}-s_{k-i+1}^{\prime})\right).\\ \end{split}$ (43) $$

Following a similar reasoning as the one applied above shows us that picking

$$ $c(\delta,d,n)>(n-1)^{2}(\delta^{-d}-1)\qquad\Longrightarrow\qquad\sum_{k=0}^{n-2}(s_{k+1}-s_{k+1}^{\prime})=0,$ (44) $$

and requires us to satisfy

$$ $\begin{split}&c\left(\sum_{i=1}^{n-2}c^{i-1}\sum_{k=i}^{n-2}\binom{k}{i}(s_{k-i+1}-s_{k-i+1}^{\prime})\right)=0\\ \Longleftrightarrow\quad&\sum_{k=1}^{n-2}\binom{k}{1}(s_{k}^{\prime}-s_{k})=c\left(\sum_{i=2}^{n-2}c^{i-2}\sum_{k=i}^{n-2}\binom{k}{i}(s_{k-i+1}-s_{k-i+1}^{\prime})\right)\\ \Longleftrightarrow\quad&\sum_{k=1}^{n-2}k(s_{k}^{\prime}-s_{k})=c\left(\sum_{i=2}^{n-2}c^{i-2}\sum_{k=i}^{n-2}\binom{k}{i}(s_{k-i+1}-s_{k-i+1}^{\prime})\right).\\ \end{split}$ (45) $$

Once again, then, by choosing

$$ $c(\delta,d,n)>\frac{(n-2)(n-1)^{2}}{2}(\delta^{-d}-1)\qquad\Longrightarrow\qquad\sum_{k=1}^{n-2}k(s_{k}-s_{k}^{\prime})=0.$ (46) $$

This reasoning can be repeated recursively: at each step $i$ of the recursion, by imposing a stricter and stricter bound on $c(\delta,d,n)$ we gain more and more conditions that the quantity ${\bm{s}}^{\prime}-{\bm{s}}$ needs to satisfy:

$$ $c(\delta,d,n)>(n-1)(\delta^{-d}-1)\sum_{k=i}^{n-2}\binom{k}{i}\qquad\Longrightarrow\qquad\sum_{k=i}^{n-2}\binom{k}{i}(s_{k-i+1}-s_{k-1+1}^{\prime})=0.$ (47) $$

Notice that, every time we increase $i=0\dots n-2$, these conditions involve one less term $s_{k-i+1}-s_{k-i+1}^{\prime}$, $k=i\dots n-2$: if we were to collect all these conditions within a single linear system, the system would have an upper-triangular structure, and hence be non-singular. This implies that for the set of $n$ independent conditions on ${\bm{s}}-{\bm{s}}^{\prime}$ to hold (we have $n-1$ in eq.˜47, plus one more in eq.˜42), the only possibility is that ${\bm{s}}\equiv{\bm{s}}^{\prime}$. Because of Lemma˜C.3, though, this also implies ${\bm{l}}\equiv{\bm{l}}^{\prime}$: we have finally reached a contradiction, and proven that indeed ${\bm{l}}\mapsto\tilde{l}_{n}$ is one-to-one, under an opportune condition on $c(\delta,d,n)$.
Such condition can be promptly recovered(^12^1212This is a consequence of some useful properties of the binomial coefficient, namely the Hockey stick identity Jones (), and the symmetry of $\binom{k}{i}$ with respect to $i$.) by eq.˜47:

$$ $\max_{i=0\dots n-2}\sum_{k=i}^{n-2}\binom{k}{i}=\max_{i=0\dots n-2}\binom{n-1}{i+1}=\binom{n-1}{\left\lceil\frac{n-1}{2}\right\rceil}.$ (48) $$

Substituting this in eq.˜47, we recover that it suffices to impose

$$ $c(\delta,d,n)>(n-1)(\delta^{-d}-1)\binom{n-1}{\left\lceil\frac{n-1}{2}\right\rceil}.$ (49) $$

∎

The next few lemmas are needed to bound the elements in the $\tilde{l}_{j}$ sequence, which in turn are used to prove property *(ii)* in Def.˜C.2.

###### Lemma C.5 .

$\tilde{l}_{j}$ in eq.˜29 is an increasing sequence.

###### Proof.

This can be proven directly: we have in fact, by definition eq.˜29,

$$ $\begin{split}\tilde{l}_{j}>\tilde{l}_{j-1}\quad\Longleftrightarrow\quad&l_{j}+cs_{j}+\sum_{i=0}^{j-2}c^{i+2}\sum_{k=i}^{j-2}\binom{k}{i}s_{k-i+1}\\ &>l_{j-1}+cs_{j-1}+\sum_{i=0}^{j-3}c^{i+2}\sum_{k=i}^{j-3}\binom{k}{i}s_{k-i+1}\\ \prescript{\text{\tiny combine sums}}{}{\quad\Longleftrightarrow\quad}&(l_{j}-l_{j-1})(1-c)+\sum_{i=0}^{j-2}c^{i+2}\binom{j-2}{i}s_{j-1-i}>0\\ \prescript{\text{\tiny$\binom{j-2}{i}\geq 1,c^{i+2}\geq c^{2}$}}{}{\quad\Longleftarrow\quad}&(l_{j}-l_{j-1})(1-c)+c^{2}\sum_{i=0}^{j-2}s_{j-1-i}>0\\ \prescript{\text{\tiny\lx@cref{creftype~refnum}{eqn::partial_sum}}}{}{\quad\Longleftrightarrow\quad}&(l_{j}-l_{j-1})(1-c)+c^{2}\sum_{i=0}^{j-2}\left(\sum_{k=1}^{n}l_{k}-l_{j-1-i}\right)>0\\ \quad\Longleftrightarrow\quad&(l_{j}-l_{j-1})(1-c)+c^{2}\left((j-1)\sum_{k=1}^{n}l_{k}-\sum_{k=1}^{j-1}l_{k}\right)>0\\ \quad\Longleftrightarrow\quad&(1-c)l_{j}+(c-1)l_{j-1}+c^{2}(j-2)\sum_{k=1}^{n}l_{k}+c^{2}\sum_{k=j}^{n}l_{k}>0\\ \quad\Longleftrightarrow\quad&(c^{2}-c+1)l_{j}+(c-1)l_{j-1}+c^{2}(j-2)\sum_{k=1}^{n}l_{k}+c^{2}\sum_{k=j+1}^{n}l_{k}>0\\ \end{split}$ (50) $$

Already with $c>1$, all the coefficients are positive (and at least one is non-zero), implying that the condition above is always satisfied and that indeed $\tilde{l}_{j}$ is an increasing sequence.
∎

###### Lemma C.6 .

Under constraint eq.˜36, each term $\tilde{l}_{j}$, $j>1$ in eq.˜29 is bounded from below by

$$ $\tilde{l}_{j}>c^{j}\delta,$ $$

and each term $\tilde{l}_{j}$, $1<j<n$ is bounded from above by

$$ $\tilde{l}_{j}<c^{j+1}\delta.$ $$

###### Proof.

We start by proving the lower bound. By definition eq.˜29, we have

$$ $\tilde{l}_{j}=l_{j}+cs_{j}+\sum_{i=0}^{j-2}c^{i+2}\sum_{k=i}^{j-2}\binom{k}{i}s_{k-i+1}=l_{j}+cs_{j}+c^{j}s_{1}+\sum_{i=0}^{j-3}c^{i+2}\sum_{k=i}^{j-2}\binom{k}{i}s_{k-i+1}.$ (51) $$

Since by assumption $l_{j}$ is an ordered sequence without repetitions, for $j>1$ we necessarily have $l_{j}>l_{1}\geq 0$, and hence $l_{j}\geq\delta$. All the other terms in eq.˜51 are non-negative, so we can safely claim that

$$ $\tilde{l}_{j}\geq\delta+c^{j}\delta>c^{j}\delta\qquad\forall j>1,$ (52) $$

which confirms the lower bound.

For the upper bound, we start again from the definition of $\tilde{l}_{j}$:

$$ $\begin{split}\tilde{l}_{j}&=l_{j}+cs_{j}+\sum_{i=0}^{j-2}c^{i+2}\sum_{k=i}^{j-2}\binom{k}{i}s_{k-i+1}\\ &<(\delta^{-d}-1)\delta+c(n-1)(\delta^{-d}-1)\delta+s_{1}\sum_{i=0}^{j-2}c^{i+2}\binom{j-1}{i+1}\\ &\leq(n-1)(\delta^{-d}-1)\binom{j-1}{\left\lceil\frac{j-1}{2}\right\rceil}\delta\sum_{i=0}^{j}c^{i}=(n-1)(\delta^{-d}-1)\binom{j-1}{\left\lceil\frac{j-1}{2}\right\rceil}\delta\frac{1-c^{j+1}}{1-c},\end{split}$ (53) $$

where we used relationship eq.˜48 and collected all $c$ terms within the sum. Notice that, for a given $a>1$ we have that

$$ $\frac{1-c^{j+1}}{1-c}\leq ac^{j},$ (54) $$

provided that $c\geq\frac{a}{a-1}$. In fact,

$$ $\begin{split}\frac{1-c^{j+1}}{1-c}\leq ac^{j}\quad\Longleftrightarrow\quad&\frac{1-c^{j+1}-ac^{j}+ac^{j+1}}{1-c}\leq 0\\ \quad\Longleftarrow\quad&\frac{1}{a-1}+\left(c-\frac{a}{a-1}\right)c^{j}\geq 0\quad\Longleftarrow\quad\frac{1}{a-1}\geq 0\\ \end{split}$ (55) $$

which is always satisfied. After substituting eq.˜54 in eq.˜53, this allows us to write

$$ $\tilde{l}_{j}<a(n-1)(\delta^{-d}-1)\binom{j-1}{\left\lceil\frac{j-1}{2}\right\rceil}\delta c^{j}.$ (56) $$

To prove that $\tilde{l}_{j}<\delta c^{j+1}$, then, it remains to show that

$$ $c\geq a(n-1)(\delta^{-d}-1)\binom{j-1}{\left\lceil\frac{j-1}{2}\right\rceil}\qquad\forall 1<j<n.$ (57) $$

Substituting condition eq.˜36 in the inequality above, we are left with proving

$$ $\binom{n-1}{\left\lceil\frac{n-1}{2}\right\rceil}\geq\max_{j=2\dots n-1}a\binom{j-1}{\left\lceil\frac{j-1}{2}\right\rceil}=a\binom{n-2}{\left\lceil\frac{n-2}{2}\right\rceil}.$ (58) $$

The outcome depends on the parity of $n$. For $n$ odd, we have

$$ $\binom{n-1}{\left\lceil\frac{n-1}{2}\right\rceil}\geq a\binom{n-2}{\left\lceil\frac{n-2}{2}\right\rceil}\qquad\Longleftrightarrow\qquad 2\frac{n-1}{n-1}\geq a,$ (59) $$

to satisfy which it suffices to pick $a=2$. This requires having $c\geq\frac{a}{a-1}=2$, which is automatically satisfied. For $n$ even, on the other hand, the binomial coefficients simplify to

$$ $\binom{n-1}{\left\lceil\frac{n-1}{2}\right\rceil}\geq a\binom{n-2}{\left\lceil\frac{n-2}{2}\right\rceil}\qquad\Longleftrightarrow\qquad 2\frac{n-1}{n}\geq a.$ (60) $$

To satisfy this, we need to pick $a=2\frac{n-1}{n}$, which requires $c\geq\frac{a}{a-1}=2\frac{n-1}{n-2}$; however, this too is automatically satisfied by eq.˜36 provided $n\geq 4$. This completes the proof.
∎

###### Lemma C.7 .

Under the constraint eq.˜36, condition eq.˜30 holds.

###### Proof.

We remind that condition eq.˜30 is necessary for the correct “functioning” of the global shift layer, and it composes of two parts.
The first part requires that $\tilde{l}_{j}<c^{n}\delta$ $\forall j<n$. Thanks to Lemma˜C.5, it suffices to show that $\tilde{l}_{n-1}<c^{n}\delta$, but this is already granted by the upper bound in Lemma˜C.6. Analogously, for the second part, we need to show that $\tilde{l}_{n}>c^{n}\delta$: for this too we can use the lower bound in Lemma˜C.6.
∎

We finally have all the ingredients to prove the main theorem of this section:

###### Theorem C.8 .

The map in eq.˜34, given by

$$ ${\bm{X}}\mapsto{\bm{q}}({\bm{X}})={\bm{u}}^{T}\Psi({\bm{X}})$ $$

represents a contextual mapping.

###### Proof.

As defined in Def.˜C.2, a contextual mapping must satisfy two conditions. The first one is that

$$ $q_{i}({\bm{X}})\neq q_{j}({\bm{X}}),\quad\forall i\neq j\qquad\text{and}\qquad\forall{\bm{X}}\in\mathbb{L}.$ (61) $$

This is directly proven by considering Lemma˜C.5: since $\tilde{l}_{j}$ is a (strictly) increasing sequence, all its elements are already distinct. The action of the last global shift layer merely translates all these elements by a same quantity, but they remain distinct nonetheless.

The second condition for a contextual mapping is given by

$$ $q_{i}({\bm{X}})\neq q_{j}({\bm{X}}^{\prime}),\quad\forall i,j\quad\text{and}\quad\forall{\bm{X}},{\bm{X}}^{\prime}\in\mathbb{L},\quad\text{with}\quad{\bm{X}}\neq{\bm{X}}^{\prime}.$ (62) $$

We prove that this holds for eq.˜34 by directly considering the difference between two components $i,j$ for different inputs:

$$ $q_{i}({\bm{X}})-q_{j}({\bm{X}}^{\prime})=\tilde{l}_{i}-\tilde{l}_{j}^{\prime}+c^{n+1}\left(\tilde{l}_{n}-\tilde{l}_{n}^{\prime}\right)=0\quad\Longleftrightarrow\quad\tilde{l}_{i}-\tilde{l}_{j}^{\prime}=c^{n+1}\left(\tilde{l}_{n}^{\prime}-\tilde{l}_{n}\right).$ (63) $$

Notice that, due to Lemma˜C.4, we have $\tilde{l}_{n}-\tilde{l}_{n}^{\prime}\neq 0$ and particularly, $|\tilde{l}_{n}-\tilde{l}_{n}^{\prime}|\geq\delta$. On the other hand, in light of the bounds in Lemma˜C.6, we have that the left-hand side $|\tilde{l}_{j}-\tilde{l}_{i}|<c^{n}\delta$. Consequently, the two sides can never cancel each other out, and the proof is complete.
∎

###### Lemma C.3 .

###### Proof.

###### Lemma C.4 .

###### Proof.

###### Lemma C.5 .

###### Proof.

###### Lemma C.6 .

###### Proof.

###### Lemma C.7 .

###### Proof.

###### Theorem C.8 .

###### Proof.

## Appendix D Lipschitzness of Sigmoid Attention

In the following, we report the proof for the recovering the Lipschitzness constant associated with $\mathrm{SigmoidAttn}$, as stated in Thm.˜3.2.

Letting $A=W_{q}^{T}W_{k}$, and calling $\sigma_{ij}=\sigma(\langle W_{q}x_{i},W_{k}x_{j}\rangle)$ and $\sigma^{\prime}_{ij}=\sigma^{\prime}(\langle W_{q}x_{i},W_{k}x_{j}\rangle)$, we find that the Jacobian of $\phi$ in the direction $(\delta_{1},\dots,\delta_{n})$ for the sample $x_{i}$ is given by:

$$ $\mathrm{Jac}_{i}=\left(\sum_{j=1}^{n}\sigma^{\prime}_{ij}x_{j}x_{j}^{T}A^{T}\right)\delta_{i}+\sum_{j=1}^{n}\left(\sigma^{\prime}_{ij}x_{j}x_{i}^{T}A+\sigma_{ij}I_{p}\right)\delta_{j},$ (64) $$

We see that this Jacobian is the sum of two terms. To control its norm, we can control each norm individually.

The first term, $\left(\sum_{j=1}^{n}\sigma^{\prime}_{ij}x_{j}x_{j}^{T}A^{T}\right)\delta_{i}$ is of the form $U_{i}\delta_{i}$ with $U_{i}$ a matrix.
Its squared-norm is therefore:

$$ $\sum_{i=1}^{n}\|U_{i}\delta_{i}\|^{2}\leq\max_{i}\|U_{i}\|_{2}^{2}\|\delta\|_{F}.$ (65) $$

Hence, its squared spectral norm is bounded by $\max_{i}\|U_{i}\|_{2}^{2}$.

We now let $\sigma^{\prime}_{\infty}$ be a bound on $n\times|\sigma^{\prime}|$;
We have:

$$ $\displaystyle\|U_{i}\|_{2}$ $\displaystyle\leq\sum_{j=1}^{n}\|\sigma^{\prime}_{ij}x_{j}x_{j}^{\top}A\|_{2}$ (66) $\displaystyle\leq\sigma^{\prime}_{\infty}\|A\|_{2}\frac{1}{n}\sum_{j=1}^{n}\|x_{j}\|^{2}$ (67) $\displaystyle\leq\sigma^{\prime}_{\infty}\|A\|_{2}\mathbb{E}[\|x_{j}\|^{2}].$ (68) $$

We see that if the points $x_{i}$ have norm $\leq R$, then the Jacobian grows at most like $R^{2}$, because it is “quadratic” in $x$.
However, we see that the quadratic term is likely to be mitigated by the $\sigma^{\prime}(a_{ij})$ term that goes to 0 if $a_{ij}$ is large.

The second term, $\sum_{j=1}^{n}\left(\sigma^{\prime}_{ij}x_{j}x_{i}^{T}A+\sigma_{ij}I_{p}\right)\delta_{j}$, is the sum of two terms. Here, too, we use the triangular inequality to control their norm individually. We get:

$$ $\displaystyle\|\sum_{j=1}^{n}\sigma_{ij}\delta_{j}\|^{2}$ $\displaystyle=\|\delta^{T}\sigma_{i}\|^{2}$ (69) $\displaystyle\leq\|\delta\|_{F}^{2}\|\sigma_{i}\|^{2},$ (70) $$

where $\sigma_{i}\in\mathbb{R}^{p}$ is the $i$-th column of $\sigma_{ij}$, and $\delta\in\mathbb{R}^{n\times p}$.
and by summing, letting $\sigma_{\infty}$ an upper bound on $n\times|\sigma(x)|$:

$$ $\sum_{i=1}^{n}\|\sum_{j=1}^{n}\sigma_{ij}\delta_{j}\|^{2}\leq\sigma_{\infty}^{2}\|\delta\|_{F}^{2}.$ (71) $$

So that $\sigma_{\infty}$ upper bounds the spectral norm of the last term.

For the final term, $\sum_{j=1}^{n}\sigma^{\prime}_{ij}x_{j}x_{i}^{T}A\delta_{j}$, define $\hat{\delta}=\delta A^{T}$. We get:

$$ $\sum_{j=1}^{n}\sigma^{\prime}_{ij}x_{j}x_{i}^{T}A\delta_{j}=\sum_{j=1}^{n}\sigma^{\prime}_{ij}\langle x_{i},\hat{\delta}_{j}\rangle x_{j}.$ (72) $$

Hence, letting $M$ the matrix of entries $M_{ij}=\sigma^{\prime}_{ij}\langle x_{i},\hat{\delta}_{j}\rangle$, we see that the previous term is simply $x^{T}M_{i}^{T}$, so that we get the upper bound on the norm of the term:

$$ $\sum_{i=1}^{n}\|x^{T}M_{i}^{T}\|^{2}\leq\|x\|_{F}^{2}\|M\|_{F}^{2}$ (73) $$

and $\|M\|_{F}^{2}=\sum_{ij}(\sigma^{\prime}_{ij})^{2}\langle x_{i},\hat{\delta}_{j}\rangle^{2}\leq\frac{1}{n^{2}}\sigma^{\prime}_{\infty}\|x\|_{F}^{2}\|A\|_{2}^{2}\|\delta\|_{F}^{2}$, giving overall:

$$ $\sqrt{\sum_{i=1}^{n}\|x^{T}M_{i}^{T}\|^{2}}\leq\sigma^{\prime}_{\infty}\|A\|_{2}\mathbb{E}[\|x_{j}\|^{2}]\|\delta\|_{F}.$ (74) $$

Notice how this quantity matches the one in eq.˜68.

Finally, summing all together gives:

$$ $\|\mathrm{Jac}\|_{2}\leq 2\sigma^{\prime}_{\infty}\|A\|_{2}\mathbb{E}[\|x_{j}\|^{2}]+\sigma_{\infty},$ (75) $$

which completes the proof.

Remark: The previous upper bound might not be tight. Indeed, intuitively, if the $x_{i}$ are large, then the term $\sigma^{\prime}_{ij}$ should be exponentially small (provided, of course, that $W_{q}x_{i}$ and $W_{k}x_{j}$ are not orthogonal), which would even remove the dependency on the variance in the sigmoid attention.

## Appendix E The Bias Term of Sigmoid Attention

One of the differences between $\mathrm{SigmoidAttn}$ and $\mathrm{SoftmaxAttn}$ is the normalization constant. In $\mathrm{SigmoidAttn}$, one way to emulate the effect of a normalization constant (which links all the elements of the input together and defines a distribution over them), is to include a bias term in the definition as proposed in eq. 3.

For an input vector $\bm{z}\in\mathbb{R}^{n}$, the output of the sigmoid with bias $b$ is

$$ $\sigma^{b}(\bm{z})_{i}:=\frac{\exp(z_{i})}{\exp(z_{i})+\exp(-b)}$ $$

Contrary to the softmax, this output cannot always sum to one because there is no normalization.
We therefore seek a value for $b$ that *approximately* normalizes $\sigma^{b}(\bm{z})$, i.e., such that $\sum_{i=1}^{n}\sigma^{b}(\bm{z})_{i}\simeq 1$.
We have

###### Proposition E.1 .

Let $\bm{z}\in\mathbb{R}^{n}$, and take $m,M\in\mathbb{R}$ such that for all $i$, it holds $m\leq z_{i}\leq M$. Then, the equation $\sum_{i=1}^{n}\sigma^{b}(\bm{z})_{i}=1$ with variable $b$ has a single solution $b^{*}$ with

$$ $-\log(n-1)-M\leq b^{*}\leq-\log(n-1)-m\kern 5.0pt.$ $$

###### Proof.

The function $\phi:b\to\sum_{i=1}^{n}\sigma^{b}(\bm{z})_{i}$ is smooth and monotonically increasing, and we have $\phi(-\log(n-1)-M)\leq 1$ and $\phi(-\log(n-1)-m)\geq 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 a 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^{\star}$, where $p_{i}^{\star}=\frac{e^{z_{i}}}{e^{z_{i}}+\sum_{j\neq i}e^{z_{j}}}$. The goal is to find the bias term $b$ such that sigmoid function with weights over all elements denoted by $p$, where $p_{i}=\sigma^{b}(\bm{z})_{i}$, approximates $p^{\star}$. Note that, as mentioned before, $p$ is not necessarily a distribution, i.e. $\sum_{i=1}^{n}p_{i}$ is not always equal to one.

In technical terms, we aim to estimate the normalizing factor $\bm{Z}=\sum_{i=1}^{n}e^{z_{i}}$. The existing approaches for estimating $\bm{Z}$ is compute-expensive for high dimensions and requires resampling methods.
Also, the optimal value of $b$ would depend on the exact values of $\bm{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 $\mathrm{SigmoidAttn}$, we assume that each element has uniform weight for the model apriori, which means that none of the elements of the input vector $\bm{z}$ has any known importance over the others. In the simplest case when softmax is a uniform distribution, we ideally want to have the same order of values for sigmoid as of softmax, which should be $\frac{1}{n}$. Therefore, we can write down the following:

$$ $\displaystyle\forall i\qquad p_{i}=\frac{1}{1+e^{-(z_{i}+b)}}\simeq\frac{1}{n}=p_{i}^{\star}$ (76) $$

Ideally, we would like to have $1+e^{-(z_{i}+b)}\simeq n$. Requiring that $p=p^{*}$ in the case where all the $z_{i}$ are 0 gives $\exp(-b)=n-1$, i.e. $b\simeq-\log(n)$ for large $n$. In the case that all the $z_{i}$ are bounded, $|z_{i}|\leq M<\infty$ for some constant $M$, then $b\simeq-(M+\log(n))\approx-\max\{M,\log(n)\}$. However, 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_{i}$ are sampled from a standard Gaussian distribution, i.e. $z_{i}\sim\mathcal{N}(0,\sigma^{2})$ where $\sigma=1$. Note that this assumption comes from the fact that $z_{i}$ in our problem is one of the elements of $\bm{Q}\bm{K}^{T}/\sqrt{d_{qk}}$, which is the sum of $d_{qk}$ random variables. Using Central Limit Theorem, we can assume that $z_{i}$ is sampled from a Gaussian distribution. The idea is to estimate $M$, such that with high probability, $|z_{i}|\leq M$, i.e. $\mathbb{P}\left(|z_{i}|>M\right)\leq\epsilon$ for a desired $\epsilon$. Therefore, we have

$$ $\displaystyle\mathbb{P}\left(|z_{i}|>M\right)=\mathbb{P}\left(|z_{i}|>\frac{M}{\sigma}\sigma\right)\leq\frac{1}{(\frac{M}{\sigma})^{2}}=\frac{\sigma^{2}}{M^{2}}\leq\epsilon,$ (77) $$

where the inequality is resulted from Chebychev’s inequality. Setting $\sigma=1$, we have $M\simeq\sqrt{1/\epsilon}$. Therefore, the order-optimal value would be $b\simeq-\max\{\sqrt{1/\epsilon},\log(n)\}$, and for long sequence length, $b\simeq-\log(n)$. For example, if we want $90\%$ accuracy in our estimation, $M\approx 3\sigma=3$, which means $b\simeq-\max\{3,\log(n)\}$.
Note that this approximation also follows the intuition that as $n$ grows, we expect the $\mathrm{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^{\star}$ and $p$ with respect to $z_{i}$ to behave similarly. We show that $b\simeq-\log(n)$ is still a good choice in this scenario. Let us derive the derivative of $\mathrm{SigmoidAttn}$ and $\mathrm{SoftmaxAttn}$ with respect to the input. We note that for any $i$, both functions can be written as $\frac{e^{z_{i}}}{e^{z_{i}}+\bm{Z}_{-i}}$ where $\bm{Z}_{-i}$ is the share of normalization factor except element $i$ of $\bm{z}$. For $\mathrm{SoftmaxAttn}$, $\bm{Z}_{-i}=\sum_{j\neq i}e^{z_{j}}$ and for $\mathrm{SigmoidAttn}$, $\bm{Z}_{-i}=e^{-b}$. Now, we have

$$ $\displaystyle\frac{\partial}{\partial z_{i}}\frac{e^{z_{i}}}{e^{z_{i}}+\bm{Z}_{-i}}=\frac{e^{z_{i}}\bm{Z}_{-i}}{\left(e^{z_{i}}+\bm{Z}_{-i}\right)^{2}}.$ (78) $$

Therefore, we have the following

$$ $\displaystyle\frac{\partial p_{i}^{\star}}{\partial z_{i}}$ $\displaystyle=p_{i}^{\star}(1-p_{i}^{\star})$ (79) $\displaystyle\frac{\partial p_{i}}{\partial z_{i}}$ $\displaystyle=p_{i}(1-p_{i}).$ (80) $$

We can see that if $p_{i}\simeq p_{i}^{\star}$, then $\frac{\partial p_{i}}{\partial z_{i}}\simeq\frac{\partial p_{i}^{\star}}{\partial z_{i}}$. So, the previous choice of bias term $b\simeq-\log(n)$ approximates the order of gradients as well. In fact, this is the only valid choice even though we have a quadratic term.

$$ $\displaystyle\frac{\partial p_{i}}{\partial z_{i}}\simeq\frac{\partial p_{i}^{\star}}{\partial z_{i}}$ $\displaystyle\quad\Longleftrightarrow\quad p_{i}^{\star}(1-p_{i}^{\star})=p_{i}(1-p_{i})$ (81) $\displaystyle\quad\Longleftrightarrow\quad\left(p_{i}-p_{i}^{\star}\right)\left(p_{i}-(1-p_{i}^{\star})\right)=0.$ (82) $$

Which means either $p_{i}\simeq p_{i}^{\star}$ or $p_{i}\simeq 1-p_{i}^{\star}$. The first one provides us with $b\simeq-\log(n)$ while the second one cannot happen since the nominator of $p_{i}$ is dependent on $z_{i}$ while the nominator of $1-p_{i}^{\star}$ is independent of $z_{i}$.

###### Proposition E.1 .

###### Proof.

## Appendix 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 a considerable performance boost in model inference over those of FlashAttention2 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 of training vision transformers (Dosovitskiy et al., ) on the ImageNet dataset (Deng et al., ).
Finally, we also provide kernel benchmarking details of FlashSigmoid implementation by taking into account ALiBi slopes (Press et al., ), which is one of the important components of $\mathrm{SigmoidAttn}$ as seen in the main text of the paper.

### F.1 Details of FlashSigmoid Algorithm

Softmax vs. Sigmoid Attention: In this subsection, we discuss the implementation details of FlashSigmoid algorithm, which is a hardware-aware implementation of $\mathrm{SigmoidAttn}$ approach.
We begin with the expressions of the forward and backward passes of softmax and sigmoid attention mechanisms.
Let $\bm{Q}$, $\bm{K}$, and $\bm{V}$ represent the query, key, and value tensors.
Then, the desired forward and backward pass expressions are reported in Tab.˜4.

**Table 4: Description of the forward and backward passes of softmax and sigmoid attention. With $\odot$, we denote Hadamard (element-wise) multiplication.**
| Softmax | Sigmoid |  |  |
| --- | --- | --- | --- |
| Forward | Backward | Forward | Backward |
| $\displaystyle\bm{S}=\frac{\bm{Q}\cdot\bm{K}^{\top}}{\sqrt{d}}$ | $\displaystyle\mathbf{d}\bm{V}=\bm{P}^{\top}\cdot\mathbf{d}\bm{O}$ | $\displaystyle\bm{S}=\frac{\bm{Q}\cdot\bm{K}^{\top}}{\sqrt{d}}$ | $\displaystyle\mathbf{d}\bm{V}=\bm{P}^{\top}\cdot\mathbf{d}\bm{O}$ |
| $\displaystyle\bm{P}=\textrm{softmax}\left(\bm{S}\right)$ | $\displaystyle\mathbf{d}\bm{P}=\mathbf{d}\bm{O}\cdot\bm{V}^{\top}$ | $\displaystyle\bm{P}=\sigma\left(\bm{S}\right)$ | $\displaystyle\mathbf{d}\bm{P}=\mathbf{d}\bm{O}\cdot\bm{V}^{\top}$ |
| $\displaystyle\bm{O}=\bm{P}\cdot\bm{V}$ | $\displaystyle\mathbf{d}\bm{S}=\bm{P}\odot\left(\mathbf{d}\bm{P}-\textrm{rowsum}\left(\mathbf{d}\bm{O}\odot\bm{O}\right)\right)$ | $\bm{O}=\bm{P}\cdot\bm{V}$ | $\displaystyle\mathbf{d}\bm{S}=\bm{P}\odot\left(1-\bm{P}\right)\odot\mathbf{d}\bm{P}$ |
|  | $\displaystyle\mathbf{d}\bm{Q}=\sqrt{d}\cdot\mathbf{d}\bm{S}\cdot\bm{K}$ |  | $\displaystyle\mathbf{d}\bm{Q}=\sqrt{d}\cdot\mathbf{d}\bm{S}\cdot\bm{K}$ |
|  | $\displaystyle\mathbf{d}\bm{K}=\sqrt{d}\cdot\mathbf{d}\bm{S}^{\top}\cdot\bm{Q}$ |  | $\displaystyle\mathbf{d}\bm{K}=\sqrt{d}\cdot\mathbf{d}\bm{S}^{\top}\cdot\bm{Q}$ |

Figure: Algorithm 1 FlashSigmoid Forward Pass

Figure: Algorithm 2 FlashSigmoid Backward Pass

The application of sigmoid and softmax activation functions, as highlighted in orange color in Tab.˜4, is the only implementation difference in the forward passes.
Similarly, the expressions for the gradients of the preactivation ($\mathbf{d}\bm{S}$), 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 as an extension of the FlashAttention2 (Dao, ) algorithm, which is a highly optimized hardware-aware implementation of $\mathrm{SoftmaxAttn}$.

Flash Attention in Brief: As pointed at in the main text, the FlashAttention (Dao et al., ) and FlashAttention2 (Dao, ) algorithms provide hardware-aware implementations of exact attention mechanism by optimizing for bottlenecks of modern accelerators (Choquette et al., ; Choquette, ).
These GPUs possess massive amounts (e.g., $\sim 80$GB) of High-Bandwidth Memory (HBM), which stores large tensors but is slow in moving the data to the accelerators.
On the other hand, they have smaller amounts (e.g., $\sim 20$ MB) of SRAM, which is often more than an 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., ).
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 is 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 FlashAttention2 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.

Fewer Tensor Allocations, Fewer Memory Accesses, Fast-Tanh: In FlashAttention and FlashAttention2, the attention mechanism is computed by splitting the attention matrix into blocks.
Since softmax activation requires a 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 it 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.e., moving the logsumexp tensor $\bm{L}\in\mathbb{R}^{n}$ on HBM in the forward and backward passes).
In addition, applying softmax in a stable manner requires tracking the row-max variable $m_{i}$ on chip, which instead is not needed for sigmoid activation.
This further helps reducing some on-chip operations and lowering register pressure in FlashSigmoid.

Moving on to the backward pass (described in Alg.˜2), FlashAttention2 requires computing $\textrm{rowsum}\left(\mathbf{d}\bm{O}\odot\bm{O}\right)$, which is needed to backpropagate the gradients of softmax attention outputs to the preactivations.
However, since sigmoid activation is applied element-wise, 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 FlashAttention and FlashAttention2 consists of partially re-computing the forward pass of attention mechanism in the backward pass to 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 in 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: $\sigma\left(x\right)=0.5\cdot\left(1+\tanh\left(0.5\cdot x\right)\right)$.
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-boost in both the forward and backward passes of FlashSigmoid, while maintaining parity in performance with a naïve implementation of sigmoid attention.
Based on these points of modification, we extend FlashAttention2 to obtain FlashSigmoid, a hardware-aware implementation of $\mathrm{SigmoidAttn}$.

### 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 in PyTorch (Paszke et al., ) that inputs query, key, and value tensors (all of shape $[\textrm{batch},\textrm{tokens},\textrm{heads},\textrm{features}]$) and passes these through a number of attention layers.
Mimicking the design of vision transformers (ViTB-16/224) (Dosovitskiy et al., ), we set the number of heads and per-head features as $12$ and $64$, respectively.
We set a batch size of $32$, and consider a $10$-layer architecture.
Then, for the number of tokens sampled from a wide range of $[64,78\mathrm{k}]$, we compute the forward and backward passes of this model.
For these computations, we measure the kernel GPU time using PyTorch’s profiler.
We carry out our experiments on both H100 (Choquette, ) and A100 (Choquette et al., ) GPUs.

Figure: (a) Inference mode kernels on H100.
Refer to caption: x7.png

Figure: (a) Inference mode kernels on A100.
Refer to caption: x9.png

Results: Figures˜10 and 11 show the GPU time comparisons of kernels in inference mode and training mode of FlashSigmoid and FlashAttention2 respectively.
We observe that we obtain a large average speed-boost for inference and a modest average speed-boost for training.
Note that the speed-ups in all the subsequent figures are obtained by averaging the performances for tokens sampled in the range of $[64,78\mathrm{k}]$.

Details of Individual Kernels: Next, we 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_fwd_kernel, and the backward pass of both these approaches is made up of three kernels, which we term bwd_dq_dk_dv, bwd_dot_do_o, and bwd_convert_dq.
In code, the real names of these kernels are as follows.

$$ fwd $\displaystyle=\textrm{flash\char 95\relax fwd\char 95\relax kernel}$ (83) bwd_dq_dk_dv $\displaystyle=\textrm{flash\char 95\relax bwd\char 95\relax dq\char 95\relax dk\char 95\relax dv\char 95\relax loop\char 95\relax seqk\char 95\relax parallel\char 95\relax kernel}$ bwd_dot_do_o $\displaystyle=\textrm{flash\char 95\relax bwd\char 95\relax dot\char 95\relax do\char 95\relax o\char 95\relax kernel}$ bwd_convert_dq $\displaystyle=\textrm{flash\char 95\relax bwd\char 95\relax convert\char 95\relax dq\char 95\relax kernel}$ $$

Here, we first provide a brief description of the tasks performed by each of these kernels; for a detailed explanation, we refer the reader to FlashAttention2 (Dao, ) paper and code.
The fwd kernel computes the full forward pass of the model as shown in Tab.˜4.
The bulk of computations of the backward pass happen in the bwd_dq_dk_dv kernel, which performs re-computation of attention matrix and reduction of key and value gradient tensors ($\mathbf{d}\bm{K}$, $\mathbf{d}\bm{V}$).
Again, the exact steps carried out in the backward pass can be checked from Tab.˜4.
The bwd_convert_dq kernel performs the reduction of query gradient tensor ($\mathbf{d}\bm{Q}$).
Finally, note that the bwd_dot_do_o kernel in FlashAttention2 performs the task of computing the $\textrm{rowsum}(\mathbf{d}\bm{O}\odot\bm{O})$ tensor along with clearing of the accumulators of query gradients ($\mathbf{d}\bm{Q}$).
Although FlashSigmoid does not require this row-sum tensor, the clearing of accumulators of query gradients is still needed.
For this reason, bwd_dot_do_o kernel also appears in the profiling of FlashSigmoid.

Performance of Individual Kernels:  Figures˜12 and 13 show the performance comparison of each flash kernel in FlashSigmoid with the corresponding kernel in FlashAttention2 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_dq_dk_dv kernel of FlashSigmoid has a modest average speed boost over FlashAttention2.
The bwd_dot_do_o kernel in FlashSigmoid is significantly faster on A100 GPUs.
Note that even though the bwd_dot_do_o kernel of FlashSigmoid appears to be slower on average on H100 GPUs, the kernel time of bwd_dot_do_o ($\sim 5$ms) is negligible compared to that of the main bwd_dq_dk_dv kernel ($\sim 5000$ms).
Thus, the combined backward pass kernel in FlashSigmoid time does not suffer from this slowdown.
Finally, note that for bwd_convert_dq, FlashSigmoid and FlashAttention2 have identical performance.
This is expected, since the task of this kernel is to reduce the gradient of the queries $\mathbf{d}\bm{Q}$, which is a common step in both the approaches and is not modified in FlashSigmoid.

Figure: fwd: $17.39\%$ faster for self-attention and $18.76\%$ for causal.
Refer to caption: x11.png

Figure: fwd: $14.33\%$ faster for self-attention and $16.92\%$ for causal.
Refer to caption: x15.png

### F.3 Speed Boosts of FlashSigmoid in Realistic Settings

In this section, we demonstrate how the performance boosts measured in Sec.˜F.2 for the individual kernels of FlashSigmoid contributes to speeding-up realistic runs with end-to-end training.

Setup: As a target experiment, we consider training a vision transformer (Dosovitskiy et al., ) on the ImageNet dataset (Deng et al., ).
We create two vision transformer model variants– one with FlashAttention2 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., ).
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\times 64$, $90\times 90$, and $100\times 100$), along with patch size of $1$ to result in different number of tokens for the underlying attention mechanism in the vision transformer model ($64\times 64=4096$, $90\times 90=8100$, and $100\times 100=10000$ tokens).
For each of these configurations, we select batch sizes so that the GPU memory utilization would be greater than $80\%$.
These considerations 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-downs due to across-nodes communications.

Results: The results of the runs on H100 nodes and A100 nodes are shown in Tab.˜5 and 6 respectively.
There, we show how the kernel GPU times for forward and backward passes vary according to 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 $\sim 8\%$ speed-up in wall-clock time of inference and $\sim 4\%$ speed-up in wall-clock time of training.

**Table 5: FlashSigmoid vs. FlashAttention2 on H100 nodes. The kernel GPU time for both the approaches is reported in milliseconds and wall-clock times is reported in seconds per epoch.**
| Tokens | Kernel GPU Time Comparison | Full Run Wall-Clock Time Comparison |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- |
| Kernels | FlashAttention2 (ms) | FlashSigmoid (ms) | Mode | FlashAttention2 (s) | FlashSigmoid (s) |  |
| 4096 | fwd | ${4.98}{\pm{0.01}}$ | ${4.17}{\pm{0.01}}\ \mathbf{\left(-16.31\%\right)}$ | Inference | ${11.17}{\pm{0.18}}$ | ${10.68}{\pm{0.18}}\ \mathbf{\left(-4.42\%\right)}$ |
| $\textrm{fwd}+\textrm{bwd}$ | ${19.58}{\pm{0.06}}$ | ${18.12}{\pm{0.04}}\ \mathbf{\left(-7.45\%\right)}$ | Training | ${1563.39}{\pm{1.30}}$ | ${1521.68}{\pm{2.27}}\ \mathbf{\left(-2.67\%\right)}$ |  |
| 8100 | fwd | ${20.46}{\pm{0.05}}$ | ${16.73}{\pm{0.05}}\ \mathbf{\left(-18.22\%\right)}$ | Inference | ${28.21}{\pm{0.18}}$ | ${25.93}{\pm{0.17}}\ \mathbf{\left(-8.06\%\right)}$ |
| $\textrm{fwd}+\textrm{bwd}$ | ${77.63}{\pm{0.13}}$ | ${72.70}{\pm{0.12}}\ \mathbf{\left(-6.35\%\right)}$ | Training | ${4282.75}{\pm{2.14}}$ | ${4129.25}{\pm{4.14}}\ \mathbf{\left(-3.58\%\right)}$ |  |
| 10000 | fwd | ${31.17}{\pm{0.07}}$ | ${25.49}{\pm{0.05}}\ \mathbf{\left(-18.20\%\right)}$ | Inference | ${38.71}{\pm{0.19}}$ | ${35.37}{\pm{0.17}}\ \mathbf{\left(-8.62\%\right)}$ |
| $\textrm{fwd}+\textrm{bwd}$ | ${117.53}{\pm{0.13}}$ | ${109.87}{\pm{0.12}}\ \mathbf{\left(-6.52\%\right)}$ | Training | ${5990.72}{\pm{2.21}}$ | ${5751.43}{\pm{5.77}}\ \mathbf{\left(-3.99\%\right)}$ |  |

**Table 6: FlashSigmoid vs. FlashAttention2 on A100 nodes. The kernel GPU time for both the approaches is reported in milliseconds and wall-clock times is reported in seconds per epoch.**
| Tokens | Kernel GPU Time Comparison | Full Run Wall-Clock Time Comparison |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- |
| Kernels | FlashAttention2 (ms) | FlashSigmoid (ms) | Mode | FlashAttention2 (s) | FlashSigmoid (s) |  |
| 4096 | fwd | ${8.32}{\pm{0.02}}$ | ${7.84}{\pm{0.03}}\ \mathbf{\left(-5.79\%\right)}$ | Inference | ${19.05}{\pm{0.22}}$ | ${18.74}{\pm{0.19}}\ \mathbf{\left(-1.65\%\right)}$ |
| $\textrm{fwd}+\textrm{bwd}$ | ${31.81}{\pm{0.08}}$ | ${31.11}{\pm{0.08}}\ \mathbf{\left(-2.19\%\right)}$ | Training | ${2795.03}{\pm{2.35}}$ | ${2769.44}{\pm{5.10}}\ \mathbf{\left(-0.92\%\right)}$ |  |
| 8100 | fwd | ${33.65}{\pm{0.09}}$ | ${27.92}{\pm{0.07}}\ \mathbf{\left(-17.04\%\right)}$ | Inference | ${47.35}{\pm{0.20}}$ | ${44.05}{\pm{0.17}}\ \mathbf{\left(-6.96\%\right)}$ |
| $\textrm{fwd}+\textrm{bwd}$ | ${128.18}{\pm{0.13}}$ | ${119.04}{\pm{0.12}}\ \mathbf{\left(-7.13\%\right)}$ | Training | ${7519.64}{\pm{4.21}}$ | ${7254.84}{\pm{12.64}}\ \mathbf{\left(-3.52\%\right)}$ |  |
| 10000 | fwd | ${51.17}{\pm{0.07}}$ | ${42.49}{\pm{0.06}}\ \mathbf{\left(-16.96\%\right)}$ | Inference | ${64.61}{\pm{0.32}}$ | ${59.55}{\pm{0.18}}\ \mathbf{\left(-7.82\%\right)}$ |
| $\textrm{fwd}+\textrm{bwd}$ | ${194.54}{\pm{0.14}}$ | ${180.59}{\pm{0.15}}\ \mathbf{\left(-7.17\%\right)}$ | Training | ${10455.64}{\pm{8.85}}$ | ${10052.04}{\pm{18.87}}\ \mathbf{\left(-3.86\%\right)}$ |  |

Connection of Wall-Clock Time Speed-Up and Kernel Speed-Up: From Tab.˜6 and 5, it 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 $\tau_{\textrm{sm}}$ and $\tau_{\sigma}$ the total kernel GPU time for softmax attention and sigmoid attention respectively.
Then, the kernel speed-up is given by $s_{\textrm{kernel}}:=1-\frac{\tau_{\sigma}}{\tau_{\textrm{sm}}}$.
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 $\tau$.
Thus, the wall-clock time speed-up of a full run with end-to-end training is $s_{\textrm{wall-clock}}:=1-\frac{\tau_{\sigma}+\tau}{\tau_{\textrm{sm}}+\tau}$.
Since we have faster sigmoid kernels, we have $\tau_{\sigma}<\tau_{\textrm{sm}}$, which in turn shows that $s_{\textrm{wall-clock}}=1-\frac{\tau_{\sigma}+\tau}{\tau_{\textrm{sm}}+\tau}<1-\frac{\tau_{\sigma}}{\tau_{\textrm{sm}}}=s_{\textrm{kernel}}$.
This explains the speed boost trends in kernel time versus full run wall-clock time for each setting in Tab.˜6 and 5.
However, in particular, if a model performs attention mechanism over large number of tokens, the attention mechanism, and hence the corresponding kernel time, starts to dominate the other computations in the network.
In that case, we see that the wall-clock time speed-boost is closer to the kernel speed-boost.
Mathematically, if $\tau_{\sigma},\tau_{\textrm{sm}}>\!\!>\tau$, we have: $\tau_{\sigma}+\tau\approx\tau_{\sigma}$, $\tau_{\textrm{sm}}+\tau\approx\tau_{\textrm{sm}}$.
Thus, $s_{\textrm{kernel}}\approx s_{\textrm{wall-clock}}$, thereby making ${s_{\textrm{wall-clock}}}/{s_{\textrm{kernel}}}\rightarrow 1$.

Figure: (a) Inference mode kernels on H100.
Refer to caption: x19.png

Figure: (a) Inference mode kernels on A100.
Refer to caption: x21.png

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 kernels and $5-10\%$ during inference of full runs).
We posit that this speed-up in inference is extremely critical as well.
Contemporary large-scale models, once trained, spend a huge portion of the rest their lifetime in inference mode (OpenAI, ).
Thus, significant performance boosts in inference mode have immense potential for saving resources in deployment of large models for inference.

### F.4 FlashSigmoid with ALiBi

It is evident from the main text of the paper that improved positional embeddings, like ALiBi (Press et al., ), can be crucial for certain tasks and data modalities.
Thus, we also provide a FlashSigmoid implementation that incorporates ALiBi.
We compare the FlashSigmoid with ALiBi implementation with the FlashAttention2 with ALiBi implementation (Dao, ).
Figures˜15 and 14 show the kernel GPU time for the forward and backward pass kernels of FlashSigmoid with ALiBi implementation versus FlashAttention2 with ALiBi implementation.
Again, we observe that FlashSigmoid kernels for inference have significant speed-up in wall-clock time over those in FlashAttention2 and the kernels for training also have modest wall-clock improvements.

### F.5 Directions for Future Work on FlashSigmoid

In this section, we discussed FlashSigmoid, a hardware-aware implementation of the $\mathrm{SigmoidAttn}$ algorithm. Then, we demonstrated via kernel benchmarking and realistic setting runs that FlashSigmoid provides significant gains in inference as well as modest gains in training of models with attention mechanism.
In 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 FlashAttention2 by adding functionality for forward and backward pass of sigmoid attention in place of the standard softmax attention.
In particular, for all FlashSigmoid results discussed so far, we inherit directly from FlashAttention2 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_{r},B_{c}$ in Alg.˜1 and 2, which are identical in FlashAttention2 and FlashSigmoid.
This choice is dictated by the need to ensure a fair comparison between the two implementations, and allows us to demonstrate the speed-up of sigmoid attention by minimizing confounders associated with parallel computations on different GPU architectures for different input shapes.

Although FlashSigmoid kernels lead to speed-ups in inference and training for both H100 and A100 GPUs, we observe that the kernel timing speed-ups on 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 of attention mechanisms should 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 FlashSigmoid to values that are different than those in FlashAttention2.

Note that the implementation of FlashAttention2 is templated according to block shapes, grid shapes, and other kernel launch parameters.

Note that FlashAttention2 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. no 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 $\textsc{FlashSigmoid}^{\dagger}$, where we update the block sizes for query and key tensors from $\left(B_{r},B_{c}\right)=\left(128,128\right)$ of FlashSigmoid to $\left(B_{r},B_{c}\right)=\left(128,64\right)$ of $\textsc{FlashSigmoid}^{\dagger}$ *only* for our input setting (template with features per head being $64$).

Figure: (a) Inference mode kernels on A100.
Refer to caption: x23.png

**Table 7: FlashAttention2 vs. FlashSigmoid vs. $\textsc{FlashSigmoid}^{\dagger}$ on A100 nodes. The kernel GPU time for all three approaches are reported in milliseconds. We observe that $\textsc{FlashSigmoid}^{\dagger}$ provides better and more uniform speed-ups across all example tokens.**
| Tokens | Kernel GPU Time Comparison |  |  |  |
| --- | --- | --- | --- | --- |
| Kernels | FlashAttention2 (ms) | FlashSigmoid (ms) | $\textsc{FlashSigmoid}^{\dagger}$ (ms) |  |
| 4096 | fwd | ${8.32}{\pm{0.02}}$ | ${7.84}{\pm{0.03}}\ \mathbf{\left(-5.79\%\right)}$ | ${7.26}{\pm{0.02}}\ \mathbf{\left(-13.21\%\right)}$ |
| $\textrm{fwd}+\textrm{bwd}$ | ${31.81}{\pm{0.08}}$ | ${31.11}{\pm{0.08}}\ \mathbf{\left(-2.19\%\right)}$ | ${30.62}{\pm{0.09}}\ \mathbf{\left(-4.03\%\right)}$ |  |
| 8100 | fwd | ${33.65}{\pm{0.09}}$ | ${27.92}{\pm{0.07}}\ \mathbf{\left(-17.04\%\right)}$ | ${28.54}{\pm{0.07}}\ \mathbf{\left(-15.50\%\right)}$ |
| $\textrm{fwd}+\textrm{bwd}$ | ${128.18}{\pm{0.13}}$ | ${119.04}{\pm{0.12}}\ \mathbf{\left(-7.13\%\right)}$ | ${119.85}{\pm{0.13}}\ \mathbf{\left(-6.81\%\right)}$ |  |
| 10000 | fwd | ${51.17}{\pm{0.07}}$ | ${42.49}{\pm{0.06}}\ \mathbf{\left(-16.96\%\right)}$ | ${43.53}{\pm{0.09}}\ \mathbf{\left(-15.32\%\right)}$ |
| $\textrm{fwd}+\textrm{bwd}$ | ${194.54}{\pm{0.14}}$ | ${180.59}{\pm{0.15}}\ \mathbf{\left(-7.17\%\right)}$ | ${181.97}{\pm{0.17}}\ \mathbf{\left(-6.87\%\right)}$ |  |
| 16384 | fwd | ${134.19}{\pm{0.12}}$ | ${125.43}{\pm{0.10}}\ \mathbf{\left(-6.53\%\right)}$ | ${116.75}{\pm{0.10}}\ \mathbf{\left(-13.40\%\right)}$ |
| $\textrm{fwd}+\textrm{bwd}$ | ${494.65}{\pm{0.28}}$ | ${482.08}{\pm{0.23}}\ \mathbf{\left(-2.54\%\right)}$ | ${474.52}{\pm{0.28}}\ \mathbf{\left(-4.48\%\right)}$ |  |

Experimentation and Results: For this variant, we perform kernel benchmarking as described in Sec.˜F.2, and report the corresponding results in Fig.˜16.
Comparing the plots for kernel timing with FlashSigmoid plots from Fig.˜11, we observe that $\textsc{FlashSigmoid}^{\dagger}$ not only provides a 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.˜7 shows the inference mode and training mode kernel speed-ups for a subset of sequence lengths under consideration.
This experiment indicates that it is possible to 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.

## Appendix G Experiments

### G.1 Extra Ablations

#### G.1.1 The Effect of Multiplicative Sequence Length Normalization

Figure: Figure 17: $b=-\ln n$.
Refer to caption: figures/attn_norm_seed1000001_softmax_vs_sigmoid.png

Wortsman et al. () notes that models trained with sigmoid or ReLU attention require scaling by the sequence length, $n^{-\alpha}\sigma(\bm{Q}\bm{K}^{T}/\sqrt{d_{qk}})\bm{V}$. We ablate this by comparing the scaled solution to the one we propose in App.˜E. We also generalize the variant proposed in (Wortsman et al., ) to variadic sequence lengths such that it works with auto-regressive (AR) training, for example for $n=3$:

$$ $\displaystyle\underbrace{\begin{bmatrix}1&1&1\\ 0.5^{-\alpha}&0.5^{-\alpha}&1\\ 0.33^{-\alpha}&0.33^{-\alpha}&0.33^{-\alpha}\\ \end{bmatrix}}_{n^{-\alpha}}\odot\underbrace{\begin{bmatrix}1&0&0\\ 1&1&0\\ 1&1&1\\ \end{bmatrix}}_{\text{Causal Mask }\bm{M}}\odot\ \sigma(\bm{Q}\bm{K}^{T}/\sqrt{d_{qk}})\bm{V}.$ (84) $$

We repeat the experiment from Fig.˜6, using ALiBi positional embeddings for all trials. We apply $\alpha=\{1,0.5\}$ AR normalization proposed in Equation˜84. While there is an observable difference in terms of the attention norm, $\lVert\sigma(\bm{Q}\bm{K}^{T}/\sqrt{d_{qk}})\bm{V}\rVert$, we find that the train NLL is slightly worse for both normalized variants (Fig.˜19 and 19) in comparison to the $b=-\ln n$ variant in Fig.˜19.

#### G.1.2 Attention Bias Stability Ablation

To validate the stabilizing effects of attention bias we repeat the experiment from Fig.˜8 and 8, keeping all of the same hyper-parameters, while enabling QK norm and LayerScale (initialized at $10^{-4}$). We train with a range of constant bias offsets, $b\in\{-15,-10,-6,-4,-1\}$ and visualize the results below in Fig.˜20.

Figure: Figure 20: Attention bias ablation.
Refer to caption: x25.png

We observe a systematic increase in stability (and lower $\mathrm{SigmoidAttn}$ NLL) for 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 ImageNet1k Top-1%

Figure: Figure 21: ImageNet1k test top-1% for $\mathrm{SoftmaxAttn}$ vs. $\mathrm{SigmoidAttn}$ using models from Fig.˜2.
Refer to caption: x26.png

Fig.˜21 reports the test linear probe results for the ViT-B/16 BYOL (Grill et al., ; Busbridge et al., ), ViT-B/16 SimCLR (Chen et al., ; Zhai et al., ) and the finetuned performance for the ViT-L/16 MAE (He et al., ) and the test top-1% results for for ViT-B/16 supervised model (Dosovitskiy et al., ).
Across these wide range of SSL and supervised learning tasks, trained with contrastive (SimCLR), EMA distillation (BYOL) and reconstructive objectives (MAE), we find that $\mathrm{SigmoidAttn}$ not only matches the training dynamics (Fig.˜2), but also the linear probe and finetuned performance of the baseline $\mathrm{SoftmaxAttn}$.

#### G.2.2 LayerScale Free Sigmoid Attention

Figure: Figure 22: A competitive $\mathrm{SigmoidAttn}$ ViT-B/16 model can be learned without LayerScale or QK norm using a large initial *learnable* scalar temperature $t=10$ and bias $b=-10$ (similar to SigLIP (Zhai et al., )): $\sigma(e^{t}[\bm{Q}\bm{K}^{T}/\sqrt{d_{qk}}]+b)\bm{V},\{b,t\}\in\mathbb{R}$. This regularizes the model, as it must move the temperature to a learnable regime. The $t=10,b=-10$ curve makes no progress in train NLL or test top-1 for $\sim$25 epochs (near max LR), but ultimately outperforms baselines.
Refer to caption: x27.png

While Fig.˜22 demonstrates the possibility of learning $\mathrm{SigmoidAttn}$ without LayerScale, it involves task specific tuning of $\{t,b\}$. We also explored gating attention from learning (through a simple multiply by zero) for $\sim$25 epochs and were able to match the $t=10,b=-10$ training curves from above. However, we opted for the LayerScale method due to its simplicity.

#### G.2.3 Sigmoid Attention vs. Attention Relaxations

Figure: Figure 23: Supervised ViT-B/16 ImageNet1k classification. We contrast $\mathrm{SigmoidAttn}$ and $\mathrm{SoftmaxAttn}$ against (a) linear attention with no activation: $\bm{Q}\bm{K}^{T}/\sqrt{d_{qk}}$ and (b) fast attention via positive orthogonal random features, used in Performer (Choromanski et al., ). $\mathrm{SigmoidAttn}$, like $\mathrm{SoftmaxAttn}$, differs from attention relaxations like Performer which uses low-rank representations of the attention matrix. $\mathrm{SigmoidAttn}$ maintains performance parity with $\mathrm{SoftmaxAttn}$, while outperforming other efficient attention variants.
Refer to caption: x28.png

#### G.2.4 Hyper-Parameters

**Table 8: $\mathrm{SigmoidAttn}$ SimCLR and BYOL ViT-B/16 hyperparameters.**
| Parameter | SimCLR | BYOL |
| --- | --- | --- |
| Attention bias | None | None |
| LayerScale Init | $10^{-4}$ | $10^{-4}$ |
| QK Norm | Yes | Yes |
| Pos Embed | SinCos | Learnable |
| Freeze Patcher | Yes | No |
| Weight init | MocoV3 (Chen et al., ) | trunc_normal(.02) |
| Normalization | LayerNorm | LayerNorm |
| LR schedule | Single Cycle Cosine | Single Cycle Cosine |
| LR warmup | 10 Epochs | 40 Epochs |
| Min LR | $1\times 10^{-6}$ | $1\times 10^{-6}$ |
| Training duration | 300 Epochs | 600 Epochs |
| Optimizer | AdamW | AdamW |
| Optimizer scaling rule | Linear | Linear |
| Base Adam ($\beta_{1},\beta_{2}$) | (0.9, 0.95) | (0.9, 0.95) |
| Base LR | $2\times 10^{-4}$ | $1\times 10^{-4}$ |
| Base batch size | 256 | 256 |
| Total batch size | 4096 | 4096 |
| Base teacher momentum | - | 0.996 |
| Weight decay | 0.1 | 0.3 |
| Weight decay skip bias | Yes | Yes |
| Numerical precision | bf16 | bf16 |
| Stochastic depth | 0.0 | 0.2 |
| Augmentation stack | SimCLR (Chen et al., ) | DINO multicrop (Caron et al., ) |
| Color Jitter Scaling | 0.5 (Chen et al., ) | 1.0 |

**Table 9: $\mathrm{SigmoidAttn}$ Supervised ViT-B/16 and MAE ViT-L/16 hyperparameters.**
| Parameter | Supervised | MAE |
| --- | --- | --- |
| Attention bias | None | $b=-\ln{n}$ |
| LayerScale Init | $10^{-4}$ | $10^{-4}$ |
| QK Norm | Yes | Yes |
| Pos Embed | Learnable | Learnable |
| Architecture | ViT-B/16 | ViT-L/16 |
| Mask Ratio | - | 0.75 |
| Freeze Patcher | No | No |
| Weight init | trunc_normal(.02) | trunc_normal(.02) |
| Normalization | LayerNorm | LayerNorm |
| LR schedule | Single Cycle Cosine | Single Cycle Cosine |
| LR warmup | 20 Epochs | 40 Epochs |
| Min LR | $1\times 10^{-6}$ | $0.0$ |
| Training duration | 300 Epochs | 400 Epochs |
| Optimizer | AdamW | AdamW |
| Optimizer scaling rule | Linear | Linear |
| Base Adam ($\beta_{1},\beta_{2}$) | (0.9, 0.95) | (0.9, 0.95) |
| Base LR | $1\times 10^{-4}$ | $1.5\times 10^{-4}$ |
| Base batch size | 256 | 256 |
| Total batch size | 4096 | 4096 |
| Weight decay | 0.3 | 0.05 |
| Weight decay skip bias | Yes | Yes |
| Numerical precision | bf16 | bf16 |
| Stochastic depth | 0.28 | 0.0 |
| Augmentation stack | RandAug (Cubuk et al., ) | RRC + HFLIP |

### G.3 Language Model

#### G.3.1 Hyper-Parameters

**Table 10: Training details for the Llama-style 1B LM training.**
| Parameter | Value |
| --- | --- |
| Params | 1B |
| Context Length | 2048 |
| Total Tokens | 300B |
| Batch size | 4M tokens |
| LR Schedule | Cosine |
| LR Warmup Steps | 5000 |
| Peak LR | 1e-2 |
| Final LR | 10% of peak |
| Optimizer | AdamW |
| Optimizer momentum | 0.9, 0.95 |
| Weight decay | 1e-4 |
| Gradient clipping | 1.0 |
| Position encoding | ALiBi |
| Q/K Norm | Applied |
| Num layers | 24 |
| Num heads | 32 |
| Hidden dim | 2048 |

Tab.˜10 shows the hyper-parameters for the final comparison. MuP-simple (Wortsman et al., ) is used, where the peak learning rate is set to 1e-2. Weight decay is decoupled, following Loshchilov & Hutter (). In addition, to confirm that applying QK-Norm does not hurt the baseline, we show training parity with and without QK-Norm in Fig.˜25.

Figure: Figure 24: 1B $\mathrm{SoftmaxAttn}$ LLM training with and without QK Norm, converging to the same loss.
Refer to caption: x29.png

Figure: Figure 26: 85M training using $\mathrm{SigmoidAttn}$ and $\mathrm{SoftmaxAttn}$ (n = 4096). Training loss matches.
Refer to caption: x31.png

#### G.3.2 Gradient Norm

While a $\mathrm{SigmoidAttn}$ based LM using aforementation hyper-parameters has a smooth loss curve, we do see more gradient norm fluctuations. See Fig.˜25, where spikes larger than $0.5$ are not visible in the $\mathrm{SoftmaxAttn}$ equivalent.

### 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-based encoder model has 255M parameters: 1D convolution of kernel 7 and stride 3 followed by CAPE positional embedding if it is used and 36 transformer blocks with pre-LayerNorm, an embedding dimension of 768, 4 heads, 3072 units in the MLP layers.
The model is trained with CTC loss and a 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 a CTC loss and a character vocabulary.
The conformer model has 104M parameters and consists of 1D convolution of kernel 7 and stride 3 followed by 16 conformer blocks with an 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., ) is used for augmentation with 2 frequency masks (max width 30) and 10 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 $\mathrm{SoftmaxAttn}$ and then change only attention to the introduced configuration of $\mathrm{SigmoidAttn}$ 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.

For the bias term $b=-\log n$ in $\mathrm{SigmoidAttn}$, we do not use max sequence length as in language model experiments. Instead, for every audio sample we use its own duration as a bias terms resulting into non-trainable bias vector for the minibatch. For experiments with sequence normalization, we also use not the max sequence length in the minibatch but rather the ground truth sample duration to properly normalize encoder attention.

**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 | pre | pre + post | pre | post |
| Dropout | 0.1 | 0.1 | 0.3 | 0.3 |
| Layer drop | 0.1 | 0.0 | 0.3 | 0.3 |
| 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 |
| SpecAugment start | 0k | 10k | 0k | 0k |
| LR Warmup Steps | 64k | 10k | 64k | 64k |
| Peak LR | 1e-3 | 2e-3 | 0.1 | 0.03 |
| LR start decay | 250k | 250k | 200k | 330k |
| LR decay step | 50k | 50k | 30k | 50k |
| Optimizer | AdamW | AdamW | Adagrad | Adagrad |
| Optimizer momentum | 0.9, 0.999 | 0.9, 0.98 | - | - |
| Weight decay | 1e-6 | 1e-6 | 0 | 0 |
| Gradient clipping | 1.0 | 0.5 | 1.0 | 1.0 |
| Position encoding | CAPE / ALiBi / RoPE | RoPE | CAPE | CAPE / ALiBi / RoPE |
| Q/K Norm $\mathrm{SoftmaxAttn}$ | Not Applied | Not Applied | Not Applied | Not Applied |
| Q/K Norm $\mathrm{SigmoidAttn}$ | Applied | Applied | Not Applied | Applied |
| Num layers | 36 | 16 | 36 / 24 / 12 | 36 |
| Num heads | 4 | 4 | 4 | 4 |

To evaluate behaviour for length generalization we use TED-LIUM v3 dataset Hernandez et al. () as its validation and test sets have longer audio duration than LibriSpeech: LibriSpeech has in average 10-15s duration, while in 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 v3 (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-10s, 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 no official adoption of it yet(^13^1313See discussion in [https://github.com/ofirpress/attention_with_linear_biases/issues/5](https://github.com/ofirpress/attention_with_linear_biases/issues/5).) for the encoder models (without causal masking), we follow the best practices found in [https://iclr-blogposts.github.io/2024/blog/alibi-mlm/](https://iclr-blogposts.github.io/2024/blog/alibi-mlm/) of nonsymmetric ALiBi with different slopes instead of symmetric version used by (Lee et al., ).

**Table 12: Word error rate (%) on LibriSpeech dev/test sets and TED-LIUM v3 (Hernandez et al., ) (“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 $\mathrm{SoftmaxAttn}$ or $\mathrm{SigmoidAttn}$ (w/ LayerScale, w/o QK norm, w/ $b=-\log n$) trained on LibriSpeech 100h data (average duration is 10-15s). Hyper-parameters can be found in Table˜11.**
| attn | # layers | dev-clean | test-clean | dev-other | test-other | ted 0-10s | ted 10-20s | ted 20-30s | ted 30s+ |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| softmax | 36 | 6.7 | 7.1 | 20.0 | 20.4 | 26.4 | 22.4 | 23.3 | 21.8 |
| sigmoid | 36 | 7.0 | 7.3 | 20.3 | 20.5 | 26.2 | 23.4 | 23.6 | 21.8 |
| $b=0$ | 36 | 6.8 | 7.1 | 19.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 |
| $b=0$ | 24 | 6.7 | 6.9 | 20.2 | 20.7 |  |  |  |  |
| softmax | 12 | 8.2 | 8.7 | 25.0 | 25.4 | 29.0 | 25.6 | 27.1 | 27.4 |
| sigmoid | 12 | 8.3 | 8.7 | 24.8 | 25.2 | 29.0 | 25.7 | 26.3 | 25.5 |
| $b=0$ | 12 | 8.7 | 8.5 | 24.4 | 24.7 |  |  |  |  |

**Table 13: Word error rate (%) on LibriSpeech dev/test sets for post-LayerNorm transformer (255M) with either $\mathrm{SoftmaxAttn}$ (w/o QK norm) or $\mathrm{SigmoidAttn}$ (by default w/ LayerScale, w/ QK norm, w/ $b=-\log n$) trained on LibriSpeech 100h data. Hyper-parameters can be found in Table˜11.**
| attn | PE | dev-clean | test-clean | dev-other | test-other |
| --- | --- | --- | --- | --- | --- |
| softmax | CAPE | 6.4 | 6.5 | 18.4 | 18.2 |
| + QK norm |  | 6.1 | 6.3 | 18.2 | 18.1 |
| sigmoid |  | 8.0 | 8.4 | 22.7 | 22.7 |
| - QK norm |  | 7.5 | 7.9 | 22.1 | 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 |
| softmax | RoPE | 6.6 | 6.9 | 18.3 | 18.5 |
| sigmoid |  | 6.8 | 7.1 | 20.8 | 20.8 |
| sigmoid ($b=-10$, learnable) |  | 8.7 | 9.4 | 23.5 | 24.0 |
| softmax | AliBi | 6.4 | 6.9 | 18.3 | 18.3 |
| sigmoid |  | 6.9 | 7.2 | 20.8 | 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., ) (“TED”, joint validation and test sets with split according to audio duration) for conformer (104M) with RoPE and with either $\mathrm{SoftmaxAttn}$ or $\mathrm{SigmoidAttn}$ (w/ LayerScale, w/ QK norm, w/ $b=-\log n$) trained on LibriSpeech 960h data (average duration is 10-15s). Hyper-parameters can be found in Table˜11.**
| attn | dev-clean | test-clean | dev-other | test-other | ted 0-10s | ted 10-20s | ted 20-30s | ted 30s+ |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| softmax | 2.2 | 2.5 | 5.4 | 5.6 | 13.0 | 11.1 | 13.2 | 7.1 |
| sigmoid | 2.3 | 2.5 | 5.6 | 5.8 | 13.5 | 10.8 | 13.3 | 10.2 |
| sigmoid ($b=-10$, learnable) | 2.4 | 2.7 | 5.8 | 5.8 | 12.9 | 11.1 | 14.1 | 54.9 |

Figure: Figure 28: ASR Transformer model (255M) training with post-LayerNorm (left) and pre-LayerNorm (right) on LibriSpeech 960h with $\mathrm{SigmoidAttn}$ (w/ bias term, $b=0$, w/o QK norm, w/ LayerScale) or with $\mathrm{SoftmaxAttn}$. Huge gradient norms and training loss spikes are observed for $\mathrm{SigmoidAttn}$ which can result in worse final model performance hence models for $\mathrm{SigmoidAttn}$ are unstable.
Refer to caption: x33.png

Figure: Figure 29: ASR Transformer model (255M) training with pre-LayerNorm on LibriSpeech 960h with $\mathrm{SigmoidAttn}$ (w/ bias term, $b=-\log n$, w/ QK norm, w/ LayerScale) and different positional embeddings CAPE, RoPE, ALiBi. The bias $b$ is able to stabilize $\mathrm{SigmoidAttn}$ training: smooth training loss and only marginal rare spikes in gradient norms are observed.
Refer to caption: x34.png

#### G.4.2 Results and Ablations

Initial investigation on post-LayerNorm and pre-LayerNorm transformers on both LibriSpeech 100h and 960h revealed that $\mathrm{SigmoidAttn}$ 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˜28. 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 $\mathrm{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 $\mathrm{SoftmaxAttn}$. Moreover, some combination of them is needed for the stable training, though w/o both of them we got the best performance for $\mathrm{SigmoidAttn}$ (still behind $\mathrm{SoftmaxAttn}$), see Table˜13.
We believe, further adaptation and deeper investigation is needed for $\mathrm{SigmoidAttn}$ and post-LayerNorm, though recent advances in machine learning do not use post-LayerNorm models due to high training instability even for $\mathrm{SoftmaxAttn}$.

Switching to pre-LayerNorm transformers and varying the depth of the models lead to stable training with $\mathrm{SigmoidAttn}$ and bias term $b=-\log n$ with few (2-5 times) spikes in the gradient norm and smooth loss. In this case, $\mathrm{SigmoidAttn}$ matches results for $\mathrm{SoftmaxAttn}$ and they both generalize to TED-LIUM data similarly, see Table˜12. If the bias term is removed, $\mathrm{SigmoidAttn}$ can still match $\mathrm{SoftmaxAttn}$ but large spikes in gradient norm and loss can occur.

Finally, we experiment with a conformer model, in Table˜14. Again, we found that bias term $b=-\log n$ stabilizes training. The learnable $b=-10$ works though we see significant gradient norm spikes while the train loss remains smooth. Besides, $b=-\log n$ generalizes well to longer sequences while learnable $b=-10$ fails to do so with RoPE for conformer. Overall, $\mathrm{SigmoidAttn}$ is able to match $\mathrm{SoftmaxAttn}$ having stable training with $b=-\log n$.

In experiments with different variants of bias term for $\mathrm{SigmoidAttn}$, the bias $b=-\log n$ is found to be the most stable (only few marginal gradient norm spikes are observed with the train loss being smooth) and it provides similar performance as $\mathrm{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 $\mathrm{SoftmaxAttn}$). This happens due to high attention weight of every token which can be biased towards zero with a bias term in $\mathrm{SigmoidAttn}$ definition.
Preliminary results to connect this to the local attention property needed at the beginning of the training for stable training failed, as local attention did not converge well at all (it is deactivated after some initial training).

To fully benefit from the improved throughput of FlashSigmoid, for the bias term $b=-\log n$ in $\mathrm{SigmoidAttn}$, we experimented with configuration when the maximum audio duration in the minibatch is used as $n$ 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 as 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 $b=-\log n$ as the final configuration.

**Table 15: Word error rate (%) on LibriSpeech dev/test sets and TED-LIUM v3 (Hernandez et al., ) (“TED”, joint validation and test sets split according to duration) for transformer (255M params) with either $\mathrm{SoftmaxAttn}$ or $\mathrm{SigmoidAttn}$ (LayerScale and QK norm are used with $b=-\log n$) trained on LibriSpeech 960h data (mean duration is 10-15s). Hyper-parameters are in Sec.˜G.4.**
| attn | PE | dev-clean | test-clean | dev-other | test-other | ted 0-10s | ted 10-20s | ted 20-30s | ted 30s+ |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| softmax | CAPE | 2.2 | 2.3 | 5.6 | 5.7 | 12.4 | 10.5 | 11.9 | 9.1 |
| sigmoid | 2.2 | 2.4 | 5.2 | 5.5 | 12.4 | 10.3 | 12.3 | 9.7 |  |
| sigmoid, $b=-\log(\max_{batch}n)$ | 2.1 | 2.3 | 5.2 | 5.3 | 12.2 | 10.6 | 12.0 | 9.3 |  |
| softmax | RoPE | 2.2 | 2.2 | 5.4 | 5.5 | 12.7 | 10.6 | 12.8 | 9.5 |
| sigmoid | 2.0 | 2.3 | 5.2 | 5.4 | 12.3 | 10.1 | 12.3 | 8.6 |  |
| sigmoid, $b=-\log(\max_{batch}n)$ | 2.1 | 2.3 | 5.0 | 5.1 | 12.3 | 10.1 | 12.1 | 10.4 |  |
| softmax | ALiBi | 2.1 | 2.2 | 5.3 | 5.4 | 12.3 | 10.7 | 12.1 | 8.6 |
| sigmoid | 2.1 | 2.3 | 5.0 | 5.1 | 12.3 | 10.5 | 12.6 | 9.1 |  |
| sigmoid, $b=-\log(\max_{batch}n)$ | 2.0 | 2.3 | 5.2 | 5.2 | 12.3 | 10.5 | 11.9 | 10.2 |  |

### G.5 Simple Experiments

#### G.5.1 k–Summation Problem Definition

Here we look at a 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 a $\mathbb{R}^{n}\rightarrow\mathbb{R}$ target function. In the first half of each input are samples from a $\mathcal{N}(0,1)$ distribution, and the second half is a $k$-hot binary vector indicating which values in the first half to sum.

The results presented here are for the $n=40$ problem with various values for $k$. Where a transformer is used, the transformer is a single layer to aid visualization. In all cases (unless noted otherwise), the optimizer is Adam with a constant learning rate of 0.001, and the training data is continuously generated to preclude over-fitting.

A few examples for $n=10$ (not drawn from $\mathcal{N}(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 0 0 1 $\rightarrow$ 5 |
| --- |
| 1 2 3 4 5 1 0 0 0 1 $\rightarrow$ 6 |
| 8 1 2 0 5 0 1 1 1 0 $\rightarrow$ 3 |
| 2 0 2 2 2 1 1 0 1 0 $\rightarrow$ 4 |

#### G.5.2 Comparison to Softmax

In Figure˜30, we see the performance of three architectures on the $k$-summation problem as $k$ increases. The sigmoid activated transformer has similar scaling to the softmax activation.

Figure: Figure 30: Final loss is shown after training convergence as k-summation problem complexity increases. The ReLU MLP has two 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 $\mathrm{SigmoidAttn}$ is applied after a learned offset initialized to -4, A+param(-4).
Refer to caption: figures/final_loss_vs_k_summation.png

#### G.5.3 Attention Evolution

In Figures˜31 and 32, forty samples are used to monitor the single head, single layer post-activation attention matrix as training progresses. In Figure˜31, 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 is that the sigmoid has fewer high magnitude values than softmax indicating a more distributed attention.

Figure: Figure 31: Softmax
Refer to caption: figures/trans_softmax_attn_fraction.png

In Figure˜32, metrics on the post-activation attention matrices are used and show comparable behavior in the first half of training.
In the second half of training, the $\mathrm{SigmoidAttn}$ can be seen to reduce in norm and in sparsity. (see following discussion of Figure˜33 for further insights).

Figure: Figure 32: Norm
Refer to caption: figures/norm_training_1_head.png

In Figure˜33, we see post-activation attention values for eight samples at training progresses. The most notable difference between the activations is, that by the end of training, the $\mathrm{SigmoidAttn}$ is less sparse in the $\mathcal{N}(0,1)$ self-attention in the upper-left quadrant. We can see that softmax tends to produce sparser values (as it is designed to) while sigmoid controls the magnitude and location of peak attention independently, leading to a less sparse attention at the end of training.

Figure: Figure 33: For 8 samples, the post-activation attentions is visualized as training progresses on the $k=1,n=40$ summation problem. The model has one head to simplify the visualization. The attention is shown in pairs for each sample with softmax attention is in black and sigmoid is in blue. A $2\times 2$ block structure is evident in both cases, resulting from each halve of the input containing different information.
Refer to caption: figures/attn_progress_by_sample_large2.png

#### G.5.4 Pair Repeat Problem

We define a synthetic task of identifying if the first two symbols in a sequence repeat. The symbols, $s_{i}$ below come from a fixed vocabulary of size $K$, and the repeat location (when present) is uniformly distributed in the sequence.

$f(s_{0},s_{1},s_{2},...,s_{N})=\begin{cases}1,&\text{ if }\exists\ n>1\mid(s_{0},s_{1})=(s_{n},s_{n+1}),\\
0&\text{ otherwise}\end{cases}$

A simple two layer transformer is trained on this problem. The model has an embedding dimension of 160, MLP ratio of 4, QK norm, and layers with eight heads. The results for different model architectures are shown in Figure Figure˜34. 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. A cosine learning rate schedule with 5% linear warmup and a maximum learning rate of 1e-3 is used with the Adam optimizer.

In this result, we see the sigmoid activation has higher data efficiency and similar fall-off in the out of distribution cases. From shorter runs, we estimate that the softmax network would fit the training with 4–5x 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 to compensate for this proposed cause in the softmax case.

Figure: Figure 34: 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 a similar number of parameters. The sigmoid transformer has a learned offset initialized to -4.
Refer to caption: figures/seq_length_results.png

## Appendix H Contributions

All authors contributed to writing this paper, designing the experiments, discussing results at each stage of the project.

##### Preliminary work

Preliminary viability of $\mathrm{SigmoidAttn}$ done by Jason Ramapuram.

##### Universal Function Approximation

Proof of UFA (Sections˜3.1 and C) sculpted by Federico Danieli.

##### Lipschitzness of Sigmoid Attention

Lipschitzness analysis (Sections˜3.2 and D) molded by Pierre Ablin.

##### FlashSigmoid

Implementation and analysis driven by Eeshan Dhekane in collaboration with Jagrit Digani (Sections˜4 and F).

##### 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 (Sections˜5.5 and G.3) driven by Floris Weers.

##### Stability Analysis

QK norm (Figure˜8), LayerScale (Figure˜8) and bias (Figure˜20) ablations crafted by Dan Busbridge using Attention Simulator. Attention Simulator written by Jason Ramapuram and used to validate norm growth (Figures˜6, 6, 6, 6, 19 and 19).

##### ASR Results

All ASR experiments (Section˜5.4) and ablations (Section˜G.4.2) are conducted by Tatiana Likhomanenko in discussions with Jason Ramapuram and Zijin Gu. Baseline ASR models code is written by Zijin Gu and Tatiana Likhomanenko. Baseline models are optimized by Zijin Gu to be close to state-of-the-art results.

##### Vision Results

All vision experiments (Sections˜5.2 and 5.3) and ablations (Figures˜9, G.2.1, 22 and 23) conducted and written by Jason Ramapuram.

##### Simple Experiments

Simple experiments to compare $\mathrm{SigmoidAttn}$ to $\mathrm{SoftmaxAttn}$, including visualizing attention evolution and simple sequence length generalization analysis (Section˜G.5) conducted by Russ Webb.

## Peer Reviews

confidence: 5
contribution: 2
initial_rating: 6
presentation: 3
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]. 
1) is it worth to  switch from original softmax attention to sigmoid attention? 
2) will training Sigmoid-Attention be more stable than original Softmax Attention?
soundness: 3
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 . 
 
The 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. 

The experimental part (Section 5) is very good. Most interesting parts 
1.  observation that training of sigmoid-attention based transformers is very sensitive to initial norm of |sigma (QK)"|v| . 
2. demonstarion that combination of layer-scale with  QK norm can help to  stabilize training in some cases
3. very solid ablation study (CV, ASR, LM) 
 
I 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) 

The paper is well writen and it is easy to read.
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.
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"?: 
- will Sigmoid-Attention  be more stable than original Softmax Attention? 
- will it help with long context? 
On the positive  side, thanks to flash-sigmoid we can see some speed-up for LLM inference (Table 3) 

More details: 
The 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? 
The 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)
 
The 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.  

The analysis of computational complexity is based on computing flops. But the real complexity should take into consideration memory access (read/write ) instead. 

A few comments related to the experimental part: 
-  it looks like the proposed remedy ( bias b=-10  only partially helps to control the  increase in the amplitude of |act * V} (graph 6) 
- ASR: sigmoid (RoPE with bias) performs badly on long audio ( Table 2) 
- LLM: Sigmoid (ALibi) underperform vs Softmax (Alibi) . No results for Sigmoid (RoPE)

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 4
questions: 1. It would be helpful if the authors could make Fig. 2 clearer.

2. It would be better if the authors conducted experiments on large-scale language models (around 7B parameters).
soundness: 3
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.

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

**Clarity:** The paper is written very well and easy to follow. 

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

2. Table 2 shows that sigmoid attention causes unstable training, whereas softmax does not. I believe a theoretical analysis of training stability is necessary.

---

confidence: 5
contribution: 2
initial_rating: 5
presentation: 3
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?

2). 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?

3). 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?

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

5). 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?
soundness: 3
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.

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

**Clarity:** The paper is well written and easy to follow with lots of details provided in the appendix. 

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

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

Also 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? 

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

Meta-review: There is a disagreement between reviewers whether the paper should be accepted. On one had the paper does a good set of experiments establishing the advantages and drawbacks of sigmoid over softmax for attention. On the other hand there is no strong evidence to prefer sigmoid over softmax. Further the discussion of theoretical results for sigmoid only caused more confusion for the reviewers. I believe the paper does a good exploration, even though doesnt definitely address the softmax alternate question, and provides valuable code that others can build on towards giving more concrete answers to this question. Hence I suggest acceptance.

Decision: Accept (Poster)