Structure confidence heads

Distogram Linear Head

EasyRaw PyTorchAbout 10 min3 public examples plus private tests

AlphaFold2 AlphaFold3 ESMFold ESMFold2 RFdiffusion

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

Your task

Produce one unnormalized score per distance bin for each residue pair. Use pair features directly for the supplied bin projection. Combine opposite pair directions using the requested symmetry, then zero pairs with absent endpoints.

Function signature

distogram_linear_head(
    activations: dict,
    parameters: dict,
    masks: dict,
    distogram_symmetry: str = 'sum',
) -> torch.Tensor

Worked example

Combine the two directed logit vectors

At the symmetry stage, suppose two valid residues have already projected logits Z[0,1]=[1,3], Z[1,0]=[2,5], and Z[0,0]=[4,6], including their projection biases. Sum symmetry returns [3,8] in both off-diagonal positions and [8,12] on (0,0). Mean symmetry instead gives [1.5,4] and [4,6]. These operations combine logits before any probability conversion; they do not average already-softmaxed probabilities.

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.
`distogram_symmetry`strDefault `'sum'`	"sum" adds the transposed pair logits; "mean" divides that symmetric sum by two.

Return

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

Return Symmetrized distance-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.
Use pair features directly; there is no hidden projection or activation.
There is no extra pair normalization before the distogram projection.
Apply W_distogram and b_distogram across the feature axis, returning one logit per supplied bin.
distogram_symmetry defaults to "sum": add logits(i,j)+logits(j,i), including doubling diagonal values. "mean" divides that sum by 2.
Zero all logits for invalid endpoint pairs after the final bias/symmetrization.

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 distogram_linear_head(
    activations: dict,
    parameters: dict,
    masks: dict,
    distogram_symmetry: str = 'sum',
) -> torch.Tensor:
    """Return a tensor with shape (B, N, N, D)."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement distogram_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.