QASSMaxMLP – Query-Aware Scalable Softmax with MLPs

Applies scaling to queries using two MLPs:

\(q_{\text{scaled}} = q \cdot \text{base\_mlp}(\log n) \cdot (1 + \tanh(\text{query\_mlp}(q)))\)

The base MLP learns length-dependent scaling while the query MLP learns query-dependent modulation. The query MLP is zero-initialised so that the initial modulation is the identity (\(1 + 0 = 1\)).

Usage

QASSMaxMLP(num_heads, head_dim, n_hidden = 64L, elementwise = FALSE)

Arguments

num_heads

integer(1) Number of attention heads.

head_dim

integer(1) Dimension of each attention head.

n_hidden

integer(1) Number of hidden units in each MLP. Default: 64L.

elementwise

logical(1) If TRUE, apply elementwise scaling per head dimension. Default: FALSE.

Value

An nn_module object (class generator).

See also

create_ssmax_layer