Genomics and RNA · INTEGRATION

Learned softmax pooling for sequence windows and pair rows

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

Reduce each window or query row to a learned weighted average of its original feature vectors. Scores determine how much each valid candidate contributes, and the supplied output matrix transforms the pooled vector. Return the resulting features, candidate weights and group-validity flags.

Function signature

softmax_attention_pooling(
    x: torch.Tensor,
    valid_mask: torch.Tensor,
    logit_weight: torch.Tensor,
    output_weight: torch.Tensor,
    mode: str = 'window',
    window: int = 2,
    remainder: str = 'drop',
) -> dict[str, torch.Tensor]

Worked example

Scores choose how much of each original vector to keep

For one group of two valid candidates, use feature vectors [0,2] and [log(3),6], shared logit_weight=[[1],[0]], and output_weight=[[0],[2]]. The scores are [0,log(3)], so candidate weights are [1/4,3/4]. The pooled original features are [(3/4)·log(3),5], since (1/4)·2+(3/4)·6=5. The final projection returns 10. This group can be one width-2 window or one query row; grouping happens before this calculation.

Parameters

Parameter	Shape and meaning
`x`torch.Tensor	$[B, L, C]$ or $[B, Q, K, C]$ Input features arranged as a sequence in window mode or as explicit query/key rows in rows mode.
`valid_mask`torch.Tensor	$[B, L]$ or $[B, K]$ or $[B, Q, K]$ True permits a position in its pooling group. Row masks select keys; they do not separately mask queries.
`logit_weight`torch.Tensor	$[C, 1]$ or $[C, C]$ Bias-free feature-to-score projection; one score channel is shared across features, or C scores weight channels separately.
`output_weight`torch.Tensor	$[C, D]$ Identity for unprojected Enformer pooling
`mode`strDefault `'window'`	Choose "window" for nonoverlapping sequence windows or "rows" to reduce the supplied within-row position axis.
`window`intDefault `2`	Number of positions per pooling window.
`remainder`strDefault `'drop'`	Window mode only: "drop" discards an incomplete final window; "pad" retains it with masked padding.

Return

dict[str, torch.Tensor]

Return exactly the named fields applicable to this operation.

Required behavior

Window mode (default): group x $[B,L,C]$ into nonoverlapping windows of size window (default $2$ ). Config remainder='drop' discards an incomplete tail; 'pad' creates a final masked window. The source Enformer setting has divisible $L$ .
Row mode (mode='rows'): x $[B,Q,K,C]$ , valid_mask $[B,K]$ or a bool tensor broadcastable to $[B,Q,K]$ , pools the $K$ axis without query masking.
For each group compute $\mathrm{logits}_{k}=x_{k}\,\,\,\mathrm{logit\_weight}[C,P]$ , where $P$ is $1$ (shared weights) or $C$ (channel-specific). Normalize logits over eligible positions only: $a=\exp(\mathrm{logit}-\max)/\mathrm{sum}\,\exp(\mathrm{logit}-\max)$ .
Return weights with trailing $[K,P]$ , $\mathrm{pooled}=\sum_{k}\,a_{k}\cdot x_{k}$ , $y=\mathrm{pooled}\,\,\,\mathrm{output\_weight}[C,D]$ , and a group-valid flag. Enformer uses $\mathrm{output\_weight}=\mathrm{identity}$ ; ESMFold2 uses $P=1$ and its learned output projection.
Missing payloads are sanitized before scoring; all-masked groups have exactly zero weights and output. Row $K$ must be positive; a dropped short window may produce zero groups. No learned biases or randomness are hidden.
Compute logits as $x\mathbin{@}logit\_weight$ , and the final output as $pooled\mathbin{@}output\_weight$ .

Why this matters

A mean treats every position equally; attention pooling learns a score for each position. The score can be shared across channels or different for every channel, while the final projection maps the pooled vector to the requested output width.

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 softmax_attention_pooling(
    x: torch.Tensor,  # (B, L, C) or (B, Q, K, C)
    valid_mask: torch.Tensor,  # (B, L) or (B, K) or (B, Q, K)
    logit_weight: torch.Tensor,  # (C, 1) or (C, C)
    output_weight: torch.Tensor,  # (C, D)
    mode: str = 'window',
    window: int = 2,
    remainder: str = 'drop',
) -> dict[str, torch.Tensor]:
    """Return dict[str, torch.Tensor]:
    
    y: (B, G, D).
    weights: (B, G, K, P).
    valid_mask: (B, G)."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement softmax_attention_pooling")

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.