Transformer encoder block, post-norm (Geron Ch 16, morie.fn grteb)

h = LN(x + MHA(x)); y = LN(h + FFN(h)).

Usage

morie_geron_transformer_encoder_block(
  x,
  mha_weights,
  ffn_weights,
  mask = NULL,
  eps = 1e-05
)

Arguments

x

(T, d_model) tokens.

mha_weights

Attention weights, optionally with gamma/beta.

ffn_weights

Feed-forward weights, optionally with gamma/beta.

mask

Optional attention mask.

eps

Layer-norm floor.

Value

List with output, attention_output, attention_weights, hidden, ffn_output.

Examples

set.seed(2)
X <- matrix(rnorm(12), 3, 4)
h1 <- matrix(rnorm(8, 0, 0.4), 4, 2)
mha <- list(WQ = list(h1, h1), WK = list(h1, h1), WV = list(h1, h1),
            WO = matrix(rnorm(16, 0, 0.4), 4, 4),
            gamma = rep(1, 4), beta = rep(0, 4))
ffn <- list(W1 = matrix(rnorm(32, 0, 0.4), 4, 8),
            W2 = matrix(rnorm(32, 0, 0.4), 8, 4),
            gamma = rep(1, 4), beta = rep(0, 4))
r <- morie_geron_transformer_encoder_block(X, mha, ffn)
str(r, max.level = 1)
#> List of 7
#>  $ output           : num [1:3, 1:4] 0.155 0.593 0.658 -1.211 -1.345 ...
#>  $ attention_output : num [1:3, 1:4] 1.224 0.301 0.716 -0.245 -0.623 ...
#>  $ attention_weights:List of 2
#>  $ hidden           : num [1:3, 1:4] 0.333 1.106 0.863 -1.462 -1.565 ...
#>  $ ffn_output       : num [1:3, 1:4] -0.172 -0.648 -0.171 0.122 -0.581 ...
#>  $ estimate         : num [1:3, 1:4] 0.155 0.593 0.658 -1.211 -1.345 ...
#>  $ n                : int 3