Attention#

class serket.nn.MultiHeadAttention(num_heads, q_features, k_features=None, v_features=None, out_features=None, *, key, q_weight_init='glorot_uniform', q_bias_init='zeros', q_dtype=<class 'jax.numpy.float32'>, k_weight_init='glorot_uniform', k_bias_init='zeros', k_dtype=<class 'jax.numpy.float32'>, v_weight_init='glorot_uniform', v_bias_init='zeros', v_dtype=<class 'jax.numpy.float32'>, out_weight_init='glorot_uniform', out_bias_init='zeros', out_dtype=<class 'jax.numpy.float32'>, drop_rate=0.0, drop_broadcast=False)[source]#

Multi-head attention module.

The module consists of linear projections for query, key, and value (q, k, v), followed by the application of scaled dot-product attention with optional masking and dropout. It supports multi-head attention where the input features are divided into multiple heads to capture different aspects of relationships between the input vectors.

Parameters:

num_heads (int) – Number of attention heads.
q_features (int) – Number of features for the query.
k_features (Optional[int]) – Number of features for the key.
v_features (Optional[int]) – Number of features for the value.
out_features (Optional[int]) – Number of features for the output.
key (Array) – Key for the random number generator.
q_weight_init (Union[Literal['he_normal', 'he_uniform', 'glorot_normal', 'glorot_uniform', 'lecun_normal', 'lecun_uniform', 'normal', 'uniform', 'ones', 'zeros', 'xavier_normal', 'xavier_uniform', 'orthogonal'], Callable[[Array, Tuple[int, ...], Union[dtype, str, Any]], Array | None]]) – Initializer for the query weight. Defaults to glorot_uniform.
q_bias_init (Union[Literal['he_normal', 'he_uniform', 'glorot_normal', 'glorot_uniform', 'lecun_normal', 'lecun_uniform', 'normal', 'uniform', 'ones', 'zeros', 'xavier_normal', 'xavier_uniform', 'orthogonal'], Callable[[Array, Tuple[int, ...], Union[dtype, str, Any]], Array | None]]) – Initializer for the query bias. Defaults to zeros. use None to disable bias.
q_dtype (Union[dtype, str, Any]) – Data type for the query. float32
k_weight_init (Union[Literal['he_normal', 'he_uniform', 'glorot_normal', 'glorot_uniform', 'lecun_normal', 'lecun_uniform', 'normal', 'uniform', 'ones', 'zeros', 'xavier_normal', 'xavier_uniform', 'orthogonal'], Callable[[Array, Tuple[int, ...], Union[dtype, str, Any]], Array | None]]) – Initializer for the key weight. Defaults to glorot_uniform.
k_bias_init (Union[Literal['he_normal', 'he_uniform', 'glorot_normal', 'glorot_uniform', 'lecun_normal', 'lecun_uniform', 'normal', 'uniform', 'ones', 'zeros', 'xavier_normal', 'xavier_uniform', 'orthogonal'], Callable[[Array, Tuple[int, ...], Union[dtype, str, Any]], Array | None]]) – Initializer for the key bias. Defaults to zeros. use None to disable bias.
k_dtype (Union[dtype, str, Any]) – Data type for the key. float32
v_weight_init (Union[Literal['he_normal', 'he_uniform', 'glorot_normal', 'glorot_uniform', 'lecun_normal', 'lecun_uniform', 'normal', 'uniform', 'ones', 'zeros', 'xavier_normal', 'xavier_uniform', 'orthogonal'], Callable[[Array, Tuple[int, ...], Union[dtype, str, Any]], Array | None]]) – Initializer for the value weight. Defaults to glorot_uniform.
v_bias_init (Union[Literal['he_normal', 'he_uniform', 'glorot_normal', 'glorot_uniform', 'lecun_normal', 'lecun_uniform', 'normal', 'uniform', 'ones', 'zeros', 'xavier_normal', 'xavier_uniform', 'orthogonal'], Callable[[Array, Tuple[int, ...], Union[dtype, str, Any]], Array | None]]) – Initializer for the value bias. Defaults to zeros. use None to disable bias.
v_dtype (Union[dtype, str, Any]) – Data type for the value. float32
out_weight_init (Union[Literal['he_normal', 'he_uniform', 'glorot_normal', 'glorot_uniform', 'lecun_normal', 'lecun_uniform', 'normal', 'uniform', 'ones', 'zeros', 'xavier_normal', 'xavier_uniform', 'orthogonal'], Callable[[Array, Tuple[int, ...], Union[dtype, str, Any]], Array | None]]) – Initializer for the output weight. Defaults to glorot_uniform.
out_bias_init (Union[Literal['he_normal', 'he_uniform', 'glorot_normal', 'glorot_uniform', 'lecun_normal', 'lecun_uniform', 'normal', 'uniform', 'ones', 'zeros', 'xavier_normal', 'xavier_uniform', 'orthogonal'], Callable[[Array, Tuple[int, ...], Union[dtype, str, Any]], Array | None]]) – Initializer for the output bias. Defaults to zeros. use None to disable bias.
out_dtype (Union[dtype, str, Any]) – Data type for the output. float32
drop_rate (float) – Dropout rate. defaults to 0.0.
drop_broadcast (bool) – Whether to broadcast the dropout mask across the batch dimension and the heads dimension. Defaults to False.

Example

>>> import serket as sk
>>> import jax.numpy as jnp
>>> import jax.random as jr
>>> batch = 3
>>> num_heads = 2
>>> q_features = 4
>>> k_features = 8
>>> v_features = 6
>>> q_length = 4
>>> kv_length = 2
>>> mask = jr.uniform(jr.key(0), (batch, num_heads, q_length, kv_length))
>>> mask = (mask > 0.5).astype(jnp.float32)
>>> k1, k2, k3, k4 = jr.split(jr.key(0), 4)
>>> q = jr.uniform(k1, (batch, q_length, q_features))
>>> k = jr.uniform(k2, (batch, kv_length, k_features))
>>> v = jr.uniform(k3, (batch, kv_length, v_features))
>>> layer = sk.nn.MultiHeadAttention(
...    num_heads,
...    q_features,
...    k_features,
...    v_features,
...    drop_rate=0.0,
...    key=k4,
... )
>>> print(layer(q, k, v, mask=mask, key=jr.key(1)).shape)
(3, 4, 4)

Note

If k_features, v_features, out_features are not specified, they are set to q_features.
To disable attention Dropout, use tree_eval() on the instantiated layer.

>>> import serket as sk
>>> layer = sk.nn.MultiHeadAttention(1, 1, key=jr.key(0))
>>> print(repr(layer.dropout))
Dropout(drop_rate=0.0, drop_axes=None)
>>> print(repr(sk.tree_eval(layer).dropout))
Identity()

Note

MultiHeadAttention supports lazy initialization, meaning that the weights and biases are not initialized until the first call to the layer. This is useful when the input shape is not known at initialization time.

To use lazy initialization, pass None as the q_features argument and use value_and_tree() to call the layer with an input of known shape.

>>> import jax.random as jr
>>> import serket as sk
>>> k1, k2, k3, k4, k5 = jr.split(jr.key(0), 5)
>>> q = jr.uniform(k1, (3, 2, 6))
>>> k = jr.uniform(k2, (3, 2, 6))
>>> v = jr.uniform(k3, (3, 2, 6))
>>> lazy = sk.nn.MultiHeadAttention(2, None, key=k4)
>>> _, material = sk.value_and_tree(lambda lazy: lazy(q, k, v, key=k4))(lazy)
>>> material(q, k, v, key=k5).shape
(3, 2, 6)

Reference:

__call__(q_input, k_input, v_input, mask=None, *, key=None)[source]#

Applies multi-head attention to the given inputs.

Parameters:

q_input (Array) – Query input. […, q_length, q_features]
k_input (Array) – Key input. […, kv_length, k_features]
v_input (Array) – Value input. […, kv_length, v_features]
mask (Optional[Array]) – Mask input. […, num_heads, q_length, kv_length] Defaults to None. for no masking.
key (Optional[Array]) – Key for the random number generator used for dropout. Defaults to None for no dropout.

Return type:

Array

Attention

Contents

Attention#