Encoder

Stack of multihead attention blocks.

Usage

Encoder(
  num_blocks,
  d_model,
  nhead,
  dim_feedforward,
  dropout = 0,
  activation = "gelu",
  norm_first = TRUE,
  bias_free_ln = FALSE,
  use_rope = FALSE,
  rope_base = 100000L,
  rope_interleaved = TRUE,
  ssmax = FALSE,
  recompute = FALSE
)

Arguments

num_blocks

Integer. Number of multihead attention blocks in the stack.

d_model

Integer. Model dimension.

nhead

Integer. Number of attention heads and should be a divisor of d_model.

dim_feedforward

Integer. Dimension of the feedforward network in each block.

dropout

Float, default 0. Dropout probability.

activation

Character string or unary function, default "gelu". The activation function used in the feedforward network.

norm_first

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

bias_free_ln

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

use_rope

Logical, default FALSE. Whether to use rotary positional encoding.

rope_base

Integer, default 100000L. A base scaling factor for rotary position encoding.

rope_interleaved

Logical, default TRUE. If TRUE, uses interleaved rotation where dimension pairs are (1,2), (3,4), etc. If FALSE, uses non-interleaved rotation where the embedding is split into first half and second half.

ssmax

Logical or character string, default FALSE. Type of scalable softmax to use. If TRUE, equivalent to "qassmax-mlp-elementwise". If FALSE, equivalent to "none". If a string, uses the specified scalable softmax type. Options include: "none", "ssmax", "ssmax-mlp", "ssmax-mlp-elementwise", "qassmax-mlp", "qassmax-mlp-elementwise".

recompute

Logical, default FALSE. If TRUE, uses gradient checkpointing to save memory at the cost of additional computation.

Value

An nn_module instance of class Encoder.

Methods

Usage

enc <- Encoder(num_blocks = 6L, d_model = 128L, nhead = 8L, dim_feedforward = 256L)
enc$forward(src, train_size = NULL)
enc$forward_with_cache(src, icl_cache, train_size, use_cache, store_cache)

See also

multihead_attention_block, RotaryEmbedding, KVCache, KVCacheEntry