SSMax – Scalable Softmax with Learnable Per-Head Scaling

Applies scaling to queries: \(q_{\text{scaled}} = q \cdot (s \cdot \log n)\), where \(s\) is a learnable per-head parameter.

Usage

SSMax(num_heads)

Arguments

num_heads

integer(1) Number of attention heads.

Value

An nn_module object (class generator).

See also

create_ssmax_layer