Protein sequence and inverse folding

ProteinMPNN node-and-edge message-passing layer

HardRaw PyTorchAbout 30 min2 public examples plus private tests

Free account with email, Google or GitHub. Your code runs in an isolated PyTorch sandbox.

Your task

Run one ProteinMPNN-style encoder update. Compute neighbor messages to update nodes, then use those updated nodes to update directed edges.

Function signature

node_edge_message_passing(
    activations: dict,
    parameters: dict,
    masks: dict,
    structure: dict,
    eps: float = 1e-05,
    aggregation: str = 'fixed',
    scale: float = 30.0,
) -> tuple[torch.Tensor, torch.Tensor]

Worked example

Messages first change the target node

Inspect the first aggregation stage for a valid target with h_v=[1,0]. Suppose its message MLP produces active messages [2,4] and [6,8]; a third slot is inactive and must contribute zero even if its pre-mask message is [100,100]. aggregation="mean" gives [4,6]. The residual sum is [5,6], with channel mean 5.5 and population variance 0.25. For eps=0.75, norm1_weight=[2,1] and norm1_bias=[0,1], v1=[−1,1.5]. The required feedforward and second normalization then produce v2; the edge pass must gather from that v2. For aggregation="fixed" with scale=4, the same active message sum would instead give aggregate [2,3].

Parameters

Parameter	Shape and meaning
`activations`dict	Supplied intermediate tensors.
`parameters`dict	Supplied model parameters, keyed as shown below.
`masks`dict	Supplied masks. Each field below defines what True includes or permits.
`structure`dict	Fields supplied together in this dictionary.
`eps`floatDefault `1e-05`	Positive value added to final-axis population variance inside each LayerNorm square root. It is not a mask/count denominator or an output offset.
`aggregation`strDefault `'fixed'`	"fixed" divides summed neighbor messages by scale; "mean" divides by the number of active neighbors, clamped to at least one.
`scale`floatDefault `30.0`	Positive fixed denominator for summed neighbor messages when aggregation="fixed". Ignored in mean mode.

Return

tuple[torch.Tensor, torch.Tensor]

Return these tensors in the stated order.

Required behavior

Implement a deterministic ProteinMPNN-shaped encoder layer with dropout disabled. h_v $[B,L,C]$ and h_e $[B,L,K,C]$ are finite floating states, $B,L,K,C\ge 1$ . $\mathrm{neighbors}[b,i,k]=j$ means source node $j$ sends a message to target $i$ through edge slot $k$ .
Indices are in $[0,L)$ ; repeated neighbors and self edges are allowed. All floating inputs share dtype/device. node_mask $[B,L]$ and edge_mask $[B,L,K]$ are bool; True means valid/enabled. An active slot requires edge_mask $[b,i,k]$ AND node_mask $[b,i]$ AND node_mask $[b,j]$ .
Each $\mathrm{MLP\_prefix}(x)$ has three affine layers: $\operatorname{GELU}(x\mathbin{@}\mathrm{prefix\_w1}+\mathrm{prefix\_b1})$ , then $\operatorname{GELU}(\ldots\mathbin{@}\mathrm{prefix\_w2}+\mathrm{prefix\_b2})$ , then $\ldots\mathbin{@}\mathrm{prefix\_w3}+\mathrm{prefix\_b3}$ .
Weights have shapes $[3C,C]$ , $[C,C]$ , $[C,C]$ ; all biases $[C]$ . GELU is exact $x\cdot (1+\operatorname{erf}(x/\sqrt{2}))/2$ . Initial messages use $\operatorname{CONCAT}(h_{v}[i], h_{e}[i,k], h_{v}[j])$ in TARGET, EDGE, SOURCE order. Zero inactive messages before summing over $K$ .
Config: aggregation is 'fixed' (default) or 'mean'; $\mathrm{scale}>0$ defaults to $30$ and is used only for fixed aggregation; $\varepsilon>0$ defaults to $10^{-5}$ . No other options. Fixed divides the neighbor sum by scale. Mean divides by $\operatorname{max}(\mathrm{number\_of\_active\_slots},1)$ , counting duplicate slots separately.
A target with no active neighbors gets a zero aggregate, but still undergoes residual normalization and its feed-forward update.
For every LN_n, compute population variance over $C$ and: $\mathrm{LN\_n}(x)=(x-\mathrm{mean\_C}(x))/\sqrt{\mathrm{mean\_C}((x-\mathrm{mean\_C}(x))^{2})+\varepsilon}\cdot \mathrm{normN\_weight}+\mathrm{normN\_bias}.$ Set $\mathrm{v1}=\mathrm{LN\_1}(h_{v}+\mathrm{aggregate})$ .
Set $f=\operatorname{GELU}(\mathrm{v1}\mathbin{@}\mathrm{ffn\_w1}+\mathrm{ffn\_b1})\mathbin{@}\mathrm{ffn\_w2}+\mathrm{ffn\_b2}$ , where ffn_w1 $[C,F]$ , ffn_b1 $[F]$ , ffn_w2 $[F,C]$ , ffn_b2 $[C]$ and $F\ge 1$ . Set $\mathrm{v2}=\mathrm{LN\_2}(\mathrm{v1}+f)$ , then zero invalid nodes.
Re-gather SOURCE and TARGET from v2, not from h_v or v1. Compute $\mathrm{e2}=\mathrm{LN\_3}(h_{e}+\mathrm{MLP\_edge}(\operatorname{CONCAT}(\mathrm{v2}[i],h_{e}[i,k],\mathrm{v2}[j])))$ . Zero every inactive edge slot after this normalization.
Return $h_{v}=\mathrm{v2}$ and $h_{e}=\mathrm{e2}$ as complete residual-updated states, not just messages or deltas. Supplied node/edge channels are both $C$ ; dropout, graph construction and amino-acid decoding are outside this block.

Why this matters

Node messages combine the target, connecting edge and source in a fixed order; the later edge update sees the completed node residual stack.

Starter code

You write the body of this function in the BioTorch workspace. Your code runs in an isolated PyTorch sandbox; Run checks the public examples and Submit grades the full test suite, including gradients.

import torch


def node_edge_message_passing(
    activations: dict,
    parameters: dict,
    masks: dict,
    structure: dict,
    eps: float = 1e-05,
    aggregation: str = 'fixed',
    scale: float = 30.0,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Return tuple[torch.Tensor, torch.Tensor]:
    
    h_v: (B, L, C).
    h_e: (B, L, K, C)."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement node_edge_message_passing")

Ready to implement it?

3 progressive hints, a solution walkthrough, and grading that checks values, shapes and gradients are waiting in the workspace.

Write and run your solution Sign in

Free account with email, Google or GitHub. Your code runs in an isolated PyTorch sandbox.

Sources and further reading

A primitive exercise, not an end-to-end reproduction of the named models. No affiliation or endorsement is implied.