Enhanced multi-head attention with RoPE, scalable softmax, and KV caching

Enhanced multi-head attention with RoPE, scalable softmax, and KV caching

Usage

multihead_attention(embed_dim, num_heads, dropout = 0, ssmax = FALSE)

Arguments

embed_dim

Integer. Model dimension (total size of each attention head combined).

num_heads

Integer. Number of attention heads.

dropout

Numeric. Dropout probability applied to attention weights (default: 0.0).

ssmax

Logical or character. Type of scalable softmax to use:

• FALSE or "none": No scaling applied

• TRUE or "qassmax-mlp-elementwise": Elementwise query-aware scaling

• Other strings: specific scalable softmax type (default: FALSE)

Value

A nn_module for multi-head attention with advanced features.

Note

The implementation always uses batch_first = TRUE, so input tensors have shape (.., seq_len, embed_dim).

References

1. Su et al., "RoFormer: Enhanced Transformer with Rotary Position Embedding"

2. Liu et al., "Scalable-Softmax Is Superior for Attention"