Swin windowed self-attention (morie.fn grswin)

Attention inside M x M windows, O(HW M^2) rather than O((HW)^2); a non-zero shift rolls the partition by -shift on both spatial axes before windowing and rolls the output back.

Usage

morie_geron_swin_window_attention(X, window_size, WQ, WK, WV, shift = 0)

Arguments

X

(H, W, d_model) array.

window_size

M dividing both H and W.

WQ, WK, WV

Projections with d_model rows.

shift

Cyclic shift in [0, M-1].

Value

List with output, window_weights, n_windows, tokens_per_window, shift.

Examples

set.seed(7)
X <- array(rnorm(4 * 4 * 3), c(4, 4, 3))
W <- diag(3)
r <- morie_geron_swin_window_attention(X, window_size = 2,
                                       WQ = W, WK = W, WV = W)
str(r, max.level = 1)
#> List of 7
#>  $ output           : num [1:4, 1:4, 1:3] 2.2255 -0.9251 -0.2448 -0.0948 -0.9328 ...
#>  $ window_weights   :List of 4
#>  $ n_windows        : int 4
#>  $ tokens_per_window: int 4
#>  $ shift            : int 0
#>  $ estimate         : num [1:4, 1:4, 1:3] 2.2255 -0.9251 -0.2448 -0.0948 -0.9328 ...
#>  $ n                : int 16