Structure confidence heads

PAE GELU 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

For each ordered pair, score the possible aligned error of point j in frame i. Apply the supplied hidden affine map, exact GELU, hidden LayerNorm, any requested PAE-specific normalization, and the bin projection. Return directed bin logits.

Function signature

pae_gelu_head(
    activations: dict,
    parameters: dict,
    masks: dict,
    norm_eps: float = 1e-05,
    normalize_pairs: bool = False,
) -> torch.Tensor

Worked example

GELU precedes hidden normalization

For one valid directed pair at the hidden stage, suppose the affine map produces [0,1]. Exact GELU gives $[0,g]$ , where

g=\frac{1+\operatorname{erf}(1/\sqrt2)}{2}\approx0.841344746.

With hidden scale [1,1], shift [0,0], and norm_eps=3, population mean is $g/2$ and variance is $g^2/4$ . The normalized vector is

\frac{[-g/2,\;g/2]}{\sqrt{3+g^2/4}}\approx[-0.236014,\;0.236014].

With normalize_pairs=False, an identity bin projection and zero output bias return these two logits. The reverse directed pair is processed independently.

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.
`normalize_pairs`boolDefault `False`	Normalize pair-confidence features when True.

Return

torch.Tensor · $[B, N, N, E]$

Return Directed aligned-error-bin logits.

Required behavior

pair is $[B,N,N,C_p]$ ; a pair is valid only when both token_mask endpoints are True. Sanitize invalid pair payloads before the head.
Project with W_pair_hidden/b_pair_hidden, apply exact erf GELU, then affine population LayerNorm with pair_hidden_scale/pair_hidden_offset.
normalize_pairs (default False) optionally applies affine population LayerNorm with pae_scale/pae_offset before the final projection.
Apply W_pae and b_pae across the feature axis, returning one logit per supplied bin.
Keep PAE directed: logits(i,j) and logits(j,i) need not agree.
Zero all logits for invalid endpoint pairs after the final bias. Every used LayerNorm places norm_eps (default 1e-5) inside the square root.
Use exact $GELU(x)=x[1+erf(x/\sqrt2)]/2$ , not its tanh approximation.

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