MSA and pair representations

Row Attention Pooling

EasyRaw PyTorchAbout 10 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 query row separately by assigning learned weights to its valid candidate positions. Average the original feature vectors with those weights and apply the supplied output projection; return row outputs, weights and validity.

Function signature

row_attention_pooling(
    x: torch.Tensor,
    valid_mask: torch.Tensor,
    logit_weight: torch.Tensor,
    output_weight: torch.Tensor,
) -> 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 query row; other query rows are pooled independently.

Parameters

Parameter	Shape and meaning
`x`torch.Tensor	$[B,Q,K,C]$ Candidate feature vectors grouped into independent rows.
`valid_mask`torch.Tensor	$[B,K]$ or $[B,Q,K]$ True includes this candidate; row-shared masks broadcast over Q.
`logit_weight`torch.Tensor	$[C, 1]$ or $[C, C]$ Bias-free score projection; one shared score or one score per input feature channel.
`output_weight`torch.Tensor	$[C, D]$ Projection from the pooled C-channel vector to D output channels.

Return

dict[str, torch.Tensor]

Return exactly the named fields applicable to this operation.

Required behavior

x is $[B,Q,K,C]$ , where each row has K positive candidate slots. valid_mask is $[B,Q,K]$ or $[B,K]$ , with the latter broadcast over rows.
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,D]$ . 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.

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 row_attention_pooling(
    x: torch.Tensor,  # (B,Q,K,C)
    valid_mask: torch.Tensor,  # (B,K) or (B,Q,K)
    logit_weight: torch.Tensor,  # (C, 1) or (C, C)
    output_weight: torch.Tensor,  # (C, D)
) -> dict[str, torch.Tensor]:
    """Return dict[str, torch.Tensor]:
    
    y: (B,Q,D).
    weights: (B,Q,K,P).
    valid_mask: (B,Q)."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement row_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.