Structure confidence heads

PAE Linear Head

EasyRaw PyTorchAbout 10 min3 public examples plus private tests

AlphaFold2OpenFoldAlphaFold3 ESMFold ESMFold2

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

Your task

For each ordered pair, produce unnormalized scores for the error of point j when evaluated in alignment frame i. Project the pair features into PAE bins, preserving the row/column direction; the two opposite pairs can have different logits.

Function signature

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

Worked example

Preserve directional pair logits

The final weight is identity and the bias is zero, so each pair vector is unchanged. Pair(0,1) gives[0,1], while pair(1,0) gives[2,1]. Do not average or sum the two directions.

import torch

activations = {
    'pair': torch.tensor([[[[1.0, 0.0], [0.0, 1.0]], [[2.0, 1.0], [1.0, 2.0]]]], dtype=torch.float64, requires_grad=True),
}
parameters = {
    'W_pae': torch.tensor([[1.0, 0.0], [0.0, 1.0]], dtype=torch.float64, requires_grad=True),
    'b_pae': torch.tensor([0.0, 0.0], dtype=torch.float64, requires_grad=True),
    'pae_scale': torch.tensor([1.0, 1.0], dtype=torch.float64, requires_grad=True),
    'pae_offset': torch.tensor([0.0, 0.0], dtype=torch.float64, requires_grad=True),
}
masks = {
    'token_mask': torch.tensor([[True, True]], dtype=torch.bool),
}
norm_eps = 3.0
normalize_pairs = False

result = pae_linear_head(activations=activations, parameters=parameters, masks=masks, norm_eps=norm_eps, normalize_pairs=normalize_pairs)

Expected return

torch.tensor([[[[1.0, 0.0], [0.0, 1.0]], [[2.0, 1.0], [1.0, 2.0]]]], 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.
`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]$ ; a pair is valid only when both token_mask endpoints are True. Sanitize invalid pair payloads before the head.
Use pair features directly; there is no hidden projection or activation.
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 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.

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_linear_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_linear_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.