Structure confidence heads

pLDDT Confidence Summary

EasyRaw PyTorchAbout 10 min3 public examples plus private tests

AlphaFold2 AlphaFold3 ESMFold ESMFold2 ESM3 RFdiffusion

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

Your task

Convert supplied per-atom pLDDT logits into expected confidence scores and average them across valid atoms and tokens.

Function signature

plddt_confidence_summary(
    plddt_logits: torch.Tensor,
    activations: dict,
    masks: dict,
    plddt_scale: float = 100.0,
) -> dict[str, torch.Tensor]

Worked example

Average atoms, even when tokens have different counts

Use two tokens with two atom slots each and centers [0.25,0.75], plddt_scale=100. Both tokens are valid; their atom masks are [True,True] and [True,False]. Supply logits [0,log(3)] and [log(3),0] for token 0, and [0,log(3)] for token 1’s valid atom (the invalid atom’s logits may be [0,0]). The valid atom probabilities are [1/4,3/4], [3/4,1/4], and [1/4,3/4], giving pLDDT scores 62.5,37.5,62.5. The per-token means are [50,62.5]; the structure mean is (62.5+37.5+62.5)/3≈54.166667, not the mean of those two token means. The invalid atom score is zero.

Parameters

Parameter	Shape and meaning
`plddt_logits`torch.Tensor	$[B,N,A,L]$ Supplied local-confidence logits, with bins on the final axis.
`activations`dict	Supplied intermediate tensors.
`masks`dict	Supplied masks. Each field below defines what True includes or permits.
`plddt_scale`floatDefault `100.0`	Scale applied to the expected local-confidence bin center.

Return

dict[str, torch.Tensor]

Return exactly the named fields applicable to this operation.

Required behavior

plddt_logits is $[B,N,A,L]$ , and plddt_centers is (L). An atom is included only when both its atom_mask and token_mask are True. Ignore masked payloads before softmax.
Apply stable softmax over bins L, take the probability-weighted bin-center mean, and multiply by plddt_scale (default 100). Invalid atom scores are exact zero.
plddt_per_token is the exact mean over valid atoms within each token. mean_plddt is the mean over ALL valid atoms in the structure, not the mean of per-token means. Empty tokens or structures return zero.

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_confidence_summary(
    plddt_logits: torch.Tensor,  # (B,N,A,L)
    activations: dict,
    masks: dict,
    plddt_scale: float = 100.0,
) -> dict[str, torch.Tensor]:
    """Return dict[str, torch.Tensor]:
    
    plddt: (B, N, A).
    plddt_per_token: (B, N).
    mean_plddt: (B,)."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement plddt_confidence_summary")

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.