Triangular gated self-attention over the AlphaFold pair representation

Algorithms 13 and 14 of the Supplementary Information to Jumper et al. (2021), pp. 19-20. Edge ij attends over the edges sharing its starting node i or its ending node j. The attention logit is a scaled query-key product modulated by a bias taken from the third edge of the triangle, which is what makes it triangular rather than ordinary attention.

Usage

Alftriat(z, wq, wk, wv, wb, wg, wo, mode = "starting")

Arguments

z

Pair representation, an n x n x cz array.

wq, wk, wv

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

wb

Per-head bias projection, nhead x cz (line 3).

wg

List of per-head gate projections, each c x cz (line 4).

wo

Output projection, cz x (nhead * c) (line 7).

mode

Either "starting" (Algorithm 13) or "ending" (Algorithm 14).

Value

A list with z, the attention array attn, estimate, n, mode and method.

References

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

Examples

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