Attention block supporting RoPE, scalable softmax, and KV caching

Attention block supporting RoPE, scalable softmax, and KV caching

Usage

multihead_attention_block(
  d_model,
  nhead,
  dim_feedforward,
  dropout = 0,
  activation = "gelu",
  norm_first = TRUE,
  bias_free_ln = FALSE,
  ssmax = FALSE
)

Arguments

d_model

Integer. Model dimension.

nhead

Integer. Number of attention heads.

dim_feedforward

Integer. Dimension of the feedforward network.

dropout

Numeric. Dropout probability (default: 0.0).

activation

Character or function. Activation function for feedforward: "relu", "gelu", or a unary callable (default: "gelu").

norm_first

Logical. If TRUE, uses pre-norm architecture (LayerNorm before attention and feedforward) (default: TRUE).

bias_free_ln

Logical. If TRUE, removes bias from all LayerNorm layers (default: FALSE).

ssmax

Logical or character. Type of scalable softmax to use in attention (default: FALSE).

Value

A nn_module for transformer encoder layer with advanced features.