Sigmoid Attention

Sigmoid attention with element-wise normalization.

Replaces the row-wise softmax of standard attention with an element-wise sigmoid followed by L1 normalization. This overcomes the zero-sum token competition inherent in softmax, allowing each token to independently determine its relevance to every other token. Achieves ~17% kernel speedup via FlashSigmoid and requires hybrid-norm stabilization for training.

Reference:

Ramapuram et al. (2024), “Sigmoid Attention: Overcoming the Zero-Sum Token Competition”, arXiv:2409.04431.

class src.model.attention.sigmoid.SigmoidAttention(*args: Any, **kwargs: Any)

Bases: Module

Sigmoid attention with element-wise normalization (Ramapuram et al. 2024).

Projects the input into query, key, and value tensors, computes scaled dot-product attention scores, applies element-wise sigmoid, and normalizes by the sum of weights per query position. Unlike softmax attention, each key-value pair’s contribution is determined independently, eliminating the zero-sum competition among tokens.

Reference:

Ramapuram et al. (2024), “Sigmoid Attention: Overcoming the Zero-Sum Token Competition”, arXiv:2409.04431.

Parameters:

config – Model configuration object with attributes hidden_size, num_heads, dropout, use_bitnet, and optionally mode ("encoder" or "decoder").

hidden_size

Dimensionality of the input and output embeddings.

num_heads

Number of parallel attention heads.

head_dim

Dimensionality of each attention head (hidden_size // num_heads).

scale

Scaling factor 1 / sqrt(head_dim) applied to dot products.

eps

Small constant for numerical stability in L1 normalization.

q_proj

Linear (or BitLinear) projection for queries.

k_proj

Linear (or BitLinear) projection for keys.

v_proj

Linear (or BitLinear) projection for values.

out_proj

Linear (or BitLinear) output projection.

dropout

Dropout layer applied to attention weights.

mode

"encoder" for bidirectional attention, "decoder" for causal (upper-triangular) masking.

Raises:

ValueError – If hidden_size is not divisible by num_heads.

forward(x: torch.Tensor, logical_layer_idx: int | None = None) → torch.Tensor

Compute sigmoid attention with L1 normalization.

Parameters:

• x – Input tensor of shape (batch_size, seq_len, hidden_size).

• logical_layer_idx – Logical layer index (unused; accepted for interface compatibility with other attention modules).

Returns:

Output tensor of shape (batch_size, seq_len, hidden_size).