Per-head multi-head attention used by the transformer blocks

Per-head multi-head attention used by the transformer blocks

Usage

morie_geron_block_multi_head_attention(Q_in, KV_in, weights, mask = NULL)

Arguments

Q_in

Query-side input.

KV_in

Key/value-side input.

weights

Named list with WQ, WK, WV (lists of per-head matrices) and WO.

mask

Optional attention mask.

Value

List with output and weights (one per head).

Examples

set.seed(7)
Q <- matrix(rnorm(12), 3, 4)
KV <- matrix(rnorm(20), 5, 4)
h1 <- matrix(rnorm(8, 0, 0.5), 4, 2)
wts <- list(WQ = list(h1, h1), WK = list(h1, h1), WV = list(h1, h1),
            WO = matrix(rnorm(16, 0, 0.5), 4, 4))
r <- morie_geron_block_multi_head_attention(Q, KV, wts)
dim(r$output)
#> [1] 3 4