Protein sequence and inverse folding

Sequence-conditioned ProteinMPNN decoder stack

HardRaw PyTorchAbout 30 min1 public example plus private tests

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

Your task

Update each residue’s feature vector using its structural neighbors and only the amino-acid identities allowed by the supplied decoding-order mask. A visible neighbor contributes its current decoder state and sequence embedding; an unseen neighbor contributes only its original encoder state. Return the final features and every intermediate layer state.

Function signature

proteinmpnn_decoder(
    activations: dict,
    parameters: dict,
    masks: dict,
    structure: dict,
    eps: float = 1e-05,
    scale: float = 30.0,
    dropout: float = 0.0,
) -> tuple[torch.Tensor, torch.Tensor]

Worked example

What one neighbor is allowed to reveal

Consider the message-input stage with one feature channel. Let the current query state be H_d[i]=2, its edge feature be 1, and neighbor j have sequence embedding 7, original encoder state H_0[j]=5, and current decoder state H_d[j]=3. If the edge is active and backward_mask=True, concatenate [2,1,7,3]. If the edge is active but backward_mask=False, concatenate [2,1,0,5]. The latter retains structural information while withholding the unseen amino-acid embedding and decoder state. These vectors then enter the same supplied message MLP.

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 stabilizer inside sqrt(population channel variance + eps) in both LayerNorm operations.
`scale`floatDefault `30.0`	Positive divisor of the summed active neighbor messages at every layer; it is independent of the neighbor count.
`dropout`floatDefault `0.0`	Probability used with supplied keep masks. Do not draw new random masks.

Return

tuple[torch.Tensor, torch.Tensor]

Return these tensors in the stated order.

Required behavior

$B,L,C,D,F \ge 1$ , $K\ge 0$ , $V\ge 2$ . Start $\mathrm{H0}=\mathrm{encoder\_nodes}$ masked to zero on invalid nodes; keep it as the immutable structural context. $E=\mathrm{encoder\_edges}$ [ $B,L,K,C$ ]; $S=\mathrm{embedding}[\mathrm{tokens}]$ has shape [ $B,L,C$ ]. All floating inputs are finite and share dtype/device.
weights have $[\mathrm{input},\mathrm{output}]$ orientation and a leading $D$ layer axis. Effective edge validity $a_{\mathrm{ijk}}=\mathrm{edge\_mask}$ AND both endpoint node masks. backward_mask is supplied from a valid strict decoding order (therefore no self/future edges); it is intersected with a, and is not recomputed here.
At layer $d$ , for neighbor $j$ of $i$ , let $s_{\mathrm{ij}}=S_{j}$ and $n_{\mathrm{ij}}=H_{d}$ , $j$ on active backward edges; otherwise $s_{\mathrm{ij}}=0$ and $n_{\mathrm{ij}}=\mathrm{H0\_j}$ . The message input is CONCAT $[H_{d},i, E_{\mathrm{ij}}, s_{\mathrm{ij}}, n_{\mathrm{ij}}]$ , exactly that $4 C$ order.
Apply affine $\mathrm{message\_w1}/\mathrm{b1} \to \mathrm{exact} \operatorname{GELU} \to \mathrm{affine} \mathrm{message\_w2}/\mathrm{b2} \to \mathrm{exact} \operatorname{GELU} \to \mathrm{affine} \mathrm{message\_w3}/\mathrm{b3}$ . Sum only active messages over $K$ and divide by $\mathrm{scale}>0$ (default $30$ ), not by degree.
Then $U=\mathrm{LayerNorm}(H_{d} + \mathrm{supplied\_dropout}(\mathrm{message\_update}), \mathrm{norm1\_scale}/\mathrm{bias})$ .
Let $f=\operatorname{GELU}(U\mathbin{@}\mathrm{ffn\_up}+\mathrm{ffn\_up\_bias})\mathbin{@}\mathrm{ffn\_down}+\mathrm{ffn\_down\_bias}$ ; $H_{\mathrm{}}(d+1)=\mathrm{LayerNorm}(U+\mathrm{supplied\_dropout}(f),\mathrm{norm2\_scale}/\mathrm{bias})$ , zeroing invalid nodes afterwards. Exact GELU is $x\cdot (1+\operatorname{erf}(x/\sqrt{2}))/2$ .
LayerNorm uses population channel variance with $\mathrm{eps}>0$ inside the square root (default $10^{-5}$ ). Config $\mathrm{dropout}=p$ in $[0,1)$ , default $0$ .
The two supplied keep masks [ $D,B,L,C$ ] define $\mathrm{dropout}(x)=\operatorname{where}(\mathrm{keep},x/(1-p),0)$ ; they are applied even if $p=0$ , so source inference corresponds to $p=0$ and all keeps True. These are the only keys.
Return $\mathrm{nodes}=H_{D}$ and $\mathrm{states}=\mathrm{stack}(\mathrm{H0},\ldots,H_{D})$ with shape [ $D+1$ , $B,L,C$ ]. No output projection, softmax, sampling, edge updates or residual outside the specified stack. $K=0$ means zero message aggregation, but the residual/norm/FFN path still executes.
$F=4C$ gives the upstream transition width; other positive $F$ is a stated teaching extension. Noncontiguous states/weights and padding must work without mutation.
The message concatenation is $[H_d[i],E_{ij},s_{ij},n_{ij}]$ . A visible neighbor uses $n_{ij}=H_d[j]$ ; an unseen neighbor uses the unchanged $H_0[j]$ .

Why this matters

Known neighbors can supply their sequence embedding and current decoder state. Future neighbors supply only the original encoder state. This preserves structural context while enforcing the supplied decoding order.

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 proteinmpnn_decoder(
    activations: dict,
    parameters: dict,
    masks: dict,
    structure: dict,
    eps: float = 1e-05,
    scale: float = 30.0,
    dropout: float = 0.0,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Return tuple[torch.Tensor, torch.Tensor]:
    
    nodes: (B, L, C).
    states: (D+1, B, L, C)."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement proteinmpnn_decoder")

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.