Neural Network Architecture
4,387,749 PARAMSEnd-to-end data flow: ESM-2 embeddings โ Pairformer โ Diffusion ScoreNetwork โ FunctionHead
Forward Pass
x = ESM2(seq) # [B, L, 1280]
s, z = Pairformer(x) # [B, L, 256], [B, L, L, 64]
coords = ScoreNetwork(x_t, t, s) # [B, L, 3]
logits = FunctionHead(s.mean(1)) # [B, 354]
Input Layer
ESM-2 650M encoder โ 1280-dim per-residue embeddings (frozen)
Sequence โ torch.Tensor[B, L, 1280] via fair-esm API
Projection
Linear(1280 โ d_single=256) ยท LayerNorm
Reduces ESM-2 dim to Pairformer working dimension
Outer Product Mean
Pair representation init: Z_ij = mean(a_i โ a_j) โ d_pair=64
Symmetric pair features for all (i,j) residue pairs
Pairformer ร 4
4 blocks: TriangleMultiplication (incoming/outgoing) + Attention + LayerNorm
Refines single (256-dim) and pair (64-dim) representations jointly
ScoreNetwork
MLP denoiser: (d_single + t_emb + 3) โ 512 โ 256 โ 3 (Cฮฑ coords)
Score function s_ฮธ(x_t, t, c) conditioned on Pairformer output c
FunctionHead
Linear(256โ512) โ ReLU โ Dropout(0.1) โ Linear(512โ354)
354 GO terms ยท multilabel BCE ยท threshold ฯ* optimised on val set
Dimension Summary
| Component | Input Dim | Output Dim | Notes |
|---|---|---|---|
| ESM-2 | AA sequence | [L, 1280] | Frozen pretrained |
| Projection | [L, 1280] | [L, 256] | Linear + LN |
| OuterProductMean | [L, 256] | [L, L, 64] | Pair init |
| Pairformer block (ร4) | [L,256],[L,L,64] | [L,256],[L,L,64] | Triangle + Attn |
| ScoreNetwork | [L, 256+t+3] | [L, 3] | Cฮฑ coordinates |
| FunctionHead | [256] | [354] | Mean-pooled single repr |