Multi-head attention: split, attend, concat, project (Geron Ch 15, morie.fn hmmha)

Usage

morie_geron_multihead_attention(Q, K, V, n_heads, W_O = NULL, mask = NULL)

Arguments

Q, K, V: Query/key/value matrices.
n_heads: Head count (must divide Q/K width).
W_O: Optional output projection.
mask: Optional mask.

Value

List with output, head_outputs, attention_weights, d_head, estimate, n, method.

Examples

Q <- matrix(c(1, 0), nrow = 1)
K <- matrix(c(1, 0, 0, 1), nrow = 2, byrow = TRUE)
V <- K
morie_geron_multihead_attention(Q, K, V, n_heads = 2)
#> $output
#>           [,1] [,2]
#> [1,] 0.7310586  0.5
#> 
#> $concat
#>           [,1] [,2]
#> [1,] 0.7310586  0.5
#> 
#> $head_outputs
#> $head_outputs[[1]]
#>           [,1]
#> [1,] 0.7310586
#> 
#> $head_outputs[[2]]
#>      [,1]
#> [1,]  0.5
#> 
#> 
#> $attention_weights
#> $attention_weights[[1]]
#>           [,1]      [,2]
#> [1,] 0.7310586 0.2689414
#> 
#> $attention_weights[[2]]
#>      [,1] [,2]
#> [1,]  0.5  0.5
#> 
#> 
#> $d_head
#> [1] 1
#> 
#> $n_heads
#> [1] 2
#> 
#> $estimate
#> [1] 0.8856899
#> 
#> $n
#> [1] 1
#> 
#> $method
#> [1] "Multi-head attention (split, attend, concat, project)"
#>