Linformer linear-complexity attention via low-rank projection

Wang, Li, Khabsa, Fang and Ma (2020), "Linformer: Self-Attention with Linear Complexity", arXiv:2006.04768, Eq 7: with projection matrices E, F (k x n) applied to the key and value layers, head = softmax(Q (E K)' / sqrt(d_k)) (F V), so the attention matrix is n x k and cost drops from O(n^2) to O(n k). With k = n and E = F = I this reduces exactly to standard scaled dot-product attention (the test anchor).

Usage

Linatt(Q, K, V, E, F_)

Arguments

Q

Query matrix (n x d_k).

K

Key matrix (n x d_k).

V

Value matrix (n x d_v).

E

Key projection (k x n).

F_

Value projection (k x n).

Value

List with output, weights, projected_K, projected_V, k, estimate, n, method.

References

Wang, S., Li, B. Z., Khabsa, M., Fang, H. and Ma, H. (2020), arXiv:2006.04768, Section 4, Eq 7.

Examples

Linatt(Q = 0.5, K = 5L, V = 5L, E = c(1, 2, 3, 4, 5, 6, 7, 8), F_ = c(1, 2, 3, 4, 5, 6, 7, 8))
#> $output
#>          [,1]
#> [1,] 39.55287
#> 
#> $weights
#>              [,1]         [,2]         [,3]         [,4]         [,5]
#> [1,] 2.304884e-08 2.807923e-07 3.420751e-06 4.167328e-05 0.0005076844
#>             [,6]       [,7]     [,8]
#> [1,] 0.006184863 0.07534705 0.917915
#> 
#> $projected_K
#>      [,1]
#> [1,]    5
#> [2,]   10
#> [3,]   15
#> [4,]   20
#> [5,]   25
#> [6,]   30
#> [7,]   35
#> [8,]   40
#> 
#> $projected_V
#>      [,1]
#> [1,]    5
#> [2,]   10
#> [3,]   15
#> [4,]   20
#> [5,]   25
#> [6,]   30
#> [7,]   35
#> [8,]   40
#> 
#> $k
#> [1] 8
#> 
#> $estimate
#> [1] 39.55287
#> 
#> $n
#> [1] 1
#> 
#> $method
#> [1] "Linformer low-rank attention (Wang et al. 2020, Eq 7)"
#>