Molecular graphs, equivariance and docking

Geometric vector perceptron (GVP)

HardRaw PyTorchAbout 30 min2 public examples plus private tests

ESM-IF1GVP-GNNProSST structure encoder

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

Your task

Update scalar and 3D vector feature channels with a geometric vector perceptron. Use vector lengths to inform scalar features and scalar gates to scale vector outputs.

Function signature

gvp_scalar_vector(
    scalars: torch.Tensor,
    vectors: torch.Tensor,
    valid: torch.Tensor,
    parameters: dict,
    eps: float,
) -> tuple[torch.Tensor, torch.Tensor]

Worked example

Vector length controls the scalar gate

Use one batch with one valid position, scalar feature [0] and vector feature [1,0,0]. Set w_vector_hidden=[[2]], w_scalar=[[1],[1]], w_vector_out=[[3]], w_gate=[[1]], both biases [0], and eps=5. The hidden vector is [2,0,0], whose stabilized length is $\sqrt{4+5}=3$ . The scalar map adds this length to the zero scalar, then applies SiLU:

s=\frac{3}{1+e^{-3}}\approx2.857722,\qquad g=\frac{1}{1+e^{-s}}\approx0.945716.

The vector output is $g[6,0,0]\approx[5.674299,0,0]$ . Return scalars with shape (1,1,1) and vectors with shape (1,1,1,3). The gate uses the activated scalar s, not its preactivation 3.

Parameters

Parameter	Shape and meaning
`scalars`torch.Tensor	$[\ldots, S]$ Scalar feature channels; these do not rotate with the coordinate system.
`vectors`torch.Tensor	$[\ldots, V, 3]$ $V$ is channel, final axis is xyz.
`valid`torch.Tensor	$[\ldots]$ `True` means `valid` item.
`parameters`dict	Supplied model parameters, keyed as shown below.
`eps`float	Required positive offset inside $\sqrt{\sum v_{\mathrm{hidden}}^2+\varepsilon}$ . A zero hidden vector therefore contributes $\sqrt{\varepsilon}$ as its scalar norm feature.

Return

tuple[torch.Tensor, torch.Tensor]

Return these tensors in the stated order.

Required behavior

Inputs have common nonempty leading shape L: scalars [L,S], vectors [L,V,3], valid [L]. S,H,T are positive; V and O may be zero. All weight matrices use input-channel-first orientation. Sanitize invalid scalar/vector payloads before arithmetic.
Compute $U[...,h,a]=\sum_v vectors[...,v,a]\,w_{vector\_hidden}[v,h]$ , then $n_h=\sqrt{\sum_a U_{ha}^{2}+eps}$ . Concatenate the S scalar features with these H norms.
Apply the scalar affine map and SiLU: $s_{out}=\operatorname{SiLU}(\operatorname{concat}(scalars,n)@w_{scalar}+b_{scalar})$ , with $\operatorname{SiLU}(z)=z\,\operatorname{sigmoid}(z)$ . Use stable arithmetic for sigmoid.
Form $gate=\operatorname{sigmoid}(s_{out}@w_{gate}+b_{gate})$ . Return $v_{out}[...,o,a]=(\sum_h U[...,h,a]\,w_{vector\_out}[h,o])\,gate[...,o]$ . The gate uses activated scalar output.
eps is required and strictly positive, added inside the squared-norm root. A zero vector has norm sqrt(eps) but produces a zero vector output. V=0 gives zero hidden vectors; O=0 gives an empty output vector-channel axis.
Return scalars [L,T] and vectors [L,O,3], preserving dtype/device. Both outputs are zero where valid=False. Invalid payloads may be NaN/Inf. There is no vector bias, XYZ-axis weight, dropout, vector normalization or residual.

Why this matters

Vector norms provide rotation-invariant scalar information; scalar gates then scale whole output vectors without privileging a Cartesian direction.

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 gvp_scalar_vector(
    scalars: torch.Tensor,  # (..., S)
    vectors: torch.Tensor,  # (..., V, 3)
    valid: torch.Tensor,  # (...)
    parameters: dict,
    eps: float,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Return tuple[torch.Tensor, torch.Tensor]:
    
    scalars: (..., T).
    vectors: (..., O, 3)."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement gvp_scalar_vector")

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.

Sources and further reading

A primitive exercise, not an end-to-end reproduction of the named models. No affiliation or endorsement is implied.