Module confidence_head 

Confidence head: predicts pLDDT, pAE, pDE, and distogram.

Architecture:

1. 4-layer Pairformer trunk refines (single, pair)
2. Four linear heads project to logit bins:
   • pLDDT: single → 50 bins → softmax → weighted mean → scalar [0,1]
   • pAE: pair → 64 bins → softmax → weighted mean → matrix [Å]
   • pDE: pair → 64 bins (predicted distance error)
   • distogram: pair → 39 bins

Weight layout (rooted at confidence_head):

trunk.blocks.{0..3}.*   — 4-layer Pairformer (8 heads, d_pair=256)
plddt_head.weight        — linear d_single → num_plddt_bins (no bias)
pae_head.weight          — linear d_pair   → num_pae_bins   (no bias)
pde_head.weight          — linear d_pair   → num_pde_bins   (no bias)
distogram_head.weight    — linear d_pair   → distogram_bins (no bias)

Structs

ConfidenceHead

4-layer Pairformer confidence head.

ConfidenceOutput

Output of the confidence head.

Functions

bins_to_scalar

Convert bin logits to a scalar via softmax + weighted average of bin centres.

plddt_from_logits

Convert pLDDT logits [B, N, 50] → scalar pLDDT [B, N] in [0, 1].