Structure confidence heads

Predicted TM Score

MediumRaw PyTorchAbout 20 min3 public examples plus private tests

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

Your task

Convert each ordered pair’s aligned-error probabilities into expected TM agreement, average those agreements for each candidate alignment row, then select a row using the stated importance weights. Return its original row score as pTM, all row scores, and whether any allowed weighted pair exists.

Function signature

predicted_tm_score(
    pae_probability: torch.Tensor,
    activations: dict,
    masks: dict,
    tm_length: str = 'mask',
    count_eps: float = 1e-08,
    tm_weighted_selection: bool = True,
) -> dict[str, torch.Tensor]

Worked example

Row selection can differ from maximum row score

Take two valid tokens with weights=[1,0.5], all pairs allowed, tm_length="mask", tm_weighted_selection=True, and count_eps=0.01. Use error centers $[0,d_0]$ , where

d_0=1.24\sqrt[3]{19-15}-1.8.

Their TM kernels are [1,1/2]. Let both pairs in row 0 have probabilities [1/4,3/4], giving pair term 0.625, and both in row 1 have [3/4,1/4], giving 0.875. Endpoint-product weights give row scores

\mathrm{row}_0=\frac{0.9375}{1.51}\approx0.620861,\qquad \mathrm{row}_1=\frac{0.65625}{0.76}\approx0.863487.

Weighted selection compares $0.620861$ with $0.5\cdot0.863487\approx0.431743$ , so it chooses row 0. Return ptm≈0.620861, with ptm_valid=True.

Parameters

Parameter	Shape and meaning
`pae_probability`torch.Tensor	$[B,N,N,L]$ Categorical aligned-error probabilities; axis i selects the alignment frame and axis j the evaluated token.
`activations`dict	Supplied intermediate tensors.
`masks`dict	Supplied masks. Each field below defines what True includes or permits.
`tm_length`strDefault `'mask'`	"mask" uses the valid-token count; "weight_floor" uses the sum of valid token weights rounded down for the TM length scale.
`count_eps`floatDefault `1e-08`	Positive offset in the declared TM/confidence weighted-count denominators, added after summing the relevant support.
`tm_weighted_selection`boolDefault `True`	Use residue weights when selecting the best aligned TM row.

Return

dict[str, torch.Tensor]

Return exactly the named fields applicable to this operation.

Required behavior

pae_probability is $[B,N,N,L]$ , with nonnegative bins summing to 1 on valid token pairs; invalid pair values are ignored. pae_centers supplies L nonnegative aligned-error bin centers. weights is a nonnegative importance weight per token; zero invalid token weights.
For a length $n$ , use $d_0=1.24(\max(n,19)-15)^{1/3}-1.8$ and $K_l=d_0^2/(d_0^2+center_l^2)$ . A pair term is the expected K under that pair’s PAE probabilities.
tm_length="mask"(default) uses valid-token count; "weight_floor" uses floor(sum of valid weights). tm_pair_mask additionally restricts scored ordered pairs.
For row i, average pair terms with weights $w_iw_j$ over allowed j, dividing by their sum plus count_eps (default $10^{-8}$ ). Empty rows score zero. tm_weighted_selection=True (default) selects the row maximizing row_score × w_i; False selects the maximum row_score. Return the UNWEIGHTED score of the selected row; ties choose the first row.
The validity flag is True iff at least one allowed pair has positive weights at both endpoints. No support returns zero scores and False, without inventing a positive denominator or fallback.
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

Predicted TM scoring converts aligned-error probabilities into expected structural similarity, then chooses the best candidate alignment frame. It averages over supported point columns within each row rather than taking a single global average.

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 predicted_tm_score(
    pae_probability: torch.Tensor,  # (B,N,N,L)
    activations: dict,
    masks: dict,
    tm_length: str = 'mask',
    count_eps: float = 1e-08,
    tm_weighted_selection: bool = True,
) -> dict[str, torch.Tensor]:
    """Return dict[str, torch.Tensor]:
    
    ptm_per_row: (B, N).
    ptm: (B,).
    ptm_valid: (B,)."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement predicted_tm_score")

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.