Genomics and RNA

Window Attention Pooling

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

Pool each consecutive sequence window by assigning learned weights to its valid positions. Use those weights to average the original feature vectors, then project the pooled vector; return the weights and window validity alongside it.

Function signature

window_attention_pooling(
    x: torch.Tensor,
    valid_mask: torch.Tensor,
    logit_weight: torch.Tensor,
    output_weight: torch.Tensor,
    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 is one consecutive width-2 sequence window; window grouping happens before this calculation.

Parameters

Parameter	Shape and meaning
`x`torch.Tensor	$[B,L,C]$ Sequence features to group into adjacent windows.
`valid_mask`torch.Tensor	$[B,L]$ True includes this sequence position.
`logit_weight`torch.Tensor	$[C, 1]$ or $[C, C]$ Bias-free scoring projection: one shared score or one score per feature channel.
`output_weight`torch.Tensor	$[C, D]$ Map pooled input channels to output channels; an identity matrix leaves pooled features unchanged.
`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

x is $[B,L,C]$ . Group consecutive positions into windows of width window (default 2, positive). remainder="drop" discards a final incomplete window; "pad" keeps it with invalid zero padding.
logit_weight is $[C,1]$ for shared attention weights or $[C,C]$ for independent weights per feature channel. Compute grouped x @ logit_weight and normalize stably over the candidate axis, excluding invalid entries.
The candidate weights multiply the ORIGINAL cleaned grouped feature vectors; sum over candidates and multiply by output_weight $[C,Cout]$ . There is no bias or residual.
All-masked groups return zero weights and zero output, with valid_mask=False. The returned mask reports whether each group had any included candidate.
Use CPU float32/float64 tensors with one shared floating dtype/device. Preserve strided layouts and do not mutate supplied tensors.
Use raw tensor operations, not a ready-made implementation of the requested primitive. Supplied parameters and random variates must not be reinitialized.

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 window_attention_pooling(
    x: torch.Tensor,  # (B,L,C)
    valid_mask: torch.Tensor,  # (B,L)
    logit_weight: torch.Tensor,  # (C, 1) or (C, C)
    output_weight: torch.Tensor,  # (C, D)
    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 window_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.