Structure confidence heads

Normalized Resolved-Atom Head

EasyRaw PyTorchAbout 10 min1 public example plus private tests

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

Your task

Predict two logits per atom slot: unresolved and resolved. Return logits only; probability conversion is a separate step.

Function signature

resolved_atom_head_normalized(
    activations: dict,
    parameters: dict,
    masks: dict,
    norm_eps: float = 1e-05,
) -> torch.Tensor

Worked example

With norm_eps=3, LayerNorm maps the single vectors to [-0.5, 0.5] and [0.5, -0.5]; the identity weight returns those unresolved/resolved logits.

import torch

activations = {
    'single': torch.tensor([[[1.0, 3.0], [2.0, 0.0]]], dtype=torch.float64, requires_grad=True),
}
parameters = {
    'W_resolved': torch.tensor([[[1.0, 0.0]], [[0.0, 1.0]]], dtype=torch.float64, requires_grad=True),
    'b_resolved': torch.tensor([[0.0, 0.0]], dtype=torch.float64, requires_grad=True),
    'resolved_scale': torch.tensor([1.0, 1.0], dtype=torch.float64, requires_grad=True),
    'resolved_offset': torch.tensor([0.0, 0.0], dtype=torch.float64, requires_grad=True),
}
masks = {
    'token_mask': torch.tensor([[True, True]], dtype=torch.bool),
    'atom_mask': torch.tensor([[[True], [True]]], dtype=torch.bool),
}
norm_eps = 3.0

result = resolved_atom_head_normalized(activations=activations, parameters=parameters, masks=masks, norm_eps=norm_eps)

Expected return

torch.tensor([[[[-0.5, 0.5]], [[0.5, -0.5]]]], dtype=torch.float64)

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

Return

torch.Tensor · $[B, N, A, 2]$

Return Unresolved/resolved atom logits.

Required behavior

single is $[B,N,C]$ ; token_mask and atom_mask mark included token/atom slots. Sanitize absent token features before the learned head.
Apply affine population LayerNorm using resolved_scale/resolved_offset and norm_eps (default 1e-5) inside the square root.
Contract W_resolved $[C,A,2]$ with the single feature axis and add b_resolved $[A,2]$ . The final class order is unresolved, resolved. Return shape $[B,N,A,2]$ , with absent token/atom logits exactly zero after bias.
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 head predicts whether an atom is represented as resolved in the reference structure. It returns two unnormalized scores per atom: unresolved first, resolved second.

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 resolved_atom_head_normalized(
    activations: dict,
    parameters: dict,
    masks: dict,
    norm_eps: float = 1e-05,
) -> torch.Tensor:
    """Return a tensor with shape (B, N, A, 2)."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement resolved_atom_head_normalized")

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.