Structure confidence heads

Distogram GELU Head

MediumRaw PyTorchAbout 20 min1 public example plus private tests

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. Apply the supplied affine hidden projection, exact erf GELU and hidden LayerNorm before the bin projection. Combine opposite pair directions using the requested symmetry, then zero pairs with absent endpoints.

Function signature

distogram_gelu_head(
    activations: dict,
    parameters: dict,
    masks: dict,
    distogram_symmetry: str = 'sum',
    norm_eps: float = 1e-05,
) -> torch.Tensor

Worked example

GELU and normalization precede logit symmetry

For the supplied identity hidden/bin projections and zero biases, take valid pair features x[0,0]=[1,0], x[0,1]=[0,1], and x[1,0]=[2,1]. Apply exact GELU, then population LayerNorm with scale one, shift zero and norm_eps=3. Writing G for GELU, define q=(G(1)/2)/sqrt(3+(G(1)/2)²)≈0.236014 and r=((G(2)-G(1))/2)/sqrt(3+((G(2)-G(1))/2)²)≈0.305933. The three projected vectors are [q,-q], [-q,q], and [r,-r]. Sum symmetry doubles the diagonal to [2q,-2q]≈[0.472028,-0.472028] and gives both off-diagonals [r-q,q-r]≈[0.069919,-0.069919]. Mean symmetry halves those results. These remain logits; no softmax is applied.

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.
`norm_eps`floatDefault `1e-05`	Positive value inside the square root of every configured LayerNorm variance.

Return

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

Return Symmetrized distance-bin logits.

Required behavior

pair is $[B,N,N,C]$ ; a pair is valid only when both token_mask endpoints are True. Sanitize invalid pair payloads before the head.
Project with W_pair_hidden/b_pair_hidden, apply exact erf GELU, then affine population LayerNorm with pair_hidden_scale/pair_hidden_offset.
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. Every used LayerNorm places norm_eps (default 1e-5) inside the square root.
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

This educational distance classifier combines a GELU regression head with a symmetric pair readout. Pair features are supplied; the task covers neither feature construction nor probability calibration.

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_gelu_head(
    activations: dict,
    parameters: dict,
    masks: dict,
    distogram_symmetry: str = 'sum',
    norm_eps: float = 1e-05,
) -> 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_gelu_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.