Genomics and RNA

Enformer genomic relative attention

HardRaw PyTorchAbout 30 min1 public example plus private tests

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

Your task

Let each genomic position read other positions using both content similarity and their relative sequence offset. Add a content score to the learned projection of the key-minus-query distance features, normalize over allowed keys, and mix their values. Return output features, attention, and the relative-score/basis diagnostics.

Function signature

genomic_relative_attention(
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    relative_weight: torch.Tensor,
    content_bias: torch.Tensor,
    position_bias: torch.Tensor,
    output_weight: torch.Tensor,
    output_bias: torch.Tensor,
    masks: dict,
    scale_query: bool = True,
    basis_per_family: int = 2,
    symmetric: bool = False,
) -> dict[str, torch.Tensor]

Worked example

Relative position changes which key is read

Consider the attention stage after the specified basis has been projected, with L=2, one head and D=V=1. Let projected relative keys for offsets [-1,0,+1] be [-log(3),0,log(3)]. Use queries [0,0], keys [0,0], content_bias=0, position_bias=1, values [2,6], all-valid masks, output weight 1 and output bias 0. For query 0, offsets [0,+1] yield logits [0,log(3)]; attention [1/4,3/4] returns (1/4)·2+(3/4)·6=5. Query 1 reads offsets [-1,0], giving logits [-log(3),0] and the same normalized weights. Relative logits change with query position even when content logits are zero.

Parameters

Parameter	Shape and meaning
`q`torch.Tensor	$[B, H, L, D]$ float[ $B,H,L,D$ ] projected queries
`k`torch.Tensor	$[B, H, L, D]$ float[ $B,H,L,D$ ] projected keys
`v`torch.Tensor	$[B, H, L, V]$ float[ $B,H,L,V$ ] projected values
`relative_weight`torch.Tensor	$[F, H, D]$ Relative-key projection
`content_bias`torch.Tensor	$[H, D]$ Added to scaled $q$ in the content term
`position_bias`torch.Tensor	$[H, D]$ Added to scaled $q$ in the relative term
`output_weight`torch.Tensor	$[H\cdot V, O]$ Concatenated-head projection
`output_bias`torch.Tensor	$[O]$ `float` $[O]$ output offset
`masks`dict	Supplied masks. Each field below defines what True includes or permits.
`scale_query`boolDefault `True`	Apply inverse-square-root query-width scaling when True.
`basis_per_family`intDefault `2`	Number of basis channels in each relative-position family.
`symmetric`boolDefault `False`	Use unsigned-only relative features when True.

Return

dict[str, torch.Tensor]

Return exactly the named fields applicable to this operation.

Required behavior

Given projected $q$ , $k$ [ $B,H,L,D$ ], $v$ [ $B,H,L,V$ ], and $L\ge 1$ , use offsets $r=-(L-1)$ .. $L-1$ , $n=\mathrm{basis\_per\_family}\ge 1$ (default $2$ ).
For $j=0$ .. $n-1$ let $f=j/\operatorname{max}(n-1,1)$ , $\mathrm{half\_life}=2^{3+f\cdot (\operatorname{log2}(L)-3)}$ , $e_{j}(r)=\operatorname{exp}(-\operatorname{log}(2)\cdot \left|r\right|/\mathrm{half\_life})$ , and $c_{j}(r)=1[\left|r\right|<2^{j+1}-1]$ .
Set $\operatorname{mean}=L/n+f\cdot (L-L/n)$ , $\mathrm{std}=L/(2n)$ , $\alpha=(\operatorname{mean}/\mathrm{std})^{2}$ , $\mathrm{rate}=\operatorname{mean}/\mathrm{std}^{2}$ .
Gamma density is $\operatorname{exp}((\alpha-1)\operatorname{log}(\left|r\right|)-\mathrm{rate}\cdot \left|r\right|-\operatorname{lgamma}(\alpha)+\alpha\cdot \operatorname{log}(\mathrm{rate}))$ , with $\mathrm{density}(0)=0$ for these $\alpha\ge 4$ .
Add $10^{-8}$ and divide each $\gamma$ feature by its maximum across the entire sampled offset grid. Concatenate $[e,c,g]$ ; unless symmetric=True, append $\operatorname{sign}(r)\cdot [e,c,g]$ . Thus $F=3n$ or $6 n$ . Project features with relative_weight[ $F,H,D$ ].
Scale q by $D^{-1/2}=1/\sqrt D$ unless scale_query=False. Add content_bias and position_bias after this scaling. The content term is (q_i+content_bias)·k_j and the relative term is (q_i+position_bias)·relative_key[j-i+L-1].
Apply masked softmax over $j$ , attend to $v$ , concatenate heads and apply output_weight plus output_bias. valid_mask[ $B,L$ ] excludes queries and keys; optional attention_mask[ $B,L,L$ ] further excludes pairs. Sanitize excluded $q/k/v$ payloads.
Any query with no allowed key returns zero weights and zero $y$ , including its output bias. Also return the fixed basis and relative logits before masking, computed from sanitized $q$ .

Why this matters

The effect of another sequence position depends on both its content and where it lies relative to the query. A learned projection of distance features supplies that positional term without assigning a separate learned vector to every absolute position.

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 genomic_relative_attention(
    q: torch.Tensor,  # (B, H, L, D)
    k: torch.Tensor,  # (B, H, L, D)
    v: torch.Tensor,  # (B, H, L, V)
    relative_weight: torch.Tensor,  # (F, H, D)
    content_bias: torch.Tensor,  # (H, D)
    position_bias: torch.Tensor,  # (H, D)
    output_weight: torch.Tensor,  # (H* V, O)
    output_bias: torch.Tensor,  # (O,)
    masks: dict,
    scale_query: bool = True,
    basis_per_family: int = 2,
    symmetric: bool = False,
) -> dict[str, torch.Tensor]:
    """Return dict[str, torch.Tensor]:
    
    y: (B, L, O).
    weights: (B, H, L, L).
    relative_logits: (B, H, L, L).
    features: (2L-1, F)."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement genomic_relative_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.

Sources and further reading

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