Triangular multiplicative update of the AlphaFold pair representation

Algorithms 11 and 12 of the Supplementary Information to Jumper et al. (2021), p. 18. Each edge ij is updated from the other two edges of every triangle it belongs to: the "outgoing" form contracts over a_ik * b_jk, the "incoming" form over a_ki * b_kj. Those are the only lines that differ between the two algorithms.

Usage

Alftrimu(z, wag, wav, wbg, wbv, wg, wo, mode = "outgoing", layernorm = TRUE)

Arguments

z

Pair representation, an n x n x cz array.

wag, wav, wbg, wbv

Gate and value projections for the left and right edges, each c x cz (line 2).

wg

Output gate projection, cz x cz (line 3).

wo

Output projection, cz x c (line 4).

mode

Either "outgoing" (Algorithm 11) or "incoming" (Algorithm 12).

layernorm

Apply the layer normalisations of lines 1 and 4. Setting it to FALSE exposes the degenerate single-channel reduction to a plain matrix product.

Value

A list with the updated pair tensor z, its mean estimate, n, mode and method.

Details

All weights are supplied by the caller, so the result is a fixed function of its arguments.

References

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

Examples

set.seed(13)
z <- array(rnorm(3 * 3 * 2), c(3, 3, 2))
w <- function() matrix(rnorm(4, sd = 0.5), 2, 2)
dim(Alftrimu(z, w(), w(), w(), w(), w(), w(), mode = "outgoing")$z)
#> [1] 3 3 2