Template pointwise attention of AlphaFold

Algorithm 17 of the Supplementary Information to Jumper et al. (2021), p. 21. For each residue pair the pair representation queries the stack of template features at that same pair, and the templates are pooled by attention. The attention runs over templates only, never across pairs, which is why it is pointwise: pair ij never sees pair kl.

Usage

Alftmpl(t, z, wq, wk, wv, wo)

Arguments

t

Template pair features, an ntempl x n x n x ct array.

z

Pair representation, an n x n x cz array.

wq

List of per-head query projections from z, c x cz.

wk, wv

Lists of per-head key and value projections from t, each c x ct.

wo

Output projection, cz x (nhead * c).

Value

A list with z, attn, estimate, n, ntempl and method.

References

Jumper et al (2021) Nature 596:583-589, Suppl. Algorithm 17

Examples

set.seed(7)
n <- 3; nt <- 2; ct <- 4; cz <- 4; c <- 2
t <- array(rnorm(nt * n * n * ct), c(nt, n, n, ct))
z <- array(rnorm(n * n * cz), c(n, n, cz))
wq <- list(matrix(rnorm(c * cz), c, cz))
wk <- list(matrix(rnorm(c * ct), c, ct)); wv <- list(matrix(rnorm(c * ct), c, ct))
wo <- matrix(rnorm(cz * c), cz, c)
dim(Alftmpl(t, z, wq, wk, wv, wo)$z)
#> [1] 3 3 4