Geometric direction-and-distance attention

Your task

Function signature

geometric_reasoning_attention(
    activations: dict,
    parameters: dict,
    masks: dict,
    structure: dict,
) -> dict[str, torch.Tensor]

Worked example

Rotate values back into each query frame

Parameters

Return

dict[str, torch.Tensor]

Required behavior

• Implement the ESM3 geometric-attention kernel from supplied local projections. rotation[B,N,3,3] maps column local vectors to the global frame and is a proper orthogonal matrix; translation[B,N,3] is its global origin.
• q_direction,k_direction,q_point,k_point each have [B,N,H,3]; values has [B,N,H,M,3]. B,N,H,M \ge 1. Learned direction_scale [H] and distance_scale [H] are unconstrained finite logits, not already positive scales.
• output_weight [H\cdot M\cdot 3,\mathrm{Cout}] and optional output_bias [\mathrm{Cout}] complete the readout, \mathrm{Cout}\ge 1. All floating values share \mathrm{float32}/64 dtype/device and have representable products, squared lengths and results.
• For residue i/head h, rotate \mathrm{qd\_i}=R_{i}\cdot q_{\mathrm{direction\_i}} and \mathrm{kd\_j}=R_{j}\cdot k_{\mathrm{direction\_j}}; transform points \mathrm{qp\_i}=R_{i}\cdot q_{\mathrm{point\_i}}+t_{i} and \mathrm{kp\_j}=R_{j}\cdot k_{\mathrm{point\_j}}+t_{j}.
• Score a_{\mathrm{ijh}}=(\operatorname{softplus}(\mathrm{direction\_scale\_h})\cdot \operatorname{dot}(\mathrm{qd\_i},\mathrm{kd\_j})-\operatorname{softplus}(\mathrm{distance\_scale\_h})\cdot \Vert \mathrm{qp\_i}-\mathrm{kp\_j}\Vert \_2)/\sqrt{3}, where \operatorname{softplus}(x)=\operatorname{log}(1+\operatorname{exp}(x)) must be evaluated stably.
• The distance is not squared and has no added \epsilon. Define its derivative at a zero displacement to be the zero vector. frame_mask[B,N] is required bool True for a valid frame. An optional bool pair_mask broadcasts to [B,H,N,N]; True allows an edge. Allowed edges intersect query/key frame validity and pair_mask.
• Normalize scores over key j with stable masked softmax; all-masked head rows return exact zeros. Rotate each local value vector with R_j only (NO translation), aggregate by attention, then rotate the result back with R_i transpose. Return vectors[B,N,H,M,3], attention[B,H,N,N] and output[B,N,\mathrm{Cout}].
• Flatten local vectors in H,M,xyz order, multiply by output_weight, and add optional output_bias. Force output to zero when no head has an allowed key so bias cannot revive an empty query. There is no scalar residual or internal Q/K/V projection/LayerNorm in this kernel.
• Accept noncontiguous inputs, finite masked padding, arbitrary per-head exclusion masks, and never mutate inputs. Empty batch/residue/head/message axes and nonorthogonal input rotations are outside the forward domain.

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 geometric_reasoning_attention(
    activations: dict,
    parameters: dict,
    masks: dict,
    structure: dict,
) -> dict[str, torch.Tensor]:
    """Return dict[str, torch.Tensor]:
    
    output: (B, N, Cout).
    attention: (B, H, N, N).
    vectors: (B, N, H, M, 3)."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement geometric_reasoning_attention")

Sources and further reading

• www.science.org
• Biohub/esm · geom_attention.py
• Biohub/esm · vqvae.py