Initial representation embeddings of AlphaFold

Algorithms 3-5 of the Supplementary Information to Jumper et al. (2021), p. 13. The pair representation starts as an outer sum of two projections of the target features plus a relative position encoding; the MSA representation starts as a projection of the MSA features with the target projection broadcast across sequences.

Usage

Alfembed(
  target_feat,
  residue_index,
  msa_feat,
  wa,
  wb,
  wrel,
  wmsa,
  wtgt,
  bins = NULL
)

Arguments

target_feat

Per-residue target features, an n x ctf matrix.

residue_index

Residue index per position; only differences matter.

msa_feat

MSA features, an s x n x cmf array.

wa, wb

The two target projections of line 1, each cz x ctf.

wrel

Relative position projection, cz x length(bins).

wmsa

MSA feature projection, cm x cmf.

wtgt

Target projection added to every MSA row, cm x ctf.

bins

Relative position bins; defaults to -32, ..., 32.

Value

A list with the initial z and m, the relative position contribution pos, estimate, n and method.

Details

The relative position encoding clips the separation at 32 residues, so nothing beyond that is distinguished. That is deliberate: it de-emphasises primary sequence distance and lets the network run on chains longer than those it was trained on.

References

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

Examples

set.seed(15)
n <- 4; ctf <- 3; cz <- 2; cm <- 2; cmf <- 3
tf <- matrix(rnorm(n * ctf), n, ctf)
msa <- array(rnorm(2 * n * cmf), c(2, n, cmf))
r <- Alfembed(tf, residue_index = 1:n, msa_feat = msa,
              wa = matrix(rnorm(cz * ctf), cz, ctf), wb = matrix(rnorm(cz * ctf), cz, ctf),
              wrel = matrix(rnorm(cz * 65), cz, 65), wmsa = matrix(rnorm(cm * cmf), cm, cmf),
              wtgt = matrix(rnorm(cm * ctf), cm, ctf))
dim(r$z); dim(r$m)
#> [1] 4 4 2
#> [1] 2 4 2