MSA row-wise and column-wise gated self-attention (AlphaFold)

Algorithms 7 and 8 of the Supplementary Information to Jumper et al. (2021), pp. 15-16. The row-wise form attends over residues within one sequence and adds a bias read off the pair representation, which is how the pair stack speaks back to the MSA stack. The column-wise form attends over sequences at a fixed residue and carries no pair bias.

Usage

Alfmsaat(m, wq, wk, wv, wg, wo, z = NULL, wb = NULL, mode = "row")

Arguments

m

MSA representation, an s x n x cm array.

wq, wk, wv, wg

Lists of per-head projection matrices, each c x cm.

wo

Output projection, cm x (nhead * c).

z

Pair representation n x n x cz; required when mode is "row".

wb

Per-head pair-bias projection nhead x cz; required when mode is "row" (line 3 of Algorithm 7).

mode

Either "row" (Algorithm 7) or "column" (Algorithm 8).

Value

A list with the MSA update m, the attention attn, estimate, n, s, mode and method.

References

Jumper et al (2021) Nature 596:583-589, Suppl. Algorithms 7-8

Examples

set.seed(3)
m <- array(rnorm(2 * 3 * 2), c(2, 3, 2)); z <- array(rnorm(3 * 3 * 2), c(3, 3, 2))
w <- function() list(matrix(rnorm(4, sd = 0.5), 2, 2))
r <- Alfmsaat(m, w(), w(), w(), w(), wo = matrix(rnorm(4), 2, 2), z = z, wb = matrix(rnorm(2), 1, 2))
dim(r$m)
#> [1] 2 3 2