Transformer decoder block, post-norm (Geron Ch 16, morie.fn grtdb)

LN(x + masked MHA) -> LN(h1 + cross(h1, enc)) -> LN(h2 + FFN); the causal mask is built in and the cross-attention is unmasked.

Usage

morie_geron_transformer_decoder_block(x, encoder_output, weights, eps = 1e-05)

Arguments

x

(T, d_model) target tokens.

encoder_output

(S, d_model) source states.

weights

Named list with self, cross and ffn blocks.

eps

Layer-norm floor.

Value

List with output, self_attention_weights, cross_attention_weights, h1, h2, causal_mask.

Examples

set.seed(3)
X <- matrix(rnorm(12), 3, 4)
E <- matrix(rnorm(20), 5, 4)
h1 <- matrix(rnorm(8, 0, 0.4), 4, 2)
mk <- function() 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))
w <- list(self = mk(), cross = mk(),
          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_decoder_block(X, E, w)
str(r, max.level = 1)
#> List of 8
#>  $ output                 : num [1:3, 1:4] -0.91 -1.052 0.524 -0.93 0.125 ...
#>  $ self_attention_weights :List of 2
#>  $ cross_attention_weights:List of 2
#>  $ h1                     : num [1:3, 1:4] -0.927 -1.049 0.995 -0.87 0.106 ...
#>  $ h2                     : num [1:3, 1:4] -0.92631 -1.06432 0.93403 -0.86418 -0.00608 ...
#>  $ causal_mask            : logi [1:3, 1:3] TRUE TRUE TRUE FALSE TRUE TRUE ...
#>  $ estimate               : num [1:3, 1:4] -0.91 -1.052 0.524 -0.93 0.125 ...
#>  $ n                      : int 3