Attention-derived protein contact prediction

Your task

Function signature

attention_contact_prediction(
    attention: torch.Tensor,
    tokens: torch.Tensor,
    token_mask: torch.Tensor,
    weight: torch.Tensor,
    bias: torch.Tensor,
    special_ids: list = [0, 1, 2],
) -> dict[str, torch.Tensor]

Worked example

Symmetrize and remove row/column background

Parameters

Return

dict[str, torch.Tensor]

Required behavior

• Inputs attention [B,A,H,T,T] is finite nonnegative floating attention (B,A,H,T\ge 1); token_mask [B,T] is bool with True means present; tokens [B,T] is integer. weight [A\cdot H] and scalar bias [] share attention dtype/device. Config special_ids is a list of integer IDs, default [0,1,2]; no other keys.
• A residue is valid exactly when token_mask is True and its token ID is not in special_ids. Mask BOTH attention endpoints before any reduction. Keep the original T axes; do not remove diagonal/self contacts.
• For each layer/head plane, let M_{\mathrm{ij}}=\mathrm{attention\_ij} on valid pairs, zero elsewhere; S=M+M^{\top} (a SUM, not a mean). Set r_{i}=\sum_{j} S_{\mathrm{ij}}, c_{j}=\sum_{i} S_{\mathrm{ij}}, Z=\sum_{\mathrm{ij}} S_{\mathrm{ij}}.
• If Z>0, C_{\mathrm{ij}}=S_{\mathrm{ij}}-r_{i}\cdot c_{j}/Z; if Z=0 define C=0 with zero attention derivative for that whole plane. No \epsilon is added to a positive APC denominator. Flatten A,H in layer-major, head-minor order to features [B,T,T,A\cdot H].
• Let z_{\mathrm{ij}}=\sum_{c} \mathrm{features\_ijc}\cdot \mathrm{weight\_c}+\mathrm{bias}; probabilities are \operatorname{sigmoid}(z). Return features, logits and probabilities with invalid pairs exactly zero, plus pair_mask [B,T,T]. Empty valid-residue sets yield finite zeros even with nonzero bias.
• The positive-mass branch is differentiable; the zero-mass extension is explicitly defined, not source 0/0 behavior. Preserve \mathrm{float32}/64 and noncontiguous inputs; do not mutate inputs.

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 attention_contact_prediction(
    attention: torch.Tensor,  # (B, A, H, T, T)
    tokens: torch.Tensor,  # (B, T)
    token_mask: torch.Tensor,  # (B, T)
    weight: torch.Tensor,  # (A* H)
    bias: torch.Tensor,  # Scalar tensor, shape ().
    special_ids: list = [0, 1, 2],
) -> dict[str, torch.Tensor]:
    """Return dict[str, torch.Tensor]:
    
    features: (B, T, T, A* H).
    logits: (B, T, T).
    probabilities: (B, T, T).
    pair_mask: (B, T, T)."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement attention_contact_prediction")

Sources and further reading

• www.science.org
• facebookresearch/esm · modules.py