Foundations

Multi-head attention with masks and additive bias

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

For each query position, score the allowed context positions separately in each head and take a probability-weighted sum of their value vectors. Combine those head summaries into the output features. attention[b,h,i,j] is the fraction assigned to context position j by query i in head h.

Function signature

multihead_attention(
    activations: dict,
    parameters: dict,
    masks: dict,
    heads: int = 1,
    causal: bool = False,
    causal_offset: int = 0,
) -> tuple[torch.Tensor, torch.Tensor]

Worked example

Two heads select different context values

Use one batch, two queries, three context positions, heads=2 and causal=False. The data below omit the sole batch axis from query, context and the masks. Omit all biases except b_o.

query   = [[1,1], [1,1]]
context = [[0,1,2,1], [1,0,4,4], [100,100,99,99]]
w_q = [[1,0], [0,1]]
w_k = [[1,0], [0,1], [0,0], [0,0]]
w_v = [[0,0], [0,0], [1,0], [0,1]]
w_o = [[1,2], [3,-1]]; b_o = [1,2]
key_mask   = [True,True,False]
query_mask = [True,False]
attention_mask (head, query, key) =
  [[[True,True,True],  [True,True,True]],
   [[False,True,True], [False,True,True]]]

Scores and probabilities. Each head has one key channel and one value channel, so sqrt(D)=1. For query 0, head 0 has included scores [0,1] and probabilities approximately [0.268941,0.731059,0]. Head 1 may use only context 1, so its probabilities are [0,1,0].
Head summaries. Head 0 returns (2+4e)/(1+e) ≈ 3.462117; head 1 returns 4. Concatenate them in head order: approximately [3.462117,4].
Output features. Multiply that vector by w_o and add b_o to obtain approximately [16.462117,4.924234].
Masked query. Query 1 has zero probabilities in both heads. Its final output is [0,0], including the suppression of b_o.

Parameters

Parameter	Shape and meaning
`activations`dict	Supplied intermediate tensors.
`parameters`dict	Supplied model parameters, keyed as shown below.
`masks`dict	Supplied masks. Each field below defines what True includes or permits.
`heads`intDefault `1`	Positive number of attention heads. Projected query/key widths divide into heads of width D; value heads may have a different width Dv.
`causal`boolDefault `False`	When True, a query at index i can use only keys j <= i + causal_offset, intersected with all supplied masks.
`causal_offset`intDefault `0`	Integer offset in the causal key bound j <= i + causal_offset; used only when causal is True.

Return

tuple[torch.Tensor, torch.Tensor]

Return these tensors in the stated order.

Required behavior

Inputs query $[B,Q,C_q]$ and context $[B,K,C_{\mathrm{ctx}}]$ have $B,Q,K\ge 1$ . $\mathrm{heads}=H$ (default $1$ ), with positive per-head key width $D$ and value width $D_v$ .
Supplied matrices are w_q $[C_q,H\cdot D]$ , w_k $[C_{\mathrm{ctx}},H\cdot D]$ , w_v $[C_{\mathrm{ctx}},H\cdot D_v]$ , w_o $[H\cdot D_v,C_{\mathrm{out}}]$ . Each optional b_q $[H\cdot D]$ , b_k $[H\cdot D]$ , b_v $[H\cdot D_v]$ , b_o $[C_{\mathrm{out}}]$ is added to its matching projection.
Split projected final channels in head-major order and form $\mathrm{scores}[b,h,i,j]=\operatorname{dot}(q[b,h,i],k[b,h,j])/\sqrt{D}+\mathrm{bias}[b,h,i,j]$ ; optional floating bias must broadcast to $[B,H,Q,K]$ .
Boolean True means allowed for every mask: optional key_mask $[B,K]$ , query_mask $[B,Q]$ , and attention_mask broadcastable to $[B,H,Q,K]$ are intersected.
For causal=True (default False), also require $j\le i+\mathrm{causal\_offset}$ (integer, default $0$ ), even for $Q\ne K$ . Normalize over keys with a stable masked softmax. An empty per-head row returns zero attention, not a uniform distribution or NaNs.
Concatenate each head's weighted value sum in head order, then apply w_o and optional b_o. If a query has no permitted key in ANY head, force the final output to zero after the output projection, so b_o cannot resurrect masked queries. $A$ query with at least one active head retains its output bias.
Return output features with shape $(B,Q,C_{out})$ and attention probabilities with shape $(B,H,Q,K)$ . Do not add a residual, dropout, normalization or rotary embeddings. $D_v$ may differ from $D$ ; $Q$ may differ from $K$ . Preserve dtype/device and accept noncontiguous tensors. Optional masks/bias are absent when unused.

Why this matters

Each head compares projected queries with projected context keys, then mixes the context values. Masks decide which comparisons exist; the output projection combines heads only after each head has formed its own mixture.

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 multihead_attention(
    activations: dict,
    parameters: dict,
    masks: dict,
    heads: int = 1,
    causal: bool = False,
    causal_offset: int = 0,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Return tuple[torch.Tensor, torch.Tensor]:
    
    output: (B, Q, Cout).
    attention: (B, H, Q, K)."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement multihead_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.