Structure confidence heads

GELU pLDDT Head

MediumRaw PyTorchAbout 20 min1 public example plus private tests

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

Your task

Produce one unnormalized score per pLDDT local-confidence bin for each atom slot of each token. Higher logits favor their corresponding bins; this head returns those scores. Apply the first affine projection, exact erf GELU, hidden-channel LayerNorm, then the atom/bin projection.

Function signature

plddt_gelu_head(
    activations: dict,
    parameters: dict,
    masks: dict,
    norm_eps: float = 1e-05,
) -> torch.Tensor

Worked example

For the first vector [1,3], the identity hidden projection followed by exact GELU gives approximately [0.841345,2.995950]. Subtract their mean to get [-a,a], where a≈1.077303. LayerNorm with norm_eps=3 returns [-a,a]/sqrt(a²+3)≈[-0.528154,0.528154]; the identity output projection preserves these bin logits. Normalizing before GELU changes the answer.

import torch

activations = {
    'single': torch.tensor([[[1.0, 3.0], [2.0, 0.0]]], dtype=torch.float64, requires_grad=True),
}
parameters = {
    'single_scale': torch.tensor([1.0, 1.0], dtype=torch.float64, requires_grad=True),
    'single_offset': torch.tensor([0.0, 0.0], dtype=torch.float64, requires_grad=True),
    'W1': torch.tensor([[1.0, 0.0], [0.0, 1.0]], dtype=torch.float64, requires_grad=True),
    'b1': torch.tensor([0.0, 0.0], dtype=torch.float64, requires_grad=True),
    'W_plddt': torch.tensor([[[1.0, 0.0]], [[0.0, 1.0]]], dtype=torch.float64, requires_grad=True),
    'b_plddt': torch.tensor([[0.0, 0.0]], dtype=torch.float64, requires_grad=True),
}
masks = {
    'token_mask': torch.tensor([[True, True]], dtype=torch.bool),
    'atom_mask': torch.tensor([[[True], [True]]], dtype=torch.bool),
}
norm_eps = 3.0

result = plddt_gelu_head(activations=activations, parameters=parameters, masks=masks, norm_eps=norm_eps)

Expected return

torch.tensor([[[[-0.5281542455415454, 0.5281542455415456]], [[0.49139578443196674, -0.49139578443196674]]]], dtype=torch.float64)

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.
`norm_eps`floatDefault `1e-05`	Positive value inside the square root of every configured LayerNorm variance.

Return

torch.Tensor · $[B, N, A, L]$

Return Predicted logits in atom-slot/bin order.

Required behavior

single is $[B,N,Cs]$ ; token_mask is $[B,N]$ , atom_mask is $[B,N,A]$ , and True means present. Sanitize invalid single-token payloads before the learned head.
Apply W1/b1 to single, then exact erf GELU, then affine population LayerNorm over the hidden width. For normalized heads use eps=norm_eps (default 1e-5) inside the population-variance square root.
Project hidden features with $W_{plddt}:[H,A,L]$ and add $b_{plddt}:[A,L]$ . Return $[B,N,A,L]$ in bin-logit order, without softmax or pLDDT scaling.
Zero every output logit whose token is absent or atom slot is absent, including after adding biases. Supply all parameters; do not initialize new layers.
Use CPU float32/float64 tensors with one shared floating dtype/device. Preserve strided layouts and do not mutate supplied tensors.
Use raw tensor operations, not a ready-made implementation of the requested primitive. Supplied parameters and random variates must not be reinitialized.
Exact GELU is $x[1+\operatorname{erf}(x/\sqrt2)]/2$ ; use it after W1/b1 and before hidden-channel LayerNorm.

Why this matters

Different heads answer different questions: pLDDT describes local accuracy, a distogram describes distance, and PAE describes error after alignment in a particular residue’s frame. Their probabilities feed separate summaries; they are not interchangeable confidence scales.

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 plddt_gelu_head(
    activations: dict,
    parameters: dict,
    masks: dict,
    norm_eps: float = 1e-05,
) -> torch.Tensor:
    """Return a tensor with shape (B, N, A, L)."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement plddt_gelu_head")

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.

Related exercises

Sources and further reading

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