Multi-head attention (Geron Ch 16, morie.fn grmha)

h heads of width d_model/h, each through the shared attention kernel, concatenated then mixed by WO.

Usage

morie_geron_multi_head_attention(Q, K, V, WQ, WK, WV, WO, h, mask = NULL)

Arguments

Q, K, V

(T, d_model) inputs.

WQ, WK, WV, WO

Projection matrices with d_model rows.

h

Number of heads dividing d_model.

mask

Optional attention mask.

Value

List with output, head_outputs, attention_weights, concat, d_head, d_model, n_heads.

Examples

set.seed(6)
Q <- matrix(rnorm(8), 2, 4); K <- Q; V <- Q
W <- diag(4)
morie_geron_multi_head_attention(Q, K, V, WQ = W, WK = W, WV = W,
                                 WO = W, h = 2)
#> $output
#>            [,1]     [,2]       [,3]       [,4]
#> [1,] -0.2593845 1.373508 0.06931878 -1.0404125
#> [2,] -0.4589616 1.563977 0.28590154  0.2495103
#> 
#> $head_outputs
#> $head_outputs[[1]]
#>            [,1]     [,2]
#> [1,] -0.2593845 1.373508
#> [2,] -0.4589616 1.563977
#> 
#> $head_outputs[[2]]
#>            [,1]       [,2]
#> [1,] 0.06931878 -1.0404125
#> [2,] 0.28590154  0.2495103
#> 
#> 
#> $attention_weights
#> $attention_weights[[1]]
#>           [,1]      [,2]
#> [1,] 0.4119658 0.5880342
#> [2,] 0.1901127 0.8098873
#> 
#> $attention_weights[[2]]
#>           [,1]      [,2]
#> [1,] 0.8687428 0.1312572
#> [2,] 0.2388443 0.7611557
#> 
#> 
#> $concat
#>            [,1]     [,2]       [,3]       [,4]
#> [1,] -0.2593845 1.373508 0.06931878 -1.0404125
#> [2,] -0.4589616 1.563977 0.28590154  0.2495103
#> 
#> $d_head
#> [1] 2
#> 
#> $d_model
#> [1] 4
#> 
#> $n_heads
#> [1] 2
#> 
#> $estimate
#>            [,1]     [,2]       [,3]       [,4]
#> [1,] -0.2593845 1.373508 0.06931878 -1.0404125
#> [2,] -0.4589616 1.563977 0.28590154  0.2495103
#> 
#> $n
#> [1] 2
#>