Skip to contents

Algorithm 6 of the Supplementary Information to Jumper et al. (2021), p. 14. One block runs, in order: MSA row attention with pair bias, MSA column attention, MSA transition, outer product mean into the pair stack, the two triangular multiplicative updates, the two triangular attentions, and the pair transition. Every sublayer is a residual update.

Usage

Alfevo(m, z, w, nblock = 1, drop = NULL)

Arguments

m

MSA representation, an s x n x cm array.

z

Pair representation, an n x n x cz array.

w

Named list of weights; see the Python mirror for the key list.

nblock

Number of blocks; a fixed count, with no early exit.

drop

Optional named list of multiplicative dropout masks.

Value

A list with m, z, the single representation s, estimate, nblock and method.

Details

All weights come from the caller and are shared across blocks, as in the published model. Dropout is applied only if the caller passes explicit masks, so inference is deterministic.

References

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

Examples

# \donttest{
# one Evoformer block at toy width (cm = cz = 2); every weight a small fixed matrix
set.seed(25)
m <- array(rnorm(8), c(2, 2, 2)); z <- array(rnorm(8), c(2, 2, 2))
nw <- function(r, c) matrix(rnorm(r * c, sd = 0.3), r, c)
w <- list(rowq = list(nw(2, 2)), rowk = list(nw(2, 2)), rowv = list(nw(2, 2)),
          rowg = list(nw(2, 2)), rowo = nw(2, 2), rowb = nw(1, 2),
          colq = list(nw(2, 2)), colk = list(nw(2, 2)), colv = list(nw(2, 2)),
          colg = list(nw(2, 2)), colo = nw(2, 2), mt1 = nw(3, 2), mt2 = nw(2, 3),
          opa = nw(2, 2), opb = nw(2, 2), opo = nw(2, 4),
          tmoag = nw(2, 2), tmoav = nw(2, 2), tmobg = nw(2, 2), tmobv = nw(2, 2),
          tmog = nw(2, 2), tmoo = nw(2, 2),
          tmiag = nw(2, 2), tmiav = nw(2, 2), tmibg = nw(2, 2), tmibv = nw(2, 2),
          tmig = nw(2, 2), tmio = nw(2, 2),
          tasq = list(nw(2, 2)), task = list(nw(2, 2)), tasv = list(nw(2, 2)),
          tasb = nw(1, 2), tasg = list(nw(2, 2)), taso = nw(2, 2),
          taeq = list(nw(2, 2)), taek = list(nw(2, 2)), taev = list(nw(2, 2)),
          taeb = nw(1, 2), taeg = list(nw(2, 2)), taeo = nw(2, 2),
          pt1 = nw(3, 2), pt2 = nw(2, 3), sout = nw(3, 2))
r <- Alfevo(m, z, w)
dim(r$s)
#> [1] 2 3
# }