Structure module, templates and recycling

Template Pair Cross-Attention

MediumRaw PyTorchAbout 20 min2 public examples plus private tests

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

Your task

For each residue pair, use its query to choose how much to read from each valid template. Normalize query–template scores across the template axis, mix the projected values per head, and return the projected pair update together with those template weights.

Function signature

template_pair_cross_attention(
    activations: dict,
    parameters: dict,
    masks: dict,
) -> tuple[torch.Tensor, torch.Tensor]

Worked example

A query favors one template

Consider one valid residue pair, one head and D=V=1 in attention mode, after the supplied projections. Let the projected query be 1, template keys [0,log(3)], and template values [2,4]; both templates are valid. Scores are [0,log(3)], giving weights [1/4,3/4] over templates. The weighted value is (1/4)·2 + (3/4)·4 = 3.5. With output weight 1 and bias 1, the pair update is 4.5. These are update features; no input residual is added.

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.

Return

tuple[torch.Tensor, torch.Tensor]

Return these tensors in the stated order.

Required behavior

template is $[B,T,N,N,Ct]$ ; combine template_mask with pair_mask broadcast across T. True means included. Invalid template payloads must not influence projections or sums.
Project query into $[B,N,N,H,D]$ , templates into keys $[B,T,N,N,H,D]$ and values $[B,T,N,N,H,V]$ . Compute scaled dot products $q\cdot k/\sqrt D$ and a stable masked softmax over T only. Return weights $[B,N,N,H,T]$ .
Aggregate values with those weights, flatten the head and value axes in head-major order, then apply w_out and b_out. There is no gate, LayerNorm, or residual.
After the final bias, zero every pair with no valid template. Template count T is positive; all-masked pairs are supported.
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

For each residue pair, its query representation attends over the available templates. Each head can choose different template weights.

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 template_pair_cross_attention(
    activations: dict,
    parameters: dict,
    masks: dict,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Return tuple[torch.Tensor, torch.Tensor]:
    
    update: (B,N,N,C_q).
    weights: (B,N,N,H,T)."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement template_pair_cross_attention")

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.