Learned structural confidence heads and model-aware ranking

Your task

Function signature

structure_confidence_readouts(
    activations: dict,
    parameters: dict,
    masks: dict,
    structure: dict,
    options: dict,
) -> dict[str, torch.Tensor]

Worked example

From logits to confidence and TM agreement

Parameters

Return

dict[str, torch.Tensor]

Required behavior

• Build learned confidence readouts from single and pair representations. Use B\ge 1 batches, N\ge 1 tokens, A\ge 1 dense atom slots, \mathrm{Cs},\mathrm{Cp} \ge 1 channels. All float inputs share \mathrm{float32}/64 and are finite except masked single/pair payloads, which must be ignored before arithmetic.
• True means present. Effective atom mask is atom_mask [b,n,a] AND token_mask [b,n]; pair mask is m[b,i,j]=\mathrm{token\_mask}[b,i] AND token_mask [b,j]. All absent logits, probabilities, expectations and summary contributions are zero. Biases must not leak through padding. Dimensions need not match checkpoint sizes.
• 1. \operatorname{LN}(x;g,b)=(x-\operatorname{mean}(x))\cdot \mathrm{rsqrt}(\operatorname{mean}((x-\operatorname{mean}(x))^{2})+\mathrm{norm\_eps})\cdot g+b, reducing channels.
• Config plddt_head defaults to 'relu_mlp': apply h=\mathrm{relu}(\operatorname{LN}(\mathrm{single})\mathbin{@}\mathrm{W1}+\mathrm{b1}), h=\mathrm{relu}(h\mathbin{@}\mathrm{W2}+\mathrm{b2}). 'linear_stack' applies those two affine maps WITHOUT either ReLU. 'direct' uses \operatorname{LN}(\mathrm{single}) itself and does not receive W1,b1,W2,b2. 'plain_linear' uses raw single without normalization or hidden layers; it has no single_scale/offset and expresses RFdiffusion's LDDTNetwork. 'gelu_regression' instead uses h=\operatorname{LN}(\operatorname{GELU}(\mathrm{single}\mathbin{@}\mathrm{W1}+\mathrm{b1})), WITHOUT initial normalization or W2,b2; its single_scale/offset have hidden width H1.
• \operatorname{GELU}(x)=\frac{x}{2}\left(1+\operatorname{erf}\frac{x}{\sqrt2}\right), not the tanh approximation. This order matters for ESM3. Then \mathrm{plddt\_logits}[b,n,a,l]=\sum_{c} h[b,n,c]\cdot W_{\mathrm{plddt}}[c,a,l]+b_{\mathrm{plddt}}[a,l].
• Dense slot weights express atom-specific heads; packed/token-to-atom conversion is a separate primitive.
• 2. Resolve probabilities with \mathrm{logits}[b,n,a,r]=\sum_{c} u[b,n,c]\cdot W_{\mathrm{resolved}}[c,a,r]+b_{\mathrm{resolved}}[a,r], r=0(\mathrm{unresolved}),1(resolved). u is raw single by default, or its independent resolved_scale/offset LN when \mathrm{normalize\_resolved}=\mathrm{true}.
• This two-class head matches AF3-style readouts; a zero \mathrm{class}-0 logit gives the sigmoid parameterization used for binary AF2 atom logits. It is not an additional \mathrm{ESMFold}-\mathrm{v1} head.
• 3. By default pair_head='linear' uses the supplied pair as z. pair_head='gelu_regression' first builds z=\operatorname{LN}(\operatorname{GELU}(\mathrm{pair}\mathbin{@}W_{\mathrm{pair\_hidden}}+b_{\mathrm{pair\_hidden}})), using pair_hidden_scale/offset.
• This covers the ESM3 pair classifier after its separate product/difference feature construction. Produce distance logits z\mathbin{@}W_{\mathrm{distogram}}+b_{\mathrm{distogram}} and symmetrize as Z_{\mathrm{ij}}+Z_{\mathrm{ji}} ('sum', default) or (Z_{\mathrm{ij}}+Z_{\mathrm{ji}})/2 ('mean').
• Project directed PAE logits with W_pae,b_pae. Project PDE with W_pde,b_pde and apply pde_symmetry='sum' (default), 'mean', or 'none'. With \mathrm{normalize\_pairs}=\mathrm{true}, apply separate pae_scale/offset and pde_scale/offset LNs to z before their respective projections; distance logits use z directly.
• PAE must remain directional. PDE predicts an error in a distance, whereas a distogram predicts the distance. Symmetry is a logit operation before softmax, not probability averaging.
• For a joint upstream classifier, its final projection columns can be partitioned into these heads; the unused direction-label outputs of ESM3's decoder are not returned here.
• 4. Compute softmax along bins using max subtraction. Return the learned logits, distogram probabilities, and categorical means against supplied centers. \mathrm{plddt}=\mathrm{plddt\_scale}\cdot \sum_{l} p_{l}\cdot \mathrm{plddt\_centers\_l} (scale default 100; use1 for a [0,1] convention); resolved_probability is class1.
• Bin supports are intentionally supplied: AF2 uses (l+0.5)/L for l=0..L-1, while original RFdiffusion reports (l+1)/L (upper-edge values), with scale1 and A=1. The per-token pLDDT averages present atom slots and mean_plddt averages all present atom slots.
• Divide counts by \operatorname{max}(\operatorname{count},1); no atoms gives zero. PAE/PDE and expected_distance use their own centers, never another head's bins.
• 5. Compute TM reductions from PAE probabilities. w=\operatorname{where}(\mathrm{token\_mask},\mathrm{weights},0), observed weights in [0,1]. L is \sum(\mathrm{token\_mask}) by default (tm_length='mask'); tm_length='weight_floor' instead uses \operatorname{floor}(\sum(w)), the AF2 reporting-length convention.
• d_0(L)=1.24\bigl(\max(L,19)-15\bigr)^{1/3}-1.8. T_{ij}=\sum_k\frac{p^{\mathrm{PAE}}_{ijk}}{1+(\texttt{pae\_centers}_k/d_0)^2}. Global pTM always uses this global-length kernel. With tm_interface_length='global' (default), iPTM uses it too.
• With 'chain_pair' (AF3), count masked-on tokens in each chain; for pair i,j use the union length of their two chains, counting a same-chain pair only once. Recompute that pair's d0 and TM expectation. These counts depend on token_mask, not fractional weights or frame-validity masks.
• Effective TM support is m AND tm_pair_mask; the latter permits missing/ill-defined alignment frames.
• For pTM use this support; for iPTM additionally require \mathrm{chain\_ids\_i}\ne \mathrm{chain\_ids\_j}. Define Q_{\mathrm{ij}}=M_{\mathrm{ij}}\cdot w_{i}\cdot w_{j} and \mathrm{row}_i=\frac{\sum_j Q_{ij}T_{ij}}{\texttt{count\_eps}+\sum_j Q_{ij}}.
• Select the FIRST maximizing \mathrm{row\_i}\cdot w_{i} by default (\mathrm{tm\_weighted\_selection}=\mathrm{true}), or row_i when false, and return the unmultiplied selected row score. Use weights1 and unweighted selection for the \mathrm{AF3}/\mathrm{ESMFold2} mask-based reduction.
• This is max-of-row-means, not a global mean or max over individual pairs. An unsupported score is 0 with its boolean \mathrm{valid}=\mathrm{false}; this replaces upstream NaN reporting with an explicit finite-value contract.
• Also return chain_pair_tm[B,C,C], with axes defined by supplied unique chain_labels [C], which must contain every observed chain ID (unused labels are allowed). Entry(c,d) uses the interface-length kernel, restricted to alignment-frame rows in chain c and point columns in d.
• chain_pair_mode='directed' (default) keeps only that direction; 'bidirectional' includes both c\to d and d\to c before reducing, giving AF3-style symmetric chain-pair scores.
• For c=d, use within-chain pairs. Return corresponding chain_pair_valid. This natural frame-row/point-column convention is the TRANSPOSE of the indices used by the pinned ESMFold2 per-chain matrix.
• 6. \mathrm{interface\_plddt}=\mathrm{true} enables the ESMFold2-style interface summary. A polymer token receives weight1 if any present OTHER-chain representative has supplied \mathrm{representative\_distances\_ij}<\mathrm{interface\_cutoff} (default 8 angstroms), otherwise zero.
• A ligand_mask token receives weight2 regardless of contact. Multiply by atom_mask and average pLDDT over weighted atom slots with denominator \mathrm{count\_eps}+\sum(\mathrm{weights}). Missing contact partners cannot create an interface. The distance-based selection is discrete; distances are supplied rather than inferred from confidence.
• Disabled interface_plddt returns zero.
• 7. ranking_mode='tm_mix' (default) returns (1-\mathrm{ranking\_mix})\cdot \mathrm{ptm}+\mathrm{ranking\_mix}\cdot \mathrm{iptm}, mixture default 0.8. ranking_mode='af3' adds the explicit AF3 reporting rule.
• For protein_mask AND token_mask residues, look up a supplied positive max_accessible_area table using residue_types; \mathrm{rASA}_i=\min\left(\frac{\texttt{accessible\_surface\_area}_i}{\texttt{max\_area}_i},1\right).
• Within each chain, collect present protein residues in original token order and smooth rASA with an odd reflected window rasa_window (default 25). Reflection excludes repeating the endpoint: for chain length M>1 an offset index k maps to r=k mod(2M-2), then r if r<M else 2M-2-r; M=1 repeats its sole value.
• Thus widths larger than a short chain are well-defined. No smoothing crosses chains or uses nonprotein/padded tokens. fraction_disordered is the fraction of present protein residues whose smoothed rASA is STRICTLY greater than rasa_cutoff (default 0.581), or 0 when none exist. Return rasa/smoothed_rasa zero outside protein sites.
• The final score is \mathrm{base}+0.5\cdot \mathrm{fraction\_disordered}-100\cdot \mathrm{has\_clash}, where base is the TM mixture when iptm_valid, otherwise pTM alone. has_clash is supplied by the separate structural-violation/clash exercise. With default 0.8 mixture this is the pinned AF3 rule; it is not a calibrated probability.
• Disabled AF3 ranking returns zero rASA/disorder diagnostics. \mathrm{norm\_eps}=10^{-5} and \mathrm{count\_eps}=10^{-8} are positive. ranking_mix and rasa_cutoff lie in [0,1]; interface_cutoff is positive; rasa_window is a positive odd integer.
• Invalid values or unknown head/symmetry/length/chain/ranking modes raise ValueError. plddt_scale is finite positive. Bin centers are finite, increasing and nonnegative; pLDDT centers lie in [0,1]. Areas are nonnegative on present protein sites; table maxima are positive; lookup IDs are valid there and ignored at missing sites.
• All dimensions/bins/hidden widths are positive; optional parameters exist only for selected variants. Empty batch/token axes raise ValueError; other shape/domain violations are outside the contract. All-zero masks, empty chain labels' support and fractional weights are supported.
• Different source bin conventions use supplied centers: \mathrm{AF2}/\mathrm{AF3} final catch-all supports versus ESMFold2 \mathrm{evenly}-\mathrm{spaced0}..32 angstrom bin centers. No no_grad shortcut: differentiate continuous calculations even where upstream reporting code detaches them.

Why this matters

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 structure_confidence_readouts(
    activations: dict,
    parameters: dict,
    masks: dict,
    structure: dict,
    options: dict,
) -> dict[str, torch.Tensor]:
    """Return dict[str, torch.Tensor]:
    
    plddt_logits: (B, N, A, L).
    resolved_logits: (B, N, A, 2).
    distogram_logits: (B, N, N, D).
    pae_logits: (B, N, N, E).
    pde_logits: (B, N, N, P).
    plddt: (B, N, A).
    plddt_per_token: (B, N).
    mean_plddt: (B,).
    resolved_probability: (B, N, A).
    distogram_probability: (B, N, N, D).
    expected_distance: (B, N, N).
    pae: (B, N, N).
    pde: (B, N, N).
    ptm_per_row: (B, N).
    iptm_per_row: (B, N).
    ptm: (B,).
    iptm: (B,).
    ptm_valid: (B,).
    iptm_valid: (B,).
    chain_pair_tm: (B, C, C).
    chain_pair_valid: (B, C, C).
    interface_plddt: (B,).
    rasa: (B, N).
    smoothed_rasa: (B, N).
    fraction_disordered: (B,).
    ranking_score: (B,)."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement structure_confidence_readouts")

Related exercises

• Builds onReLU pLDDT Head
• Builds onLinear-Stack pLDDT Head
• Builds onNormalized Linear pLDDT Head
• Builds onGELU pLDDT Head
• Builds onLinear pLDDT Head
• Builds onResolved-Atom Head
• Builds onNormalized Resolved-Atom Head
• Builds onDistogram Linear Head
• Builds onDistogram GELU Head
• Builds onPAE Linear Head
• Builds onPAE GELU Head
• Builds onPDE Linear Head
• Builds onPDE GELU Head
• Builds onpLDDT Confidence Summary
• Builds onPredicted TM Score
• Builds onInterface TM Score
• Builds onDirected Chain-Pair TM Scores
• Builds onBidirectional Chain-Pair TM Scores
• Builds onInterface pLDDT Summary
• Builds onRelative Accessible Surface Area
• Builds onTM Ranking Score
• Builds onAlphaFold 3 Ranking Score

Sources and further reading

• www.nature.com
• www.nature.com
• google-deepmind/alphafold · modules.py
• google-deepmind/alphafold · confidence.py
• facebookresearch/esm · esmfold.py
• facebookresearch/esm · esmfold.py
• google-deepmind/alphafold3 · confidence_head.py
• google-deepmind/alphafold3 · confidences.py
• google-deepmind/alphafold3 · confidences.py
• google-deepmind/alphafold3 · confidences.py
• huggingface/transformers · modeling_esmfold2.py
• Biohub/esm · vqvae.py
• Biohub/esm · vqvae.py
• RosettaCommons/RFdiffusion · AuxiliaryPredictor.py
• RosettaCommons/RFdiffusion · RoseTTAFoldModel.py