Skip to contents

Invariant point attention (AlphaFold structure module)

Source: R/Alfipa.R
Alfipa.Rd

Algorithm 22 of the Supplementary Information to Jumper et al. (2021), p. 28. Attention over residues whose logit carries, besides the usual query-key term and a pair bias, a squared distance between query and key points placed in each residue's own rigid frame. Because those points are compared in the global frame and the output points are mapped back through the inverse of the receiving frame, the operation is invariant under a global rigid motion of all frames (supplement eq. 3-6).

Usage

Alfipa(s, z, frames, wq, wk, wv, wqp, wkp, wvp, wb, gamma, wo)

Arguments

s

Single representation, an n x cs matrix.

z

Pair representation, an n x n x cz array.

frames

List of backbone frames, each list(R = 3x3, t = 3).

wq, wk, wv

Lists of per-head scalar projections, each c x cs.

wqp, wkp

Lists over heads of lists over points of 3 x cs projections (line 2).

wvp

Lists over heads of lists over points of 3 x cs projections (line 3).

wb

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

gamma

Per-head scalar for the point term of line 7, passed after the model's softplus so nothing is implicit.

wo

Output projection of line 11.

Value

A list with the update s, attn, the local output points, estimate, n and method.

References

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

Examples

# \donttest{
w <- function(r, c, k) matrix(sin(seq_len(r * c) * k + k), r, c)
rot <- function(a) { q <- c(1, a) / sqrt(1 + sum(a^2))
  rbind(c(1 - 2 * (q[3]^2 + q[4]^2), 2 * (q[2] * q[3] - q[1] * q[4]), 2 * (q[2] * q[4] + q[1] * q[3])),
        c(2 * (q[2] * q[3] + q[1] * q[4]), 1 - 2 * (q[2]^2 + q[4]^2), 2 * (q[3] * q[4] - q[1] * q[2])),
        c(2 * (q[2] * q[4] - q[1] * q[3]), 2 * (q[3] * q[4] + q[1] * q[2]), 1 - 2 * (q[2]^2 + q[3]^2))) }
s <- rbind(c(0.2, 0.5), c(-0.3, 0.1), c(0.4, -0.6))
z <- array(cos(seq_len(18) * 0.41), c(3, 3, 2))
fr <- list(list(R = rot(c(0.1, 0, 0)), t = c(0, 0, 0)), list(R = rot(c(0, 0.2, 0)), t = c(1, 0, 0)),
           list(R = rot(c(0, 0, 0.3)), t = c(0, 1, 1)))
r <- Alfipa(s, z, fr, list(w(2, 2, 0.1)), list(w(2, 2, 0.2)), list(w(2, 2, 0.3)),
            list(list(w(3, 2, 0.4))), list(list(w(3, 2, 0.5))), list(list(w(3, 2, 0.6))),
            wb = w(1, 2, 0.7), gamma = 0.9, wo = w(2, 8, 0.8))
round(r$attn[[1]], 3)
#>       [,1]  [,2]  [,3]
#> [1,] 0.437 0.357 0.206
#> [2,] 0.444 0.417 0.139
#> [3,] 0.408 0.270 0.322
# }